import collections
import contextlib
import dataclasses
import functools
import inspect
import os
import re
from itertools import chain, count
from typing import Any, Dict, List, Optional, Set, Tuple, Union
import sympy
from sympy import Expr
import torch
from torch._dynamo.utils import counters, dynamo_timed
from torch._inductor.codecache import get_cpp_wrapper_cubin_path_name
from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols, SymTypes
from torch.fx.node import _get_qualified_name
from torch.utils._sympy.singleton_int import SingletonInt
from .. import codecache, config, ir
from ..codecache import CudaKernelParamCache
from ..ir import ComputedBuffer, InputBuffer, ReinterpretView
from ..triton_heuristics import grid as default_grid
from ..utils import (
cache_on_self,
get_benchmark_name,
LineContext,
sympy_product,
sympy_str,
)
from ..virtualized import V
from .common import CodeGen, DeferredLine, IndentedBuffer, PythonPrinter
from .triton_utils import config_of, signature_to_meta
pexpr = PythonPrinter().doprint
def buffer_reuse_key(node: ir.Buffer):
return (
node.get_device(),
node.get_dtype(),
# NB: this is symbolic so that we don't try to reuse a buffer
# for s0 for s1, just because they happen to share the same
# size hint
sympy_str(V.graph.sizevars.simplify(node.layout.storage_size())),
)
def is_int(s: str):
# Cpp code gen adds L at the end of ints
# Lets remove it for checking whether we have an int or not
if s and s[-1] == "L":
s = s[:-1]
try:
int(s)
except ValueError:
return False
except TypeError:
return False
return True
def is_float(s: str):
try:
float(s)
except ValueError:
return False
return True
def convert_arg_type(python_type: str):
from .cpp import CONTAINER_PYTHON_TO_CPP, PYTHON_TO_CPP
if python_type == "Tensor":
# Conversions rules follow https://github.com/pytorch/pytorch/tree/main/aten/src/ATen/native#func
return f"at::{python_type} const&"
if python_type in PYTHON_TO_CPP:
return PYTHON_TO_CPP[python_type]
# Convert args of container types e.g. Optional[*]
for py_container, cpp_container in CONTAINER_PYTHON_TO_CPP.items():
container_match = re.findall(py_container + r"\[([a-zA-Z_]+)]", python_type)
if len(container_match) == 1:
contained_type = container_match[0]
assert (
contained_type in PYTHON_TO_CPP
), f"unsupported {py_container} type in convert_arg_type: {contained_type}"
cpp_contained_type = PYTHON_TO_CPP[contained_type]
return f"{cpp_container}<{cpp_contained_type}>"
raise AssertionError(f"unsupport python_type: {python_type}")
def convert_return_type(python_type: str):
# TODO: support alias
python_to_cpp = {
"Tensor": "at::Tensor",
"List[Tensor]": "std::vector<at::Tensor>",
}
cpp_type = python_to_cpp.get(python_type, None)
assert cpp_type is not None, f"NYI return type: {python_type}"
return cpp_type
def get_cpp_op_schema(kernel):
# use x.real_type instead of x.type so that we get ScalarType instead of int
arg_types = [repr(x.real_type) for x in kernel._schema.arguments]
arg_names = [x.name for x in kernel._schema.arguments]
returns = [repr(x.real_type) for x in kernel._schema.returns]
num_returns = len(returns)
assert num_returns > 0, "must have at least one return value"
if num_returns == 1:
cpp_return_value = convert_return_type(returns[0])
elif num_returns > 1:
tuple_returns = ", ".join([convert_return_type(r) for r in returns])
cpp_return_value = f"std::tuple<{tuple_returns}>"
cpp_arg_type = [
f"{convert_arg_type(arg_type)} {arg_name}"
for arg_type, arg_name in zip(arg_types, arg_names)
]
return f"{cpp_return_value}({', '.join(cpp_arg_type)})"
def user_defined_kernel_grid_fn_code(name, configs, grids):
output = IndentedBuffer()
fn_name = f"grid_wrapper_for_{name}"
output.writeline(f"def {fn_name}(meta):")
with output.indent():
if len(grids) == 1:
output.writeline(f"return {grids[0]}")
else:
assert len(grids) == len(configs)
for grid, c in zip(grids, configs):
guards = [f"meta['{name}'] == {val}" for name, val in c.kwargs.items()]
guards = " and ".join(guards)
output.writeline(f"if {guards}: return {grid}")
return fn_name, output.getvalue()
@dataclasses.dataclass
class SymbolicCallArg:
inner: Any
# the original symbolic expression represented by inner
inner_expr: sympy.Expr
def __str__(self):
return str(self.inner)
class MemoryPlanningState:
def __init__(self):
super().__init__()
self.reuse_pool: Dict[Any, List[FreeIfNotReusedLine]] = collections.defaultdict(
list
)
def __contains__(self, key):
return bool(self.reuse_pool.get(key, None))
def pop(self, key) -> "FreeIfNotReusedLine":
item = self.reuse_pool[key].pop()
assert not item.is_reused
return item
def push(self, key, item: "FreeIfNotReusedLine"):
assert not item.is_reused
self.reuse_pool[key].append(item)
@dataclasses.dataclass
class EnterCudaDeviceContextManagerLine:
device_idx: int
last_seen_device_guard_index: Optional[int]
def codegen(self, code: IndentedBuffer, device_cm_stack: contextlib.ExitStack):
if V.graph.cpp_wrapper:
code.writeline("\n")
if V.graph.aot_mode:
# In AOT mode, we have a stream provided as a param. A stream is
# associated with a device, so we never expect the device to change.
# CUDAStreamGuard sets the stream and the device.
if self.last_seen_device_guard_index is None:
if config.aot_inductor.abi_compatible:
code.writeline(
"AOTICudaStreamGuard stream_guard(stream, this->device_idx_);"
)
else:
code.writeline(
"at::cuda::CUDAStreamGuard stream_guard("
+ "at::cuda::getStreamFromExternal(stream, this->device_idx_));"
)
else:
assert (
self.last_seen_device_guard_index == self.device_idx
), "AOTInductor only supports running on one CUDA device"
else:
if self.last_seen_device_guard_index is None:
code.writeline(
f"at::cuda::CUDAGuard device_guard({self.device_idx});"
)
else:
code.writeline(f"device_guard.set_index({self.device_idx});")
else:
# Note _DeviceGuard has less overhead than device, but only accepts
# integers
code.writeline(f"with torch.cuda._DeviceGuard({self.device_idx}):")
device_cm_stack.enter_context(code.indent())
code.writeline(
f"torch.cuda.set_device({self.device_idx}) # no-op to ensure context"
)
class ExitCudaDeviceContextManagerLine:
def codegen(self, code: IndentedBuffer, device_cm_stack: contextlib.ExitStack):
if not V.graph.cpp_wrapper:
device_cm_stack.close()
@dataclasses.dataclass
class MemoryPlanningLine:
wrapper: "WrapperCodeGen"
def plan(self, state: MemoryPlanningState) -> "MemoryPlanningLine":
"""First pass to find reuse"""
return self
def codegen(self, code: IndentedBuffer):
"""Second pass to output code"""
pass
def __str__(self):
"""
Emits a string representation that fits on one line.
"""
args: List[str] = []
for field in dataclasses.fields(self):
if field.name == "wrapper":
continue
val = getattr(self, field.name)
args.append(
f"{field.name}={val.get_name() if field.type is ir.Buffer else val}"
)
return f"{type(self).__name__}({', '.join(args)})"
@dataclasses.dataclass
class AllocateLine(MemoryPlanningLine):
node: ir.Buffer
def plan(self, state: MemoryPlanningState):
if self.node.get_name() in V.graph.removed_buffers:
return NullLine(self.wrapper)
# try to reuse a recently freed buffer
key = buffer_reuse_key(self.node)
if config.allow_buffer_reuse and key in state:
free_line = state.pop(key)
free_line.is_reused = True
return ReuseLine(self.wrapper, free_line.node, self.node)
return self
def codegen(self, code: IndentedBuffer):
assert self.node.get_name() not in V.graph.removed_buffers
line = self.wrapper.make_buffer_allocation(self.node)
code.writeline(line)
@dataclasses.dataclass
class FreeIfNotReusedLine(MemoryPlanningLine):
node: ir.Buffer
is_reused: bool = False
def plan(self, state: MemoryPlanningState):
if isinstance(self.node.layout, (ir.AliasedLayout, ir.MultiOutputLayout)):
return self
assert not self.is_reused
if self.node.get_name() in V.graph.removed_buffers:
return NullLine(self.wrapper)
if config.allow_buffer_reuse:
state.push(buffer_reuse_key(self.node), self)
return self
def codegen(self, code: IndentedBuffer):
assert self.node.get_name() not in V.graph.removed_buffers
if not self.is_reused:
code.writeline(self.wrapper.make_buffer_free(self.node))
@dataclasses.dataclass
class ReuseLine(MemoryPlanningLine):
node: ir.Buffer
reused_as: ir.Buffer
delete_old: bool = True
def plan(self, state: MemoryPlanningState):
if self.node.get_name() in V.graph.removed_buffers:
assert self.reused_as.get_name() in V.graph.removed_buffers
return NullLine(self.wrapper)
assert self.reused_as.get_name() not in V.graph.removed_buffers
return self
def codegen(self, code: IndentedBuffer):
assert self.node.get_name() not in V.graph.removed_buffers
assert self.reused_as.get_name() not in V.graph.removed_buffers
code.writeline(
self.wrapper.make_buffer_reuse(self.node, self.reused_as, self.delete_old)
)
class NullLine(MemoryPlanningLine):
pass
class WrapperCodeGen(CodeGen):
"""
Generate outer wrapper in Python that calls the kernels.
"""
def __init__(self):
super().__init__()
self._names_iter = count()
self.header = IndentedBuffer()
self.prefix = IndentedBuffer()
self.wrapper_call = IndentedBuffer()
self.src_to_kernel = {}
self.kenel_numel_expr = set()
self.lines = []
self.declare = ""
self.ending = ""
self.open_bracket = "["
self.closed_bracket = "]"
self.comment = "#"
self.namespace = ""
self.none_str = "None"
self.size = "size()"
self.stride = "stride()"
self.last_seen_device_guard_index = None
self.supports_intermediate_hooks = True
self.expr_printer = pexpr
self.cached_thread_locals = set()
self.user_defined_kernel_cache: Dict[Tuple[Any, ...], str] = {}
self.unbacked_symbol_decls = set()
self.write_header()
self.write_prefix()
if not V.graph.aot_mode:
for name, hashed in V.graph.constant_reprs.items():
# include a hash so our code cache puts different constants into different files
self.write_constant(name, hashed)
self.allocated = set()
self.freed: Set[str] = set()
# maps from reusing buffer to reused buffer
self.reuses = dict()
self.write_get_raw_stream = functools.lru_cache(None)( # type: ignore[assignment]
self.write_get_raw_stream
)
@functools.lru_cache(None)
def add_import_once(line):
self.header.writeline(line)
self.add_import_once = add_import_once
self._metas = {}
def write_constant(self, name, hashed):
self.header.writeline(f"{name} = None # {hashed}")
def write_header(self):
self.header.splice(
f"""
from ctypes import c_void_p, c_long
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty, empty_strided
from {codecache.__name__} import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
aten = torch.ops.aten
inductor_ops = torch.ops.inductor
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
alloc_from_pool = torch.ops.inductor._alloc_from_pool
reinterpret_tensor = torch.ops.inductor._reinterpret_tensor
async_compile = AsyncCompile()
"""
)
@cache_on_self
def write_triton_header_once(self):
self.header.splice(
"""
import triton
import triton.language as tl
from torch._inductor.triton_heuristics import grid, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_cuda_stream
"""
)
def add_meta_once(self, meta):
meta = repr(meta)
if meta not in self._metas:
var = f"meta{len(self._metas)}"
self._metas[meta] = var
self.header.writeline(f"{var} = {meta}")
return self._metas[meta]
@cache_on_self
def get_output_refs(self):
return [x.codegen_reference(self.wrapper_call) for x in V.graph.graph_outputs]
def mark_output_type(self):
return
def codegen_input_size_asserts(self):
for name, buf in V.graph.graph_inputs.items():
if isinstance(buf, sympy.Expr):
continue
# comparing strides for 0 size tensor is tricky. Ignore them for now.
if sympy_product(buf.get_size()) == 0:
continue
size = self.codegen_shape_tuple(buf.get_size())
stride = self.codegen_shape_tuple(buf.get_stride())
self.prefix.writeline(f"assert_size_stride({name}, {size}, {stride})")
def write_prefix(self):
self.prefix.splice(
"""
async_compile.wait(globals())
del async_compile
def call(args):
"""
)
with self.prefix.indent():
if config.triton.debug_sync_graph:
self.prefix.writeline("torch.cuda.synchronize()")
inp_len = len(V.graph.graph_inputs.keys())
if inp_len != 0:
lhs = f"{', '.join(V.graph.graph_inputs.keys())}{'' if inp_len != 1 else ','}"
self.prefix.writeline(f"{lhs} = args")
self.prefix.writeline("args.clear()")
self.codegen_inputs(self.prefix, V.graph.graph_inputs)
if config.size_asserts:
self.codegen_input_size_asserts()
def write_get_raw_stream(self, index):
self.write_triton_header_once()
name = f"stream{index}"
self.writeline(f"{name} = get_cuda_stream({index})")
return name
def next_kernel_suffix(self):
return f"{next(self._names_iter)}"
def codegen_device_guard_enter(self, device_idx):
self.writeline(
EnterCudaDeviceContextManagerLine(
device_idx, self.last_seen_device_guard_index
)
)
self.last_seen_device_guard_index = device_idx
def codegen_device_guard_exit(self):
self.writeline(ExitCudaDeviceContextManagerLine())
def generate_return(self, output_refs):
if output_refs:
self.wrapper_call.writeline("return (" + ", ".join(output_refs) + ", )")
else:
self.wrapper_call.writeline("return ()")
def generate_end(self, result):
return
def generate_fallback_kernel(self, fallback_kernel, args):
self.generate_extern_kernel_alloc(fallback_kernel, args)
def generate_extern_kernel_alloc(self, extern_kernel, args):
ending = self.ending
if config.memory_planning and "view_as_complex" in str(extern_kernel.kernel):
# view operation fallbacks cause issues since inductor
# doesn't know the memory is still needed and might reuse it.
ending = f".clone(){ending}"
output_name = extern_kernel.get_name()
origin_node = extern_kernel.get_origin_node()
kernel_name = extern_kernel.codegen_kernel_name()
self.writeline(
f"{self.declare}{output_name} = {kernel_name}({', '.join(args)}){ending}"
)
if (
self.supports_intermediate_hooks
and config.generate_intermediate_hooks
and origin_node is not None
):
counters["inductor"]["intermediate_hooks"] += 1
self.writeline(
f"run_intermediate_hooks({origin_node.name!r}, {output_name})"
)
def generate_extern_kernel_out(self, output_view, codegen_reference, args, kernel):
if output_view:
args.append(f"out={output_view.codegen_reference()}")
else:
args.append(f"out={codegen_reference}")
self.writeline(f"{kernel}({', '.join(args)})")
def generate_user_defined_triton_kernel(self, kernel_name, grid, configs, args):
grid, code = user_defined_kernel_grid_fn_code(kernel_name, configs, grid)
# Must happen after free symbols are already codegened
with self.prefix.indent():
self.prefix.splice(code)
stream_name = self.write_get_raw_stream(V.graph.scheduler.current_device.index)
self.writeline(
f"{kernel_name}.run({', '.join(args)}, grid={grid}, stream={stream_name})"
)
def generate_scatter_fallback(
self, output, inputs, kernel, fn, src_is_tensor, reduce, kwargs
):
line = f"{kernel}({','.join(map(str, inputs))}"
if kernel == "aten.scatter_":
if reduce:
line += f", reduce={repr(reduce)}"
else:
line += ", ".join([""] + kwargs)
line += f"){self.ending}"
self.writeline(line)
def generate_extern_kernel_alloc_and_find_schema_if_needed(
self,
name,
kernel,
codegen_args,
cpp_op_schema,
cpp_kernel_key,
cpp_kernel_overload_name="",
op_overload=None,
raw_args=None,
outputs=None,
):
self.writeline(f"{name} = {kernel}({', '.join(codegen_args)})")
def generate_inf_and_nan_checker(self, node):
# TODO: Add check for python too.
pass
@dynamo_timed
def generate(self, is_inference):
if config.profile_bandwidth:
self.write_triton_header_once()
result = IndentedBuffer()
result.splice(self.header)
with contextlib.ExitStack() as stack:
stack.enter_context(self.wrapper_call.indent())
if config.profiler_mark_wrapper_call:
self.generate_profiler_mark_wrapper_call(stack)
if config.profile_bandwidth:
self.generate_start_graph()
# We disable planning during training because it presently increases peak memory consumption.
if is_inference and config.memory_planning:
self.memory_plan()
else:
self.memory_plan_reuse()
device_cm_stack = contextlib.ExitStack()
for line in self.lines:
if isinstance(line, MemoryPlanningLine):
line.codegen(self.wrapper_call)
elif isinstance(
line,
(
EnterCudaDeviceContextManagerLine,
ExitCudaDeviceContextManagerLine,
),
):
line.codegen(self.wrapper_call, device_cm_stack)
else:
self.wrapper_call.writeline(line)
output_refs = self.get_output_refs()
self.mark_output_type()
if config.triton.debug_sync_graph:
self.wrapper_call.writeline("torch.cuda.synchronize()")
if config.profile_bandwidth:
self.generate_end_graph()
self.generate_return(output_refs)
self.append_precomputed_sizes_to_prefix()
self.finalize_prefix()
result.splice(self.prefix)
with result.indent():
result.splice(self.wrapper_call)
self.generate_end(result)
self.add_benchmark_harness(result)
return result.getvaluewithlinemap()
def memory_plan(self):
from .memory_planning import MemoryPlanner
self.lines = MemoryPlanner(self).plan(self.lines)
def memory_plan_reuse(self):
out_names = V.graph.get_output_names()
while (
self.lines
and isinstance(self.lines[-1], MemoryPlanningLine)
# TODO: this seems legit, NullLine has no node
and self.lines[-1].node.name not in out_names # type: ignore[attr-defined]
):
# these lines will be pointless
self.lines.pop()
# codegen allocations in two passes
planning_state = MemoryPlanningState()
for i in range(len(self.lines)):
if isinstance(self.lines[i], MemoryPlanningLine):
self.lines[i] = self.lines[i].plan(planning_state)
def codegen_input_size_var_decl(self, code: IndentedBuffer, name):
code.writeline(f"{self.declare}{name}_size = {name}.{self.size}{self.ending}")
def codegen_input_stride_var_decl(self, code: IndentedBuffer, name):
code.writeline(
f"{self.declare}{name}_stride = {name}.{self.stride}{self.ending}"
)
def codegen_inputs(
self, code: IndentedBuffer, graph_inputs: Dict[str, ir.TensorBox]
):
"""Assign all symbolic shapes to locals"""
@functools.lru_cache(None)
def sizeof(name):
self.codegen_input_size_var_decl(code, name)
return f"{name}_size"
@functools.lru_cache(None)
def strideof(name):
self.codegen_input_stride_var_decl(code, name)
return f"{name}_stride"
# Assign all symbolic shapes needed to local variables
needed = V.graph.sizevars.free_symbols()
def is_expr(x):
return isinstance(x[1], sympy.Expr)
graph_inputs_expr = list(filter(is_expr, graph_inputs.items()))
graph_inputs_tensors = list(
filter(lambda x: not is_expr(x), graph_inputs.items())
)
for name, shape in graph_inputs_expr:
shape = V.graph.sizevars.simplify(shape)
if shape in needed:
needed.remove(shape)
code.writeline(f"{self.declare}{shape} = {name}{self.ending}")
for name, value in graph_inputs_tensors:
shapes = value.get_size()
for dim, shape in enumerate(shapes):
shape = V.graph.sizevars.simplify(shape)
if shape in needed:
needed.remove(shape)
code.writeline(
f"{self.declare}{shape} = {sizeof(name)}[{dim}]{self.ending}"
)
for name, value in graph_inputs_tensors:
shapes = value.get_stride()
for dim, shape in enumerate(shapes):
shape = V.graph.sizevars.simplify(shape)
if shape in needed:
needed.remove(shape)
code.writeline(
f"{self.declare}{shape} = {strideof(name)}[{dim}]{self.ending}"
)
def append_precomputed_sizes_to_prefix(self):
with self.prefix.indent():
for sym, expr in V.graph.sizevars.inv_precomputed_replacements.items():
self.prefix.writeline(
f"{self.declare}{sym} = {self.expr_printer(expr)}{self.ending}"
)
def finalize_prefix(self):
pass
def codegen_python_sizevar(self, x: Expr) -> str:
return pexpr(V.graph.sizevars.simplify(x))
def codegen_sizevar(self, x: Expr) -> str:
return self.codegen_python_sizevar(x)
def codegen_tuple_access(self, basename: str, name: str, index: str) -> str:
return f"{basename}[{index}]"
def codegen_python_shape_tuple(self, shape: Tuple[Expr, ...]) -> str:
parts = list(map(self.codegen_python_sizevar, shape))
if len(parts) == 0:
return "()"
if len(parts) == 1:
return f"({parts[0]}, )"
return f"({', '.join(parts)})"
def codegen_shape_tuple(self, shape: Tuple[Expr, ...]) -> str:
return self.codegen_python_shape_tuple(shape)
def codegen_alloc_from_pool(self, name, offset, dtype, shape, stride) -> str:
return "alloc_from_pool({})".format(
", ".join(
[
name,
pexpr(offset), # bytes not numel
str(dtype),
self.codegen_shape_tuple(shape),
self.codegen_shape_tuple(stride),
]
)
)
def codegen_reinterpret_view(self, data, size, stride, offset, writer) -> str:
size = self.codegen_shape_tuple(size)
stride = self.codegen_shape_tuple(stride)
offset = self.codegen_sizevar(offset)
return f"reinterpret_tensor({data.get_name()}, {size}, {stride}, {offset})"
def codegen_device_copy(self, src, dst):
self.writeline(f"{dst}.copy_({src})")
def codegen_multi_output(self, name, value):
self.writeline(f"{self.declare}{name} = {value}{self.ending}")
def benchmark_compiled_module(self, output):
def add_fake_input(name, shape, stride, device, dtype):
output.writeline(
f"{name} = rand_strided("
f"{self.codegen_python_shape_tuple(shape)}, "
f"{self.codegen_python_shape_tuple(stride)}, "
f"device='{device}', dtype={dtype})"
)
def add_expr_input(name, val):
output.writeline(f"{name} = {val}")
output.writelines(
["", "", "def benchmark_compiled_module(times=10, repeat=10):"]
)
with output.indent():
output.splice(
"""
from torch._dynamo.testing import rand_strided
from torch._inductor.utils import print_performance
""",
strip=True,
)
for name, value in V.graph.constants.items():
# all the constants are global variables, that's why we need
# these 'global var_name' lines
output.writeline(f"global {name}")
add_fake_input(
name, value.size(), value.stride(), value.device, value.dtype
)
for name, value in V.graph.graph_inputs.items():
if isinstance(value, sympy.Symbol) and isinstance(
V.graph.sizevars.var_to_val.get(value, None), SingletonInt
):
# Inductor should only work with dense -> dense graph, and
# SingletonInts belong to metadata that should only live on
# the subclass.
continue
if isinstance(value, sympy.Expr): # Don't need to add symbolic
add_expr_input(name, V.graph.sizevars.size_hint(value))
else:
shape = [V.graph.sizevars.size_hint(x) for x in value.get_size()]
stride = [V.graph.sizevars.size_hint(x) for x in value.get_stride()]
add_fake_input(
name, shape, stride, value.get_device(), value.get_dtype()
)
call_str = f"call([{', '.join(V.graph.graph_inputs.keys())}])"
output.writeline(f"fn = lambda: {call_str}")
output.writeline("return print_performance(fn, times=times, repeat=repeat)")
def add_benchmark_harness(self, output):
"""
Append a benchmark harness to generated code for debugging
"""
if not config.benchmark_harness:
return
self.benchmark_compiled_module(output)
output.writelines(["", "", 'if __name__ == "__main__":'])
with output.indent():
output.writelines(
[
"from torch._inductor.wrapper_benchmark import compiled_module_main",
f"compiled_module_main('{get_benchmark_name()}', benchmark_compiled_module)",
]
)
def define_kernel(
self, name: str, kernel: str, metadata: Optional[str] = None, cuda=True
):
metadata_comment = f"{metadata}\n" if metadata else ""
self.header.splice(f"\n\n{metadata_comment}{name} = {kernel}")
def define_user_defined_triton_kernel(self, kernel, configs, kwargs):
original_name = kernel.__name__
# Distinguish between different functions using function id
cache_key = [id(kernel.fn)]
for arg in kwargs.values():
if isinstance(arg, (ir.Buffer, ir.ReinterpretView)):
cache_key.append(arg.get_dtype())
elif len(configs) > 0:
# We need to key on non tensor arg only in autotune mode
cache_key.append(arg)
cache_key = tuple(cache_key)
if cache_key in self.user_defined_kernel_cache:
return self.user_defined_kernel_cache[cache_key]
name = f"{original_name}_{len(self.user_defined_kernel_cache)}"
# Add to the cache for the next use
self.user_defined_kernel_cache[cache_key] = name
compile_wrapper = IndentedBuffer()
compile_wrapper.writeline(f"async_compile.triton({original_name!r}, '''")
compile_wrapper.splice(
"""
import triton
import triton.language as tl
from torch._inductor.utils import instance_descriptor
from torch._inductor.triton_heuristics import user_autotune
""",
strip=True,
)
compile_wrapper.newline()
from .common import SizeArg, TensorArg
signature: List[Union[TensorArg, SizeArg]] = []
constants = {}
for key, arg in kwargs.items():
idx = kernel.arg_names.index(key)
if idx in kernel.constexprs:
constants[key] = arg
continue
if isinstance(arg, (ir.Buffer, ir.ReinterpretView)):
signature.append(
TensorArg(
key,
arg.codegen_reference(),
arg.get_dtype(),
# For ReinterpretView, we do not want to check alignment
not isinstance(arg, ReinterpretView),
)
)
else:
signature.append(SizeArg(key, arg))
index_dtype = "tl.int32"
inductor_meta = {
"kernel_name": name,
}
triton_meta = {
"signature": signature_to_meta(signature, size_dtype=index_dtype),
"device": V.graph.scheduler.current_device.index,
"device_type": V.graph.scheduler.current_device.type,
"constants": constants,
"configs": [config_of(signature)],
}
configs = [
{
"kwargs": config.kwargs,
"num_warps": config.num_warps,
"num_stages": config.num_stages,
}
for config in configs
]
compile_wrapper.splice(
f"""
@user_autotune(
configs={configs!r},
inductor_meta={inductor_meta!r},
triton_meta={triton_meta!r},
filename=__file__
)
@triton.jit
"""
)
compile_wrapper.splice(kernel.src, strip=True)
# Also include any possible kernel being called indirectly
from triton import JITFunction
symbols_included = {original_name}
def traverse(cur_kernel):
for symbol_name in cur_kernel.fn.__code__.co_names:
if symbol_name in symbols_included:
continue
if symbol_name in cur_kernel.fn.__globals__:
symbol = cur_kernel.fn.__globals__[symbol_name]
if isinstance(symbol, JITFunction):
compile_wrapper.newline()
compile_wrapper.writeline("@triton.jit")
compile_wrapper.splice(symbol.src, strip=True)
symbols_included.add(symbol_name)
traverse(symbol)
elif isinstance(symbol, (int, str, bool)):
compile_wrapper.newline()
compile_wrapper.writeline(f"{symbol_name} = {symbol!r}")
symbols_included.add(symbol_name)
traverse(kernel)
compile_wrapper.writeline("''')")
_, lineno = inspect.getsourcelines(kernel.fn)
srcfile = inspect.getsourcefile(kernel.fn)
metadata = f"# Original path: {srcfile}:{lineno}"
self.define_kernel(
name,
compile_wrapper.getvalue(),
metadata,
)
return name
def generate_numel_expr(self, kernel_name: str, tree):
expr = f"{kernel_name}_{tree.prefix}numel"
if expr not in self.kenel_numel_expr:
self.kenel_numel_expr.add(expr)
self.writeline(
f"{self.declare}{expr} = {self.expr_printer(tree.numel)}{self.ending}"
)
else:
self.writeline(f"{expr} = {self.expr_printer(tree.numel)}{self.ending}")
# We can get symbolic expressions here, like s0*64
# It is fine to have them here, but we need to handle them correctly as their own type
# This is tricky to do, so we wrap in a custom type, distinct from scalars, but also from sympy*
# scalars as well.
# This is handled in `generate_args_decl` which has a correct comment of: TODO: only works for
# constant now, need type info. I agree, this needs type info, and while this is not true type info
# it suffices as a type hint for the purposes of producing the correct code for this type.
return SymbolicCallArg(expr, tree.numel)
def wrap_kernel_call(self, name, call_args):
return f"{name}({', '.join(call_args)}){self.ending}"
def generate_profiler_mark_wrapper_call(self, stack):
self.wrapper_call.writeline("from torch.profiler import record_function")
self.wrapper_call.writeline(
f"with record_function('graph_{V.graph.graph_id}_inductor_wrapper_call'):"
)
stack.enter_context(self.wrapper_call.indent())
def generate_start_graph(self):
self.wrapper_call.writeline("start_graph()")
def generate_end_graph(self):
self.wrapper_call.writeline("end_graph()")
def generate_default_grid(self, name: str, grid_args: List[Any]):
return grid_args
def generate_kernel_call(
self,
name,
call_args,
grid=None,
device_index=None,
cuda=True,
triton=True,
):
"""
Generates kernel call code.
cuda: Defines whether the backend is GPU. Otherwise the backend is CPU.
triton: Defines whether the GPU backend uses Triton for codegen.
Otherwise it uses the CUDA language for codegen.
Only valid when cuda == True.
"""
if cuda:
call_args_str = ", ".join(pexpr(item) for item in call_args)
stream_name = self.write_get_raw_stream(
V.graph.scheduler.current_device.index
)
if triton:
grid_str = ", ".join(pexpr(item) for item in grid)
self.writeline(
f"{name}.run({call_args_str}, grid=grid({grid_str}), stream={stream_name})"
)
else:
stream_ptr = f"c_void_p({stream_name})"
self.writeline(f"{name}.{name}({call_args_str}, {stream_ptr})")
else:
self.writeline(self.wrap_kernel_call(name, call_args))
def writeline(self, line):
self.lines.append(line)
def enter_context(self, ctx):
self.lines.append(LineContext(ctx))
def val_to_cpp_arg_str(self, type_, val, is_legacy_abi) -> str:
raise NotImplementedError()
def val_to_arg_str(self, s):
if isinstance(s, SymTypes):
return pexpr(sympy.expand(repr(s)))
elif isinstance(s, sympy.Expr):
return pexpr(s)
elif isinstance(s, (tuple, list)):
@dataclasses.dataclass
class Shim:
ref: Any
def __repr__(self):
return self.ref
return repr(type(s)(Shim(self.val_to_arg_str(a)) for a in s))
elif isinstance(s, torch._ops.OpOverload):
return _get_qualified_name(s)
elif isinstance(s, (ComputedBuffer, InputBuffer, ReinterpretView)):
return s.codegen_reference()
else:
return repr(s)
# The following methods are for memory management
def make_buffer_allocation(self, buffer):
device = buffer.get_device()
dtype = buffer.get_dtype()
shape = tuple(buffer.get_size())
stride = tuple(buffer.get_stride())
return self.make_allocation(buffer.get_name(), device, dtype, shape, stride)
def make_allocation(self, name, device, dtype, shape, stride):
try:
expected = tuple(ir.make_contiguous_strides_for(shape))
except Exception: # cannot determine truth value of Relational
expected = None
if stride == expected:
return (
f"{name} = empty("
f"{self.codegen_shape_tuple(shape)}, "
f"device='{device.type}', dtype={dtype})"
)
else:
return (
f"{name} = empty_strided("
f"{self.codegen_shape_tuple(shape)}, "
f"{self.codegen_shape_tuple(stride)}, "
f"device='{device.type}', dtype={dtype})"
)
def make_tensor_alias(self, new_name, old_name, comment=""):
return f"{self.declare}{new_name} = {old_name}{self.ending} {self.comment} {comment}"
def make_buffer_free(self, buffer):
return f"del {buffer.get_name()}"
def make_free_by_names(self, names_to_del: List[str]):
return f"del {', '.join(name for name in names_to_del)}"
def codegen_exact_buffer_reuse(self, old_name: str, new_name: str, del_line: str):
return f"{self.declare}{new_name} = {old_name}{del_line}{self.ending} {self.comment} reuse"
def make_buffer_reuse(self, old, new, delete_old: bool):
assert old.get_dtype() == new.get_dtype()
old_name = old.get_name()
new_name = new.get_name()
del_line = ";"
if old_name not in V.graph.get_output_names() and delete_old:
del_line = f"; {self.make_buffer_free(old)}"
if old.get_size() == new.get_size() and old.get_stride() == new.get_stride():
if old_name in self.cached_thread_locals:
self.cached_thread_locals.add(new_name)
return self.codegen_exact_buffer_reuse(old_name, new_name, del_line)
reinterpret_view = self.codegen_reinterpret_view(
old, new.get_size(), new.get_stride(), 0, self.wrapper_call
)
if reinterpret_view in self.cached_thread_locals:
self.cached_thread_locals.add(new_name)
return f"{self.declare}{new_name} = {reinterpret_view}{del_line} {self.comment} reuse"
def codegen_deferred_allocation(self, name, layout):
self.writeline(
DeferredLine(
name,
f"{self.declare}{name} = {layout.view.codegen_reference()}{self.ending} {self.comment} alias",
)
)
def codegen_allocation(self, buffer):
assert (
buffer.get_workspace_size() == 0
), "Only support zero workspace size for now!"
name = buffer.get_name()
if name in V.graph.removed_buffers or name in self.allocated:
return
self.allocated.add(name)
if isinstance(
buffer,
(ir.ExternKernelAlloc, ir.MultiOutput),
):
return
layout = buffer.get_layout()
if isinstance(layout, ir.MutationLayout):
return
if isinstance(layout, ir.AliasedLayout):
assert isinstance(
layout.view, ir.ReinterpretView
), f"unexpected {type(layout.view)}: {layout.view}"
self.codegen_allocation(layout.view.data)
self.codegen_deferred_allocation(name, layout)
return
self.writeline(AllocateLine(self, buffer))
def codegen_free(self, buffer):
assert (
buffer.get_workspace_size() == 0
), "Only support zero workspace size for now!"
name = buffer.get_name()
# can be freed but not reused
if isinstance(buffer, ir.InputBuffer):
self.writeline(self.make_buffer_free(buffer))
return
if not self.can_reuse(buffer):
return
self.freed.add(name)
self.writeline(FreeIfNotReusedLine(self, buffer))
def can_reuse(self, input_buffer, output_buffer=None):
name = input_buffer.get_name()
if (
name in V.graph.removed_buffers
or name in V.graph.graph_inputs
or name in V.graph.constants
or name in V.graph.never_reuse_buffers
or name in self.freed
):
return False
return True
def did_reuse(self, buffer, reused_buffer):
# Check whether a given buffer was reused by a possible reuser in the wrapper codegen
# Can be consulted from inside ir codegen, e.g. to determine whether a copy is needed
return (
buffer.get_name() in self.reuses
and self.reuses[buffer.get_name()] == reused_buffer.get_name()
)
def codegen_inplace_reuse(self, input_buffer, output_buffer):
assert buffer_reuse_key(input_buffer) == buffer_reuse_key(output_buffer)
self.codegen_allocation(input_buffer)
self.freed.add(input_buffer.get_name())
self.allocated.add(output_buffer.get_name())
self.reuses[output_buffer.get_name()] = input_buffer.get_name()
self.writeline(ReuseLine(self, input_buffer, output_buffer))
def codegen_unbacked_symbol_decl(self, symbol):
name = str(symbol)
if name in self.unbacked_symbol_decls:
return name
else:
# When in CppWrapperCodeGen, we should only generate the declaration once
self.unbacked_symbol_decls.add(name)
return self.declare + name
class CppWrapperCodeGen(WrapperCodeGen):
"""
Generates cpp wrapper for running on CPU and calls cpp kernels
"""
def __init__(self):
super().__init__()
self.declare = "auto "
self.ending = ";"
self.open_bracket = "{"
self.closed_bracket = "}"
self.comment = "//"
self.namespace = "at::"
self.none_str = "at::Tensor()"
self.extern_call_ops = set()
self.size = "sizes()"
self.stride = "strides()"
self.call_func_name = "inductor_entry_cpp"
self.cuda = False
self.supports_intermediate_hooks = False
self.outputs_need_copy = set()
self.kernel_callsite_id = count()
self.int_array_id = count() # for int array local variable declarations
self.declared_int_array_vars = set()
self.tmp_tensor_id = count() # for tmp tensor local variable declarations
self.arg_var_id = count()
self.used_cached_dtypes = set()
from .cpp import cexpr, CppPrinter
self.expr_printer = cexpr
# CppPrinter sometimes calls at::native functions which causes problems in
# the ABI-compatible mode. Currently we are hitting this problem when codegen
# Grid computation expressions, but we my need to fix other size computation
# as well.
class GridExprCppPrinter(CppPrinter):
def _print_FloorDiv(self, expr):
x, div = expr.args
x = self.paren(self.doprint(x))
div = self.paren(self.doprint(div))
assert expr.is_integer, "Expect integers in GridExprPrinter"
return f"({x}/{div})"
self.grid_expr_printer = GridExprCppPrinter().doprint
def generate_kernel_call(
self,
name,
call_args,
grid=None,
device_index=None,
cuda=True,
triton=True,
):
"""
Generates kernel call code.
cuda: Defines whether the backend is GPU. Otherwise the backend is CPU.
triton: Defines whether the GPU backend uses Triton for codegen.
Otherwise it uses the CUDA language for codegen.
Only valid when cuda == True.
"""
if cuda:
return super().generate_kernel_call(
name, call_args, grid, device_index, cuda, triton
)
else:
if V.graph.aot_mode and config.aot_inductor.abi_compatible:
from .cpp import DTYPE_TO_CPP
new_args = []
for arg in call_args:
var_name = f"var_{next(self.arg_var_id)}"
self.writeline(f"void *{var_name}{self.ending}")
self.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr({arg}, &{var_name}));"
)
dtype = V.graph.get_dtype(arg)
cpp_dtype = DTYPE_TO_CPP[dtype]
new_args.append(f"({cpp_dtype}*)({var_name})")
self.writeline(self.wrap_kernel_call(name, new_args))
else:
self.writeline(self.wrap_kernel_call(name, call_args))
def write_constant(self, name, hashed):
# include a hash so our code cache gives different constants different files
self.header.writeline(f"// {name} {hashed}")
def write_header(self):
if V.graph.aot_mode:
with open(
os.path.join(os.path.dirname(__file__), "aoti_runtime", "interface.cpp")
) as f:
self.header.splice(f.read())
else:
self.header.splice(
"""
import torch
from torch._inductor.codecache import CppWrapperCodeCache
cpp_wrapper_src = (
'''
"""
)
if config.aot_inductor.abi_compatible:
self.header.splice("#include <torch/csrc/inductor/aoti_torch/c/shim.h>")
else:
self.header.splice(
"""
#include <ATen/ATen.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <ATen/native/BinaryOps.h>
#include <torch/csrc/inductor/aoti_torch/tensor_converter.h>
#include <torch/csrc/inductor/inductor_ops.h>
#define reinterpret_tensor torch::inductor::_reinterpret_tensor
#define alloc_from_pool torch::inductor::_alloc_from_pool
"""
)
self.header.splice("#include <c10/util/generic_math.h>")
from .memory_planning import ALIGN_BYTES
# Round up to the nearest multiple of ALIGN_BYTES
# ALIGN_BYTES must be a power of 2
self.header.splice(
f"""
[[maybe_unused]] static int64_t align(int64_t nbytes) {{
return (nbytes + {ALIGN_BYTES} - 1) & -{ALIGN_BYTES};
}}
"""
)
def mark_output_type(self):
# mark output type to unwrap tensor back to python scalar
from ..ir import ShapeAsConstantBuffer
output_is_tensor = dict()
for idx, x in enumerate(V.graph.graph_outputs):
if isinstance(x, ShapeAsConstantBuffer):
output_is_tensor[idx] = False
else:
output_is_tensor[idx] = True
self.output_is_tensor = output_is_tensor
def write_prefix(self):
if V.graph.aot_mode:
self.prefix.writeline("namespace torch {")
self.prefix.writeline("namespace aot_inductor {")
def write_input_output_info(
self,
info_kind: str,
idx: int,
name: str,
):
self.prefix.writeline(f"""{info_kind}[{idx}].name = "{name}";""")
def write_wrapper_decl(self):
inputs_len = len(V.graph.graph_inputs.keys())
if V.graph.aot_mode:
self.prefix.splice(
"""
void AOTInductorModel::run_impl(
AtenTensorHandle*
input_handles, // array of input AtenTensorHandle; handles
// are stolen; the array itself is borrowed
AtenTensorHandle*
output_handles, // array for writing output AtenTensorHandle; handles
// will be stolen by the caller; the array itself is
// borrowed
DeviceStreamType stream,
AOTIProxyExecutorHandle proxy_executor
) {
"""
)
else:
self.prefix.splice(
f"""std::vector<at::Tensor> {self.call_func_name}(const std::vector<at::Tensor>& inputs) {{"""
)
with self.prefix.indent():
# assign inputs and outputs in both cases so the later codegen can be simplified
if V.graph.aot_mode:
if config.aot_inductor.abi_compatible:
self.prefix.splice(
"""
auto inputs = steal_from_raw_handles_to_raii_handles(input_handles, num_inputs());
"""
)
else:
# This looks dumb, but can avoid creating two versions of code in the AOTInductor runtime.
self.prefix.splice(
"""
auto inputs = alloc_tensors_by_stealing_from_handles(input_handles, num_inputs());
"""
)
else:
self.prefix.splice(
"""
py::gil_scoped_release release;
"""
)
if inputs_len != 0:
for idx, input_key in enumerate(V.graph.graph_inputs.keys()):
# unwrap input tensor back to scalar
if isinstance(V.graph.graph_inputs[input_key], sympy.Expr):
from ..graph import may_get_constant_buffer_dtype
from .cpp import DTYPE_TO_CPP
dtype = may_get_constant_buffer_dtype(
V.graph.graph_inputs[input_key]
)
assert (
dtype is not None
), "Fails to get the dtype of the sympy.Expr"
cpp_dtype = DTYPE_TO_CPP[dtype]
assert (
not config.aot_inductor.abi_compatible
), "Need to add .item support for abi_compatible AOTInductor codegen"
self.prefix.writeline(
f"{cpp_dtype} {input_key} = inputs[{idx}].item<{cpp_dtype}>();"
)
else:
self.prefix.writeline(
f"auto {input_key} = std::move(inputs[{idx}]);"
)
assert all(
isinstance(v, torch.Tensor) for v in list(V.graph.constants.values())
), "Expect all constants to be Tensor"
for idx, constants_key in enumerate(V.graph.constants.keys()):
if V.graph.aot_mode:
# Weights are stored in constants_ and owned by RAIIAtenTensorHandle there.
# Don't call std::move here because it will cause constants_ to lose the ownership.
if config.aot_inductor.abi_compatible:
self.prefix.writeline(
f"""auto {constants_key} = constants_.at({idx});"""
)
else:
self.prefix.writeline(
f"auto {constants_key} = *tensor_handle_to_tensor_pointer("
+ f"""constants_.at({idx}));"""
)
else:
# Append constants as inputs to the graph
constants_idx = inputs_len + idx
self.prefix.writeline(
f"auto {constants_key} = inputs[{constants_idx}];"
)
self.codegen_inputs(self.prefix, V.graph.graph_inputs)
if V.graph.aot_mode:
self.prefix.writeline("inputs.clear();")
self.prefix.writeline(
"auto& kernels = *dynamic_cast<AOTInductorModelKernels*>(this->kernels_.get());"
)
def codegen_input_size_var_decl(self, code: IndentedBuffer, name):
if config.aot_inductor.abi_compatible:
code.writeline(f"int64_t* {name}_size;")
code.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_sizes({name}, &{name}_size));"
)
else:
super().codegen_input_size_var_decl(code, name)
def codegen_input_stride_var_decl(self, code: IndentedBuffer, name):
if config.aot_inductor.abi_compatible:
code.writeline(f"int64_t* {name}_stride;")
code.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_strides({name}, &{name}_stride));"
)
else:
super().codegen_input_stride_var_decl(code, name)
def codegen_model_kernels(self):
self.prefix.writeline("namespace {")
self.prefix.writeline(
"class AOTInductorModelKernels : public AOTInductorModelKernelsBase {"
)
self.prefix.writeline(" public:")
for kernel in chain(
self.src_to_kernel.values(), self.user_defined_kernel_cache.values()
):
self.prefix.writeline(f" CUfunction {kernel}{{nullptr}};")
self.prefix.writeline("};")
self.prefix.writeline("} // namespace")
def codegen_model_constructor(self):
"""
// Generated code example
AOTInductorModel::AOTInductorModel()
: AOTInductorModelBase(4, 1) {
inputs_info_[0].name = "input0";
inputs_info_[0].dtype = "torch.float16";
...
constants_info_[0].name = "L__self___weight";
constants_info_[0].dtype = at::kFloat;
constants_info_[0].offset = 0;
constants_info_[0].data_size = 8192;
constants_info_[0].shape = {64, 32};
constants_info_[0].stride = {32, 1};
...
outputs_info_[0].name = "output0";
outputs_info_[0].dtype = "torch.float16";
}
"""
num_inputs = len(V.graph.graph_inputs)
num_outputs = len(V.graph.graph_outputs)
num_constants = len(V.graph.constants)
self.prefix.splice(
f"""
AOTInductorModel::AOTInductorModel(std::shared_ptr<ConstantMap> constants_map, std::optional<std::string> cubin_dir)
: AOTInductorModelBase({num_inputs}, {num_outputs}, {num_constants}, cubin_dir) {{
"""
)
with self.prefix.indent():
for idx, (name, inp) in enumerate(V.graph.graph_inputs.items()):
assert not isinstance(
inp, sympy.Expr
), f"input {name=} cannot be symbolic"
self.write_input_output_info("inputs_info_", idx, name)
for idx, (name, tensor) in enumerate(V.graph.constants.items()):
assert isinstance(tensor, torch.Tensor)
self.prefix.writeline(f"""constants_info_[{idx}].name = "{name}";""")
self.prefix.writeline(
f"constants_info_[{idx}].dtype = static_cast<int32_t>({self.codegen_dtype(tensor.dtype)});"
)
self.prefix.writeline(
f"constants_info_[{idx}].offset = {tensor.storage_offset()};"
)
self.prefix.writeline(
f"constants_info_[{idx}].data_size = {tensor.untyped_storage().nbytes()};"
)
size_str = ", ".join([str(s) for s in tensor.size()])
self.prefix.writeline(f"constants_info_[{idx}].shape = {{{size_str}}};")
stride_str = ", ".join([str(s) for s in tensor.stride()])
self.prefix.writeline(
f"constants_info_[{idx}].stride = {{{stride_str}}};"
)
self.prefix.writeline("update_constants_map(std::move(constants_map));")
def escape_string(x):
return (
x.replace("\\", "\\\\")
.replace('"', '\\"')
.replace("\n", "\\n")
.replace("\t", "\\t")
)
self.prefix.writeline(
f'in_spec_ = "{escape_string(config.aot_inductor.serialized_in_spec)}";'
)
self.prefix.writeline(
f'out_spec_ = "{escape_string(config.aot_inductor.serialized_out_spec)}";'
)
for idx, output in enumerate(V.graph.graph_outputs):
assert not isinstance(
output, sympy.Expr
), f"output {name=} cannot be symbolic"
name = f"output{idx}"
self.write_input_output_info("outputs_info_", idx, name)
self.prefix.writeline(
"this->kernels_ = std::make_unique<AOTInductorModelKernels>();"
)
self.prefix.writeline("}")
def generate(self, is_inference):
if V.graph.aot_mode:
self.codegen_model_kernels()
self.codegen_model_constructor()
self.write_wrapper_decl()
return super().generate(is_inference)
def finalize_prefix(self):
cached_dtypes_buffer = IndentedBuffer()
if config.aot_inductor.abi_compatible:
for dtype in self.used_cached_dtypes:
cached_dtypes_buffer.writeline(f"CACHE_TORCH_DTYPE({dtype});")
cached_dtypes_buffer.splice(self.prefix)
self.prefix = cached_dtypes_buffer
def define_kernel(
self, name: str, kernel: str, metadata: Optional[str] = None, cuda=False
):
self.header.splice(f"\n{kernel}\n")
def generate_return(self, output_refs):
if V.graph.aot_mode:
cst_names = V.graph.constants.keys()
for idx, output in enumerate(output_refs):
if output in cst_names:
# In some rare cases where we return a constant, we
# have to return a copy of this constant, because
# (1) constants are not owned by the Model instance
# (2) constants remain the same cross inference runs,
# assuming they are not updated at runtime
# Basically, we cannot release or transfer the ownership
# of any origianl constant to the user.
if config.aot_inductor.abi_compatible:
self.wrapper_call.writeline(
f"aoti_torch_clone({output}, &output_handles[{idx}]);"
)
else:
self.wrapper_call.writeline(
f"output_handles[{idx}] = reinterpret_cast<AtenTensorHandle>("
+ f"new at::Tensor(std::move({output}.clone())));"
)
else:
if config.aot_inductor.abi_compatible:
if output in self.cached_thread_locals:
self.wrapper_call.writeline(
f"aoti_torch_new_uninitialized_tensor(&output_handles[{idx}]);"
)
self.wrapper_call.writeline(
f"aoti_torch_assign_tensors({output}, output_handles[{idx}]);"
)
else:
self.wrapper_call.writeline(
f"output_handles[{idx}] = {output}.release();"
)
else:
self.wrapper_call.writeline(
f"output_handles[{idx}] = reinterpret_cast<AtenTensorHandle>("
+ f"new at::Tensor({output}));"
)
else:
self.wrapper_call.writeline(f"return {{{', '.join(output_refs)}}};\n}}")
def generate_end(self, result):
if V.graph.aot_mode:
result.writeline("} // AOTInductorModel::run_impl")
result.writeline("} // namespace aot_inductor")
result.writeline("} // namespace torch")
return
result.writeline("'''\n)")
# get the hash of the wrapper code to name the extension
wrapper_call_hash = codecache.code_hash(result.getvalue())
result.splice(
f"""
module = CppWrapperCodeCache.load(cpp_wrapper_src, '{self.call_func_name}', '{wrapper_call_hash}', {self.cuda})
"""
)
# unwrap output tensor back to python scalar
if all(x for x in self.output_is_tensor.values()):
# If no ShapeAsConstantBuffer in the output, directly return the output as tensors
return_str = "return f(args_tensor)"
else:
outputs = [
f"outputs[{i}]" if self.output_is_tensor[i] else f"outputs[{i}].item()"
for i in range(len(V.graph.graph_outputs))
]
outputs_str = f"[{', '.join(outputs)}]"
return_str = f"""
outputs = f(args_tensor)
return {outputs_str}
"""
args_str = "args_tensor = [arg if isinstance(arg, torch.Tensor) else torch.tensor(arg) for arg in args]"
if V.graph.constants:
# Append constants to the input args for cpp wrapper.
# Python wrapper directly gets the value inside the wrapper call
# as a global variable passed when calling exec(code, mod.__dict__, mod.__dict__).
# For cpp wrapper, we need to pass this python value to the inductor_entry_cpp function explicitly.
assert all(
isinstance(v, torch.Tensor) for v in list(V.graph.constants.values())
), "Expect all constants to be Tensor"
constants_str = f"[{', '.join(V.graph.constants.keys())}]"
args_str += f"""
constants_tensor = {constants_str}
args_tensor.extend(constants_tensor)
"""
# Wrap the func to support setting result._boxed_call = True
result.splice(
f"""
def _wrap_func(f):
def g(args):
{args_str}
{return_str}
return g
call = _wrap_func(module.{self.call_func_name})
"""
)
def generate_c_shim_extern_kernel_call(self, kernel, args):
# In the abi_compatible mode, we call fallback aten ops through a C shim layer
kernel_tokens = kernel.split("::")
kernel_suffix = kernel_tokens[-1]
if kernel_suffix == "call":
kernel_suffix = kernel_tokens[-2]
shim_fn = f"aoti_torch_{kernel_suffix}"
self.writeline(f"AOTI_TORCH_ERROR_CODE_CHECK({shim_fn}({', '.join(args)}));")
def generate_c_shim_extern_kernel_alloc(self, extern_kernel, args):
# registered output buffer name
name = extern_kernel.name
output_handle_name = f"{name}_handle"
self.writeline(f"AtenTensorHandle {output_handle_name};")
output_arg = f"&{output_handle_name}"
self.generate_c_shim_extern_kernel_call(
extern_kernel.codegen_kernel_name(), args + [output_arg]
)
self.writeline(f"RAIIAtenTensorHandle {name}({output_handle_name});")
def generate_extern_kernel_alloc(self, extern_kernel, args):
if V.graph.aot_mode and config.aot_inductor.abi_compatible:
self.generate_c_shim_extern_kernel_alloc(extern_kernel, args)
else:
super().generate_extern_kernel_alloc(extern_kernel, args)
def generate_c_shim_fallback_kernel(self, fallback_kernel, args):
output_args = []
output_raii_handles = []
output_name_base = fallback_kernel.get_name()
for idx, output in enumerate(fallback_kernel.outputs):
if isinstance(output, ir.MultiOutput):
name = f"{output.get_name()}"
output_handle_name = f"{name}_handle"
if output.indices:
assert (
output.indices[0][1] == idx
), f"expected {output.indices[0][1]=} == {idx=} for {output_name_base=}"
self.writeline(f"AtenTensorHandle {output_handle_name};")
output_args.append(f"&{output_handle_name}")
output_raii_handles.append(
f"RAIIAtenTensorHandle {name}({output_handle_name});"
)
elif isinstance(output, int):
output_name = f"{output_name_base}_{idx}"
self.writeline(f"int64_t {output_name} = {output};")
output_args.append(f"&{output_name}")
elif output is None:
output_args.append("nullptr")
else:
raise NotImplementedError("unsupported type of {output=}")
args = args + output_args
assert (
fallback_kernel.abi_compatible_kernel is not None
), f"abi_compatible_kernel is None for {fallback_kernel.kernel=}"
self.generate_c_shim_extern_kernel_call(
fallback_kernel.abi_compatible_kernel, args
)
for raii_handle in output_raii_handles:
self.writeline(raii_handle)
def generate_fallback_kernel(self, fallback_kernel, args):
if V.graph.aot_mode and config.aot_inductor.abi_compatible:
self.generate_c_shim_fallback_kernel(fallback_kernel, args)
else:
super().generate_fallback_kernel(fallback_kernel, args)
def generate_extern_kernel_out(self, output_view, codegen_reference, args, kernel):
if output_view:
output_as_strided = f"{output_view.codegen_reference()}"
output_name = f"{output_view.get_name()}_as_strided"
self.writeline(f"auto {output_name} = {output_as_strided};")
args.insert(0, output_name)
else:
args.insert(0, f"{codegen_reference}")
if V.graph.aot_mode and config.aot_inductor.abi_compatible:
self.generate_c_shim_extern_kernel_call(kernel, args)
else:
self.writeline(self.wrap_kernel_call(kernel, args))
def generate_user_defined_triton_kernel(self, kernel_name, grid, configs, args):
assert len(grid) != 0
if len(grid) == 1:
grid_decision = grid[0]
else:
meta = CudaKernelParamCache.get(kernel_name)
assert meta is not None
grid_decision = None
for i, c in enumerate(configs):
if all(arg == meta["meta"][key] for key, arg in c.kwargs.items()):
grid_decision = grid[i]
break
assert grid_decision is not None
self.generate_kernel_call(
kernel_name,
args,
grid=grid_decision,
device_index=V.graph.scheduler.current_device.index,
cuda=True,
triton=True,
)
def generate_scatter_fallback(
self, output, inputs, kernel, fn, src_is_tensor, reduce, kwargs
):
# TODO: support other overload for cpp wrapper and remove the below assertions
if V.graph.aot_mode and config.aot_inductor.abi_compatible:
# call the ABI shim function instead of the ATen one
kernel = kernel.replace("at::", "aoti_torch_")
line = f"{kernel}({output}, {','.join(map(str, inputs))}"
if fn == "aten.scatter_":
if src_is_tensor:
if reduce:
line += f", {V.graph.wrapper_code.val_to_arg_str(reduce)}"
else:
assert (
reduce is None
), "Expect reduce to be None for aten.scatter_ with scalar src"
else:
line += f", {','.join(kwargs)}"
line += f"){self.ending}"
self.writeline(line)
def add_benchmark_harness(self, output):
if V.graph.aot_mode:
return
super().add_benchmark_harness(output)
def codegen_sizevar(self, x: Expr) -> str:
return self.expr_printer(V.graph.sizevars.simplify(x))
def codegen_tuple_access(self, basename: str, name: str, index: str) -> str:
if V.graph.aot_mode and config.aot_inductor.abi_compatible:
# in the abi_compatible mode, outputs are returned via arguments
return name
else:
return f"std::get<{index}>({basename})"
def codegen_shape_tuple(self, shape: Tuple[Expr, ...]) -> str:
parts = list(map(self.codegen_sizevar, shape))
if len(parts) == 0:
return "{}"
if len(parts) == 1:
return f"{{{parts[0]}, }}"
return f"{{{', '.join(parts)}}}"
def is_statically_known_int(self, x):
try:
val = V.graph._shape_env._maybe_evaluate_static(x)
int(x)
return True
except Exception:
return False
def is_statically_known_list_of_ints(self, lst):
return all(isinstance(self.is_statically_known_int(x), int) for x in lst)
def can_prove_buffer_has_static_shape(self, buffer):
return self.is_statically_known_list_of_ints(buffer.get_size())
def can_cache_buffer_in_thread_local(self, buffer):
# We are gated off on CUDA because this is intended to reduce overhead in
# overhead-bound CPU use case.
return (
not self.cuda
and config.allow_buffer_reuse
and self.can_prove_buffer_has_static_shape(buffer)
)
def make_buffer_free(self, buffer):
return (
""
if isinstance(buffer.get_layout(), ir.MultiOutputLayout)
or (V.graph.aot_mode and self.can_cache_buffer_in_thread_local(buffer))
else f"{buffer.get_name()}.reset();"
)
def make_free_by_names(self, names_to_del: List[str]):
return " ".join(f"{name}.reset();" for name in names_to_del)
def codegen_exact_buffer_reuse(self, old_name: str, new_name: str, del_line: str):
if config.aot_inductor.abi_compatible:
return f"auto {new_name} = std::move({old_name}); // reuse"
else:
return super().codegen_exact_buffer_reuse(old_name, new_name, del_line)
def generate_profiler_mark_wrapper_call(self, stack):
self.wrapper_call.writeline(
'RECORD_FUNCTION("inductor_wrapper_call", c10::ArrayRef<c10::IValue>());'
)
def write_triton_header_once(self):
pass
def generate_start_graph(self):
pass
def generate_end_graph(self):
pass
def generate_inf_and_nan_checker(self, nodes):
for buf in nodes.get_names():
# TODO: Add buf name directly into check_inf_and_nan.
self.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_check_inf_and_nan({buf}));"
)
def codegen_device(self, device):
if config.aot_inductor.abi_compatible:
return f"cached_torch_device_type_{device.type},{device.index if device.index else 0}"
else:
from .cpp import DEVICE_TO_ATEN
return (
f"c10::Device({DEVICE_TO_ATEN[device.type]}, {device.index})"
if device.index is not None
else f"{DEVICE_TO_ATEN[device.type]}"
)
def codegen_dtype(self, dtype):
if config.aot_inductor.abi_compatible:
dtype_str = str(dtype).split(".")[-1]
self.used_cached_dtypes.add(dtype_str)
return f"cached_torch_dtype_{dtype_str}"
else:
from .cpp import DTYPE_TO_ATEN
return DTYPE_TO_ATEN[dtype]
@functools.lru_cache(None)
def codegen_int_array_var(self, int_array: str, writer=None):
# Because the memory planning is done in two passes (see the implementation
# of self.generate), the writeline behavior is different in the two passes.
# As a result, the emitted int array declarations may appear in a later
# position of the generated code, so the second pass codegen should not
# reuse int array declarations generated in the first pass
if writer is None:
# The first pass codegen uses `self` as the writer
writer = self
var = f"int_array_{next(self.int_array_id)}"
if var not in self.declared_int_array_vars:
self.declared_int_array_vars.add(var)
writer.writeline(f"int64_t {var}[] = {int_array};")
return var
def make_buffer_allocation(self, buffer):
return self.make_allocation(
buffer.get_name(),
buffer.get_device(),
buffer.get_dtype(),
buffer.get_size(),
buffer.get_stride(),
self.can_cache_buffer_in_thread_local(buffer),
)
def make_allocation(
self, name, device, dtype, shape, stride, can_cache_buffer_in_thread_local=False
):
device = self.codegen_device(device)
dtype = self.codegen_dtype(dtype)
size = self.codegen_shape_tuple(shape)
stride = self.codegen_shape_tuple(stride)
if config.aot_inductor.abi_compatible:
device_type, device_id = device.split(",")
args = [
str(len(shape)),
self.codegen_int_array_var(size, self.wrapper_call),
self.codegen_int_array_var(stride, self.wrapper_call),
dtype,
device_type,
"this->device_idx_" if V.graph.aot_mode else device_id,
f"&{name}_handle",
]
def gen_alloc(wrapper_call, name, args):
wrapper_call.writeline(f"AtenTensorHandle {name}_handle;")
wrapper_call.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided({', '.join(args)}));"
)
if can_cache_buffer_in_thread_local:
self.cached_thread_locals.add(name)
self.wrapper_call.writeline(
f"thread_local RAIIAtenTensorHandle {name}_handle = ([&] {{"
)
with self.wrapper_call.indent():
gen_alloc(self.wrapper_call, name, args)
self.wrapper_call.writeline(f"return {name}_handle;")
self.wrapper_call.writeline("})();")
return f"AtenTensorHandle {name}({name}_handle.get());"
else:
gen_alloc(self.wrapper_call, name, args)
return f"RAIIAtenTensorHandle {name}({name}_handle);"
if V.graph.aot_mode and device.startswith("c10::Device("):
tensor_device = f"{device.split(',')[0]}, this->device_idx_)"
else:
tensor_device = device
return (
f"{self.declare}{name} = {self.namespace}empty_strided("
f"{size}, {stride}, at::TensorOptions({tensor_device}).dtype({dtype})){self.ending}"
)
def codegen_alloc_from_pool(self, name, offset, dtype, shape, stride) -> str:
if config.aot_inductor.abi_compatible:
size = self.codegen_shape_tuple(shape)
stride = self.codegen_shape_tuple(stride)
tmp_name = f"tmp_tensor_handle_{next(self.tmp_tensor_id)}"
args = [
name,
pexpr(offset), # bytes not numel
self.codegen_dtype(dtype),
str(len(shape)),
self.codegen_int_array_var(size, self.wrapper_call),
self.codegen_int_array_var(stride, self.wrapper_call),
f"&{tmp_name}",
]
self.wrapper_call.writeline(f"AtenTensorHandle {tmp_name};")
self.wrapper_call.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch__alloc_from_pool({', '.join(args)}));"
)
return f"RAIIAtenTensorHandle({tmp_name})"
return "alloc_from_pool({})".format(
", ".join(
[
name,
pexpr(offset), # bytes not numel
self.codegen_dtype(dtype),
self.codegen_shape_tuple(shape),
self.codegen_shape_tuple(stride),
]
)
)
def codegen_reinterpret_view(
self, data, size_list, stride_list, offset, writer
) -> str:
dim = str(len(size_list))
size = self.codegen_shape_tuple(size_list)
stride = self.codegen_shape_tuple(stride_list)
offset = self.codegen_sizevar(offset)
if config.aot_inductor.abi_compatible:
tmp_name = f"tmp_tensor_handle_{next(self.tmp_tensor_id)}"
# Because the memory planning is done in two passes (see the implementation
# of self.generate), the writeline behavior is different in the two passes.
if writer is None:
writer = self
args = [
f"{data.get_name()}",
dim,
self.codegen_int_array_var(size, writer),
self.codegen_int_array_var(stride, writer),
offset,
f"&{tmp_name}",
]
def gen_reinterpret_call(writer, args):
writer.writeline(f"AtenTensorHandle {tmp_name};")
writer.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch__reinterpret_tensor({', '.join(args)}));"
)
if (
self.can_cache_buffer_in_thread_local(data)
and self.is_statically_known_list_of_ints(size_list)
and self.is_statically_known_list_of_ints(stride_list)
):
self.cached_thread_locals.add(tmp_name)
writer.writeline(
f"thread_local RAIIAtenTensorHandle {tmp_name}_handle = ([&] {{"
)
if hasattr(writer, "indent"):
indent = writer.indent()
else:
indent = contextlib.nullcontext()
with indent:
gen_reinterpret_call(writer, args)
writer.writeline(f"return {tmp_name};")
writer.writeline("})();")
writer.writeline(
f"AtenTensorHandle {tmp_name}({tmp_name}_handle.get());"
)
return tmp_name
gen_reinterpret_call(writer, args)
# NB, the return handle here represents a temporary tensor, which will be automatically
# released.
# Here's a sample usage in the cpp wrapper code:
# ```
# aoti_torch_addmm_out(
# buf1,
# arg1_1,
# RAIIAtenTensorHandle(tmp_tensor_handle_0),
# buf0,
# 1L,
# 1L));
# ```
# RAIIAtenTensorHandle(tmp_tensor_handle_0) will be released after the call to addmm_out.
# This could be problematic when it's used in a different pattern, for example:
# ````
# AtenTensorHandle tensor_args[] = {RAIIAtenTensorHandle(tmp_tensor_handle_2), buf5, buf6};
# aoti_torch_proxy_executor_call_function(..., tensor_args);
# ````
# RAIIAtenTensorHandle(tmp_tensor_handle_2) will be invalid when it's used in the latter
# kernel call.
#
# This is solved by updating the proxy_executor invocation to
# ```
# aoti_torch_proxy_executor_call_function(...,
# std::vector<AtenTensorHandle>{
# RAIIAtenTensorHandle(tmp_tensor_handle_2), buf5, buf6
# }.data()
# );
# ```
return f"RAIIAtenTensorHandle({tmp_name})"
else:
args = [data.get_name(), size, stride, offset]
return f"reinterpret_tensor({', '.join(args)})"
def codegen_device_copy(self, src, dst):
if config.aot_inductor.abi_compatible:
self.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_tensor_copy_({src}, {dst}));"
)
else:
self.writeline(f"{dst}.copy_({src});")
def codegen_multi_output(self, name, value):
# in the abi_compatible mode, outputs are retrieved by passing
# output pointers, so we skip its codegen here.
if not config.aot_inductor.abi_compatible:
super().codegen_multi_output(name, value)
def generate_extern_kernel_args_decl_if_needed(
self, op_overload, raw_args, output_args
):
arg_types = [x.real_type for x in op_overload._schema.arguments]
return_types = [x.type for x in op_overload._schema.returns]
new_tensor_args = []
new_int_args = []
def fill_args(arg, arg_type):
static_arg_types = (
torch.FloatType,
torch.BoolType,
torch.StringType,
torch.Type,
torch.DeviceObjType,
)
inductor_tensor_buffers = (
ir.Buffer,
ir.ReinterpretView,
)
if isinstance(arg_type, torch.TensorType):
assert isinstance(arg, inductor_tensor_buffers), f"got {type(arg)}"
new_tensor_args.append(f"{arg.codegen_reference()}")
elif isinstance(arg_type, torch.IntType):
# int
new_int_args.append(str(arg))
elif isinstance(arg_type, torch.SymIntType):
# SymInt
new_int_args.append(str(arg))
elif isinstance(arg_type, torch.NumberType):
# Scalar of type int
assert isinstance(arg, (int, float, bool))
# Only treat int Scalar as dynamic
if isinstance(arg, int):
new_int_args.append(str(arg))
elif isinstance(arg_type, torch.ListType):
assert isinstance(arg, (list, tuple))
# List[Tensor]
if isinstance(arg_type.getElementType(), torch.TensorType):
new_tensor_args.extend([f"{a.codegen_reference()}" for a in arg])
# List[Optional[Tensor]]
elif isinstance(
arg_type.getElementType(), torch.OptionalType
) and isinstance(
arg_type.getElementType().getElementType(), torch.TensorType
):
new_tensor_args.extend(
[f"{a.codegen_reference()}" for a in arg if a is not None]
)
# List [int] or List[SymInt]
elif isinstance(
arg_type.getElementType(), (torch.IntType, torch.SymIntType)
):
new_int_args.extend([str(a) for a in arg])
# List[Scalar]
elif isinstance(arg_type.getElementType(), torch.NumberType):
# Only treat int Scalar as dynamic
is_int_type = [isinstance(a, int) for a in arg]
if any(is_int_type):
assert all(
is_int_type
), "AOTInductor only supports int scalars of the same type"
new_int_args.extend([str(a) for a in arg])
else:
assert isinstance(
arg_type.getElementType(), static_arg_types # type: ignore[arg-type]
), f"Fall through arguments must be one of static_arg_types, got {type(arg_type)}"
else:
assert isinstance(
arg_type, static_arg_types # type: ignore[arg-type]
), f"Fall through arguments must be one of static_arg_types, got {type(arg_type)}"
for arg, arg_type in zip(raw_args, arg_types):
if arg is not None:
if isinstance(arg_type, torch.OptionalType):
fill_args(arg, arg_type.getElementType())
else:
fill_args(arg, arg_type)
def fill_output_arg(arg, return_type):
if isinstance(return_type, torch.TensorType):
self.writeline(f"AtenTensorHandle {arg}_handle; // output buffer")
self.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&{arg}_handle));"
)
self.writeline(f"RAIIAtenTensorHandle {arg}({arg}_handle);")
new_tensor_args.append(f"{arg}")
elif isinstance(return_type, torch.SymIntType):
raise NotImplementedError("NYI support for return type: SymInt")
elif isinstance(return_type, torch.ListType) and isinstance(
return_type.getElementType(), torch.SymIntType
):
raise NotImplementedError("NYI support for return type: List[SymInt]")
else:
raise AssertionError(f"Unsupported return type found: {return_type}")
# TODO: Only support tensor(s) returns for now, SymInt is not implemented yet
for return_type in return_types:
if isinstance(return_type, (torch.TensorType)):
pass
elif isinstance(return_type, torch.OptionalType):
assert isinstance(return_type.getElementType(), torch.TensorType)
elif isinstance(return_type, torch.ListType):
assert isinstance(return_type.getElementType(), torch.TensorType)
else:
raise NotImplementedError(
f"return type {return_type} is not yet supported."
)
for output_arg in output_args:
assert output_arg is not None, "Optional return types are not yet supported"
if isinstance(output_arg, (list, tuple)):
for out in output_arg:
fill_output_arg(out, torch.TensorType.get())
else:
fill_output_arg(output_arg, torch.TensorType.get())
return new_tensor_args, new_int_args
def generate_extern_kernel_alloc_and_find_schema_if_needed(
self,
name,
kernel,
codegen_args,
cpp_op_schema,
cpp_kernel_key,
cpp_kernel_overload_name="",
op_overload=None,
raw_args=None,
outputs=None,
):
if config.is_fbcode():
assert op_overload is not None
assert raw_args is not None
assert outputs is not None
return self.generate_extern_kernel_alloc_and_find_schema_if_needed_fbcode(
name,
cpp_kernel_key,
op_overload,
raw_args,
outputs,
)
else:
return self.generate_extern_kernel_alloc_and_find_schema_if_needed_oss(
name,
kernel,
codegen_args,
cpp_op_schema,
cpp_kernel_key,
cpp_kernel_overload_name,
)
def generate_extern_kernel_alloc_and_find_schema_if_needed_oss(
self,
name,
kernel,
codegen_args,
cpp_op_schema,
cpp_kernel_key,
cpp_kernel_overload_name="",
):
if cpp_kernel_key not in self.extern_call_ops:
self.writeline(
f"static auto op_{cpp_kernel_key} = c10::Dispatcher::singleton()"
)
self.writeline(
f'\t.findSchemaOrThrow("{kernel}", "{cpp_kernel_overload_name}")'
)
self.writeline(f"\t.typed<{cpp_op_schema}>();")
self.extern_call_ops.add(cpp_kernel_key)
self.writeline(
f"auto {name} = op_{cpp_kernel_key}.call({', '.join(codegen_args)});"
)
def generate_extern_kernel_alloc_and_find_schema_if_needed_fbcode(
self,
name,
cpp_kernel_key,
op_overload,
raw_args, # contains both args and flatten kwargs
outputs,
):
def extract_output_name(out):
assert out is not None, "None, i.e. optional output is not supported"
if isinstance(out, ir.MultiOutput):
return out.get_name()
elif isinstance(out, (list, tuple)):
return type(out)(extract_output_name(o) for o in out)
else:
raise AssertionError(f"Unexpected output: {type(out)}")
# output_args has the same pytree structure as outputs
output_args = extract_output_name(outputs)
if isinstance(output_args, str):
output_args = [output_args]
(
tensor_call_args,
int_call_args,
) = self.generate_extern_kernel_args_decl_if_needed(
op_overload, raw_args, output_args
)
tensor_call_args_str = ", ".join(tensor_call_args)
int_call_args_str = ", ".join(int_call_args)
extern_kernel_node_index = len(V.graph.extern_kernel_nodes) - 1
self.writeline(
f"aoti_torch_proxy_executor_call_function(proxy_executor, "
f"{extern_kernel_node_index}, "
f"{len(int_call_args)}, "
f"std::vector<int64_t>{{{int_call_args_str}}}.data(), "
f"{len(tensor_call_args)}, "
f"std::vector<AtenTensorHandle>{{{tensor_call_args_str}}}.data());"
)
self.extern_call_ops.add(cpp_kernel_key)
def val_to_cpp_arg_str(self, type_, val, is_legacy_abi) -> str:
if (
config.aot_inductor.abi_compatible
and not is_legacy_abi
and isinstance(type_, torch.OptionalType)
):
if val is None:
return "0" # nullptr is not available in C
if isinstance(val, (bool, int, str, float)):
var_name = f"var_{next(self.arg_var_id)}"
self.writeline(f"auto {var_name} = {self.val_to_arg_str(val)};")
return f"&{var_name}"
if not isinstance(type_.getElementType(), torch.TensorType):
return f"&{self.val_to_arg_str(val)}"
return self.val_to_arg_str(val)
def val_to_arg_str(self, val) -> str:
if val is None:
# When None is passed as an argument, it represents an optional that does not contain a value.
if config.aot_inductor.abi_compatible:
return "0" # nullptr is not available in C
return "c10::nullopt"
elif isinstance(val, bool):
if config.aot_inductor.abi_compatible:
return "1" if val else "0"
else:
return "true" if val else "false"
elif isinstance(val, int):
return f"{val}L"
elif isinstance(val, str):
return f'"{val}"'
elif isinstance(val, (ComputedBuffer, InputBuffer, ReinterpretView)):
return val.codegen_reference()
elif isinstance(val, torch.device):
return self.codegen_device(val)
elif isinstance(val, torch.dtype):
return self.codegen_dtype(val)
elif isinstance(val, float) and val in [float("inf"), float("-inf")]:
if val == float("inf"):
return "std::numeric_limits<float>::infinity()"
else:
return "-std::numeric_limits<float>::infinity()"
elif isinstance(val, (list, tuple)):
# FIXME handle embedded optional types?
result = f"{{{', '.join(self.val_to_arg_str(x) for x in val)}}}"
if config.aot_inductor.abi_compatible:
# Need to pass the array length because we can't use std::vector
return f"{self.codegen_int_array_var(result)}, {len(val)}"
else:
return result
else:
return repr(val)
class CudaWrapperCodeGen(CppWrapperCodeGen):
"""
Generates cpp wrapper for running on GPU and calls CUDA kernels
"""
def __init__(self):
super().__init__()
self.grid_id = count()
self.cuda = True
def write_header(self):
super().write_header()
self.header.splice("#include <filesystem>")
if not config.aot_inductor.abi_compatible:
self.header.splice(
"""
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
"""
)
self.header.splice(
"""
#define CUDA_DRIVER_CHECK(EXPR) \\
do { \\
CUresult code = EXPR; \\
const char *msg; \\
cuGetErrorString(code, &msg); \\
if (code != CUDA_SUCCESS) { \\
throw std::runtime_error( \\
std::string("CUDA driver error: ") + \\
std::string(msg)); \\
} \\
} while (0);
namespace {
struct Grid {
Grid(uint32_t x, uint32_t y, uint32_t z)
: grid_x(x), grid_y(y), grid_z(z) {}
uint32_t grid_x;
uint32_t grid_y;
uint32_t grid_z;
bool is_non_zero() {
return grid_x > 0 && grid_y > 0 && grid_z > 0;
}
};
} // anonymous namespace
static inline CUfunction loadKernel(
std::string filePath,
const std::string &funcName,
uint32_t sharedMemBytes,
const std::optional<std::string> &cubinDir = std::nullopt) {
if (cubinDir) {
std::filesystem::path p1{*cubinDir};
std::filesystem::path p2{filePath};
filePath = (p1 / p2.filename()).string();
}
CUmodule mod;
CUfunction func;
CUDA_DRIVER_CHECK(cuModuleLoad(&mod, filePath.c_str()));
CUDA_DRIVER_CHECK(cuModuleGetFunction(&func, mod, funcName.c_str()));
if (sharedMemBytes > 0) {
CUDA_DRIVER_CHECK(cuFuncSetAttribute(
func,
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
sharedMemBytes
))
}
return func;
}
static inline void launchKernel(
CUfunction func,
uint32_t gridX,
uint32_t gridY,
uint32_t gridZ,
uint32_t numWarps,
uint32_t sharedMemBytes,
void* args[],
cudaStream_t stream) {
CUDA_DRIVER_CHECK(cuLaunchKernel(
func, gridX, gridY, gridZ, 32*numWarps, 1, 1, sharedMemBytes, stream, args, nullptr
));
}
"""
)
def write_get_raw_stream(self, index):
name = f"stream{index}"
self.writeline(
f"cudaStream_t {name} = at::cuda::getCurrentCUDAStream({index});"
)
return name
def define_kernel(
self, name: str, kernel: str, metadata: Optional[str] = None, cuda=True
):
if not cuda:
return super().define_kernel(name, kernel, metadata, cuda)
def generate(self, is_inference):
self.prefix.writeline("\n")
if not V.graph.aot_mode:
for kernel in chain(
self.src_to_kernel.values(), self.user_defined_kernel_cache.values()
):
self.prefix.writeline(f"static CUfunction {kernel} = nullptr;")
self.prefix.writeline("\n")
return super().generate(is_inference)
@functools.lru_cache(None)
def generate_load_kernel_once(
self, name: str, mangled_name: str, cubin_path: str, shared_mem: int
):
if V.graph.aot_mode:
self.writeline(f"if (kernels.{name} == nullptr) {{")
self.writeline(
f""" kernels.{name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem}, this->cubin_dir_);"""
)
self.writeline("}")
else:
self.writeline(f"if ({name} == nullptr) {{")
self.writeline(
f""" {name} = loadKernel("{cubin_path}", "{mangled_name}", {shared_mem});"""
)
self.writeline("}")
def generate_args_decl(self, call_args):
dynamic_symbols = V.graph.sizevars.free_symbols()
# TODO: only works for constant now, need type info
new_args = []
for arg in call_args:
var_name = f"var_{next(self.arg_var_id)}"
if isinstance(arg, (sympy.Integer, sympy.Symbol, SymbolicCallArg)):
self.writeline(f"auto {var_name} = {arg};")
elif isinstance(arg, sympy.Expr):
self.writeline(f"auto {var_name} = {self.expr_printer(arg)};")
elif is_int(arg):
self.writeline(f"int {var_name} = {arg};")
elif is_float(arg):
self.writeline(f"float {var_name} = {arg};")
elif any(str(arg) == s.name for s in dynamic_symbols):
self.writeline(f"auto {var_name} = {arg};")
elif arg == "nullptr":
self.writeline(f"auto {var_name} = nullptr;")
elif arg == "c10::nullopt":
self.writeline(f"auto {var_name} = c10::nullopt;")
else:
if config.aot_inductor.abi_compatible:
self.writeline(f"CUdeviceptr {var_name};")
self.writeline(
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_data_ptr({arg}, reinterpret_cast<void**>(&{var_name})));"
)
else:
self.writeline(
f"CUdeviceptr {var_name} = reinterpret_cast<CUdeviceptr>({arg}.data_ptr());"
)
new_args.append(f"&{var_name}")
return ", ".join(new_args)
def generate_default_grid(self, name: str, grid: List[Any], cuda: bool = True):
"""
Generate grid configs for launching a CUDA kernel using the grid
function from triton_heuristics.
"""
if not cuda:
return grid
assert isinstance(grid, list), f"expected {grid=} to be a list"
grid = [e.inner_expr if isinstance(e, SymbolicCallArg) else e for e in grid]
grid_fn = default_grid(*grid)
params = CudaKernelParamCache.get(name)
assert (
params is not None
), f"cuda kernel parameters for {name} should already exist at this moment"
block_cfg = {
"XBLOCK": params["x_block"],
"YBLOCK": params["y_block"],
"ZBLOCK": params["z_block"],
}
return grid_fn(block_cfg)
def generate_kernel_call(
self, name, call_args, grid=None, device_index=None, cuda=True, triton=True
):
if not cuda:
# Even in CudaWrapperCodeGen, we may see cpp kernels
return super().generate_kernel_call(
name, call_args, grid, device_index, cuda, triton
)
params = CudaKernelParamCache.get(name)
assert (
params is not None
), f"cuda kernel parameters for {name} should already exist at this moment"
mangled_name = params.get("mangled_name", None)
assert mangled_name is not None, "missing mangled_name"
cubin_path = params.get(get_cpp_wrapper_cubin_path_name(), None)
assert cubin_path is not None and os.path.exists(
cubin_path
), f"cubin file should already exist at this moment: {cubin_path}"
shared_mem = params.get("shared_mem", 0)
self.generate_load_kernel_once(name, mangled_name, cubin_path, shared_mem)
call_args = self.generate_args_decl(call_args)
kernel_args_var = f"kernel_args_var_{next(self.kernel_callsite_id)}"
self.writeline(f"void* {kernel_args_var}[] = {{{call_args}}};")
stream = (
"stream" if V.graph.aot_mode else self.write_get_raw_stream(device_index)
)
grid_name = f"{name}_grid_{next(self.grid_id)}"
assert isinstance(
grid, (list, tuple)
), f"expected grid to be a list or tuple but got: {grid=}"
grid = [V.graph.sizevars.simplify(item) for item in grid]
grid_has_unbacked_symbols = any(free_unbacked_symbols(item) for item in grid)
grid_args = [self.grid_expr_printer(item) for item in grid]
grid_args_str = ", ".join(grid_args)
self.writeline(f"Grid {grid_name} = Grid({grid_args_str});")
if grid_has_unbacked_symbols:
self.writeline(f"if ({grid_name}.is_non_zero()) {{")
kernel_var_name = f"kernels.{name}" if V.graph.aot_mode else name
self.writeline(
"launchKernel({}, {}, {}, {}, {}, {}, {}, {});".format(
kernel_var_name,
f"{grid_name}.grid_x",
f"{grid_name}.grid_y",
f"{grid_name}.grid_z",
params["num_warps"],
params["shared_mem"],
kernel_args_var,
stream,
)
)
if grid_has_unbacked_symbols:
self.writeline("}")