blob: 4c8ed41560b49cf7d5659c6a4a069f537ccd8a2a [file] [log] [blame]
import logging
import operator
import os
import re
import sys
import time
from typing import Dict, List, Optional, Set
import sympy
import torch
import torch._logging
import torch.fx
from torch._decomp import get_decompositions
from torch._dynamo.utils import dynamo_timed
from torch.fx.experimental.symbolic_shapes import (
magic_methods,
method_to_operator,
ShapeEnv,
SymTypes,
)
from torch.utils._mode_utils import no_dispatch
from .._dynamo import config as dynamo_config
from . import config, ir
from .codegen.wrapper import CppWrapperCodeGen, CudaWrapperCodeGen, WrapperCodeGen
from .exc import (
LoweringException,
MissingOperatorWithDecomp,
MissingOperatorWithoutDecomp,
)
from .ir import Constant, FixedLayout, InputBuffer, Pointwise, Reduction, TensorBox
from .lowering import (
FALLBACK_ALLOW_LIST,
fallback_handler,
fallback_node_due_to_unsupported_type,
layout_constraints,
lowerings,
make_fallback,
needs_realized_inputs,
unsupported_output_tensor,
)
from .sizevars import SizeVarAllocator
from .utils import (
convert_shape_to_inductor,
gather_origins,
get_dtype_size,
sympy_product,
)
from .virtualized import V
log = logging.getLogger(__name__)
output_code_log = torch._logging.getArtifactLogger(__name__, "output_code")
def supported_dtype_of_cpp_wrapper(dtype, cuda):
supported_dtype = {
torch.float32,
torch.float64,
torch.int64,
torch.int32,
torch.int16,
torch.int8,
torch.uint8,
torch.bool,
torch.bfloat16,
# torch.float16, # TODO: implement this
}
if cuda:
supported_dtype.add(torch.float16)
return dtype in supported_dtype
def may_get_constant_buffer_dtype(constant_buffer):
assert isinstance(
constant_buffer, (sympy.Symbol, sympy.core.numbers.Integer)
), "get_constant_buffer_dtype only supports input of sympy.Symbol or sympy.core.numbers.Integer"
if isinstance(constant_buffer, sympy.core.numbers.Integer):
return torch.int64
if constant_buffer.is_integer:
return torch.int64
elif constant_buffer.is_float:
return torch.float32
else:
return None
def is_magic_method(op):
magic_ops = {method_to_operator(m) for m in magic_methods}
return op in magic_ops
class GraphLowering(torch.fx.Interpreter):
def symbolic_sizes_strides(self, ex: torch.Tensor):
"""
Support dynamic shapes and dynamic strides by assigning variables
to each dimension. We duck-shape tensors, so if two tensors
have the same size they get assigned the same symbolic variable.
"""
if self.reuse_shape_env:
return convert_shape_to_inductor(ex.size()), convert_shape_to_inductor(
ex.stride()
)
else:
from torch._dynamo.source import ConstantSource
# TODO: this should not be needed once #93059 lands
# https://github.com/pytorch/pytorch/pull/94031#discussion_r1096044816
# TODO: make a dedicated UnknownSource for this?
# NB: This is using the legacy default behavior from
# create_symbolic_sizes_strides_storage_offset but we hope we can
# just delete this entirely
source = ConstantSource(
f"__inductor_unknown_tensor_{len(self._shape_env.var_to_val)}"
)
(
size,
stride,
_,
) = self._shape_env.create_symbolic_sizes_strides_storage_offset(
ex,
source,
)
size = [i.node.expr if isinstance(i, torch.SymInt) else i for i in size]
stride = [i.node.expr if isinstance(i, torch.SymInt) else i for i in stride]
return size, stride
def static_sizes_strides(self, ex: torch.Tensor):
"""
Primarily used to weights
"""
size = [sympy.Integer(i) for i in ex.size()]
stride = [sympy.Integer(i) for i in ex.stride()]
return size, stride
def __init__(
self,
gm: torch.fx.GraphModule,
shape_env=None,
num_static_inputs=None,
graph_id=None,
cpp_wrapper=False,
aot_mode=False,
):
super().__init__(gm)
self.extra_traceback = False # we do our own error wrapping
if shape_env is None:
shape_env = ShapeEnv()
self.reuse_shape_env = False
else:
self._shape_env = shape_env
self.reuse_shape_env = True
self._shape_env = shape_env
self.sizevars = SizeVarAllocator(shape_env)
self.graph_inputs: Dict[str, TensorBox] = {}
self.graph_inputs_original: Dict[str, InputBuffer] = {}
self.graph_outputs: Optional[List[ir.IRNode]] = None
self.device_types: Set[str] = set()
self.device_idxs: Set[int] = set()
self.cuda = False
self.buffers: List[ir.ComputedBuffer] = []
self.constants: Dict[str, torch.Tensor] = {}
self.removed_buffers: Set[str] = set()
self.inplaced_to_remove: Set[str] = set()
self.wrapper_code: Optional[WrapperCodeGen] = None
self.num_static_inputs = num_static_inputs
self.mutated_inputs: Set[str] = set()
self.unaligned_buffers: Set[str] = set()
self.randomness_offset = sympy.Integer(0)
self.randomness_seeds: List[str] = []
self.name_to_buffer: Dict[str, ir.ComputedBuffer] = {}
self.creation_time = time.time()
self.name = "GraphLowering"
self.cpp_wrapper = cpp_wrapper
self.aot_mode = aot_mode
self.graph_id = graph_id
self.scheduler = None
self._warned_fallback = {"aten.convolution_backward"}
def warn_fallback(self, name):
if name not in self._warned_fallback:
self._warned_fallback.add(name)
log.info("Using FallbackKernel: %s", name)
def add_device_idx(self, idx: Optional[int]):
if idx is not None:
self.device_idxs.add(idx)
@property
def fake_mode(self):
return V.fake_mode
def get_buffer(self, buffer_name: str):
if buffer_name in self.name_to_buffer:
return self.name_to_buffer[buffer_name]
if buffer_name in self.graph_inputs:
return self.graph_inputs[buffer_name]
return None
def get_dtype(self, buffer_name: str):
if buffer_name in self.constants:
return self.constants[buffer_name].dtype
if buffer_name in self.name_to_buffer:
return self.name_to_buffer[buffer_name].get_dtype()
if buffer_name in self.graph_inputs:
return self.graph_inputs[buffer_name].get_dtype()
m = re.match(r"as_strided\(([a-zA-Z0-9_]+),", buffer_name)
if m:
return self.get_dtype(m.group(1))
raise KeyError(f"could not find {buffer_name}")
def get_numel(self, buffer_name: str):
from .ir import MultiOutputLayout
if buffer_name in self.constants:
return self.constants[buffer_name].numel()
if buffer_name in self.name_to_buffer:
buf = self.name_to_buffer[buffer_name]
if isinstance(getattr(buf, "layout", None), MultiOutputLayout):
return 1
return buf.get_numel()
if buffer_name in self.graph_inputs:
return self.graph_inputs[buffer_name].get_numel()
raise KeyError(f"could not find {buffer_name}")
def random_seed_buffer(self, device: torch.device):
"""
Return a device-unique 1-element tensor storing our RNG seed.
This will get initialized at the start of each graph in
`wrapper.py`.
Note this is only used by cuda backends. The CPU backend handles
RNG seeds as a sizevar.
"""
name = f"seed_{device.type}_{device.index}"
if name not in self.constants:
self.constants[name] = torch.zeros((), device=device, dtype=torch.int64)
self.randomness_seeds.append(name)
return ir.RandSeedBuffer(
name=name,
layout=ir.FixedLayout(
device=device,
dtype=torch.int64,
size=[],
stride=[],
),
)
def increment_randomness_offset(self, numel):
"""
A global counter of how many random numbers we have handed out so far.
"""
offset = self.randomness_offset
self.randomness_offset = offset + numel
return offset
@dynamo_timed
def run(self, *args):
return super().run(*args)
def disable_cpp_wrapper(self, cond):
self.cpp_wrapper = False
log.debug("Set cpp_wrapper to False due to %s", cond)
def register_buffer(self, buffer: ir.ComputedBuffer):
name = f"buf{len(self.buffers)}"
self.buffers.append(buffer)
self.name_to_buffer[name] = buffer
return name
def realize_users_of(self, name: str):
"""
When a buffer is mutated we need to make sure all the reads to
the old version are realized before the mutation happens.
"""
assert isinstance(name, str)
def visit(value):
if isinstance(value, (list, tuple)):
return [visit(x) for x in value]
if isinstance(value, ir.IRNode):
if value.is_user_of(name):
value.realize()
return value
for key, value in self.env.items():
try:
visit(value)
except Exception:
log.warning("error in realize_users_of", exc_info=True)
def add_tensor_constant(self, data):
def allocate():
for name, value in self.constants.items():
if (
data.size() == value.size()
and data.stride() == value.stride()
and data.dtype == value.dtype
and data.device == value.device
and torch.eq(data, value).all()
):
return name
name = f"constant{len(self.constants)}"
self.constants[name] = data
return name
return TensorBox.create(
ir.ConstantBuffer(
allocate(),
FixedLayout(data.device, data.dtype, *self.static_sizes_strides(data)),
)
)
def constant_name(self, name: str, device_override: torch.device):
"""
We AOT copy constants to the devices they are needed on.
If device_override doesn't match the constant's device, then
copy it and return a different name.
"""
if self.constants[name].device == device_override or device_override is None:
return name
alt_name = f"{name}_{device_override.type}{device_override.index or 0}"
if alt_name not in self.constants:
self.constants[alt_name] = self.constants[name].to(device_override)
return alt_name
def placeholder(self, target: str, args, kwargs):
example = super().placeholder(target, args, kwargs)
if isinstance(example, SymTypes):
expr = example.node.expr
self.graph_inputs[target] = expr
return expr
elif isinstance(example, (int, bool, float)):
expr = sympy.sympify(example)
self.graph_inputs[target] = expr
return expr
assert isinstance(example, torch.Tensor), example
# todo(chilli): We can remove the last check once we turn buffers into
# static shape tensors. That's a hack to workaround Inductor believing
# the buffer should be static but us passing in a fake tensor with
# symbolic shapes.
if (
config.static_weight_shapes
and (
len(self.graph_inputs) < self.num_static_inputs
or not dynamo_config.dynamic_shapes
)
and not example._has_symbolic_sizes_strides
):
# the first N inputs are weights
sizes, strides = self.static_sizes_strides(example)
else:
sizes, strides = self.symbolic_sizes_strides(example)
# TODO(jansel): handle input aliasing
tensor = TensorBox.create(
InputBuffer(
target,
FixedLayout(example.device, example.dtype, sizes, strides),
)
)
self.graph_inputs[target] = tensor
self.graph_inputs_original[target] = tensor.data.data
self.device_types.add(example.device.type)
self.add_device_idx(example.device.index)
return tensor
def call_function(self, target, args, kwargs):
if target is operator.getitem and isinstance(args[0], (list, tuple)):
return super().call_function(target, args, kwargs)
if hasattr(target, "_inductor_lowering_function"):
# passthrough lowerings from .pattern_matcher
return target(*args, **kwargs)
if target not in lowerings:
base_name = target.name().split(".")[0]
if base_name in FALLBACK_ALLOW_LIST:
make_fallback(target)
elif config.implicit_fallbacks:
error = (
MissingOperatorWithDecomp
if get_decompositions([target])
else MissingOperatorWithoutDecomp
)
log.info(
"Creating implicit fallback for:\n%s",
error.operator_str(target, args, kwargs),
)
make_fallback(target)
elif get_decompositions([target]):
# There isn't a good way to dynamically patch this in
# since AOT Autograd already ran. The error message tells
# the user how to fix it.
raise MissingOperatorWithDecomp(target, args, kwargs)
else:
raise MissingOperatorWithoutDecomp(target, args, kwargs)
try:
out = lowerings[target](*args, **kwargs)
return out
except Exception as e:
raise LoweringException(e, target, args, kwargs).with_traceback(
e.__traceback__
) from None
def get_attr(self, target, args, kwargs):
# this is a constant
value = getattr(self.module, target)
if unsupported_output_tensor(value):
return self.add_tensor_constant(value)
with no_dispatch():
if value.shape == ():
return Constant(value.item(), value.dtype, value.device)
if len(value.shape) == 1 and value.shape[0] <= 8:
# tensor lowering has constant inlining logic
from .lowering import tensor
return tensor(value.tolist(), dtype=value.dtype, device=value.device)
return self.add_tensor_constant(value)
def call_module(self, target, args, kwargs):
raise AssertionError()
def call_method(self, target, args, kwargs):
raise AssertionError()
def output(self, target, args, kwargs):
result = super().output(target, args, kwargs)
assert isinstance(result, (tuple, list)), type(result)
assert all(
isinstance(
x,
(
TensorBox,
ir.Constant,
type(None),
ir.ConstantBuffer,
sympy.Expr,
sympy.Rel,
int,
),
)
for x in result
), result
self.graph_outputs = [ir.ExternKernel.realize_input(x) for x in result]
for name, value in self.graph_inputs.items():
assert isinstance(value, (TensorBox, sympy.Expr))
if not isinstance(value, TensorBox):
continue
value.realize()
assert isinstance(value, TensorBox)
value = value.data
assert isinstance(value, ir.StorageBox)
value_storage_box = value
value = value.data
if not isinstance(value, InputBuffer) or value.get_name() != name:
# one of our inputs was mutated, need to turn that into a copy
ir.MutationLayout.realize_into(value, self.graph_inputs_original[name])
# replace output with mutated input
try:
ind = self.graph_outputs.index(value_storage_box)
self.graph_outputs[ind] = self.graph_inputs_original[name]
except ValueError:
pass
self.finalize()
def finalize(self):
for buf in self.buffers:
buf.decide_layout()
def run_node(self, n: torch.fx.Node):
origins = {n}
if n.op == "call_function":
args, kwargs = self.fetch_args_kwargs_from_env(n)
origins |= gather_origins(args, kwargs)
with ir.IRNode.current_origins(origins):
if (
n.op == "call_function"
and n.target is not operator.getitem
and fallback_node_due_to_unsupported_type(n)
):
result = fallback_handler(n.target, add_to_fallback_set=False)(
*args, **kwargs
)
elif n.op == "call_function" and n.target in layout_constraints:
args, kwargs = layout_constraints[n.target](n, *args, **kwargs)
result = self.call_function(n.target, args, kwargs)
elif is_magic_method(n.target):
if isinstance(n.meta["val"], torch.SymInt):
result = n.meta["val"].node.expr
else:
result = super().run_node(n)
else:
result = super().run_node(n)
# require the same stride order for dense outputs,
# 1. user-land view() will not throw because inductor
# output different strides than eager
# long term the solution is to make view() always succeed
# with infallible strides.
# 2: as_strided ops, we need make sure its input has same size/stride with
# eager model to align with eager behavior.
as_strided_ops = [
torch.ops.aten.as_strided.default,
torch.ops.aten.as_strided_.default,
torch.ops.aten.as_strided_scatter.default,
]
if any(
user.op == "output" or user.target in as_strided_ops for user in n.users
) and isinstance(n.meta["val"], torch.Tensor):
strides = n.meta["val"].stride()
dense = torch._prims_common.is_non_overlapping_and_dense(n.meta["val"])
# requiring a stride order for a non-dense output wouldn't
# recreate the same strides, and would fail with view, defer for now.
if dense and len(strides):
result = ir.ExternKernel.require_stride_order(
result, ir.get_stride_order(strides)
)
# Realize if (1) any user need inputs realized, or (2) there is
# already too many reads and rematerializing can be bad.
num_users = len(set(n.users))
if num_users > 1 and isinstance(result, TensorBox):
for user in n.users:
if user.target in needs_realized_inputs:
result.realize_hint()
# This inclusion is somewhat controversial (from
# discussion between Horace, Natalia, and Elias).
# Currently, it's not very clear why this is helpful.
# The general idea here is that even though a node may
# have FlexibleLayout, we still often *treat* it as if
# it was contiguous. This appears to sometimes result in
# suboptimal behavior.
#
# When we do a better job selecting layout, we should
# revisit this.
need_fixed_layout = [
torch.ops.aten.convolution.default,
torch.ops.aten.convolution_backward.default,
torch.ops.aten.mm.default,
torch.ops.aten._int_mm.default,
]
if torch._C.has_mkldnn:
need_fixed_layout += [
torch.ops.mkldnn._convolution_pointwise.default,
torch.ops.mkldnn._convolution_pointwise.binary,
torch.ops.mkldnn._convolution_pointwise_.binary,
torch.ops.mkldnn._convolution_transpose_pointwise.default,
torch.ops.mkldnn._linear_pointwise.default,
torch.ops.mkldnn._linear_pointwise.binary,
]
if torch._C.has_mkl:
need_fixed_layout += [torch.ops.mkl._mkl_linear.default]
if user.target in need_fixed_layout:
result = ir.ExternKernel.require_stride_order(
result, ir.get_stride_order(n.meta["val"].stride())
)
if user.op == "output":
if isinstance(result.data.data, (Pointwise, Reduction)):
result.realize()
# TODO(jansel): introduce a store vs inline choice
result.mark_reuse(len(n.users))
# Realize if the IRNode already has accumulated lots of reads
if isinstance(result, TensorBox) and result.has_exceeded_max_reads():
# Prevent excessive accumulation in a computed buffer, when
# there are multiple branches each with small number of memory
# reads, but they converge to a user.
result.realize_hint()
# This is not complete, but it doesn't have to be: origin_node
# tracking is best effort. The logic here critically relies on direct
# TensorBox -> StorageBox denoting a non-view; we don't bother trying
# to get views to work. Feel free to add any extra cases as needed.
#
# Note: we can't YOLO tree_map over this result, because if there are
# buffers or a view involved, we might not be able to validly assign
# the origin_node here.
if isinstance(result, TensorBox) and isinstance(result.data, ir.StorageBox):
if isinstance(result.data.data, ir.Loops):
result.data.data.origin_node = n
elif isinstance(result.data.data, ir.Buffer):
result.data.data.origin_node = n
if isinstance(result.data.data, ir.ComputedBuffer) and isinstance(
result.data.data.data, ir.Loops
):
result.data.data.data.origin_node = n
# Not really multi-output, can straightforwardly recurse in
elif (
isinstance(result.data.data, ir.MultiOutput)
and not result.data.data.indices
):
if isinstance(result.data.data.inputs[0], ir.Buffer):
result.data.data.inputs[0].origin_node = n
return result
def check_cpp_codegen_disabled(self):
if config.disable_cpp_codegen:
self.disable_cpp_wrapper("cpp codegen disabled")
def check_platform(self):
if sys.platform != "linux":
self.disable_cpp_wrapper("platform not linux")
def check_input_for_cpp_buffer(self):
for _, value in self.graph_inputs.items():
dtype = None
if isinstance(value, TensorBox):
dtype = value.get_dtype()
elif isinstance(value, (sympy.Symbol, sympy.core.numbers.Integer)):
dtype = may_get_constant_buffer_dtype(value)
if not supported_dtype_of_cpp_wrapper(dtype, self.cuda):
self.disable_cpp_wrapper("unsupported inputs dtype")
def check_constant_for_cpp_buffer(self):
if self.constants:
self.disable_cpp_wrapper("Constants")
def check_cpp_wrapper(self):
self.check_cpp_codegen_disabled()
self.check_platform()
self.check_input_for_cpp_buffer()
self.check_constant_for_cpp_buffer()
def init_wrapper_code(self):
self.cuda = "cuda" in self.device_types
if self.cpp_wrapper:
self.check_cpp_wrapper()
# Re-check self.cpp_wrapper because it might be disabled due to failed checking
if self.cuda:
assert self.cpp_wrapper, "CudaWrapperCodeGen hit unsupported case"
if self.cpp_wrapper:
self.wrapper_code = (
CudaWrapperCodeGen() if self.cuda else CppWrapperCodeGen()
)
return
self.wrapper_code = WrapperCodeGen()
def codegen(self):
from .scheduler import Scheduler
self.init_wrapper_code()
self.scheduler = Scheduler(self.buffers)
assert self.scheduler is not None # mypy can't figure this out
self.scheduler.codegen()
assert self.wrapper_code is not None
return self.wrapper_code.generate()
def count_bytes(self):
from .scheduler import FusedSchedulerNode, NopKernelSchedulerNode, Scheduler
scheduler = Scheduler(self.buffers)
def get_read_write_buffers_sizes(node):
if isinstance(node, NopKernelSchedulerNode):
return 0
reads = {dep.name for dep in node.read_writes.reads}
writes = {dep.name for dep in node.read_writes.writes}
def is_materialized(buf):
buf_uses = {user.node for user in scheduler.name_to_node[buf].users}
return len(buf_uses - set(node.snodes)) > 0
if isinstance(node, FusedSchedulerNode):
removed_buffers = {dep for dep in writes if not is_materialized(dep)}
writes = writes - removed_buffers
reads = reads - removed_buffers
node_bytes = 0
for buf in reads | writes:
if buf in self.name_to_buffer:
buf = self.name_to_buffer[buf]
elif buf in self.graph_inputs:
buf = self.graph_inputs[buf]
else:
continue
node_bytes += V.graph.sizevars.size_hint(
sympy_product(buf.get_size())
) * get_dtype_size(buf.get_dtype())
return node_bytes
total_bytes = 0
node_counts = []
for node in scheduler.nodes:
num_bytes = get_read_write_buffers_sizes(node)
node_counts.append((node, num_bytes // 4))
total_bytes += num_bytes
return total_bytes, node_counts
@dynamo_timed
def compile_to_module(self):
from .codecache import PyCodeCache
code, linemap = self.codegen()
mod = PyCodeCache.load(code, linemap=linemap)
for name, value in self.constants.items():
setattr(mod, name, value)
log.debug("Output code written to: %s", mod.__file__)
output_code_log.debug("Output code: \n%s", code)
if config.benchmark_kernel:
print(f"Compiled module path: {mod.__file__}", file=sys.stderr)
V.debug.output_code(mod.__file__)
V.debug.rename(os.path.splitext(mod.__file__)[0] + ".debug")
return mod
def compile_to_fn(self):
if self.aot_mode:
from .codecache import AotCodeCache
code, linemap = self.codegen()
output_code_log.debug("Output code: \n%s", code)
return AotCodeCache.compile(self, code, cuda=self.cuda)
else:
return self.compile_to_module().call
def get_output_names(self):
assert self.graph_outputs is not None
return [
node.get_name()
for node in self.graph_outputs
if not isinstance(node, ir.NoneAsConstantBuffer)
and not isinstance(node, ir.ShapeAsConstantBuffer)
]
def is_unspec_arg(self, name: str):
# dynamo wraps unspec variable as 0d CPU tensor,
# need to convert to scalar during codegen (triton only)
return (
name in self.graph_inputs.keys()
and self.graph_inputs[name].get_numel() == 1
and self.graph_inputs[name].get_device().type == "cpu"
)