Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,10 @@
|
||||
"""torchgen
|
||||
|
||||
This module contains codegeneration utilities for PyTorch. It is used to
|
||||
build PyTorch from source, but may also be used for out-of-tree projects
|
||||
that extend PyTorch.
|
||||
|
||||
Note well that we provide no BC guarantees for torchgen. If you're interested
|
||||
in using torchgen and want the PyTorch team to be aware, please reach out
|
||||
on GitHub.
|
||||
"""
|
||||
@@ -0,0 +1,206 @@
|
||||
# Be extra careful when you edit this file, because it affects AOTInductor ABI compatibility. See
|
||||
# https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436
|
||||
# for details.
|
||||
#
|
||||
# The inductor_fallback_ops list is based on the fallback ops from torch/_inductor/lowering.py.
|
||||
#
|
||||
# Generally speaking, it is ok to add a new op to the list, but you need to run
|
||||
# `python torchgen/gen.py --update-aoti-c-shim` in order to regenerate C shim header files.
|
||||
# But it is NOT ok to remove an existing fallback op from the list, since that will break
|
||||
# some existing AOTInductor-compiled models.
|
||||
#
|
||||
# A fallback op version defaults to 1. If you want to extend an existing fallback op by adding
|
||||
# a new argument with a default value, while it is fine in the Python world, it will be BC-breaking
|
||||
# when generating C shim. Thus you need to bump up the version number of that fallback op by
|
||||
# updating the entry in the inductor_fallback_ops list, adding a new version number with a list
|
||||
# of new arguments, and then run `python torchgen/gen.py --update-aoti-c-shim` to regenerate.
|
||||
|
||||
inductor_fallback_ops: dict[str, dict[str, list[str]]] = {
|
||||
"aten._adaptive_avg_pool2d_backward.default": {},
|
||||
"aten._adaptive_avg_pool2d.default": {},
|
||||
"aten._adaptive_avg_pool3d_backward.default": {},
|
||||
"aten._adaptive_avg_pool3d.default": {},
|
||||
"aten._addmm_activation.default": {},
|
||||
"aten._cdist_backward.default": {},
|
||||
"aten._cdist_forward.default": {},
|
||||
"aten._cudnn_rnn.default": {},
|
||||
"aten._dyn_quant_matmul_4bit.default": {},
|
||||
"aten._dyn_quant_pack_4bit_weight.default": {},
|
||||
"aten._efficient_attention_backward.default": {},
|
||||
"aten._efficient_attention_forward.default": {},
|
||||
"aten._efficientzerotensor.default": {},
|
||||
"aten._embedding_bag_dense_backward.default": {},
|
||||
"aten._embedding_bag_forward_only.default": {},
|
||||
"aten._embedding_bag_per_sample_weights_backward.default": {},
|
||||
"aten._embedding_bag.default": {},
|
||||
"aten._fft_c2c.default": {},
|
||||
"aten._fft_r2c.default": {},
|
||||
"aten._flash_attention_backward.default": {},
|
||||
"aten._flash_attention_forward.default": {"v2": ["block_table", "num_splits"]},
|
||||
"aten._flash_attention_forward_no_dropout_inplace.default": {"v2": ["num_splits"]},
|
||||
"aten._flash_attention_forward.quantized": {},
|
||||
"aten._fused_moving_avg_obs_fq_helper_functional.default": {},
|
||||
"aten._fused_moving_avg_obs_fq_helper.default": {},
|
||||
"aten._fused_rms_norm.default": {},
|
||||
"aten._grouped_mm.default": {},
|
||||
"aten._histogramdd_from_bin_cts.default": {},
|
||||
"aten._int_mm.out": {},
|
||||
"aten._pdist_backward.default": {},
|
||||
"aten._pdist_forward.default": {},
|
||||
"aten._scaled_dot_product_attention_math_for_mps.default": {"v2": ["enable_gqa"]},
|
||||
"aten._scaled_dot_product_cudnn_attention_backward.default": {},
|
||||
"aten._scaled_dot_product_cudnn_attention.default": {},
|
||||
"aten._scaled_dot_product_efficient_attention_backward.default": {},
|
||||
"aten._scaled_dot_product_efficient_attention.default": {},
|
||||
"aten._scaled_dot_product_flash_attention_backward.default": {},
|
||||
"aten._scaled_dot_product_flash_attention_for_cpu_backward.default": {},
|
||||
"aten._scaled_dot_product_flash_attention_for_cpu.default": {},
|
||||
"aten._scaled_dot_product_flash_attention.default": {},
|
||||
"aten._scaled_dot_product_flash_attention.quantized": {},
|
||||
"aten._scaled_dot_product_fused_attention_overrideable_backward.default": {},
|
||||
"aten._scaled_dot_product_fused_attention_overrideable.default": {},
|
||||
"aten._scaled_mm.default": {},
|
||||
"aten._scaled_grouped_mm.default": {},
|
||||
"aten._scaled_mm.out": {},
|
||||
"aten._segment_reduce_backward.default": {},
|
||||
"aten._thnn_fused_lstm_cell.default": {},
|
||||
"aten._to_sparse.default": {},
|
||||
"aten._trilinear.default": {},
|
||||
"aten._weight_int4pack_mm.default": {},
|
||||
"aten._weight_int8pack_mm.default": {},
|
||||
"aten.abs.default": {},
|
||||
"aten.adaptive_max_pool2d_backward.default": {},
|
||||
"aten.adaptive_max_pool2d.default": {},
|
||||
"aten.adaptive_max_pool3d_backward.default": {},
|
||||
"aten.adaptive_max_pool3d.default": {},
|
||||
"aten.add.Scalar": {},
|
||||
"aten.add.Tensor": {},
|
||||
"aten.addbmm.default": {},
|
||||
"aten.addmm.out": {},
|
||||
"aten.addmv.default": {},
|
||||
"aten.angle.default": {},
|
||||
"aten.avg_pool2d_backward.default": {},
|
||||
"aten.avg_pool2d.default": {},
|
||||
"aten.avg_pool3d_backward.default": {},
|
||||
"aten.avg_pool3d.default": {},
|
||||
"aten.baddbmm.out": {},
|
||||
"aten.bernoulli_.float": {},
|
||||
"aten.bernoulli_.Tensor": {},
|
||||
"aten.bmm.out": {},
|
||||
"aten.bucketize.Tensor": {},
|
||||
"aten.cat.default": {},
|
||||
"aten.cholesky_inverse.default": {},
|
||||
"aten.cholesky_solve.default": {},
|
||||
"aten.convolution_backward.default": {},
|
||||
"aten.convolution.default": {},
|
||||
"aten.cummax.default": {},
|
||||
"aten.cummin.default": {},
|
||||
"aten.cumprod.default": {},
|
||||
"aten.cumsum.default": {},
|
||||
"aten.exponential.default": {},
|
||||
"aten.fill_.Scalar": {},
|
||||
"aten.fractional_max_pool2d_backward.default": {},
|
||||
"aten.fractional_max_pool2d.default": {},
|
||||
"aten.fractional_max_pool3d_backward.default": {},
|
||||
"aten.fractional_max_pool3d.default": {},
|
||||
"aten.gcd.default": {},
|
||||
"aten.geqrf.default": {},
|
||||
"aten.grid_sampler_2d_backward.default": {},
|
||||
"aten.hann_window.default": {},
|
||||
"aten.histc.default": {},
|
||||
"aten.histogram.bin_ct": {},
|
||||
"aten.index_put.default": {},
|
||||
"aten.index_reduce.default": {},
|
||||
"aten.index.Tensor": {},
|
||||
"aten.kthvalue.default": {},
|
||||
"aten.logcumsumexp.default": {},
|
||||
"aten.lu_unpack.default": {},
|
||||
"aten.masked_scatter_backward.default": {},
|
||||
"aten.masked_scatter.default": {},
|
||||
"aten.masked_select.default": {},
|
||||
"aten.max_pool2d_with_indices_backward.default": {},
|
||||
"aten.max_pool2d_with_indices.default": {},
|
||||
"aten.max_pool3d_with_indices_backward.default": {},
|
||||
"aten.max_pool3d_with_indices.default": {},
|
||||
"aten.max_unpool2d.default": {},
|
||||
"aten.max_unpool3d.default": {},
|
||||
"aten.median.default": {},
|
||||
"aten.mm.out": {},
|
||||
"aten.mm.dtype_out": {},
|
||||
"aten.mode.default": {},
|
||||
"aten.mul.Scalar": {},
|
||||
"aten.mul.Tensor": {},
|
||||
"aten.nanmedian.default": {},
|
||||
"aten.narrow.default": {},
|
||||
"aten.native_dropout.default": {},
|
||||
"aten.nonzero.default": {},
|
||||
"aten.nonzero_static.default": {},
|
||||
"aten.normal_functional.default": {},
|
||||
"aten.ormqr.default": {},
|
||||
"aten.pad.default": {},
|
||||
"aten.permute.default": {},
|
||||
"aten.polar.default": {},
|
||||
"aten.pow.Scalar": {},
|
||||
"aten.pow.Tensor_Scalar": {},
|
||||
"aten.pow.Tensor_Tensor": {},
|
||||
"aten.rand.default": {},
|
||||
"aten.rand.generator": {},
|
||||
"aten.randint.default": {},
|
||||
"aten.randint.generator": {},
|
||||
"aten.randint.low_out": {},
|
||||
"aten.randint.low": {},
|
||||
"aten.randn.default": {},
|
||||
"aten.randn.generator": {},
|
||||
"aten.randperm.default": {},
|
||||
"aten.rand_like.default": {},
|
||||
"aten.rand_like.generator": {},
|
||||
"aten.randint_like.default": {},
|
||||
"aten.randint_like.low_dtype": {},
|
||||
"aten.randn_like.default": {},
|
||||
"aten.randn_like.generator": {},
|
||||
"aten.repeat_interleave.Tensor": {},
|
||||
"aten.replication_pad1d_backward.default": {},
|
||||
"aten.replication_pad2d_backward.default": {},
|
||||
"aten.reshape.default": {},
|
||||
"aten.resize_.default": {},
|
||||
"aten.resize_as_.default": {},
|
||||
"aten.scatter_reduce.two_out": {},
|
||||
"aten.scatter.src_out": {},
|
||||
"aten.scatter.value_out": {},
|
||||
"aten.searchsorted.Scalar": {},
|
||||
"aten.searchsorted.Tensor": {},
|
||||
"aten.segment_reduce.default": {},
|
||||
"aten.set_.source_Tensor": {},
|
||||
"aten.slice.Tensor": {},
|
||||
"aten.soft_margin_loss_backward.default": {},
|
||||
"aten.sort.default": {},
|
||||
"aten.sort.stable": {},
|
||||
"aten.squeeze.dim": {},
|
||||
"aten.to_sparse.default": {},
|
||||
"aten.topk.default": {},
|
||||
"aten.triangular_solve.default": {},
|
||||
"aten.uniform.default": {},
|
||||
"aten.upsample_bicubic2d_backward.default": {},
|
||||
"aten.upsample_linear1d_backward.default": {},
|
||||
"aten.upsample_trilinear3d_backward.default": {},
|
||||
"aten.view_as_complex.default": {},
|
||||
"aten.view_as_real.default": {},
|
||||
"aten.view.dtype": {},
|
||||
"aten._weight_int4pack_mm_with_scales_and_zeros.default": {},
|
||||
}
|
||||
|
||||
# `python torchgen/gen.py --update-aoti-c-shim` will automatically generate
|
||||
# c_shim_aten.{h/cpp} based on the list below.
|
||||
# Operators in this list are intended to be used in torch/csrc/stable/ops.h
|
||||
# Unlike other c_shims, operators in this file do not bypass the dispatcher.
|
||||
# The same BC rules apply as inductor_fallback_ops.
|
||||
aten_shimified_ops: dict[str, dict[str, list[str]]] = {
|
||||
"aten.fill_.Scalar": {},
|
||||
"aten.pad.default": {},
|
||||
"aten.narrow.default": {},
|
||||
"aten.amax.default": {},
|
||||
"aten.new_empty.default": {},
|
||||
"aten.new_zeros.default": {},
|
||||
"aten.full.default": {},
|
||||
"aten.subtract.Tensor": {},
|
||||
}
|
||||
@@ -0,0 +1,892 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import cast, TYPE_CHECKING
|
||||
|
||||
from torchgen import local
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import BaseCType, Binding, NamedCType, tensorListT
|
||||
from torchgen.model import (
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
NativeFunction,
|
||||
NativeFunctionsViewGroup,
|
||||
SchemaKind,
|
||||
Type,
|
||||
)
|
||||
from torchgen.utils import IDENT_REGEX
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# Represents a saved attribute involved in backward calculation.
|
||||
# Note that it can be a derived property of an input argument, e.g.:
|
||||
# we could save `other.scalar_type()` instead of the entire `other` tensor.
|
||||
@dataclass(frozen=True)
|
||||
class SavedAttribute:
|
||||
# The NamedCType holds the updated name and cpp type of the attribute
|
||||
# for the name, Suffix is appended if it's derived property, e.g.: `other_scalar_type`
|
||||
nctype: NamedCType
|
||||
|
||||
# The expression to read the derived property at save time, e.g.:
|
||||
# `other.scalar_type()`.
|
||||
expr: str
|
||||
|
||||
|
||||
# Represents a backward formula that calculates derivatives for one
|
||||
# or more tensors.
|
||||
@dataclass(frozen=True)
|
||||
class Derivative:
|
||||
# The formula string (legit C++ expression).
|
||||
# Note that expressions against input arguments have been replaced with the
|
||||
# corresponding saved attributes.
|
||||
# E.g.:
|
||||
# raw formula: `mul_tensor_backward(grad, self, other.scalar_type())`
|
||||
# here: `mul_tensor_backward(grad, self, other_scalar_type)`
|
||||
formula: str
|
||||
|
||||
# The formula string before input argument replacement
|
||||
original_formula: str
|
||||
|
||||
# Names of the arguments for which this formula calculates derivatives.
|
||||
var_names: tuple[str, ...]
|
||||
|
||||
# Saved inputs that are referenced by the formula.
|
||||
saved_inputs: tuple[SavedAttribute, ...]
|
||||
|
||||
# Saved outputs that are referenced by the formula.
|
||||
saved_outputs: tuple[SavedAttribute, ...]
|
||||
|
||||
# Gradients that are referenced by name in the formula.
|
||||
named_gradients: set[str]
|
||||
|
||||
|
||||
# Represents a forward formula that calculates forward derivatives
|
||||
# for one tensor.
|
||||
@dataclass(frozen=True)
|
||||
class ForwardDerivative:
|
||||
# The formula string (legit C++ expression).
|
||||
# Note that special keywords such as "linear" or "element_wise" have been
|
||||
# replaced by the automatically generated formula.
|
||||
formula: str
|
||||
|
||||
# Name of the output arguments for which this formula calculates forward
|
||||
# derivatives
|
||||
var_names: tuple[str, ...]
|
||||
|
||||
# Type of the output arguments for which this formula calculates forward
|
||||
# derivatives
|
||||
var_types: tuple[Type, ...]
|
||||
|
||||
# Inputs for which the forward derivatives are required for this formula
|
||||
required_inputs_fw_grad: tuple[str, ...] | None
|
||||
|
||||
# Inputs for which the primal is required for this formula
|
||||
required_inputs_primal: tuple[str, ...] | None
|
||||
|
||||
# Flag to specify if this formula requires the original value of self
|
||||
# This is only used by inplace operations
|
||||
required_original_self_value: bool
|
||||
|
||||
# If this formula is specified in derivatives.yaml or if we are reusing the
|
||||
# out of place formula for inplace
|
||||
is_reusing_outplace_formula: bool
|
||||
|
||||
|
||||
# Represents differentiability info for a NativeFunction.
|
||||
@dataclass(frozen=True)
|
||||
class DifferentiabilityInfo:
|
||||
# The base name read from derivatives.yaml.
|
||||
name: str
|
||||
|
||||
# The matching native function.
|
||||
#
|
||||
# There can be multiple NativeFunction having the same base name:
|
||||
# - different overloads with different types of input arguments;
|
||||
# - in-place/out/functional variants of the same function;
|
||||
#
|
||||
# We first use the schema string (under the 'name' key) in derivatives.yaml
|
||||
# to find the NativeFunction having the same schema string.
|
||||
# Then we find the in-place/out/functional variants of the matching function.
|
||||
# Among these variants, we choose the one having the same name as the
|
||||
# derivatives.yaml entry. If there is no exact match, then we choose the
|
||||
# in-place variant.
|
||||
# TODO: maybe the logic to search for all variants is no longer necessary?
|
||||
func: NativeFunction
|
||||
|
||||
# The name of the generated autograd function.
|
||||
# It's set only if we will calculate a derivative, i.e.
|
||||
# 'args_with_derivatives' is not empty.
|
||||
op: str | None
|
||||
|
||||
# The derivatives formulae for this function.
|
||||
# Note that the length of this sequence is the number of differentiable inputs
|
||||
derivatives: Sequence[Derivative]
|
||||
|
||||
# The forward derivatives formulae for this function.
|
||||
# Note that the length of this sequence is the number of differentiable outputs
|
||||
forward_derivatives: Sequence[ForwardDerivative]
|
||||
|
||||
# The union of 'saved_inputs' of all 'derivatives'.
|
||||
all_saved_inputs: Sequence[SavedAttribute]
|
||||
|
||||
# The union of 'saved_outputs' of all 'derivatives'.
|
||||
all_saved_outputs: Sequence[SavedAttribute]
|
||||
|
||||
# All named gradients that are available for use, in the same
|
||||
# order as in the grads vector.
|
||||
available_named_gradients: Sequence[str]
|
||||
|
||||
# The named gradients that are used in any of the derivatives.
|
||||
# Invariant: all(name in available_named_gradients for name in used_named_gradients)
|
||||
used_named_gradients: set[str]
|
||||
|
||||
# The function's input arguments for which it calculates derivatives.
|
||||
# It's the union of 'var_names' of all 'derivatives', sorted by the
|
||||
# argument order in the function schema.
|
||||
args_with_derivatives: Sequence[Binding]
|
||||
|
||||
# Names of arguments whose derivative formula is 'non_differentiable'.
|
||||
non_differentiable_arg_names: Sequence[str]
|
||||
|
||||
# Raw data read from derivatives.yaml.
|
||||
output_differentiability: list[bool] | None
|
||||
|
||||
# output_differentiability in derivatives.yaml can be a list of
|
||||
# conditions that express if the output is differentiable. In this case,
|
||||
# the number of conditions must match the number of outputs
|
||||
# (NB: we only support one condition right now).
|
||||
# output_differentiability gets populated with True for each condition,
|
||||
# while output_differentiability_conditions gets populated with the conditions
|
||||
output_differentiability_conditions: list[str] | None
|
||||
|
||||
@property
|
||||
def has_derivatives(self) -> bool:
|
||||
return len(self.args_with_derivatives) > 0
|
||||
|
||||
# Generates a new DifferentiabilityInfo using the exact same set of derivative information,
|
||||
# but with a new operator name.
|
||||
# This is used when generating "copy" variants of view ops,
|
||||
# which are able to use the exact same derivative formula as the original view op
|
||||
# See Note [Codegen'd {view}_copy Operators]
|
||||
def create_view_copy_from_view_derivative(
|
||||
self, g: NativeFunctionsViewGroup
|
||||
) -> DifferentiabilityInfo | None:
|
||||
if g.view_copy is None:
|
||||
return None
|
||||
f = g.view_copy
|
||||
|
||||
name_split_by_period = self.name.split(".", maxsplit=2)
|
||||
# Append a "_copy" to the base name of the operator (but keep the overload name the same)
|
||||
view_copy_name = f"{name_split_by_period[0]}_copy." + ".".join(
|
||||
name_split_by_period[1:]
|
||||
)
|
||||
view_copy_op_name = None if self.op is None else f"{self.op}_copy"
|
||||
|
||||
return DifferentiabilityInfo(
|
||||
# Use the "_copy" version of name/func/op
|
||||
name=view_copy_name,
|
||||
func=f,
|
||||
op=view_copy_op_name,
|
||||
# But keep all derivative info the same
|
||||
derivatives=self.derivatives,
|
||||
forward_derivatives=self.forward_derivatives,
|
||||
all_saved_inputs=self.all_saved_inputs,
|
||||
all_saved_outputs=self.all_saved_outputs,
|
||||
available_named_gradients=self.available_named_gradients,
|
||||
used_named_gradients=self.used_named_gradients,
|
||||
args_with_derivatives=self.args_with_derivatives,
|
||||
non_differentiable_arg_names=self.non_differentiable_arg_names,
|
||||
output_differentiability=self.output_differentiability,
|
||||
output_differentiability_conditions=self.output_differentiability_conditions,
|
||||
)
|
||||
|
||||
|
||||
def uses_ident(info: DifferentiabilityInfo | None, ident: str) -> bool:
|
||||
if info is None:
|
||||
return False
|
||||
for derivative in info.derivatives:
|
||||
formula = derivative.formula
|
||||
if re.search(IDENT_REGEX.format(ident), formula):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def uses_retain_variables(info: DifferentiabilityInfo | None) -> bool:
|
||||
return uses_ident(info, "retain_variables")
|
||||
|
||||
|
||||
def uses_single_grad(info: DifferentiabilityInfo | None) -> bool:
|
||||
return uses_ident(info, "grad")
|
||||
|
||||
|
||||
# Represents a differentiable `Argument`.
|
||||
# How is it different from the `Argument` type?
|
||||
# - It's processed Arguments which are differentiable and only used in the
|
||||
# context of the autograd codegen;
|
||||
# - It can represent SelfArgument or regular Argument but not TensorOptionsArgument;
|
||||
@dataclass(frozen=True)
|
||||
class DifferentiableInput:
|
||||
name: str
|
||||
type: Type
|
||||
|
||||
# TODO: only to keep it byte-for-byte compatible with the old codegen, should remove.
|
||||
cpp_type: str
|
||||
|
||||
|
||||
# Represents a differentiable `Return`.
|
||||
# How it it different from the `Return` type?
|
||||
# - The name in `Return` is optional. Here it is always populated using the same
|
||||
# `cpp.return_names()` method.
|
||||
# TODO: some cpp naming logic (e.g. resolving name conflict) might be irrelevant?
|
||||
# - It's processed Returns which are differentiable, in compliance with the
|
||||
# `output_differentiability` field defined in derivatives.yaml (if specified),
|
||||
# and are only used in the context of the autograd codegen;
|
||||
@dataclass(frozen=True)
|
||||
class DifferentiableOutput:
|
||||
name: str
|
||||
type: Type
|
||||
|
||||
# TODO: only to keep it byte-for-byte compatible with the old codegen, should remove.
|
||||
cpp_type: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NativeFunctionWithDifferentiabilityInfo:
|
||||
func: NativeFunction
|
||||
info: dict[str, DifferentiabilityInfo] | None
|
||||
fw_derivatives: dict[str, Sequence[ForwardDerivative]] | None
|
||||
|
||||
|
||||
# TODO: Update comment below since it is out of date.
|
||||
def dispatch_strategy(fn: NativeFunctionWithDifferentiabilityInfo) -> str:
|
||||
"""How are we going to call the underlying implementation of a
|
||||
declaration? There are two strategies:
|
||||
- use_derived: we want to call the implementation on CPUDoubleType
|
||||
(or a similar, derived Type instance). Because these derived
|
||||
instances deal in Tensors, not Variables (it's a completely different
|
||||
object, so it doesn't dispatch back to VariableType), code on
|
||||
this dispatch path needs to wrap/unwrap tensors. If the
|
||||
derived implementation takes and returns tensors, the
|
||||
implementation is usually differentiable (although we also use
|
||||
the derived dispatch path for non-differentiable functions
|
||||
that we still want to dispatch on the derived Type instance;
|
||||
e.g., size())
|
||||
- use_type: we want to call the implementation on Type, because
|
||||
it is implemented concretely, and the functions it invokes will
|
||||
get dispatched back to VariableType (which will ensure that they
|
||||
are differentiable.)
|
||||
"""
|
||||
# fn is derived as long as any of its per-key differentiability infos
|
||||
# has_derivatives. dispatch_strategy() is used to guard generation of fns in VariableType
|
||||
# and ADInplaceOrViewType. We want to generate these functions as long as a
|
||||
# derivative is defined for ANY dispatch key.
|
||||
if fn.func.is_abstract or (
|
||||
fn.info is not None and any(info.has_derivatives for info in fn.info.values())
|
||||
):
|
||||
# If the function is abstract (not implemented on at::Type), we must
|
||||
# call the implementation on the derived type with unpacked tensors.
|
||||
|
||||
# If the function has a derivative specified and is concrete, we could
|
||||
# call either implementation. We prefer the calling the derived
|
||||
# type's implementation with unpacked tensors because it is more
|
||||
# performant in some cases: any internal calls to other ATen functions
|
||||
# won't have the history tracked.
|
||||
|
||||
# If the function has a type dispatched argument (i.e. is a factory),
|
||||
# we prefer calling the derived type's implementation both because it is
|
||||
# more performant and to ensure factory functions return tensors with _version
|
||||
# of 0 (probably not strictly necessary, but nice to have to keeps versions simple
|
||||
# to understand.
|
||||
|
||||
return "use_derived"
|
||||
else:
|
||||
# If the function is concrete (we don't have to override it) and we
|
||||
# didn't declare it in derivatives.yaml, we'll assume that it is
|
||||
# actually implemented out of differentiable functions. (This
|
||||
# assumption might not hold, but then you'll see gradcheck fail.)
|
||||
return "use_type"
|
||||
|
||||
|
||||
def is_foreach_func(f: NativeFunction) -> bool:
|
||||
return f.func.name.name.base.startswith("_foreach_")
|
||||
|
||||
|
||||
# note(crcrpar): Most foreach functions can reference an out-place `torch` function whose schema kind
|
||||
# is functional for their backward derivatives (and forward derivatives in the future), i.e.,
|
||||
# they would find such one in `functional_info_by_signature`. There however are some exceptions:
|
||||
_foreach_with_inplace_ref = {"_foreach_zero_"}
|
||||
_foreach_with_tensor_overload = {
|
||||
"_foreach_add.Tensor",
|
||||
"_foreach_mul.Tensor",
|
||||
"_foreach_div.Tensor",
|
||||
}
|
||||
# The following do not support the alpha kwarg, which the nonforeach versions support.
|
||||
_skip_argument_len_check = {
|
||||
"_foreach_add.Scalar",
|
||||
"_foreach_add_.Scalar",
|
||||
"_foreach_add.ScalarList",
|
||||
"_foreach_add_.ScalarList",
|
||||
"_foreach_sub.Scalar",
|
||||
"_foreach_sub_.Scalar",
|
||||
"_foreach_sub.ScalarList",
|
||||
"_foreach_sub_.ScalarList",
|
||||
}
|
||||
|
||||
|
||||
# Checks if `function_schema` is a native, non-foreach function which `f`, a foreach function
|
||||
# reference to generate derivatives.
|
||||
def is_reference_for_foreach(
|
||||
f: NativeFunction,
|
||||
function_schema: FunctionSchema,
|
||||
) -> bool:
|
||||
return (
|
||||
f.func.name.name.base.split("_foreach_")[-1] == function_schema.name.name.base
|
||||
and (
|
||||
not function_schema.name.name.inplace
|
||||
or str(f.func.name) in _foreach_with_inplace_ref
|
||||
)
|
||||
and (
|
||||
str(f.func.name) in _skip_argument_len_check
|
||||
or len(f.func.arguments.flat_non_out)
|
||||
== len(function_schema.arguments.flat_non_out)
|
||||
)
|
||||
and all(
|
||||
ref_arg.type in (arg.type, getattr(arg.type, "elem", None))
|
||||
for arg, ref_arg in zip(
|
||||
f.func.arguments.flat_non_out,
|
||||
function_schema.arguments.flat_non_out,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# TODO(crcrpar): Avoid hard coding "Default" ideally.
|
||||
def gen_foreach_derivativeinfo(
|
||||
foreach_function: NativeFunction,
|
||||
functional_info_by_signature: dict[
|
||||
FunctionSchema, dict[str, DifferentiabilityInfo]
|
||||
],
|
||||
non_functional_info_by_signature: dict[
|
||||
FunctionSchema, dict[str, DifferentiabilityInfo]
|
||||
],
|
||||
dispatch_key: str = "Default",
|
||||
) -> tuple[DifferentiabilityInfo | None, bool]:
|
||||
"""Generate DifferentiabilityInfo for out-place foreach function, return the existing one for in-place.
|
||||
|
||||
The second return value indicates whether the info is generated in this function.
|
||||
"""
|
||||
ref_diff_info: DifferentiabilityInfo | None = None
|
||||
|
||||
for function_schema, diff_info in functional_info_by_signature.items():
|
||||
if not is_reference_for_foreach(foreach_function, function_schema):
|
||||
continue
|
||||
ref_diff_info = diff_info[dispatch_key]
|
||||
if ref_diff_info is not None:
|
||||
break
|
||||
# note(crcrpar): It seems like `zero`'s info isn't available in functional_info_by_signature
|
||||
# while the info of `zero_` is in non_functional_info_by_signature
|
||||
if (
|
||||
ref_diff_info is None
|
||||
and foreach_function.func.kind() == SchemaKind.inplace
|
||||
and str(foreach_function.func.name) in _foreach_with_inplace_ref
|
||||
):
|
||||
for function_schema, diff_info in non_functional_info_by_signature.items():
|
||||
if not is_reference_for_foreach(foreach_function, function_schema):
|
||||
continue
|
||||
ref_diff_info = diff_info[dispatch_key]
|
||||
if ref_diff_info is not None:
|
||||
break
|
||||
if ref_diff_info is None:
|
||||
return None, False
|
||||
|
||||
# non out-place uses the existing Derivative.
|
||||
if foreach_function.func.kind() == SchemaKind.inplace:
|
||||
return ref_diff_info, False
|
||||
|
||||
map_refarg2foreacharg, map_name2arg = {}, {}
|
||||
for i, (arg, ref_arg) in enumerate(
|
||||
zip(
|
||||
foreach_function.func.arguments.flat_non_out,
|
||||
function_schema.arguments.flat_non_out,
|
||||
)
|
||||
):
|
||||
map_refarg2foreacharg[ref_arg.name] = arg.name
|
||||
map_name2arg[arg.name] = arg
|
||||
|
||||
all_saved_inputs, all_saved_outputs, all_var_names = [], [], []
|
||||
modified_derivative_formulas = []
|
||||
for i, derivative in enumerate(ref_diff_info.derivatives):
|
||||
modified_formula = derivative.formula.replace("grad", "grads[i]").replace(
|
||||
"result", "result[i]"
|
||||
)
|
||||
saved_inputs, saved_outputs = [], []
|
||||
# note(crcrpar): This context seems necessary to call `cpp.argument_type`
|
||||
with local.parametrize(
|
||||
use_const_ref_for_mutable_tensors=foreach_function.use_const_ref_for_mutable_tensors,
|
||||
use_ilistref_for_tensor_lists=foreach_function.part_of_structured_group,
|
||||
):
|
||||
for ref_input in derivative.saved_inputs:
|
||||
ref_input_jit_name = ref_input.expr.split(".")[0]
|
||||
mapped_name = map_refarg2foreacharg[ref_input_jit_name]
|
||||
if isinstance(map_name2arg[mapped_name].type, ListType):
|
||||
mapped_expr = mapped_name + "[i]"
|
||||
else:
|
||||
mapped_expr = mapped_name
|
||||
new_expr = ref_input.expr.replace(ref_input_jit_name, mapped_expr)
|
||||
modified_formula = modified_formula.replace(
|
||||
cast(str, ref_input.nctype.name), new_expr
|
||||
)
|
||||
|
||||
nctype = cpp.argument_type(map_name2arg[mapped_name], binds=mapped_name)
|
||||
canonical_nctype = NamedCType(
|
||||
nctype.name, nctype.type.remove_const_ref()
|
||||
)
|
||||
saved_inputs.append(
|
||||
SavedAttribute(nctype=canonical_nctype, expr=mapped_name)
|
||||
)
|
||||
for ref_output in derivative.saved_outputs:
|
||||
if ref_output.nctype.name == "result":
|
||||
saved_outputs.append(
|
||||
SavedAttribute(
|
||||
nctype=NamedCType(
|
||||
name="result", type=BaseCType(tensorListT)
|
||||
),
|
||||
expr="result",
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("")
|
||||
var_names = [map_refarg2foreacharg[var] for var in derivative.var_names]
|
||||
all_var_names.extend(var_names)
|
||||
all_saved_inputs.extend(saved_inputs)
|
||||
all_saved_outputs.extend(saved_outputs)
|
||||
modified_derivative = Derivative(
|
||||
formula=modified_formula,
|
||||
original_formula=derivative.formula,
|
||||
var_names=tuple(var_names),
|
||||
saved_inputs=tuple(saved_inputs),
|
||||
saved_outputs=tuple(saved_outputs),
|
||||
named_gradients=set(),
|
||||
)
|
||||
modified_derivative_formulas.append(modified_derivative)
|
||||
|
||||
with local.parametrize(
|
||||
use_const_ref_for_mutable_tensors=foreach_function.use_const_ref_for_mutable_tensors,
|
||||
use_ilistref_for_tensor_lists=foreach_function.part_of_structured_group,
|
||||
):
|
||||
args_with_derivatives = [
|
||||
Binding(
|
||||
name=arg.name,
|
||||
nctype=cpp.argument_type(arg, binds=arg.name),
|
||||
argument=arg,
|
||||
default=None,
|
||||
)
|
||||
for arg in foreach_function.func.arguments.flat_non_out
|
||||
if arg.name in all_var_names
|
||||
]
|
||||
|
||||
forward_derivatives: list[ForwardDerivative] = []
|
||||
fw_derivative: ForwardDerivative
|
||||
for fw_derivative in ref_diff_info.forward_derivatives:
|
||||
var_names: list[str] = list(fw_derivative.var_names) # type: ignore[no-redef]
|
||||
var_types: list[Type] = list(fw_derivative.var_types)
|
||||
required_inputs_fw_grad: list[str] = []
|
||||
required_inputs_primal: list[str] = []
|
||||
if fw_derivative.required_inputs_fw_grad is not None:
|
||||
required_inputs_fw_grad = list(fw_derivative.required_inputs_fw_grad)
|
||||
if fw_derivative.required_inputs_primal:
|
||||
required_inputs_primal = list(fw_derivative.required_inputs_primal)
|
||||
modified_formula = fw_derivative.formula
|
||||
|
||||
# Foreach's result is TensorList
|
||||
if "result" in modified_formula:
|
||||
modified_formula = fw_derivative.formula.replace("result", "result[i]")
|
||||
|
||||
for foreach_arg, ref_arg in zip(
|
||||
foreach_function.func.arguments.flat_non_out,
|
||||
ref_diff_info.func.func.arguments.flat_non_out,
|
||||
):
|
||||
# Modify reference forward formula
|
||||
if (
|
||||
isinstance(foreach_arg.type, ListType)
|
||||
and not foreach_arg.type.is_tensor_like()
|
||||
):
|
||||
# Assuming ScalarList
|
||||
modified_formula = modified_formula.replace(
|
||||
ref_arg.name, foreach_arg.name + "[i]"
|
||||
)
|
||||
elif foreach_arg.type.is_tensor_like():
|
||||
# Assuming TensorList / Tensor
|
||||
if not (
|
||||
isinstance(foreach_arg.type, ListType)
|
||||
or (
|
||||
foreach_arg.type == BaseType(BaseTy.Tensor)
|
||||
and str(foreach_function.func.name)
|
||||
in _foreach_with_tensor_overload
|
||||
)
|
||||
):
|
||||
raise AssertionError(
|
||||
f"{foreach_function.func.name}, {foreach_arg.type}"
|
||||
)
|
||||
for suffix in ("_p", "_t"):
|
||||
curr_expr = ref_arg.name + suffix
|
||||
if curr_expr in modified_formula:
|
||||
new_expr = foreach_arg.name + suffix
|
||||
modified_formula = modified_formula.replace(curr_expr, new_expr)
|
||||
else:
|
||||
# Assuming Scalar
|
||||
if foreach_arg.name != ref_arg.name:
|
||||
modified_formula = modified_formula.replace(
|
||||
ref_arg.name, foreach_arg.name
|
||||
)
|
||||
|
||||
# note(crcrpar): there should exist a cooler way...
|
||||
for i, name in enumerate(var_names):
|
||||
if name == ref_arg.name:
|
||||
var_names[i] = foreach_arg.name
|
||||
var_types[i] = foreach_arg.type
|
||||
for i, name in enumerate(required_inputs_fw_grad):
|
||||
if name == ref_arg.name:
|
||||
required_inputs_fw_grad[i] = foreach_arg.name
|
||||
for i, name in enumerate(required_inputs_primal):
|
||||
if name == ref_arg.name:
|
||||
required_inputs_primal[i] = foreach_arg.name
|
||||
forward_derivatives.append(
|
||||
ForwardDerivative(
|
||||
formula=modified_formula,
|
||||
var_names=tuple(var_names),
|
||||
var_types=tuple(var_types),
|
||||
required_inputs_fw_grad=tuple(required_inputs_fw_grad),
|
||||
required_inputs_primal=tuple(required_inputs_primal),
|
||||
required_original_self_value=fw_derivative.required_original_self_value,
|
||||
is_reusing_outplace_formula=fw_derivative.is_reusing_outplace_formula,
|
||||
)
|
||||
)
|
||||
|
||||
return (
|
||||
DifferentiabilityInfo(
|
||||
name=foreach_function.func.name.name.base,
|
||||
func=foreach_function,
|
||||
op=f"Foreach{ref_diff_info.op}{foreach_function.func.name.overload_name}",
|
||||
derivatives=modified_derivative_formulas,
|
||||
forward_derivatives=forward_derivatives,
|
||||
all_saved_inputs=tuple(set(all_saved_inputs)),
|
||||
all_saved_outputs=tuple(set(all_saved_outputs)),
|
||||
available_named_gradients=(),
|
||||
used_named_gradients=set(),
|
||||
args_with_derivatives=args_with_derivatives,
|
||||
non_differentiable_arg_names=[],
|
||||
output_differentiability=None,
|
||||
output_differentiability_conditions=None,
|
||||
),
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
def match_differentiability_info(
|
||||
native_functions: list[NativeFunction],
|
||||
differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],
|
||||
) -> list[NativeFunctionWithDifferentiabilityInfo]:
|
||||
"""Sets the "derivative" key on declarations to matching autograd function
|
||||
In-place functions will use the out-of-place derivative definition if there
|
||||
is no in-place specific derivative.
|
||||
"""
|
||||
|
||||
functional_info_by_signature = {
|
||||
schema.signature(strip_default=True): info_dict
|
||||
for schema, info_dict in differentiability_infos.items()
|
||||
if schema.kind() == SchemaKind.functional
|
||||
}
|
||||
non_functional_info_by_signature = {
|
||||
schema.signature(strip_default=True): info_dict
|
||||
for schema, info_dict in differentiability_infos.items()
|
||||
if schema.kind() != SchemaKind.functional
|
||||
}
|
||||
|
||||
def find_info(
|
||||
f: NativeFunction,
|
||||
) -> tuple[dict[str, DifferentiabilityInfo] | None, bool]:
|
||||
# Don't bother matching info to generated out= variants
|
||||
if "generated" in f.tags and f.func.kind() == SchemaKind.out:
|
||||
return None, False
|
||||
|
||||
# (1) Check for an exact match
|
||||
if f.func in differentiability_infos:
|
||||
return differentiability_infos[f.func], True
|
||||
|
||||
# (2) If no exact match, check if the out-of-place variant
|
||||
# of this operator has a match.
|
||||
# i.e mul() for mul_() or mul_out()
|
||||
# note(crcrpar): Check foreach or not because in-place foreach functions use backward defined for the existing
|
||||
# native functions instead of the out-place counterparts.
|
||||
f_sig = f.func.signature(strip_default=True)
|
||||
if f_sig in functional_info_by_signature and not is_foreach_func(f):
|
||||
return functional_info_by_signature[f_sig], False
|
||||
|
||||
# (3) Some operators have a derivative explicitly defined for the mutable
|
||||
# variant, but get a code-generated out-of-place variant which does *not*
|
||||
# come with a derivative formula.
|
||||
# For the generated out-of-place variant, use the mutable variant's formula
|
||||
# if it exists.
|
||||
if "generated" in f.tags and f_sig in non_functional_info_by_signature:
|
||||
info_dict = non_functional_info_by_signature[f_sig]
|
||||
# See https://github.com/pytorch/pytorch/pull/76320/files#r874816389
|
||||
if any(
|
||||
any("self" in str(input.nctype.name) for input in info.all_saved_inputs)
|
||||
for info in info_dict.values()
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Attempted to convert a derivative formula for a mutable operator "
|
||||
f'to be used automatically by its functional variant ("{str(f.func)}"). '
|
||||
"This is not currently supported (we'd need to fix up the formula in the codegen)."
|
||||
)
|
||||
return info_dict, False
|
||||
|
||||
# (4) Generate derivative information of foreach functions if none is defined in `derivatives.yaml`
|
||||
if is_foreach_func(f):
|
||||
if f.func in differentiability_infos:
|
||||
raise AssertionError(
|
||||
f"Foreach function {f.func.name} already has differentiability info"
|
||||
)
|
||||
diff_info, is_generated = gen_foreach_derivativeinfo(
|
||||
f,
|
||||
functional_info_by_signature,
|
||||
non_functional_info_by_signature,
|
||||
)
|
||||
if diff_info is None:
|
||||
return None, False
|
||||
# TODO(crcrpar): Avoid hard coding "Default" ideally.
|
||||
diff_info_dict = {"Default": diff_info}
|
||||
if is_generated:
|
||||
differentiability_infos[f.func] = diff_info_dict
|
||||
functional_info_by_signature[f.func] = diff_info_dict
|
||||
return diff_info_dict, is_generated
|
||||
|
||||
return None, False
|
||||
|
||||
result: list[NativeFunctionWithDifferentiabilityInfo] = []
|
||||
for f in native_functions:
|
||||
info_dict, is_exact_match = find_info(f)
|
||||
|
||||
# Currently, the '.strides()' to 'strides_or_error' replacement does not support
|
||||
# 'self' derivatives of an inplace function, so we must check for this case.
|
||||
if f.func.kind() == SchemaKind.inplace and (info_dict is not None):
|
||||
for info in info_dict.values():
|
||||
for derivative in info.derivatives:
|
||||
if "self" in derivative.var_names:
|
||||
for saved_input in derivative.saved_inputs:
|
||||
if "strides_or_error" in saved_input.expr:
|
||||
raise AssertionError(
|
||||
"Calling '.strides()' in the 'self' derivative formula of an "
|
||||
f"in-place function is not supported: {f.func}"
|
||||
)
|
||||
|
||||
if not info_dict:
|
||||
result.append(
|
||||
NativeFunctionWithDifferentiabilityInfo(
|
||||
func=f, info=None, fw_derivatives=None
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
fw_derivative_dict: dict[str, Sequence[ForwardDerivative]] = {}
|
||||
for key, info in info_dict.items():
|
||||
if not info.forward_derivatives:
|
||||
fw_derivative_dict[key] = []
|
||||
continue
|
||||
|
||||
forward_derivatives = info.forward_derivatives
|
||||
|
||||
# For functions that have a single def for out-of-place and inplace (like abs())
|
||||
if f.func.kind() == SchemaKind.inplace:
|
||||
# For inplace functions there is a little bit of work to do:
|
||||
# 1) Validate the formula and make sure the input that is modified in not used:
|
||||
# - If there is a formula for the inplace variant of the function (is_exact_match == True) then
|
||||
# we make sure that the original value of the input that is being modified inplace (self_p) is
|
||||
# not used in the formula. Note that the formula can use "original_self_p" here and that would
|
||||
# trigger a clone of the original input.
|
||||
# - If we are reusing the out of place formula (is_exact_match == False) then we replace every
|
||||
# occurrence of self_p and self_t by original_self_p and original_self_t. These will be
|
||||
# populated by cloned version of the original input (either the clone done by the backward AD
|
||||
# logic if self is also used in a backward formula or a special clone that we add).
|
||||
# 2) At this point, there cannot be a self_p in the formula.
|
||||
# 3) Change "result" into "self_p" as by design, in the inplace function codegen, the result is
|
||||
# simply called self (as it is modified inplace).
|
||||
# 4) Update the required primals data in case it used to contain "result" but should now contain
|
||||
# "self"
|
||||
# 5) If it is not an exact match, the user formula is not modifying the existing forward grad
|
||||
# inplace as it should. So add some code that makes sure that we do so if the forward grad
|
||||
# already exists.
|
||||
|
||||
if len(info.forward_derivatives) != 1:
|
||||
raise AssertionError(
|
||||
"Only single output inplace should exist, "
|
||||
f"got {len(info.forward_derivatives)}"
|
||||
)
|
||||
fw_info = info.forward_derivatives[0]
|
||||
formula = fw_info.formula
|
||||
|
||||
def replace_self_with_original_self(formula: str, postfix: str) -> str:
|
||||
def repl(m: re.Match[str]) -> str:
|
||||
return f"{m.group(1)}original_self{postfix}{m.group(2)}"
|
||||
|
||||
return re.sub(IDENT_REGEX.format(f"self{postfix}"), repl, formula)
|
||||
|
||||
if re.search(IDENT_REGEX.format("self_p"), formula):
|
||||
if is_exact_match:
|
||||
# For manually defined formulas, don't allow the original value to be used
|
||||
raise RuntimeError(
|
||||
f'The formula for "{f.func.name}" is using the original value of self '
|
||||
"that is being modified inplace. This would lead to wrong forward gradients. "
|
||||
'Please use "result" in the formula only.'
|
||||
)
|
||||
else:
|
||||
# When the original formula is out of place, we save a clone of the primal
|
||||
# value to be able to access this value if needed
|
||||
# replace "self_p"/"self_t" from the formula by "original_self_p"/"original_self_t"
|
||||
formula = replace_self_with_original_self(formula, "_p")
|
||||
formula = replace_self_with_original_self(formula, "_t")
|
||||
|
||||
# replace "result" from the formula by "self_p"
|
||||
def repl(m: re.Match[str]) -> str:
|
||||
return f"{m.group(1)}self_p{m.group(2)}"
|
||||
|
||||
formula = re.sub(IDENT_REGEX.format("result"), repl, formula)
|
||||
|
||||
required_primals = fw_info.required_inputs_primal
|
||||
if re.search(IDENT_REGEX.format("self_p"), formula):
|
||||
required_primals = (
|
||||
required_primals + ("self",) if required_primals else ("self",)
|
||||
)
|
||||
|
||||
if not is_exact_match:
|
||||
# NOTE [In-place forward AD formula Optimization]
|
||||
#
|
||||
# This optimization transforms the formula to directly do inplace, i.e.
|
||||
# instead of self_t.copy_(self_t.op()) we do self_t.op_() when the following are met:
|
||||
#
|
||||
# 1) the formula satisfies the pattern: "self_t.op(*args)"
|
||||
# 2) "op" in (1) needs to be the same as the op the derivative is for
|
||||
#
|
||||
# (2) may seem too strict, but currently the only ops that satisfy (1) also satisfy (2)
|
||||
# If there is a need, we can relax (2) to allow any op that has an in-place variant
|
||||
is_single_method_on_self_t = False
|
||||
directly_do_inplace = False
|
||||
op_name: str | None = None
|
||||
between_parens: str | None = None
|
||||
match = re.fullmatch(r"self_t.([\w]*)\((.*)\)", formula)
|
||||
if match:
|
||||
op_name, between_parens = match.group(1), match.group(2)
|
||||
|
||||
# We want to...
|
||||
# Match: self_t.op1(other_p.op2(arg))
|
||||
# Avoid: self_t.op1(args) + self_t.op2(args)
|
||||
# Avoid: self_t.op1(other_p.op2(arg)) + self_t.op2(args)
|
||||
def check_parens_nest_level_gt_zero(s: str) -> bool:
|
||||
level = 1
|
||||
for ch in s:
|
||||
if ch == ")":
|
||||
level -= 1
|
||||
if level == 0:
|
||||
return False
|
||||
if ch == "(":
|
||||
level += 1
|
||||
return True
|
||||
|
||||
is_single_method_on_self_t = check_parens_nest_level_gt_zero(
|
||||
between_parens
|
||||
)
|
||||
directly_do_inplace = (
|
||||
is_single_method_on_self_t and op_name == info.name
|
||||
)
|
||||
|
||||
if directly_do_inplace:
|
||||
if op_name is None:
|
||||
raise AssertionError("op_name must be non-None for inplace")
|
||||
if between_parens is None:
|
||||
raise AssertionError(
|
||||
"between_parens must be non-None for inplace"
|
||||
)
|
||||
formula = f"self_t_raw.defined() ? self_t_raw.{op_name}_({between_parens}) : {formula}"
|
||||
else:
|
||||
# Make sure that the forward grad is modified inplace when the original formula
|
||||
# is out of place
|
||||
formula = f"self_t_raw.defined() ? self_t_raw.copy_({formula}) : {formula}"
|
||||
|
||||
required_original_self_value = bool(
|
||||
re.search(IDENT_REGEX.format("original_self_p"), formula)
|
||||
) or bool(re.search(IDENT_REGEX.format("original_self_t"), formula))
|
||||
|
||||
forward_derivatives = [
|
||||
ForwardDerivative(
|
||||
formula=formula,
|
||||
var_names=("self",),
|
||||
var_types=fw_info.var_types,
|
||||
required_inputs_fw_grad=fw_info.required_inputs_fw_grad,
|
||||
required_inputs_primal=required_primals,
|
||||
required_original_self_value=required_original_self_value,
|
||||
is_reusing_outplace_formula=not is_exact_match,
|
||||
),
|
||||
]
|
||||
|
||||
fw_derivative_dict[key] = forward_derivatives
|
||||
|
||||
result.append(
|
||||
NativeFunctionWithDifferentiabilityInfo(
|
||||
func=f, info=info_dict, fw_derivatives=fw_derivative_dict
|
||||
)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def is_differentiable(
|
||||
name: str, type: Type, info: DifferentiabilityInfo | None
|
||||
) -> bool:
|
||||
return type.is_tensor_like() and (
|
||||
info is None or name not in info.non_differentiable_arg_names
|
||||
)
|
||||
|
||||
|
||||
def gen_differentiable_outputs(
|
||||
fn: NativeFunctionWithDifferentiabilityInfo, key: str = "Default"
|
||||
) -> list[DifferentiableOutput]:
|
||||
f = fn.func
|
||||
info = fn.info[key] if fn.info else None
|
||||
outputs: list[DifferentiableOutput] = [
|
||||
DifferentiableOutput(
|
||||
name=name,
|
||||
type=ret.type,
|
||||
cpp_type=cpp.return_type(ret, symint=True).cpp_type(),
|
||||
)
|
||||
for name, ret in zip(cpp.return_names(f), f.func.returns)
|
||||
]
|
||||
output_differentiability = info.output_differentiability if info else None
|
||||
if output_differentiability is not None:
|
||||
if len(output_differentiability) != len(outputs):
|
||||
raise RuntimeError(
|
||||
f"The length of output_differentiability ({len(output_differentiability)}), "
|
||||
f"does not match the number of outputs ({len(outputs)})."
|
||||
)
|
||||
differentiable_outputs: list[DifferentiableOutput] = []
|
||||
if False in output_differentiability and f.func.kind() == SchemaKind.inplace:
|
||||
raise RuntimeError(
|
||||
"output_differentiability=False for inplace operation (version_counter won't get updated)"
|
||||
)
|
||||
for differentiable, output in zip(output_differentiability, outputs):
|
||||
if differentiable:
|
||||
differentiable_outputs.append(output)
|
||||
return differentiable_outputs
|
||||
candidate_differentiable_outputs = list(
|
||||
filter(lambda r: is_differentiable(r.name, r.type, info), outputs)
|
||||
)
|
||||
if uses_single_grad(info):
|
||||
return candidate_differentiable_outputs[:1]
|
||||
else:
|
||||
return candidate_differentiable_outputs
|
||||
@@ -0,0 +1,475 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from torchgen import local
|
||||
from torchgen.api.types import (
|
||||
ArgName,
|
||||
ArrayCType,
|
||||
ArrayRefCType,
|
||||
BaseCType,
|
||||
BaseTypeToCppMapping,
|
||||
Binding,
|
||||
boolT,
|
||||
ConstRefCType,
|
||||
CType,
|
||||
dimnameListT,
|
||||
intArrayRefT,
|
||||
iTensorListRefT,
|
||||
ListCType,
|
||||
longT,
|
||||
MutRefCType,
|
||||
NamedCType,
|
||||
OptionalCType,
|
||||
optionalIntArrayRefT,
|
||||
optionalSymIntArrayRefT,
|
||||
scalarT,
|
||||
SpecialArgName,
|
||||
symIntArrayRefT,
|
||||
SymIntT,
|
||||
tensorListT,
|
||||
tensorOptionsT,
|
||||
tensorT,
|
||||
TupleCType,
|
||||
VectorCType,
|
||||
voidT,
|
||||
)
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
Arguments,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
NativeFunction,
|
||||
OptionalType,
|
||||
Return,
|
||||
SelfArgument,
|
||||
TensorOptionsArguments,
|
||||
Type,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# This file describes the translation of JIT schema to the public C++
|
||||
# API, which is what people use when they call functions like at::add.
|
||||
#
|
||||
# Prominent characteristics of the C++ API:
|
||||
#
|
||||
# - dtype, layout, device and pin_memory are collected into
|
||||
# a single C++ type TensorOptions (the native functions API
|
||||
# also has this, but tensor options is really most relevant
|
||||
# for the C++ API; it makes calling kwarg factory functions
|
||||
# pleasant)
|
||||
#
|
||||
# - defaulting lives here (in fact, the dispatcher is completely
|
||||
# oblivious of defaults!)
|
||||
#
|
||||
# BTW: policy on name collisions: we try not to have types with
|
||||
# collisions, but functions are fair game to collide
|
||||
|
||||
|
||||
def name(
|
||||
func: FunctionSchema,
|
||||
*,
|
||||
faithful_name_for_out_overloads: bool = False,
|
||||
symint_overload: bool = False,
|
||||
) -> str:
|
||||
name = str(func.name.name)
|
||||
if symint_overload:
|
||||
name += "_symint"
|
||||
if func.is_out_fn():
|
||||
if faithful_name_for_out_overloads:
|
||||
name += "_outf"
|
||||
else:
|
||||
name += "_out"
|
||||
|
||||
return name
|
||||
|
||||
|
||||
# Translation of "value types" in JIT schema to C++ API type. Value
|
||||
# types look the same no matter if they are argument types or return
|
||||
# types. Returns None if the type in question is not a value type.
|
||||
def valuetype_type(
|
||||
t: Type,
|
||||
*,
|
||||
binds: ArgName,
|
||||
mutable: bool = True,
|
||||
symint: bool = False,
|
||||
) -> NamedCType | None:
|
||||
if isinstance(t, BaseType):
|
||||
if t.name in (BaseTy.Tensor, BaseTy.Scalar):
|
||||
return None
|
||||
elif str(t) == "SymInt":
|
||||
if symint:
|
||||
return NamedCType(binds, BaseCType(SymIntT))
|
||||
else:
|
||||
return NamedCType(binds, BaseCType(longT))
|
||||
# All other BaseType currently map directly to BaseCppTypes.
|
||||
return NamedCType(binds, BaseCType(BaseTypeToCppMapping[t.name]))
|
||||
elif isinstance(t, OptionalType):
|
||||
elem = valuetype_type(t.elem, binds=binds, mutable=mutable, symint=symint)
|
||||
if elem is None:
|
||||
return None
|
||||
return NamedCType(binds, OptionalCType(elem.type))
|
||||
elif isinstance(t, ListType):
|
||||
if str(t.elem) == "bool":
|
||||
if t.size is None:
|
||||
raise AssertionError("bool ListType must have a size")
|
||||
return NamedCType(binds, ArrayCType(BaseCType(boolT), t.size))
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
# Translation of types occurring in JIT arguments to a C++ argument type.
|
||||
# If remove_non_owning_ref_types is set, we'll guarantee that the output CType is not a non-owning reference type.
|
||||
# For example, we'll return std::vector<int> instead of IntArrayRef.
|
||||
# See Note [translation from C++ reference to value types]
|
||||
def argumenttype_type(
|
||||
t: Type,
|
||||
*,
|
||||
mutable: bool,
|
||||
binds: ArgName,
|
||||
remove_non_owning_ref_types: bool = False,
|
||||
symint: bool = False,
|
||||
) -> NamedCType:
|
||||
# If it's a value type, do the value type translation
|
||||
r = valuetype_type(
|
||||
t,
|
||||
binds=binds,
|
||||
mutable=mutable,
|
||||
symint=symint,
|
||||
)
|
||||
if r is not None:
|
||||
return r
|
||||
|
||||
if isinstance(t, BaseType):
|
||||
if t.name == BaseTy.Tensor:
|
||||
if mutable and not local.use_const_ref_for_mutable_tensors():
|
||||
return NamedCType(binds, MutRefCType(BaseCType(tensorT)))
|
||||
else:
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(tensorT)))
|
||||
elif t.name == BaseTy.Scalar:
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))
|
||||
else:
|
||||
raise AssertionError(f"base type should have been value type {t}")
|
||||
elif isinstance(t, OptionalType):
|
||||
if str(t.elem) == "Tensor":
|
||||
if mutable and not local.use_const_ref_for_mutable_tensors():
|
||||
return NamedCType(
|
||||
binds, MutRefCType(BaseCType(tensorT))
|
||||
) # TODO: fix this discrepancy
|
||||
else:
|
||||
return NamedCType(
|
||||
binds, ConstRefCType(OptionalCType(BaseCType(tensorT)))
|
||||
)
|
||||
elif str(t.elem) == "Scalar":
|
||||
return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT))))
|
||||
elif isinstance(t.elem, ListType) and str(t.elem.elem) == "int":
|
||||
return NamedCType(binds, BaseCType(optionalIntArrayRefT))
|
||||
elif isinstance(t.elem, ListType) and str(t.elem.elem) == "SymInt":
|
||||
if symint:
|
||||
return NamedCType(binds, BaseCType(optionalSymIntArrayRefT))
|
||||
else:
|
||||
return NamedCType(binds, BaseCType(optionalIntArrayRefT))
|
||||
elem = argumenttype_type(t.elem, mutable=mutable, binds=binds, symint=symint)
|
||||
return NamedCType(binds, OptionalCType(elem.type))
|
||||
elif isinstance(t, ListType):
|
||||
# TODO: remove these special cases, ArrayRef fallthrough works fine
|
||||
if str(t.elem) == "int":
|
||||
if remove_non_owning_ref_types:
|
||||
return NamedCType(binds, VectorCType(BaseCType(longT)))
|
||||
else:
|
||||
return NamedCType(binds, BaseCType(intArrayRefT))
|
||||
if str(t.elem) == "SymInt":
|
||||
if remove_non_owning_ref_types:
|
||||
if symint:
|
||||
return NamedCType(binds, VectorCType(BaseCType(SymIntT)))
|
||||
else:
|
||||
return NamedCType(binds, VectorCType(BaseCType(longT)))
|
||||
else:
|
||||
if symint:
|
||||
return NamedCType(binds, BaseCType(symIntArrayRefT))
|
||||
else:
|
||||
return NamedCType(binds, BaseCType(intArrayRefT))
|
||||
if str(t.elem) == "Tensor":
|
||||
if local.use_ilistref_for_tensor_lists():
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(iTensorListRefT)))
|
||||
else:
|
||||
return NamedCType(binds, BaseCType(tensorListT))
|
||||
elif str(t.elem) == "Scalar":
|
||||
return NamedCType(binds, ArrayRefCType(BaseCType(scalarT)))
|
||||
elif str(t.elem) == "Dimname":
|
||||
return NamedCType(binds, BaseCType(dimnameListT))
|
||||
elif str(t.elem) == "Tensor?":
|
||||
return NamedCType(
|
||||
binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT))))
|
||||
)
|
||||
elem = argumenttype_type(t.elem, mutable=mutable, binds=binds, symint=symint)
|
||||
return NamedCType(binds, ArrayRefCType(elem.type))
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
# Translate a JIT argument into its C++ type
|
||||
def argument_type(a: Argument, *, binds: ArgName, symint: bool = False) -> NamedCType:
|
||||
return argumenttype_type(a.type, mutable=a.is_write, symint=symint, binds=binds)
|
||||
|
||||
|
||||
# Translation of a (non-multi) return type from JIT to C++
|
||||
# N.B: returntype_type returns a CType, not a NamedCType.
|
||||
# This is mostly because of the mismatch between return types and return names.
|
||||
# e.g. a function with a return type of 'void' has 0 return names,
|
||||
# and a function with a return type of 'std::tuple' has >1 return name.
|
||||
def returntype_type(t: Type, *, mutable: bool, symint: bool = False) -> CType:
|
||||
# placeholder is ignored
|
||||
# NB: symint is ALWAYS respected for return types. So symint argument
|
||||
# here is IGNORED
|
||||
r = valuetype_type(t, binds="__placeholder__", mutable=mutable, symint=True)
|
||||
if r is not None:
|
||||
return r.type
|
||||
|
||||
if isinstance(t, BaseType):
|
||||
if t.name == BaseTy.Tensor:
|
||||
if mutable:
|
||||
if local.use_const_ref_for_mutable_tensors():
|
||||
return ConstRefCType(BaseCType(tensorT))
|
||||
else:
|
||||
return MutRefCType(BaseCType(tensorT))
|
||||
else:
|
||||
# Note [Tensor Copy Returns]
|
||||
# Currently, we use "Argument.is_write" to determine
|
||||
# whether or not Tensor return types should be copies or references.
|
||||
# If that ever changes, take a look at other locations of this note!
|
||||
return BaseCType(tensorT)
|
||||
elif t.name == BaseTy.Scalar:
|
||||
return BaseCType(scalarT)
|
||||
elif isinstance(t, ListType):
|
||||
if mutable:
|
||||
raise AssertionError(
|
||||
"Native functions should never return a mutable tensor list. "
|
||||
"They should return void."
|
||||
)
|
||||
elem = returntype_type(t.elem, mutable=False)
|
||||
if t.size is not None:
|
||||
raise AssertionError(f"fixed size list returns not supported: {t}")
|
||||
return VectorCType(elem)
|
||||
elif isinstance(t, OptionalType):
|
||||
elem = returntype_type(t.elem, mutable=mutable)
|
||||
if str(t.elem) == "Tensor":
|
||||
return OptionalCType(elem)
|
||||
|
||||
raise AssertionError(f"unrecognized return type {t}")
|
||||
|
||||
|
||||
# Translation of a single return to its C++ type
|
||||
def return_type(r: Return, *, symint: bool = False) -> CType:
|
||||
return returntype_type(r.type, mutable=r.is_write, symint=symint)
|
||||
|
||||
|
||||
# Translation of a full (possibly multi) return from JIT to its C++ type
|
||||
def returns_type(rs: Sequence[Return], *, symint: bool = False) -> CType:
|
||||
if len(rs) == 0:
|
||||
return BaseCType(voidT)
|
||||
elif len(rs) == 1:
|
||||
return return_type(rs[0], symint=symint)
|
||||
else:
|
||||
return TupleCType([return_type(r, symint=symint) for r in rs])
|
||||
|
||||
|
||||
def return_names(f: NativeFunction, *, fallback_name: str = "result") -> Sequence[str]:
|
||||
returns: list[str] = []
|
||||
for i, r in enumerate(f.func.returns):
|
||||
# If we have an inplace function, the return argument is
|
||||
# implicitly named self.
|
||||
# TODO: Consider incorporating this into the data model
|
||||
if f.func.name.name.inplace:
|
||||
if i != 0:
|
||||
raise AssertionError("illegal inplace function with multiple returns")
|
||||
name = "self"
|
||||
# If we are out function, the name is the name of the
|
||||
# corresponding output function (r.name will get recorded
|
||||
# in field_name later.)
|
||||
elif f.func.is_out_fn():
|
||||
name = f.func.arguments.out[i].name
|
||||
# If the return argument is explicitly named...
|
||||
elif r.name:
|
||||
name_conflict = any(
|
||||
r.name == a.name for a in f.func.schema_order_arguments()
|
||||
)
|
||||
if name_conflict and not f.func.is_out_fn():
|
||||
name = f"{r.name}_return"
|
||||
else:
|
||||
name = r.name
|
||||
# If there is no explicit name and no fallback name was passed in, we just name the output result,
|
||||
# unless it's a multi-return, in which case it's result0,
|
||||
# result1, etc (zero-indexed)
|
||||
else:
|
||||
name = fallback_name if len(f.func.returns) == 1 else f"{fallback_name}{i}"
|
||||
returns.append(name)
|
||||
return returns
|
||||
|
||||
|
||||
JIT_TO_CPP_DEFAULT = {
|
||||
"False": "false",
|
||||
"True": "true",
|
||||
"None": "::std::nullopt", # UGH this one is type directed
|
||||
"Mean": "at::Reduction::Mean",
|
||||
"[]": "{}",
|
||||
"contiguous_format": "c10::MemoryFormat::Contiguous",
|
||||
"long": "at::kLong",
|
||||
}
|
||||
|
||||
|
||||
# Convert a JIT default into C++ expression representing the default
|
||||
def default_expr(d: str, t: Type, *, symint: bool) -> str:
|
||||
if d == "None" and str(t) == "Tensor?":
|
||||
return "{}"
|
||||
if isinstance(t, BaseType) and t.name is BaseTy.str:
|
||||
# Schema allows single quotes but C++ needs double
|
||||
if len(d) >= 2 and d[0] == "'" and d[-1] == "'":
|
||||
s = ""
|
||||
i = 1
|
||||
while i + 1 < len(d):
|
||||
if d[i] != "\\":
|
||||
if d[i] == '"':
|
||||
s += '\\"'
|
||||
else:
|
||||
s += d[i]
|
||||
i += 1
|
||||
else:
|
||||
if d[i + 1] == "'":
|
||||
s += "'"
|
||||
else:
|
||||
s += d[i : i + 2]
|
||||
i += 2
|
||||
|
||||
return f'"{s}"'
|
||||
|
||||
if isinstance(t, OptionalType):
|
||||
if d == "None":
|
||||
return "::std::nullopt"
|
||||
|
||||
return default_expr(d, t.elem, symint=symint)
|
||||
|
||||
if isinstance(t, ListType):
|
||||
if d.startswith("[") and d.endswith("]"):
|
||||
return "{" + d[1:-1] + "}"
|
||||
elif symint and d.isdigit() and str(t.elem) == "SymInt":
|
||||
return f"c10::SymInt({d})"
|
||||
elif t.size is None:
|
||||
# NOTE: Sized lists can have scalar defaults
|
||||
raise ValueError(f"Expected a list default '[...]' but found: '{d}'")
|
||||
|
||||
return JIT_TO_CPP_DEFAULT.get(d, d)
|
||||
|
||||
|
||||
# Convert an argument into its C++ API form
|
||||
|
||||
|
||||
def argument(
|
||||
a: Argument | TensorOptionsArguments | SelfArgument,
|
||||
*,
|
||||
cpp_no_default_args: set[str],
|
||||
method: bool,
|
||||
faithful: bool,
|
||||
symint: bool = False,
|
||||
has_tensor_options: bool,
|
||||
) -> list[Binding]:
|
||||
def sub_argument(
|
||||
a: Argument | TensorOptionsArguments | SelfArgument,
|
||||
) -> list[Binding]:
|
||||
return argument(
|
||||
a,
|
||||
cpp_no_default_args=cpp_no_default_args,
|
||||
method=method,
|
||||
faithful=faithful,
|
||||
symint=symint,
|
||||
has_tensor_options=has_tensor_options,
|
||||
)
|
||||
|
||||
if isinstance(a, Argument):
|
||||
binds: ArgName
|
||||
if a.name == "memory_format" and has_tensor_options:
|
||||
binds = SpecialArgName.possibly_redundant_memory_format
|
||||
else:
|
||||
binds = a.name
|
||||
default: str | None = None
|
||||
if a.name not in cpp_no_default_args and a.default is not None:
|
||||
default = default_expr(a.default, a.type, symint=symint)
|
||||
return [
|
||||
Binding(
|
||||
nctype=argument_type(a, binds=binds, symint=symint),
|
||||
name=a.name,
|
||||
default=default,
|
||||
argument=a,
|
||||
)
|
||||
]
|
||||
elif isinstance(a, TensorOptionsArguments):
|
||||
if faithful:
|
||||
return (
|
||||
sub_argument(a.dtype)
|
||||
+ sub_argument(a.layout)
|
||||
+ sub_argument(a.device)
|
||||
+ sub_argument(a.pin_memory)
|
||||
)
|
||||
else:
|
||||
default = None
|
||||
# Enforced by NativeFunction.__post_init__
|
||||
if "options" in cpp_no_default_args:
|
||||
raise AssertionError("'options' should not be in cpp_no_default_args")
|
||||
if all(x.default == "None" for x in a.all()):
|
||||
default = "{}"
|
||||
elif a.dtype.default == "long":
|
||||
default = "at::kLong" # TODO: this is wrong
|
||||
return [
|
||||
Binding(
|
||||
nctype=NamedCType("options", BaseCType(tensorOptionsT)),
|
||||
name="options",
|
||||
default=default,
|
||||
argument=a,
|
||||
)
|
||||
]
|
||||
elif isinstance(a, SelfArgument):
|
||||
if method:
|
||||
# Caller is responsible for installing implicit this in context!
|
||||
return []
|
||||
else:
|
||||
return sub_argument(a.argument)
|
||||
else:
|
||||
assert_never(a)
|
||||
|
||||
|
||||
def arguments(
|
||||
arguments: Arguments,
|
||||
*,
|
||||
faithful: bool,
|
||||
symint: bool = False,
|
||||
method: bool,
|
||||
cpp_no_default_args: set[str],
|
||||
) -> list[Binding]:
|
||||
args: list[Argument | TensorOptionsArguments | SelfArgument] = []
|
||||
if faithful:
|
||||
args.extend(arguments.non_out)
|
||||
args.extend(arguments.out)
|
||||
else:
|
||||
args.extend(arguments.out)
|
||||
args.extend(arguments.non_out)
|
||||
return [
|
||||
r.no_default() if faithful else r
|
||||
for a in args
|
||||
for r in argument(
|
||||
a,
|
||||
faithful=faithful,
|
||||
symint=symint,
|
||||
method=method,
|
||||
has_tensor_options=arguments.tensor_options is not None,
|
||||
cpp_no_default_args=cpp_no_default_args,
|
||||
)
|
||||
]
|
||||
@@ -0,0 +1,125 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from typing import TYPE_CHECKING
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import ArgName, Binding, CType, NamedCType
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
FunctionSchema,
|
||||
Return,
|
||||
SelfArgument,
|
||||
TensorOptionsArguments,
|
||||
Type,
|
||||
)
|
||||
from torchgen.utils import concatMap
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# This file describes the translation of JIT schema to the dispatcher
|
||||
# API, the *unboxed* calling convention by which invocations through
|
||||
# the dispatcher are made. Historically, the dispatcher API matched
|
||||
# the C++ API, but with the establishment of the boxed API, we've
|
||||
# made changes to the dispatcher API to so that the unboxed API
|
||||
# better aligns with the boxed API. The dispatcher API hooks heavily
|
||||
# into our template based boxing/unboxing machinery, so changes
|
||||
# to this convention will usually need template updates too.
|
||||
#
|
||||
# Prominent characteristics of the dispatcher API:
|
||||
#
|
||||
# - dtype, layout, device and pin_memory are represented as separate
|
||||
# arguments.
|
||||
#
|
||||
|
||||
|
||||
def name(func: FunctionSchema) -> str:
|
||||
return cpp.name(func)
|
||||
|
||||
|
||||
def argumenttype_type(
|
||||
t: Type,
|
||||
*,
|
||||
mutable: bool,
|
||||
binds: ArgName,
|
||||
remove_non_owning_ref_types: bool = False,
|
||||
symint: bool = True,
|
||||
) -> NamedCType:
|
||||
# This is a faux amis. If it makes sense in the future to add
|
||||
# more special cases here, or invert things so cpp.argument_type
|
||||
# calls this, or just completely inline the function, please do
|
||||
# it.
|
||||
return cpp.argumenttype_type(
|
||||
t,
|
||||
mutable=mutable,
|
||||
binds=binds,
|
||||
symint=symint,
|
||||
remove_non_owning_ref_types=remove_non_owning_ref_types,
|
||||
)
|
||||
|
||||
|
||||
def argument_type(
|
||||
a: Argument,
|
||||
*,
|
||||
binds: ArgName,
|
||||
remove_non_owning_ref_types: bool = False,
|
||||
symint: bool = True,
|
||||
) -> NamedCType:
|
||||
return argumenttype_type(
|
||||
a.type,
|
||||
mutable=a.is_write,
|
||||
binds=binds,
|
||||
remove_non_owning_ref_types=remove_non_owning_ref_types,
|
||||
symint=symint,
|
||||
)
|
||||
|
||||
|
||||
def returns_type(rs: Sequence[Return], *, symint: bool = True) -> CType:
|
||||
# At present, there is no difference. But there could be!
|
||||
return cpp.returns_type(rs, symint=symint)
|
||||
|
||||
|
||||
def jit_arguments(func: FunctionSchema) -> list[Argument]:
|
||||
def to_argument(
|
||||
a: Argument | TensorOptionsArguments | SelfArgument,
|
||||
) -> list[Argument]:
|
||||
if isinstance(a, Argument):
|
||||
return [a]
|
||||
elif isinstance(a, SelfArgument):
|
||||
return [a.argument]
|
||||
elif isinstance(a, TensorOptionsArguments):
|
||||
return [a.dtype, a.layout, a.device, a.pin_memory]
|
||||
else:
|
||||
assert_never(a)
|
||||
|
||||
return list(
|
||||
concatMap(
|
||||
to_argument,
|
||||
itertools.chain(
|
||||
func.arguments.positional, func.arguments.kwarg_only, func.arguments.out
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def argument(
|
||||
a: Argument, *, remove_non_owning_ref_types: bool = False, symint: bool = True
|
||||
) -> Binding:
|
||||
return Binding(
|
||||
nctype=argument_type(
|
||||
a,
|
||||
binds=a.name,
|
||||
remove_non_owning_ref_types=remove_non_owning_ref_types,
|
||||
symint=symint,
|
||||
),
|
||||
name=a.name,
|
||||
argument=a,
|
||||
)
|
||||
|
||||
|
||||
def arguments(func: FunctionSchema, *, symint: bool = True) -> list[Binding]:
|
||||
return [argument(a, symint=symint) for a in jit_arguments(func)]
|
||||
@@ -0,0 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from torchgen.api import dispatcher
|
||||
from torchgen.api.types import (
|
||||
BaseCppType,
|
||||
BaseCType,
|
||||
Binding,
|
||||
boolT,
|
||||
ConstRefCType,
|
||||
CType,
|
||||
longT,
|
||||
NamedCType,
|
||||
tensorT,
|
||||
)
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
NativeFunction,
|
||||
NativeFunctionsViewGroup,
|
||||
)
|
||||
|
||||
|
||||
# This file describes the translation of JIT schema to API's used
|
||||
# when creating `ViewMeta` specializations that are used by the functionalization pass.
|
||||
# These API's mostly follow the dispatcher API, with one difference:
|
||||
# - While the forward function just directly calls into the at::_ops API
|
||||
# (following the dispatcher convention), the logic here for the reverse function
|
||||
# is responsible for generating both the call-site, and the declarations
|
||||
# (which are implemented manually in the at::functionalization::impl namespace).
|
||||
|
||||
# Define some specific lambda input arguments.
|
||||
base_binding = Binding(
|
||||
name="base",
|
||||
nctype=NamedCType(name="base", type=ConstRefCType(BaseCType(tensorT))),
|
||||
argument=Argument(
|
||||
name="base", type=BaseType(BaseTy.Tensor), default=None, annotation=None
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
|
||||
has_symbolic_inputs_binding = Binding(
|
||||
name="has_symbolic_inputs",
|
||||
nctype=NamedCType(name="has_symbolic_inputs", type=BaseCType(boolT)),
|
||||
argument=Argument(
|
||||
name="has_symbolic_inputs",
|
||||
type=BaseType(BaseTy.bool),
|
||||
default=None,
|
||||
annotation=None,
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
mutated_view_binding = Binding(
|
||||
name="mutated_view",
|
||||
nctype=NamedCType(name="mutated_view", type=ConstRefCType(BaseCType(tensorT))),
|
||||
argument=Argument(
|
||||
name="base", type=BaseType(BaseTy.Tensor), default=None, annotation=None
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
out_index_binding = Binding(
|
||||
name="out_index",
|
||||
nctype=NamedCType(name="out_index", type=BaseCType(longT)),
|
||||
argument=Argument(
|
||||
name="out_index", type=BaseType(BaseTy.int), default=None, annotation=None
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
reapply_views_binding = Binding(
|
||||
name="reapply_views",
|
||||
nctype=NamedCType(name="reapply_views", type=BaseCType(boolT)),
|
||||
argument=Argument(
|
||||
name="reapply_views", type=BaseType(BaseTy.bool), default=None, annotation=None
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
|
||||
InverseReturnModeT = BaseCppType("at::functionalization", "InverseReturnMode")
|
||||
inverse_return_mode_binding = Binding(
|
||||
name="inverse_return_mode",
|
||||
nctype=NamedCType(name="inverse_return_mode", type=BaseCType(InverseReturnModeT)),
|
||||
argument=Argument(
|
||||
name="inverse_return_mode",
|
||||
# NB: not actually a bool but it doesn't matter because this isn't used
|
||||
type=BaseType(BaseTy.bool),
|
||||
default=None,
|
||||
annotation=None,
|
||||
),
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
# Name of the `ViewMeta` specialization class created.
|
||||
def classname(func: FunctionSchema, with_namespace: bool = False) -> str:
|
||||
namespace = "at::functionalization::" if with_namespace else ""
|
||||
return f"{namespace}{func.name.unambiguous_name()}_ViewMeta"
|
||||
|
||||
|
||||
# Name of the operation called inside the `forward`/`reverse` implementations.
|
||||
def name(
|
||||
g: NativeFunctionsViewGroup,
|
||||
*,
|
||||
is_reverse: bool,
|
||||
include_namespace: bool,
|
||||
reapply_views: bool | None = None,
|
||||
) -> str:
|
||||
if reapply_views is None:
|
||||
# reapply_views is only important for the fwd lambda,
|
||||
# since we always plumb the runtime "reapply_views" argument into the reverse function.
|
||||
if not is_reverse:
|
||||
raise AssertionError("reapply_views can only be None for reverse")
|
||||
if is_reverse:
|
||||
return reverse_name(g.view, include_namespace)
|
||||
# in the forward case, we just directly call into the at::_ops API (so we always need the namespace)
|
||||
if not include_namespace:
|
||||
raise AssertionError("include_namespace must be True for forward")
|
||||
if g.view_copy is None:
|
||||
raise AssertionError("view_copy must be non-None for forward")
|
||||
api_name = (
|
||||
g.view.func.name.unambiguous_name()
|
||||
if reapply_views
|
||||
else g.view_copy.func.name.unambiguous_name()
|
||||
)
|
||||
return f"at::_ops::{api_name}::call"
|
||||
|
||||
|
||||
def reverse_name(f: NativeFunction, include_namespace: bool) -> str:
|
||||
# for the reverse: we plumb the "reapply_views" flag into that function and support
|
||||
# both copy and non-copy variants. (We could avoid doing that, but that would require
|
||||
# writing out twice as many view inverse functions).
|
||||
api_name = f.func.name.unambiguous_name()
|
||||
# in the reverse case, we codegen both the call-sites (which need the full namespace) and the declarations (which don't)
|
||||
if include_namespace:
|
||||
return f"at::functionalization::FunctionalInverses::{api_name}_inverse"
|
||||
else:
|
||||
return f"{api_name}_inverse"
|
||||
|
||||
|
||||
def returns_type(func: FunctionSchema) -> CType:
|
||||
# Assertion: all view ops return tensor-like outputs
|
||||
if len(func.returns) < 1:
|
||||
raise AssertionError("Expected at least one return value")
|
||||
for ret in func.returns:
|
||||
if not ret.type.is_tensor_like():
|
||||
raise AssertionError(f"Expected tensor-like return type, got {ret.type}")
|
||||
# However, the return type of the lambda is always an individual tensor.
|
||||
# For multi-tensor outputs, each tensor needs to be tracked individually.
|
||||
return BaseCType(tensorT)
|
||||
|
||||
|
||||
# Checks whether `func` might return more than one value.
|
||||
def is_multi_output(func: FunctionSchema) -> bool:
|
||||
return len(func.returns) > 1 or (
|
||||
len(func.returns) == 1 and func.returns[0].type.is_list_like() is not None
|
||||
)
|
||||
|
||||
|
||||
# `ViewMeta` specialization constructor parameters.
|
||||
def base_ctor_arguments(func: FunctionSchema) -> list[Binding]:
|
||||
# All specializations are parematerized by `has_symbolic_inputs` flag.
|
||||
arguments = [has_symbolic_inputs_binding]
|
||||
|
||||
# If `func` might return more than 1 value, we also parameterize this specialization
|
||||
# with the output index.
|
||||
if is_multi_output(func):
|
||||
arguments.append(out_index_binding)
|
||||
|
||||
return arguments
|
||||
|
||||
|
||||
# `ViewMeta` specialized class' constructor arguments.
|
||||
#
|
||||
# Values needed specifically by this specialization, that the base class does not need.
|
||||
# Same as the class' attributes, but non-owning.
|
||||
def extra_ctor_arguments(func: FunctionSchema) -> list[Binding]:
|
||||
return attributes(func, owning=False)
|
||||
|
||||
|
||||
# `ViewMeta` specialized class' non-static member data.
|
||||
#
|
||||
# Essential data for calling the instance's `forward` and `reverse functions. You can
|
||||
# think of them as values that should be captured from the functionalization kernel.
|
||||
def attributes(func: FunctionSchema, owning: bool = True) -> list[Binding]:
|
||||
args = func.arguments.flat_all
|
||||
if args[0].type != BaseType(BaseTy.Tensor):
|
||||
raise AssertionError(f"Expected first arg to be Tensor, got {args[0].type}")
|
||||
return [
|
||||
reapply_views_binding,
|
||||
inverse_return_mode_binding,
|
||||
*[dispatcher.argument(a, remove_non_owning_ref_types=owning) for a in args[1:]],
|
||||
]
|
||||
|
||||
|
||||
def op_arguments(func: FunctionSchema, is_reverse: bool) -> list[Binding]:
|
||||
args = func.arguments.flat_all
|
||||
if args[0].type != BaseType(BaseTy.Tensor):
|
||||
raise AssertionError(f"Expected first arg to be Tensor, got {args[0].type}")
|
||||
non_self_args = args[1:]
|
||||
# The forward lambda calls the at::_ops API, while the reverse lambda calls the view inverse API.
|
||||
# Both of these follow the dispatcher API.
|
||||
non_self_bindings = [dispatcher.argument(a) for a in non_self_args]
|
||||
if not is_reverse:
|
||||
# the forward lambda swaps out the original tensor argument with the lambd arg "base"
|
||||
return [base_binding] + non_self_bindings
|
||||
else:
|
||||
# the reverse lambda does the same, but with an additional "mutated_view" arg
|
||||
# additionally, we have a calling convention: for view ops that return multiple tensor outputs
|
||||
# their corresponding view_inverse function takes in an additional index argument.
|
||||
if is_multi_output(func):
|
||||
return [
|
||||
base_binding,
|
||||
mutated_view_binding,
|
||||
inverse_return_mode_binding,
|
||||
out_index_binding,
|
||||
] + non_self_bindings
|
||||
else:
|
||||
return [
|
||||
base_binding,
|
||||
mutated_view_binding,
|
||||
inverse_return_mode_binding,
|
||||
] + non_self_bindings
|
||||
@@ -0,0 +1,470 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from torchgen.api.types import (
|
||||
BaseCppType,
|
||||
BaseCType,
|
||||
boolT,
|
||||
CType,
|
||||
deviceT,
|
||||
doubleT,
|
||||
generatorT,
|
||||
layoutT,
|
||||
ListCType,
|
||||
longT,
|
||||
memoryFormatT,
|
||||
NamedCType,
|
||||
OptionalCType,
|
||||
scalarT,
|
||||
scalarTypeT,
|
||||
stringT,
|
||||
SymIntT,
|
||||
VectorCType,
|
||||
)
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
OperatorName,
|
||||
OptionalType,
|
||||
Return,
|
||||
TensorOptionsArguments,
|
||||
Type,
|
||||
)
|
||||
|
||||
|
||||
_valueT: BaseCppType | None = None
|
||||
|
||||
|
||||
# A ValueT is an IR type which represents the computation of a Tensor. In other
|
||||
# words, a PyTorch user will do operations on lazy tensors, and each output lazy
|
||||
# tensor internally tracks a ValueT representing the IR node that would have
|
||||
# actually produced the value of this tensor for real.
|
||||
#
|
||||
# This is configurable because different lazy tensor backends (LTC vs XLA) will
|
||||
# have different IR representations. (Though, arguably, after unification they
|
||||
# shouldn't!)
|
||||
def getValueT() -> BaseCppType:
|
||||
global _valueT
|
||||
if not _valueT:
|
||||
raise NotImplementedError(
|
||||
"The value type needs to be set with setValueT() in run_gen_lazy_tensor()"
|
||||
)
|
||||
|
||||
return _valueT
|
||||
|
||||
|
||||
def setValueT(val: BaseCppType) -> None:
|
||||
global _valueT
|
||||
_valueT = val
|
||||
|
||||
|
||||
# this is a bad hack. I need to refactor the data model to represent each arg in the schema as an object,
|
||||
# making it easier to represent special properties of an arg.
|
||||
tensorListValueT = BaseCppType("torch::lazy", "Value")
|
||||
|
||||
|
||||
def process_ir_type(
|
||||
typ: Type, properties: LazyIrProperties, *, symint: bool
|
||||
) -> BaseCType | VectorCType | OptionalCType | ListCType:
|
||||
"""
|
||||
This function takes a type from NativeFunctions and converts it for use with
|
||||
lazy tensor codegen.
|
||||
|
||||
Type conversion for lazy currently consists of
|
||||
(1) changing at::Tensors into lazy::Values
|
||||
(2) wrapping everything in a BaseCType
|
||||
(3) making cpp-reference types into cpp-value types (e.g. vector instead of IntArrayRef)
|
||||
|
||||
(1) converts at::Tensors to lazy::Values (which wrap lazy::Nodes, with which Lazy IR represents tensors.)
|
||||
There is special handling for Optional[Tensor] or list[Tensor], etc- hence 'tensor-like'
|
||||
|
||||
This is incomplete- there are assertions in places that it's expected to need to add
|
||||
more types as the codegen is used with more operators.
|
||||
"""
|
||||
if isinstance(typ, BaseType):
|
||||
if typ.name == BaseTy.Tensor:
|
||||
return BaseCType(getValueT())
|
||||
elif typ.name == BaseTy.Scalar:
|
||||
if properties.TreatScalarsAsConstants:
|
||||
return BaseCType(scalarT)
|
||||
# at::scalar has special handling,
|
||||
# and is wrapped in an lazy::Value just like at::tensor
|
||||
return BaseCType(getValueT())
|
||||
elif typ.name == BaseTy.ScalarType:
|
||||
return BaseCType(scalarTypeT)
|
||||
elif typ.name == BaseTy.int:
|
||||
return BaseCType(longT)
|
||||
elif typ.name == BaseTy.SymInt:
|
||||
if symint:
|
||||
return BaseCType(getValueT())
|
||||
else:
|
||||
return BaseCType(longT)
|
||||
elif typ.name == BaseTy.bool:
|
||||
return BaseCType(boolT)
|
||||
elif typ.name == BaseTy.float:
|
||||
return BaseCType(doubleT)
|
||||
elif typ.name == BaseTy.str:
|
||||
return BaseCType(stringT)
|
||||
elif typ.name == BaseTy.Device:
|
||||
return BaseCType(deviceT)
|
||||
elif typ.name == BaseTy.Generator:
|
||||
return BaseCType(generatorT)
|
||||
elif typ.name == BaseTy.Layout:
|
||||
return BaseCType(layoutT)
|
||||
elif typ.name == BaseTy.MemoryFormat:
|
||||
return BaseCType(memoryFormatT)
|
||||
else:
|
||||
raise AssertionError(f"TODO add support for type {repr(typ)}")
|
||||
elif isinstance(typ, OptionalType):
|
||||
return OptionalCType(process_ir_type(typ.elem, properties, symint=symint))
|
||||
elif isinstance(typ, ListType):
|
||||
if str(typ.elem) == "Tensor?":
|
||||
# TODO(whc) is this actually correct? or should it use a Vector like above
|
||||
return ListCType(OptionalCType(BaseCType(getValueT())))
|
||||
elif str(typ.elem) == "Tensor":
|
||||
# this is a TensorList which comes in from GetTensorList as a Value
|
||||
return BaseCType(tensorListValueT)
|
||||
elif typ.elem == BaseType(BaseTy.SymInt):
|
||||
# TODO: return a value type. The problem here is analogous to
|
||||
# the problem with tensorListValueT: if you have SymInt[] you
|
||||
# cannot conveniently save the list of Value directly, as nodes
|
||||
# expect to save values as a vector for ALL arguments. So you
|
||||
# need a separate IR node that represents all of the size nodes
|
||||
# assembled into a list. I'm not an LTC dev so I don't want to
|
||||
# figure it out right now. Y'all figure it out...
|
||||
return VectorCType(BaseCType(longT))
|
||||
|
||||
else:
|
||||
return VectorCType(process_ir_type(typ.elem, properties, symint=symint))
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(typ)}")
|
||||
|
||||
|
||||
# TODO: Determining this based off of CType is bad; this should be computed
|
||||
# from Type directly; then the same logic as process_ir_type can be used
|
||||
#
|
||||
# Invariant: passed typ should be an *owning* CType (e.g., we will report
|
||||
# that ArrayRef<Value> is NOT a value type)
|
||||
def isValueType(typ: CType, properties: LazyIrProperties | None = None) -> bool:
|
||||
"""
|
||||
Given a type, determine if it is a Value-like type. This is equivalent to
|
||||
being Tensor-like, but assumes the type has already been transformed.
|
||||
"""
|
||||
if isinstance(typ, BaseCType):
|
||||
# I am regretting my naming conventions, but now we are wrapping at::scalar in
|
||||
# lazy value, while preserving other 'scalar' types as scalars in the IR
|
||||
treat_scalars_as_constants = properties and properties.TreatScalarsAsConstants
|
||||
return (
|
||||
typ.type == getValueT()
|
||||
or (typ.type == scalarT and not treat_scalars_as_constants)
|
||||
or typ.type == SymIntT
|
||||
)
|
||||
elif typ == VectorCType(BaseCType(SymIntT)):
|
||||
# TODO: report True for this
|
||||
return False
|
||||
elif isinstance(typ, (OptionalCType, ListCType, VectorCType)):
|
||||
return isValueType(typ.elem, properties)
|
||||
return False
|
||||
|
||||
|
||||
def isSymIntType(typ: Type) -> bool:
|
||||
return isinstance(typ, BaseType) and typ.name == BaseTy.SymInt
|
||||
|
||||
|
||||
def isWrappedScalarType(typ: Type) -> bool:
|
||||
"""
|
||||
Given a type, determine if it is a c10::scalar which we will wrap in a lazy Value.
|
||||
Since we literally change the type from scalarT to valueT, information is lost.
|
||||
This function helps build a list of wrapped scalars to save that information
|
||||
"""
|
||||
if isinstance(typ, BaseType):
|
||||
# I am regretting my naming conventions, but now we are wrapping at::scalar in
|
||||
# lazy value, while preserving other 'scalar' types as scalars in the IR
|
||||
return typ.name == BaseTy.Scalar
|
||||
elif isinstance(typ, (OptionalType, ListType)):
|
||||
return isWrappedScalarType(typ.elem)
|
||||
return False
|
||||
|
||||
|
||||
# TODO: dedupe with Type.is_generator_like
|
||||
def isGeneratorType(typ: Type) -> bool:
|
||||
if isinstance(typ, BaseType):
|
||||
return typ.name == BaseTy.Generator
|
||||
elif isinstance(typ, (OptionalType)):
|
||||
return isGeneratorType(typ.elem)
|
||||
return False
|
||||
|
||||
|
||||
# This class caches a few derived properties computed from an Argument
|
||||
# and LazyIrProperties
|
||||
class LazyArgument:
|
||||
name: str
|
||||
orig_type: Type
|
||||
lazy_type_: CType | None
|
||||
is_wrapped_scalar: bool
|
||||
is_generator: bool
|
||||
# TODO: this is lies, it is false for symint list
|
||||
is_symint_or_list: bool
|
||||
|
||||
# Whether or not we are treating this as symint or not
|
||||
symint: bool
|
||||
|
||||
# true if this argument is or contains a lazy IR value
|
||||
is_lazy_value: bool
|
||||
|
||||
def __init__(
|
||||
self, arg: Argument, properties: LazyIrProperties, *, symint: bool
|
||||
) -> None:
|
||||
self.name = arg.name
|
||||
self.orig_type = arg.type
|
||||
self.symint = symint
|
||||
self.is_optional = isinstance(arg.type, OptionalType)
|
||||
self.is_generator = isGeneratorType(arg.type)
|
||||
self.lazy_type_ = process_ir_type(arg.type, properties, symint=symint)
|
||||
self.is_wrapped_scalar = isWrappedScalarType(arg.type)
|
||||
self.is_symint_or_list = symint and (
|
||||
isSymIntType(arg.type)
|
||||
or (isinstance(arg.type, OptionalType) and isSymIntType(arg.type.elem))
|
||||
# TODO: lists of symints are not currently treated as value types
|
||||
# or (isinstance(arg.type, ListType) and isSymIntType(arg.type.elem))
|
||||
)
|
||||
|
||||
self.is_lazy_value = isValueType(self.lazy_type, properties)
|
||||
|
||||
@property
|
||||
def lazy_type(self) -> CType:
|
||||
if self.lazy_type_ is None:
|
||||
raise AssertionError(
|
||||
f"Attempted to access lazy_type for invalid argument {self.name}"
|
||||
)
|
||||
return self.lazy_type_
|
||||
|
||||
|
||||
class LazyIrProperties:
|
||||
"""Collection of properties for an IR node
|
||||
|
||||
The property groups are listed below. Each group is mutually
|
||||
exclusive, meaning that only one property from each group can be True
|
||||
at any one time. The properties can be accessed as if they were normal
|
||||
attributes. The mutual exclusivity is automatically handled.
|
||||
"""
|
||||
|
||||
Properties: tuple[tuple[str, ...], ...] = (
|
||||
(
|
||||
"ShapePrecompute", # Assume shape has been precomputed
|
||||
"ShapeCompute", # Need to compute the shape on construction
|
||||
"ShapeCache", # Utilize the shape cache to defer computation
|
||||
),
|
||||
(
|
||||
"Lower", # Codegen full lower function
|
||||
"LowerDeclOnly", # Codegen only lower function declaration
|
||||
),
|
||||
(
|
||||
"CanBeReused", # Codegen full reuse function
|
||||
"CanBeReusedDeclOnly", # Codegen only reuse function declaration
|
||||
),
|
||||
(
|
||||
"CreateFn", # Codegen full create function
|
||||
"CreateFnDeclOnly", # Codegen only create function declaration
|
||||
),
|
||||
(
|
||||
"TreatScalarsAsConstants", # Treat Scalars as constants instead of handling like values
|
||||
),
|
||||
)
|
||||
|
||||
def __init__(self, *default_properties: str) -> None:
|
||||
properties: dict[tuple[str, ...], str | None] = dict.fromkeys(
|
||||
LazyIrProperties.Properties
|
||||
)
|
||||
self.__dict__["properties"] = properties
|
||||
for p in default_properties:
|
||||
setattr(self, p, True)
|
||||
|
||||
def __getattr__(self, key: str) -> Any:
|
||||
properties = self.__dict__["properties"]
|
||||
for values in LazyIrProperties.Properties:
|
||||
if key in values:
|
||||
return properties[values] == key
|
||||
|
||||
return self.__getattribute__(key)
|
||||
|
||||
def __setattr__(self, key: str, value: Any) -> Any:
|
||||
properties = self.__dict__["properties"]
|
||||
for values in LazyIrProperties.Properties:
|
||||
if key in values:
|
||||
properties[values] = key if value else None
|
||||
return value
|
||||
|
||||
raise KeyError(f"Invalid property: {key}")
|
||||
|
||||
|
||||
# Inspired by a FunctionSchema object, a LazyIrSchema holds the schema of a Lazy IR node.
|
||||
# Unlike a FunctionSchema, it has no round-trippable string form (relating to the YAML),
|
||||
# but carries type information from a native FunctionSchema modified for use with IR nodes,
|
||||
# and preserving original argument names.
|
||||
#
|
||||
# TODO: This is not idiomatic with how other torchgen APIs transform on schema.
|
||||
class LazyIrSchema:
|
||||
# The name of the operator this function schema describes.
|
||||
name: OperatorName
|
||||
|
||||
positional_args: tuple[LazyArgument, ...]
|
||||
keyword_args: tuple[LazyArgument, ...]
|
||||
|
||||
# TODO: Need to handle collisions with argument names at some point
|
||||
returns: tuple[Return, ...]
|
||||
|
||||
# if this schema has a Generator arg, list its orig ctype/name but don't
|
||||
# build a LazyArgument since lazy IR doesn't support it
|
||||
generator_arg: NamedCType | None = None
|
||||
|
||||
# original function schema
|
||||
func: FunctionSchema
|
||||
|
||||
# Whether or not we are code-genning for SymInt or not
|
||||
symint: bool
|
||||
|
||||
properties: LazyIrProperties = LazyIrProperties(
|
||||
# default properties
|
||||
"ShapePrecompute",
|
||||
"Lower",
|
||||
"CanBeReused",
|
||||
)
|
||||
opkind: str | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
func: FunctionSchema,
|
||||
properties: LazyIrProperties | None = None,
|
||||
*,
|
||||
symint: bool,
|
||||
) -> None:
|
||||
if properties:
|
||||
self.properties = properties
|
||||
|
||||
self.func = func
|
||||
self.symint = symint
|
||||
positional_args: list[LazyArgument] = []
|
||||
for arg_field in ["pre_self_positional", "self_arg", "post_self_positional"]:
|
||||
if arg_field == "self_arg" and func.arguments.self_arg is not None:
|
||||
arg = func.arguments.self_arg.argument
|
||||
positional_args.append(
|
||||
LazyArgument(arg, self.properties, symint=symint)
|
||||
)
|
||||
elif getattr(func.arguments, arg_field) is not None:
|
||||
positional_args.extend(
|
||||
LazyArgument(arg, self.properties, symint=symint)
|
||||
for arg in getattr(func.arguments, arg_field)
|
||||
)
|
||||
self.positional_args = tuple(positional_args)
|
||||
|
||||
keyword_args: list[LazyArgument] = []
|
||||
for arg_field in [
|
||||
"pre_tensor_options_kwarg_only",
|
||||
"tensor_options",
|
||||
"post_tensor_options_kwarg_only",
|
||||
"out",
|
||||
]:
|
||||
curr_args = getattr(func.arguments, arg_field)
|
||||
if curr_args is not None:
|
||||
if isinstance(curr_args, TensorOptionsArguments):
|
||||
curr_args = curr_args.all()
|
||||
for arg in curr_args:
|
||||
if isGeneratorType(arg.type):
|
||||
if self.generator_arg is not None:
|
||||
raise AssertionError(
|
||||
"We expect there is only one generator arg"
|
||||
)
|
||||
self.generator_arg = NamedCType(
|
||||
arg.name,
|
||||
arg.type, # type:ignore[arg-type]
|
||||
)
|
||||
keyword_args.extend(
|
||||
LazyArgument(arg, self.properties, symint=symint)
|
||||
for arg in curr_args
|
||||
)
|
||||
self.keyword_args = tuple(keyword_args)
|
||||
self.name = func.name
|
||||
self.returns = func.returns
|
||||
|
||||
@property
|
||||
def node_name(self) -> str:
|
||||
"""
|
||||
Return camel-case version of op in node.
|
||||
|
||||
Note: This function also appends any `overload_name` in the operation.
|
||||
For example, if the op is `bitwise_and.Tensor`, the returned name
|
||||
will be `BitwiseAndTensor`.
|
||||
"""
|
||||
op_name = f"{self.name.name}_{self.name.overload_name}".lower()
|
||||
return "".join(word.capitalize() or "" for word in op_name.split("_"))
|
||||
|
||||
@property
|
||||
def aten_name(self) -> str:
|
||||
return str(self.name.name)
|
||||
|
||||
@property
|
||||
def base_name(self) -> str:
|
||||
return f"{self.name.name.base}"
|
||||
|
||||
def filtered_args(
|
||||
self,
|
||||
positional: bool = True,
|
||||
keyword: bool = True,
|
||||
values: bool = True,
|
||||
scalars: bool = True,
|
||||
generator: bool = True,
|
||||
) -> list[LazyArgument]:
|
||||
# This function maintains the sorted order of arguments but provides different filtered views.
|
||||
# Some parts of the code care about kwargs vs args (TS lowerings),
|
||||
# other parts care about whether they need to wrap the arg in a lazy value or leave it alone.
|
||||
# Generators are special cased, as they are needed for fallback/shape-inference but not supported
|
||||
# in TS lowerings and therefore also omitted from lazy IR.
|
||||
args: list[LazyArgument] = []
|
||||
if positional:
|
||||
args.extend(self.positional_args)
|
||||
if keyword:
|
||||
args.extend(self.keyword_args)
|
||||
|
||||
if values and scalars and generator:
|
||||
return args
|
||||
elif values and scalars:
|
||||
return [a for a in args if not a.is_generator]
|
||||
elif values:
|
||||
return [a for a in args if a.is_lazy_value]
|
||||
elif scalars:
|
||||
return [
|
||||
a
|
||||
for a in args
|
||||
if not a.is_lazy_value and (generator or not a.is_generator)
|
||||
]
|
||||
|
||||
return []
|
||||
|
||||
@property
|
||||
def positional_values(self) -> list[LazyArgument]:
|
||||
return self.filtered_args(
|
||||
positional=True, keyword=False, values=True, scalars=False
|
||||
)
|
||||
|
||||
@property
|
||||
def positional_scalars(self) -> list[LazyArgument]:
|
||||
return self.filtered_args(
|
||||
positional=True, keyword=False, values=False, scalars=True
|
||||
)
|
||||
|
||||
@property
|
||||
def keyword_values(self) -> list[LazyArgument]:
|
||||
return self.filtered_args(
|
||||
positional=False, keyword=True, values=True, scalars=False
|
||||
)
|
||||
|
||||
@property
|
||||
def keyword_scalars(self) -> list[LazyArgument]:
|
||||
return self.filtered_args(
|
||||
positional=False, keyword=True, values=False, scalars=True
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from torchgen.model import NativeFunctionsGroup
|
||||
|
||||
|
||||
# Follows dispatcher calling convention, but:
|
||||
# - Mutable arguments not allowed. Meta functions are always
|
||||
# written in functional form. Look at FunctionSchema.signature()
|
||||
# - No tensor returns; instead we return a TensorMeta describing
|
||||
# the tensor in question
|
||||
|
||||
|
||||
def name(g: NativeFunctionsGroup) -> str:
|
||||
# use the overload name from the functional version
|
||||
return str(g.functional.func.name).replace(".", "_")
|
||||
@@ -0,0 +1,159 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from torchgen import local
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import (
|
||||
ArgName,
|
||||
BaseCType,
|
||||
Binding,
|
||||
boolT,
|
||||
ConstRefCType,
|
||||
CType,
|
||||
deviceT,
|
||||
layoutT,
|
||||
ListCType,
|
||||
MutRefCType,
|
||||
NamedCType,
|
||||
OptionalCType,
|
||||
scalarT,
|
||||
scalarTypeT,
|
||||
tensorT,
|
||||
)
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
FunctionSchema,
|
||||
Return,
|
||||
SelfArgument,
|
||||
TensorOptionsArguments,
|
||||
Type,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# This file describes the translation of JIT schema to the native functions API.
|
||||
# This looks a lot like the C++ API (which makes historical sense, because the
|
||||
# idea was you wrote native functions to implement functions in the C++ API),
|
||||
# but over time we have evolved the C++ API without actually changing our
|
||||
# native:: kernels. The intention is to make native API and dispatcher API
|
||||
# line up as closely as possible, since this results in the least overhead
|
||||
# (no translation is needed from dispatcher API to native API).
|
||||
#
|
||||
# NB: this is symint aware, you will get the non-SymInt variant for some
|
||||
# dispatch entries and SymInt for others.
|
||||
|
||||
|
||||
def name(func: FunctionSchema) -> str:
|
||||
name = str(func.name.name)
|
||||
# TODO: delete this!
|
||||
if func.is_out_fn():
|
||||
name += "_out"
|
||||
if func.name.overload_name:
|
||||
name += f"_{func.name.overload_name}"
|
||||
return name
|
||||
|
||||
|
||||
def argumenttype_type(
|
||||
t: Type, *, mutable: bool, binds: ArgName, symint: bool
|
||||
) -> NamedCType:
|
||||
if str(t) == "Tensor?":
|
||||
tensor_type: OptionalCType = OptionalCType(BaseCType(tensorT))
|
||||
if mutable and not local.use_const_ref_for_mutable_tensors():
|
||||
return NamedCType(binds, MutRefCType(tensor_type))
|
||||
else:
|
||||
return NamedCType(binds, ConstRefCType(tensor_type))
|
||||
elif str(t) == "Tensor?[]":
|
||||
return NamedCType(
|
||||
binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT))))
|
||||
)
|
||||
elif str(t) == "Scalar":
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))
|
||||
elif str(t) == "Scalar?":
|
||||
return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT))))
|
||||
return cpp.argumenttype_type(t, mutable=mutable, binds=binds, symint=symint)
|
||||
|
||||
|
||||
def returns_type(rs: Sequence[Return], *, symint: bool) -> CType:
|
||||
return cpp.returns_type(rs, symint=symint)
|
||||
|
||||
|
||||
def argument_type(a: Argument, *, binds: ArgName, symint: bool) -> NamedCType:
|
||||
return argumenttype_type(a.type, mutable=a.is_write, binds=binds, symint=symint)
|
||||
|
||||
|
||||
def argument(
|
||||
a: Argument | SelfArgument | TensorOptionsArguments,
|
||||
*,
|
||||
is_out: bool,
|
||||
symint: bool,
|
||||
) -> list[Binding]:
|
||||
# Ideally, we NEVER default native functions. However, there are a number
|
||||
# of functions that call native:: directly and rely on the defaulting
|
||||
# existing. So for BC, we generate defaults for non-out variants (but not
|
||||
# for out variants, where it is impossible to generate an appropriate
|
||||
# default)
|
||||
should_default = not is_out
|
||||
if isinstance(a, Argument):
|
||||
default: str | None = None
|
||||
if should_default and a.default is not None:
|
||||
default = cpp.default_expr(a.default, a.type, symint=symint)
|
||||
return [
|
||||
Binding(
|
||||
nctype=argument_type(a, binds=a.name, symint=symint),
|
||||
name=a.name,
|
||||
default=default,
|
||||
argument=a,
|
||||
)
|
||||
]
|
||||
elif isinstance(a, SelfArgument):
|
||||
# Erase SelfArgument from the distinction
|
||||
return argument(a.argument, is_out=is_out, symint=symint)
|
||||
elif isinstance(a, TensorOptionsArguments):
|
||||
default = None
|
||||
if should_default:
|
||||
default = "{}"
|
||||
# TODO: Not sure why the arguments assigned here are for
|
||||
# TensorOptionsArguments and not the constituent pieces. It seems
|
||||
# to matter
|
||||
return [
|
||||
Binding(
|
||||
nctype=NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT))),
|
||||
name="dtype",
|
||||
default=default,
|
||||
argument=a,
|
||||
),
|
||||
Binding(
|
||||
nctype=NamedCType("layout", OptionalCType(BaseCType(layoutT))),
|
||||
name="layout",
|
||||
default=default,
|
||||
argument=a,
|
||||
),
|
||||
Binding(
|
||||
nctype=NamedCType("device", OptionalCType(BaseCType(deviceT))),
|
||||
name="device",
|
||||
default=default,
|
||||
argument=a,
|
||||
),
|
||||
Binding(
|
||||
nctype=NamedCType("pin_memory", OptionalCType(BaseCType(boolT))),
|
||||
name="pin_memory",
|
||||
default=default,
|
||||
argument=a,
|
||||
),
|
||||
]
|
||||
else:
|
||||
assert_never(a)
|
||||
|
||||
|
||||
def arguments(func: FunctionSchema, *, symint: bool) -> list[Binding]:
|
||||
args: list[Argument | TensorOptionsArguments | SelfArgument] = []
|
||||
args.extend(func.arguments.non_out)
|
||||
args.extend(func.arguments.out)
|
||||
return [
|
||||
r for arg in args for r in argument(arg, symint=symint, is_out=func.is_out_fn())
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,158 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import (
|
||||
ArgName,
|
||||
ArrayRefCType,
|
||||
BaseCType,
|
||||
Binding,
|
||||
ConstRefCType,
|
||||
dimnameListT,
|
||||
intArrayRefT,
|
||||
iOptTensorListRefT,
|
||||
iTensorListRefT,
|
||||
NamedCType,
|
||||
OptionalCType,
|
||||
optionalIntArrayRefT,
|
||||
optionalScalarRefT,
|
||||
optionalTensorRefT,
|
||||
scalarT,
|
||||
tensorT,
|
||||
)
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
ListType,
|
||||
NativeFunctionsGroup,
|
||||
OptionalType,
|
||||
SelfArgument,
|
||||
TensorOptionsArguments,
|
||||
Type,
|
||||
)
|
||||
|
||||
|
||||
# This file describes the translation of JIT schema to the structured functions API.
|
||||
# This is similar to native API, but a number of historical problems with native
|
||||
# API have been fixed.
|
||||
|
||||
|
||||
# Translation of types occurring in JIT arguments to a C++ argument type.
|
||||
# NB: For now, mutable doesn't do anything; but it could if we make
|
||||
# some more nominal types
|
||||
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> NamedCType:
|
||||
# If it's a value type, do the value type translation
|
||||
# NB: structured kernels ALWAYS have symint off, since they involve actual
|
||||
# kernels that require real ints. The one exception is the
|
||||
# CompositeExplicitAutograd and the meta function (which could
|
||||
# hypothetically be SymInt), but for simplicity we plan for these to just
|
||||
# be handled in Python
|
||||
r = cpp.valuetype_type(t, symint=False, binds=binds, mutable=mutable)
|
||||
if r is not None:
|
||||
return r
|
||||
|
||||
if isinstance(t, BaseType):
|
||||
if t.name == BaseTy.Tensor:
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(tensorT)))
|
||||
elif t.name == BaseTy.Scalar:
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))
|
||||
else:
|
||||
raise AssertionError(f"base type should have been value type {t}")
|
||||
elif isinstance(t, OptionalType):
|
||||
if t.elem == BaseType(BaseTy.Tensor):
|
||||
return NamedCType(binds, BaseCType(optionalTensorRefT))
|
||||
elif t.elem == BaseType(BaseTy.Scalar):
|
||||
return NamedCType(binds, BaseCType(optionalScalarRefT))
|
||||
elif isinstance(t.elem, ListType) and str(t.elem.elem) == "int":
|
||||
return NamedCType(binds, BaseCType(optionalIntArrayRefT))
|
||||
elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)
|
||||
return NamedCType(binds, OptionalCType(elem.type))
|
||||
elif isinstance(t, ListType):
|
||||
if t.elem == BaseType(BaseTy.Tensor):
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(iTensorListRefT)))
|
||||
elif t.elem == OptionalType(BaseType(BaseTy.Tensor)):
|
||||
return NamedCType(binds, BaseCType(iOptTensorListRefT))
|
||||
# TODO: delete these special cases; see torchgen.api.cpp--these
|
||||
# must be changed in tandem, but there are problems; see
|
||||
# https://github.com/pytorch/pytorch/pull/51485
|
||||
elif str(t.elem) == "int":
|
||||
return NamedCType(binds, BaseCType(intArrayRefT))
|
||||
elif str(t.elem) == "Dimname":
|
||||
return NamedCType(binds, BaseCType(dimnameListT))
|
||||
elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)
|
||||
return NamedCType(binds, ArrayRefCType(elem.type))
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
def argument_type(a: Argument, *, binds: ArgName) -> NamedCType:
|
||||
return argumenttype_type(a.type, mutable=a.is_write, binds=binds)
|
||||
|
||||
|
||||
# returns_type intentionally omitted, because structured kernels never "return";
|
||||
# instead, they always indirectly report their outputs (in the case of a meta
|
||||
# function, by calling set_output; in the case of an impl function, by writing
|
||||
# directly into the provided out argument).
|
||||
|
||||
|
||||
# Structured kernels are never defaulted
|
||||
def argument(a: Argument | SelfArgument | TensorOptionsArguments) -> list[Binding]:
|
||||
if isinstance(a, Argument):
|
||||
return [
|
||||
Binding(
|
||||
nctype=argument_type(a, binds=a.name),
|
||||
name=a.name,
|
||||
default=None,
|
||||
argument=a,
|
||||
)
|
||||
]
|
||||
elif isinstance(a, SelfArgument):
|
||||
return argument(a.argument)
|
||||
elif isinstance(a, TensorOptionsArguments):
|
||||
raise AssertionError("structured kernels don't support TensorOptions yet")
|
||||
else:
|
||||
assert_never(a)
|
||||
|
||||
|
||||
def impl_arguments(g: NativeFunctionsGroup) -> list[Binding]:
|
||||
args: list[Argument | TensorOptionsArguments | SelfArgument] = []
|
||||
|
||||
if g.out.precomputed:
|
||||
# A list of parameters for the impl function with
|
||||
# certain parameters replaced with precomputed counterparts
|
||||
# as specified in native_functions.yaml.
|
||||
non_out_args_replaced: list[
|
||||
Argument | TensorOptionsArguments | SelfArgument
|
||||
] = []
|
||||
for a in g.out.func.arguments.non_out:
|
||||
if isinstance(a, Argument) and a.name in g.out.precomputed.replace:
|
||||
# If a is in precompute.replace, append the parameters
|
||||
# that should replace it onto non_out_args_replaced.
|
||||
non_out_args_replaced.extend(g.out.precomputed.replace[a.name])
|
||||
else:
|
||||
# If not, push a as it is.
|
||||
non_out_args_replaced.append(a)
|
||||
|
||||
args.extend(non_out_args_replaced)
|
||||
# g.out.precomputed.add is the list of parameters that are added
|
||||
# without replacement after the non out args and just before the out args
|
||||
args.extend(g.out.precomputed.add)
|
||||
else:
|
||||
args.extend(g.out.func.arguments.non_out)
|
||||
|
||||
args.extend(g.out.func.arguments.out)
|
||||
return [r for arg in args for r in argument(arg)]
|
||||
|
||||
|
||||
def meta_arguments(g: NativeFunctionsGroup) -> list[Binding]:
|
||||
args: list[Argument | TensorOptionsArguments | SelfArgument] = []
|
||||
args.extend(g.functional.func.arguments.non_out)
|
||||
return [r for arg in args for r in argument(arg)]
|
||||
|
||||
|
||||
def out_arguments(g: NativeFunctionsGroup) -> list[Binding]:
|
||||
args: list[Argument | TensorOptionsArguments | SelfArgument] = []
|
||||
args.extend(g.out.func.arguments.out)
|
||||
return [r for arg in args for r in argument(arg)]
|
||||
@@ -0,0 +1,437 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import NoReturn, TYPE_CHECKING
|
||||
|
||||
from torchgen.api.types import (
|
||||
ArrayRefCType,
|
||||
BaseCType,
|
||||
Binding,
|
||||
boolT,
|
||||
ConstRefCType,
|
||||
deviceT,
|
||||
Expr,
|
||||
intArrayRefT,
|
||||
iOptTensorListRefT,
|
||||
layoutT,
|
||||
ListCType,
|
||||
longT,
|
||||
memoryFormatT,
|
||||
MutRefCType,
|
||||
NamedCType,
|
||||
opmath_t,
|
||||
OptionalCType,
|
||||
optionalIntArrayRefT,
|
||||
optionalScalarRefT,
|
||||
optionalSymIntArrayRefT,
|
||||
optionalTensorRefT,
|
||||
scalar_t,
|
||||
scalarT,
|
||||
scalarTypeT,
|
||||
SpecialArgName,
|
||||
symIntArrayRefT,
|
||||
SymIntT,
|
||||
tensorOptionsT,
|
||||
tensorT,
|
||||
VectorCType,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# This file implements a small program synthesis engine that implements
|
||||
# conversions between one API to another.
|
||||
#
|
||||
# The key data type in this file in NamedCType, short for Named C++ semantic type. A NamedCType
|
||||
# represents a C++ type, plus semantic information about what it represents.
|
||||
# For example, consider the argument "bool pin_memory"; its normal C++ type is
|
||||
# "bool", but its C++ semantic type also keeps track that this represents a
|
||||
# "pin_memory"; you can't just use a random other boolean in a context where you
|
||||
# need a "pin_memory"!
|
||||
#
|
||||
# The translator takes a list of needed NamedCTypes, and then figures out how
|
||||
# to construct expressions with these NamedCTypes from the given bindings. Many
|
||||
# of these expressions are trivial (I need a Tensor other; there's a Tensor
|
||||
# other scope); others are more nontrivial and may require packing/unpacking.
|
||||
# Some examples of non-trivial action:
|
||||
#
|
||||
# - Need the "dtype" binding? Well, maybe "dtype" isn't available
|
||||
# in the context, instead, "options" is, and you need to extract
|
||||
# it from there. (Gather)
|
||||
#
|
||||
# - Need the "context" binding? Well, maybe "context" isn't available
|
||||
# in the context, and you need to construct it from "dtype", "device",
|
||||
# etc. (Scatter)
|
||||
#
|
||||
# - Need the "memory_format" binding? Well, actually, it's available
|
||||
# from both "memory_format" and "options", so you had better make sure
|
||||
# they are consistent. (Join)
|
||||
|
||||
options_ctype = NamedCType("options", ConstRefCType(BaseCType(tensorOptionsT)))
|
||||
|
||||
out_tensor_ctype = NamedCType("out", ConstRefCType(BaseCType(tensorT)))
|
||||
|
||||
longVec_ctype = VectorCType(BaseCType(longT))
|
||||
longSymVec_ctype = VectorCType(BaseCType(SymIntT))
|
||||
optionalLongVec_ctype = OptionalCType(VectorCType(BaseCType(longT)))
|
||||
optionalScalar_ctype = OptionalCType(BaseCType(scalarT))
|
||||
optionalTensor_ctype = OptionalCType(BaseCType(tensorT))
|
||||
|
||||
|
||||
class UnsatError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
# Given a set of in-scope bindings and a set of target bindings, synthesize
|
||||
# a list of expressions that uses only the in-scope bindings (bindings) that
|
||||
# have all of the types of goals. You may want to use this function if
|
||||
# you're generating code for a function like:
|
||||
#
|
||||
# void f({args}) {
|
||||
# g({exprs}); // g is a different API
|
||||
# }
|
||||
#
|
||||
# and you need to generate "exprs".
|
||||
#
|
||||
# Typically, a list of Bindings is convenient to get (you usually call something
|
||||
# like arguments() to get them); but technically you only need less information:
|
||||
# for 'bindings' an (un-ordered) list of Exprs is sufficient; similarly, for
|
||||
# 'goals', an (ordered) list of NamedCType goals is sufficient. If you are doing
|
||||
# something more complicated, e.g., tracking the set of bindings in a context,
|
||||
# you may find using these smaller types more convenient.
|
||||
def translate(
|
||||
bindings: Sequence[Expr | Binding],
|
||||
goals: Sequence[NamedCType | Binding],
|
||||
*,
|
||||
method: bool = False,
|
||||
allow_expensive_conversions: bool = False,
|
||||
) -> list[Expr]:
|
||||
binding_exprs: list[Expr] = []
|
||||
for b in bindings:
|
||||
if isinstance(b, Binding):
|
||||
binding_exprs.append(
|
||||
Expr(
|
||||
expr=b.name,
|
||||
type=b.nctype,
|
||||
)
|
||||
)
|
||||
else:
|
||||
binding_exprs.append(b)
|
||||
|
||||
goal_ctypes: list[NamedCType] = []
|
||||
for g in goals:
|
||||
if isinstance(g, Binding):
|
||||
goal_ctypes.append(g.nctype)
|
||||
else:
|
||||
goal_ctypes.append(g)
|
||||
|
||||
# Add all the bindings to the context
|
||||
ctx: dict[NamedCType, str] = {}
|
||||
for b in binding_exprs:
|
||||
ctx[b.type] = b.expr
|
||||
|
||||
# While we're at it, do some simple forward inference, looking through
|
||||
# constructors.
|
||||
#
|
||||
# NB: When should you do forward inference versus backward inference?
|
||||
# The general idea:
|
||||
#
|
||||
# - Backward inference WHEN the goal gets smaller
|
||||
# - Forward inference WHEN the hypothesis gets smaller
|
||||
#
|
||||
# This helps ensure termination: backward inference starts with a goal
|
||||
# and tries to make it simpler and simpler until it's trivial; if the
|
||||
# goal can grow in size, we blow up to a really huge goal size.
|
||||
# Similarly, with forward inference we take hypotheses and decompose
|
||||
# them into simpler hypotheses; if hypotheses could expand in size,
|
||||
# we also have potential nontermination. (In the code below, forward
|
||||
# inference is only ever carried out at a single step, but you could
|
||||
# imagine repeated application of forward inference being profitable.)
|
||||
#
|
||||
# A good starting point in the literature for exploring more about proof
|
||||
# search are these lecture notes
|
||||
# https://www.cs.cmu.edu/~fp/courses/oregon-m10/04-focusing.pdf
|
||||
#
|
||||
# TODO: My kingdom for a pattern matcher
|
||||
# https://www.python.org/dev/peps/pep-0634/
|
||||
#
|
||||
# TODO: This could get us in recomputation trouble if b.expr is nontrivial.
|
||||
# Fix this by implementing some sort of sharing so that if multiple
|
||||
# goals share the same expression, we only compute it once. This seems
|
||||
# to matter in practice as compiler is often unwilling to CSE nontrivial
|
||||
# expressions like scalar.to<scalar_t>()
|
||||
t = b.type
|
||||
if (
|
||||
isinstance(t, ConstRefCType)
|
||||
and isinstance(t.elem, OptionalCType)
|
||||
and isinstance(t.elem.elem, BaseCType)
|
||||
and str(t.elem.elem.type) == "at::Tensor"
|
||||
):
|
||||
ctx[NamedCType(t.elem.elem.name, ConstRefCType(BaseCType(tensorT)))] = (
|
||||
f"({b.expr}.has_value() ? *{b.expr} : at::Tensor())"
|
||||
)
|
||||
|
||||
if t.type == ConstRefCType(OptionalCType(BaseCType(tensorT))):
|
||||
ctx[NamedCType(t.name, BaseCType(optionalTensorRefT))] = (
|
||||
f"(({b.expr}.has_value() && (*{b.expr}).defined()) ? at::OptionalTensorRef(*{b.expr}) : at::OptionalTensorRef())"
|
||||
)
|
||||
|
||||
if t.type == ConstRefCType(BaseCType(scalarT)):
|
||||
ctx[NamedCType(t.name, BaseCType(opmath_t))] = f"({b.expr}).to<opmath_t>()"
|
||||
|
||||
if t.type == ConstRefCType(OptionalCType(BaseCType(scalarT))):
|
||||
ctx[NamedCType(t.name, BaseCType(optionalScalarRefT))] = (
|
||||
f"({b.expr}.has_value() ? at::OptionalScalarRef(&({b.expr}.value())) : at::OptionalScalarRef())"
|
||||
)
|
||||
|
||||
if t.type == BaseCType(scalar_t):
|
||||
ctx[NamedCType(t.name, BaseCType(opmath_t))] = (
|
||||
f"static_cast<opmath_t>({b.expr})"
|
||||
)
|
||||
|
||||
# [Note: IOptTensorListRef]
|
||||
if t.type == ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT)))):
|
||||
ctx[NamedCType(t.name, BaseCType(iOptTensorListRefT))] = (
|
||||
f"at::IOptTensorListRef({b.expr})"
|
||||
)
|
||||
|
||||
# Add implicit bindings if the generated code is inside a Tensor method
|
||||
if method:
|
||||
ctx[NamedCType("self", MutRefCType(BaseCType(tensorT)))] = (
|
||||
"const_cast<Tensor&>(*this)"
|
||||
)
|
||||
ctx[NamedCType("self", ConstRefCType(BaseCType(tensorT)))] = (
|
||||
"const_cast<Tensor&>(*this)"
|
||||
)
|
||||
# This is better! Byte-for-byte compat
|
||||
# ctx[NamedCType("self", ConstRefCType(BaseCType(tensorT)))] = "*this"
|
||||
|
||||
def unsat(goal: NamedCType) -> NoReturn:
|
||||
ctx_desc = "\n".join(
|
||||
f" {t.cpp_type()} {t.name}; // {e}" for t, e in ctx.items()
|
||||
)
|
||||
raise UnsatError(
|
||||
f"""
|
||||
Failed to synthesize the expression "{goal.cpp_type()} {goal.name}".
|
||||
When I failed, the following bindings were available in the context:
|
||||
|
||||
{ctx_desc}
|
||||
|
||||
This probably means there is a missing rule in the rules of torchgen.api.translate.
|
||||
Check this module for more information.
|
||||
"""
|
||||
)
|
||||
|
||||
# A shitty backtracking search implementation. It's shitty because it
|
||||
# does backtracking via stack (bad idea!) and for the most part tries to
|
||||
# avoid backtracking. In particular, if
|
||||
# direct=True, we won't try to do any fancy synthesis, just trivial
|
||||
# conversions (e.g., "T a" is OK for "const T& a"). So all of the
|
||||
# existing rules in this function simply try to solve immediately,
|
||||
# and bail if things don't work out.
|
||||
def solve(goal: NamedCType, *, direct: bool) -> str:
|
||||
def direct_solve(goal: NamedCType) -> str:
|
||||
return solve(goal, direct=True)
|
||||
|
||||
if goal in ctx:
|
||||
# Trivial
|
||||
return ctx[goal]
|
||||
|
||||
# const & is satisfied with mutable &
|
||||
if isinstance(goal.type, ConstRefCType):
|
||||
try:
|
||||
# WARNING: not strictly decreasing; be careful not
|
||||
# to add a direct conversion that goes satisfies
|
||||
# mutable& with const&
|
||||
return solve(
|
||||
NamedCType(goal.name, MutRefCType(goal.type.elem)), direct=direct
|
||||
)
|
||||
except UnsatError:
|
||||
pass
|
||||
|
||||
# mutable & is satisfied with value
|
||||
if isinstance(goal.type, MutRefCType):
|
||||
try:
|
||||
return solve(NamedCType(goal.name, goal.type.elem), direct=direct)
|
||||
except UnsatError:
|
||||
pass
|
||||
|
||||
# TODO: These are referentially equal, shouldn't have to do this;
|
||||
# ensuring we don't use type synonym IntArrayRef in codegen would
|
||||
# help
|
||||
if goal.type == ArrayRefCType(BaseCType(longT)):
|
||||
return solve(NamedCType(goal.name, BaseCType(intArrayRefT)), direct=direct)
|
||||
|
||||
if direct:
|
||||
unsat(goal)
|
||||
|
||||
# For now, all of these rules are mutually exclusive.
|
||||
if goal == NamedCType("memory_format", OptionalCType(BaseCType(memoryFormatT))):
|
||||
memory_format = direct_solve(
|
||||
NamedCType(
|
||||
SpecialArgName.possibly_redundant_memory_format,
|
||||
OptionalCType(BaseCType(memoryFormatT)),
|
||||
)
|
||||
)
|
||||
# No need to join "memory_format" and "options" if the target API takes "options" directly.
|
||||
# Otherwise it will cause the redundant memory_format error.
|
||||
if options_ctype in goal_ctypes:
|
||||
return memory_format
|
||||
try:
|
||||
options = direct_solve(options_ctype)
|
||||
return f"c10::impl::check_tensor_options_and_extract_memory_format({options}, {memory_format})"
|
||||
except UnsatError:
|
||||
return memory_format
|
||||
elif goal == NamedCType("options", BaseCType(tensorOptionsT)):
|
||||
dtype = direct_solve(
|
||||
NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT)))
|
||||
)
|
||||
pin_memory = direct_solve(
|
||||
NamedCType("pin_memory", OptionalCType(BaseCType(boolT)))
|
||||
)
|
||||
device = direct_solve(
|
||||
NamedCType("device", OptionalCType(BaseCType(deviceT)))
|
||||
)
|
||||
layout = direct_solve(
|
||||
NamedCType("layout", OptionalCType(BaseCType(layoutT)))
|
||||
)
|
||||
return f"TensorOptions().dtype({dtype}).layout({layout}).device({device}).pinned_memory({pin_memory})"
|
||||
|
||||
elif goal == NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT))):
|
||||
try:
|
||||
options = direct_solve(options_ctype)
|
||||
return f"c10::optTypeMetaToScalarType({options}.dtype_opt())"
|
||||
except UnsatError:
|
||||
out_tensor = direct_solve(out_tensor_ctype)
|
||||
return f"{out_tensor}.scalar_type()"
|
||||
|
||||
elif goal == NamedCType("layout", OptionalCType(BaseCType(layoutT))):
|
||||
try:
|
||||
options = direct_solve(options_ctype)
|
||||
return f"{options}.layout_opt()"
|
||||
except UnsatError:
|
||||
out_tensor = direct_solve(out_tensor_ctype)
|
||||
return f"{out_tensor}.layout()"
|
||||
|
||||
elif goal == NamedCType("device", OptionalCType(BaseCType(deviceT))):
|
||||
try:
|
||||
options = direct_solve(options_ctype)
|
||||
return f"{options}.device_opt()"
|
||||
except UnsatError:
|
||||
out_tensor = direct_solve(out_tensor_ctype)
|
||||
return f"{out_tensor}.device()"
|
||||
|
||||
elif goal == NamedCType("pin_memory", OptionalCType(BaseCType(boolT))):
|
||||
try:
|
||||
options = direct_solve(options_ctype)
|
||||
return f"{options}.pinned_memory_opt()"
|
||||
except UnsatError:
|
||||
# If we're calling a factory op from its out= variant,
|
||||
# We don't actually care about the value of pin_memory.
|
||||
out_tensor = direct_solve(out_tensor_ctype)
|
||||
return "::std::nullopt"
|
||||
|
||||
# We can always do translations from value types to reference types, like vector<int> -> IntArrayRef
|
||||
elif goal.type == BaseCType(intArrayRefT):
|
||||
try:
|
||||
return direct_solve(NamedCType(goal.name, longVec_ctype))
|
||||
except UnsatError:
|
||||
# We can also go SymIntArrayRef -> IntArrayRef
|
||||
symIntArrayRef_type = direct_solve(
|
||||
NamedCType(goal.name, BaseCType(symIntArrayRefT))
|
||||
)
|
||||
return f"C10_AS_INTARRAYREF_SLOW({symIntArrayRef_type})"
|
||||
elif goal.type == BaseCType(symIntArrayRefT):
|
||||
try:
|
||||
r = direct_solve(NamedCType(goal.name, BaseCType(intArrayRefT)))
|
||||
return f"c10::fromIntArrayRefSlow({r})"
|
||||
except UnsatError:
|
||||
return direct_solve(NamedCType(goal.name, longSymVec_ctype))
|
||||
elif goal.type == BaseCType(SymIntT):
|
||||
return direct_solve(NamedCType(goal.name, BaseCType(longT)))
|
||||
elif goal.type == OptionalCType(BaseCType(SymIntT)):
|
||||
argname = direct_solve(
|
||||
NamedCType(goal.name, OptionalCType(BaseCType(longT)))
|
||||
)
|
||||
return f"{argname}.has_value() ? ::std::make_optional(c10::SymInt(*{argname})) : ::std::nullopt"
|
||||
elif goal.type == BaseCType(longT):
|
||||
symInt_type = direct_solve(NamedCType(goal.name, BaseCType(SymIntT)))
|
||||
return f"{symInt_type}.guard_int(__FILE__, __LINE__)"
|
||||
elif goal.type == OptionalCType(BaseCType(longT)):
|
||||
argname = direct_solve(
|
||||
NamedCType(goal.name, OptionalCType(BaseCType(SymIntT)))
|
||||
)
|
||||
return f"{argname}.has_value() ? ::std::make_optional({argname}->guard_int(__FILE__, __LINE__)) : ::std::nullopt"
|
||||
elif goal.type == BaseCType(optionalIntArrayRefT):
|
||||
try:
|
||||
return direct_solve(NamedCType(goal.name, optionalLongVec_ctype))
|
||||
except UnsatError:
|
||||
argname = direct_solve(
|
||||
NamedCType(goal.name, BaseCType(optionalSymIntArrayRefT))
|
||||
)
|
||||
return f"{argname}.has_value() ? ::std::make_optional(C10_AS_INTARRAYREF_SLOW(*{argname})) : ::std::nullopt"
|
||||
elif goal.type == BaseCType(optionalSymIntArrayRefT):
|
||||
# TODO: You might also want to solve this from longSymVec_ctype or
|
||||
# an optional version of it
|
||||
argname = direct_solve(
|
||||
NamedCType(goal.name, BaseCType(optionalIntArrayRefT))
|
||||
)
|
||||
return f"{argname}.has_value() ? ::std::make_optional(c10::fromIntArrayRefSlow(*{argname})) : ::std::nullopt"
|
||||
elif goal.type == BaseCType(optionalScalarRefT):
|
||||
return direct_solve(NamedCType(goal.name, optionalScalar_ctype))
|
||||
elif goal.type == BaseCType(optionalTensorRefT):
|
||||
return direct_solve(NamedCType(goal.name, optionalTensor_ctype))
|
||||
|
||||
# Note [translation from C++ reference to value types]
|
||||
# The below cases are all for when we have an argument with a reference type,
|
||||
# and a corresponding goal with a value type.
|
||||
# These are needed when we populate the inputs to a lambda capture and we need
|
||||
# to guarantee the lifetime of each captured argument.
|
||||
# We guard it with an explicit kwarg because converting to a value type is expensive
|
||||
# (O(n)) to convert from IntArrayRef to vector<int>),
|
||||
# so the caller of translate() should be explicit that they need it.
|
||||
if allow_expensive_conversions:
|
||||
if goal.type == VectorCType(BaseCType(longT)):
|
||||
intArrayRef_ctype = NamedCType(goal.name, BaseCType(intArrayRefT))
|
||||
argname = direct_solve(intArrayRef_ctype)
|
||||
return f"{argname}.vec()"
|
||||
if goal.type == VectorCType(BaseCType(SymIntT)):
|
||||
symIntArrayRef_ctype = NamedCType(goal.name, BaseCType(symIntArrayRefT))
|
||||
argname = direct_solve(symIntArrayRef_ctype)
|
||||
return f"{argname}.vec()"
|
||||
elif goal.type == OptionalCType(VectorCType(BaseCType(longT))):
|
||||
optionalIntArrayRef_ctype = NamedCType(
|
||||
goal.name, BaseCType(optionalIntArrayRefT)
|
||||
)
|
||||
argname = direct_solve(optionalIntArrayRef_ctype)
|
||||
return f"{argname}.has_value() ? ::std::make_optional({argname}->vec()) : ::std::nullopt"
|
||||
elif goal.type == OptionalCType(BaseCType(scalarT)):
|
||||
optionalScalarRef_ctype = NamedCType(
|
||||
goal.name, BaseCType(optionalScalarRefT)
|
||||
)
|
||||
argname = direct_solve(optionalScalarRef_ctype)
|
||||
return f"{argname}.has_value() ? ::std::make_optional({argname}) : ::std::nullopt"
|
||||
elif goal.type == OptionalCType(BaseCType(scalarT)):
|
||||
optionalTensorRef_ctype = NamedCType(
|
||||
goal.name, BaseCType(optionalTensorRefT)
|
||||
)
|
||||
argname = direct_solve(optionalTensorRef_ctype)
|
||||
return f"{argname}.has_value() ? ::std::make_optional({argname}) : ::std::nullopt"
|
||||
# Technically, we also need to handle cases of C++ containers holding reference types.
|
||||
# But there currently aren't any ops that require lambda capture codegen
|
||||
# With arguments like ::std::vector<IntArrayRef>.
|
||||
# If that changes, we'll have to add the translation here.
|
||||
|
||||
# We allow const casting on tensors, since const-correctness is a bit broken for at::Tensor.
|
||||
# We could probably generalize this to non-tensor types too.
|
||||
if goal.type == MutRefCType(BaseCType(tensorT)):
|
||||
const_ref_tensor_ctype = NamedCType(
|
||||
goal.name, ConstRefCType(BaseCType(tensorT))
|
||||
)
|
||||
argname = direct_solve(const_ref_tensor_ctype)
|
||||
return f"const_cast<Tensor&>({argname})"
|
||||
|
||||
unsat(goal)
|
||||
|
||||
return [Expr(solve(g, direct=False), g) for g in goal_ctypes]
|
||||
@@ -0,0 +1,5 @@
|
||||
from torchgen.api.types.types import *
|
||||
from torchgen.api.types.types_base import *
|
||||
|
||||
|
||||
from torchgen.api.types.signatures import * # usort: skip
|
||||
@@ -0,0 +1,358 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from torchgen.api.types.types_base import Binding, CType, Expr
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator, Sequence
|
||||
|
||||
from torchgen.model import (
|
||||
BackendIndex,
|
||||
FunctionSchema,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
NativeFunctionsViewGroup,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CppSignature:
|
||||
"""
|
||||
A CppSignature represents a single overload in the C++ API. For
|
||||
any given function schema, there may be multiple CppSignatures
|
||||
corresponding to it, based on how we desugar to C++. See also
|
||||
CppSignatureGroup.
|
||||
"""
|
||||
|
||||
# The schema this signature is derived from
|
||||
func: FunctionSchema
|
||||
|
||||
# Is this a C++ signature for a method, i.e. Tensor::my_op(...)?
|
||||
method: bool
|
||||
|
||||
# Is this a faithful C++ signature (i.e. following the JIT schema) or a convenience API
|
||||
# (i.e. with a potential TensorOptions argument and out arguments in the front)
|
||||
faithful: bool
|
||||
|
||||
# Is this a symint C++ signature. For BC reasons, functions that take
|
||||
# SymInts still present as int64_t in C++, and the SymInt variant is
|
||||
# offered at a different overload name
|
||||
#
|
||||
# NB: If a function RETURNS a SymInt, this is ALWAYS false
|
||||
symint: bool
|
||||
|
||||
# The set of C++ arguments which should not have defaults applied to them
|
||||
cpp_no_default_args: set[str]
|
||||
|
||||
# Is this a fallback C++ binding? Fallback bindings are enabled by
|
||||
# manual_cpp_binding: True and are alternate, non-public API that
|
||||
# lets manual C++ binding implementers access the binding that would
|
||||
# have been automatically generated
|
||||
fallback_binding: bool = False
|
||||
|
||||
# Return the unpacked argument structure of this signature,
|
||||
# discarding information about which arguments are semantically
|
||||
# related to each other.
|
||||
def arguments(self) -> Sequence[Binding]:
|
||||
return cpp.arguments(
|
||||
self.func.arguments,
|
||||
faithful=self.faithful,
|
||||
symint=self.symint,
|
||||
method=self.method,
|
||||
cpp_no_default_args=self.cpp_no_default_args,
|
||||
)
|
||||
|
||||
def name(self, *, suppress_symint_suffix: bool = False) -> str:
|
||||
n = cpp.name(
|
||||
self.func,
|
||||
faithful_name_for_out_overloads=self.faithful,
|
||||
symint_overload=False if suppress_symint_suffix else self.symint,
|
||||
)
|
||||
if self.fallback_binding:
|
||||
n = f"__dispatch_{n}"
|
||||
return n
|
||||
|
||||
# Render the C++ declaration for this signature
|
||||
def decl(
|
||||
self,
|
||||
*,
|
||||
name: str | None = None,
|
||||
prefix: str = "",
|
||||
is_redispatching_fn: bool = False,
|
||||
suppress_symint_suffix: bool = False,
|
||||
) -> str:
|
||||
returns_type = cpp.returns_type(
|
||||
self.func.returns, symint=self.symint
|
||||
).cpp_type()
|
||||
cpp_args = [a.decl() for a in self.arguments()]
|
||||
if is_redispatching_fn:
|
||||
cpp_args = ["c10::DispatchKeySet dispatchKeySet"] + cpp_args
|
||||
cpp_args_str = ", ".join(cpp_args)
|
||||
if name is None:
|
||||
name = prefix + self.name(suppress_symint_suffix=suppress_symint_suffix)
|
||||
return f"{returns_type} {name}({cpp_args_str})"
|
||||
|
||||
# Render the C++ definition for this signature, not including
|
||||
# the body (with curly braces)
|
||||
def defn(
|
||||
self,
|
||||
*,
|
||||
name: str | None = None,
|
||||
prefix: str = "",
|
||||
is_redispatching_fn: bool = False,
|
||||
) -> str:
|
||||
returns_type = cpp.returns_type(
|
||||
self.func.returns, symint=self.symint
|
||||
).cpp_type()
|
||||
cpp_args = [a.defn() for a in self.arguments()]
|
||||
if is_redispatching_fn:
|
||||
cpp_args = ["c10::DispatchKeySet dispatchKeySet"] + cpp_args
|
||||
cpp_args_str = ", ".join(cpp_args)
|
||||
if name is None:
|
||||
name = prefix + self.name()
|
||||
return f"{returns_type} {name}({cpp_args_str})"
|
||||
|
||||
def ptr_type(self) -> str:
|
||||
args_types_str = ", ".join(a.type for a in self.arguments())
|
||||
return f"{cpp.returns_type(self.func.returns, symint=self.symint).cpp_type()} (*)({args_types_str})"
|
||||
|
||||
# Return the C++ function type, e.g., something like int(bool)
|
||||
def type(self) -> str:
|
||||
args_types_str = ", ".join(a.type for a in self.arguments())
|
||||
return f"{cpp.returns_type(self.func.returns, symint=self.symint).cpp_type()} ({args_types_str})"
|
||||
|
||||
|
||||
# Represents group of all CppSignatures associated with a
|
||||
# FunctionSchema. Right now, that's the regular, user-visible
|
||||
# signature, as well as a "faithful" signature which doesn't
|
||||
# have grouping.
|
||||
@dataclass(frozen=True)
|
||||
class CppSignatureGroup:
|
||||
func: FunctionSchema
|
||||
signature: CppSignature
|
||||
faithful_signature: CppSignature | None
|
||||
symint_signature: CppSignature | None
|
||||
symint_faithful_signature: CppSignature | None
|
||||
|
||||
def most_faithful_signature(self) -> CppSignature:
|
||||
if self.faithful_signature:
|
||||
return self.faithful_signature
|
||||
else:
|
||||
return self.signature
|
||||
|
||||
def signatures(self, *, symint: bool = True) -> Iterator[CppSignature]:
|
||||
yield self.signature
|
||||
if self.faithful_signature:
|
||||
yield self.faithful_signature
|
||||
if symint:
|
||||
if self.symint_signature:
|
||||
yield self.symint_signature
|
||||
if self.symint_faithful_signature:
|
||||
yield self.symint_faithful_signature
|
||||
|
||||
@staticmethod
|
||||
def from_native_function(
|
||||
f: NativeFunction, *, method: bool, fallback_binding: bool = False
|
||||
) -> CppSignatureGroup:
|
||||
func = f.func
|
||||
|
||||
def make_sig(*, faithful: bool, symint: bool) -> CppSignature:
|
||||
return CppSignature(
|
||||
func=func,
|
||||
faithful=faithful,
|
||||
symint=symint,
|
||||
method=method,
|
||||
fallback_binding=fallback_binding,
|
||||
cpp_no_default_args=f.cpp_no_default_args,
|
||||
)
|
||||
|
||||
def make_sigs(*, symint: bool) -> tuple[CppSignature, CppSignature | None]:
|
||||
faithful_signature: CppSignature | None = None
|
||||
if func.arguments.tensor_options is not None or len(func.arguments.out) > 0:
|
||||
faithful_signature = make_sig(faithful=True, symint=symint)
|
||||
signature = make_sig(faithful=False, symint=symint)
|
||||
return signature, faithful_signature
|
||||
|
||||
signature, faithful_signature = make_sigs(symint=False)
|
||||
symint_signature: CppSignature | None = None
|
||||
symint_faithful_signature: CppSignature | None = None
|
||||
if func.has_symint():
|
||||
symint_signature, symint_faithful_signature = make_sigs(symint=True)
|
||||
|
||||
return CppSignatureGroup(
|
||||
func=func,
|
||||
signature=signature,
|
||||
faithful_signature=faithful_signature,
|
||||
symint_signature=symint_signature,
|
||||
symint_faithful_signature=symint_faithful_signature,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DispatcherSignature:
|
||||
# The schema this signature is derived from
|
||||
func: FunctionSchema
|
||||
|
||||
# Allows you to prepend an arbitrary prefix to the signature name.
|
||||
# This is useful for parts of the codegen that generate wrappers around kernels,
|
||||
# and need to avoid naming collisions.
|
||||
prefix: str = ""
|
||||
|
||||
symint: bool = True
|
||||
|
||||
def arguments(self) -> list[Binding]:
|
||||
return dispatcher.arguments(self.func, symint=self.symint)
|
||||
|
||||
def name(self) -> str:
|
||||
return self.prefix + dispatcher.name(self.func)
|
||||
|
||||
def decl(self, name: str | None = None) -> str:
|
||||
args_str = ", ".join(a.decl() for a in self.arguments())
|
||||
if name is None:
|
||||
name = self.name()
|
||||
return f"{self.returns_type().cpp_type()} {name}({args_str})"
|
||||
|
||||
def defn(
|
||||
self, name: str | None = None, *, is_redispatching_fn: bool = False
|
||||
) -> str:
|
||||
args = [a.defn() for a in self.arguments()]
|
||||
if is_redispatching_fn:
|
||||
args = ["c10::DispatchKeySet dispatchKeySet"] + args
|
||||
args_str = ", ".join(args)
|
||||
if name is None:
|
||||
name = self.name()
|
||||
return f"{self.returns_type().cpp_type()} {name}({args_str})"
|
||||
|
||||
def exprs(self) -> list[Expr]:
|
||||
return [Expr(a.name, a.nctype) for a in self.arguments()]
|
||||
|
||||
def returns_type(self) -> CType:
|
||||
return dispatcher.returns_type(self.func.returns, symint=self.symint)
|
||||
|
||||
def ptr_type(self) -> str:
|
||||
dispatcher_args_types_str = ", ".join(a.type for a in self.arguments())
|
||||
return f"{self.returns_type().cpp_type()} (*)({dispatcher_args_types_str})"
|
||||
|
||||
# Return the C++ function type, e.g., something like int(bool)
|
||||
def type(self) -> str:
|
||||
dispatcher_args_types_str = ", ".join(a.type for a in self.arguments())
|
||||
return f"{self.returns_type().cpp_type()} ({dispatcher_args_types_str})"
|
||||
|
||||
@staticmethod
|
||||
def from_schema(
|
||||
func: FunctionSchema, *, prefix: str = "", symint: bool = True
|
||||
) -> DispatcherSignature:
|
||||
return DispatcherSignature(func, prefix, symint)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NativeSignature:
|
||||
# The schema this signature is derived from
|
||||
func: FunctionSchema
|
||||
|
||||
symint: bool
|
||||
|
||||
prefix: str = ""
|
||||
|
||||
def name(self) -> str:
|
||||
return self.prefix + native.name(self.func)
|
||||
|
||||
def decl(self, name: str | None = None) -> str:
|
||||
args_str = ", ".join(a.decl() for a in self.arguments())
|
||||
if name is None:
|
||||
name = self.name()
|
||||
return f"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} {name}({args_str})"
|
||||
|
||||
def defn(self, name: str | None = None) -> str:
|
||||
args_str = ", ".join(a.defn() for a in self.arguments())
|
||||
if name is None:
|
||||
name = self.name()
|
||||
return f"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} {name}({args_str})"
|
||||
|
||||
def ptr_type(self) -> str:
|
||||
# don't include defaults in type signature!
|
||||
args_str = ", ".join(a.defn() for a in self.arguments())
|
||||
return f"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} (*)({args_str})"
|
||||
|
||||
def arguments(self) -> list[Binding]:
|
||||
return native.arguments(self.func, symint=self.symint)
|
||||
|
||||
def returns_type(self) -> CType:
|
||||
return native.returns_type(self.func.returns, symint=self.symint)
|
||||
|
||||
def dispatcher_exprs(self) -> list[Expr]:
|
||||
return translate.translate(
|
||||
self.arguments(), dispatcher.arguments(self.func), method=False
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ViewInverseSignature:
|
||||
g: NativeFunctionsViewGroup
|
||||
|
||||
def name(self) -> str:
|
||||
return functionalization.reverse_name(self.g.view, include_namespace=False)
|
||||
|
||||
def decl(self) -> str:
|
||||
return_type = functionalization.returns_type(self.g.view.func)
|
||||
decls = [
|
||||
a.decl()
|
||||
for a in functionalization.op_arguments(self.g.view.func, is_reverse=True)
|
||||
]
|
||||
return f"static {return_type.cpp_type()} {self.name()}({', '.join(decls)});"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StructuredImplSignature:
|
||||
g: NativeFunctionsGroup
|
||||
name: str
|
||||
|
||||
def defn(self, name: str | None = None) -> str:
|
||||
args_str = ", ".join(a.defn() for a in self.arguments())
|
||||
return f"TORCH_IMPL_FUNC({self.name})({args_str})"
|
||||
|
||||
def arguments(self) -> list[Binding]:
|
||||
return structured.impl_arguments(self.g)
|
||||
|
||||
|
||||
# Helper functions
|
||||
|
||||
|
||||
def kernel_signature(
|
||||
f: NativeFunction, backend_index: BackendIndex, *, prefix: str = ""
|
||||
) -> NativeSignature | DispatcherSignature:
|
||||
# Note [External Backends Follow Dispatcher API]
|
||||
# Kernel signatures for in-tree backends follow the "native" API,
|
||||
# while kernels for out-of-tree backends follow the dispatcher API.
|
||||
# See the comments in `native.py` for details, but historically there have been
|
||||
# some small differences in schema convention between them and the Dispatcher API.
|
||||
# Any differences that require translating between the two will results in a runtime cost,
|
||||
# so we'd like to keep the differences as small as possible.
|
||||
# With external backends, we'd like to enforce that they write their kernels with schemas
|
||||
# that match the Dispatcher API directly, if they can.
|
||||
meta = backend_index.get_kernel(f)
|
||||
symint = meta is not None and meta.supports_symint()
|
||||
if symint:
|
||||
if not f.func.has_symint():
|
||||
raise AssertionError(
|
||||
f"attempted to define symint kernel for {backend_index.dispatch_key} "
|
||||
"without SymInt in schema"
|
||||
)
|
||||
if backend_index.external:
|
||||
return DispatcherSignature.from_schema(f.func, prefix=prefix, symint=symint)
|
||||
else:
|
||||
return NativeSignature(f.func, prefix=prefix, symint=symint)
|
||||
|
||||
|
||||
# Functions only, no types
|
||||
from torchgen.api import (
|
||||
cpp,
|
||||
dispatcher,
|
||||
functionalization,
|
||||
native,
|
||||
structured,
|
||||
translate,
|
||||
)
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
Where should I add a new type? `types_base.py` vs `types.py`
|
||||
|
||||
This file defines data model classes for torchgen typing system, as well as some base types such as int32_t.
|
||||
|
||||
`types.py` defines ATen Tensor type and some c10 types, along with signatures that use these types.
|
||||
|
||||
The difference between these two files, is `types_base.py` should be implementation-agnostic, meaning it shouldn't
|
||||
contain any type definition that is tight to a specific C++ library (e.g., ATen), so that it can be easily reused
|
||||
if we want to generate code for another C++ library.
|
||||
|
||||
Add new types to `types.py` if these types are ATen/c10 related.
|
||||
Add new types to `types_base.py` if they are basic and not attached to ATen/c10.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from torchgen.api.types.types_base import (
|
||||
BaseCppType,
|
||||
BaseCType,
|
||||
boolT,
|
||||
byteT,
|
||||
charT,
|
||||
CType,
|
||||
doubleT,
|
||||
floatT,
|
||||
int32T,
|
||||
longT,
|
||||
shortT,
|
||||
)
|
||||
from torchgen.model import BaseTy, ScalarType
|
||||
|
||||
|
||||
TENSOR_LIST_LIKE_CTYPES = [
|
||||
"at::TensorList",
|
||||
"const c10::List<::std::optional<at::Tensor>> &",
|
||||
"const at::ITensorListRef &",
|
||||
]
|
||||
|
||||
|
||||
halfT = BaseCppType("at", "Half")
|
||||
complexHalfT = BaseCppType(
|
||||
"c10", "complex<c10::Half>"
|
||||
) # stuffing template param here is an abuse
|
||||
complexFloatT = BaseCppType("c10", "complex<float>")
|
||||
complexDoubleT = BaseCppType("c10", "complex<double>")
|
||||
bfloat16T = BaseCppType("at", "BFloat16")
|
||||
float8_e5m2T = BaseCppType("at", "Float8_e5m2")
|
||||
float8_e5m2fnuzT = BaseCppType("at", "Float8_e5m2fnuz")
|
||||
float8_e4m3fnT = BaseCppType("at", "Float8_e4m3fn")
|
||||
float8_e4m3fnuzT = BaseCppType("at", "Float8_e4m3fnuz")
|
||||
float8_e8m0fnuT = BaseCppType("at", "Float8_e8m0fnu")
|
||||
stringT = BaseCppType("c10", "string_view")
|
||||
generatorT = BaseCppType("at", "Generator")
|
||||
scalarTypeT = BaseCppType("at", "ScalarType")
|
||||
tensorT = BaseCppType("at", "Tensor")
|
||||
optionalTensorRefT = BaseCppType("at", "OptionalTensorRef")
|
||||
tensorListT = BaseCppType("at", "TensorList")
|
||||
iTensorListRefT = BaseCppType("at", "ITensorListRef")
|
||||
iOptTensorListRefT = BaseCppType("at", "IOptTensorListRef")
|
||||
dimnameT = BaseCppType("at", "Dimname")
|
||||
dimnameListT = BaseCppType("at", "DimnameList")
|
||||
dimVectorT = BaseCppType("at", "DimVector")
|
||||
layoutT = BaseCppType("at", "Layout")
|
||||
deviceT = BaseCppType("at", "Device")
|
||||
deviceIndexT = BaseCppType("at", "DeviceIndex")
|
||||
scalarT = BaseCppType("at", "Scalar")
|
||||
optionalScalarRefT = BaseCppType("at", "OptionalScalarRef")
|
||||
memoryFormatT = BaseCppType("at", "MemoryFormat")
|
||||
qschemeT = BaseCppType("at", "QScheme")
|
||||
storageT = BaseCppType("at", "Storage")
|
||||
streamT = BaseCppType("at", "Stream")
|
||||
intArrayRefT = BaseCppType("at", "IntArrayRef")
|
||||
optionalIntArrayRefT = BaseCppType("at", "OptionalIntArrayRef")
|
||||
optionalSymIntArrayRefT = BaseCppType("at", "OptionalSymIntArrayRef")
|
||||
tensorOptionsT = BaseCppType("at", "TensorOptions")
|
||||
typeAndSizeT = BaseCppType("torch::autograd::generated", "TypeAndSize")
|
||||
tensorGeometryT = BaseCppType("at", "TensorGeometry")
|
||||
SymIntT = BaseCppType("c10", "SymInt")
|
||||
SymBoolT = BaseCppType("c10", "SymBool")
|
||||
symIntArrayRefT = BaseCppType("c10", "SymIntArrayRef")
|
||||
|
||||
# Types representing template parameters. Technically, we probably shouldn't
|
||||
# represent them this way in codegen, but it was pretty convenient.
|
||||
scalar_t = BaseCppType("", "scalar_t")
|
||||
opmath_t = BaseCppType("", "opmath_t")
|
||||
|
||||
ScalarTypeToCppMapping: dict[ScalarType, BaseCppType] = {
|
||||
ScalarType.Byte: byteT,
|
||||
ScalarType.Char: charT,
|
||||
ScalarType.Short: shortT,
|
||||
ScalarType.Int: int32T,
|
||||
ScalarType.Long: longT,
|
||||
ScalarType.Half: halfT,
|
||||
ScalarType.Float: floatT,
|
||||
ScalarType.Double: doubleT,
|
||||
ScalarType.ComplexHalf: complexHalfT,
|
||||
ScalarType.ComplexFloat: complexFloatT,
|
||||
ScalarType.ComplexDouble: complexDoubleT,
|
||||
ScalarType.Bool: boolT,
|
||||
ScalarType.Float8_e5m2: float8_e5m2T,
|
||||
ScalarType.Float8_e5m2fnuz: float8_e5m2fnuzT,
|
||||
ScalarType.Float8_e4m3fn: float8_e4m3fnT,
|
||||
ScalarType.Float8_e4m3fnuz: float8_e4m3fnuzT,
|
||||
ScalarType.Float8_e8m0fnu: float8_e8m0fnuT,
|
||||
}
|
||||
|
||||
BaseTypeToCppMapping: dict[BaseTy, BaseCppType] = {
|
||||
BaseTy.int: longT,
|
||||
BaseTy.float: doubleT,
|
||||
BaseTy.bool: boolT,
|
||||
BaseTy.str: stringT,
|
||||
BaseTy.Generator: generatorT,
|
||||
BaseTy.ScalarType: scalarTypeT,
|
||||
BaseTy.Tensor: tensorT,
|
||||
BaseTy.Dimname: dimnameT,
|
||||
BaseTy.DimVector: dimVectorT,
|
||||
BaseTy.Layout: layoutT,
|
||||
BaseTy.Device: deviceT,
|
||||
BaseTy.DeviceIndex: deviceIndexT,
|
||||
BaseTy.Scalar: scalarT,
|
||||
BaseTy.MemoryFormat: memoryFormatT,
|
||||
BaseTy.QScheme: qschemeT,
|
||||
BaseTy.Storage: storageT,
|
||||
BaseTy.Stream: streamT,
|
||||
BaseTy.SymInt: SymIntT,
|
||||
BaseTy.SymBool: SymBoolT,
|
||||
}
|
||||
|
||||
# CTypes encode C++ type structure as needed for translation.
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OptionalCType(CType):
|
||||
elem: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
# Do not pass `strip_ref` recursively.
|
||||
return f"::std::optional<{self.elem.cpp_type()}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return OptionalCType(self.elem.remove_const_ref())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ListCType(CType):
|
||||
elem: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
# Do not pass `strip_ref` recursively.
|
||||
return f"c10::List<{self.elem.cpp_type()}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return ListCType(self.elem.remove_const_ref())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArrayRefCType(CType):
|
||||
elem: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
# Do not pass `strip_ref` recursively.
|
||||
return f"at::ArrayRef<{self.elem.cpp_type()}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return ArrayRefCType(self.elem.remove_const_ref())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorizedCType(CType):
|
||||
# This template is explicitly specialized, so the only valid
|
||||
# elems are those we have specializations for (e.g., float, double, ...)
|
||||
# scalar_t is also a common argument here (when we are codegen in
|
||||
# a templated context)
|
||||
elem: BaseCType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
return f"at::vec::Vectorized<{self.elem.cpp_type()}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return self
|
||||
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
Where should I add a new type? `types_base.py` vs `types.py`
|
||||
|
||||
This file defines data model classes for torchgen typing system, as well as some base types such as int32_t.
|
||||
|
||||
`types.py` defines ATen Tensor type and some c10 types, along with signatures that use these types.
|
||||
|
||||
The difference between these two files, is `types_base.py` should be implementation-agnostic, meaning it shouldn't
|
||||
contain any type definition that is tight to a specific C++ library (e.g., ATen), so that it can be easily reused
|
||||
if we want to generate code for another C++ library.
|
||||
|
||||
Add new types to `types.py` if these types are ATen/c10 related.
|
||||
Add new types to `types_base.py` if they are basic and not attached to ATen/c10.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import auto, Enum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torchgen.model import Argument, SelfArgument, TensorOptionsArguments
|
||||
|
||||
|
||||
# An ArgName is just the str name of the argument in schema;
|
||||
# but in some special circumstances, we may add a little extra
|
||||
# context. The Enum SpecialArgName covers all of these cases;
|
||||
# grep for their construction sites to see when they can occur.
|
||||
|
||||
|
||||
class SpecialArgName(Enum):
|
||||
possibly_redundant_memory_format = auto()
|
||||
|
||||
|
||||
ArgName = str | SpecialArgName
|
||||
|
||||
|
||||
# This class shouldn't be created directly; instead, use/create one of the singletons below.
|
||||
@dataclass(frozen=True)
|
||||
class BaseCppType:
|
||||
ns: str | None
|
||||
name: str
|
||||
|
||||
def __str__(self) -> str:
|
||||
if self.ns is None or self.ns == "":
|
||||
return self.name
|
||||
return f"{self.ns}::{self.name}"
|
||||
|
||||
|
||||
# The set of all non-templated, valid, fully-qualified names of C++ types that are used in the codegen.
|
||||
# Templated types get their own dataclass, mainly to make namespace parsing easier.
|
||||
byteT = BaseCppType("", "uint8_t")
|
||||
charT = BaseCppType("", "int8_t")
|
||||
shortT = BaseCppType("", "int16_t")
|
||||
# It would be more symmetric for this to be called intT, but it easy to mix
|
||||
# this up with JIT int (which is int64_t in C++), so we intentionally don't
|
||||
# define intT to make it obvious when you've stuffed it up
|
||||
int32T = BaseCppType("", "int32_t")
|
||||
longT = BaseCppType("", "int64_t")
|
||||
doubleT = BaseCppType("", "double")
|
||||
floatT = BaseCppType("", "float")
|
||||
boolT = BaseCppType("", "bool")
|
||||
voidT = BaseCppType("", "void")
|
||||
|
||||
|
||||
class CType(ABC):
|
||||
@abstractmethod
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def remove_const_ref(self) -> CType:
|
||||
return self
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BaseCType(CType):
|
||||
type: BaseCppType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
return str(self.type)
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return self
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ConstRefCType(CType):
|
||||
elem: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
if strip_ref:
|
||||
return self.elem.cpp_type(strip_ref=strip_ref)
|
||||
return f"const {self.elem.cpp_type()} &"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return self.elem.remove_const_ref()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorCType(CType):
|
||||
elem: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
# Do not pass `strip_ref` recursively.
|
||||
return f"::std::vector<{self.elem.cpp_type()}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return VectorCType(self.elem.remove_const_ref())
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArrayCType(CType):
|
||||
elem: CType
|
||||
size: int
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
# Do not pass `strip_ref` recursively.
|
||||
return f"::std::array<{self.elem.cpp_type()},{self.size}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return ArrayCType(self.elem.remove_const_ref(), self.size)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TupleCType(CType):
|
||||
elems: list[CType]
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
# Do not pass `strip_ref` recursively.
|
||||
return f"::std::tuple<{','.join([e.cpp_type() for e in self.elems])}>"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return TupleCType([e.remove_const_ref() for e in self.elems])
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MutRefCType(CType):
|
||||
elem: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
if strip_ref:
|
||||
return self.elem.cpp_type(strip_ref=strip_ref)
|
||||
return f"{self.elem.cpp_type()} &"
|
||||
|
||||
def remove_const_ref(self) -> CType:
|
||||
return self.elem.remove_const_ref()
|
||||
|
||||
|
||||
# A NamedCType is short for Named C++ semantic type. A NamedCType represents a C++ type, plus
|
||||
# semantic information about what it represents. For example, consider the
|
||||
# argument "bool pin_memory"; its normal C++ type is "bool", but its C++
|
||||
# semantic type also keeps track that this represents a "pin_memory"; you can't
|
||||
# just use a random other boolean in a context where you need a "pin_memory"!
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NamedCType:
|
||||
name: ArgName
|
||||
type: CType
|
||||
|
||||
def cpp_type(self, *, strip_ref: bool = False) -> str:
|
||||
return self.type.cpp_type(strip_ref=strip_ref)
|
||||
|
||||
def remove_const_ref(self) -> NamedCType:
|
||||
return NamedCType(self.name, self.type.remove_const_ref())
|
||||
|
||||
def with_name(self, name: str) -> NamedCType:
|
||||
return NamedCType(name, self.type)
|
||||
|
||||
|
||||
# A binding represents any C++ binding site for a formal parameter.
|
||||
# We don't distinguish between binding sites for different APIs;
|
||||
# instead, all of the important distinctions are encoded in CType,
|
||||
# which you can use to figure out if a given Binding is appropriate
|
||||
# for use in another context. (See torchgen.api.translate)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Binding:
|
||||
name: str
|
||||
nctype: NamedCType
|
||||
argument: Argument | TensorOptionsArguments | SelfArgument
|
||||
# TODO: maybe don't represent default here
|
||||
default: str | None = None
|
||||
|
||||
def rename(self, name: str) -> Binding:
|
||||
return Binding(
|
||||
name=name,
|
||||
nctype=self.nctype,
|
||||
argument=self.argument,
|
||||
default=self.default,
|
||||
)
|
||||
|
||||
@property
|
||||
def type(self) -> str:
|
||||
return self.nctype.cpp_type()
|
||||
|
||||
def no_default(self) -> Binding:
|
||||
return Binding(
|
||||
name=self.name,
|
||||
nctype=self.nctype,
|
||||
default=None,
|
||||
argument=self.argument,
|
||||
)
|
||||
|
||||
def decl(self, *, func_ptr_cast: bool = False) -> str:
|
||||
mb_default = ""
|
||||
if self.default is not None:
|
||||
mb_default = f"={self.default}"
|
||||
|
||||
# casting only needs to know the type
|
||||
if func_ptr_cast:
|
||||
return f"{self.type}"
|
||||
else:
|
||||
return f"{self.type} {self.name}{mb_default}"
|
||||
|
||||
def defn(self) -> str:
|
||||
return f"{self.type} {self.name}"
|
||||
|
||||
def with_name(self, name: str) -> Binding:
|
||||
return Binding(
|
||||
name=name, nctype=self.nctype, argument=self.argument, default=self.default
|
||||
)
|
||||
|
||||
|
||||
# An Expr is a C++ expression. It has a C++ string representing its syntax,
|
||||
# as well as a CType saying what it provides.
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Expr:
|
||||
expr: str
|
||||
type: NamedCType
|
||||
@@ -0,0 +1,211 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torchgen.api.types as api_types
|
||||
from torchgen.api import cpp, structured
|
||||
from torchgen.api.types import (
|
||||
ArgName,
|
||||
BaseCppType,
|
||||
BaseCType,
|
||||
Binding,
|
||||
ConstRefCType,
|
||||
CType,
|
||||
NamedCType,
|
||||
scalarT,
|
||||
)
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
DispatchKey,
|
||||
FunctionSchema,
|
||||
NativeFunctionsGroup,
|
||||
Type,
|
||||
)
|
||||
|
||||
|
||||
def schema_kernel_name(func: FunctionSchema, dispatch_key: DispatchKey) -> str:
|
||||
if not func.is_out_fn():
|
||||
raise AssertionError("ufunc.kernel_name should only be invoked on out schemas")
|
||||
return f"ufunc_{func.name.name}_{dispatch_key}"
|
||||
|
||||
|
||||
def kernel_name(g: NativeFunctionsGroup, dispatch_key: DispatchKey) -> str:
|
||||
return schema_kernel_name(g.out.func, dispatch_key)
|
||||
|
||||
|
||||
# Tensors are omitted (as they are stored in TensorIterator), everything else is
|
||||
# passed along (technically, we can pass tensors along too, it just wastes
|
||||
# argument registers)
|
||||
#
|
||||
# NB: used for CPU only
|
||||
def dispatchstub_type(t: Type, *, binds: ArgName) -> NamedCType | None:
|
||||
# Dispatch stubs are always plain ints
|
||||
r = cpp.valuetype_type(t, binds=binds, symint=False)
|
||||
if r is not None:
|
||||
return r
|
||||
|
||||
if t == BaseType(BaseTy.Scalar):
|
||||
return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))
|
||||
elif t == BaseType(BaseTy.Tensor):
|
||||
return None
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
def opmath_type(scalar_t: BaseCppType) -> BaseCppType:
|
||||
if scalar_t == api_types.scalar_t:
|
||||
return api_types.opmath_t
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
# NB: Tensors in constructor are stored in opmath_t, not scalar_t
|
||||
# because Tensor in constructor = its a scalar tensor partially applied =
|
||||
# it can be higher precision and we want to compute in that higher precision
|
||||
#
|
||||
# NB: CUDA only
|
||||
def ufunctor_ctor_type(t: Type, *, binds: ArgName, scalar_t: BaseCppType) -> NamedCType:
|
||||
r = cpp.valuetype_type(t, binds=binds, symint=False)
|
||||
if r is not None:
|
||||
return r
|
||||
|
||||
if t == BaseType(BaseTy.Scalar):
|
||||
return NamedCType(binds, BaseCType(opmath_type(scalar_t)))
|
||||
elif t == BaseType(BaseTy.Tensor):
|
||||
return NamedCType(binds, BaseCType(opmath_type(scalar_t)))
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
# Only Tensors ever get passed directly to operator()
|
||||
#
|
||||
# NB: CUDA only
|
||||
# (Actually, this works for CPU too)
|
||||
def ufunctor_apply_type(
|
||||
t: Type, *, binds: ArgName, scalar_t: BaseCppType
|
||||
) -> NamedCType:
|
||||
if t == BaseType(BaseTy.Tensor):
|
||||
return NamedCType(binds, BaseCType(scalar_t))
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
# The actual ufunc template function the user writes. Everything here
|
||||
# is done in the computation type. compute_t is opmath_t in CUDA and scalar_t
|
||||
# in CPU
|
||||
def ufunc_type(t: Type, *, binds: ArgName, compute_t: CType) -> NamedCType:
|
||||
r = cpp.valuetype_type(t, binds=binds, symint=False)
|
||||
if r is not None:
|
||||
return r
|
||||
|
||||
if t == BaseType(BaseTy.Scalar):
|
||||
return NamedCType(binds, compute_t)
|
||||
elif t == BaseType(BaseTy.Tensor):
|
||||
return NamedCType(binds, compute_t)
|
||||
else:
|
||||
raise AssertionError(f"unrecognized type {repr(t)}")
|
||||
|
||||
|
||||
def ufunctor_ctor_argument(a: Argument, scalar_t: BaseCppType) -> Binding:
|
||||
return Binding(
|
||||
nctype=ufunctor_ctor_type(a.type, binds=a.name, scalar_t=scalar_t),
|
||||
name=a.name,
|
||||
default=None,
|
||||
argument=a,
|
||||
)
|
||||
|
||||
|
||||
def ufunctor_apply_argument(a: Argument, scalar_t: BaseCppType) -> Binding:
|
||||
return Binding(
|
||||
nctype=ufunctor_apply_type(a.type, binds=a.name, scalar_t=scalar_t),
|
||||
name=a.name,
|
||||
default=None,
|
||||
argument=a,
|
||||
)
|
||||
|
||||
|
||||
def ufunc_argument(a: Argument, compute_t: CType) -> Binding:
|
||||
return Binding(
|
||||
nctype=ufunc_type(a.type, binds=a.name, compute_t=compute_t),
|
||||
name=a.name,
|
||||
default=None,
|
||||
argument=a,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UfunctorBindings:
|
||||
ctor: list[Binding]
|
||||
apply: list[Binding]
|
||||
|
||||
|
||||
# ufunctors are a CUDA-only concept representing functors that take some of
|
||||
# their arguments on a host-side constructor, and the rest in the device-side
|
||||
# apply. E.g.,
|
||||
#
|
||||
# template <typename scalar_t>
|
||||
# struct CUDAFunctorOnSelf_add {
|
||||
# using opmath_t = at::opmath_type<scalar_t>;
|
||||
# opmath_t other_;
|
||||
# opmath_t alpha_;
|
||||
# CUDAFunctorOnSelf_add(opmath_t other, opmath_t alpha) : other_(other), alpha_(alpha) {}
|
||||
# __device__ scalar_t operator()(scalar_t self) {
|
||||
# return ufunc::add(static_cast<opmath_t>(self), other_, alpha_);
|
||||
# }
|
||||
# };
|
||||
#
|
||||
# The ctor refers to the constructor CUDAFunctorOnSelf_add, while apply refers
|
||||
# to the operator() definition
|
||||
def ufunctor_arguments(
|
||||
g: NativeFunctionsGroup, *, scalar_tensor_idx: int | None, scalar_t: BaseCppType
|
||||
) -> UfunctorBindings:
|
||||
ctor = []
|
||||
apply = []
|
||||
for a in g.functional.func.arguments.flat_non_out:
|
||||
if a.type.is_tensor_like():
|
||||
if scalar_tensor_idx == 0:
|
||||
# put it in the ctor anyway
|
||||
ctor.append(ufunctor_ctor_argument(a, scalar_t=scalar_t))
|
||||
scalar_tensor_idx = None
|
||||
else:
|
||||
if scalar_tensor_idx is not None:
|
||||
scalar_tensor_idx -= 1
|
||||
apply.append(ufunctor_apply_argument(a, scalar_t=scalar_t))
|
||||
else:
|
||||
ctor.append(ufunctor_ctor_argument(a, scalar_t=scalar_t))
|
||||
if scalar_tensor_idx is not None:
|
||||
raise AssertionError("scalar_tensor_idx should be None at end of processing")
|
||||
return UfunctorBindings(ctor=ctor, apply=apply)
|
||||
|
||||
|
||||
# ufuncs are the inner loop template functions that you wrote in ufunc/add.h
|
||||
# which do the actual computation in question. E.g.,
|
||||
#
|
||||
# template <typename T>
|
||||
# C10_HOST_DEVICE T add(T self, T other, T alpha) __ubsan_ignore_undefined__ {
|
||||
# return self + alpha * other;
|
||||
# }
|
||||
#
|
||||
# In this file, we refer to T as compute_t which is bound by caller
|
||||
def ufunc_arguments(g: NativeFunctionsGroup, *, compute_t: CType) -> list[Binding]:
|
||||
return [
|
||||
ufunc_argument(a, compute_t=compute_t)
|
||||
for a in g.functional.func.arguments.flat_non_out
|
||||
]
|
||||
|
||||
|
||||
# Stubs are the DispatchStub trampolines that CPU kernels use to get to their
|
||||
# vectorized versions. E.g.,
|
||||
#
|
||||
# using structured_binary_fn_alpha = void(*)(TensorIteratorBase&, const Scalar& alpha);
|
||||
# DECLARE_DISPATCH(structured_binary_fn_alpha, add_stub);
|
||||
def stub_arguments(g: NativeFunctionsGroup) -> list[Binding]:
|
||||
# stubs drop all tensor arguments (they are implicit in the TensorIterator
|
||||
# argument and keep everything else)
|
||||
return [
|
||||
r
|
||||
for a in g.out.func.arguments.flat_non_out
|
||||
if not a.type.is_tensor_like()
|
||||
for r in structured.argument(a)
|
||||
]
|
||||
@@ -0,0 +1,241 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import Binding, CppSignatureGroup, CType
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
ListType,
|
||||
NativeFunction,
|
||||
OptionalType,
|
||||
Type,
|
||||
)
|
||||
|
||||
|
||||
# This file generates the code for unboxing wrappers, i.e., the glue logic to unbox a boxed operator and convert the
|
||||
# ivalues from stack to correct arguments to the unboxed kernel, based on corresponding JIT schema. This codegen is
|
||||
# an alternative way to generate unboxing wrappers similar to the existing C++ metaprogramming approach but gets the
|
||||
# job done statically. These generated unboxing wrappers will be useful under the scenario where we need to register
|
||||
# a fixed set of operators known at compile time and thus can save some time in runtime initialization phase.
|
||||
#
|
||||
# Here's an example on how the codegen works:
|
||||
#
|
||||
# - Function Schema (source of truth)
|
||||
#
|
||||
# aten::empty.names(int[] size, *, Dimname[]? names,
|
||||
# ScalarType? dtype=None, Layout? layout=None,
|
||||
# Device? device=None, bool? pin_memory=None,
|
||||
# MemoryFormat? memory_format=None) -> Tensor
|
||||
# - Argument Conversion
|
||||
# Generates C++ code to convert an ivalue (from stack) to its underlying C++ type.
|
||||
# - int[] size
|
||||
# ```cpp
|
||||
# const c10::List<c10::IValue> size_list_in = (std::move(peek(stack, 0, 7))).toList();
|
||||
#
|
||||
# std::vector<int64_t> size_vec;
|
||||
# for (c10::IValue size_elem: size_list_in) {
|
||||
# int64_t size_base = size_elem.to<int64_t>();
|
||||
# size_vec.push_back(size_base);
|
||||
# }
|
||||
# at::ArrayRef<int64_t> size_list_out(size_vec);
|
||||
# ~~~~~~~~~~~~~ <-- The converted argument from ivalues in the stack.
|
||||
# Will be passed to unboxed kernel.
|
||||
# ```
|
||||
# - Dimname[]? names
|
||||
# ```cpp
|
||||
# ::std::optional<c10::IValue> names_opt = (std::move(peek(stack, 1, 7))).toOptional<c10::IValue>();
|
||||
# ::std::optional<at::ArrayRef<at::Dimname>> names_opt_out;
|
||||
# if (names_opt.has_value()) {
|
||||
# ~~~~~~~~~~~ <-- Unwrapping optional shell
|
||||
# const c10::IValue names_opt_in = names_opt.value();
|
||||
# const c10::List<c10::IValue> names_list_in = names_opt_in.toList();
|
||||
#
|
||||
# std::vector<at::Dimname> names_vec;
|
||||
# for (c10::IValue names_elem: names_list_in) {
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~ <-- Unrolling list, then convert elements one by one.
|
||||
# at::Dimname names_base = names_elem.to<at::Dimname>();
|
||||
# names_vec.push_back(names_base);
|
||||
# }
|
||||
# at::ArrayRef<at::Dimname> names_list_out(names_vec);
|
||||
#
|
||||
# names_opt_out = ::std::optional<at::ArrayRef<at::Dimname>>(names_list_out);
|
||||
# } else {
|
||||
# names_opt_out = ::std::optional<at::ArrayRef<at::Dimname>>();
|
||||
# }
|
||||
# ```
|
||||
# - ScalarType? dtype (similarly for the rest of the arguments)
|
||||
# ```cpp
|
||||
# ::std::optional<c10::IValue> dtype_opt = (std::move(peek(stack, 2, 7))).toOptional<c10::IValue>();
|
||||
# ::std::optional<at::ScalarType> dtype_opt_out;
|
||||
# if (dtype_opt.has_value()) {
|
||||
# const c10::IValue dtype_opt_in = dtype_opt.value();
|
||||
# at::ScalarType dtype_base = dtype_opt_in.to<at::ScalarType>();
|
||||
# ~~~~~~~~~~~~~~~~~~~~ <-- For base types, convert ivalue to it
|
||||
# directly using ".to<T>()" API.
|
||||
# dtype_opt_out = ::std::optional<at::ScalarType>(dtype_base);
|
||||
# } else {
|
||||
# dtype_opt_out = ::std::optional<at::ScalarType>();
|
||||
# }
|
||||
# ```
|
||||
#
|
||||
# - Unboxed Kernel Call
|
||||
# ```cpp
|
||||
# auto result_ = torch::empty(
|
||||
# size_list_out,
|
||||
# names_opt_out,
|
||||
# options,
|
||||
# memory_format_opt_out
|
||||
# );
|
||||
# ```
|
||||
#
|
||||
# - Push Result Back to Stack
|
||||
# ```cpp
|
||||
# drop(stack, 7);
|
||||
# pack(stack, std::move(result_));
|
||||
# ```
|
||||
connector = "\n\t"
|
||||
|
||||
|
||||
# Return unboxing function name for a NativeFunction
|
||||
def name(f: NativeFunction) -> str:
|
||||
return f.func.name.unambiguous_name()
|
||||
|
||||
|
||||
# Convert all the arguments in a NativeFunction to C++ code
|
||||
def convert_arguments(f: NativeFunction) -> tuple[list[Binding], list[str]]:
|
||||
# we need the 'self' argument so method needs to be False
|
||||
args = (
|
||||
CppSignatureGroup.from_native_function(f, method=False)
|
||||
.most_faithful_signature()
|
||||
.arguments()
|
||||
)
|
||||
code_list = [
|
||||
f"c10::IValue {args[i].name} = std::move(peek(stack, {i}, {len(args)}));"
|
||||
for i in range(len(args))
|
||||
] + [""]
|
||||
binding_list = []
|
||||
for arg in args:
|
||||
# expecting only Argument
|
||||
if not isinstance(arg.argument, Argument):
|
||||
raise Exception( # noqa: TRY002
|
||||
f"Unexpected argument type, expecting `Argument` but got {arg}"
|
||||
)
|
||||
argument: Argument = arg.argument
|
||||
unboxed_name, _, code, decl = argumenttype_ivalue_convert(
|
||||
argument.type,
|
||||
argument.name,
|
||||
mutable=argument.is_write,
|
||||
)
|
||||
code_list.extend(decl)
|
||||
code_list.extend(code)
|
||||
binding_list.append(arg.with_name(unboxed_name))
|
||||
return binding_list, code_list
|
||||
|
||||
|
||||
# Takes in the type, name and mutability corresponding to an argument, and generates a tuple of:
|
||||
# (1) the C++ code necessary to unbox the argument
|
||||
# (2) A Binding corresponding to the newly created unboxed variable, including variable name and its CType
|
||||
def argumenttype_ivalue_convert(
|
||||
t: Type, arg_name: str, *, mutable: bool = False
|
||||
) -> tuple[str, CType, list[str], list[str]]:
|
||||
# Unboxing is for mobile, which doesn't care about SymInts
|
||||
ctype = cpp.argumenttype_type(
|
||||
t=t, mutable=mutable, binds=arg_name, symint=False
|
||||
).type
|
||||
|
||||
if isinstance(t, BaseType):
|
||||
out_name = f"{arg_name}_base"
|
||||
code, decl = _gen_code_base_type(
|
||||
arg_name=arg_name, out_name=out_name, ctype=ctype
|
||||
)
|
||||
elif isinstance(t, OptionalType):
|
||||
out_name = f"{arg_name}_opt_out"
|
||||
code, decl = _gen_code_optional_type(
|
||||
arg_name=arg_name,
|
||||
out_name=out_name,
|
||||
t=t,
|
||||
ctype=ctype,
|
||||
)
|
||||
elif isinstance(t, ListType):
|
||||
out_name = f"{arg_name}_list_out"
|
||||
code, decl = _gen_code_list_type(
|
||||
arg_name=arg_name,
|
||||
out_name=out_name,
|
||||
t=t,
|
||||
ctype=ctype,
|
||||
)
|
||||
else:
|
||||
raise Exception(f"Cannot handle type {t}. arg_name: {arg_name}") # noqa: TRY002
|
||||
return out_name, ctype, code, decl
|
||||
|
||||
|
||||
def _gen_code_base_type(
|
||||
arg_name: str, out_name: str, ctype: CType
|
||||
) -> tuple[list[str], list[str]]:
|
||||
return [
|
||||
f"{ctype.cpp_type(strip_ref=True)} {out_name} = {arg_name}.to<{ctype.cpp_type(strip_ref=True)}>();"
|
||||
], []
|
||||
|
||||
|
||||
def _gen_code_optional_type(
|
||||
arg_name: str, out_name: str, t: OptionalType, ctype: CType
|
||||
) -> tuple[list[str], list[str]]:
|
||||
in_name = f"{arg_name}_opt_in"
|
||||
res_name, _, res_code, decl = argumenttype_ivalue_convert(t.elem, in_name)
|
||||
return (
|
||||
f"""
|
||||
auto {arg_name}_opt = {arg_name}.toOptional<c10::IValue>();
|
||||
{ctype.cpp_type(strip_ref=True)} {out_name};
|
||||
if ({arg_name}_opt.has_value()) {{
|
||||
const c10::IValue {in_name} = {arg_name}_opt.value();
|
||||
{connector.join(res_code)}
|
||||
{out_name} = {ctype.cpp_type(strip_ref=True)}({res_name});
|
||||
}} else {{
|
||||
{out_name} = {ctype.cpp_type(strip_ref=True)}();
|
||||
}}
|
||||
""".split("\n"),
|
||||
decl,
|
||||
)
|
||||
|
||||
|
||||
def _gen_code_list_type(
|
||||
arg_name: str, out_name: str, t: ListType, ctype: CType
|
||||
) -> tuple[list[str], list[str]]:
|
||||
in_name = f"{arg_name}_list_in"
|
||||
elem_name = f"{arg_name}_elem"
|
||||
code = [f"const c10::List<c10::IValue> {in_name} = {arg_name}.toList();"]
|
||||
res_name, res_ctype, res_code, decl = argumenttype_ivalue_convert(t.elem, elem_name)
|
||||
# handle list type with size, e.g., bool[4]
|
||||
if isinstance(t.elem, BaseType) and t.elem.name == BaseTy.bool and t.size:
|
||||
code.extend(
|
||||
f"""
|
||||
{ctype.cpp_type(strip_ref=True)} {out_name} = as_array<{res_ctype.cpp_type(strip_ref=True)}, {t.size}>({in_name});
|
||||
""".split("\n")
|
||||
)
|
||||
# we have to use c10::List for optional element. e.g., Tensor?[] -> c10::List<::std::optional<at::Tensor>>
|
||||
elif isinstance(t.elem, OptionalType):
|
||||
code.extend(
|
||||
f"""
|
||||
{ctype.cpp_type(strip_ref=True)} {out_name};
|
||||
for (c10::IValue {elem_name}: {in_name}) {{
|
||||
{connector.join(res_code)}
|
||||
{out_name}.push_back({res_name});
|
||||
}}
|
||||
""".split("\n")
|
||||
)
|
||||
else:
|
||||
# use ArrayRef as default.
|
||||
vec_name = arg_name + "_vec"
|
||||
# need to bring vector instantiation out of scope so that ArrayRef has valid data
|
||||
decl.append(f"std::vector<{res_ctype.cpp_type(strip_ref=True)}> {vec_name};")
|
||||
code.extend(
|
||||
f"""
|
||||
for (c10::IValue {elem_name}: {in_name}) {{
|
||||
{connector.join(res_code)}
|
||||
{vec_name}.push_back({res_name});
|
||||
}}
|
||||
{ctype.cpp_type(strip_ref=True)} {out_name}({vec_name});
|
||||
""".split("\n")
|
||||
)
|
||||
return code, decl
|
||||
@@ -0,0 +1,109 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import re
|
||||
import textwrap
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
|
||||
# match $identifier or ${identifier} and replace with value in env
|
||||
# If this identifier is at the beginning of whitespace on a line
|
||||
# and its value is a list then it is treated as
|
||||
# block substitution by indenting to that depth and putting each element
|
||||
# of the list on its own line
|
||||
# if the identifier is on a line starting with non-whitespace and a list
|
||||
# then it is comma separated ${,foo} will insert a comma before the list
|
||||
# if this list is not empty and ${foo,} will insert one after.
|
||||
|
||||
|
||||
class CodeTemplate:
|
||||
substitution_str = r"(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})"
|
||||
substitution = re.compile(substitution_str, re.MULTILINE)
|
||||
|
||||
pattern: str
|
||||
filename: str
|
||||
|
||||
@staticmethod
|
||||
def from_file(filename: str) -> CodeTemplate:
|
||||
with open(filename) as f:
|
||||
return CodeTemplate(f.read(), filename)
|
||||
|
||||
def __init__(self, pattern: str, filename: str = "") -> None:
|
||||
self.pattern = pattern
|
||||
self.filename = filename
|
||||
|
||||
def substitute(
|
||||
self, env: Mapping[str, object] | None = None, **kwargs: object
|
||||
) -> str:
|
||||
if env is None:
|
||||
env = {}
|
||||
|
||||
def lookup(v: str) -> object:
|
||||
if env is None:
|
||||
raise AssertionError("env must be non-None")
|
||||
return kwargs[v] if v in kwargs else env[v]
|
||||
|
||||
def indent_lines(indent: str, v: Sequence[object]) -> str:
|
||||
content = "\n".join(
|
||||
itertools.chain.from_iterable(str(e).splitlines() for e in v)
|
||||
)
|
||||
content = textwrap.indent(content, prefix=indent)
|
||||
# Remove trailing whitespace on each line
|
||||
return "\n".join(map(str.rstrip, content.splitlines())).rstrip()
|
||||
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
indent = match.group(1)
|
||||
key = match.group(2)
|
||||
comma_before = ""
|
||||
comma_after = ""
|
||||
if key[0] == "{":
|
||||
key = key[1:-1]
|
||||
if key[0] == ",":
|
||||
comma_before = ", "
|
||||
key = key[1:]
|
||||
if key[-1] == ",":
|
||||
comma_after = ", "
|
||||
key = key[:-1]
|
||||
v = lookup(key)
|
||||
if indent is not None:
|
||||
if not isinstance(v, list):
|
||||
v = [v]
|
||||
return indent_lines(indent, v)
|
||||
elif isinstance(v, list):
|
||||
middle = ", ".join([str(x) for x in v])
|
||||
if len(v) == 0:
|
||||
return middle
|
||||
return comma_before + middle + comma_after
|
||||
else:
|
||||
return str(v)
|
||||
|
||||
return self.substitution.sub(replace, self.pattern)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
c = CodeTemplate(
|
||||
"""\
|
||||
int foo($args) {
|
||||
|
||||
$bar
|
||||
$bar
|
||||
$a+$b
|
||||
}
|
||||
int commatest(int a${,stuff})
|
||||
int notest(int a${,empty,})
|
||||
"""
|
||||
)
|
||||
print(
|
||||
c.substitute(
|
||||
args=["hi", 8],
|
||||
bar=["what", 7],
|
||||
a=3,
|
||||
b=4,
|
||||
stuff=["things...", "others"],
|
||||
empty=[],
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,134 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
from typing import Any, TYPE_CHECKING, TypeVar
|
||||
|
||||
import torchgen.local as local
|
||||
from torchgen.model import (
|
||||
BackendIndex,
|
||||
DispatchKey,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
NativeFunctionsViewGroup,
|
||||
)
|
||||
from torchgen.utils import context, S, T
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterator
|
||||
|
||||
|
||||
# Helper functions for defining generators on things in the model
|
||||
|
||||
F = TypeVar(
|
||||
"F",
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
NativeFunctionsViewGroup,
|
||||
NativeFunction | NativeFunctionsGroup,
|
||||
NativeFunction | NativeFunctionsViewGroup,
|
||||
)
|
||||
|
||||
F2 = TypeVar(
|
||||
"F2",
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
NativeFunction | None,
|
||||
bool,
|
||||
str,
|
||||
)
|
||||
|
||||
F3 = TypeVar("F3", tuple[NativeFunction, Any], list[NativeFunction])
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def native_function_manager(
|
||||
g: NativeFunctionsGroup | NativeFunctionsViewGroup | NativeFunction,
|
||||
) -> Iterator[None]:
|
||||
if isinstance(g, NativeFunctionsGroup):
|
||||
# By default, we associate all errors with structured native functions
|
||||
# with the out variant. In some cases, it might be better to have
|
||||
# a more specific place to hang things; if so, use
|
||||
# native_function_manager again on the inside
|
||||
f = g.out
|
||||
elif isinstance(g, NativeFunctionsViewGroup):
|
||||
# We associate errors with the view operator
|
||||
f = g.view
|
||||
else:
|
||||
f = g
|
||||
with context(lambda: f"in native_functions.yaml line {f.loc}:\n {f.func}"):
|
||||
with local.parametrize(
|
||||
use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors,
|
||||
use_ilistref_for_tensor_lists=f.part_of_structured_group,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
# Given a function that operates on NativeFunction, wrap it into a new function
|
||||
# that sets some appropriate context managers for that native function.
|
||||
# YOU MUST WRAP FUNCTIONS IN THIS for calls to api modules to be sound
|
||||
# (you will get an error if we try to access the local variables without having
|
||||
# set them).
|
||||
def with_native_function(func: Callable[[F], T]) -> Callable[[F], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(f: F) -> T:
|
||||
with native_function_manager(f):
|
||||
return func(f)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def with_native_function_and(func: Callable[[F, F2], T]) -> Callable[[F, F2], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(f: F, f2: F2) -> T:
|
||||
# The first native_function is assumed to be the one with the appropriate context.
|
||||
with native_function_manager(f):
|
||||
return func(f, f2)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def method_with_native_function(func: Callable[[S, F], T]) -> Callable[[S, F], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(slf: S, f: F) -> T:
|
||||
with native_function_manager(f):
|
||||
return func(slf, f)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def method_with_nested_native_function(
|
||||
func: Callable[[S, F3], T],
|
||||
) -> Callable[[S, F3], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(slf: S, f: F3) -> T:
|
||||
with native_function_manager(f[0]):
|
||||
return func(slf, f)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# Convenience decorator for functions that explicitly take in a BackendIndex,
|
||||
# instead of indirectly taking one in as a closure
|
||||
def with_native_function_and_index(
|
||||
func: Callable[[F, BackendIndex], T],
|
||||
) -> Callable[[F, BackendIndex], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(f: F, backend_index: BackendIndex) -> T:
|
||||
with native_function_manager(f):
|
||||
return func(f, backend_index)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# Convenience decorator for functions that explicitly take in a Dict of BackendIndices
|
||||
def with_native_function_and_indices(
|
||||
func: Callable[[F, dict[DispatchKey, BackendIndex]], T],
|
||||
) -> Callable[[F, dict[DispatchKey, BackendIndex]], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(f: F, backend_indices: dict[DispatchKey, BackendIndex]) -> T:
|
||||
with native_function_manager(f):
|
||||
return func(f, backend_indices)
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,19 @@
|
||||
from torchgen.dest.lazy_ir import (
|
||||
generate_non_native_lazy_ir_nodes as generate_non_native_lazy_ir_nodes,
|
||||
GenLazyIR as GenLazyIR,
|
||||
GenLazyNativeFuncDefinition as GenLazyNativeFuncDefinition,
|
||||
GenLazyShapeInferenceDefinition as GenLazyShapeInferenceDefinition,
|
||||
)
|
||||
from torchgen.dest.native_functions import (
|
||||
compute_native_function_declaration as compute_native_function_declaration,
|
||||
)
|
||||
from torchgen.dest.register_dispatch_key import (
|
||||
gen_registration_headers as gen_registration_headers,
|
||||
gen_registration_helpers as gen_registration_helpers,
|
||||
RegisterDispatchKey as RegisterDispatchKey,
|
||||
)
|
||||
from torchgen.dest.ufunc import (
|
||||
compute_ufunc_cpu as compute_ufunc_cpu,
|
||||
compute_ufunc_cpu_kernel as compute_ufunc_cpu_kernel,
|
||||
compute_ufunc_cuda as compute_ufunc_cuda,
|
||||
)
|
||||
@@ -0,0 +1,718 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from abc import ABC
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torchgen.api.dispatcher as dispatcher
|
||||
from torchgen.api.lazy import (
|
||||
getValueT,
|
||||
isValueType,
|
||||
LazyArgument,
|
||||
LazyIrProperties,
|
||||
LazyIrSchema,
|
||||
tensorListValueT,
|
||||
)
|
||||
from torchgen.api.translate import translate
|
||||
from torchgen.api.types import (
|
||||
BaseCType,
|
||||
Binding,
|
||||
deviceT,
|
||||
DispatcherSignature,
|
||||
kernel_signature,
|
||||
NativeSignature,
|
||||
OptionalCType,
|
||||
VectorCType,
|
||||
)
|
||||
from torchgen.context import method_with_native_function
|
||||
from torchgen.dest.lazy_ts_lowering import ts_lowering_body
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BackendIndex,
|
||||
BackendMetadata,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
)
|
||||
|
||||
|
||||
def node_ctor_arg_rvalue_string(arg: LazyArgument) -> str:
|
||||
"""
|
||||
Given a LazyArgument,
|
||||
generate a c++ string for materializing an rvalue of that arg for passing into
|
||||
a lazy Node constructor.
|
||||
"""
|
||||
|
||||
# TODO: Matching on CType seems wrong; should be matching on Type
|
||||
if isValueType(arg.lazy_type):
|
||||
if isinstance(arg.lazy_type, BaseCType):
|
||||
if arg.is_wrapped_scalar:
|
||||
return f"node_{arg.name}"
|
||||
elif arg.lazy_type.type is tensorListValueT:
|
||||
return f"lazy_{arg.name}_tensorlist"
|
||||
elif arg.is_symint_or_list:
|
||||
return f"GetSymIntValue({arg.name})"
|
||||
return f"lazy_{arg.name}->GetIrValue()"
|
||||
elif isinstance(arg.lazy_type, OptionalCType):
|
||||
if arg.is_symint_or_list:
|
||||
# TODO: I don't understand when you should put lazy_ in the name
|
||||
# or not
|
||||
return f"{arg.name} ? std::make_optional(GetSymIntValue(*{arg.name})) : ::std::nullopt"
|
||||
elif arg.is_wrapped_scalar:
|
||||
return f"node_{arg.name}"
|
||||
return (
|
||||
f"lazy_{arg.name} ? "
|
||||
f"std::make_optional(lazy_{arg.name}->GetIrValue()) : "
|
||||
"::std::nullopt"
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"TODO not sure if there are other valid types to handle here ({arg.lazy_type})"
|
||||
)
|
||||
else:
|
||||
# NB: this is here because right now we aren't treating SymInt[] as a
|
||||
# value type; when we do this needs to move above
|
||||
# NB: we cannot test arg.lazy_type as we've already specified it is an
|
||||
# int64_t and so we cannot distinguish between SymInt and int64_t
|
||||
if isinstance(arg.orig_type, ListType) and arg.orig_type.elem == BaseType(
|
||||
BaseTy.SymInt
|
||||
):
|
||||
if arg.symint:
|
||||
return f"GetSymIntArrayRefValue({arg.name})"
|
||||
else:
|
||||
return f"std::vector<int64_t>({arg.name}.begin(), {arg.name}.end())"
|
||||
elif isinstance(arg.lazy_type, VectorCType) and isinstance(
|
||||
arg.lazy_type.elem, BaseCType
|
||||
):
|
||||
return f"std::vector<{arg.lazy_type.elem.type}>({arg.name}.begin(), {arg.name}.end())"
|
||||
elif (
|
||||
isinstance(arg.lazy_type, OptionalCType)
|
||||
and isinstance(arg.lazy_type.elem, VectorCType)
|
||||
and isinstance(arg.lazy_type.elem.elem, BaseCType)
|
||||
):
|
||||
return f"torch::lazy::ToOptionalVector<{arg.lazy_type.elem.elem.type}>({arg.name})"
|
||||
else:
|
||||
return f"{arg.name}"
|
||||
|
||||
|
||||
def node_ctor_inputs(schema: LazyIrSchema) -> str:
|
||||
"""
|
||||
Produce a formatted string with the arguments as passed into the constructor of a node class.
|
||||
"""
|
||||
node_ctor_values = [
|
||||
node_ctor_arg_rvalue_string(arg) for arg in schema.filtered_args()
|
||||
]
|
||||
return ", ".join(node_ctor_values)
|
||||
|
||||
|
||||
def gen_fallback_code(
|
||||
schema: LazyIrSchema,
|
||||
sig: DispatcherSignature | NativeSignature,
|
||||
overload_name: str,
|
||||
) -> str:
|
||||
"""
|
||||
Generate code that falls back to eager conditioned on a predicate
|
||||
"""
|
||||
dispatcher_sig = DispatcherSignature.from_schema(schema.func)
|
||||
exprs = translate(sig.arguments(), dispatcher_sig.arguments())
|
||||
fallback_args = ",\n ".join([a.expr for a in exprs])
|
||||
if len(overload_name):
|
||||
aten_op_str = f"ATEN_OP2({schema.aten_name}, {overload_name})"
|
||||
else:
|
||||
aten_op_str = f"ATEN_OP({schema.aten_name})"
|
||||
return f"""
|
||||
if (force_eager_fallback({aten_symbol(schema)})) {{
|
||||
return at::native::call_fallback_fn_symint<<c_eager_fallback, {aten_op_str}>::call(
|
||||
{fallback_args}
|
||||
);
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
def aten_symbol(schema: LazyIrSchema) -> str:
|
||||
missing_interned_strings = {
|
||||
"sigmoid_backward",
|
||||
}
|
||||
if schema.aten_name in missing_interned_strings:
|
||||
return f'c10::Symbol::fromQualString("aten::{schema.aten_name}")'
|
||||
|
||||
if not schema.aten_name.startswith("at::"):
|
||||
return f"at::aten::{schema.aten_name}"
|
||||
else:
|
||||
return schema.aten_name
|
||||
|
||||
|
||||
# converts all tensor-like arguments to meta tensors. Returns:
|
||||
# (1) a string containing all of the logic that does the conversions.
|
||||
# (2) a context, to be used by translate(), with all of the relevant bindings.
|
||||
def convert_to_meta_tensors(sig: DispatcherSignature) -> tuple[str, list[Binding]]:
|
||||
context: list[Binding] = []
|
||||
unwrapped_tensor_args: list[str] = []
|
||||
for arg in sig.arguments():
|
||||
if isinstance(arg.argument, Argument) and arg.argument.type.is_tensor_like():
|
||||
unwrapped_name = f"{arg.name}_meta"
|
||||
unwrapped_tensor_args.append(
|
||||
f"auto {unwrapped_name} = to_meta({arg.name});"
|
||||
)
|
||||
context.append(arg.with_name(unwrapped_name))
|
||||
else:
|
||||
context.append(arg)
|
||||
unwrap_tensor_args_str = "\n ".join(unwrapped_tensor_args)
|
||||
return unwrap_tensor_args_str, context
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GenLazyIR(ABC):
|
||||
backend_index: BackendIndex
|
||||
backend_name: str
|
||||
node_base: str
|
||||
use_lazy_shape: bool
|
||||
|
||||
@method_with_native_function
|
||||
def __call__(self, f: NativeFunctionsGroup | NativeFunction) -> list[str]:
|
||||
func = f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func
|
||||
metadata = self.backend_index.get_kernel(
|
||||
f.functional if isinstance(f, NativeFunctionsGroup) else f
|
||||
)
|
||||
schema = LazyIrSchema(
|
||||
func, symint=metadata is not None and metadata.supports_symint()
|
||||
)
|
||||
return self.gen(schema)
|
||||
|
||||
# there is no lowering functionality generated unless this IR base class is subclassed and
|
||||
# implemented as a backend-specific node
|
||||
def lowering_function(self, schema: LazyIrSchema) -> str:
|
||||
return ""
|
||||
|
||||
def create_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:
|
||||
return ""
|
||||
|
||||
def can_be_reused_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:
|
||||
return f"""bool CanBeReused({node_ctor_args}) const {{
|
||||
return false;
|
||||
}}"""
|
||||
|
||||
def node_base_ctor_call(self, schema: LazyIrSchema) -> str:
|
||||
value_args = schema.filtered_args(values=True, scalars=False)
|
||||
# backends can customize the way the node base class constructor is called,
|
||||
# as long as all of its arguments can be generated from information available from the schema
|
||||
base_ctor_value_args_list = []
|
||||
for arg in value_args:
|
||||
if isinstance(arg.lazy_type, (BaseCType, VectorCType)):
|
||||
base_ctor_value_args_list.append(f"{arg.name}")
|
||||
elif isinstance(arg.lazy_type, OptionalCType):
|
||||
base_ctor_value_args_list.append(f"{arg.name}.value_or(kNullValue)")
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"Unsupported type ({arg.lazy_type}) - add support if necessary"
|
||||
)
|
||||
base_ctor_value_args = ", ".join(base_ctor_value_args_list)
|
||||
|
||||
scalar_args = schema.filtered_args(values=False, scalars=True)
|
||||
|
||||
# Shape construction.
|
||||
# Conditionally build shape depending on specified shape property
|
||||
if schema.properties.ShapePrecompute:
|
||||
shape_ctor_arg = "std::move(shapes),"
|
||||
elif schema.properties.ShapeCompute:
|
||||
shape_args = [a.name for a in value_args]
|
||||
shape_args.extend(a.name for a in scalar_args)
|
||||
shape_ctor_arg = f"compute_shape_{schema.name}({', '.join(shape_args)}),"
|
||||
elif schema.properties.ShapeCache:
|
||||
shape_args = [f"operand({i})" for i in range(len(value_args))]
|
||||
shape_args.extend(a.name for a in scalar_args)
|
||||
shape_ctor_arg = f"[&](){{ return compute_shape_{schema.name}({', '.join(shape_args)})[0]; }},"
|
||||
else:
|
||||
shape_ctor_arg = ""
|
||||
|
||||
scalar_hashes = ", ".join(f"{a.name}" for a in scalar_args)
|
||||
|
||||
return f"""{self.node_base}(
|
||||
{schema.node_name}::ClassOpKind(),
|
||||
OpList{{{base_ctor_value_args}}},
|
||||
{shape_ctor_arg}
|
||||
/* num_outputs */ {len(schema.returns)},
|
||||
torch::lazy::MHash({scalar_hashes}))"""
|
||||
|
||||
def gen(self, schema: LazyIrSchema) -> list[str]:
|
||||
opkind = schema.opkind or aten_symbol(schema)
|
||||
|
||||
# for now, we just want one IR class decl and soon after also the method defs
|
||||
# and we use the functional version not out/inplace.
|
||||
all_args = schema.filtered_args()
|
||||
scalar_args = schema.filtered_args(values=False, scalars=True)
|
||||
|
||||
ctor_args = [f"const {i.lazy_type.cpp_type()}& {i.name}" for i in all_args]
|
||||
reuse_ctor_args = ", ".join(ctor_args)
|
||||
if self.use_lazy_shape and schema.properties.ShapePrecompute:
|
||||
ctor_args.append("std::vector<torch::lazy::Shape>&& shapes")
|
||||
node_ctor_args = ", ".join(ctor_args)
|
||||
|
||||
scalar_initializers = ",\n ".join(
|
||||
[
|
||||
# This code is just special casing the mapping from string_view -> strings
|
||||
f"{a.name}({a.name}.has_value() ? ::std::make_optional(std::string(*{a.name})) : ::std::nullopt)"
|
||||
if a.lazy_type.cpp_type() == "::std::optional<c10::string_view>"
|
||||
else f"{a.name}({a.name})"
|
||||
for a in scalar_args
|
||||
]
|
||||
)
|
||||
if len(scalar_initializers):
|
||||
scalar_initializers = f",\n {scalar_initializers}"
|
||||
scalar_decls = "\n ".join(
|
||||
[
|
||||
f"std::string {a.name};"
|
||||
if a.lazy_type.cpp_type() == "c10::string_view"
|
||||
else f"::std::optional<std::string> {a.name};"
|
||||
if a.lazy_type.cpp_type() == "::std::optional<c10::string_view>"
|
||||
else f"{a.lazy_type.cpp_type()} {a.name};"
|
||||
for a in scalar_args
|
||||
]
|
||||
)
|
||||
optional_values = [
|
||||
arg.name
|
||||
for arg in schema.filtered_args(values=True, scalars=False)
|
||||
if isinstance(arg.lazy_type, OptionalCType)
|
||||
]
|
||||
has_optional_decls = "\n ".join(
|
||||
[f"bool has_{value}: 1;" for value in optional_values]
|
||||
)
|
||||
has_optional_defs = "\n ".join(
|
||||
[f"has_{value} = !!{value};" for value in optional_values]
|
||||
)
|
||||
members_to_string = []
|
||||
for arg in scalar_args:
|
||||
if isinstance(arg.lazy_type, OptionalCType):
|
||||
value = f"{arg.name}.value()"
|
||||
if arg.is_generator:
|
||||
value = '"torch.Generator()"'
|
||||
members_to_string.append(
|
||||
f"""if ({arg.name}.has_value()) {{
|
||||
ss << ", {arg.name}=" << {value};
|
||||
}} else {{
|
||||
ss << ", {arg.name}=null";
|
||||
}}"""
|
||||
)
|
||||
else:
|
||||
members_to_string.append(f'ss << ", {arg.name}=" << {arg.name};')
|
||||
members_to_string_str = "\n ".join(members_to_string)
|
||||
|
||||
return [
|
||||
f"""\
|
||||
class {schema.node_name} : public {self.node_base} {{
|
||||
public:
|
||||
static torch::lazy::OpKind ClassOpKind() {{
|
||||
return torch::lazy::OpKind({opkind});
|
||||
}}
|
||||
|
||||
{schema.node_name}({node_ctor_args})
|
||||
: {self.node_base_ctor_call(schema)}{scalar_initializers}
|
||||
{{
|
||||
{has_optional_defs}
|
||||
}}
|
||||
|
||||
std::string ToString() const override {{
|
||||
std::stringstream ss;
|
||||
ss << {self.node_base}::ToString();
|
||||
{members_to_string_str}
|
||||
return ss.str();
|
||||
}}
|
||||
|
||||
{self.create_function(schema, reuse_ctor_args)}
|
||||
|
||||
{self.can_be_reused_function(schema, reuse_ctor_args)}
|
||||
|
||||
{self.lowering_function(schema)}
|
||||
|
||||
{scalar_decls}
|
||||
{has_optional_decls}
|
||||
|
||||
}};
|
||||
|
||||
""",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GenTSLazyIR(GenLazyIR):
|
||||
def lowering_function(self, schema: LazyIrSchema) -> str:
|
||||
signature = """
|
||||
torch::lazy::TSOpVector Lower(
|
||||
std::shared_ptr<torch::jit::GraphFunction> function,
|
||||
torch::lazy::TSLoweringContext* loctx) const override"""
|
||||
|
||||
if schema.properties.LowerDeclOnly:
|
||||
return f"{signature};"
|
||||
elif schema.properties.Lower:
|
||||
return f"""{signature} {{
|
||||
{ts_lowering_body(schema)}
|
||||
}}
|
||||
"""
|
||||
else:
|
||||
return ""
|
||||
|
||||
def create_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:
|
||||
signature = f"static NodePtr Create({node_ctor_args})"
|
||||
if schema.properties.CreateFnDeclOnly:
|
||||
return f"{signature};"
|
||||
elif not schema.properties.CreateFn:
|
||||
return ""
|
||||
return f"""{signature} {{
|
||||
return ReuseOrMakeNode<{schema.node_name}>(data);
|
||||
}}"""
|
||||
|
||||
def can_be_reused_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:
|
||||
signature = f"bool CanBeReused({node_ctor_args}) const"
|
||||
if schema.properties.CanBeReusedDeclOnly:
|
||||
return f"{signature};"
|
||||
elif not schema.properties.CanBeReused:
|
||||
return ""
|
||||
value_comparison = []
|
||||
for arg in itertools.chain(schema.positional_values, schema.keyword_values):
|
||||
if isinstance(arg.lazy_type, OptionalCType):
|
||||
value_comparison.append(
|
||||
f"nullable_operand(i++) == {arg.name}.value_or(kNullValue)"
|
||||
)
|
||||
else:
|
||||
value_comparison.append(f"operand(i++) == {arg.name}")
|
||||
for arg in itertools.chain(schema.positional_scalars, schema.keyword_scalars):
|
||||
if isinstance(arg.lazy_type, OptionalCType):
|
||||
value_comparison.append(
|
||||
f"((!this->{arg.name}&&!{arg.name}) || (this->{arg.name}&&{arg.name} && *(this->{arg.name}) == *{arg.name}))"
|
||||
)
|
||||
else:
|
||||
value_comparison.append(f"this->{arg.name} == {arg.name}")
|
||||
value_comparison_str = " &&\n ".join(value_comparison)
|
||||
|
||||
return f"""{signature} {{
|
||||
size_t i = 0;
|
||||
return ({value_comparison_str});
|
||||
}}"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GenLazyNativeFuncDefinition:
|
||||
class_method_name: str
|
||||
backend_index: BackendIndex
|
||||
tensor_class: str
|
||||
gen_forced_fallback_code: bool
|
||||
backend_namespace: str
|
||||
get_tensorlist: str
|
||||
get_tensor_or_wrap_number: str
|
||||
try_get_tensor: str
|
||||
metrics_counter: str
|
||||
create_tensor: str
|
||||
create_from_first_tensor: bool
|
||||
create_aten_from_ltc_tensor: str
|
||||
tuple_aten_from_ltc_tensors: str
|
||||
lazy_tensor_ptr: str
|
||||
get_device_fn: str
|
||||
|
||||
def lazy_tensor_decls(self, func: NativeFunction, schema: LazyIrSchema) -> str:
|
||||
value_args = schema.filtered_args(values=True, scalars=False)
|
||||
# Generates lazy_{name} variables for LazyTensors wrapping input tensors
|
||||
lazy_tensor_decls: list[str] = []
|
||||
for arg in value_args:
|
||||
if arg.is_wrapped_scalar:
|
||||
if isinstance(arg.lazy_type, OptionalCType):
|
||||
lazy_tensor_decls.append(
|
||||
f"""auto node_{arg.name} = {arg.name} ?
|
||||
std::make_optional(torch::lazy::LazyGraphExecutor::Get()->
|
||||
GetIrValueForScalarFromCodegen(*{arg.name}, *common_device)):
|
||||
::std::nullopt;"""
|
||||
)
|
||||
else:
|
||||
lazy_tensor_decls.append(
|
||||
f"""auto node_{arg.name} = torch::lazy::LazyGraphExecutor::Get()->
|
||||
GetIrValueForScalarFromCodegen({arg.name}, *common_device);"""
|
||||
)
|
||||
elif arg.is_symint_or_list:
|
||||
continue # values are extracted in isValueType
|
||||
elif isinstance(arg.lazy_type, BaseCType):
|
||||
if arg.lazy_type.type is tensorListValueT:
|
||||
lazy_tensor_decls.append(
|
||||
f"auto lazy_{arg.name}_tensorlist = "
|
||||
f"{self.backend_namespace}::{self.get_tensorlist}({arg.name});"
|
||||
)
|
||||
else:
|
||||
lazy_tensor_decls.append(
|
||||
f"{self.lazy_tensor_ptr} lazy_{arg.name} = "
|
||||
f"{self.backend_namespace}::{self.get_tensor_or_wrap_number}({arg.name}, *common_device);"
|
||||
)
|
||||
elif isinstance(arg.lazy_type, OptionalCType):
|
||||
if arg.lazy_type.elem != BaseCType(getValueT()):
|
||||
raise AssertionError(
|
||||
f"Expected OptionalCType elem to be {BaseCType(getValueT())}, "
|
||||
f"got {arg.lazy_type.elem}"
|
||||
)
|
||||
# TODO(alanwaketan): Maybe we want to apply GetLtcTensorOrCreateForWrappedNumber here, but hold it
|
||||
# until we encounter a real world example.
|
||||
lazy_tensor_decls.append(
|
||||
f"{self.lazy_tensor_ptr} lazy_{arg.name} = "
|
||||
f"{self.backend_namespace}::{self.try_get_tensor}({arg.name}.value_or(at::Tensor()));"
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"TODO not sure if there are other valid types to handle here ({arg.lazy_type})"
|
||||
)
|
||||
return ("\n ").join(lazy_tensor_decls)
|
||||
|
||||
def force_eager_fallback(
|
||||
self,
|
||||
func: NativeFunction,
|
||||
schema: LazyIrSchema,
|
||||
metadata: BackendMetadata,
|
||||
sig: DispatcherSignature | NativeSignature,
|
||||
) -> str:
|
||||
if self.gen_forced_fallback_code:
|
||||
return gen_fallback_code(
|
||||
schema, sig, overload_name=func.func.name.overload_name
|
||||
)
|
||||
return ""
|
||||
|
||||
def metrics(self, func: NativeFunction, schema: LazyIrSchema) -> str:
|
||||
return f"{self.metrics_counter};"
|
||||
|
||||
def get_device(self, func: NativeFunction, schema: LazyIrSchema) -> str:
|
||||
value_args = schema.filtered_args(values=True, scalars=False)
|
||||
scalar_args = schema.filtered_args(values=False, scalars=True)
|
||||
value_types_names = [f"{a.name}" for a in value_args if not a.is_wrapped_scalar]
|
||||
optional_device = OptionalCType(BaseCType(deviceT))
|
||||
optional_devices = [
|
||||
a.name for a in scalar_args if a.lazy_type == optional_device
|
||||
]
|
||||
if len(value_types_names) == 0 and len(optional_devices) == 0:
|
||||
raise AssertionError("Expected at least one Value or Device type")
|
||||
get_device_str = (
|
||||
f"{self.get_device_fn}({', '.join(value_types_names + optional_devices)})"
|
||||
)
|
||||
return f"""auto common_device = {get_device_str};
|
||||
TORCH_INTERNAL_ASSERT(common_device);
|
||||
"""
|
||||
|
||||
def shape_inference(self, func: NativeFunction, schema: LazyIrSchema) -> str:
|
||||
metadata = self.backend_index.get_kernel(func)
|
||||
if metadata is None:
|
||||
raise AssertionError(f"No kernel metadata found for {func.func.name}")
|
||||
all_args = schema.filtered_args()
|
||||
returns_length = len(schema.returns)
|
||||
# call the meta kernel if it exists, to compute output shape/dtype for our IR
|
||||
# Note [Generated LTC Shape Functions]
|
||||
# LTC uses meta tensors from core to do shape inference when possible, and otherwise
|
||||
# we generate a shape function declaration that needs to be manually implemented.
|
||||
# How do we detect which ops are eligible to use meta tensors?
|
||||
# In general we should be able to use meta tensors not just on structured operators,
|
||||
# but also on composite operators that are implemented in terms of structured kernels.
|
||||
# We don't currently have a way of knowing at codegen time which ops are implemented that way.
|
||||
# This is the case for all view and view_copy operators however, so we're going to
|
||||
# use them specifically for all of the view_copy ops (instead of manually writing shape rules for all of them).
|
||||
is_view_copy_op = "view_copy" in func.tags
|
||||
is_structured = func.structured or func.structured_delegate is not None
|
||||
if is_structured or is_view_copy_op:
|
||||
meta_out = """
|
||||
std::vector<torch::lazy::Shape> shapes{torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};"""
|
||||
if returns_length > 1:
|
||||
|
||||
def this_shape(i: int) -> str:
|
||||
return f"torch::lazy::Shape(std::get<{i}>(out_meta).scalar_type(), std::get<{i}>(out_meta).sizes().vec())"
|
||||
|
||||
shapes_str = ",".join([this_shape(i) for i in range(returns_length)])
|
||||
meta_out = "std::vector<torch::lazy::Shape> shapes{" + shapes_str + "};"
|
||||
|
||||
# Convert tensor args to the meta device and call it.
|
||||
# (We can't pass in the input tensors directly, because they are "functional wrappers".
|
||||
# If any of the meta kernels call a tensor op and redispatch, we don't want to hit the functionalize kernels.)
|
||||
# Even at::meta:: functions might redispatch, e.g. if they call into view ops.
|
||||
dispatcher_sig = DispatcherSignature.from_schema(func.func)
|
||||
meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig)
|
||||
meta_call_args = [
|
||||
e.expr
|
||||
for e in translate(
|
||||
meta_call_ctx, dispatcher_sig.arguments(), method=False
|
||||
)
|
||||
]
|
||||
if is_view_copy_op:
|
||||
# view_copy ops always have a CompositeExplicitAutogradNonFunctional kernel
|
||||
if not func.has_composite_explicit_autograd_non_functional_kernel:
|
||||
raise AssertionError(
|
||||
f"view_copy op {func.func.name} must have "
|
||||
"CompositeExplicitAutogradNonFunctional kernel"
|
||||
)
|
||||
dispatch_ns = "compositeexplicitautogradnonfunctional"
|
||||
else:
|
||||
dispatch_ns = "meta"
|
||||
aten_name = schema.aten_name
|
||||
# TODO: this is trolling
|
||||
if func.func.has_symint() and metadata.supports_symint():
|
||||
aten_name += "_symint"
|
||||
shape_str = f"""\
|
||||
{meta_conversion_str}
|
||||
auto out_meta = at::{dispatch_ns}::{aten_name}({", ".join(meta_call_args)});
|
||||
{meta_out}"""
|
||||
else:
|
||||
shape_sig = ComputeShapeSignature(
|
||||
metadata.kernel, func, symint=metadata.supports_symint()
|
||||
)
|
||||
shape_str = f"""
|
||||
auto shapes = {shape_sig.shape_call};"""
|
||||
|
||||
shape_str += f"""
|
||||
TORCH_INTERNAL_ASSERT(shapes.size() == {returns_length});"""
|
||||
|
||||
# Calculating which dimensions are symbolic
|
||||
func_schema_str = "aten::" + str(func.func)
|
||||
shape_str += f"""
|
||||
if(torch::lazy::symbolicShapeEnabled()){{
|
||||
std::vector<torch::jit::IValue> inputs = {{ {", ".join(str(a.name) for a in all_args)} }};
|
||||
const char* schema_str = "{func_schema_str}";
|
||||
applySymbolicShapesOnLT(schema_str, inputs, shapes);
|
||||
}}
|
||||
"""
|
||||
return shape_str
|
||||
|
||||
def build_ir_node(self, func: NativeFunction, schema: LazyIrSchema) -> str:
|
||||
node_ctor_input_str = node_ctor_inputs(schema)
|
||||
return f"""torch::lazy::NodePtr node = torch::lazy::ReuseNode<{schema.node_name}>({node_ctor_input_str});
|
||||
if (!node) {{
|
||||
{self.shape_inference(func, schema)}
|
||||
node = torch::lazy::MakeNode<{schema.node_name}>({node_ctor_input_str}, std::move(shapes));
|
||||
CacheNode(node);
|
||||
}}
|
||||
"""
|
||||
|
||||
def create_lazy_tensor(self, first_tensor_name: str | None = None) -> str:
|
||||
# xla uses an instance method for tensor creation, for the time being
|
||||
if self.create_from_first_tensor:
|
||||
# TODO(whc) remove this if XLA switches to using static method for creation
|
||||
if first_tensor_name is None:
|
||||
raise AssertionError("Requires first tensor to create lazy tensor")
|
||||
return f"{first_tensor_name}.{self.create_tensor}"
|
||||
return f"{self.backend_namespace}::{self.create_tensor}"
|
||||
|
||||
def return_aten_tensor(self, func: NativeFunction, schema: LazyIrSchema) -> str:
|
||||
returns_length = len(schema.returns)
|
||||
value_args = schema.filtered_args(values=True, scalars=False)
|
||||
value_types_names = [f"{a.name}" for a in value_args if not a.is_wrapped_scalar]
|
||||
first_tensor_name = value_types_names[0] if len(value_types_names) > 0 else None
|
||||
bridge_str = f"""auto result = {self.create_aten_from_ltc_tensor}(
|
||||
{self.create_lazy_tensor(first_tensor_name)}(std::move(node), *common_device));"""
|
||||
|
||||
if returns_length > 1:
|
||||
if len(value_types_names) == 0:
|
||||
raise AssertionError(
|
||||
"Code below assumes there is at least one tensor arg"
|
||||
)
|
||||
bridge_str = f"""std::vector<{self.lazy_tensor_ptr}> lazy_tensors;
|
||||
for (int i = 0; i < {returns_length}; i++) {{
|
||||
lazy_tensors.push_back({self.create_lazy_tensor(first_tensor_name)}({getValueT()}(node, i), *common_device));
|
||||
}}
|
||||
auto result = {self.tuple_aten_from_ltc_tensors}<{returns_length}>(lazy_tensors);"""
|
||||
|
||||
if schema.name.name.inplace or func.func.is_out_fn():
|
||||
if returns_length != 1:
|
||||
raise AssertionError(
|
||||
"We assumed there was no such case where an op is an in-place variant "
|
||||
f"and has tuple outputs, but got tuple of len {returns_length}."
|
||||
)
|
||||
bridge_str = f"""lazy_{first_tensor_name}->SetInPlaceIrValue(node);
|
||||
auto& result = {first_tensor_name};"""
|
||||
|
||||
bridge_str += """
|
||||
return result;"""
|
||||
return bridge_str
|
||||
|
||||
@method_with_native_function
|
||||
def __call__(self, func: NativeFunction) -> list[str]:
|
||||
sig = kernel_signature(func, self.backend_index)
|
||||
metadata = self.backend_index.get_kernel(func)
|
||||
if metadata is None:
|
||||
raise AssertionError(f"No kernel metadata found for {func.func.name}")
|
||||
schema = LazyIrSchema(func.func, symint=metadata.supports_symint())
|
||||
return [
|
||||
f"""\
|
||||
{sig.decl(name=f"{self.class_method_name}::{metadata.kernel}")} {{
|
||||
{self.force_eager_fallback(func, schema, metadata, sig)}
|
||||
{self.metrics(func, schema)}
|
||||
{self.get_device(func, schema)}
|
||||
{self.lazy_tensor_decls(func, schema)}
|
||||
{self.build_ir_node(func, schema)}
|
||||
{self.return_aten_tensor(func, schema)}
|
||||
}}\n
|
||||
"""
|
||||
]
|
||||
|
||||
|
||||
class ComputeShapeSignature:
|
||||
"""
|
||||
Here we use the base name as the suffix of the signature to avoid generating for in-place variants.
|
||||
"""
|
||||
|
||||
def __init__(self, kernel_name: str, f: NativeFunction, *, symint: bool) -> None:
|
||||
self.__schema = LazyIrSchema(f.func, symint=symint)
|
||||
self.__dispatch_args = ", ".join(
|
||||
[a.decl() for a in dispatcher.arguments(f.func, symint=symint)]
|
||||
)
|
||||
self.__call_args = ", ".join(
|
||||
[f"{arg.name}" for arg in self.__schema.filtered_args(generator=True)]
|
||||
)
|
||||
self.__kernel_name = kernel_name
|
||||
|
||||
def __decl_suffix(self) -> str:
|
||||
return f"{self.__kernel_name}({self.__dispatch_args})"
|
||||
|
||||
def __call_suffix(self) -> str:
|
||||
return f"{self.__kernel_name}({self.__call_args})"
|
||||
|
||||
@property
|
||||
def shape_decl(self) -> str:
|
||||
return f"TORCH_API std::vector<torch::lazy::Shape> compute_shape_{self.__decl_suffix()}"
|
||||
|
||||
@property
|
||||
def shape_call(self) -> str:
|
||||
return f"torch::lazy::compute_shape_{self.__call_suffix()}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GenLazyShapeInferenceDefinition:
|
||||
backend_index: BackendIndex
|
||||
tensor_class: str
|
||||
|
||||
@method_with_native_function
|
||||
def __call__(self, f: NativeFunction) -> list[str]:
|
||||
metadata = self.backend_index.get_kernel(f)
|
||||
if metadata is None:
|
||||
raise AssertionError(f"No kernel metadata found for {f.func.name}")
|
||||
|
||||
# See Note [Generated LTC Shape Functions]
|
||||
is_view_copy_op = "view_copy" in f.tags
|
||||
is_structured = f.structured or f.structured_delegate is not None
|
||||
if is_structured or is_view_copy_op:
|
||||
return []
|
||||
else:
|
||||
shape_sig = ComputeShapeSignature(
|
||||
metadata.kernel, f, symint=metadata.supports_symint()
|
||||
)
|
||||
return ["\n".join([f"{shape_sig.shape_decl};"])]
|
||||
|
||||
|
||||
def generate_non_native_lazy_ir_nodes(
|
||||
non_native: list[dict[str, Any]], gen_lazy_ir: GenLazyIR
|
||||
) -> list[str]:
|
||||
"""Generate the non-native lazy IR node classes"""
|
||||
nodes = []
|
||||
for op in non_native:
|
||||
# Set default properties for Non-Native IRs
|
||||
properties = LazyIrProperties("ShapeCache", "CanBeReused", "LowerDeclOnly")
|
||||
for p in op.get("properties", []):
|
||||
setattr(properties, p, True)
|
||||
|
||||
# non-native is assumed to want symint bindings if you wrote symint
|
||||
schema = LazyIrSchema(FunctionSchema.parse(op["func"]), properties, symint=True)
|
||||
schema.opkind = op.get("opkind")
|
||||
nodes.append(gen_lazy_ir.gen(schema)[0])
|
||||
|
||||
return nodes
|
||||
@@ -0,0 +1,48 @@
|
||||
from torchgen.api.lazy import LazyArgument, LazyIrSchema
|
||||
from torchgen.api.types import OptionalCType
|
||||
|
||||
|
||||
def ts_lowering_body(schema: LazyIrSchema) -> str:
|
||||
# for now, we just want one IR class decl and soon after also the method defs
|
||||
# and we use the functional version not out/inplace.
|
||||
emplace_arguments = []
|
||||
|
||||
def get_value(arg: LazyArgument) -> str:
|
||||
if isinstance(arg.lazy_type, OptionalCType):
|
||||
return f"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr"
|
||||
return "loctx->GetOutputOp(operand(i++))"
|
||||
|
||||
for arg in schema.positional_args:
|
||||
if arg.is_lazy_value:
|
||||
emplace_arguments.append(get_value(arg))
|
||||
continue
|
||||
emplace_arguments.append(f'"{arg.name}", {arg.name}')
|
||||
|
||||
emplace_arguments_str = "\n ".join(
|
||||
[f"arguments.emplace_back({a});" for a in emplace_arguments]
|
||||
)
|
||||
emplace_kwarg_values = [
|
||||
f'"{arg.name}", {get_value(arg)}' for arg in schema.keyword_values
|
||||
]
|
||||
emplace_kwarg_scalars = [
|
||||
f'"{arg.name}", {arg.name}' for arg in schema.keyword_scalars
|
||||
]
|
||||
emplace_kwarguments = "\n ".join(
|
||||
[
|
||||
f"kwarguments.emplace_back({a});"
|
||||
for a in emplace_kwarg_values + emplace_kwarg_scalars
|
||||
]
|
||||
)
|
||||
return f"""\
|
||||
std::vector<torch::jit::NamedValue> arguments;
|
||||
std::vector<torch::jit::NamedValue> kwarguments;
|
||||
arguments.reserve({len(emplace_arguments)});
|
||||
kwarguments.reserve({len(emplace_kwarg_values + emplace_kwarg_scalars)});
|
||||
size_t i = 0;
|
||||
{emplace_arguments_str}
|
||||
{emplace_kwarguments}
|
||||
torch::lazy::TSOpVector {schema.aten_name}_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);
|
||||
TORCH_CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)});
|
||||
|
||||
return {schema.aten_name}_out;
|
||||
"""
|
||||
@@ -0,0 +1,84 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torchgen.api.meta as meta
|
||||
import torchgen.api.structured as structured
|
||||
from torchgen.api.types import kernel_signature
|
||||
from torchgen.context import with_native_function_and_index
|
||||
from torchgen.model import BackendIndex, NativeFunction, NativeFunctionsGroup
|
||||
from torchgen.utils import mapMaybe
|
||||
|
||||
|
||||
def torch_api_key_word_prefix(bankend_index: BackendIndex) -> str:
|
||||
if bankend_index.external:
|
||||
return ""
|
||||
|
||||
# Although Intel GPU ATen library is out-of-tree, it still utilizes torchgen to produce structured
|
||||
# kernels. Regarding these produced structured kernels, they should be visible for the Intel GPU ATen
|
||||
# library. Therefore, we need to add "TORCH_XPU_API" prefix to these structured kernels,
|
||||
# rather than "TORCH_API". Because the semantic of "TORCH_API" is "hidden" for out-of-tree backends.
|
||||
# For other in-tree backends like cpu and cuda, they still use "TORCH_API" prefix with "visible" semantic.
|
||||
device_torch_api_key_word_mapping = {
|
||||
"XPU": "TORCH_XPU_API",
|
||||
}
|
||||
|
||||
return (
|
||||
device_torch_api_key_word_mapping.get(
|
||||
bankend_index.dispatch_key.name, "TORCH_API"
|
||||
)
|
||||
+ " "
|
||||
)
|
||||
|
||||
|
||||
@with_native_function_and_index
|
||||
def gen_unstructured(f: NativeFunction, backend_index: BackendIndex) -> str | None:
|
||||
sig = kernel_signature(f, backend_index)
|
||||
metadata = backend_index.get_kernel(f)
|
||||
if metadata is None:
|
||||
return None
|
||||
if "legacy::" in metadata.kernel:
|
||||
return None
|
||||
else:
|
||||
prefix = "static" if backend_index.external else "TORCH_API"
|
||||
return f"{prefix} {sig.decl(name=metadata.kernel)};"
|
||||
|
||||
|
||||
@with_native_function_and_index
|
||||
def gen_structured(g: NativeFunctionsGroup, backend_index: BackendIndex) -> list[str]:
|
||||
meta_name = meta.name(g)
|
||||
out_args = structured.impl_arguments(g)
|
||||
metadata = backend_index.get_kernel(g)
|
||||
if metadata is None:
|
||||
return []
|
||||
prefix = torch_api_key_word_prefix(backend_index)
|
||||
return [
|
||||
f"""\
|
||||
struct {prefix}structured_{metadata.kernel} : public at::meta::structured_{meta_name} {{
|
||||
void impl({", ".join(a.decl() for a in out_args)});
|
||||
}};
|
||||
"""
|
||||
]
|
||||
|
||||
|
||||
# Generates NativeFunctions.h, a list of forward declarations of all
|
||||
# actual kernel definitions we keep in aten/src/ATen/native/
|
||||
@with_native_function_and_index
|
||||
def compute_native_function_declaration(
|
||||
g: NativeFunctionsGroup | NativeFunction, backend_index: BackendIndex
|
||||
) -> list[str]:
|
||||
metadata = backend_index.get_kernel(g)
|
||||
if isinstance(g, NativeFunctionsGroup):
|
||||
if metadata is not None and metadata.structured:
|
||||
if backend_index.external:
|
||||
# Structured hasn't been tested with external backends yet.
|
||||
raise AssertionError(
|
||||
"Structured external backend functions are not implemented yet."
|
||||
)
|
||||
else:
|
||||
return gen_structured(g, backend_index)
|
||||
else:
|
||||
return list(
|
||||
mapMaybe(lambda f: gen_unstructured(f, backend_index), g.functions())
|
||||
)
|
||||
else:
|
||||
x = gen_unstructured(g, backend_index)
|
||||
return [] if x is None else [x]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,561 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torchgen.api.ufunc as ufunc
|
||||
from torchgen.api.translate import translate
|
||||
from torchgen.api.types import (
|
||||
BaseCType,
|
||||
Binding,
|
||||
CType,
|
||||
Expr,
|
||||
NamedCType,
|
||||
opmath_t,
|
||||
scalar_t,
|
||||
StructuredImplSignature,
|
||||
VectorizedCType,
|
||||
)
|
||||
from torchgen.context import with_native_function
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
DispatchKey,
|
||||
NativeFunctionsGroup,
|
||||
ScalarType,
|
||||
UfuncKey,
|
||||
)
|
||||
from torchgen.utils import OrderedSet
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from torchgen.api.ufunc import UfunctorBindings
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
#
|
||||
# CUDA STUFF
|
||||
#
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
# NB: not bothering to generate dispatch stub forward declaration in header,
|
||||
# we can just paste it wherever necessary
|
||||
|
||||
# TODO: use BackendIndex
|
||||
# dispatch_key: DispatchKey # only CPU/CUDA right now
|
||||
|
||||
|
||||
# Represents functors for implementing CUDA ufuncs.
|
||||
# Functors are templated by scalar_t because when USERS instantiate functors
|
||||
# they are templated. A functor looks something like this:
|
||||
#
|
||||
# template <typename scalar_t>
|
||||
# struct CUDAFunctorOnSelf_add {
|
||||
# using opmath_t = at::opmath_type<scalar_t>;
|
||||
# opmath_t other_;
|
||||
# opmath_t alpha_;
|
||||
# CUDAFunctorOnSelf_add(opmath_t other, opmath_t alpha)
|
||||
# : other_(other), alpha_(alpha) {}
|
||||
# __device__ scalar_t operator()(scalar_t self) {
|
||||
# return ufunc::add(static_cast<opmath_t>(self), other_, alpha_);
|
||||
# }
|
||||
# };
|
||||
#
|
||||
@dataclass(frozen=True)
|
||||
class UfunctorSignature:
|
||||
g: NativeFunctionsGroup
|
||||
scalar_tensor_idx: int | None
|
||||
name: str
|
||||
|
||||
def arguments(self) -> UfunctorBindings:
|
||||
return ufunc.ufunctor_arguments(
|
||||
self.g, scalar_tensor_idx=self.scalar_tensor_idx, scalar_t=scalar_t
|
||||
)
|
||||
|
||||
def fields(self) -> list[Binding]:
|
||||
# fields are renamed to have a trailing underscore, as is conventional
|
||||
return [b.rename(f"{b.name}_") for b in self.arguments().ctor]
|
||||
|
||||
def returns_type(self) -> CType:
|
||||
# TODO: don't hardcode; return type will be inferred based on tags on
|
||||
# the native function
|
||||
return BaseCType(scalar_t)
|
||||
|
||||
def decl_fields(self) -> str:
|
||||
return "\n".join(f"{f.type} {f.name};" for f in self.fields())
|
||||
|
||||
def inline_defn_ctor(self) -> str:
|
||||
args_str = ", ".join(a.decl() for a in self.arguments().ctor)
|
||||
# NB: hypothetically could do this with translate but the
|
||||
# transition here is very regular
|
||||
init_str = ", ".join(f"{a.name}_({a.name})" for a in self.arguments().ctor)
|
||||
return f"{self.name}({args_str}) : {init_str} {{}}"
|
||||
|
||||
def decl_apply(self) -> str:
|
||||
args_str = ", ".join(a.decl() for a in self.arguments().apply)
|
||||
return f"{self.returns_type().cpp_type()} operator()({args_str}) const"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UfuncSignature:
|
||||
g: NativeFunctionsGroup
|
||||
name: str
|
||||
compute_t: CType
|
||||
|
||||
def arguments(self) -> list[Binding]:
|
||||
return ufunc.ufunc_arguments(self.g, compute_t=self.compute_t)
|
||||
|
||||
def call(self, ctx: Sequence[Binding | Expr]) -> str:
|
||||
return f"{self.name}({', '.join(a.expr for a in translate(ctx, self.arguments()))})"
|
||||
|
||||
|
||||
# steps:
|
||||
# 1. take the functional signature
|
||||
# 2. use api.ufunc to convert it to template signature. this establishes
|
||||
# the type of the template function
|
||||
# 3. use api.ufunc (II) to generate a split struct / operator() signature.
|
||||
# this establish context in which we call the template signature
|
||||
#
|
||||
# StructuredImplSignature context
|
||||
# ~> functor constructor sig
|
||||
#
|
||||
# Functor constructor context
|
||||
# ~> functor fields sig
|
||||
#
|
||||
# Functor apply context (functor fields + functor apply sig)
|
||||
# ~> template sig
|
||||
#
|
||||
|
||||
|
||||
def eligible_for_binary_scalar_specialization(g: NativeFunctionsGroup) -> bool:
|
||||
num_tensors = sum(
|
||||
1 for a in g.functional.func.arguments.flat_non_out if a.type.is_tensor_like()
|
||||
)
|
||||
return num_tensors == 2
|
||||
|
||||
|
||||
def compute_ufunc_cuda_functors(
|
||||
g: NativeFunctionsGroup,
|
||||
) -> tuple[dict[ScalarType, dict[UfuncKey, UfunctorSignature]], str]:
|
||||
# First, build the functors.
|
||||
ufunctor_sigs: dict[ScalarType, dict[UfuncKey, UfunctorSignature]] = {}
|
||||
ufunctors: list[str] = []
|
||||
loops = g.out.ufunc_inner_loop
|
||||
scalar_tensor_idx_lookup = {
|
||||
UfuncKey.CUDAFunctorOnSelf: 1,
|
||||
UfuncKey.CUDAFunctorOnOther: 0,
|
||||
UfuncKey.CUDAFunctor: None,
|
||||
}
|
||||
if eligible_for_binary_scalar_specialization(g):
|
||||
keys = [
|
||||
UfuncKey.CUDAFunctorOnSelf,
|
||||
UfuncKey.CUDAFunctorOnOther,
|
||||
UfuncKey.CUDAFunctor,
|
||||
]
|
||||
else:
|
||||
keys = [UfuncKey.CUDAFunctor]
|
||||
for k in [UfuncKey.CUDAFunctorOnSelf, UfuncKey.CUDAFunctorOnOther]:
|
||||
if k in loops:
|
||||
raise AssertionError(f"cannot use {k} on non-binary function")
|
||||
for k in keys:
|
||||
# If the key was directly defined, skip functor codegen; we assume the
|
||||
# user already done it for us
|
||||
if k in loops:
|
||||
ufunctor_sig = UfunctorSignature(
|
||||
g, scalar_tensor_idx=scalar_tensor_idx_lookup[k], name=loops[k].name
|
||||
)
|
||||
for dtype in loops[k].supported_dtypes:
|
||||
ufunctor_sigs.setdefault(dtype, {})[k] = ufunctor_sig
|
||||
continue
|
||||
|
||||
# Note [ScalarOnly and Generic must match names for CUDA]
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# Otherwise, look in ANY of the generic entries. For simplicity of
|
||||
# codegen, both ScalarOnly and Generic are defined, the ufunc name
|
||||
# must match (if they didn't match, we'd have to generate distinct
|
||||
# functors per dtype, which is awful, so we're not going to do it unless
|
||||
# someone really forces us to)
|
||||
ufunc_name = None
|
||||
supported_dtypes: OrderedSet[ScalarType] = OrderedSet()
|
||||
for lk in [UfuncKey.ScalarOnly, UfuncKey.Generic]:
|
||||
if lk not in loops:
|
||||
continue
|
||||
if ufunc_name is None:
|
||||
ufunc_name = loops[lk].name
|
||||
else:
|
||||
# See Note [ScalarOnly and Generic must match names for CUDA]
|
||||
if ufunc_name != loops[lk].name:
|
||||
raise AssertionError(
|
||||
"ScalarOnly and Generic must have same ufunc name"
|
||||
)
|
||||
supported_dtypes |= loops[lk].supported_dtypes
|
||||
if ufunc_name is None:
|
||||
raise AssertionError("ufunc_name must be non-None")
|
||||
|
||||
name = f"{k}_{ufunc_name}"
|
||||
ufunctor_sig = UfunctorSignature(
|
||||
g, scalar_tensor_idx=scalar_tensor_idx_lookup[k], name=name
|
||||
)
|
||||
for dtype in supported_dtypes:
|
||||
ufunctor_sigs.setdefault(dtype, {})[k] = ufunctor_sig
|
||||
|
||||
ufunc_sig = UfuncSignature(
|
||||
g, name=f"ufunc::{ufunc_name}", compute_t=BaseCType(opmath_t)
|
||||
)
|
||||
apply_ctx = ufunctor_sig.fields() + ufunctor_sig.arguments().apply
|
||||
ufunctors.append(
|
||||
f"""
|
||||
template <typename scalar_t>
|
||||
struct {ufunctor_sig.name} {{
|
||||
using opmath_t = at::opmath_type<scalar_t>;
|
||||
{ufunctor_sig.decl_fields()}
|
||||
{ufunctor_sig.inline_defn_ctor()}
|
||||
__device__ {ufunctor_sig.decl_apply()} {{
|
||||
return {ufunc_sig.call(apply_ctx)};
|
||||
}}
|
||||
}};
|
||||
"""
|
||||
)
|
||||
|
||||
return ufunctor_sigs, "\n".join(ufunctors)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BinaryScalarSpecializationConfig:
|
||||
scalar_idx: int
|
||||
ctor_tensor: str
|
||||
ufunc_key: UfuncKey
|
||||
|
||||
|
||||
BinaryScalarSpecializationConfigs = [
|
||||
BinaryScalarSpecializationConfig(
|
||||
scalar_idx=0,
|
||||
ctor_tensor="self",
|
||||
ufunc_key=UfuncKey.CUDAFunctorOnOther,
|
||||
),
|
||||
BinaryScalarSpecializationConfig(
|
||||
scalar_idx=1,
|
||||
ctor_tensor="other",
|
||||
ufunc_key=UfuncKey.CUDAFunctorOnSelf,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def compute_ufunc_cuda_dtype_body(
|
||||
g: NativeFunctionsGroup,
|
||||
dtype: ScalarType,
|
||||
inner_loops: dict[UfuncKey, UfunctorSignature],
|
||||
parent_ctx: Sequence[Binding],
|
||||
) -> str:
|
||||
body = "using opmath_t = at::opmath_type<scalar_t>;"
|
||||
body += "if (false) {}\n" # for ease of codegen
|
||||
for config in BinaryScalarSpecializationConfigs:
|
||||
if config.ufunc_key not in inner_loops:
|
||||
continue
|
||||
ufunctor_sig = inner_loops[config.ufunc_key]
|
||||
scalar_idx = config.scalar_idx + 1
|
||||
# Make a copy and at the same time widen the type (not permissible
|
||||
# without copy; we don't want to mutate the input argument anyway)
|
||||
ctx: list[Expr | Binding] = list(parent_ctx)
|
||||
ctx.append(
|
||||
Expr(
|
||||
expr=f"iter.scalar_value<opmath_t>({scalar_idx})",
|
||||
type=NamedCType(config.ctor_tensor, BaseCType(opmath_t)),
|
||||
)
|
||||
)
|
||||
ufunctor_ctor_exprs_str = ", ".join(
|
||||
a.expr for a in translate(ctx, ufunctor_sig.arguments().ctor)
|
||||
)
|
||||
|
||||
# NB: ufunctor must be allocated before iter.remove_operand is called,
|
||||
# as it relies on iter
|
||||
body += f"""\
|
||||
else if (iter.is_cpu_scalar({scalar_idx})) {{
|
||||
{ufunctor_sig.name}<scalar_t> ufunctor({ufunctor_ctor_exprs_str});
|
||||
iter.remove_operand({scalar_idx});
|
||||
gpu_kernel(iter, ufunctor);
|
||||
}}"""
|
||||
|
||||
ufunctor_sig = inner_loops[UfuncKey.CUDAFunctor]
|
||||
ufunctor_ctor_exprs_str = ", ".join(
|
||||
a.expr for a in translate(parent_ctx, ufunctor_sig.arguments().ctor)
|
||||
)
|
||||
body += f"""
|
||||
else {{
|
||||
gpu_kernel(iter, {ufunctor_sig.name}<scalar_t>({ufunctor_ctor_exprs_str}));
|
||||
}}
|
||||
"""
|
||||
return body
|
||||
|
||||
|
||||
@with_native_function
|
||||
def compute_ufunc_cuda(g: NativeFunctionsGroup) -> str:
|
||||
# First, build the functors, indexing them by dtype
|
||||
ufunctor_sigs, ufunctors = compute_ufunc_cuda_functors(g)
|
||||
|
||||
# Next, build the conditionals
|
||||
sig = StructuredImplSignature(g, ufunc.kernel_name(g, DispatchKey.CUDA))
|
||||
dtype_cases = []
|
||||
for dtype, inner_ufunc_sigs in ufunctor_sigs.items():
|
||||
dtype_cases.append(
|
||||
f"""
|
||||
AT_DISPATCH_CASE(at::ScalarType::{dtype},
|
||||
[&]() {{
|
||||
{compute_ufunc_cuda_dtype_body(g, dtype, inner_ufunc_sigs, sig.arguments())}
|
||||
}}
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
dtype_cases_str = "\n".join(dtype_cases)
|
||||
|
||||
stub_sig = StubSignature(g)
|
||||
|
||||
return f"""
|
||||
{ufunctors}
|
||||
|
||||
{stub_sig.type_defn()};
|
||||
{stub_sig.dispatch_decl()}
|
||||
|
||||
{stub_sig.kernel_defn()} {{
|
||||
AT_DISPATCH_SWITCH(iter.common_dtype(), "{sig.name}",
|
||||
{dtype_cases_str}
|
||||
);
|
||||
}}
|
||||
REGISTER_DISPATCH({stub_sig.name}, &{stub_sig.kernel_name})
|
||||
|
||||
{sig.defn()} {{
|
||||
{stub_sig.direct_call(sig.arguments())};
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
#
|
||||
# CPU STUFF
|
||||
#
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StubSignature:
|
||||
g: NativeFunctionsGroup
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"{str(self.g.functional.func.name.name)}_stub"
|
||||
|
||||
@property
|
||||
def kernel_name(self) -> str:
|
||||
return f"{str(self.g.functional.func.name.name)}_kernel"
|
||||
|
||||
@property
|
||||
def type_name(self) -> str:
|
||||
return f"{str(self.g.functional.func.name.name)}_fn"
|
||||
|
||||
def arguments(self) -> list[Binding]:
|
||||
return ufunc.stub_arguments(self.g)
|
||||
|
||||
def type(self) -> str:
|
||||
cpp_args = self.arguments()
|
||||
return f"void(*)(TensorIteratorBase&, {', '.join(a.type for a in cpp_args)})"
|
||||
|
||||
def dispatch_decl(self) -> str:
|
||||
return f"DECLARE_DISPATCH({self.type_name}, {self.name})"
|
||||
|
||||
def dispatch_defn(self) -> str:
|
||||
return f"DEFINE_DISPATCH({self.name})"
|
||||
|
||||
def kernel_defn(self) -> str:
|
||||
return f"void {self.kernel_name}(TensorIteratorBase& iter, {', '.join(a.defn() for a in self.arguments())})"
|
||||
|
||||
def type_defn(self) -> str:
|
||||
return f"using {self.type_name} = {self.type()}"
|
||||
|
||||
# must be called from context where this is TensorIteratorBase*
|
||||
def call(self, ctx: Sequence[Binding]) -> str:
|
||||
return f"{self.name}(device_type(), *this, {', '.join(a.expr for a in translate(ctx, self.arguments()))})"
|
||||
|
||||
# used in CUDA to skip the unnecessary dynamic dispatch
|
||||
def direct_call(self, ctx: Sequence[Binding]) -> str:
|
||||
return f"{self.kernel_name}(*this, {', '.join(a.expr for a in translate(ctx, self.arguments()))})"
|
||||
|
||||
|
||||
@with_native_function
|
||||
def compute_ufunc_cpu(g: NativeFunctionsGroup) -> str:
|
||||
stub_sig = StubSignature(g)
|
||||
sig = StructuredImplSignature(g, ufunc.kernel_name(g, DispatchKey.CPU))
|
||||
|
||||
return f"""
|
||||
{stub_sig.type_defn()};
|
||||
{stub_sig.dispatch_decl()}
|
||||
{stub_sig.dispatch_defn()};
|
||||
|
||||
{sig.defn()} {{
|
||||
{stub_sig.call(sig.arguments())};
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
def compute_ufunc_cpu_dtype_body(
|
||||
g: NativeFunctionsGroup,
|
||||
dtype: ScalarType,
|
||||
inner_loops: dict[UfuncKey, UfuncSignature],
|
||||
parent_ctx: Sequence[Binding],
|
||||
) -> str:
|
||||
if UfuncKey.CPUScalar not in inner_loops:
|
||||
raise AssertionError(f"{dtype}, {inner_loops.keys()}")
|
||||
if not inner_loops.keys() <= {UfuncKey.CPUScalar, UfuncKey.CPUVector}:
|
||||
raise AssertionError(
|
||||
f"inner_loops keys must be subset of CPUScalar/CPUVector, got {inner_loops.keys()}"
|
||||
)
|
||||
scalar_loop = inner_loops[UfuncKey.CPUScalar]
|
||||
vec_loop = None
|
||||
if UfuncKey.CPUVector in inner_loops:
|
||||
vec_loop = inner_loops[UfuncKey.CPUVector]
|
||||
|
||||
# NB: We DON'T use translate here, because translate is
|
||||
# incapable of CSE'ing the scalar accesses in case it is also
|
||||
# used by Vectorized; also, the unpacking here is very simple
|
||||
# and only affects Scalar; everything else is implicitly captured
|
||||
# by the lambda
|
||||
|
||||
# Setup scalar in scope
|
||||
body = []
|
||||
ctx = []
|
||||
for b in parent_ctx:
|
||||
if isinstance(b.argument, Argument) and b.argument.type != BaseType(
|
||||
BaseTy.Scalar
|
||||
):
|
||||
continue
|
||||
body.append(f"auto _s_{b.name} = {b.name}.to<scalar_t>();")
|
||||
ctx.append(Expr(f"_s_{b.name}", NamedCType(b.nctype.name, BaseCType(scalar_t))))
|
||||
if vec_loop is not None:
|
||||
for b in parent_ctx:
|
||||
if isinstance(b.argument, Argument) and b.argument.type != BaseType(
|
||||
BaseTy.Scalar
|
||||
):
|
||||
continue
|
||||
body.append(
|
||||
f"auto _v_{b.name} = at::vec::Vectorized<scalar_t>(_s_{b.name});"
|
||||
)
|
||||
ctx.append(
|
||||
Expr(
|
||||
f"_v_{b.name}",
|
||||
NamedCType(b.nctype.name, VectorizedCType(BaseCType(scalar_t))),
|
||||
)
|
||||
)
|
||||
|
||||
# Setup lambda signature
|
||||
# NB: simplified version of ufunctor_arguments
|
||||
scalar_bindings = []
|
||||
vec_bindings = []
|
||||
for a in g.functional.func.arguments.flat_non_out:
|
||||
if not a.type.is_tensor_like():
|
||||
continue
|
||||
if a.type != BaseType(BaseTy.Tensor):
|
||||
raise AssertionError(f"Expected Tensor type, got {a.type}")
|
||||
scalar_bindings.append(
|
||||
Binding(
|
||||
name=a.name,
|
||||
nctype=NamedCType(a.name, BaseCType(scalar_t)),
|
||||
argument=a,
|
||||
)
|
||||
)
|
||||
if vec_loop is not None:
|
||||
vec_bindings.append(
|
||||
Binding(
|
||||
name=a.name,
|
||||
nctype=NamedCType(a.name, VectorizedCType(BaseCType(scalar_t))),
|
||||
argument=a,
|
||||
)
|
||||
)
|
||||
|
||||
def with_ctx(b: Sequence[Binding]) -> list[Expr | Binding]:
|
||||
r: list[Expr | Binding] = []
|
||||
r.extend(ctx)
|
||||
r.extend(b)
|
||||
return r
|
||||
|
||||
body_str = "\n".join(body)
|
||||
if vec_loop is not None:
|
||||
return f"""
|
||||
{body_str}
|
||||
cpu_kernel_vec(iter,
|
||||
[=]({", ".join(b.decl() for b in scalar_bindings)}) {{ return {scalar_loop.call(with_ctx(scalar_bindings))}; }},
|
||||
[=]({", ".join(b.decl() for b in vec_bindings)}) {{ return {vec_loop.call(with_ctx(vec_bindings))}; }}
|
||||
);
|
||||
"""
|
||||
else:
|
||||
return f"""
|
||||
{body_str}
|
||||
cpu_kernel(iter,
|
||||
[=]({", ".join(b.decl() for b in scalar_bindings)}) {{ return {scalar_loop.call(with_ctx(scalar_bindings))}; }}
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
@with_native_function
|
||||
def compute_ufunc_cpu_kernel(g: NativeFunctionsGroup) -> str:
|
||||
stub_sig = StubSignature(g)
|
||||
|
||||
# Reindex the ufunc by dtypes; processing generic/scalaronly as well
|
||||
loops = g.out.ufunc_inner_loop
|
||||
ufunc_sigs: dict[ScalarType, dict[UfuncKey, UfuncSignature]] = {}
|
||||
for k in [UfuncKey.CPUScalar, UfuncKey.CPUVector]:
|
||||
lks = []
|
||||
# ORDER MATTERS: this specifies overriding precedence
|
||||
if k in loops: # should happen rarely
|
||||
lks.append(k)
|
||||
if UfuncKey.ScalarOnly in loops and k is UfuncKey.CPUScalar:
|
||||
lks.append(UfuncKey.ScalarOnly)
|
||||
if UfuncKey.Generic in loops:
|
||||
lks.append(UfuncKey.Generic)
|
||||
# TODO: don't hardcode ufunc:: namespace here, should be centralized smh
|
||||
for lk in lks:
|
||||
for dtype in loops[lk].supported_dtypes:
|
||||
compute_t: CType
|
||||
if k is UfuncKey.CPUScalar:
|
||||
compute_t = BaseCType(scalar_t)
|
||||
elif k is UfuncKey.CPUVector:
|
||||
compute_t = VectorizedCType(BaseCType(scalar_t))
|
||||
else:
|
||||
raise AssertionError
|
||||
inner_ufunc_sigs = ufunc_sigs.setdefault(dtype, {})
|
||||
if k not in inner_ufunc_sigs:
|
||||
inner_ufunc_sigs[k] = UfuncSignature(
|
||||
g, name=f"ufunc::{loops[lk].name}", compute_t=compute_t
|
||||
)
|
||||
|
||||
# Build the conditionals
|
||||
dtype_cases = []
|
||||
for dtype, inner_ufunc_sigs in ufunc_sigs.items():
|
||||
dtype_cases.append(
|
||||
f"""
|
||||
AT_DISPATCH_CASE(at::ScalarType::{dtype},
|
||||
[&]() {{
|
||||
{compute_ufunc_cpu_dtype_body(g, dtype, inner_ufunc_sigs, stub_sig.arguments())}
|
||||
}}
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
dtype_cases_str = "\n".join(dtype_cases)
|
||||
return f"""
|
||||
namespace {{
|
||||
|
||||
{stub_sig.kernel_defn()} {{
|
||||
AT_DISPATCH_SWITCH(iter.common_dtype(), "{stub_sig.name}",
|
||||
{dtype_cases_str}
|
||||
);
|
||||
}}
|
||||
|
||||
}} // anonymous namespace
|
||||
|
||||
{stub_sig.type_defn()};
|
||||
{stub_sig.dispatch_decl()}
|
||||
REGISTER_DISPATCH({stub_sig.name}, &{stub_sig.kernel_name})
|
||||
"""
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,783 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import difflib
|
||||
import os
|
||||
import textwrap
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from torchgen.aoti.fallback_ops import aten_shimified_ops, inductor_fallback_ops
|
||||
from torchgen.api.types import DispatcherSignature
|
||||
from torchgen.api.types.signatures import CppSignature, CppSignatureGroup
|
||||
from torchgen.context import method_with_native_function
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BackendIndex,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
DispatchKey,
|
||||
FunctionSchema,
|
||||
is_cuda_dispatch_key,
|
||||
ListType,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
OperatorName,
|
||||
OptionalType,
|
||||
Type,
|
||||
Variant,
|
||||
)
|
||||
from torchgen.utils import FileManager, mapMaybe
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
base_type_to_c_type = {
|
||||
BaseTy.Tensor: "AtenTensorHandle",
|
||||
BaseTy.bool: "int32_t", # Use int to pass bool
|
||||
BaseTy.int: "int64_t",
|
||||
BaseTy.SymInt: "int64_t", # Inductor-generated code won't see a SymInt
|
||||
BaseTy.Scalar: "double", # Use double to pass both integer and floating point
|
||||
BaseTy.float: "double", # TODO: how about other floating point types?
|
||||
BaseTy.str: "const char*",
|
||||
BaseTy.DeviceIndex: "int32_t",
|
||||
BaseTy.Layout: "int32_t", # Represent enum as int
|
||||
BaseTy.MemoryFormat: "int32_t", # Represent enum as int
|
||||
BaseTy.ScalarType: "int32_t", # Represent enum as int
|
||||
BaseTy.Generator: "AtenGeneratorHandle",
|
||||
}
|
||||
|
||||
base_type_to_aten_type = {
|
||||
BaseTy.Tensor: "at::Tensor",
|
||||
BaseTy.bool: "bool",
|
||||
BaseTy.int: "int64_t",
|
||||
BaseTy.SymInt: "c10::SymInt",
|
||||
BaseTy.Scalar: "c10::Scalar",
|
||||
BaseTy.float: "double",
|
||||
BaseTy.str: "::std::string_view",
|
||||
BaseTy.DeviceIndex: "c10::DeviceIndex",
|
||||
BaseTy.Layout: "c10::Layout",
|
||||
BaseTy.MemoryFormat: "c10::MemoryFormat",
|
||||
BaseTy.ScalarType: "c10::ScalarType",
|
||||
BaseTy.Generator: "at::Generator",
|
||||
}
|
||||
|
||||
base_type_to_callsite_expr = {
|
||||
BaseTy.Tensor: "resolve_tensor_dispatch_flags",
|
||||
BaseTy.bool: "",
|
||||
BaseTy.int: "",
|
||||
BaseTy.SymInt: "",
|
||||
BaseTy.Scalar: "",
|
||||
BaseTy.float: "",
|
||||
BaseTy.str: "",
|
||||
BaseTy.DeviceIndex: "static_cast<c10::DeviceIndex>",
|
||||
BaseTy.Layout: "static_cast<c10::Layout>",
|
||||
BaseTy.MemoryFormat: "static_cast<c10::MemoryFormat>",
|
||||
BaseTy.ScalarType: "static_cast<c10::ScalarType>",
|
||||
BaseTy.Generator: "*generator_handle_to_generator_pointer",
|
||||
}
|
||||
|
||||
|
||||
# convert args to C types, names in declarations, and expressions in function bodies
|
||||
def convert_arg_type_and_name(
|
||||
typ: Type,
|
||||
name: str,
|
||||
is_write: bool = False,
|
||||
) -> tuple[list[str], list[str], list[str], list[str]]:
|
||||
if isinstance(typ, BaseType):
|
||||
if typ.name in base_type_to_c_type:
|
||||
if typ.name == BaseTy.Tensor and is_write:
|
||||
# For output tensors, our normal call to resolve_tensor_dispatch_flags
|
||||
# results in an rvalue tensor, which can't be passed to at::Tensor&.
|
||||
# Override this case specifically.
|
||||
callsite_expr = [f"*tensor_handle_to_tensor_pointer({name})"]
|
||||
else:
|
||||
callsite_expr = [
|
||||
f"{base_type_to_callsite_expr[typ.name]}({name})"
|
||||
if base_type_to_callsite_expr[typ.name]
|
||||
else name
|
||||
]
|
||||
|
||||
return (
|
||||
[base_type_to_c_type[typ.name]],
|
||||
[name],
|
||||
[base_type_to_aten_type[typ.name]],
|
||||
callsite_expr,
|
||||
)
|
||||
elif typ.name == BaseTy.Device:
|
||||
return (
|
||||
["int32_t", "int32_t"],
|
||||
[name, name + "_index_"],
|
||||
["c10::Device"],
|
||||
[
|
||||
f"c10::Device(static_cast<c10::DeviceType>({name}), static_cast<c10::DeviceIndex>({name}_index_))"
|
||||
],
|
||||
)
|
||||
else:
|
||||
# TODO: BaseTy.Dimname, etc.
|
||||
raise NotImplementedError(f"TODO: add support for arg type {repr(typ)}")
|
||||
elif isinstance(typ, OptionalType):
|
||||
c_types, names, aten_types, callsite_exprs = convert_arg_type_and_name(
|
||||
typ.elem, name
|
||||
)
|
||||
j = 0 # index for names
|
||||
new_aten_types = []
|
||||
new_callsite_exprs = []
|
||||
for aten_type in aten_types:
|
||||
# Use pointer to denote optional type
|
||||
c_types[j] = c_types[j] + "*"
|
||||
if aten_type.startswith("c10::ArrayRef<"):
|
||||
# ArrayRef is passed as pointer + size, but no need to add "*" to the size argument
|
||||
new_aten_types.append(f"::std::optional<{aten_type}>")
|
||||
base_type = aten_type[len("c10::ArrayRef<") : -1]
|
||||
new_callsite_exprs.append(
|
||||
f"pointer_to_optional_list<{base_type}>({names[j]}, {names[j + 1]})"
|
||||
)
|
||||
j += 2
|
||||
elif aten_type == "c10::Device":
|
||||
# Device is passed as device_type + device_index
|
||||
new_aten_types.append("::std::optional<c10::Device>")
|
||||
new_callsite_exprs.append(
|
||||
f"pointer_to_optional_device({names[j]}, {names[j + 1]})"
|
||||
)
|
||||
j += 2
|
||||
elif aten_type == "at::Tensor":
|
||||
new_aten_types.append(f"::std::optional<{aten_type}>")
|
||||
new_callsite_exprs.append(f"resolve_tensor_dispatch_flags({names[j]})")
|
||||
j += 1
|
||||
else:
|
||||
new_aten_types.append(f"::std::optional<{aten_type}>")
|
||||
new_callsite_exprs.append(
|
||||
f"pointer_to_optional<{aten_type}>({names[j]})"
|
||||
)
|
||||
j += 1
|
||||
|
||||
return (
|
||||
c_types,
|
||||
names,
|
||||
new_aten_types,
|
||||
new_callsite_exprs,
|
||||
)
|
||||
elif isinstance(typ, ListType):
|
||||
# Need to explicitly pass the list as pointer + length
|
||||
c_types, names, aten_types, _ = convert_arg_type_and_name(typ.elem, name)
|
||||
if len(c_types) != 1:
|
||||
raise AssertionError(f"ListType with unsupported element type {repr(typ)}")
|
||||
|
||||
# The list content should never be modified
|
||||
c_types[0] = f"const {c_types[0]}*"
|
||||
c_types.append("int64_t")
|
||||
name = names[0]
|
||||
names.append(name + "_len_")
|
||||
|
||||
atype = aten_types[0]
|
||||
callsite_exprs = []
|
||||
if atype == "bool":
|
||||
# no converter from std::vector<bool> to c10::ArrayRef<bool>
|
||||
# construct std::array<bool, N> instead
|
||||
if typ.size is None:
|
||||
raise AssertionError("bool ListType must have a size")
|
||||
callsite_exprs.append(f"pointer_to_list<{typ.size}>({name})")
|
||||
elif atype == "at::Tensor" and not is_write:
|
||||
callsite_exprs.append(
|
||||
f"resolve_tensor_list_dispatch_flags({name}, {name}_len_)"
|
||||
)
|
||||
elif atype == "::std::optional<at::Tensor>":
|
||||
# convert from std::vector<::std::optional<at::Tensor>> to c10::List<::std::optional<at::Tensor>>
|
||||
callsite_exprs.append(
|
||||
f"c10::List<{atype}>(c10::ArrayRef<{atype}>(resolve_tensor_list_dispatch_flags({name}, {name}_len_)))"
|
||||
)
|
||||
else:
|
||||
callsite_exprs.append(f"pointer_to_list<{atype}>({name}, {name}_len_)")
|
||||
|
||||
aten_types = [f"c10::ArrayRef<{t}>" for t in aten_types]
|
||||
return (
|
||||
c_types,
|
||||
names,
|
||||
aten_types,
|
||||
callsite_exprs,
|
||||
)
|
||||
raise NotImplementedError(f"Argument type {repr(typ)} not supported!")
|
||||
|
||||
|
||||
def zip_type_and_name(types: list[str], names: list[str]) -> list[str]:
|
||||
return [typ + " " + name for typ, name in zip(types, names)]
|
||||
|
||||
|
||||
# Generate argument declarations and callsite expressions
|
||||
def gen_arguments(
|
||||
flat_arguments: Sequence[Argument], skipped_args: set[str]
|
||||
) -> tuple[list[str], list[str]]:
|
||||
types: list[str] = []
|
||||
new_names: list[str] = []
|
||||
callsite_exprs: list[str] = []
|
||||
for arg in flat_arguments:
|
||||
if arg.name in skipped_args:
|
||||
# Pass the arg's schema default when available (e.g. "false" for
|
||||
# a bool arg with default=False), so non-optional args with defaults
|
||||
# can be versioned too. Fall back to std::nullopt for optional args
|
||||
# with no default (matches historical behavior).
|
||||
if arg.default is not None:
|
||||
from torchgen.api.cpp import default_expr
|
||||
|
||||
callsite_exprs.append(default_expr(arg.default, arg.type, symint=False))
|
||||
else:
|
||||
callsite_exprs.append("std::nullopt")
|
||||
continue
|
||||
new_types, names, _, new_callsite_exprs = convert_arg_type_and_name(
|
||||
arg.type, arg.name, arg.is_write
|
||||
)
|
||||
types.extend(new_types)
|
||||
new_names.extend(names)
|
||||
callsite_exprs.extend(new_callsite_exprs)
|
||||
return zip_type_and_name(types, new_names), callsite_exprs
|
||||
|
||||
|
||||
# Return values are passed out as pointer arguments because all the C shim functions
|
||||
# are expected to return AOTITorchError.
|
||||
# Generate returns as declarations and callsite expressions
|
||||
def gen_returns(schema: FunctionSchema) -> tuple[list[str], list[str]]:
|
||||
types = []
|
||||
names = []
|
||||
for idx, ret in enumerate(schema.returns):
|
||||
names.append(f"ret{idx}")
|
||||
if isinstance(ret.type, BaseType) and ret.type.name in base_type_to_c_type:
|
||||
types.append(base_type_to_c_type[ret.type.name] + "*")
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"TODO: add support for return type {repr(ret.type)}"
|
||||
)
|
||||
|
||||
def convert_return(typ: BaseType, val: str) -> str:
|
||||
if typ.name == BaseTy.Tensor:
|
||||
return f"new_tensor_handle(std::move({val}))"
|
||||
elif typ.name == BaseTy.SymInt:
|
||||
return f"{val}.expect_int()"
|
||||
elif typ.name == BaseTy.Scalar:
|
||||
return f"{val}.toDouble()"
|
||||
else:
|
||||
return val
|
||||
|
||||
ret_pointer_can_be_null = False
|
||||
unambiguous_name = schema.name.unambiguous_name()
|
||||
for name in (
|
||||
"_functional_sym_constrain_range",
|
||||
"_scaled_dot_product_cudnn_attention",
|
||||
"_scaled_dot_product_efficient_attention_backward",
|
||||
"_scaled_dot_product_efficient_attention",
|
||||
"_scaled_dot_product_flash_attention",
|
||||
"_scaled_dot_product_fused_attention_overrideable",
|
||||
"_thhn_fused_lstm_cell_backward_impl",
|
||||
"convolution_backward",
|
||||
"grid_sampler_2d_backward",
|
||||
"grid_sampler_3d_backward",
|
||||
"linear_backward",
|
||||
):
|
||||
if name in unambiguous_name:
|
||||
ret_pointer_can_be_null = True
|
||||
break
|
||||
|
||||
callsite_exprs: list[str] = []
|
||||
for idx, ret in enumerate(schema.returns):
|
||||
tmp = "tmp_result" if len(names) == 1 else f"std::get<{idx}>(tmp_result)"
|
||||
if not isinstance(ret.type, BaseType):
|
||||
raise AssertionError(f"Expected BaseType for return, got {type(ret.type)}")
|
||||
rval = convert_return(ret.type, tmp)
|
||||
if ret_pointer_can_be_null:
|
||||
callsite_exprs.append(f"if ({names[idx]}) {{ *{names[idx]} = {rval}; }}")
|
||||
else:
|
||||
callsite_exprs.append(f"*{names[idx]} = {rval};")
|
||||
|
||||
return zip_type_and_name(types, names), callsite_exprs
|
||||
|
||||
|
||||
# gen.py generates header first and then src, so caching the result here to avoid duplicate work
|
||||
declaration_definition_cache: dict[tuple[str, str, str], tuple[str, str]] = {}
|
||||
|
||||
|
||||
def gen_declaration_and_definition(
|
||||
schema: FunctionSchema,
|
||||
device: str,
|
||||
backend_call: str,
|
||||
version_info: dict[str, list[str]],
|
||||
) -> tuple[str, str]:
|
||||
base_name = schema.name.unambiguous_name()
|
||||
|
||||
global declaration_definition_cache
|
||||
if (base_name, device, backend_call) in declaration_definition_cache:
|
||||
return declaration_definition_cache[(base_name, device, backend_call)]
|
||||
|
||||
# Check the validity of version_info. The format should look like
|
||||
# {"v2" : ["new_arg1"], "v3": ["new_arg2, new_arg3"]}.
|
||||
indexed_version_info: dict[int, list[str]] = {1: []}
|
||||
for ver_str, new_args in sorted(version_info.items()):
|
||||
if not ver_str.startswith("v"):
|
||||
raise AssertionError(
|
||||
f"Version number for {base_name} is {ver_str}, not starting with 'v'"
|
||||
)
|
||||
try:
|
||||
ver_id = int(ver_str[1:])
|
||||
except ValueError as e:
|
||||
raise AssertionError(
|
||||
f"Version number for {base_name} is {ver_str}, not a valid integer after 'v'"
|
||||
) from e
|
||||
if ver_id in indexed_version_info:
|
||||
raise AssertionError(f"{ver_str} for {base_name} has already been defined")
|
||||
indexed_version_info[ver_id] = new_args
|
||||
|
||||
declarations: list[str] = []
|
||||
definitions: list[str] = []
|
||||
skipped_args: set[str] = set()
|
||||
|
||||
for ver_id, new_args in sorted(indexed_version_info.items(), reverse=True):
|
||||
# Iterate in the reverse order, so the latest version of an op will get generated first
|
||||
# with all the arguments included, while a set of to-be-trimmed args is carried down
|
||||
# to generate earlier version of the op.
|
||||
func_name = base_name if ver_id == 1 else f"{base_name}_v{ver_id}"
|
||||
if schema.is_out_fn():
|
||||
# out_variant has out arguments in the front, and it's ok to ignore return values
|
||||
# because C shim functions only return AOTITorchError
|
||||
args, callsite_exprs = gen_arguments(
|
||||
[*schema.arguments.out, *schema.arguments.flat_non_out], skipped_args
|
||||
)
|
||||
ret_assignments: list[str] = []
|
||||
else:
|
||||
args, callsite_exprs = gen_arguments(
|
||||
schema.arguments.flat_all, skipped_args
|
||||
)
|
||||
# ignore return values for inplace ops
|
||||
ret_declarations, ret_assignments = (
|
||||
([], []) if schema.name.name.inplace else gen_returns(schema)
|
||||
)
|
||||
args.extend(ret_declarations)
|
||||
|
||||
declaration = textwrap.dedent(
|
||||
f"AOTITorchError aoti_torch_{device}_{func_name}({', '.join(args)})"
|
||||
)
|
||||
|
||||
tmp_result = "auto tmp_result = " if ret_assignments else ""
|
||||
indent = "\t\t"
|
||||
ret_assignments_str = (
|
||||
"\n".join(indent + r for r in ret_assignments) if ret_assignments else ""
|
||||
)
|
||||
definition = (
|
||||
textwrap.dedent(f"""
|
||||
{declaration} {{
|
||||
AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({{
|
||||
{tmp_result}{backend_call}(
|
||||
{", ".join(callsite_exprs)}
|
||||
);
|
||||
""")
|
||||
+ ret_assignments_str
|
||||
+ textwrap.dedent("""
|
||||
});
|
||||
}
|
||||
""")
|
||||
)
|
||||
skipped_args.update(new_args)
|
||||
declarations.append(f"AOTI_TORCH_EXPORT {declaration};")
|
||||
definitions.append(definition)
|
||||
|
||||
declaration_definition_cache[(base_name, device, backend_call)] = (
|
||||
"\n".join(declarations),
|
||||
"\n".join(definitions),
|
||||
)
|
||||
return declaration_definition_cache[(base_name, device, backend_call)]
|
||||
|
||||
|
||||
def gen_static_dispatch_backend_call_signature(
|
||||
sig: CppSignature | DispatcherSignature,
|
||||
f: NativeFunction,
|
||||
) -> CppSignature:
|
||||
sig = DispatcherSignature.from_schema(f.func)
|
||||
cpp_sigs = CppSignatureGroup.from_native_function(
|
||||
f, method=False, fallback_binding=False
|
||||
)
|
||||
if sig.symint and f.func.has_symint():
|
||||
cpp_sig = cpp_sigs.symint_signature
|
||||
else:
|
||||
cpp_sig = cpp_sigs.signature
|
||||
if cpp_sig is None:
|
||||
raise AssertionError(f"No cpp signature found for {f.func.name}")
|
||||
return cpp_sig
|
||||
|
||||
|
||||
def gen_static_dispatch_backend_call(
|
||||
f: NativeFunction,
|
||||
backend_index: BackendIndex | None = None,
|
||||
) -> str:
|
||||
sig = DispatcherSignature.from_schema(f.func)
|
||||
cpp_sig = gen_static_dispatch_backend_call_signature(sig, f)
|
||||
|
||||
if backend_index is None:
|
||||
# Check if this is a symint function and if the function only has method variants
|
||||
if sig.symint and f.func.has_symint():
|
||||
has_function_variant = Variant.function in f.variants
|
||||
|
||||
if not has_function_variant:
|
||||
# Functions with both function and method variants can use the at::{*}_symint version
|
||||
# (e.g., narrow -> at::narrow_symint), BUT
|
||||
# Method-only functions with symint parameters should use at::symint:: namespace
|
||||
# Remove the _symint suffix since at::symint:: namespace uses the base name
|
||||
# (e.g., new_empty -> at::symint::new_empty<c10::SymInt>)
|
||||
base_name = cpp_sig.name()
|
||||
base_name = base_name.removesuffix("_symint") # Remove "_symint" suffix
|
||||
return f"at::symint::{base_name}<c10::SymInt>"
|
||||
|
||||
return f"at::{cpp_sig.name()}"
|
||||
else:
|
||||
return f"at::{backend_index.dispatch_key.lower()}::{cpp_sig.name()}"
|
||||
|
||||
|
||||
def get_backend_index_for_aoti(
|
||||
func: NativeFunction,
|
||||
func_group_mapping: dict[OperatorName, NativeFunctionsGroup],
|
||||
dispatch_key: DispatchKey | None,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
extend_aoti_c_shim: bool,
|
||||
) -> BackendIndex | None:
|
||||
backend_index = None
|
||||
|
||||
if dispatch_key is None:
|
||||
return backend_index
|
||||
|
||||
if backend_indices[dispatch_key].has_kernel(func) or (
|
||||
func.structured_delegate is not None
|
||||
and func.structured_delegate in func_group_mapping
|
||||
and backend_indices[dispatch_key].has_kernel(
|
||||
func_group_mapping[func.structured_delegate]
|
||||
)
|
||||
):
|
||||
backend_index = backend_indices[dispatch_key]
|
||||
else:
|
||||
# for the extend out-of-tree kernels, we don't need to
|
||||
# duplicatly create C shim wrappers for other dispatch keys
|
||||
if extend_aoti_c_shim:
|
||||
return backend_index
|
||||
|
||||
elif backend_indices[DispatchKey.CompositeExplicitAutograd].has_kernel(func):
|
||||
# We need to create C shim wrappers for CompositeExplicitAutograd kernels
|
||||
backend_index = backend_indices[DispatchKey.CompositeExplicitAutograd]
|
||||
elif backend_indices[
|
||||
DispatchKey.CompositeExplicitAutogradNonFunctional
|
||||
].has_kernel(func):
|
||||
# We need to create C shim wrappers for CompositeExplicitAutogradNonFunctional kernels
|
||||
backend_index = backend_indices[
|
||||
DispatchKey.CompositeExplicitAutogradNonFunctional
|
||||
]
|
||||
elif backend_indices[DispatchKey.CompositeImplicitAutograd].has_kernel(func):
|
||||
backend_index = backend_indices[DispatchKey.CompositeImplicitAutograd]
|
||||
|
||||
return backend_index
|
||||
|
||||
|
||||
def get_header_for_aoti(
|
||||
func: NativeFunction,
|
||||
func_group_mapping: dict[OperatorName, NativeFunctionsGroup],
|
||||
dispatch_key: DispatchKey | None,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
extend_aoti_c_shim: bool,
|
||||
) -> str | None:
|
||||
backend_index = get_backend_index_for_aoti(
|
||||
func, func_group_mapping, dispatch_key, backend_indices, extend_aoti_c_shim
|
||||
)
|
||||
if backend_index is None:
|
||||
if dispatch_key is None:
|
||||
return f"#include <ATen/ops/{func.root_name}.h>"
|
||||
return None
|
||||
|
||||
return f"#include <ATen/ops/{func.root_name}_{backend_index.dispatch_key.lower()}_dispatch.h>"
|
||||
|
||||
|
||||
def get_fallback_op_name(func: NativeFunction) -> str:
|
||||
return (
|
||||
f"{func.namespace}.{func.func.name.name}.{func.func.name.overload_name}"
|
||||
if func.func.name.overload_name
|
||||
else f"{func.namespace}.{func.func.name.name}.default"
|
||||
)
|
||||
|
||||
|
||||
def gen_c_shim(
|
||||
func: NativeFunction,
|
||||
version_info: dict[str, list[str]],
|
||||
func_group_mapping: dict[OperatorName, NativeFunctionsGroup],
|
||||
dispatch_key: DispatchKey | None,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
header: bool,
|
||||
extend_aoti_c_shim: bool,
|
||||
) -> str | None:
|
||||
backend_index = get_backend_index_for_aoti(
|
||||
func, func_group_mapping, dispatch_key, backend_indices, extend_aoti_c_shim
|
||||
)
|
||||
if backend_index is None and dispatch_key is not None:
|
||||
return None
|
||||
|
||||
schema = func.func
|
||||
device = "aten" if dispatch_key is None else dispatch_key.lower()
|
||||
backend_call = gen_static_dispatch_backend_call(
|
||||
func,
|
||||
backend_index,
|
||||
)
|
||||
|
||||
try:
|
||||
if header:
|
||||
declaration, _ = gen_declaration_and_definition(
|
||||
schema, device, backend_call, version_info
|
||||
)
|
||||
return declaration
|
||||
else:
|
||||
_, definition = gen_declaration_and_definition(
|
||||
schema, device, backend_call, version_info
|
||||
)
|
||||
return definition
|
||||
|
||||
except NotImplementedError:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShimGenerator:
|
||||
inductor_fallback_ops: dict[str, dict[str, list[str]]]
|
||||
func_group_mapping: dict[OperatorName, NativeFunctionsGroup]
|
||||
dispatch_key: DispatchKey | None
|
||||
backend_indices: dict[DispatchKey, BackendIndex]
|
||||
header: bool # True to generate .h and False to generate .cpp
|
||||
extend_aoti_c_shim: bool
|
||||
|
||||
@method_with_native_function
|
||||
def __call__(
|
||||
self,
|
||||
func: NativeFunction,
|
||||
) -> str | None:
|
||||
version_info = self.inductor_fallback_ops[get_fallback_op_name(func)]
|
||||
result = gen_c_shim(
|
||||
func,
|
||||
version_info,
|
||||
self.func_group_mapping,
|
||||
self.dispatch_key,
|
||||
self.backend_indices,
|
||||
self.header,
|
||||
self.extend_aoti_c_shim,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def gen_aoti_c_shim(
|
||||
native_functions: Sequence[NativeFunction],
|
||||
inductor_fallback_ops: dict[str, dict[str, list[str]]],
|
||||
func_group_mapping: dict[OperatorName, NativeFunctionsGroup],
|
||||
dispatch_key: DispatchKey | None,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
header: bool,
|
||||
extend_aoti_c_shim: bool,
|
||||
includes: str = "",
|
||||
) -> str:
|
||||
body = "\n".join(
|
||||
list(
|
||||
mapMaybe(
|
||||
ShimGenerator(
|
||||
inductor_fallback_ops,
|
||||
func_group_mapping,
|
||||
dispatch_key,
|
||||
backend_indices,
|
||||
header,
|
||||
extend_aoti_c_shim,
|
||||
),
|
||||
native_functions,
|
||||
)
|
||||
)
|
||||
)
|
||||
device = "aten" if dispatch_key is None else dispatch_key.lower()
|
||||
include_device_functions = (
|
||||
"#include <ATen/Functions.h>"
|
||||
if dispatch_key is None
|
||||
else f"#include <ATen/{str(dispatch_key)}Functions.h>"
|
||||
)
|
||||
aten_warning = (
|
||||
(
|
||||
"\n\n// This file corresponds to the aten_shimified_ops list in torchgen/aoti/fallback_ops.py\n"
|
||||
)
|
||||
if dispatch_key is None
|
||||
else ""
|
||||
)
|
||||
warning = """
|
||||
|
||||
// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.
|
||||
// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details"""
|
||||
|
||||
if header:
|
||||
return (
|
||||
warning
|
||||
+ aten_warning
|
||||
+ textwrap.dedent("""
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/csrc/inductor/aoti_torch/c/shim.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
""")
|
||||
+ body
|
||||
+ textwrap.dedent("""
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // extern "C"
|
||||
#endif
|
||||
""")
|
||||
)
|
||||
else:
|
||||
return (
|
||||
warning
|
||||
+ aten_warning
|
||||
+ textwrap.dedent(f"""
|
||||
|
||||
#include <torch/csrc/inductor/aoti_torch/generated/{"extend/" if extend_aoti_c_shim else ""}c_shim_{device}.h>
|
||||
#include <torch/csrc/inductor/aoti_torch/utils.h>
|
||||
|
||||
#ifndef AT_PER_OPERATOR_HEADERS
|
||||
{include_device_functions}
|
||||
#include <ATen/CompositeExplicitAutogradFunctions.h>
|
||||
#include <ATen/CompositeExplicitAutogradNonFunctionalFunctions.h>
|
||||
#include <ATen/CompositeImplicitAutogradFunctions.h>
|
||||
#else
|
||||
""")
|
||||
+ includes
|
||||
+ textwrap.dedent("""
|
||||
#endif // AT_PER_OPERATOR_HEADERS
|
||||
|
||||
using namespace torch::aot_inductor;
|
||||
|
||||
""")
|
||||
+ body
|
||||
)
|
||||
|
||||
|
||||
def gen_aoti_c_shim_files(
|
||||
aoti_fm: FileManager,
|
||||
aoti_backends: set[DispatchKey | None],
|
||||
native_functions: Sequence[NativeFunction],
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
structured_native_functions: Sequence[NativeFunctionsGroup],
|
||||
extra_cuda_headers: str,
|
||||
extend_aoti_c_shim: bool,
|
||||
update_aoti_c_shim: bool,
|
||||
) -> None:
|
||||
structured_func_group_dict = {}
|
||||
for func_group in structured_native_functions:
|
||||
for func in func_group.functions():
|
||||
if func.structured_delegate is not None:
|
||||
structured_func_group_dict[func.structured_delegate] = func_group
|
||||
break
|
||||
|
||||
for dispatch_key in aoti_backends:
|
||||
# Use aten_shimified_ops for the aten backend, inductor_fallback_ops for others
|
||||
fallback_ops_dict = (
|
||||
aten_shimified_ops if dispatch_key is None else inductor_fallback_ops
|
||||
)
|
||||
fallbacks = {}
|
||||
for func in native_functions:
|
||||
op_name = get_fallback_op_name(func)
|
||||
if op_name in fallback_ops_dict:
|
||||
fallbacks[op_name] = func
|
||||
fallback_native_functions = tuple(
|
||||
value for _, value in sorted(fallbacks.items())
|
||||
)
|
||||
|
||||
# Use "aten" as the device name when dispatch_key is Generic
|
||||
device_name = "aten" if dispatch_key is None else dispatch_key.lower()
|
||||
|
||||
# header files were checked in for ABI-compatibility checking
|
||||
header_file_name = f"c_shim_{device_name}.h"
|
||||
new_header = gen_aoti_c_shim(
|
||||
fallback_native_functions,
|
||||
fallback_ops_dict,
|
||||
structured_func_group_dict,
|
||||
dispatch_key,
|
||||
backend_indices,
|
||||
header=True,
|
||||
extend_aoti_c_shim=extend_aoti_c_shim,
|
||||
includes="",
|
||||
)
|
||||
if update_aoti_c_shim:
|
||||
aoti_fm.write(
|
||||
header_file_name,
|
||||
lambda: new_header,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
with open(
|
||||
os.path.join(aoti_fm.install_dir, header_file_name)
|
||||
) as old_file:
|
||||
old_header = old_file.read()
|
||||
|
||||
if old_header != new_header:
|
||||
diff = "\n".join(
|
||||
difflib.unified_diff(
|
||||
old_header.splitlines(),
|
||||
new_header.splitlines(),
|
||||
fromfile="expected",
|
||||
tofile="actual",
|
||||
lineterm="",
|
||||
)
|
||||
)
|
||||
|
||||
raise RuntimeError(f"""
|
||||
The generated AOTInductor C shim header files have unexpectedly changed. This
|
||||
indicates an AOTInductor fallback operator ABI backward compatibility breakage!!!
|
||||
Only in a limited number of situations, this is allowed:
|
||||
|
||||
1. You added a fallback op to the inductor_fallback_ops list in torchgen/aoti/fallback_ops.py.
|
||||
If that's the case, run `python torchgen/gen.py --update-aoti-c-shim` to add a new entry to
|
||||
existing C shim header files.
|
||||
|
||||
2. You added a new default argument to an existing fallback op. This is clearly a BC breaking
|
||||
change in the AOTInductor land. You need to annotate the new default argument in
|
||||
torchgen/aoti/fallback_ops.py, and then run `python torchgen/gen.py --update-aoti-c-shim` to
|
||||
update the C shim header files by creating different versions of the fallback op. See
|
||||
https://github.com/pytorch/pytorch/pull/154848 as an example.
|
||||
|
||||
{diff}
|
||||
""")
|
||||
except FileNotFoundError:
|
||||
print(
|
||||
f"{os.path.join(aoti_fm.install_dir, header_file_name)} not found"
|
||||
)
|
||||
|
||||
# cpp files are always generated on-the-fly
|
||||
def headers_for_aoti() -> str:
|
||||
headers = []
|
||||
for func in fallback_native_functions:
|
||||
header = get_header_for_aoti(
|
||||
func,
|
||||
structured_func_group_dict,
|
||||
dispatch_key,
|
||||
backend_indices,
|
||||
extend_aoti_c_shim=extend_aoti_c_shim,
|
||||
)
|
||||
if header is not None:
|
||||
headers.append(header)
|
||||
return "\n".join(sorted(set(headers)))
|
||||
|
||||
extra_headers = (
|
||||
extra_cuda_headers
|
||||
if dispatch_key is not None and is_cuda_dispatch_key(dispatch_key)
|
||||
else ""
|
||||
)
|
||||
|
||||
aoti_fm.write(
|
||||
f"c_shim_{device_name}.cpp",
|
||||
lambda: gen_aoti_c_shim(
|
||||
fallback_native_functions,
|
||||
fallback_ops_dict,
|
||||
structured_func_group_dict,
|
||||
dispatch_key,
|
||||
backend_indices,
|
||||
header=False,
|
||||
extend_aoti_c_shim=extend_aoti_c_shim,
|
||||
includes=headers_for_aoti() + "\n" + extra_headers,
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,635 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
from collections import Counter, defaultdict, namedtuple
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import yaml
|
||||
|
||||
import torchgen.api.dispatcher as dispatcher
|
||||
import torchgen.dest as dest
|
||||
from torchgen.api.types import DispatcherSignature
|
||||
from torchgen.code_template import CodeTemplate
|
||||
from torchgen.context import native_function_manager
|
||||
from torchgen.gen import get_grouped_native_functions, parse_native_yaml
|
||||
from torchgen.model import (
|
||||
BackendIndex,
|
||||
BackendMetadata,
|
||||
DispatchKey,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
OperatorName,
|
||||
)
|
||||
from torchgen.selective_build.selector import SelectiveBuilder
|
||||
from torchgen.utils import concatMap, context, FileManager, NamespaceHelper, Target
|
||||
from torchgen.yaml_utils import YamlLoader
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key.
|
||||
# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping)
|
||||
ParsedExternalYaml = namedtuple(
|
||||
"ParsedExternalYaml",
|
||||
["backend_key", "autograd_key", "class_name", "cpp_namespace", "backend_indices"],
|
||||
)
|
||||
|
||||
|
||||
def parse_backend_yaml(
|
||||
backend_yaml_path: str,
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
) -> ParsedExternalYaml:
|
||||
native_functions_map: dict[OperatorName, NativeFunction] = {
|
||||
f.func.name: f
|
||||
for f in concatMap(
|
||||
lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()),
|
||||
grouped_native_functions,
|
||||
)
|
||||
}
|
||||
|
||||
with open(backend_yaml_path) as f:
|
||||
yaml_values = yaml.load(f, Loader=YamlLoader)
|
||||
if not isinstance(yaml_values, dict):
|
||||
raise AssertionError(
|
||||
f"Expected yaml_values to be a dict, got {type(yaml_values)}"
|
||||
)
|
||||
|
||||
valid_keys = [
|
||||
"backend",
|
||||
"class_name",
|
||||
"cpp_namespace",
|
||||
"extra_headers",
|
||||
"supported",
|
||||
"autograd",
|
||||
"full_codegen",
|
||||
"non_native",
|
||||
"ir_gen",
|
||||
"symint",
|
||||
]
|
||||
|
||||
backend = yaml_values.pop("backend", None)
|
||||
if backend is None:
|
||||
raise AssertionError('You must provide a value for "backend"')
|
||||
|
||||
class_name = yaml_values.pop("class_name", None)
|
||||
|
||||
cpp_namespace = yaml_values.pop("cpp_namespace", None)
|
||||
if cpp_namespace is None:
|
||||
raise AssertionError('You must provide a value for "cpp_namespace"')
|
||||
|
||||
# Mostly just defaulting to false to stick with LazyTensor convention.
|
||||
use_out_as_primary = yaml_values.pop("use_out_as_primary", False)
|
||||
if not isinstance(use_out_as_primary, bool):
|
||||
raise AssertionError(
|
||||
f"You must provide either True or False for use_out_as_primary. Provided: {use_out_as_primary}"
|
||||
)
|
||||
|
||||
use_device_guard = yaml_values.pop("device_guard", False)
|
||||
if not isinstance(use_device_guard, bool):
|
||||
raise AssertionError(
|
||||
f"You must provide either True or False for device_guard. Provided: {use_device_guard}"
|
||||
)
|
||||
|
||||
supported = yaml_values.pop("supported", [])
|
||||
if supported is None:
|
||||
supported = [] # Allow an empty list of supported ops
|
||||
if not isinstance(supported, list):
|
||||
raise AssertionError(
|
||||
f'expected "supported" to be a list, but got: {supported} (of type {type(supported)})'
|
||||
)
|
||||
|
||||
symint = yaml_values.pop("symint", [])
|
||||
if symint is None:
|
||||
symint = [] # Allow an empty list of symint ops
|
||||
if not isinstance(symint, list):
|
||||
raise AssertionError(
|
||||
f'expected "symint" to be a list, but got: {symint} (of type {type(symint)})'
|
||||
)
|
||||
symint_set = set(symint)
|
||||
|
||||
supported_autograd = yaml_values.pop("autograd", [])
|
||||
if not isinstance(supported_autograd, list):
|
||||
raise AssertionError(
|
||||
f'expected "autograd" to be a list, but got: {supported_autograd}'
|
||||
)
|
||||
|
||||
# full_codegen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
|
||||
full_codegen = yaml_values.pop("full_codegen", [])
|
||||
supported.extend(full_codegen)
|
||||
|
||||
# non_native is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
|
||||
yaml_values.pop("non_native", {})
|
||||
|
||||
# ir_gen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
|
||||
yaml_values.pop("ir_gen", {})
|
||||
|
||||
if len(yaml_values.keys()) != 0:
|
||||
raise AssertionError(
|
||||
f"{backend_yaml_path} contains unexpected keys: {', '.join(yaml_values.keys())}. "
|
||||
f"Only the following keys are supported: {', '.join(valid_keys)}"
|
||||
)
|
||||
|
||||
def create_backend_index(
|
||||
backend_ops: list[str],
|
||||
symint_ops: set[str],
|
||||
dispatch_key: DispatchKey,
|
||||
*,
|
||||
use_out_as_primary: bool,
|
||||
use_device_guard: bool,
|
||||
) -> BackendIndex:
|
||||
metadata: dict[OperatorName, BackendMetadata] = {}
|
||||
for op in backend_ops:
|
||||
op_name = OperatorName.parse(op)
|
||||
if op_name not in native_functions_map:
|
||||
raise AssertionError(f"Found an invalid operator name: {op_name}")
|
||||
# See Note [External Backends Follow Dispatcher API]
|
||||
kernel_name = dispatcher.name(native_functions_map[op_name].func)
|
||||
if op in symint_ops:
|
||||
kernel_name += "_symint"
|
||||
# TODO: allow structured external backends later.
|
||||
m = BackendMetadata(
|
||||
kernel=kernel_name, structured=False, cpp_namespace=cpp_namespace
|
||||
)
|
||||
metadata[op_name] = m
|
||||
return BackendIndex(
|
||||
dispatch_key=dispatch_key,
|
||||
use_out_as_primary=use_out_as_primary,
|
||||
external=True,
|
||||
device_guard=use_device_guard,
|
||||
index=metadata,
|
||||
)
|
||||
|
||||
backend_key: DispatchKey | None = None
|
||||
if len(supported) > 0:
|
||||
with context(
|
||||
lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.'
|
||||
):
|
||||
backend_key = DispatchKey.parse(backend)
|
||||
|
||||
backend_idx = create_backend_index(
|
||||
supported,
|
||||
symint_set,
|
||||
backend_key,
|
||||
use_out_as_primary=use_out_as_primary,
|
||||
use_device_guard=use_device_guard,
|
||||
)
|
||||
if backend_key in backend_indices:
|
||||
raise AssertionError(f"Duplicate backend key: {backend_key}")
|
||||
backend_indices[backend_key] = backend_idx
|
||||
|
||||
autograd_key: DispatchKey | None = None
|
||||
if len(supported_autograd) > 0:
|
||||
with context(
|
||||
lambda: f'The "autograd" key was specified, which indicates that you would like to override \
|
||||
the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.'
|
||||
):
|
||||
autograd_key = DispatchKey.parse(f"Autograd{backend}")
|
||||
|
||||
autograd_idx = create_backend_index(
|
||||
supported_autograd,
|
||||
symint_set,
|
||||
autograd_key,
|
||||
use_out_as_primary=use_out_as_primary,
|
||||
use_device_guard=use_device_guard,
|
||||
)
|
||||
if autograd_key in backend_indices:
|
||||
raise AssertionError(f"Duplicate autograd key: {autograd_key}")
|
||||
backend_indices[autograd_key] = autograd_idx
|
||||
|
||||
for g in grouped_native_functions:
|
||||
if isinstance(g, NativeFunction):
|
||||
forward_kernels = (
|
||||
[]
|
||||
if backend_key is None
|
||||
else [
|
||||
m
|
||||
for m in [backend_indices[backend_key].get_kernel(g)]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
backward_kernels = (
|
||||
[]
|
||||
if autograd_key is None
|
||||
else [
|
||||
m
|
||||
for m in [backend_indices[autograd_key].get_kernel(g)]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
else:
|
||||
forward_kernels = (
|
||||
[]
|
||||
if backend_key is None
|
||||
else [
|
||||
m
|
||||
for m in [
|
||||
backend_indices[backend_key].get_kernel(f)
|
||||
for f in g.functions()
|
||||
]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
backward_kernels = (
|
||||
[]
|
||||
if autograd_key is None
|
||||
else [
|
||||
m
|
||||
for m in [
|
||||
backend_indices[autograd_key].get_kernel(f)
|
||||
for f in g.functions()
|
||||
]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
|
||||
forward_kernels = [f for f in forward_kernels if f is not None]
|
||||
backward_kernels = [f for f in backward_kernels if f is not None]
|
||||
if not (len(forward_kernels) == 0 or len(backward_kernels) == 0):
|
||||
raise AssertionError(
|
||||
f"Currently, all variants of an op must either be registered to a backend key, "
|
||||
f"or to a backend's autograd key. They cannot be mix and matched. "
|
||||
f"If this is something you need, feel free to create an issue! "
|
||||
f'{forward_kernels[0].kernel} is listed under "supported", '
|
||||
f'but {backward_kernels[0].kernel} is listed under "autograd".'
|
||||
)
|
||||
|
||||
return ParsedExternalYaml(
|
||||
backend_key, autograd_key, class_name, cpp_namespace, backend_indices
|
||||
)
|
||||
|
||||
|
||||
def error_on_missing_kernels(
|
||||
native_functions: Sequence[NativeFunction],
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
backend_key: DispatchKey,
|
||||
autograd_key: DispatchKey | None,
|
||||
class_name: str,
|
||||
kernel_defn_file_path: str,
|
||||
full_codegen: list[OperatorName] | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
with open(kernel_defn_file_path) as f:
|
||||
backend_defns = f.read()
|
||||
except OSError as e:
|
||||
raise AssertionError(
|
||||
f"Unable to read from the specified impl_path file: {kernel_defn_file_path}"
|
||||
) from e
|
||||
|
||||
if full_codegen is None:
|
||||
full_codegen = []
|
||||
|
||||
indices = [backend_indices[backend_key].index] + (
|
||||
[] if autograd_key is None else [backend_indices[autograd_key].index]
|
||||
)
|
||||
# Quick mapping from each OperatorName used by the external backend
|
||||
# to its backend kernel name
|
||||
expected_backend_op_names: dict[OperatorName, str] = dict(
|
||||
list(
|
||||
concatMap(
|
||||
lambda index: [
|
||||
(op_name, metadata.kernel) for op_name, metadata in index.items()
|
||||
],
|
||||
indices,
|
||||
)
|
||||
)
|
||||
)
|
||||
expected_backend_native_funcs: list[NativeFunction] = [
|
||||
f
|
||||
for f in native_functions
|
||||
if f.func.name in expected_backend_op_names and f.func.name not in full_codegen
|
||||
]
|
||||
expected_backend_kernel_name_counts: dict[str, list[NativeFunction]] = defaultdict(
|
||||
list
|
||||
)
|
||||
for native_f in expected_backend_native_funcs:
|
||||
expected_backend_kernel_name_counts[
|
||||
expected_backend_op_names[native_f.func.name]
|
||||
].append(native_f)
|
||||
|
||||
# This just looks for lines containing "foo(", and assumes that the kernel foo has been implemented.
|
||||
# It might cause false negatives (we won't catch all cases), but that's ok - if we catch a missing kernel
|
||||
# here, then we get a nicer error message. If we miss it, you get a linker error.
|
||||
kernel_defn_regex = rf"(.*){class_name}::\s*([\w\d]*)\("
|
||||
actual_backend_kernel_name_counts = Counter(
|
||||
# A bit unwieldy (this could probably be moved into regex),
|
||||
# but we don't want to include kernel names that come from function calls,
|
||||
# like "return torch_xla::XLANativeFunctions::empty_strided_symint(...)".
|
||||
# Easy check is to ignore any lines with colons before the class name.
|
||||
[
|
||||
y
|
||||
for (x, y) in re.findall(kernel_defn_regex, backend_defns)
|
||||
if not x.endswith(":")
|
||||
]
|
||||
)
|
||||
|
||||
missing_kernels_err_msg = ""
|
||||
for expected_name, funcs in expected_backend_kernel_name_counts.items():
|
||||
expected_overload_count = len(funcs)
|
||||
actual_overload_count = actual_backend_kernel_name_counts[expected_name]
|
||||
if expected_overload_count != actual_overload_count:
|
||||
|
||||
def create_decl(f: NativeFunction) -> str:
|
||||
with native_function_manager(f):
|
||||
return DispatcherSignature.from_schema(f.func).decl()
|
||||
|
||||
expected_schemas_str = "\n".join([create_decl(f) for f in funcs])
|
||||
missing_kernels_err_msg += f"""
|
||||
{class_name} is missing a kernel definition for {expected_name}. We found {actual_overload_count} kernel(s) with that name,
|
||||
but expected {expected_overload_count} kernel(s). The expected function schemas for the missing operator are:
|
||||
{expected_schemas_str}
|
||||
|
||||
"""
|
||||
if missing_kernels_err_msg != "":
|
||||
raise AssertionError(missing_kernels_err_msg)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Generate backend stub files")
|
||||
parser.add_argument(
|
||||
"-s",
|
||||
"--source-yaml",
|
||||
"--source_yaml",
|
||||
help="path to source yaml file containing operator external definitions",
|
||||
)
|
||||
parser.add_argument("-o", "--output-dir", "--output_dir", help="output directory")
|
||||
parser.add_argument(
|
||||
"--dry-run", "--dry_run", type=bool, default=False, help="output directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--impl-path",
|
||||
"--impl_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="path to the source C++ file containing kernel definitions",
|
||||
)
|
||||
options = parser.parse_args()
|
||||
|
||||
run(options.source_yaml, options.output_dir, options.dry_run, options.impl_path)
|
||||
|
||||
|
||||
def gen_dispatchkey_nativefunc_headers(
|
||||
fm: FileManager,
|
||||
class_name: str,
|
||||
cpp_namespace: str,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
backend_dispatch_key: DispatchKey,
|
||||
autograd_dispatch_key: DispatchKey | None,
|
||||
backend_name: str = "",
|
||||
) -> None:
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must not be None")
|
||||
generated_comment = (
|
||||
"Autogenerated file by gen_backend_stubs.py. Do not edit directly!"
|
||||
)
|
||||
|
||||
# Convert to a set first to remove duplicate kernel names.
|
||||
# Backends are allowed to repeat kernel names; only generate the declaration once!
|
||||
# Sort for deterministic output.
|
||||
backend_declarations = sorted(
|
||||
set(
|
||||
concatMap(
|
||||
lambda f: dest.compute_native_function_declaration(
|
||||
f, backend_indices[backend_dispatch_key]
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
)
|
||||
autograd_declarations = sorted(
|
||||
set(
|
||||
concatMap(
|
||||
lambda f: []
|
||||
if autograd_dispatch_key is None
|
||||
else dest.compute_native_function_declaration(
|
||||
f, backend_indices[autograd_dispatch_key]
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
ns_helper = NamespaceHelper(cpp_namespace)
|
||||
fm.write_with_template(
|
||||
f"{backend_dispatch_key}NativeFunctions.h",
|
||||
"DispatchKeyNativeFunctions.h",
|
||||
lambda: {
|
||||
"generated_comment": generated_comment,
|
||||
"namespace_prologue": ns_helper.prologue,
|
||||
"class_name": class_name,
|
||||
"namespace_epilogue": ns_helper.epilogue,
|
||||
"dispatch_declarations": backend_declarations + autograd_declarations,
|
||||
"BackendName": backend_name,
|
||||
"DispatchKey": backend_dispatch_key,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def gen_dispatcher_registrations(
|
||||
fm: FileManager,
|
||||
output_dir: str,
|
||||
class_name: str,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
backend_dispatch_key: DispatchKey,
|
||||
dispatch_key: DispatchKey,
|
||||
selector: SelectiveBuilder,
|
||||
# build_in_tree is true for lazy TS backend and affects include paths, not used for external backends
|
||||
build_in_tree: bool = False,
|
||||
per_operator_headers: bool = False,
|
||||
backend_name: str = "",
|
||||
eager_registration: bool = True,
|
||||
) -> None:
|
||||
headers = [
|
||||
f"{output_dir}/{backend_dispatch_key}NativeFunctions.h",
|
||||
]
|
||||
if build_in_tree:
|
||||
external_backend_headers_str = "\n".join(f"#include <{h}>" for h in headers)
|
||||
else:
|
||||
external_backend_headers_str = "\n".join(f'#include "{h}"' for h in headers)
|
||||
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must not be None")
|
||||
backend_index = backend_indices[dispatch_key]
|
||||
|
||||
dispatch_registrations_body = list(
|
||||
concatMap(
|
||||
dest.RegisterDispatchKey(
|
||||
backend_index,
|
||||
Target.REGISTRATION,
|
||||
selector,
|
||||
rocm=False,
|
||||
symint=True,
|
||||
class_method_name=f"{class_name}",
|
||||
skip_dispatcher_op_registration=False,
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
newline = "\n"
|
||||
ns_helper = NamespaceHelper(namespace_str="at")
|
||||
deferred_dispatch_registrations = ""
|
||||
static_init_dispatch_registrations = ""
|
||||
if eager_registration:
|
||||
static_template = CodeTemplate(
|
||||
"""\
|
||||
TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) {
|
||||
$dispatch_registrations_body
|
||||
}"""
|
||||
)
|
||||
static_init_dispatch_registrations = static_template.substitute(
|
||||
dispatch_key=dispatch_key,
|
||||
dispatch_registrations_body=dispatch_registrations_body,
|
||||
)
|
||||
else:
|
||||
deferred_template = CodeTemplate(
|
||||
"""\
|
||||
TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions();
|
||||
TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
|
||||
static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key);
|
||||
$dispatch_registrations_body
|
||||
}"""
|
||||
)
|
||||
deferred_dispatch_registrations = deferred_template.substitute(
|
||||
backend_name=backend_name,
|
||||
dispatch_key=dispatch_key,
|
||||
dispatch_registrations_body=dispatch_registrations_body,
|
||||
)
|
||||
|
||||
fm.write_with_template(
|
||||
f"Register{dispatch_key}.cpp",
|
||||
"RegisterDispatchKey.cpp",
|
||||
lambda: {
|
||||
"extra_cuda_headers": "",
|
||||
"external_backend_headers": external_backend_headers_str,
|
||||
"ops_headers": "#include <ATen/Functions.h>"
|
||||
if not per_operator_headers
|
||||
else "",
|
||||
"DispatchKey": dispatch_key,
|
||||
"dispatch_namespace": dispatch_key.lower(),
|
||||
"dispatch_headers": dest.gen_registration_headers(
|
||||
backend_index, per_operator_headers=per_operator_headers, rocm=False
|
||||
),
|
||||
"dispatch_helpers": dest.gen_registration_helpers(backend_index),
|
||||
"dispatch_definitions": fm.substitute_with_template(
|
||||
"RegisterDispatchDefinitions.ini",
|
||||
lambda: {
|
||||
"ns_prologue": ns_helper.prologue,
|
||||
"ns_epilogue": ns_helper.epilogue,
|
||||
"static_init_dispatch_registrations": static_init_dispatch_registrations,
|
||||
"deferred_dispatch_registrations": deferred_dispatch_registrations,
|
||||
"dispatch_namespace": dispatch_key.lower(),
|
||||
"dispatch_namespaced_definitions": "",
|
||||
"dispatch_anonymous_definitions": list(
|
||||
concatMap(
|
||||
dest.RegisterDispatchKey(
|
||||
backend_index,
|
||||
Target.ANONYMOUS_DEFINITION,
|
||||
selector,
|
||||
rocm=False,
|
||||
symint=True,
|
||||
class_method_name=f"{class_name}",
|
||||
skip_dispatcher_op_registration=False,
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
),
|
||||
},
|
||||
).split(newline),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def run(
|
||||
source_yaml: str, output_dir: str, dry_run: bool, impl_path: str | None = None
|
||||
) -> None:
|
||||
# Assumes that this file lives at torchgen/gen_backend_stubs.py
|
||||
root = Path(__file__).absolute().parent.parent
|
||||
common_dir = os.path.join(root, "aten/src") # Assumes root is pytorch_root
|
||||
if not os.path.exists(common_dir): # This file is out-of-tree.
|
||||
common_dir = os.path.join(root, "torchgen/packaged")
|
||||
|
||||
template_dir = os.path.join(common_dir, "ATen/templates")
|
||||
|
||||
def make_file_manager(install_dir: str) -> FileManager:
|
||||
return FileManager(
|
||||
install_dir=install_dir, template_dir=template_dir, dry_run=dry_run
|
||||
)
|
||||
|
||||
fm = make_file_manager(output_dir)
|
||||
|
||||
native_yaml_path = os.path.join(common_dir, "ATen/native/native_functions.yaml")
|
||||
tags_yaml_path = os.path.join(common_dir, "ATen/native/tags.yaml")
|
||||
parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
|
||||
native_functions, backend_indices = (
|
||||
parsed_yaml.native_functions,
|
||||
parsed_yaml.backend_indices,
|
||||
)
|
||||
grouped_native_functions = get_grouped_native_functions(native_functions)
|
||||
parsed_backend_yaml = parse_backend_yaml(
|
||||
source_yaml, grouped_native_functions, backend_indices
|
||||
)
|
||||
backend_key = parsed_backend_yaml.backend_key
|
||||
autograd_key = parsed_backend_yaml.autograd_key
|
||||
cpp_namespace = parsed_backend_yaml.cpp_namespace
|
||||
class_name = parsed_backend_yaml.class_name
|
||||
backend_indices = parsed_backend_yaml.backend_indices
|
||||
|
||||
selector = SelectiveBuilder.get_nop_selector()
|
||||
|
||||
if backend_key is None:
|
||||
# This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet.
|
||||
return
|
||||
|
||||
if class_name is None:
|
||||
# class_name is an optional argument to backend yaml file.
|
||||
# if specified it allows an external backend to override
|
||||
# the name of the class that all generated kernel definitions live under.
|
||||
# if not specified, its value is given as native_function_class_name.
|
||||
class_name = backend_indices[backend_key].native_function_class_name()
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must not be None")
|
||||
|
||||
if impl_path is not None:
|
||||
error_on_missing_kernels(
|
||||
native_functions,
|
||||
backend_indices,
|
||||
backend_key,
|
||||
autograd_key,
|
||||
class_name,
|
||||
impl_path,
|
||||
)
|
||||
|
||||
gen_dispatchkey_nativefunc_headers(
|
||||
fm,
|
||||
class_name,
|
||||
cpp_namespace,
|
||||
backend_indices,
|
||||
grouped_native_functions,
|
||||
backend_key,
|
||||
autograd_key,
|
||||
)
|
||||
|
||||
for dispatch_key in (
|
||||
[backend_key] if autograd_key is None else [backend_key, autograd_key]
|
||||
):
|
||||
gen_dispatcher_registrations(
|
||||
fm,
|
||||
output_dir,
|
||||
class_name,
|
||||
backend_indices,
|
||||
grouped_native_functions,
|
||||
backend_key,
|
||||
dispatch_key,
|
||||
selector,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,593 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from collections import namedtuple
|
||||
from pathlib import Path
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import yaml
|
||||
|
||||
import torchgen.dest as dest
|
||||
from torchgen.api.lazy import setValueT
|
||||
from torchgen.api.types import BaseCppType
|
||||
from torchgen.dest.lazy_ir import GenLazyIR, GenLazyNativeFuncDefinition, GenTSLazyIR
|
||||
from torchgen.gen import get_grouped_native_functions, parse_native_yaml
|
||||
from torchgen.gen_backend_stubs import (
|
||||
error_on_missing_kernels,
|
||||
gen_dispatcher_registrations,
|
||||
gen_dispatchkey_nativefunc_headers,
|
||||
parse_backend_yaml,
|
||||
)
|
||||
from torchgen.model import NativeFunction, NativeFunctionsGroup, OperatorName
|
||||
from torchgen.selective_build.selector import SelectiveBuilder
|
||||
from torchgen.utils import FileManager, NamespaceHelper
|
||||
from torchgen.yaml_utils import YamlLoader
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterable, Iterator, Sequence
|
||||
|
||||
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
#
|
||||
# Lazy Tensor Codegen
|
||||
#
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
# Overview
|
||||
# ~~~~~~~~
|
||||
#
|
||||
# This codegen script builds on existing data models and helpers used
|
||||
# by all ATen backends, and adds new functionality specific to lazy
|
||||
# tensor backends.
|
||||
#
|
||||
# Inputs:
|
||||
# - <backend>_native_functions.yaml: controls which operators are
|
||||
# supported by the backend.
|
||||
#
|
||||
# Outputs:
|
||||
# (for all backends)
|
||||
# <DispatchKey>Ir.h defines Lazy IR classes to be constructed during tracing
|
||||
# - opt-in: also generate 'lowering' methods for the TorchScript backend only
|
||||
# <DispatchKey>NativeFunctions.cpp defines implementations of native functions which perform lazy tracing
|
||||
# - opt-in: 'full_codegen' section of backend yaml; 'supported' section omits these implementations
|
||||
# <DispatchKey>NativeFunctions.h declares implementations of native functions for both 'supported' and 'full_codegen'
|
||||
# ops
|
||||
#
|
||||
# Register<DispatchKey>.cpp registers all op implementations with the dispatcher
|
||||
# RegisterAutograd<DispatchKey>.cpp registers all autograd implementations with the dispatcher
|
||||
#
|
||||
# Validation Helpers:
|
||||
# - Shape Inference: errs if any ops in backend yaml require shape inference not provided by meta kernels or
|
||||
# implementations in torch/csrc/lazy/core/shape_inference.*
|
||||
# - native function impls: errs if any 'supported' ops do not have an implementation defined in the backend
|
||||
# (non-codegen) implementation file
|
||||
#
|
||||
#
|
||||
# About the Data Model
|
||||
# ~~~~~~~~~~~~~~~~~~~~
|
||||
#
|
||||
# Modeled after ATen codegen, the first step is to parse yaml and build a data model for the operators
|
||||
# we care about. In this case, the <backend>_native_functions yaml defines a subset of the core operators
|
||||
# (defined in more detail in the main native_functions.yaml), which will be supported by your backend.
|
||||
# Backends can list ops in two categories:
|
||||
# - `supported` ops require hand-implementations but still get codegenned declarations and registrations
|
||||
# - `full_codegen` ops get implementations (and IR classes) generated too
|
||||
#
|
||||
# Each native function is modeled as an object with a schema, and each schema has objects representing their
|
||||
# arguments. Much of the codegen is manipulation of the arguments and their types. For example, lazy tensor
|
||||
# backends need to transform 'at::Tensor' arguments into 'lazy::Value' objects, as well as replacing reference
|
||||
# types (stringref) with actual string objects, and this is done by manipulating the data model objects.
|
||||
# - see api/lazy.py for the lazy data model
|
||||
#
|
||||
# Once the data model is set up, the rest of this script processes a number of templates for output CPP file
|
||||
# and fills in the template values using helpers in `dest/lazy_ir.py` and `dest/lazy_ts_lowering.py`. These
|
||||
# helpers mostly iterate over functions and their arguments, outputting different c++ snippets.
|
||||
#
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
|
||||
|
||||
|
||||
# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key.
|
||||
# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping, full_codegen)
|
||||
ParsedExternalYaml = namedtuple(
|
||||
"ParsedExternalYaml",
|
||||
["backend_key", "autograd_key", "cpp_namespace", "backend_indices", "full_codegen"],
|
||||
)
|
||||
|
||||
|
||||
def parse_native_functions_keys(
|
||||
backend_yaml_path: str,
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
) -> tuple[list[OperatorName], list[Any], list[OperatorName]]:
|
||||
with open(backend_yaml_path) as f:
|
||||
yaml_values = yaml.load(f, Loader=YamlLoader)
|
||||
if not isinstance(yaml_values, dict):
|
||||
raise AssertionError(f"Expected dict from YAML, got {type(yaml_values)}")
|
||||
|
||||
full_codegen = yaml_values.pop("full_codegen", [])
|
||||
non_native = yaml_values.pop("non_native", [])
|
||||
ir_gen = yaml_values.pop("ir_gen", [])
|
||||
if not isinstance(full_codegen, list):
|
||||
raise AssertionError(
|
||||
f"Expected full_codegen to be list, got {type(full_codegen)}"
|
||||
)
|
||||
if not isinstance(non_native, list):
|
||||
raise AssertionError(f"Expected non_native to be list, got {type(non_native)}")
|
||||
if not isinstance(ir_gen, list):
|
||||
raise AssertionError(f"Expected ir_gen to be list, got {type(ir_gen)}")
|
||||
full_codegen_opnames = [OperatorName.parse(name) for name in full_codegen]
|
||||
ir_gen_opnames = [OperatorName.parse(name) for name in ir_gen]
|
||||
return full_codegen_opnames, non_native, ir_gen_opnames
|
||||
|
||||
|
||||
def validate_shape_inference_header(
|
||||
shape_inference_hdr: str, expected_shape_infr_decls: list[str]
|
||||
) -> None:
|
||||
try:
|
||||
with open(shape_inference_hdr) as f:
|
||||
shape_infr_decls = f.read()
|
||||
shape_infr_decl_lines = set(shape_infr_decls.split("\n"))
|
||||
except OSError as e:
|
||||
raise AssertionError(
|
||||
f"Unable to read from the specified shape_inference_hdr file: {shape_inference_hdr}"
|
||||
) from e
|
||||
|
||||
# TODO(whc) add a check for shape inference functions that have meta kernels implement and should be retired.
|
||||
|
||||
missing_decls = [
|
||||
decl for decl in expected_shape_infr_decls if decl not in shape_infr_decl_lines
|
||||
]
|
||||
if missing_decls:
|
||||
raise Exception( # noqa: TRY002
|
||||
f"""Missing shape inference function.\n
|
||||
Please add declare this function in {shape_inference_hdr}:\n
|
||||
and implement it in the corresponding shape_inference.cpp file.\n
|
||||
{os.linesep.join(missing_decls)}"""
|
||||
)
|
||||
|
||||
|
||||
# Some helper functions for the codegen.
|
||||
def get_ltc_helper_fns() -> str:
|
||||
return """\
|
||||
at::Tensor to_meta(const at::Tensor& tensor) {
|
||||
// undefined tensors can't be converted to the meta device, since they don't have sizes/strides
|
||||
if (!tensor.defined()) return tensor;
|
||||
auto out = at::native::empty_strided_meta_symint(tensor.sym_sizes(), tensor.sym_strides(), \
|
||||
/*dtype=*/tensor.scalar_type(), /*layout=*/tensor.layout(), \
|
||||
/*device=*/c10::Device(c10::kMeta), /*pin_memory=*/std::nullopt);
|
||||
// needs to handle wrapped numbers, so dtype promotion works properly.
|
||||
if (tensor.unsafeGetTensorImpl()->is_wrapped_number()) {
|
||||
out.unsafeGetTensorImpl()->set_wrapped_number(true);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
std::optional<at::Tensor> to_meta(const std::optional<at::Tensor>& tensor) {
|
||||
if (tensor.has_value()) {
|
||||
return to_meta(*tensor);
|
||||
}
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> to_meta(at::ITensorListRef t_list) {
|
||||
std::vector<at::Tensor> outs;
|
||||
outs.reserve(t_list.size());
|
||||
for (const auto& tensor : t_list) {
|
||||
outs.push_back(to_meta(tensor));
|
||||
}
|
||||
return outs;
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
class default_args:
|
||||
node_base: str = "Node"
|
||||
node_base_hdr: str | None = None
|
||||
shape_inference_hdr: str = "torch/csrc/lazy/core/shape_inference.h"
|
||||
tensor_class: str = "torch::lazy::LazyTensor"
|
||||
tensor_class_hdr: str = "torch/csrc/lazy/core/tensor.h"
|
||||
lazy_ir_generator: type[GenLazyIR] = GenLazyIR
|
||||
native_func_definition_generator: type[GenLazyNativeFuncDefinition] = (
|
||||
GenLazyNativeFuncDefinition
|
||||
)
|
||||
backend_name: str = "TorchScript"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Generate Lazy Tensor backend files")
|
||||
parser.add_argument(
|
||||
"-s",
|
||||
"--source-yaml",
|
||||
"--source_yaml",
|
||||
help="path to source yaml file containing operator external definitions",
|
||||
)
|
||||
parser.add_argument("-o", "--output-dir", "--output_dir", help="output directory")
|
||||
parser.add_argument(
|
||||
"--dry-run", "--dry_run", type=bool, default=False, help="output directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--impl-path",
|
||||
"--impl_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="path to the source C++ file containing kernel definitions",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gen-ts-lowerings",
|
||||
"--gen_ts_lowerings",
|
||||
action="store_true",
|
||||
help="Generate TorchScript lowerings in addition to Lazy IR and NativeFunctions",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--node-base",
|
||||
"--node_base",
|
||||
type=str,
|
||||
default=default_args.node_base,
|
||||
help="Name of backend specific custom Lazy IR Node base class",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--node-base-hdr",
|
||||
"--node_base_hdr",
|
||||
type=str,
|
||||
default=default_args.node_base_hdr,
|
||||
help="Path to header file defining custom Lazy IR Node base class",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shape-inference-hdr",
|
||||
"--shape_inference_hdr",
|
||||
type=str,
|
||||
default=default_args.shape_inference_hdr,
|
||||
help="Path to header file defining custom Lazy shape inference functions",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-class",
|
||||
"--tensor_class",
|
||||
type=str,
|
||||
default=default_args.tensor_class,
|
||||
help="Name of backend specific custom Lazy Tensor class",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-class-hdr",
|
||||
"--tensor_class_hdr",
|
||||
type=str,
|
||||
default=default_args.tensor_class_hdr,
|
||||
help="Path to header file defining custom Lazy Tensor class",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--backend-name",
|
||||
"--backend_name",
|
||||
type=str,
|
||||
default=default_args.backend_name,
|
||||
help="Name of the backend to generate",
|
||||
)
|
||||
options = parser.parse_args()
|
||||
|
||||
# Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py
|
||||
torch_root = Path(__file__).absolute().parents[2]
|
||||
aten_path = str(torch_root / "aten" / "src" / "ATen")
|
||||
lazy_ir_generator: type[GenLazyIR] = default_args.lazy_ir_generator
|
||||
if options.gen_ts_lowerings:
|
||||
lazy_ir_generator = GenTSLazyIR
|
||||
native_func_definition_generator: type[GenLazyNativeFuncDefinition] = (
|
||||
default_args.native_func_definition_generator
|
||||
)
|
||||
|
||||
run_gen_lazy_tensor(
|
||||
aten_path,
|
||||
options.source_yaml,
|
||||
options.output_dir,
|
||||
options.dry_run,
|
||||
options.impl_path,
|
||||
options.node_base,
|
||||
options.node_base_hdr,
|
||||
options.tensor_class,
|
||||
options.tensor_class_hdr,
|
||||
options.shape_inference_hdr,
|
||||
lazy_ir_generator,
|
||||
native_func_definition_generator,
|
||||
options.backend_name,
|
||||
)
|
||||
|
||||
|
||||
def run_gen_lazy_tensor(
|
||||
aten_path: str,
|
||||
source_yaml: str,
|
||||
output_dir: str,
|
||||
dry_run: bool,
|
||||
impl_path: str | None,
|
||||
node_base: str = default_args.node_base,
|
||||
node_base_hdr: str | None = default_args.node_base_hdr,
|
||||
tensor_class: str = default_args.tensor_class,
|
||||
tensor_class_hdr: str = default_args.tensor_class_hdr,
|
||||
shape_inference_hdr: str = default_args.shape_inference_hdr,
|
||||
lazy_ir_generator: type[GenLazyIR] = default_args.lazy_ir_generator,
|
||||
native_func_definition_generator: type[
|
||||
GenLazyNativeFuncDefinition
|
||||
] = default_args.native_func_definition_generator,
|
||||
# build_in_tree is true for TS backend and affects include paths
|
||||
build_in_tree: bool = False,
|
||||
# per_operator_headers changes whether ATen/Functions.h or individual operator headers are used
|
||||
# it must match how ATen was built
|
||||
per_operator_headers: bool = False,
|
||||
backend_name: str = default_args.backend_name,
|
||||
gen_forced_fallback_code: bool = False,
|
||||
use_lazy_shape: bool = True,
|
||||
# the following arguments are temporary customization points for xla backend migration.
|
||||
# do not rely on them otherwise, they should be removed once migration is complete
|
||||
backend_namespace: str = "torch::lazy",
|
||||
get_tensorlist: str = "GetTensorList",
|
||||
get_tensor_or_wrap_number: str = "GetLtcTensorOrCreateForWrappedNumber",
|
||||
try_get_tensor: str = "TryGetLtcTensor",
|
||||
metrics_counter: str = 'TORCH_LAZY_FN_COUNTER("lazy::")',
|
||||
create_tensor: str = "LazyTensor::Create",
|
||||
create_from_first_tensor: bool = False,
|
||||
create_aten_from_ltc_tensor: str = "torch::lazy::CreateAtenFromLtcTensor",
|
||||
tuple_aten_from_ltc_tensors: str = "torch::lazy::TupleAtenFromLtcTensors",
|
||||
lazy_value_class: str = "torch::lazy::Value",
|
||||
lazy_tensor_ptr: str = "LazyTensorPtr",
|
||||
get_device_fn: str = "torch::lazy::GetBackendDevice",
|
||||
) -> None:
|
||||
lv_tokens = lazy_value_class.split("::")
|
||||
lv_class = lv_tokens[-1]
|
||||
lv_ns = "::".join(lv_tokens[:-1])
|
||||
setValueT(BaseCppType(lv_ns, lv_class))
|
||||
template_dir = os.path.join(aten_path, "templates")
|
||||
|
||||
def make_file_manager(install_dir: str) -> FileManager:
|
||||
return FileManager(
|
||||
install_dir=install_dir, template_dir=template_dir, dry_run=dry_run
|
||||
)
|
||||
|
||||
fm = make_file_manager(output_dir)
|
||||
|
||||
native_yaml_path = os.path.join(aten_path, "native/native_functions.yaml")
|
||||
tags_yaml_path = os.path.join(aten_path, "native/tags.yaml")
|
||||
parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
|
||||
native_functions, backend_indices = (
|
||||
parsed_yaml.native_functions,
|
||||
parsed_yaml.backend_indices,
|
||||
)
|
||||
grouped_native_functions = get_grouped_native_functions(native_functions)
|
||||
|
||||
def sort_native_function(f: NativeFunctionsGroup | NativeFunction) -> str:
|
||||
"""
|
||||
We sort the native function because of the note in concat_map_codegen.
|
||||
TODO(alanwaketan): Remove this sorting hack once all ops are grouped properly.
|
||||
"""
|
||||
func = f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func
|
||||
return str(func.name.name)
|
||||
|
||||
grouped_native_functions = sorted(
|
||||
grouped_native_functions, key=sort_native_function
|
||||
)
|
||||
|
||||
parsed_backend_yaml = parse_backend_yaml(
|
||||
source_yaml, grouped_native_functions, backend_indices
|
||||
)
|
||||
backend_key = parsed_backend_yaml.backend_key
|
||||
autograd_key = parsed_backend_yaml.autograd_key
|
||||
cpp_namespace = parsed_backend_yaml.cpp_namespace
|
||||
backend_indices = parsed_backend_yaml.backend_indices
|
||||
# the following 3 keys are all processed differently
|
||||
# for full_codegen, we generate IR, kernels, etc
|
||||
# for ir_gen, we generate only IR
|
||||
# non_native is used to register kernels not declared in
|
||||
# native_functions.yaml
|
||||
full_codegen, non_native, ir_gen = parse_native_functions_keys(
|
||||
source_yaml, grouped_native_functions
|
||||
)
|
||||
|
||||
def concat_map_codegen(
|
||||
func: Callable[[NativeFunction], Sequence[str]],
|
||||
xs: Iterable[NativeFunctionsGroup | NativeFunction],
|
||||
ops_list: list[OperatorName] = full_codegen,
|
||||
) -> Iterator[str]:
|
||||
"""
|
||||
We code-gen for the functional variant, which is all we need for IR classes/lowerings/shape inferences, but we
|
||||
only code-gen additional entries for the inplace variant for the native functions.
|
||||
"""
|
||||
|
||||
for x in xs:
|
||||
fs = list(x.functions()) if isinstance(x, NativeFunctionsGroup) else [x]
|
||||
for f in fs:
|
||||
if f.func.name in ops_list:
|
||||
yield from func(f)
|
||||
|
||||
selector = SelectiveBuilder.get_nop_selector()
|
||||
|
||||
if backend_key is None:
|
||||
raise AssertionError("backend_key must be non-None")
|
||||
class_name = backend_indices[backend_key].native_function_class_name()
|
||||
|
||||
if impl_path is not None:
|
||||
error_on_missing_kernels(
|
||||
native_functions,
|
||||
backend_indices,
|
||||
backend_key,
|
||||
autograd_key,
|
||||
class_name,
|
||||
impl_path,
|
||||
full_codegen,
|
||||
)
|
||||
|
||||
""" Validate Shape Inference Definitions
|
||||
|
||||
Generated lazy native functions all perform shape inference, by first using a meta:: kernel
|
||||
if available for that op, and otherwise using a 'compute_shape_{op}' function instead. The generator
|
||||
knows the call signature for compute_shape_{op} because it matches the nativefunction (and meta::) signature,
|
||||
so it just has to check whether the op is structured and generate a call for one or the other. It's up to the dev
|
||||
to supply the missing compute_shape_{op} function, but the codegen at least warns you about this and provides
|
||||
the expected signature which can be copy-pasted into shape_inference.h.
|
||||
|
||||
compute_shape_{op} functions are handwritten and should be replaced over time as ops get ported
|
||||
to structured kernels.
|
||||
|
||||
See torch/csrc/lazy/core/shape_inference.cpp #READ THIS! for more information.
|
||||
"""
|
||||
if shape_inference_hdr is not None:
|
||||
expected_shape_infr_decls = list(
|
||||
concat_map_codegen(
|
||||
dest.GenLazyShapeInferenceDefinition(
|
||||
backend_indices[backend_key], tensor_class
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
|
||||
validate_shape_inference_header(shape_inference_hdr, expected_shape_infr_decls)
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must be non-None")
|
||||
|
||||
# Generate nativefunction declarations
|
||||
# Note, eager registrations is set to False for the lazy TS backend as another LTC backend
|
||||
# may want to register their own lazy kernels instead of registering the TS ones.
|
||||
# The registration will lazily happen when init_ts_backend is called.
|
||||
gen_dispatchkey_nativefunc_headers(
|
||||
fm,
|
||||
class_name,
|
||||
cpp_namespace,
|
||||
backend_indices,
|
||||
grouped_native_functions,
|
||||
backend_key,
|
||||
autograd_key,
|
||||
backend_name,
|
||||
)
|
||||
|
||||
# Generate Dispatcher registrations which hook up the nativefunctions
|
||||
for dispatch_key in (
|
||||
[backend_key] if autograd_key is None else [backend_key, autograd_key]
|
||||
):
|
||||
gen_dispatcher_registrations(
|
||||
fm,
|
||||
output_dir,
|
||||
class_name,
|
||||
backend_indices,
|
||||
grouped_native_functions,
|
||||
backend_key,
|
||||
dispatch_key,
|
||||
selector,
|
||||
build_in_tree=build_in_tree,
|
||||
per_operator_headers=per_operator_headers,
|
||||
backend_name=backend_name,
|
||||
eager_registration=False,
|
||||
)
|
||||
|
||||
# Generate native function impls that build IR nodes
|
||||
ns_helper = NamespaceHelper(cpp_namespace)
|
||||
fm.write_with_template(
|
||||
f"{backend_key}NativeFunctions.cpp",
|
||||
"DispatchKeyNativeFunctions.cpp",
|
||||
lambda: {
|
||||
"includes": [
|
||||
f"#include <{path}>"
|
||||
for path in [
|
||||
tensor_class_hdr,
|
||||
shape_inference_hdr,
|
||||
"ATen/Functions.h",
|
||||
"ATen/native/TensorConversions.h",
|
||||
"ATen/NativeFunctions.h",
|
||||
"ATen/CompositeExplicitAutogradNonFunctionalFunctions.h",
|
||||
"ATen/MetaFunctions.h",
|
||||
"ATen/Operators.h",
|
||||
"ATen/native/CPUFallback.h",
|
||||
"torch/csrc/lazy/core/ir_builder.h",
|
||||
"torch/csrc/lazy/core/lazy_graph_executor.h",
|
||||
"torch/csrc/lazy/core/metrics.h",
|
||||
"torch/csrc/lazy/core/shape.h",
|
||||
f"{output_dir}/{backend_key}NativeFunctions.h",
|
||||
f"{output_dir}/LazyIr.h",
|
||||
]
|
||||
+ (
|
||||
["torch/csrc/lazy/ts_backend/ts_eager_fallback.h"]
|
||||
if gen_forced_fallback_code
|
||||
else []
|
||||
)
|
||||
],
|
||||
"helper_fns": get_ltc_helper_fns(),
|
||||
"native_functions_include": "",
|
||||
"namespace_prologue": ns_helper.prologue,
|
||||
"namespace_epilogue": ns_helper.epilogue,
|
||||
"native_function_definitions": list(
|
||||
concat_map_codegen(
|
||||
native_func_definition_generator(
|
||||
f"{backend_key}NativeFunctions",
|
||||
backend_indices[backend_key],
|
||||
tensor_class,
|
||||
gen_forced_fallback_code,
|
||||
backend_namespace,
|
||||
get_tensorlist,
|
||||
get_tensor_or_wrap_number,
|
||||
try_get_tensor,
|
||||
metrics_counter,
|
||||
create_tensor,
|
||||
create_from_first_tensor,
|
||||
create_aten_from_ltc_tensor,
|
||||
tuple_aten_from_ltc_tensors,
|
||||
lazy_tensor_ptr,
|
||||
get_device_fn,
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
),
|
||||
},
|
||||
)
|
||||
# Generate IR node classes
|
||||
lazy_ir_obj = lazy_ir_generator(
|
||||
backend_indices[backend_key], backend_name, node_base, use_lazy_shape
|
||||
)
|
||||
|
||||
fm.write_with_template(
|
||||
"LazyIr.h",
|
||||
"LazyIr.h",
|
||||
lambda: {
|
||||
"lazy_ir_sysinc": [
|
||||
f"#include <{path}>"
|
||||
for path in [
|
||||
"ATen/core/Formatting.h",
|
||||
"c10/core/ScalarType.h",
|
||||
"torch/csrc/lazy/core/hash.h",
|
||||
"torch/csrc/lazy/core/ir.h",
|
||||
"torch/csrc/lazy/core/shape.h",
|
||||
"optional",
|
||||
"vector",
|
||||
]
|
||||
],
|
||||
"lazy_ir_inc": [f'#include "{node_base_hdr}"']
|
||||
if node_base_hdr is not None
|
||||
else [],
|
||||
"ir_declarations": list(
|
||||
concat_map_codegen(
|
||||
lazy_ir_obj, grouped_native_functions, full_codegen + ir_gen
|
||||
)
|
||||
),
|
||||
"namespace_prologue": ns_helper.prologue,
|
||||
"namespace_epilogue": ns_helper.epilogue,
|
||||
},
|
||||
)
|
||||
|
||||
# Generate Non Native IR Node classes
|
||||
fm.write_with_template(
|
||||
"LazyNonNativeIr.h",
|
||||
"LazyNonNativeIr.h",
|
||||
lambda: {
|
||||
"lazy_non_native_ir_inc": [
|
||||
f"#include <{path}>"
|
||||
for path in [
|
||||
"torch/csrc/lazy/core/ir.h",
|
||||
"torch/csrc/lazy/core/ir_builder.h",
|
||||
"torch/csrc/lazy/core/internal_ops/ltc_ops.h",
|
||||
"torch/csrc/lazy/core/shape_inference.h",
|
||||
]
|
||||
+ ([node_base_hdr] if node_base_hdr else [])
|
||||
if path
|
||||
],
|
||||
"non_native_ir_nodes": dest.generate_non_native_lazy_ir_nodes(
|
||||
non_native, lazy_ir_obj
|
||||
),
|
||||
"namespace_prologue": ns_helper.prologue,
|
||||
"namespace_epilogue": ns_helper.epilogue,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,98 @@
|
||||
from typing import Any
|
||||
|
||||
from torchgen.model import (
|
||||
Annotation,
|
||||
Argument,
|
||||
Arguments,
|
||||
BaseOperatorName,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
CustomClassType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
OperatorName,
|
||||
Return,
|
||||
)
|
||||
|
||||
|
||||
# Note: These aren't actually used in torchgen, they're some utilities for generating a schema
|
||||
# from real arguments. For example, this is used to generate HigherOrderOperators' schema since
|
||||
# their schemas can vary for different instances of the same HOP.
|
||||
|
||||
|
||||
class TypeGen:
|
||||
convert_to_base_ty = {
|
||||
int: BaseTy.int,
|
||||
float: BaseTy.float,
|
||||
str: BaseTy.str,
|
||||
bool: BaseTy.bool,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_example(obj: Any) -> BaseType | ListType | CustomClassType:
|
||||
import torch
|
||||
|
||||
if isinstance(obj, torch.fx.GraphModule):
|
||||
return BaseType(BaseTy.GraphModule)
|
||||
elif isinstance(obj, torch.Tensor):
|
||||
return BaseType(BaseTy.Tensor)
|
||||
elif isinstance(obj, torch.SymInt):
|
||||
return BaseType(BaseTy.SymInt)
|
||||
elif isinstance(obj, torch.SymBool):
|
||||
return BaseType(BaseTy.SymBool)
|
||||
elif isinstance(obj, torch.ScriptObject):
|
||||
return CustomClassType(obj._type().name()) # type: ignore[attr-defined]
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
if len(obj) == 0:
|
||||
raise AssertionError("list/tuple must be non-empty")
|
||||
all_base_tys = [TypeGen.from_example(x) for x in obj]
|
||||
if len(set(all_base_tys)) > 1:
|
||||
raise RuntimeError(
|
||||
f"Cannot generate schema for a sequence of args of heterogeneous types: {all_base_tys}. "
|
||||
"Consider unpacking the argument and give proper names to them if possible "
|
||||
"instead of using *args."
|
||||
)
|
||||
return ListType(all_base_tys[0], len(obj))
|
||||
tp = type(obj)
|
||||
if tp not in TypeGen.convert_to_base_ty:
|
||||
raise RuntimeError(f"unsupported type {tp}")
|
||||
return BaseType(TypeGen.convert_to_base_ty[tp])
|
||||
|
||||
|
||||
class ReturnGen:
|
||||
@staticmethod
|
||||
def from_example(
|
||||
name: str | None, obj: Any, annotation: Annotation | None
|
||||
) -> Return:
|
||||
return Return(name, TypeGen.from_example(obj), annotation)
|
||||
|
||||
|
||||
class ArgumentGen:
|
||||
@staticmethod
|
||||
def from_example(
|
||||
name: str, obj: Any, default: str | None, annotation: Annotation | None
|
||||
) -> Argument:
|
||||
return Argument(
|
||||
name, TypeGen.from_example(obj), default=default, annotation=annotation
|
||||
)
|
||||
|
||||
|
||||
class FunctionSchemaGen:
|
||||
@staticmethod
|
||||
def from_example(
|
||||
op_name: str,
|
||||
example_inputs: tuple[tuple[str, Any], ...],
|
||||
example_outputs: tuple[Any, ...],
|
||||
) -> FunctionSchema:
|
||||
args = []
|
||||
for name, inp in example_inputs:
|
||||
args.append(ArgumentGen.from_example(name, inp, None, None))
|
||||
# ignore the annotations and other attributes for now, we could add more when needed.
|
||||
arguments = Arguments(
|
||||
tuple(), None, tuple(args), tuple(), None, tuple(), tuple()
|
||||
)
|
||||
returns = tuple(
|
||||
ReturnGen.from_example(None, out, None) for out in example_outputs
|
||||
)
|
||||
op_name = OperatorName(BaseOperatorName(op_name, False, False, False), "")
|
||||
return FunctionSchema(op_name, arguments, returns)
|
||||
@@ -0,0 +1,277 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import textwrap
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from torchgen.api.translate import translate
|
||||
from torchgen.api.types import DispatcherSignature
|
||||
from torchgen.context import method_with_native_function
|
||||
from torchgen.model import (
|
||||
Argument,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
FunctionSchema,
|
||||
ListType,
|
||||
NativeFunction,
|
||||
OptionalType,
|
||||
Return,
|
||||
SchemaKind,
|
||||
Type,
|
||||
)
|
||||
from torchgen.utils import mapMaybe
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
def is_tensor(typ: Type) -> bool:
|
||||
return isinstance(typ, BaseType) and typ.name == BaseTy.Tensor
|
||||
|
||||
|
||||
def is_optional_tensor(typ: Type) -> bool:
|
||||
return isinstance(typ, OptionalType) and is_tensor(typ.elem)
|
||||
|
||||
|
||||
def is_tensor_list(typ: Type) -> bool:
|
||||
return isinstance(typ, ListType) and is_tensor(typ.elem)
|
||||
|
||||
|
||||
def unwrap_tensor(name: str, cur_level_var: str) -> list[str]:
|
||||
result = f"""\
|
||||
auto [{name}_value, {name}_bdim] = unwrapTensorAtLevel({name}, {cur_level_var});"""
|
||||
return textwrap.dedent(result).split("\n")
|
||||
|
||||
|
||||
def unwrap_optional_tensor(name: str, cur_level_var: str) -> list[str]:
|
||||
result = f"""\
|
||||
std::optional<Tensor> {name}_value;
|
||||
std::optional<int64_t> {name}_bdim;
|
||||
if ({name}) {{
|
||||
std::tie({name}_value, {name}_bdim) = unwrapTensorAtLevel({name}.value(), {cur_level_var});
|
||||
}}"""
|
||||
return textwrap.dedent(result).split("\n")
|
||||
|
||||
|
||||
def gen_unwraps(
|
||||
flat_arguments: Sequence[Argument], cur_level_var: str
|
||||
) -> tuple[str, list[str]]:
|
||||
arg_names = [a.name for a in flat_arguments]
|
||||
arg_types = [a.type for a in flat_arguments]
|
||||
|
||||
tensors = [name for typ, name in zip(arg_types, arg_names) if is_tensor(typ)]
|
||||
optional_tensors = [
|
||||
name for typ, name in zip(arg_types, arg_names) if is_optional_tensor(typ)
|
||||
]
|
||||
|
||||
unwraps = []
|
||||
for tensor in tensors:
|
||||
unwraps += unwrap_tensor(tensor, cur_level_var)
|
||||
|
||||
for opt_tensor in optional_tensors:
|
||||
unwraps += unwrap_optional_tensor(opt_tensor, cur_level_var)
|
||||
unwrap_code = "\n".join(unwraps)
|
||||
|
||||
unwrapped_arg_list = []
|
||||
for arg in arg_names:
|
||||
if arg in tensors or arg in optional_tensors:
|
||||
unwrapped_arg_list += [f"{arg}_value", f"{arg}_bdim"]
|
||||
else:
|
||||
unwrapped_arg_list.append(arg)
|
||||
return unwrap_code, unwrapped_arg_list
|
||||
|
||||
|
||||
def gen_case_where_all_bdims_are_none(
|
||||
outer_sig: DispatcherSignature, schema: FunctionSchema, cur_level_var: str
|
||||
) -> str:
|
||||
conditions = []
|
||||
flat_args = schema.arguments.flat_all
|
||||
for arg in flat_args:
|
||||
if not arg.type.is_tensor_like():
|
||||
continue
|
||||
conditions.append(f"!isBatchedAtLevel({arg.name}, {cur_level_var})")
|
||||
|
||||
sig = DispatcherSignature.from_schema(schema)
|
||||
translated_args = ", ".join(
|
||||
e.expr for e in translate(outer_sig.arguments(), sig.arguments())
|
||||
)
|
||||
return f"""\
|
||||
if ({" && ".join(conditions)}) {{
|
||||
return at::_ops::{sig.func.name.unambiguous_name()}::call({translated_args});
|
||||
}}"""
|
||||
|
||||
|
||||
def gen_returns(
|
||||
returns: tuple[Return, ...], cur_level_var: str, results_var: str
|
||||
) -> str:
|
||||
idx = 0
|
||||
wrapped_returns = []
|
||||
for ret in returns:
|
||||
if is_tensor(ret.type):
|
||||
wrapped_returns.append(
|
||||
f"makeBatched(std::get<{idx}>({results_var}), std::get<{idx + 1}>({results_var}), {cur_level_var})"
|
||||
)
|
||||
idx += 2
|
||||
elif is_tensor_list(ret.type):
|
||||
wrapped_returns.append(
|
||||
f"makeBatchedVector(std::get<{idx}>({results_var}), std::get<{idx + 1}>({results_var}), {cur_level_var})"
|
||||
)
|
||||
idx += 2
|
||||
else:
|
||||
wrapped_returns.append(f"std::get<{idx}>({results_var})")
|
||||
idx += 1
|
||||
if len(wrapped_returns) == 1:
|
||||
result = f"return {wrapped_returns[0]};"
|
||||
else:
|
||||
result = f"return std::make_tuple({', '.join(wrapped_returns)});"
|
||||
return result
|
||||
|
||||
|
||||
def accepts_at_least_one_tensor_input(schema: FunctionSchema) -> bool:
|
||||
return any(a.type.is_tensor_like() for a in schema.arguments.flat_all)
|
||||
|
||||
|
||||
def is_mutated_arg(argument: Argument) -> bool:
|
||||
return argument.annotation is not None and argument.annotation.is_write
|
||||
|
||||
|
||||
def gen_vmap_inplace_plumbing(native_function: NativeFunction) -> str | None:
|
||||
# Assumptions:
|
||||
# - only one argument is being modified in-place
|
||||
# - the argument that is being modified in-place is the first argument
|
||||
# - all returns are either Tensor, tuple of Tensor, or TensorList
|
||||
schema = native_function.func
|
||||
sig = DispatcherSignature.from_schema(schema)
|
||||
returns = schema.returns
|
||||
|
||||
# Check assumptions. If these are invalid we return None
|
||||
# and punt the work to handle them to the future.
|
||||
if schema.kind() != SchemaKind.inplace:
|
||||
raise AssertionError(f"Expected inplace schema, got {schema.kind()}")
|
||||
if not is_mutated_arg(schema.arguments.flat_all[0]):
|
||||
return None
|
||||
if len([arg for arg in schema.arguments.flat_all if is_mutated_arg(arg)]) != 1:
|
||||
return None
|
||||
|
||||
# Only support cases where all returns are Tensors or vector<Tensor>
|
||||
if len(returns) == 0:
|
||||
return None
|
||||
if not all(is_tensor(ret.type) or is_tensor_list(ret.type) for ret in returns):
|
||||
return None
|
||||
if not accepts_at_least_one_tensor_input(schema):
|
||||
return None
|
||||
|
||||
cur_level_var = "cur_level"
|
||||
|
||||
unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)
|
||||
bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)
|
||||
|
||||
return f"""\
|
||||
template <typename batch_rule_t, batch_rule_t batch_rule>
|
||||
{sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{
|
||||
c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
|
||||
auto maybe_layer = maybeCurrentDynamicLayer();
|
||||
vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
|
||||
int64_t {cur_level_var} = maybe_layer->layerId();
|
||||
{textwrap.indent(bdims_all_none_case, " ")}
|
||||
{textwrap.indent(unwraps, " ")}
|
||||
batch_rule({", ".join(unwrapped_arg_list)});
|
||||
return {schema.arguments.flat_all[0].name};
|
||||
}}"""
|
||||
|
||||
|
||||
def gen_vmap_plumbing_no_returns(native_function: NativeFunction) -> str:
|
||||
schema = native_function.func
|
||||
sig = DispatcherSignature.from_schema(schema)
|
||||
cur_level_var = "cur_level"
|
||||
|
||||
unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)
|
||||
bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)
|
||||
|
||||
return f"""\
|
||||
template <typename batch_rule_t, batch_rule_t batch_rule>
|
||||
{sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{
|
||||
c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
|
||||
auto maybe_layer = maybeCurrentDynamicLayer();
|
||||
vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
|
||||
int64_t {cur_level_var} = maybe_layer->layerId();
|
||||
{textwrap.indent(bdims_all_none_case, " ")}
|
||||
{textwrap.indent(unwraps, " ")}
|
||||
batch_rule({", ".join(unwrapped_arg_list)});
|
||||
}}"""
|
||||
|
||||
|
||||
def gen_vmap_plumbing(native_function: NativeFunction) -> str | None:
|
||||
schema = native_function.func
|
||||
sig = DispatcherSignature.from_schema(schema)
|
||||
returns = schema.returns
|
||||
|
||||
# Only support cases where all returns are Tensors or vector<Tensor>
|
||||
if not accepts_at_least_one_tensor_input(schema):
|
||||
return None
|
||||
if len(returns) == 0:
|
||||
return gen_vmap_plumbing_no_returns(native_function)
|
||||
return_symint_overrides = [
|
||||
"_scaled_dot_product_flash_attention",
|
||||
"_scaled_dot_product_cudnn_attention",
|
||||
"_scaled_dot_product_flash_attention_quantized",
|
||||
]
|
||||
if (
|
||||
not all(ret.type.is_tensor_like() for ret in returns)
|
||||
and schema.name.unambiguous_name() not in return_symint_overrides
|
||||
):
|
||||
return None
|
||||
# in-place views need special handling
|
||||
if "inplace_view" in native_function.tags:
|
||||
return None
|
||||
|
||||
if schema.kind() == SchemaKind.inplace:
|
||||
return gen_vmap_inplace_plumbing(native_function)
|
||||
|
||||
# Don't support these (mutable, out, scratch)
|
||||
if schema.kind() != SchemaKind.functional:
|
||||
return None
|
||||
|
||||
results_var = "results"
|
||||
cur_level_var = "cur_level"
|
||||
|
||||
unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)
|
||||
bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)
|
||||
|
||||
wrapped_returns = gen_returns(returns, cur_level_var, results_var)
|
||||
return f"""\
|
||||
template <typename batch_rule_t, batch_rule_t batch_rule>
|
||||
{sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{
|
||||
c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
|
||||
auto maybe_layer = maybeCurrentDynamicLayer();
|
||||
vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
|
||||
int64_t {cur_level_var} = maybe_layer->layerId();
|
||||
{textwrap.indent(bdims_all_none_case, " ")}
|
||||
{textwrap.indent(unwraps, " ")}
|
||||
auto {results_var} = batch_rule({", ".join(unwrapped_arg_list)});
|
||||
{wrapped_returns}
|
||||
}}"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ComputeBatchRulePlumbing:
|
||||
@method_with_native_function
|
||||
def __call__(self, f: NativeFunction) -> str | None:
|
||||
result = gen_vmap_plumbing(f)
|
||||
return result
|
||||
|
||||
|
||||
def gen_all_vmap_plumbing(native_functions: Sequence[NativeFunction]) -> str:
|
||||
body = "\n".join(list(mapMaybe(ComputeBatchRulePlumbing(), native_functions)))
|
||||
return f"""
|
||||
#pragma once
|
||||
#include <ATen/Operators.h>
|
||||
#include <ATen/functorch/PlumbingHelper.h>
|
||||
|
||||
namespace at {{ namespace functorch {{
|
||||
|
||||
{body}
|
||||
|
||||
}}}} // namespace at::functorch
|
||||
"""
|
||||
@@ -0,0 +1,64 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
|
||||
# Simple dynamic scoping implementation. The name "parametrize" comes
|
||||
# from Racket.
|
||||
#
|
||||
# WARNING WARNING: LOOKING TO EDIT THIS FILE? Think carefully about
|
||||
# why you need to add a toggle to the global behavior of code
|
||||
# generation. The parameters here should really only be used
|
||||
# for "temporary" situations, where we need to temporarily change
|
||||
# the codegen in some cases because we cannot conveniently update
|
||||
# all call sites, and are slated to be eliminated once all call
|
||||
# sites are eliminated. If you don't have a plan for how to get there,
|
||||
# DON'T add a new entry here.
|
||||
|
||||
|
||||
class Locals(threading.local):
|
||||
use_const_ref_for_mutable_tensors: bool | None = None
|
||||
use_ilistref_for_tensor_lists: bool | None = None
|
||||
|
||||
|
||||
_locals = Locals()
|
||||
|
||||
|
||||
def use_const_ref_for_mutable_tensors() -> bool:
|
||||
if _locals.use_const_ref_for_mutable_tensors is None:
|
||||
raise AssertionError(
|
||||
"need to initialize local.use_const_ref_for_mutable_tensors with "
|
||||
"local.parametrize"
|
||||
)
|
||||
return _locals.use_const_ref_for_mutable_tensors
|
||||
|
||||
|
||||
def use_ilistref_for_tensor_lists() -> bool:
|
||||
if _locals.use_ilistref_for_tensor_lists is None:
|
||||
raise AssertionError(
|
||||
"need to initialize local.use_ilistref_for_tensor_lists with local.parametrize"
|
||||
)
|
||||
return _locals.use_ilistref_for_tensor_lists
|
||||
|
||||
|
||||
@contextmanager
|
||||
def parametrize(
|
||||
*, use_const_ref_for_mutable_tensors: bool, use_ilistref_for_tensor_lists: bool
|
||||
) -> Iterator[None]:
|
||||
old_use_const_ref_for_mutable_tensors = _locals.use_const_ref_for_mutable_tensors
|
||||
old_use_ilistref_for_tensor_lists = _locals.use_ilistref_for_tensor_lists
|
||||
try:
|
||||
_locals.use_const_ref_for_mutable_tensors = use_const_ref_for_mutable_tensors
|
||||
_locals.use_ilistref_for_tensor_lists = use_ilistref_for_tensor_lists
|
||||
yield
|
||||
finally:
|
||||
_locals.use_const_ref_for_mutable_tensors = (
|
||||
old_use_const_ref_for_mutable_tensors
|
||||
)
|
||||
_locals.use_ilistref_for_tensor_lists = old_use_ilistref_for_tensor_lists
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,662 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import string
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torchgen.api.dispatcher as dispatcher
|
||||
from torchgen.api.translate import translate
|
||||
from torchgen.api.types import Binding, DispatcherSignature, Expr
|
||||
from torchgen.context import with_native_function
|
||||
from torchgen.model import (
|
||||
Annotation,
|
||||
Argument,
|
||||
BackendIndex,
|
||||
BackendMetadata,
|
||||
BaseOperatorName,
|
||||
BaseTy,
|
||||
BaseType,
|
||||
DEFAULT_KERNEL_NAMESPACE,
|
||||
DeviceCheckType,
|
||||
DispatchKey,
|
||||
FunctionSchema,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
OperatorName,
|
||||
Return,
|
||||
SchemaKind,
|
||||
Variant,
|
||||
)
|
||||
from torchgen.utils import concatMap
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# See Note: [Out ops with functional variants that don't get grouped properly]
|
||||
OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY = [
|
||||
# This has a functional variant, but it's currently marked private.
|
||||
# This function should be marked private as well (*_backward ops aren't exposed to python anyway).
|
||||
"adaptive_avg_pool3d_backward.grad_input",
|
||||
# There's a functional variant, _slow_conv2d_backward.output_mask, that isn't grouped properly.
|
||||
# Maybe we can kill this operator in favor of convolution_backward?
|
||||
"_slow_conv2d_backward.grad_input",
|
||||
]
|
||||
|
||||
|
||||
# See Note: [Mutable ops that cannot get an out variant]
|
||||
MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT = [
|
||||
# should be out=?
|
||||
"_cummax_helper",
|
||||
# should be out=?
|
||||
"_cummin_helper",
|
||||
]
|
||||
|
||||
# All of these operators don't have any tensor like returns
|
||||
FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT = [
|
||||
"_assert_async", # no return
|
||||
"_assert_async.msg", # no return
|
||||
"_assert_tensor_metadata", # no return
|
||||
"_cslt_sparse_mm_search", # returns an int
|
||||
"_assert_scalar", # no return
|
||||
"_dimI", # returns an int
|
||||
"_dimV", # returns an int
|
||||
"_has_same_storage_numel", # returns a boolean
|
||||
"_linalg_check_errors", # no return
|
||||
"_local_scalar_dense", # returns a Scalar
|
||||
"_nested_tensor_from_mask_left_aligned", # returns a boolean
|
||||
"_nnz", # returns an int
|
||||
"_use_cudnn_ctc_loss", # returns a boolean
|
||||
"_use_cudnn_ctc_loss.Tensor", # returns a boolean
|
||||
"_use_miopen_ctc_loss", # returns a boolean
|
||||
"_use_miopen_ctc_loss.Tensor", # returns a boolean
|
||||
"_validate_compressed_sparse_indices", # no return
|
||||
"allclose", # returns a boolean
|
||||
"dense_dim", # returns an int
|
||||
"equal", # returns a boolean
|
||||
"is_coalesced", # returns an boolean
|
||||
"is_pinned", # returns a boolean
|
||||
"is_same_size", # returns a boolean
|
||||
"is_set_to", # returns a boolean
|
||||
"q_per_channel_axis", # returns an int
|
||||
"q_scale", # returns a float
|
||||
"q_zero_point", # returns an int
|
||||
"qscheme", # returns a QScheme
|
||||
"record_stream", # no return
|
||||
"sparse_dim", # returns an int
|
||||
"sym_constrain_range", # no return
|
||||
"sym_constrain_range_for_size", # no return
|
||||
"_nested_tensor_storage_offsets", # returns a vector of ints
|
||||
"_chunk_grad_outputs_efficient_attention", # returns a bool
|
||||
"_fused_sdp_choice", # returns an int
|
||||
"_print", # no return
|
||||
"_sink_tokens", # no return
|
||||
"_nested_get_ragged_idx", # returns an int
|
||||
]
|
||||
|
||||
INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY = [
|
||||
# polygamma and polygamma.out both exist, but have a
|
||||
# pre-self arg (while polygamma_ does not)
|
||||
# We should either fix this schema so it can be grouped properly,
|
||||
# or allow the codegen to generate new functional/out= NativeFunctions for this op
|
||||
# (which would require changing its overload name to prevent overload ambiguity).
|
||||
"polygamma_"
|
||||
]
|
||||
|
||||
|
||||
# Groups "similar" NativeFunctions together
|
||||
# example add.Tensor, add_.Tensor, add.out
|
||||
# "similar" NativeFunctions are all expected to have an identical `signature()`,
|
||||
# But have differing SchemaKinds.
|
||||
def pre_group_native_functions(
|
||||
native_functions: Sequence[NativeFunction],
|
||||
) -> dict[FunctionSchema, dict[SchemaKind, NativeFunction]]:
|
||||
pre_grouped_native_functions: dict[
|
||||
FunctionSchema, dict[SchemaKind, NativeFunction]
|
||||
] = defaultdict(dict)
|
||||
for f in native_functions:
|
||||
d = pre_grouped_native_functions[f.func.signature()]
|
||||
if f.func.kind() in d:
|
||||
raise AssertionError(f"Duplicate schema kind {f.func.kind()} for {f.func}")
|
||||
d[f.func.kind()] = f
|
||||
return pre_grouped_native_functions
|
||||
|
||||
|
||||
# Returns the out variant overload name given a base function overload name
|
||||
def get_expected_out_variant_overload_name(overload_name: str | None) -> str:
|
||||
return "out" if not overload_name else f"{overload_name}_out"
|
||||
|
||||
|
||||
# Helper function: given an inplace FunctionSchema, generate its corresponding out= variant
|
||||
# Example before:
|
||||
# _add_relu_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)
|
||||
# Example after:
|
||||
# _add_relu.Scalar_out(Tensor self, Scalar other, Scalar alpha=1, *, Tensor(a!) out)
|
||||
def self_to_out_signature(func: FunctionSchema) -> FunctionSchema:
|
||||
# Generating an out= schema from an inplace schema.
|
||||
if func.kind() != SchemaKind.inplace:
|
||||
raise AssertionError(f"Expected inplace schema, got {func.kind()}")
|
||||
if func.arguments.self_arg is None:
|
||||
raise AssertionError("Expected self_arg to be non-None")
|
||||
# The new out= schema has:
|
||||
# - a new out argument with the same type as "func" (but with a mutable annotation)
|
||||
# - The returns (if any) now alias the out= argument instead of "func"
|
||||
# - an "out" overload name
|
||||
return FunctionSchema(
|
||||
name=func.name.remove_inplace().with_overload(
|
||||
get_expected_out_variant_overload_name(func.name.overload_name)
|
||||
),
|
||||
arguments=func.arguments.remove_self_annotation().with_out_args(
|
||||
[
|
||||
Argument(
|
||||
name="out",
|
||||
type=func.arguments.self_arg.argument.type,
|
||||
default=None,
|
||||
annotation=func.arguments.self_arg.argument.annotation,
|
||||
)
|
||||
]
|
||||
),
|
||||
returns=func.returns,
|
||||
)
|
||||
|
||||
|
||||
# Helper function: given a functional FunctionSchema, generate its corresponding out= variant
|
||||
# Example before:
|
||||
# _to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None,
|
||||
# bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor
|
||||
# Example after:
|
||||
# _to_copy._out(Tensor self, *, bool non_blocking=False, MemoryFormat? memory_format=None,
|
||||
# Tensor(a!) out) -> Tensor(a!)
|
||||
def functional_to_out_signature(func: FunctionSchema) -> FunctionSchema:
|
||||
# Generating an out= schema from a functional schema.
|
||||
if func.kind() != SchemaKind.functional:
|
||||
raise AssertionError(f"Expected functional schema, got {func.kind()}")
|
||||
|
||||
new_returns, new_out_args = generate_out_args_from_schema(func)
|
||||
# The new out= schema has:
|
||||
# - one or more new out argument(s) with the same type as returns (but with a mutable annotation)
|
||||
# - The returns now alias the out= arguments
|
||||
# - an "_out" overload name
|
||||
return FunctionSchema(
|
||||
name=func.name.with_overload(
|
||||
get_expected_out_variant_overload_name(func.name.overload_name)
|
||||
),
|
||||
arguments=func.arguments.signature().with_out_args(
|
||||
new_out_args,
|
||||
),
|
||||
returns=tuple(new_returns),
|
||||
)
|
||||
|
||||
|
||||
# Helper function: given a function schema, generate corresponding out arguments, also the updated return annotations.
|
||||
def generate_out_args_from_schema(
|
||||
func: FunctionSchema,
|
||||
) -> tuple[list[Return], list[Argument]]:
|
||||
# More of a sanity check - our existing restrictions on schemas should enforce that
|
||||
# mutable schema kinds never return their mutable arguments.
|
||||
if any(r.annotation is not None and r.annotation.is_write for r in func.returns):
|
||||
raise AssertionError("Mutable schema kinds should not return mutable arguments")
|
||||
|
||||
tensorlike_rets = [r for r in func.returns if r.type.is_tensor_like()]
|
||||
if len(tensorlike_rets) == 0:
|
||||
raise AssertionError("Expected at least one tensor-like return")
|
||||
|
||||
used_annotations = concatMap(
|
||||
lambda a: [] if a.annotation is None else a.annotation.alias_set,
|
||||
func.arguments.flat_all,
|
||||
)
|
||||
valid_annotations = [x for x in string.ascii_lowercase if x not in used_annotations]
|
||||
|
||||
all_rets_are_tensors = all(r.type == BaseType(BaseTy.Tensor) for r in func.returns)
|
||||
|
||||
new_out_args: list[Argument] = []
|
||||
# The end result of new_returns is that:
|
||||
# - If every return is a plain tensor, then the new returns == the old returns, but with the out= alias annotations added.
|
||||
# - Otherwise, none of the out arguments show up in the returns (and we're only left with non-tensor-like returns, if any).
|
||||
new_returns: list[Return] = []
|
||||
for i, r in enumerate(func.returns):
|
||||
if r.type.is_tensor_like():
|
||||
new_out = Argument(
|
||||
name="out" if len(func.returns) == 1 else f"out{i}",
|
||||
type=r.type,
|
||||
default=None,
|
||||
annotation=Annotation.parse(f"{valid_annotations[i]}!"),
|
||||
)
|
||||
new_out_args.append(new_out)
|
||||
if all_rets_are_tensors:
|
||||
# The convention for out= schemas is that they only return their out arguments
|
||||
# if the return is a plain Tensor (or if it's a tuple of plain Tensors)
|
||||
new_ret = Return(
|
||||
name=None, type=new_out.type, annotation=new_out.annotation
|
||||
)
|
||||
new_returns.append(new_ret)
|
||||
else:
|
||||
new_returns.append(r)
|
||||
return new_returns, new_out_args
|
||||
|
||||
|
||||
# Helper function: given a mutable FunctionSchema, generate its corresponding out= variant
|
||||
# Example before:
|
||||
# _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) # noqa: B950
|
||||
# Example after:
|
||||
# _fused_moving_avg_obs_fq_helper._out(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False, *, Tensor(e!) out0, Tensor(f!) out1) -> (Tensor(e!), Tensor(f!)) # noqa: B950
|
||||
def mutable_to_out_signature(func: FunctionSchema) -> FunctionSchema:
|
||||
# Generating an out= schema from a mutable schema.
|
||||
if func.kind() != SchemaKind.mutable:
|
||||
raise AssertionError(f"Expected mutable schema, got {func.kind()}")
|
||||
# The new out= schema has:
|
||||
# - Any non-aliased tensor-like returns are converted to mutable, aliased out= arguments
|
||||
# (if the argument is a tensor then we also return it for method chaining,
|
||||
# otherwise we return nothing)
|
||||
# - an "out" overload name
|
||||
#
|
||||
# Note that:
|
||||
# (1) This also means that we can *only* generate an out= variant from a mutable schema
|
||||
# if the mutable schema has at least one tensor-like non-aliasing return.
|
||||
# (2) The generated out= variant still has mutable positional arguments,
|
||||
# but if necessary we could probably add another out= variant that also
|
||||
# functionalizes the mutable arguments (a functional_out variant)
|
||||
|
||||
new_returns, new_out_args = generate_out_args_from_schema(func)
|
||||
|
||||
return FunctionSchema(
|
||||
name=func.name.remove_inplace().with_overload(
|
||||
get_expected_out_variant_overload_name(func.name.overload_name)
|
||||
),
|
||||
arguments=func.arguments.with_out_args(new_out_args),
|
||||
returns=tuple(new_returns),
|
||||
)
|
||||
|
||||
|
||||
# This function, given function of one SchemaKind, as well as a target SchemaKind,
|
||||
# generates a new NativeFunction with the same properties, but using the target SchemaKind.
|
||||
# We only actually generate functions for either functional or out= SchemaKinds.
|
||||
# This function returns a tuple, with:
|
||||
# - The generated NativeFunction
|
||||
# - a dictionary of `BackendIndex` objects, describing which dispatch keys
|
||||
# we will generate kernels for, for the new NativeFunction.
|
||||
# Details are in the function, but we only generate composite kernels (in some cases) today.
|
||||
def generate_function(
|
||||
f: NativeFunction, k: SchemaKind
|
||||
) -> tuple[NativeFunction, dict[DispatchKey, dict[OperatorName, BackendMetadata]]]:
|
||||
from torchgen.api import cpp
|
||||
|
||||
if k == SchemaKind.functional:
|
||||
if f.func.kind() == SchemaKind.functional:
|
||||
raise AssertionError("Cannot generate functional from functional schema")
|
||||
# The new "functional" NativeFunction has:
|
||||
# - any mutable arguments have been converted into (immutable) returns.
|
||||
# (if a mutable argument was not also a return, it gets converted to one)
|
||||
# - "_functional" appended to the base name, ONLY IF this op has a mutable variant.
|
||||
# See Note [Overload Ambiguity With Functional Variants]
|
||||
# The default grouping logic in signature() actually already does this,
|
||||
# so we can piggy-back off it (but we still want return names)
|
||||
func = f.func.signature(keep_return_names=True).with_name(
|
||||
OperatorName(
|
||||
name=BaseOperatorName(
|
||||
base=f.func.name.name.base,
|
||||
inplace=False,
|
||||
dunder_method=f.func.name.name.dunder_method,
|
||||
# See Note [Overload Ambiguity With Functional Variants]
|
||||
functional_overload=f.func.kind() == SchemaKind.mutable,
|
||||
),
|
||||
overload_name=f.func.name.overload_name,
|
||||
)
|
||||
)
|
||||
elif k == SchemaKind.out:
|
||||
# We generate out= ops mostly just so that we can pair up NativeFunctions into groups easily,
|
||||
# but at least today, there is no good reason to actually use them.
|
||||
# we'll generate a dispatcher entry for them, but won't actually register any kernels for them.
|
||||
if f.func.kind() == SchemaKind.inplace:
|
||||
func = self_to_out_signature(f.func)
|
||||
elif f.func.kind() == SchemaKind.mutable:
|
||||
func = mutable_to_out_signature(f.func)
|
||||
elif f.func.kind() == SchemaKind.functional:
|
||||
func = functional_to_out_signature(f.func)
|
||||
else:
|
||||
raise AssertionError(
|
||||
"We only bother generating out= functions from either inplace or mutable or functional variants"
|
||||
)
|
||||
else:
|
||||
raise AssertionError(
|
||||
"We currently only generate either functional or out= NativeFunctions"
|
||||
)
|
||||
|
||||
# Generated kernel naming convention for out: <op_name>_<overload_name>. The reason for this is to
|
||||
# disambiguate operator with the same name but different overload name, e.g., `randn.names_out` and
|
||||
# `randn.generator_with_names_out`.
|
||||
kernel_name = (
|
||||
func.name.unambiguous_name()
|
||||
if func.kind() == SchemaKind.out
|
||||
else cpp.name(func)
|
||||
)
|
||||
if f.func.has_symint():
|
||||
kernel_name += "_symint"
|
||||
backend_metadata = {
|
||||
DispatchKey.CompositeExplicitAutograd: {
|
||||
func.name: BackendMetadata(
|
||||
kernel=kernel_name,
|
||||
structured=False,
|
||||
cpp_namespace=DEFAULT_KERNEL_NAMESPACE,
|
||||
)
|
||||
}
|
||||
}
|
||||
tags = {"generated"} | set(
|
||||
f.tags & {"nondeterministic_seeded", "view_copy", "pt2_compliant_tag"}
|
||||
)
|
||||
|
||||
return (
|
||||
NativeFunction(
|
||||
func=func,
|
||||
use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors,
|
||||
# These generated fn's aren't meant to be user friendly- don't generate methods.
|
||||
variants={Variant.function},
|
||||
structured=False,
|
||||
structured_delegate=None,
|
||||
structured_inherits=None,
|
||||
precomputed=None,
|
||||
autogen=[],
|
||||
ufunc_inner_loop={},
|
||||
manual_kernel_registration=False,
|
||||
manual_cpp_binding=False,
|
||||
python_module=None,
|
||||
category_override=None,
|
||||
device_guard=False,
|
||||
device_check=DeviceCheckType.NoCheck,
|
||||
loc=f.loc,
|
||||
cpp_no_default_args=set(),
|
||||
is_abstract=f.is_abstract,
|
||||
has_composite_implicit_autograd_kernel=False,
|
||||
has_composite_implicit_autograd_nested_tensor_kernel=False,
|
||||
has_composite_explicit_autograd_kernel=True,
|
||||
has_composite_explicit_autograd_non_functional_kernel=False,
|
||||
# Every generated NativeFunction gets a "generated" tag, so it's easy to tell
|
||||
# which NativeFunction objects did not come directly from native_functions.yaml.
|
||||
tags=tags,
|
||||
namespace=f.namespace,
|
||||
),
|
||||
backend_metadata,
|
||||
)
|
||||
|
||||
|
||||
# This function is responsible for adding generated NativeFunctions which don't appear
|
||||
# explicitly in the codegen.
|
||||
# You can inspect the full list of NativeFunctions yourself with the torchgen package, by running
|
||||
# torchgen.parse_native_yaml("aten/src/ATen/native/native_functions.yaml", "aten/src/ATen/native/tags.yaml")
|
||||
# (Maybe we should make a friendly API for this)
|
||||
#
|
||||
# Note: this function *mutates* its two inputs,
|
||||
# adding the new NativeFunctions / BackendMetadata to them
|
||||
def add_generated_native_functions(
|
||||
rs: list[NativeFunction],
|
||||
indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]],
|
||||
) -> None:
|
||||
# The main code for generating new NativeFunctions
|
||||
# First we group of NativeFunctions by schema kind,
|
||||
# then we detect which ones are missing and generate them.
|
||||
pre_grouped_native_functions = pre_group_native_functions(rs)
|
||||
for d in pre_grouped_native_functions.values():
|
||||
has_functional = SchemaKind.functional in d
|
||||
has_inplace = SchemaKind.inplace in d
|
||||
has_mutable = SchemaKind.mutable in d
|
||||
has_out = SchemaKind.out in d
|
||||
is_core = any("core" in variant.tags for variant in d.values())
|
||||
|
||||
# We automatically generate a few native functions that don't exist in the yaml, for a few reasons:
|
||||
# (1) If an operator has an inplace/out= variant but no functional variant, we can generate
|
||||
# a simple functional variant that the functionalization pass can consume.
|
||||
# (2) If an operator has an inplace or functional but no out= variant, we generate an out=
|
||||
# variant, mostly so we can easily pair up functions into NativeFunctionsGroup,
|
||||
# while maintaining the constraint that the out= variant is "required".
|
||||
if has_mutable or has_inplace or has_out or has_functional:
|
||||
# Don't bother generating functions trio's for native functions that bypass the dispatcher.
|
||||
are_manual = all(f.manual_cpp_binding for f in d.values())
|
||||
# Don't bother generating functional + out= variants for view operators
|
||||
# set_ is technically an inplace_view, but for now it is treated
|
||||
# as a normal inplace op in the codegen
|
||||
has_view_ops = any(
|
||||
f.is_view_op and str(f.func.name.name) != "set_" for f in d.values()
|
||||
)
|
||||
# Don't generate the other variants for non-core CompositeImplicitAutograd operators.
|
||||
# We could probably do this, but the main benefit of generating the function triplets
|
||||
# is for transforms that need them, and transforms don't need to act directly
|
||||
# on CompositeImplicitAutograd operators (since we let them decompose).
|
||||
are_composite_implicit = all(
|
||||
f.has_composite_implicit_autograd_kernel for f in d.values()
|
||||
)
|
||||
if are_manual or has_view_ops or are_composite_implicit and not is_core:
|
||||
continue
|
||||
if has_out and len(d.values()) == 1:
|
||||
# Note: [Out ops with functional variants that don't get grouped properly]
|
||||
# In theory we could validly have an out= operator in native_functions.yaml
|
||||
# that has no other variants.
|
||||
# But today, all of the operators where that's the case actually do have
|
||||
# functional variants, that we are just unable to pair up properly.
|
||||
# I think banning this all together is probably safer
|
||||
# (you can always add a functional variant yourself if you want to add a new out= operator).
|
||||
#
|
||||
# We should probably fix the existing cases; this check is to prevent us from adding more over time.
|
||||
if (
|
||||
str(d[SchemaKind.out].func.name)
|
||||
not in OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Found an out= operator that we could not find any other variants of: {str(d[SchemaKind.out].func)}"
|
||||
)
|
||||
continue
|
||||
|
||||
# Some inplace ops that have problematic schemas (that we should fix), which prevent us
|
||||
# from generating out= and functional variants
|
||||
if (
|
||||
has_inplace
|
||||
and str(d[SchemaKind.inplace].func.name)
|
||||
in INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY
|
||||
):
|
||||
continue
|
||||
|
||||
base_fn = (
|
||||
d[SchemaKind.mutable]
|
||||
if has_mutable
|
||||
else d[SchemaKind.inplace]
|
||||
if has_inplace
|
||||
else d[SchemaKind.out]
|
||||
if has_out
|
||||
else d[SchemaKind.functional]
|
||||
)
|
||||
|
||||
# Note: [Mutable ops that cannot get an out variant]
|
||||
# We can only generate an out= variant if either:
|
||||
# - the original function has tensor-like returns (since we can convert them to out kwargs)
|
||||
# - or it's inplace (since we can convert `self` to an out kwarg)
|
||||
# There are only two functions that don't fit this criteria today though,
|
||||
# and they both look like they should be fixed to be out= variants,
|
||||
# so if feels safer to ban this schema all-together
|
||||
base_fn_valid = base_fn.func.kind() == SchemaKind.inplace or any(
|
||||
r.type.is_tensor_like() for r in base_fn.func.returns
|
||||
)
|
||||
# Note: [Loosen the assertion that all functional should have out variant]
|
||||
# By design all functional operators should have our variants. The needs_out check
|
||||
# is loosening this requirement, changing it to only generate out variant if there's
|
||||
# an `autogen` block in the native function, in the long run it should be removed.
|
||||
# FIXME: Remove this after figuring out CI job failures related to min, max, mean
|
||||
needs_out = any("out" in str(op_name) for op_name in base_fn.autogen)
|
||||
gets_out_variant = not has_out and base_fn_valid and needs_out
|
||||
if not has_out and not base_fn_valid:
|
||||
if (
|
||||
str(base_fn.func.name)
|
||||
not in MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT
|
||||
and str(base_fn.func.name)
|
||||
not in FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT
|
||||
):
|
||||
raise AssertionError(
|
||||
f"""Found an operator that we could not generate an out= variant for: {str(base_fn.func)}.
|
||||
This type of operators don't have tensor-like return, making it difficult to generate a proper out= variant. If
|
||||
out= variant is not needed, please add the function name into FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT list."""
|
||||
)
|
||||
|
||||
# Generate an out= variant
|
||||
if gets_out_variant:
|
||||
fn, metadata = generate_function(base_fn, SchemaKind.out)
|
||||
d[SchemaKind.out] = fn
|
||||
BackendIndex.grow_index(indices, metadata)
|
||||
rs.append(fn)
|
||||
|
||||
# Generate a functional variant, but only do it if the operator got an out= variant
|
||||
# (Functional variants are only useful if we can group up the variants,
|
||||
# which we can only do if they have an out= variant)
|
||||
if not has_functional and (has_out or gets_out_variant):
|
||||
fn, metadata = generate_function(base_fn, SchemaKind.functional)
|
||||
d[SchemaKind.functional] = fn
|
||||
BackendIndex.grow_index(indices, metadata)
|
||||
rs.append(fn)
|
||||
|
||||
|
||||
def return_str(rets: tuple[Return, ...], names: list[str]) -> str:
|
||||
if len(rets) != len(names):
|
||||
raise AssertionError(
|
||||
f"Returns and names length mismatch: {len(rets)} vs {len(names)}"
|
||||
)
|
||||
if len(rets) == 0:
|
||||
return ""
|
||||
elif len(rets) == 1:
|
||||
return f"return {names[0]};"
|
||||
else:
|
||||
return f"return {dispatcher.returns_type(rets).cpp_type()}({', '.join(names)});"
|
||||
|
||||
|
||||
# Given a function, and the name of a variable corresponding to the output of that function,
|
||||
# gather up all of the individual returns that are not aliased
|
||||
def gather_nonaliased_inner_rets(func: FunctionSchema, out_var: str) -> list[str]:
|
||||
aliased_rets = func.aliased_return_names()
|
||||
non_aliased_names = []
|
||||
is_out_var_a_tuple = len(func.returns) > 1
|
||||
for i, r in enumerate(aliased_rets):
|
||||
if r is None:
|
||||
non_aliased_names.append(
|
||||
f"std::get<{i}>({out_var})" if is_out_var_a_tuple else out_var
|
||||
)
|
||||
return non_aliased_names
|
||||
|
||||
|
||||
# Generates functional kernels in terms of their inplace.mutable counterparts.
|
||||
# We only do this for "generated" NativeFunctions
|
||||
@with_native_function
|
||||
def gen_composite_functional_kernel(g: NativeFunctionsGroup) -> str | None:
|
||||
# We should only be generating these for code-generated NativeFunctions
|
||||
if "generated" not in g.functional.tags:
|
||||
return None
|
||||
# And we always write the kernel for a generated op in terms of a non-generated op.
|
||||
if g.inplace is not None and "generated" not in g.inplace.tags:
|
||||
target_f = g.inplace
|
||||
elif g.mutable is not None and "generated" not in g.mutable.tags:
|
||||
target_f = g.mutable
|
||||
else:
|
||||
# We should be guaranteed to have a valid inplace/mutable variant to call into.
|
||||
# See Note: [Mutable Ops Not Using Functionalization]
|
||||
raise AssertionError(str(g.functional.func))
|
||||
|
||||
sig = DispatcherSignature(g.functional.func)
|
||||
target_sig = DispatcherSignature(target_f.func)
|
||||
|
||||
context: list[Binding | Expr] = []
|
||||
clone_mutable_inputs = []
|
||||
cloned_return_names = []
|
||||
# We can't just directly pass all of the arguments from the functional op into the mutating op.
|
||||
# We need to check for which inputs to the mutating operator are mutable,
|
||||
# and clone those inputs first.
|
||||
for a_curr, a_tgt in zip(
|
||||
dispatcher.jit_arguments(g.functional.func),
|
||||
dispatcher.jit_arguments(target_f.func),
|
||||
):
|
||||
if a_tgt.annotation is not None and a_tgt.annotation.is_write:
|
||||
clone_mutable_inputs.append(
|
||||
f"auto {a_curr.name}_clone = clone_arg({a_curr.name});"
|
||||
)
|
||||
context.append(
|
||||
Expr(
|
||||
expr=f"{a_curr.name}_clone",
|
||||
type=dispatcher.argument_type(a_curr, binds=a_curr.name),
|
||||
)
|
||||
)
|
||||
# Invariant: mutable arguments on the inner mutable op are always returns on the functional op.
|
||||
cloned_return_names.append(f"{a_curr.name}_clone")
|
||||
else:
|
||||
context.append(dispatcher.argument(a_curr))
|
||||
exprs = ", ".join([e.expr for e in translate(context, target_sig.arguments())])
|
||||
|
||||
out_name = "output"
|
||||
maybe_assign = f"auto {out_name} = " if len(target_f.func.returns) > 0 else ""
|
||||
inner_return_names = gather_nonaliased_inner_rets(target_f.func, out_name)
|
||||
ret_str = return_str(
|
||||
g.functional.func.returns, inner_return_names + cloned_return_names
|
||||
)
|
||||
|
||||
clone_mutable_inputs_str = "\n".join(clone_mutable_inputs)
|
||||
return f"""
|
||||
{sig.defn(name=sig.name() + ("_symint" if g.out.func.has_symint() else ""))} {{
|
||||
{clone_mutable_inputs_str}
|
||||
{maybe_assign}at::_ops::{target_f.func.name.unambiguous_name()}::call({exprs});
|
||||
{ret_str}
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
# Generates out= kernels in terms of their functional counterparts.
|
||||
# We only do this for "generated" NativeFunctions
|
||||
@with_native_function
|
||||
def gen_composite_out_kernel(g: NativeFunctionsGroup) -> str | None:
|
||||
# We should only be generating these for code-generated NativeFunctions
|
||||
if "generated" not in g.out.tags:
|
||||
return None
|
||||
# And we always write the kernel for the out= op in terms of the functional.
|
||||
# Note that the functional op might have also been generated, but we don't have to
|
||||
# worry about cycles, because the generated functional kernels are always implemented
|
||||
# in terms of non-generated kernels (see gen_composite_functional_kernel).
|
||||
|
||||
sig = DispatcherSignature(g.out.func)
|
||||
target_sig = DispatcherSignature(g.functional.func)
|
||||
|
||||
exprs = ", ".join(
|
||||
[e.expr for e in translate(sig.arguments(), target_sig.arguments())]
|
||||
)
|
||||
|
||||
copy_outs = []
|
||||
out_name = "tmp_output"
|
||||
for i, out_arg in enumerate(g.out.func.arguments.out):
|
||||
functional_return_name = (
|
||||
out_name
|
||||
if len(g.functional.func.returns) == 1
|
||||
else f"std::get<{i}>({out_name})"
|
||||
)
|
||||
copy_outs.append(
|
||||
f"""\
|
||||
resize_out_helper({out_arg.name}, {functional_return_name});
|
||||
copy_arg({out_arg.name}, {functional_return_name});"""
|
||||
)
|
||||
|
||||
rets = []
|
||||
# For each return arg in the calling (out=) operator,
|
||||
# If it corresponds to an aliased input, return the input.
|
||||
# Otherwise, return the corresponding output from calling the functional operator.
|
||||
for i, ret_name in enumerate(g.out.func.aliased_return_names()):
|
||||
if ret_name is not None:
|
||||
rets.append(ret_name)
|
||||
else:
|
||||
functional_return_name = (
|
||||
out_name
|
||||
if len(g.functional.func.returns) == 1
|
||||
else f"std::get<{i}>({out_name})"
|
||||
)
|
||||
rets.append(functional_return_name)
|
||||
|
||||
copy_outs_str = "\n".join(copy_outs)
|
||||
|
||||
# Kernel name needs to follow the naming convention defined in `generate_function()`
|
||||
return f"""
|
||||
{sig.defn(name=g.out.func.name.unambiguous_name() + ("_symint" if g.out.func.has_symint() else ""))} {{
|
||||
auto {out_name} = at::_ops::{g.functional.func.name.unambiguous_name()}::call({exprs});
|
||||
{copy_outs_str}
|
||||
{return_str(g.out.func.returns, rets)}
|
||||
}}
|
||||
"""
|
||||
+390
@@ -0,0 +1,390 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from operator import itemgetter
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.jit.generate_bytecode import generate_upgraders_bytecode
|
||||
from torchgen.code_template import CodeTemplate
|
||||
from torchgen.operator_versions.gen_mobile_upgraders_constant import (
|
||||
MOBILE_UPGRADERS_HEADER_DESCRIPTION,
|
||||
)
|
||||
|
||||
|
||||
class ByteCode(Enum):
|
||||
instructions = 1
|
||||
constants = 2
|
||||
types = 3
|
||||
operators = 4
|
||||
register_size = 5
|
||||
|
||||
|
||||
EXCLUDED_OP_SET = [
|
||||
"aten::full.names",
|
||||
"aten::full.out",
|
||||
"aten::full",
|
||||
]
|
||||
|
||||
EXCLUE_UPGRADER_SET = ["full_0_4", "full_out_0_4"]
|
||||
|
||||
ONE_INSTRUCTION = CodeTemplate(
|
||||
"""
|
||||
Instruction{OpCode::${operator_name}, ${X}, ${N}},"""
|
||||
)
|
||||
|
||||
INSTRUCTION_LIST = CodeTemplate(
|
||||
"""std::vector<Instruction>({
|
||||
${instruction_list}
|
||||
}), // instructions list"""
|
||||
)
|
||||
|
||||
ONE_CONSTANT = CodeTemplate(
|
||||
"""
|
||||
c10::IValue(${constant}),"""
|
||||
)
|
||||
|
||||
CONSTANT_LIST = CodeTemplate(
|
||||
"""std::vector<c10::IValue>({
|
||||
${constant_list}
|
||||
}), // constants list"""
|
||||
)
|
||||
|
||||
CONSTANTS_LIST_EMPTY = """std::vector<c10::IValue>(), // constants list"""
|
||||
|
||||
ONE_TYPE = CodeTemplate("""c10::parseType("${type_str}"),""")
|
||||
|
||||
TYPE_LIST = CodeTemplate(
|
||||
"""std::vector<c10::TypePtr>({
|
||||
${type_list}
|
||||
}), // types list"""
|
||||
)
|
||||
|
||||
TYPE_LIST_EMPTY = """std::vector<c10::TypePtr>(), // types list"""
|
||||
|
||||
ONE_OPERATOTR_STRING = CodeTemplate(
|
||||
"""
|
||||
OperatorString({"${operator_name}", "${overload_name}", ${num_of_args}}),"""
|
||||
)
|
||||
|
||||
OPERATOR_STRING_LIST = CodeTemplate(
|
||||
"""
|
||||
std::vector<OperatorString>({
|
||||
${operator_string_list}
|
||||
}), // operators list"""
|
||||
)
|
||||
|
||||
ONE_UPGRADER_FUNCTION = CodeTemplate(
|
||||
"""
|
||||
mobile::Function::registerFunc(
|
||||
"${upgrader_name}",
|
||||
${instruction_list},
|
||||
${constant_list},
|
||||
${type_list},
|
||||
${register_size}
|
||||
)"""
|
||||
)
|
||||
|
||||
ONE_UPGRADER_SRC = CodeTemplate(
|
||||
"""
|
||||
ByteCodeFunctionWithOperator({
|
||||
${bytecode_function},
|
||||
${operator_string_list}
|
||||
}),"""
|
||||
)
|
||||
|
||||
|
||||
ONE_UPGRADER_IN_VERSION_MAP = CodeTemplate(
|
||||
"""Upgrader({${upgrader_min_version}, ${upgrader_max_version}, "${upgrader_name}", ${bytecode_func_index}})"""
|
||||
) # noqa: E501
|
||||
|
||||
ONE_OPERATOR_IN_VERSION_MAP = CodeTemplate(
|
||||
"""
|
||||
{std::string("${operator_name}"),
|
||||
std::vector<Upgrader>({
|
||||
${upgrader_list_in_version_map}
|
||||
})},"""
|
||||
)
|
||||
|
||||
|
||||
OPERATOR_VERSION_MAP = CodeTemplate(
|
||||
"""
|
||||
const std::unordered_map<std::string, std::vector<Upgrader>>
|
||||
getOperatorVersionMapForMobile() {
|
||||
static std::unordered_map<std::string, std::vector<Upgrader>>
|
||||
operatorVersionMapForMobile({
|
||||
${operator_list_in_version_map}
|
||||
});
|
||||
return operatorVersionMapForMobile;
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
UPGRADER_CPP_SRC = CodeTemplate(
|
||||
MOBILE_UPGRADERS_HEADER_DESCRIPTION
|
||||
+ """
|
||||
#include <caffe2/serialize/versions.h>
|
||||
#include <torch/csrc/jit/mobile/type_parser.h>
|
||||
#include <torch/csrc/jit/mobile/upgrader_mobile.h>
|
||||
|
||||
namespace torch {
|
||||
namespace jit {
|
||||
|
||||
// clang-format off
|
||||
|
||||
// From operator_versions_map
|
||||
${operator_version_map}
|
||||
|
||||
const std::vector<ByteCodeFunctionWithOperator>& getUpgraderBytecodeList() {
|
||||
auto generate_upgrader_bytecode_list = []() {
|
||||
std::vector<ByteCodeFunctionWithOperator> upgrader_function_list({
|
||||
${upgrader_bytecode}
|
||||
});
|
||||
for (const auto& upgrader_function : upgrader_function_list) {
|
||||
for (const auto& op : upgrader_function.operators) {
|
||||
upgrader_function.function.append_operator(
|
||||
op.name,
|
||||
op.overload_name,
|
||||
op.num_specified_args);
|
||||
}
|
||||
}
|
||||
return upgrader_function_list;
|
||||
};
|
||||
static std::vector<ByteCodeFunctionWithOperator> upgraderBytecodeList =
|
||||
generate_upgrader_bytecode_list();
|
||||
return upgraderBytecodeList;
|
||||
}
|
||||
|
||||
// clang-format on
|
||||
|
||||
} // namespace jit
|
||||
} // namespace torch
|
||||
"""
|
||||
)
|
||||
|
||||
UPGRADER_MOBILE_FILE_NAME = "upgrader_mobile.cpp"
|
||||
|
||||
UPGRADER_ELEMENT = CodeTemplate(
|
||||
"""\
|
||||
Upgrader({${min_version}, ${max_version}, ${operator_name}, ${index}}),
|
||||
"""
|
||||
)
|
||||
|
||||
PER_OPERATOR_UPGRADER_LIST = CodeTemplate(
|
||||
"""\
|
||||
{
|
||||
std::string(${operator_name}),
|
||||
std::vector<Upgrader>({${upgrader_list}});
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def construct_instruction(instruction_list_from_yaml: list[Any]) -> str:
|
||||
instruction_list_part = [
|
||||
ONE_INSTRUCTION.substitute(
|
||||
operator_name=instruction[0],
|
||||
X=instruction[1],
|
||||
N=instruction[2],
|
||||
)
|
||||
for instruction in instruction_list_from_yaml
|
||||
]
|
||||
return INSTRUCTION_LIST.substitute(
|
||||
instruction_list="".join(instruction_list_part).lstrip("\n")
|
||||
)
|
||||
|
||||
|
||||
def construct_constants(constants_list_from_yaml: list[Any]) -> str:
|
||||
constants_list_part = []
|
||||
for constant_from_yaml in constants_list_from_yaml:
|
||||
convert_constant = None
|
||||
if isinstance(constant_from_yaml, str):
|
||||
# Add quotes if it's string
|
||||
convert_constant = f'"{constant_from_yaml}"'
|
||||
elif isinstance(constant_from_yaml, bool):
|
||||
convert_constant = "true" if constant_from_yaml else "false"
|
||||
elif constant_from_yaml is None:
|
||||
convert_constant = ""
|
||||
elif isinstance(constant_from_yaml, int):
|
||||
convert_constant = str(constant_from_yaml)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The type of {constant_from_yaml} is {type(constant_from_yaml)}. "
|
||||
"Please add change in construct_constants function in gen_mobile_upgraders.py."
|
||||
)
|
||||
constants_list_part.append(ONE_CONSTANT.substitute(constant=convert_constant))
|
||||
if len(constants_list_part) == 0:
|
||||
return CONSTANTS_LIST_EMPTY
|
||||
return CONSTANT_LIST.substitute(
|
||||
constant_list="".join(constants_list_part).lstrip("\n")
|
||||
)
|
||||
|
||||
|
||||
def construct_operators(operator_list_from_yaml: list[Any]) -> str:
|
||||
operator_list_part = [
|
||||
ONE_OPERATOTR_STRING.substitute(
|
||||
operator_name=operator[0],
|
||||
overload_name=operator[1],
|
||||
num_of_args=operator[2],
|
||||
)
|
||||
for operator in operator_list_from_yaml
|
||||
]
|
||||
return OPERATOR_STRING_LIST.substitute(
|
||||
operator_string_list="".join(operator_list_part).lstrip("\n")
|
||||
)
|
||||
|
||||
|
||||
def construct_types(types_tr_list_from_yaml: list[Any]) -> str:
|
||||
types_tr_list_part = [
|
||||
ONE_TYPE.substitute(type_str=types_tr) for types_tr in types_tr_list_from_yaml
|
||||
]
|
||||
if len(types_tr_list_part) == 0:
|
||||
return TYPE_LIST_EMPTY
|
||||
return TYPE_LIST.substitute(type_list="".join(types_tr_list_part).lstrip("\n"))
|
||||
|
||||
|
||||
def construct_register_size(register_size_from_yaml: int) -> str:
|
||||
if not isinstance(register_size_from_yaml, int):
|
||||
raise ValueError(
|
||||
f"Input register size is {register_size_from_yaml} and"
|
||||
"it's type is {type(register_size_from_yaml)}. An int type is expected."
|
||||
)
|
||||
return str(register_size_from_yaml)
|
||||
|
||||
|
||||
def construct_version_maps(
|
||||
upgrader_bytecode_function_to_index_map: dict[str, Any],
|
||||
) -> str:
|
||||
version_map = torch._C._get_operator_version_map()
|
||||
sorted_version_map_ = sorted(version_map.items(), key=itemgetter(0)) # type: ignore[no-any-return]
|
||||
sorted_version_map = dict(sorted_version_map_)
|
||||
|
||||
operator_list_in_version_map_part = []
|
||||
for op_name in sorted_version_map:
|
||||
upgraders_in_version_map_part = []
|
||||
# TODO: remove the skip after these two operators schemas are fixed
|
||||
if op_name in EXCLUDED_OP_SET:
|
||||
continue
|
||||
upgrader_ranges = torch._C._get_upgrader_ranges(op_name)
|
||||
upgrader_entries = sorted_version_map[op_name]
|
||||
if len(upgrader_ranges) != len(upgrader_entries):
|
||||
raise AssertionError(
|
||||
f"upgrader_ranges and upgrader_entries length mismatch for {op_name}: "
|
||||
f"{len(upgrader_ranges)} != {len(upgrader_entries)}"
|
||||
)
|
||||
for idx, upgrader_entry in enumerate(upgrader_entries):
|
||||
upgrader_name = upgrader_entry.upgrader_name
|
||||
bytecode_function_index = upgrader_bytecode_function_to_index_map[
|
||||
upgrader_name
|
||||
]
|
||||
upgraders_in_version_map_part.append(
|
||||
ONE_UPGRADER_IN_VERSION_MAP.substitute(
|
||||
upgrader_min_version=upgrader_ranges[idx].min_version,
|
||||
upgrader_max_version=upgrader_ranges[idx].max_version,
|
||||
upgrader_name=upgrader_name,
|
||||
bytecode_func_index=bytecode_function_index,
|
||||
)
|
||||
)
|
||||
operator_list_in_version_map_part.append(
|
||||
ONE_OPERATOR_IN_VERSION_MAP.substitute(
|
||||
operator_name=op_name,
|
||||
upgrader_list_in_version_map="".join(upgraders_in_version_map_part),
|
||||
)
|
||||
)
|
||||
return OPERATOR_VERSION_MAP.substitute(
|
||||
operator_list_in_version_map="".join(operator_list_in_version_map_part).lstrip(
|
||||
"\n"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_upgrader_bytecode_function_to_index_map(
|
||||
upgrader_dict: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
upgrader_bytecode_function_to_index_map = {}
|
||||
index = 0
|
||||
for upgrader_bytecode in upgrader_dict:
|
||||
for upgrader_name in upgrader_bytecode:
|
||||
if upgrader_name in EXCLUE_UPGRADER_SET:
|
||||
continue
|
||||
upgrader_bytecode_function_to_index_map[upgrader_name] = index
|
||||
index += 1
|
||||
return upgrader_bytecode_function_to_index_map
|
||||
|
||||
|
||||
def write_cpp(cpp_path: str, upgrader_dict: list[dict[str, Any]]) -> None:
|
||||
upgrader_bytecode_function_to_index_map = (
|
||||
get_upgrader_bytecode_function_to_index_map(upgrader_dict)
|
||||
)
|
||||
version_map_src = construct_version_maps(upgrader_bytecode_function_to_index_map)
|
||||
all_upgrader_src_string = []
|
||||
for upgrader_bytecode in upgrader_dict:
|
||||
for upgrader_name, bytecode in upgrader_bytecode.items():
|
||||
# TODO: remove the skip after these two operators schemas are fixed
|
||||
if upgrader_name in EXCLUE_UPGRADER_SET:
|
||||
continue
|
||||
instruction_list_str = ""
|
||||
constant_list_str = ""
|
||||
type_list_str = ""
|
||||
register_size_str = ""
|
||||
operator_list_str = ""
|
||||
for table_name, contents in bytecode.items():
|
||||
element = ByteCode[table_name]
|
||||
if element is ByteCode.instructions:
|
||||
instruction_list_str = construct_instruction(contents)
|
||||
elif element is ByteCode.constants:
|
||||
constant_list_str = construct_constants(contents)
|
||||
elif element is ByteCode.operators:
|
||||
operator_list_str = construct_operators(contents)
|
||||
elif element is ByteCode.types:
|
||||
type_list_str = construct_types(contents)
|
||||
elif element is ByteCode.register_size:
|
||||
register_size_str = construct_register_size(contents)
|
||||
|
||||
one_upgrader_function_string = ONE_UPGRADER_FUNCTION.substitute(
|
||||
upgrader_name=upgrader_name,
|
||||
instruction_list=instruction_list_str,
|
||||
constant_list=constant_list_str,
|
||||
type_list=type_list_str,
|
||||
register_size=register_size_str,
|
||||
)
|
||||
one_upgrader_src_string = ONE_UPGRADER_SRC.substitute(
|
||||
bytecode_function=one_upgrader_function_string.lstrip("\n"),
|
||||
operator_string_list=operator_list_str.lstrip("\n"),
|
||||
)
|
||||
all_upgrader_src_string.append(one_upgrader_src_string)
|
||||
|
||||
upgrader_file_content = UPGRADER_CPP_SRC.substitute(
|
||||
operator_version_map=version_map_src,
|
||||
upgrader_bytecode="".join(all_upgrader_src_string).lstrip("\n"),
|
||||
)
|
||||
print("writing file to : ", cpp_path + "/" + UPGRADER_MOBILE_FILE_NAME)
|
||||
with open(os.path.join(cpp_path, UPGRADER_MOBILE_FILE_NAME), "wb") as out_file:
|
||||
out_file.write(upgrader_file_content.encode("utf-8"))
|
||||
|
||||
|
||||
def sort_upgrader(upgrader_list: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
sorted_upgrader_list = sorted(
|
||||
upgrader_list, key=lambda one_upgrader: next(iter(one_upgrader))
|
||||
)
|
||||
return sorted_upgrader_list
|
||||
|
||||
|
||||
def main() -> None:
|
||||
upgrader_list = generate_upgraders_bytecode()
|
||||
sorted_upgrader_list = sort_upgrader(upgrader_list)
|
||||
for up in sorted_upgrader_list:
|
||||
print("after sort upgrader : ", next(iter(up)))
|
||||
|
||||
pytorch_dir = Path(__file__).resolve().parents[2]
|
||||
upgrader_path = pytorch_dir / "torch" / "csrc" / "jit" / "mobile"
|
||||
write_cpp(str(upgrader_path), sorted_upgrader_list)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
MOBILE_UPGRADERS_HEADER_DESCRIPTION = """/**
|
||||
* @generated
|
||||
* This is an auto-generated file. Please do not modify it by hand.
|
||||
* To re-generate, please run:
|
||||
* cd ~/pytorch && python torchgen/operator_versions/gen_mobile_upgraders.py
|
||||
*/
|
||||
"""
|
||||
+16269
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,105 @@
|
||||
# This yaml file contains all the possible tags that can be defined in `tags` in `native_functions.yaml`
|
||||
|
||||
- tag: inplace_view
|
||||
desc: |
|
||||
This tag indicates if an operator *only* modifies the tensor metadata
|
||||
- tag: pt2_compliant_tag
|
||||
desc: |
|
||||
This tag indicates if the operator is guaranteed to
|
||||
work with the PT2 compilation APIs (torch.compile,
|
||||
torch.export, etc). If you add this tag to an
|
||||
operator, please use
|
||||
`torch.testing._internal.optest.opcheck` to test that
|
||||
the operator has been registered correctly and
|
||||
works with torch.compile
|
||||
- tag: view_copy
|
||||
desc: |
|
||||
This tag indicates operators that are *_copy* variants
|
||||
of view/aliasing operators. If an operator has a view_copy tag,
|
||||
then it should have the name {op}_copy, where {op} is a view operator.
|
||||
- tag: dynamic_output_shape
|
||||
desc: |
|
||||
This tag indicates if an operator's output's shape depends on input Tensor
|
||||
data.
|
||||
- tag: data_dependent_output
|
||||
desc: |
|
||||
Operator has a non-Tensor output whose value is dependent on the data
|
||||
of Tensor inputs. Among other things, this implies that this operator
|
||||
cannot be run with meta tensor (since data is not available), nor
|
||||
can it be symbolically traced.
|
||||
- tag: generated
|
||||
desc: |
|
||||
This tag indicates that the operator doesn't have an explicit entry in
|
||||
native_functions.yaml, and instead was generated automatically by the codegen.
|
||||
- tag: nondeterministic_seeded
|
||||
desc: |
|
||||
This tag indicates if an operator is nondeterministically seeded
|
||||
(i.e., is random) such that the operator intentionally produces
|
||||
different results when run twice on the same inputs, but this randomness
|
||||
is controlled by a Generator which, if reseeded would give you the
|
||||
same result.
|
||||
- tag: nondeterministic_bitwise
|
||||
desc: |
|
||||
This tag indicates if an operator doesn't guarantee bitwise equivalence
|
||||
across different runs of an operator with identical inputs.
|
||||
- tag: needs_exact_strides
|
||||
desc: |
|
||||
This tag indicates that the operator should be passed Tensors following
|
||||
the same strides as observed in eager when compiled in inductor.
|
||||
Only one of {needs_exact_strides, needs_contiguous_strides, needs_fixed_stride_order, flexible_layout}
|
||||
can apply; if multiple are assigned then we assume the most restrictive one.
|
||||
- tag: needs_contiguous_strides
|
||||
desc: |
|
||||
This tag indicates that the operator should be passed contiguous Tensors.
|
||||
Failure to do so will result in undefined behavior.
|
||||
- tag: needs_fixed_stride_order
|
||||
desc: |
|
||||
This tag indicates that the operator should be passed Tensors following
|
||||
the same stride permutation as observed in eager when compiled in inductor.
|
||||
Only one of {needs_exact_strides, needs_contiguous_strides, needs_fixed_stride_order, flexible_layout}
|
||||
can apply; if multiple are assigned then we assume the most restrictive one.
|
||||
- tag: flexible_layout
|
||||
desc: |
|
||||
This tag indicates that the custom operator can accept inputs with varying
|
||||
strides/storage_offset and that when compiled, Inductor is allowed to change
|
||||
the strides/storage_offset of inputs to the custom operator.
|
||||
Only one of {needs_exact_strides, needs_contiguous_strides, needs_fixed_stride_order, flexible_layout}
|
||||
can apply; if multiple are assigned then we assume the most restrictive one.
|
||||
|
||||
# NOTE [Core ATen Ops]
|
||||
- tag: core
|
||||
desc: |
|
||||
Core aten ops is a subset of aten ops that remains after aten-to-aten decomposition and
|
||||
functionalization pass. Core aten ops are fully functional and adhere to single static
|
||||
assignment (SSA): this implies there will be no `inplace` or `_out` variants in this opset.
|
||||
This opset is designed to serve as the functional IR to interface with compiler backends.
|
||||
In contrast to primTorch, core aten opset doesn't decompose ops into explicit
|
||||
type promotion and broadcasting ops.
|
||||
Core aten ops is also effectively the opset produced by torchdynamo.export(aten_graph=True),
|
||||
and thus can be used as an opset for export purpose.
|
||||
- tag: pointwise
|
||||
desc: |
|
||||
Pointwise operators are operators where each element of the output is computed only by accessing
|
||||
the corresponding element of all the broadcasted inputs. The output shape will be the broadcasted
|
||||
shape of the inputs.
|
||||
- tag: maybe_aliasing_or_mutating
|
||||
desc: |
|
||||
For some ops, we can't statically determine whether the op is functional or not. Note that this is only
|
||||
relevant to CIA ops that decompose before functionalization/autograd. It is useful to
|
||||
know this information for export as we would want to decompose these ops as they are unsafe to be
|
||||
preserved.
|
||||
- tag: cudagraph_unsafe
|
||||
desc: |
|
||||
This operator does not support cudagraphs. The presence of this tag on an operator will cause
|
||||
Inductor to split the graph around this operator. Note that operators without this tag may still
|
||||
not support CUDAGraphs. Inductor may have other hardcoded lists around that.
|
||||
- tag: reduction
|
||||
desc: |
|
||||
This tag indicates that an operator performs a reduction operation, computing aggregate values
|
||||
(sum, mean, max, min, etc.) across one or more dimensions of the input tensor(s).
|
||||
- tag: out_variant
|
||||
desc: |
|
||||
This tag indicates that the operator is an out variant of a functional
|
||||
operator. This tag only applies to custom ops. Out variant operators
|
||||
write their results to pre-allocated output tensors (the out args)
|
||||
rather than allocating new tensors.
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
#include <ATen/core/ATenOpList.h>
|
||||
|
||||
#include <string>
|
||||
#include <cstring>
|
||||
#include <utility>
|
||||
#include <unordered_set>
|
||||
#include <ATen/core/operator_name.h>
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace {
|
||||
struct OpNameEquals final {
|
||||
bool operator()(const std::pair<const char*, const char*>& lhs, const std::pair<const char*, const char*>& rhs) const {
|
||||
return 0 == strcmp(lhs.first, rhs.first) && 0 == strcmp(lhs.second, rhs.second);
|
||||
}
|
||||
};
|
||||
|
||||
struct OpNameHash final {
|
||||
size_t operator()(const std::pair<const char*, const char*>& p) const {
|
||||
// use std::hash<std::string> because std::hash<const char*> would hash pointers and not pointed-to strings
|
||||
return std::hash<std::string>()(p.first) ^ (~ std::hash<std::string>()(p.second));
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
bool is_custom_op(const c10::OperatorName& opName) {
|
||||
static std::unordered_set<std::pair<const char*, const char*>, OpNameHash, OpNameEquals> ops {
|
||||
${aten_ops}
|
||||
{"", ""}
|
||||
};
|
||||
return ops.count(std::make_pair(
|
||||
opName.name.c_str(), opName.overload_name.c_str())) == 0;
|
||||
}
|
||||
}
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/InferSize.h>
|
||||
#include <ATen/Tensor.h>
|
||||
#include <ATen/native/Resize.h>
|
||||
|
||||
#ifndef AT_PER_OPERATOR_HEADERS
|
||||
#include <ATen/Operators.h>
|
||||
#else
|
||||
#include <ATen/ops/clone.h>
|
||||
$ops_headers
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
namespace native {
|
||||
|
||||
// This file contains a number of kernels for aten functions that are fully code-generated.
|
||||
// TODO: rename this file to something more generic.
|
||||
|
||||
namespace {
|
||||
at::Tensor clone_arg(const at::Tensor& t) {
|
||||
return t.clone();
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> clone_arg(const at::TensorList& t_list) {
|
||||
std::vector<at::Tensor> out(t_list.size());
|
||||
for (const auto& i : c10::irange(t_list.size())) {
|
||||
out[i] = t_list[i].clone();
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// duped with gen_resize_out_helper from structured kernels
|
||||
void copy_arg(const at::Tensor& dst, const at::Tensor& src) {
|
||||
TORCH_CHECK(src.dtype() == dst.dtype(),
|
||||
"Expected out tensor to have dtype ", src.dtype(), ", but got ", dst.dtype(), " instead");
|
||||
TORCH_CHECK(src.device() == dst.device(),
|
||||
"Expected out tensor to have device ", src.device(), ", but got ", dst.device(), " instead");
|
||||
dst.copy_(src);
|
||||
}
|
||||
|
||||
void copy_arg(const at::TensorList& dst, const at::TensorList& src) {
|
||||
TORCH_INTERNAL_ASSERT(dst.size() == src.size());
|
||||
for (const auto& i : c10::irange(dst.size())) {
|
||||
copy_arg(dst[i], src[i]);
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: this doesn't handle restriding empty tensors correctly; see
|
||||
// gen_resize_out_helper for the correct algorithm
|
||||
|
||||
void resize_out_helper(const at::Tensor& dst, const at::Tensor& src) {
|
||||
at::native::resize_output(dst, src.sizes());
|
||||
}
|
||||
|
||||
void resize_out_helper(const at::TensorList& dst, const at::TensorList& src) {
|
||||
TORCH_INTERNAL_ASSERT(dst.size() == src.size());
|
||||
for (const auto& i : c10::irange(dst.size())) {
|
||||
at::native::resize_output(dst[i], src[i].sizes());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
${CompositeViewCopyKernel_Definitions}
|
||||
|
||||
${GeneratedCompositeFunctional_Definitions}
|
||||
|
||||
${GeneratedCompositeOut_Definitions}
|
||||
|
||||
} // namespace native
|
||||
} // namespace at
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
#pragma once
|
||||
// ${generated_comment}
|
||||
|
||||
// NB: The implementing C++ file is RegisterDispatchKey.cpp
|
||||
|
||||
// The only #includes we need are for custom classes that have defaults in the C++ API
|
||||
#include <c10/core/MemoryFormat.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
|
||||
// Forward declarations of any types needed in the operator signatures.
|
||||
// We can't directly include these classes because it will cause circular include dependencies.
|
||||
// This file is included by TensorBody.h, which defines the Tensor class.
|
||||
#include <ATen/core/ATen_fwd.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace ${dispatch_namespace} {
|
||||
|
||||
${dispatch_namespaced_declarations}
|
||||
|
||||
} // namespace ${dispatch_namespace}
|
||||
} // namespace at
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
#include <ATen/core/TensorBody.h>
|
||||
|
||||
// TODO Undo all logic introduced for Note [Avoiding Include Cycles In Static Dispatch]
|
||||
// Code introduced to avoid cyclic dependency in static dispatch is no longer
|
||||
// needed as static dispatch logic is moved from TensorBody.h, which caused cycles in the first place,
|
||||
// to Operators.cpp for supporting multiple backends with multiple kernels.
|
||||
//
|
||||
// Note [Avoiding Include Cycles In Static Dispatch]
|
||||
// In order to avoid #include cycles in the static dispatch build, we've carefully split out
|
||||
// the static function definition files into {DispatchKey}Functions.h and {DispatchKey}Functions_inl.h.
|
||||
//
|
||||
// Without this split, the include cycle looks like TensorBody.h -> CPUFunctions.h -> TensorBody.h.
|
||||
// - TensorBody.h #includes CPUFunctions.h in the static dispatch build, because the tensor methods
|
||||
// all need to call into the fastpath C++ API defined in CPUFunctions.h. The methods are also all
|
||||
// directly inlined into TensorBody.h.
|
||||
// - CPUFunctions.h #includes TensorBody.h because it contains function declarations for the entire C++ API,
|
||||
// which include functions that have defaultable std::optional<Tensor> arguments.
|
||||
// That requires knowing the full Tensor class definition.
|
||||
//
|
||||
// We break the cycle by doing the following:
|
||||
// - Split out CPUFunction.h into two files: CPUFunctions.h and CPUFunctions_inl.h
|
||||
// - CPUFunction.h is a dummy file that just includes the Tensor class and includes CPUFunctions_inl.,
|
||||
// - CPUFunctions_inl.h includes everything else
|
||||
// - (only in the static dispatch build) TensorBody.h makes sure to finish defining the Tensor class,
|
||||
// and then it includes CPUFunctions_inl.h.
|
||||
// - All other files that want the cpu fastpath functions can include CPUFunctions.h directly.
|
||||
// - This also means that static dispatch build, CPUFunctions.h only needs to
|
||||
// #include TensorBody.h, and it will automatically bring in CPUFunctions_inl.h.
|
||||
${inline_headers}
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
#pragma once
|
||||
// ${generated_comment}
|
||||
|
||||
// NB: The implementing C++ file is RegisterDispatchKey.cpp
|
||||
|
||||
// The only #includes we need are for custom classes that have defaults in the C++ API
|
||||
#include <c10/core/MemoryFormat.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
|
||||
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
||||
#error This change adds a dependency on all pytorch operators, meaning the \
|
||||
file will need to be re-compiled every time an operator is changed or added. \
|
||||
Consider including a specific operator from \
|
||||
<ATen/ops/{my_operator}_${dispatch_namespace}_dispatch.h>. \
|
||||
See NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
||||
#endif
|
||||
|
||||
${DispatchKeyFunctions_inl_includes}
|
||||
|
||||
|
||||
${dispatch_namespaced_declarations}
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
// ${generated_comment}
|
||||
${includes}
|
||||
${native_functions_include}
|
||||
|
||||
namespace {
|
||||
${helper_fns}
|
||||
} // namespace
|
||||
|
||||
${namespace_prologue}
|
||||
|
||||
${native_function_definitions}
|
||||
|
||||
${namespace_epilogue}
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
#pragma once
|
||||
|
||||
// an external backend might generate file within its code tree
|
||||
// and check all the source files within the tree with clang-format.
|
||||
// so, disable it since the backend might have a different config.
|
||||
// clang-format off
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/Tensor.h>
|
||||
|
||||
${namespace_prologue}
|
||||
|
||||
struct ${class_name} {
|
||||
|
||||
${dispatch_declarations}
|
||||
|
||||
};
|
||||
${namespace_epilogue}
|
||||
@@ -0,0 +1,27 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/Context.h>
|
||||
#include <ATen/DeviceGuard.h>
|
||||
#include <ATen/TensorUtils.h>
|
||||
#include <ATen/TracerMode.h>
|
||||
#include <ATen/core/Generator.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <optional>
|
||||
#include <string_view>
|
||||
|
||||
${static_dispatch_ops_headers}
|
||||
|
||||
${operator_includes}
|
||||
|
||||
namespace at {
|
||||
|
||||
${function_definitions}
|
||||
|
||||
}
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/FunctionalStorageImpl.h>
|
||||
#include <ATen/Tensor.h>
|
||||
|
||||
namespace at {
|
||||
namespace functionalization {
|
||||
|
||||
struct FunctionalInverses {
|
||||
|
||||
${view_inverse_declarations}
|
||||
|
||||
// NB: These are not generated! They're manually implemented in the template.
|
||||
// TODO: Change codegen to generate these. See the following link:
|
||||
// https://github.com/pytorch/pytorch/blob/main/torchgen/model.py#L2583-L2585
|
||||
static at::Tensor chunk_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int64_t mutated_view_idx, int chunks, int dim);
|
||||
static at::Tensor narrow_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int dim, c10::SymInt start, c10::SymInt length);
|
||||
|
||||
};
|
||||
}
|
||||
}
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
#include <array>
|
||||
|
||||
#include <ATen/Functions.h>
|
||||
#include <ATen/Utils.h>
|
||||
#include <c10/core/Allocator.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
Tensor TensorMaker::make_tensor() {
|
||||
AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove.
|
||||
tracer::impl::NoTracerDispatchMode tracer_guard{};
|
||||
|
||||
check_size_nonnegative(sizes_);
|
||||
|
||||
TORCH_CHECK_VALUE(
|
||||
!deleter_ || !ctx_,
|
||||
"The deleter and context arguments are mutually exclusive.");
|
||||
|
||||
if (device_ == std::nullopt) {
|
||||
device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type());
|
||||
}
|
||||
|
||||
if (opts_.device().has_index()) {
|
||||
// clang-format off
|
||||
TORCH_CHECK_VALUE(
|
||||
opts_.device() == *device_,
|
||||
"Specified device ", opts_.device(), " does not match device of data ", *device_);
|
||||
// clang-format on
|
||||
}
|
||||
|
||||
std::size_t size_bytes = computeStorageSize();
|
||||
|
||||
DataPtr data_ptr{};
|
||||
if (deleter_) {
|
||||
data_ptr = makeDataPtrFromDeleter();
|
||||
} else {
|
||||
data_ptr = makeDataPtrFromContext();
|
||||
}
|
||||
|
||||
TORCH_CHECK(!resizeable_ || allocator_ != nullptr, "Must specify an allocator with allocator() if you want to use resizeable_storage()");
|
||||
Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr), /*allocator=*/allocator_, /*resizable=*/resizeable_};
|
||||
|
||||
Tensor tensor = detail::make_tensor<TensorImpl>(
|
||||
std::move(storage), opts_.computeDispatchKey(), opts_.dtype());
|
||||
|
||||
TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl();
|
||||
if (strides_) {
|
||||
tensor_impl->set_sizes_and_strides(sizes_, *strides_);
|
||||
} else {
|
||||
tensor_impl->set_sizes_contiguous(sizes_);
|
||||
}
|
||||
if (storage_offset_) {
|
||||
tensor_impl->set_storage_offset(*storage_offset_);
|
||||
}
|
||||
|
||||
tensor_impl->set_requires_grad(opts_.requires_grad());
|
||||
|
||||
return tensor;
|
||||
}
|
||||
|
||||
std::size_t TensorMaker::computeStorageSize() const noexcept {
|
||||
std::size_t itemsize = opts_.dtype().itemsize();
|
||||
|
||||
if (strides_) {
|
||||
auto storage_size = detail::computeStorageNbytes(sizes_, *strides_, itemsize);
|
||||
if (storage_offset_) {
|
||||
storage_size += storage_offset_.value() * itemsize;
|
||||
}
|
||||
return storage_size;
|
||||
}
|
||||
|
||||
std::size_t size = 1;
|
||||
for (std::int64_t s : sizes_) {
|
||||
size *= static_cast<std::size_t>(s);
|
||||
}
|
||||
auto storage_size = size * itemsize;
|
||||
if (storage_offset_) {
|
||||
storage_size += storage_offset_.value() * itemsize;
|
||||
}
|
||||
return storage_size;
|
||||
}
|
||||
|
||||
inline DataPtr TensorMaker::makeDataPtrFromDeleter() noexcept {
|
||||
return InefficientStdFunctionContext::makeDataPtr(data_, std::move(deleter_), *device_);
|
||||
}
|
||||
|
||||
inline DataPtr TensorMaker::makeDataPtrFromContext() noexcept {
|
||||
return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_};
|
||||
}
|
||||
|
||||
IntArrayRef TensorMaker::makeTempSizes() const noexcept {
|
||||
static std::int64_t zeros[5] = {0, 0, 0, 0, 0};
|
||||
if (opts_.has_memory_format()) {
|
||||
MemoryFormat format = *opts_.memory_format_opt();
|
||||
if (format == MemoryFormat::ChannelsLast) {
|
||||
return IntArrayRef(zeros, 4);
|
||||
}
|
||||
if (format == MemoryFormat::ChannelsLast3d) {
|
||||
return IntArrayRef(zeros, 5);
|
||||
}
|
||||
}
|
||||
return IntArrayRef(zeros, 1);
|
||||
}
|
||||
|
||||
} // namespace at
|
||||
+143
@@ -0,0 +1,143 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#ifdef TORCH_ASSERT_NO_OPERATORS
|
||||
#error This change adds a dependency on native_functions.yaml, \
|
||||
meaning the file will need to be re-compiled every time an operator \
|
||||
is changed or added. Consider if your change would be better placed in \
|
||||
another file, or if a more specific header might achieve the same goal. \
|
||||
See NOTE: [Tensor vs. TensorBase]
|
||||
#endif
|
||||
|
||||
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
||||
#error This change adds a dependency on all pytorch operators, meaning the \
|
||||
file will need to be re-compiled every time an operator is changed or added. \
|
||||
Consider including a specific operator from <ATen/ops/{my_operator}.h> and \
|
||||
see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
||||
#endif
|
||||
|
||||
// NOTE: [TORCH_ASSERT_ONLY_METHOD_OPERATORS]
|
||||
//
|
||||
// In ATen, certain generated headers files include the definitions of
|
||||
// every single operator in PyTorch. Unfortunately this means every
|
||||
// time an operator signature is updated or changed in
|
||||
// native_functions.yaml, you (and every other PyTorch developer) need
|
||||
// to recompile every source file that includes any of these headers.
|
||||
//
|
||||
// To break up these header dependencies, and improve incremental
|
||||
// build times for all PyTorch developers. These headers are split
|
||||
// into per-operator headers in the `ATen/ops` folder. This limits
|
||||
// incremental builds to only changes to methods of `Tensor`, or files
|
||||
// that use the specific operator being changed. With `at::sum` as an
|
||||
// example, you should include
|
||||
//
|
||||
// <ATen/ops/sum.h> // instead of ATen/Functions.h
|
||||
// <ATen/ops/sum_native.h> // instead of ATen/NativeFunctions.h
|
||||
// <ATen/ops/sum_ops.h> // instead of ATen/Operators.h
|
||||
// <ATen/ops/sum_cpu_dispatch.h> // instead of ATen/CPUFunctions.h
|
||||
//
|
||||
// However, even if you're careful to use this in your own code.
|
||||
// `Functions.h` might be included indirectly through another header
|
||||
// without you realising. To avoid this, you can add
|
||||
//
|
||||
// #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
//
|
||||
// to the top of your source file. This way any time the non-specific
|
||||
// headers are included, the compiler will error out.
|
||||
//
|
||||
// Also, be aware that `ops` are not available in all build
|
||||
// configurations (namely fb-internal) so you must guard these
|
||||
// includes with `#ifdef AT_PER_OPERATOR_HEADERS`. e.g.
|
||||
//
|
||||
// #ifndef AT_PER_OPERATOR_HEADERS
|
||||
// #include <ATen/Functions.h>
|
||||
// #else
|
||||
// #include <ATen/ops/sum.h>
|
||||
// #endif
|
||||
|
||||
#include <ATen/Context.h>
|
||||
#include <ATen/DeviceGuard.h>
|
||||
#include <ATen/TensorUtils.h>
|
||||
#include <ATen/TracerMode.h>
|
||||
#include <ATen/core/Generator.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <optional>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
|
||||
#include <ATen/ops/from_blob.h>
|
||||
#include <ATen/ops/tensor.h>
|
||||
|
||||
${Functions_includes}
|
||||
|
||||
namespace at {
|
||||
|
||||
${Functions_declarations}
|
||||
|
||||
// Special C++ only overloads for std()-like functions (See gh-40287)
|
||||
// These are needed because int -> bool conversion takes precedence over int -> IntArrayRef
|
||||
// So, for example std(0) would select the std(unbiased=False) overload
|
||||
inline Tensor var(const Tensor& self, int dim) {
|
||||
return at::var(self, IntArrayRef{dim});
|
||||
}
|
||||
inline std::tuple<Tensor, Tensor> var_mean(const Tensor& self, int dim) {
|
||||
return at::var_mean(self, IntArrayRef{dim});
|
||||
}
|
||||
inline Tensor std(const Tensor& self, int dim) {
|
||||
return at::std(self, IntArrayRef{dim});
|
||||
}
|
||||
inline std::tuple<Tensor, Tensor> std_mean(const Tensor& self, int dim) {
|
||||
return at::std_mean(self, IntArrayRef{dim});
|
||||
}
|
||||
|
||||
inline int64_t numel(const Tensor& tensor) {
|
||||
return tensor.numel();
|
||||
}
|
||||
|
||||
inline int64_t size(const Tensor& tensor, int64_t dim) {
|
||||
return tensor.size(dim);
|
||||
}
|
||||
|
||||
inline int64_t stride(const Tensor& tensor, int64_t dim) {
|
||||
return tensor.stride(dim);
|
||||
}
|
||||
|
||||
inline bool is_complex(const Tensor& tensor) {
|
||||
return tensor.is_complex();
|
||||
}
|
||||
|
||||
inline bool is_floating_point(const Tensor& tensor) {
|
||||
return tensor.is_floating_point();
|
||||
}
|
||||
|
||||
inline bool is_signed(const Tensor& tensor) {
|
||||
return tensor.is_signed();
|
||||
}
|
||||
|
||||
inline bool is_inference(const Tensor& tensor) {
|
||||
return tensor.is_inference();
|
||||
}
|
||||
|
||||
inline bool _is_zerotensor(const Tensor& tensor) {
|
||||
return tensor._is_zerotensor();
|
||||
}
|
||||
|
||||
inline bool is_conj(const Tensor& tensor) {
|
||||
return tensor.is_conj();
|
||||
}
|
||||
|
||||
inline Tensor conj(const Tensor& tensor) {
|
||||
return tensor.conj();
|
||||
}
|
||||
|
||||
inline bool is_neg(const Tensor& tensor) {
|
||||
return tensor.is_neg();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
#pragma once
|
||||
|
||||
// This file contains autogenerated LazyTensor IR nodes
|
||||
${lazy_ir_sysinc}
|
||||
${lazy_ir_inc}
|
||||
|
||||
${namespace_prologue}
|
||||
using at::operator<<;
|
||||
|
||||
// kNullValue is used to contribute a static hash value any time
|
||||
// a node has an Optional<Value> input that is nullopt. It is important
|
||||
// to differentiate between HASH(std::nullopt, something) and HASH(something, std::nullopt),
|
||||
// and using kNullValue in the hash function in the order of arguments
|
||||
// serves this purpose.
|
||||
static const torch::lazy::Value kNullValue = torch::lazy::Value();
|
||||
|
||||
${ir_declarations}
|
||||
|
||||
${namespace_epilogue}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
#pragma once
|
||||
|
||||
${lazy_non_native_ir_inc}
|
||||
|
||||
// This file contains autogenerated LazyTensor Non Native IR nodes
|
||||
|
||||
${namespace_prologue}
|
||||
|
||||
${non_native_ir_nodes}
|
||||
|
||||
${namespace_epilogue}
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#ifdef TORCH_ASSERT_NO_OPERATORS
|
||||
#error This change adds a dependency on native_functions.yaml, \
|
||||
meaning the file will need to be re-compiled every time an operator \
|
||||
is changed or added. Consider if your change would be better placed in \
|
||||
another file, or if a more specific header might achieve the same goal. \
|
||||
See NOTE: [Tensor vs. TensorBase]
|
||||
#endif
|
||||
|
||||
// Forward declarations of any types needed in the operator signatures.
|
||||
// We can't directly include these classes because it will cause circular include dependencies.
|
||||
// This file is included by TensorBody.h, which defines the Tensor class.
|
||||
#include <ATen/core/ATen_fwd.h>
|
||||
|
||||
${MethodOperators_includes}
|
||||
|
||||
namespace at {
|
||||
namespace _ops {
|
||||
${MethodOperators_declarations}
|
||||
} // namespace _ops
|
||||
} // namespace at
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <optional>
|
||||
#include <c10/core/QScheme.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
${extra_includes}
|
||||
|
||||
${native_function_declarations}
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#ifdef TORCH_ASSERT_NO_OPERATORS
|
||||
#error This change adds a dependency on native_functions.yaml, \
|
||||
meaning the file will need to be re-compiled every time an operator \
|
||||
is changed or added. Consider if your change would be better placed in \
|
||||
another file, or if a more specific header might achieve the same goal. \
|
||||
See NOTE: [Tensor vs. TensorBase]
|
||||
#endif
|
||||
|
||||
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
||||
#error This change adds a dependency on all pytorch operators, meaning the \
|
||||
file will need to be re-compiled every time an operator is changed or added. \
|
||||
Consider including a specific operator from <ATen/ops/{my_operator}_native.h> \
|
||||
and see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
||||
#endif
|
||||
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <optional>
|
||||
#include <c10/core/QScheme.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
${NativeFunctions_includes}
|
||||
|
||||
${NativeFunctions_declarations}
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <optional>
|
||||
#include <c10/core/QScheme.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
#include <ATen/TensorIterator.h>
|
||||
#include <ATen/TensorMeta.h>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
namespace at {
|
||||
namespace meta {
|
||||
|
||||
${meta_function_declarations}
|
||||
|
||||
} // namespace native
|
||||
} // namespace at
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <ATen/core/IListRef.h>
|
||||
#include <ATen/TensorMeta.h>
|
||||
#include <ATen/TensorIterator.h>
|
||||
|
||||
${NativeMetaFunctions_includes}
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace meta {
|
||||
|
||||
${NativeMetaFunctions_declarations}
|
||||
|
||||
} // namespace meta
|
||||
} // namespace at
|
||||
@@ -0,0 +1,19 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <string_view>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
// Forward declarations of any types needed in the operator signatures.
|
||||
// We can't directly include these classes because it will cause circular include dependencies.
|
||||
// This file is included by TensorBody.h, which defines the Tensor class.
|
||||
#include <ATen/core/ATen_fwd.h>
|
||||
|
||||
namespace at {
|
||||
namespace _ops {
|
||||
|
||||
${declarations}
|
||||
|
||||
}} // namespace at::_ops
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
#include <ATen/Tensor.h>
|
||||
#include <ATen/core/dispatch/Dispatcher.h>
|
||||
|
||||
// ${generated_comment}
|
||||
// NOTE See [Sharded File] comment in VariableType
|
||||
|
||||
#ifndef AT_PER_OPERATOR_HEADERS
|
||||
#include <ATen/Operators.h>
|
||||
#else
|
||||
${operator_headers}
|
||||
#endif
|
||||
|
||||
${static_dispatch_extra_headers}
|
||||
|
||||
namespace at { namespace _ops {
|
||||
|
||||
${definitions}
|
||||
|
||||
}} // namespace at::_ops
|
||||
@@ -0,0 +1,74 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#ifdef TORCH_ASSERT_NO_OPERATORS
|
||||
#error This change adds a dependency on native_functions.yaml, \
|
||||
meaning the file will need to be re-compiled every time an operator \
|
||||
is changed or added. Consider if your change would be better placed in \
|
||||
another file, or if a more specific header might achieve the same goal. \
|
||||
See NOTE: [Tensor vs. TensorBase]
|
||||
#endif
|
||||
|
||||
#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
||||
#error This change adds a dependency on all pytorch operators, meaning the \
|
||||
file will need to be re-compiled every time an operator is changed or added. \
|
||||
Consider including a specific operator from <ATen/ops/{my_operator}_ops.h> \
|
||||
and see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS].
|
||||
#endif
|
||||
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <c10/core/SymIntArrayRef.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <c10/core/QScheme.h>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
${Operators_includes}
|
||||
|
||||
// Extension writers: do you write wrapper functions? Are you frustrated with
|
||||
// resolving overloads of operators? Are you frustrated with dealing with
|
||||
// pointer-to-methods and resolving overloads of pointer-to-methods?? Look no
|
||||
// further, this is the utility for you.
|
||||
//
|
||||
// Given an operator schema: aten::op.overload(...
|
||||
//
|
||||
// Use ATEN_FN2(op, overload) to get a *function* version of the operator
|
||||
// that is guaranteed to not be overloaded. This means that you can safely
|
||||
// decltype(&ATEN_FN2(op, overload)) it. NB: the 2 means this macro takes 2 args.
|
||||
//
|
||||
// Given an operator schema without an overload name: aten::op(...
|
||||
//
|
||||
// Use ATEN_FN(op) to get an unambiguous *function* version of the operator.
|
||||
//
|
||||
// There is some interesting behavior for out= operations.
|
||||
// ATEN_FN2(sin, out) gives a function that is *faithful* to the schema;
|
||||
// that is, the order of arguments is exactly what it looks like in the schema.
|
||||
|
||||
#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload::call
|
||||
#define ATEN_FN(op_name) at::_ops::op_name::call
|
||||
|
||||
// Separately, ATEN_OP(op) and ATEN_OP2(op, overload) define a class containing compile-time
|
||||
// metadata about a given aten operator.
|
||||
// Notable data on the class includes:
|
||||
// - ATEN_OP2(add, Tensor)::name // returns the string name: "add"
|
||||
// - ATEN_OP2(add, Tensor)::overload_name // returns the string overload name: "Tensor"
|
||||
// - ATEN_OP2(add, Tensor)::schema // returns the C++ schema type: at::Tensor (const at::Tensor &, const at::Tensor &, const at::Scalar &)
|
||||
// - ATEN_OP2(add, Tensor)::schema_str // returns the string jit type: "add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor"
|
||||
|
||||
#define ATEN_OP2(op_name, overload) at::_ops::op_name##_##overload
|
||||
#define ATEN_OP(op_name) at::_ops::op_name
|
||||
|
||||
// WARNING: Please do not call any of the ops in the _ops namespace directly.
|
||||
// Use the ATEN_FN macros. We do not guarantee stability of the naming
|
||||
// scheme for the functions in at::_ops
|
||||
|
||||
// See Note [The ATen Operators API] for details of the at::_ops namespace
|
||||
|
||||
namespace at {
|
||||
namespace _ops {
|
||||
${Operators_declarations}
|
||||
} // namespace _ops
|
||||
} // namespace at
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/RedispatchFunctions.h>
|
||||
#include <ATen/Functions.h>
|
||||
|
||||
#include <ATen/core/dispatch/Dispatcher.h>
|
||||
#include <ATen/core/op_registration/adaption.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace redispatch {
|
||||
${function_redispatch_definitions}
|
||||
} // namespace redispatch
|
||||
|
||||
} // namespace at
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#ifdef TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
#error This change adds a dependency on all pytorch operators, meaning the \
|
||||
file will need to be re-compiled every time an operator is changed or added. \
|
||||
Consider using the at::_ops::{name}::redispatch() interface by including \
|
||||
the specific operator from <ATen/ops/{my_operator}_ops.h>
|
||||
#endif
|
||||
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <ATen/Tensor.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <ATen/core/Generator.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <ATen/DeviceGuard.h>
|
||||
#include <c10/core/TensorOptions.h>
|
||||
#include <ATen/core/Reduction.h>
|
||||
#include <optional>
|
||||
#include <ATen/TensorUtils.h>
|
||||
#include <ATen/Context.h>
|
||||
#include <ATen/TracerMode.h>
|
||||
#include <ATen/Operators.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace redispatch {
|
||||
${function_redispatch_definitions}
|
||||
} // namespace redispatch
|
||||
|
||||
}
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
// We register ops with a higher priority dispatch key (BackendSelect) than the usual backend-specific keys (e.g. CPU)
|
||||
// which makes calls to the factory functions dispatch to here.
|
||||
// We then 'manually' compute a lower-priority to re-dispatch to (e.g. CPU) to get to the eventually correct backend.
|
||||
// ${generated_comment}
|
||||
|
||||
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
#include <ATen/core/Tensor.h>
|
||||
#include <ATen/core/dispatch/DispatchKeyExtractor.h>
|
||||
#include <torch/library.h>
|
||||
|
||||
#ifndef AT_PER_OPERATOR_HEADERS
|
||||
#include <ATen/Operators.h>
|
||||
#else
|
||||
|
||||
${ops_headers}
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace {
|
||||
|
||||
${backend_select_method_definitions}
|
||||
|
||||
TORCH_LIBRARY_IMPL(aten, BackendSelect, m) {
|
||||
${backend_select_function_registrations};
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // at
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
#include <torch/csrc/jit/runtime/operator.h>
|
||||
#include <torch/csrc/jit/runtime/custom_operator.h>
|
||||
#include <torch/csrc/jit/runtime/register_ops_utils.h>
|
||||
|
||||
#include <ATen/UnboxingFunctions.h>
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
// NOTE [Sharded File]: This file is generated in a sharded fashion to speed up
|
||||
// incremental rebuilds. See the comment at the top of
|
||||
// templates/VariableType.cpp for an analogous, in-depth discussion.
|
||||
//
|
||||
// Generated by tools/jit/gen_unboxing.py. This file registers all ATen ops into JIT op registry instead of c10
|
||||
// dispatcher. JIT op registry only takes boxed kernels, so we are calling unboxing functions in UnboxingFunctions.h
|
||||
// to cast arguments into C++ types (instead of IValue) and delegate to unboxed kernels.
|
||||
|
||||
namespace torch { namespace jit {
|
||||
|
||||
using autograd::Variable;
|
||||
using autograd::variable_list;
|
||||
using at::Scalar;
|
||||
using at::ScalarType;
|
||||
using at::Tensor;
|
||||
using at::TensorOptions;
|
||||
using at::DeviceGuard;
|
||||
|
||||
using ::c10::fmap;
|
||||
using ::c10::filter;
|
||||
|
||||
namespace {
|
||||
|
||||
RegisterOperators reg({
|
||||
|
||||
// Generated operators
|
||||
${unboxed_ops}
|
||||
});
|
||||
|
||||
} // anon namespace
|
||||
|
||||
|
||||
}} // namespace torch::jit
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
${ns_prologue}
|
||||
|
||||
// NB: TORCH_LIBRARY_IMPL must be in an anonymous namespace to avoid
|
||||
// ambiguity with conflicting identifiers that may have been defined in
|
||||
// at namespace already.
|
||||
namespace {
|
||||
|
||||
${dispatch_anonymous_definitions}
|
||||
|
||||
${static_init_dispatch_registrations}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
${deferred_dispatch_registrations}
|
||||
|
||||
namespace ${dispatch_namespace} {
|
||||
|
||||
${dispatch_namespaced_definitions}
|
||||
|
||||
} // namespace ${dispatch_namespace}
|
||||
|
||||
${ns_epilogue}
|
||||
+51
@@ -0,0 +1,51 @@
|
||||
// an external backend might generate file within its code tree
|
||||
// and check all the source files within the tree with clang-format.
|
||||
// so, disable it since the backend might have a different config.
|
||||
// clang-format off
|
||||
|
||||
// NOTE: This condition is true for all PyTorch internal libraries, it
|
||||
// just excludes external projects such as torch_xla which
|
||||
// reuse some of the PyTorch codegen machinery.
|
||||
#if defined(CAFFE2_BUILD_MAIN_LIB) || \
|
||||
defined(TORCH_CUDA_BUILD_MAIN_LIB) || \
|
||||
defined(TORCH_XPU_BUILD_MAIN_LIB)
|
||||
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
#endif
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <c10/core/TensorImpl.h>
|
||||
#include <c10/core/Allocator.h>
|
||||
#include <ATen/DeviceGuard.h>
|
||||
#include <ATen/NamedTensorUtils.h>
|
||||
#include <ATen/Utils.h>
|
||||
#include <ATen/WrapDimUtils.h>
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <c10/util/ExclusivelyOwned.h>
|
||||
#include <c10/util/Half.h>
|
||||
#include <c10/core/UndefinedTensorImpl.h>
|
||||
#include <optional>
|
||||
#include <ATen/Tensor.h>
|
||||
#include <ATen/native/Resize.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include <ATen/Config.h>
|
||||
#include <ATen/core/op_registration/adaption.h>
|
||||
#include <torch/library.h>
|
||||
$extra_cuda_headers
|
||||
$external_backend_headers
|
||||
$dispatch_headers
|
||||
$ops_headers
|
||||
|
||||
namespace at {
|
||||
namespace {
|
||||
$dispatch_helpers
|
||||
} // namespace
|
||||
} // namespace at
|
||||
|
||||
// See template file RegisterDispatchDefinitions.ini
|
||||
$dispatch_definitions
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/core/LegacyTypeDispatch.h>
|
||||
#include <ATen/EmptyTensor.h>
|
||||
#include <ATen/FunctionalTensorWrapper.h>
|
||||
#include <ATen/ViewMetaClasses.h>
|
||||
#include <ATen/MemoryOverlap.h>
|
||||
#include <torch/library.h>
|
||||
|
||||
#include <c10/util/env.h>
|
||||
#ifndef AT_PER_OPERATOR_HEADERS
|
||||
#include <ATen/Operators.h>
|
||||
#include <ATen/NativeFunctions.h>
|
||||
#else
|
||||
// needed for the meta tensor calls to get stride info in functionalization
|
||||
#include <ATen/ops/empty_strided_native.h>
|
||||
// needed for special handling of copy_().
|
||||
// See Note [functionalizating copy_() and not preserving strides]
|
||||
#include <ATen/ops/to_ops.h>
|
||||
#include <ATen/ops/expand_copy_ops.h>
|
||||
|
||||
$ops_headers
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
namespace functionalization {
|
||||
|
||||
// This keyset is used by functionalization when it calls into meta kernels
|
||||
// to accurately propagate stride metadata.
|
||||
// Exclude any modes: the purpose of calling into meta kernels is only as an implementation
|
||||
// detail to perform shape inference, and we don't want any modal keys to run.
|
||||
// Specifically, we want to prevent functionalization and Python modes from running.
|
||||
constexpr auto exclude_keys_for_meta_dispatch =
|
||||
c10::functorch_transforms_ks |
|
||||
c10::DispatchKeySet({
|
||||
c10::DispatchKey::FuncTorchDynamicLayerBackMode,
|
||||
c10::DispatchKey::FuncTorchDynamicLayerFrontMode,
|
||||
c10::DispatchKey::Python,
|
||||
c10::DispatchKey::PreDispatch,
|
||||
|
||||
});
|
||||
|
||||
// Helper around at::has_internal_overlap.
|
||||
// The ATen util is used in hot-path eager mode: it's always fast,
|
||||
// but might return TOO_HARD sometimes.
|
||||
// During functionalization, we're ok taking a bit longer
|
||||
// to detect memory overlap.
|
||||
inline bool has_internal_overlap_helper(const at::Tensor t) {
|
||||
auto has_overlap = at::has_internal_overlap(t);
|
||||
if (has_overlap == at::MemOverlap::Yes) return true;
|
||||
if (has_overlap == at::MemOverlap::No) return false;
|
||||
return false;
|
||||
}
|
||||
|
||||
|
||||
inline Tensor to_meta(const Tensor& t) {
|
||||
if (!t.defined()) return t;
|
||||
return at::native::empty_strided_meta_symint(t.sym_sizes(), t.sym_strides(),
|
||||
/*dtype=*/t.scalar_type(), /*layout=*/t.layout(),
|
||||
/*device=*/c10::Device(kMeta), /*pin_memory=*/std::nullopt);
|
||||
}
|
||||
|
||||
inline std::optional<Tensor> to_meta(const std::optional<Tensor>& t) {
|
||||
if (t.has_value()) {
|
||||
return to_meta(*t);
|
||||
}
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
inline std::vector<Tensor> to_meta(at::ITensorListRef t_list) {
|
||||
std::vector<Tensor> outputs;
|
||||
outputs.reserve(t_list.size());
|
||||
for (const auto& tensor : t_list) {
|
||||
outputs.push_back(to_meta(tensor));
|
||||
}
|
||||
return outputs;
|
||||
}
|
||||
|
||||
inline c10::List<Tensor> to_meta(const c10::List<Tensor>& t_list) {
|
||||
c10::List<Tensor> outputs;
|
||||
outputs.reserve(t_list.size());
|
||||
for (const auto i : c10::irange(t_list.size())) {
|
||||
outputs.push_back(to_meta(t_list[i]));
|
||||
}
|
||||
return outputs;
|
||||
}
|
||||
|
||||
inline c10::List<::std::optional<Tensor>> to_meta(const c10::List<::std::optional<Tensor>>& t_list) {
|
||||
c10::List<::std::optional<Tensor>> outputs;
|
||||
outputs.reserve(t_list.size());
|
||||
for (const auto i : c10::irange(t_list.size())) {
|
||||
outputs.push_back(to_meta(t_list[i]));
|
||||
}
|
||||
return outputs;
|
||||
}
|
||||
|
||||
static bool disable_meta_reference() {
|
||||
static auto env = c10::utils::get_env("TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE");
|
||||
return env == "1";
|
||||
}
|
||||
|
||||
|
||||
${func_definitions}
|
||||
|
||||
} // namespace functionalization
|
||||
|
||||
namespace {
|
||||
|
||||
TORCH_LIBRARY_IMPL(aten, Functionalize, m) {
|
||||
${func_registrations};
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
} // namespace at
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
// ${generated_comment}
|
||||
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
#include <torch/library.h>
|
||||
|
||||
namespace at {
|
||||
TORCH_LIBRARY(aten, m) {
|
||||
${aten_schema_registrations};
|
||||
// Distributed Ops
|
||||
// Implementations located in torch/csrc/jit/runtime/register_distributed_ops.cpp
|
||||
m.def("get_gradients(int context_id) -> Dict(Tensor, Tensor)");
|
||||
}
|
||||
${schema_registrations}
|
||||
} // namespace at
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
// This file contains all native_functions that can be registered to
|
||||
// and the schema string that they should be registered with
|
||||
|
||||
${registration_declarations}
|
||||
+758
@@ -0,0 +1,758 @@
|
||||
#pragma once
|
||||
|
||||
#ifdef TORCH_ASSERT_NO_OPERATORS
|
||||
#error This change adds a dependency on native_functions.yaml, \
|
||||
meaning the file will need to be re-compiled every time an operator \
|
||||
is changed or added. Consider if your change would be better placed in \
|
||||
another file, or if a more specific header might achieve the same goal. \
|
||||
See NOTE: [Tensor vs. TensorBase]
|
||||
#endif
|
||||
|
||||
#include <c10/core/Device.h>
|
||||
#include <c10/core/Layout.h>
|
||||
#include <c10/core/MemoryFormat.h>
|
||||
#include <c10/core/QScheme.h>
|
||||
#include <c10/core/Stream.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/core/ScalarTypeToTypeMeta.h>
|
||||
#include <c10/core/Storage.h>
|
||||
#include <c10/core/TensorImpl.h>
|
||||
#include <c10/core/UndefinedTensorImpl.h>
|
||||
#include <c10/core/WrapDimMinimal.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <c10/util/ExclusivelyOwned.h>
|
||||
#include <c10/util/Deprecated.h>
|
||||
#include <c10/util/MaybeOwned.h>
|
||||
#include <optional>
|
||||
#include <c10/util/OptionalArrayRef.h>
|
||||
#include <c10/util/intrusive_ptr.h>
|
||||
#include <c10/macros/Export.h>
|
||||
#include <c10/macros/Macros.h>
|
||||
#include <ATen/core/CheckMemoryFormat.h>
|
||||
#include <ATen/core/DeprecatedTypePropertiesRegistry.h>
|
||||
#include <ATen/core/DeprecatedTypeProperties.h>
|
||||
#include <ATen/core/NamedTensor.h>
|
||||
#include <ATen/core/QuantizerBase.h>
|
||||
#include <c10/core/SymInt.h>
|
||||
#include <ATen/core/TensorAccessor.h>
|
||||
#include <ATen/core/TensorBase.h>
|
||||
|
||||
|
||||
#include <ATen/MethodOperators.h>
|
||||
|
||||
namespace c10{
|
||||
template<class T> class List;
|
||||
template<class T> class IListRef;
|
||||
}
|
||||
namespace at {
|
||||
struct Generator;
|
||||
struct Type;
|
||||
class DeprecatedTypeProperties;
|
||||
class Tensor;
|
||||
} // namespace at
|
||||
namespace at {
|
||||
namespace indexing {
|
||||
struct TensorIndex;
|
||||
} // namespace indexing
|
||||
} // namespace at
|
||||
|
||||
namespace torch { namespace autograd {
|
||||
|
||||
struct Node;
|
||||
|
||||
}} // namespace torch::autograd
|
||||
|
||||
namespace at {
|
||||
|
||||
class OptionalTensorRef;
|
||||
class TensorRef;
|
||||
class Tensor;
|
||||
using TensorList = ArrayRef<Tensor>;
|
||||
using ITensorList = c10::IListRef<Tensor>;
|
||||
|
||||
using Stream = c10::Stream;
|
||||
|
||||
// Tensor is a "generic" object holding a pointer to the underlying TensorImpl object, which
|
||||
// has an embedded reference count. In this way, Tensor is similar to boost::intrusive_ptr.
|
||||
//
|
||||
// For example:
|
||||
//
|
||||
// void func(Tensor a) {
|
||||
// Tensor b = a;
|
||||
// ...
|
||||
// }
|
||||
//
|
||||
// In this example, when we say Tensor b = a, we are creating a new object that points to the
|
||||
// same underlying TensorImpl, and bumps its reference count. When b goes out of scope, the
|
||||
// destructor decrements the reference count by calling release() on the TensorImpl it points to.
|
||||
// The existing constructors, operator overloads, etc. take care to implement the correct semantics.
|
||||
//
|
||||
// Note that Tensor can also be NULL, i.e. it is not associated with any underlying TensorImpl, and
|
||||
// special care must be taken to handle this.
|
||||
class TORCH_API Tensor: public TensorBase {
|
||||
protected:
|
||||
// Create a Tensor with a +0 reference count. Special care must be
|
||||
// taken to avoid decrementing this reference count at destruction
|
||||
// time. Intended to support MaybeOwnedTraits<Tensor>.
|
||||
explicit Tensor(unsafe_borrow_t, const TensorBase& rhs): TensorBase(unsafe_borrow_t{}, rhs) {}
|
||||
friend MaybeOwnedTraits<Tensor>;
|
||||
friend OptionalTensorRef;
|
||||
friend TensorRef;
|
||||
|
||||
public:
|
||||
Tensor() = default;
|
||||
// This constructor should not be used by end users and is an implementation
|
||||
// detail invoked by autogenerated code.
|
||||
explicit Tensor(
|
||||
c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> tensor_impl)
|
||||
: TensorBase(std::move(tensor_impl)) {}
|
||||
Tensor(const Tensor &tensor) = default;
|
||||
Tensor(Tensor &&tensor) = default;
|
||||
|
||||
// Implicitly move-constructible from TensorBase, but must be explicit to increase refcount
|
||||
explicit Tensor(const TensorBase &base): TensorBase(base) {}
|
||||
/*implicit*/ Tensor(TensorBase &&base): TensorBase(std::move(base)) {}
|
||||
|
||||
// Creates a new wrapper from TensorImpl. Intentionally a free method because
|
||||
// it should be used with care. Checks necessary invariants
|
||||
static Tensor wrap_tensor_impl(
|
||||
c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> tensor_impl) {
|
||||
return TensorBase::wrap_tensor_impl(std::move(tensor_impl));
|
||||
}
|
||||
|
||||
Tensor contiguous(MemoryFormat memory_format=MemoryFormat::Contiguous) const {
|
||||
return TensorBase::contiguous(memory_format);
|
||||
}
|
||||
|
||||
Tensor conj() const {
|
||||
if (!this->is_complex()) {
|
||||
return *this;
|
||||
}
|
||||
|
||||
C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")
|
||||
switch (this->layout()) {
|
||||
case at::kSparse:
|
||||
case at::kSparseCsr:
|
||||
case at::kSparseCsc:
|
||||
case at::kSparseBsr:
|
||||
case at::kSparseBsc:
|
||||
return this->conj_physical();
|
||||
default:
|
||||
return this->_conj();
|
||||
}
|
||||
C10_DIAGNOSTIC_POP()
|
||||
}
|
||||
|
||||
// Aliased by Dimname overloads, so need explicit using
|
||||
using TensorBase::size;
|
||||
using TensorBase::sym_size;
|
||||
using TensorBase::stride;
|
||||
|
||||
/// Should be used if *this can reasonably be expected to be contiguous and
|
||||
/// performance is important.
|
||||
/// Compared to contiguous, it saves a reference count
|
||||
/// increment/decrement if *this is already contiguous, at the cost
|
||||
/// in all cases of an extra pointer of stack usage, an extra branch
|
||||
/// to access, and an extra branch at destruction time.
|
||||
c10::MaybeOwned<Tensor> expect_contiguous(MemoryFormat memory_format=MemoryFormat::Contiguous) const &;
|
||||
|
||||
// Use .contiguous() instead. Trying to borrow from a prvalue Tensor
|
||||
// will only lead to trouble and dangling references.
|
||||
c10::MaybeOwned<Tensor> expect_contiguous(MemoryFormat memory_format=MemoryFormat::Contiguous) && = delete;
|
||||
|
||||
// The following overloads are very intriguing. Consider the following
|
||||
// program:
|
||||
//
|
||||
// x[1] = 3;
|
||||
//
|
||||
// We would expect that the first entry of x is written to 3. But how can we
|
||||
// actually achieve this? x[1] evaluates to a tensor...
|
||||
//
|
||||
// The answer is, using a ref-qualifier. x[1] is an rvalue, which cannot be
|
||||
// (profitably) assigned to in the traditional sense, so we overload
|
||||
// assignment to mean, "Actually, copy 3 into the tensor data." This is done
|
||||
// with an rvalue-reference ref-qualified overload (the methods with && at the
|
||||
// end of their type.)
|
||||
//
|
||||
// There's one more fly in the ointment: We also want
|
||||
//
|
||||
// Tensor x = y;
|
||||
//
|
||||
// to work, and we want it NOT to copy. So we need a traditional operator=
|
||||
// overload. But we MUST specify a mutable lvalue ref-qualifier, to
|
||||
// disambiguate the traditional overload from the rvalue-reference
|
||||
// ref-qualified overload. Otherwise, it will be ambiguous, because
|
||||
// a non ref-qualified method is eligible for all situations.
|
||||
|
||||
// Unfortunately, we have to write these constructors out manually
|
||||
// to work around an MSVC bug:
|
||||
// error C2580: 'at::Tensor &at::Tensor::operator =(const at::Tensor &) &':
|
||||
// multiple versions of a defaulted special member functions are not allowed
|
||||
// Tensor& operator=(const Tensor&) & = default;
|
||||
// Tensor& operator=(Tensor&&) & = default;
|
||||
|
||||
// Also MSVC will wrongly issue the following warning with the aforementioned fix
|
||||
// warning C4522: 'at::Tensor': multiple assignment operators specified
|
||||
// Let's just skip the warning.
|
||||
//
|
||||
// TODO: temporarily disabled
|
||||
|
||||
Tensor& operator=(const TensorBase& x) & noexcept {
|
||||
impl_ = x.getIntrusivePtr();
|
||||
return *this;
|
||||
}
|
||||
Tensor& operator=(TensorBase&& x) & noexcept {
|
||||
impl_ = x.unsafeReleaseIntrusivePtr();
|
||||
return *this;
|
||||
}
|
||||
|
||||
Tensor& operator=(const Tensor &x) & noexcept {
|
||||
return operator=(static_cast<const TensorBase&>(x));
|
||||
}
|
||||
Tensor& operator=(Tensor &&x) & noexcept {
|
||||
return operator=(static_cast<TensorBase&&>(x));
|
||||
}
|
||||
|
||||
Tensor& operator=(const Scalar &v) && {
|
||||
return fill_(v);
|
||||
}
|
||||
Tensor& operator=(const Tensor &rhs) && {
|
||||
return copy_(rhs);
|
||||
}
|
||||
|
||||
// NOLINTNEXTLINE(performance-noexcept-move-constructor)
|
||||
Tensor& operator=(Tensor&& rhs) && {
|
||||
return copy_(rhs);
|
||||
}
|
||||
|
||||
C10_DEPRECATED_MESSAGE("Tensor.type() is deprecated. Instead use Tensor.options(), which in many cases (e.g. in a constructor) is a drop-in replacement. If you were using data from type(), that is now available from Tensor itself, so instead of tensor.type().scalar_type(), use tensor.scalar_type() instead and instead of tensor.type().backend() use tensor.device().")
|
||||
DeprecatedTypeProperties & type() const {
|
||||
return globalDeprecatedTypePropertiesRegistry().getDeprecatedTypeProperties(
|
||||
dispatchKeyToBackend(legacyExtractDispatchKey(key_set())),
|
||||
scalar_type());
|
||||
}
|
||||
|
||||
Tensor toType(ScalarType t) const {
|
||||
return to(options().dtype(t), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
// TODO: Deprecate me
|
||||
Tensor toBackend(Backend b) const {
|
||||
return to(options().device(backendToDeviceType(b)).layout(layout_from_backend(b)), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
C10_DEPRECATED_MESSAGE("Tensor.is_variable() is deprecated; everything is a variable now. (If you want to assert that variable has been appropriately handled already, use at::impl::variable_excluded_from_dispatch())")
|
||||
bool is_variable() const noexcept {
|
||||
return !at::impl::variable_excluded_from_dispatch();
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
C10_DEPRECATED_MESSAGE("Tensor.data<T>() is deprecated. Please use Tensor.data_ptr<T>() instead.")
|
||||
T * data() const {
|
||||
return data_ptr<T>();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T item() const;
|
||||
|
||||
template<typename T, size_t N, template <typename U> class PtrTraits = DefaultPtrTraits, typename index_t = int64_t>
|
||||
C10_DEPRECATED_MESSAGE("packed_accessor is deprecated, use packed_accessor32 or packed_accessor64 instead")
|
||||
GenericPackedTensorAccessor<T,N,PtrTraits,index_t> packed_accessor() const & {
|
||||
return generic_packed_accessor<T,N,PtrTraits,index_t>();
|
||||
}
|
||||
template<typename T, size_t N, template <typename U> class PtrTraits = DefaultPtrTraits, typename index_t = int64_t>
|
||||
C10_DEPRECATED_MESSAGE("packed_accessor is deprecated, use packed_accessor32 or packed_accessor64 instead")
|
||||
GenericPackedTensorAccessor<T,N,PtrTraits,index_t> packed_accessor() && = delete;
|
||||
|
||||
Tensor operator~() const {
|
||||
return bitwise_not();
|
||||
}
|
||||
Tensor operator-() const {
|
||||
return neg();
|
||||
}
|
||||
Tensor& operator+=(const Tensor & other) {
|
||||
return add_(other);
|
||||
}
|
||||
Tensor& operator+=(const Scalar & other) {
|
||||
return add_(other);
|
||||
}
|
||||
Tensor& operator-=(const Tensor & other) {
|
||||
return sub_(other);
|
||||
}
|
||||
Tensor& operator-=(const Scalar & other) {
|
||||
return sub_(other);
|
||||
}
|
||||
Tensor& operator*=(const Tensor & other) {
|
||||
return mul_(other);
|
||||
}
|
||||
Tensor& operator*=(const Scalar & other) {
|
||||
return mul_(other);
|
||||
}
|
||||
Tensor& operator/=(const Tensor & other) {
|
||||
return div_(other);
|
||||
}
|
||||
Tensor& operator/=(const Scalar & other) {
|
||||
return div_(other);
|
||||
}
|
||||
Tensor& operator&=(const Tensor & other) {
|
||||
return bitwise_and_(other);
|
||||
}
|
||||
Tensor& operator|=(const Tensor & other) {
|
||||
return bitwise_or_(other);
|
||||
}
|
||||
Tensor& operator^=(const Tensor & other) {
|
||||
return bitwise_xor_(other);
|
||||
}
|
||||
Tensor operator[](const Scalar & index) const {
|
||||
if (!index.isIntegral(false)) {
|
||||
TORCH_CHECK_INDEX(false, "Can only index tensors with integral scalars");
|
||||
}
|
||||
return this->operator[](index.toLong());
|
||||
}
|
||||
Tensor operator[](const Tensor & index) const {
|
||||
// These properties are checked in the Scalar constructor, but we already
|
||||
// check them here to provide more useful diagnostics for the user.
|
||||
if (!index.defined()) {
|
||||
TORCH_CHECK_INDEX(false, "Can only index with tensors that are defined");
|
||||
}
|
||||
if (index.dim() != 0) {
|
||||
TORCH_CHECK_INDEX(false,
|
||||
"Can only index with tensors that are scalars (zero-dim)");
|
||||
}
|
||||
// The Scalar(Tensor) constructor is explicit, so we need to call it.
|
||||
return this->operator[](index.item());
|
||||
}
|
||||
Tensor operator[](int64_t index) const {
|
||||
return select(0, index);
|
||||
}
|
||||
|
||||
Tensor index(ArrayRef<at::indexing::TensorIndex> indices) const;
|
||||
Tensor index(std::initializer_list<at::indexing::TensorIndex> indices) const;
|
||||
|
||||
Tensor & index_put_(ArrayRef<at::indexing::TensorIndex> indices, Tensor const & rhs);
|
||||
Tensor & index_put_(ArrayRef<at::indexing::TensorIndex> indices, const Scalar& v);
|
||||
Tensor & index_put_(std::initializer_list<at::indexing::TensorIndex> indices, Tensor const & rhs);
|
||||
Tensor & index_put_(std::initializer_list<at::indexing::TensorIndex> indices, const Scalar& v);
|
||||
|
||||
Tensor cpu() const {
|
||||
return to(options().device(c10::DeviceType::CPU), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
// TODO: The Python version also accepts arguments
|
||||
Tensor cuda() const {
|
||||
return to(options().device(c10::DeviceType::CUDA), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
Tensor hip() const {
|
||||
return to(options().device(c10::DeviceType::HIP), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
Tensor ve() const {
|
||||
return to(options().device(c10::DeviceType::VE), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
Tensor vulkan() const {
|
||||
return to(options().device(c10::DeviceType::Vulkan), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
Tensor metal() const {
|
||||
return to(options().device(c10::DeviceType::Metal), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
Tensor meta() const {
|
||||
return to(options().device(c10::DeviceType::Meta), /*non_blocking*/ false, /*copy*/ false);
|
||||
}
|
||||
|
||||
// ~~~~~ Autograd API ~~~~~
|
||||
|
||||
/// \fn bool is_leaf() const;
|
||||
///
|
||||
/// All Tensors that have `requires_grad()` which is ``false`` will be leaf Tensors by convention.
|
||||
///
|
||||
/// For Tensors that have `requires_grad()` which is ``true``, they will be leaf Tensors if they were
|
||||
/// created by the user. This means that they are not the result of an operation and so
|
||||
/// `grad_fn()` is `nullptr`.
|
||||
///
|
||||
/// Only leaf Tensors will have their `grad()` populated during a call to `backward()`.
|
||||
/// To get `grad()` populated for non-leaf Tensors, you can use `retain_grad()`.
|
||||
///
|
||||
/// Example:
|
||||
/// @code
|
||||
/// auto a = torch::rand(10, torch::requires_grad());
|
||||
/// std::cout << a.is_leaf() << std::endl; // prints `true`
|
||||
///
|
||||
/// auto b = torch::rand(10, torch::requires_grad()).to(torch::kCUDA);
|
||||
/// std::cout << b.is_leaf() << std::endl; // prints `false`
|
||||
/// // b was created by the operation that cast a cpu Tensor into a cuda Tensor
|
||||
///
|
||||
/// auto c = torch::rand(10, torch::requires_grad()) + 2;
|
||||
/// std::cout << c.is_leaf() << std::endl; // prints `false`
|
||||
/// // c was created by the addition operation
|
||||
///
|
||||
/// auto d = torch::rand(10).cuda();
|
||||
/// std::cout << d.is_leaf() << std::endl; // prints `true`
|
||||
/// // d does not require gradients and so has no operation creating it (that is tracked by the autograd engine)
|
||||
///
|
||||
/// auto e = torch::rand(10).cuda().requires_grad_();
|
||||
/// std::cout << e.is_leaf() << std::endl; // prints `true`
|
||||
/// // e requires gradients and has no operations creating it
|
||||
///
|
||||
/// auto f = torch::rand(10, torch::device(torch::kCUDA).requires_grad(true));
|
||||
/// std::cout << f.is_leaf() << std::endl; // prints `true`
|
||||
/// // f requires grad, has no operation creating it
|
||||
/// @endcode
|
||||
|
||||
/// \fn void backward(const Tensor & gradient={}, std::optional<bool> retain_graph=std::nullopt, bool create_graph=false, std::optional<TensorList> inputs=std::nullopt) const;
|
||||
///
|
||||
/// Computes the gradient of current tensor with respect to graph leaves.
|
||||
///
|
||||
/// The graph is differentiated using the chain rule. If the tensor is
|
||||
/// non-scalar (i.e. its data has more than one element) and requires
|
||||
/// gradient, the function additionally requires specifying ``gradient``.
|
||||
/// It should be a tensor of matching type and location, that contains
|
||||
/// the gradient of the differentiated function w.r.t. this Tensor.
|
||||
///
|
||||
/// This function accumulates gradients in the leaves - you might need to
|
||||
/// zero them before calling it.
|
||||
///
|
||||
/// \param gradient Gradient w.r.t. the
|
||||
/// tensor. If it is a tensor, it will be automatically converted
|
||||
/// to a Tensor that does not require grad unless ``create_graph`` is True.
|
||||
/// None values can be specified for scalar Tensors or ones that
|
||||
/// don't require grad. If a None value would be acceptable then
|
||||
/// this argument is optional.
|
||||
/// \param retain_graph If ``false``, the graph used to compute
|
||||
/// the grads will be freed. Note that in nearly all cases setting
|
||||
/// this option to True is not needed and often can be worked around
|
||||
/// in a much more efficient way. Defaults to the value of
|
||||
/// ``create_graph``.
|
||||
/// \param create_graph If ``true``, graph of the derivative will
|
||||
/// be constructed, allowing to compute higher order derivative
|
||||
/// products. Defaults to ``false``.
|
||||
/// \param inputs Inputs w.r.t. which the gradient will be accumulated into
|
||||
/// ``at::Tensor::grad``. All other Tensors will be ignored. If not
|
||||
/// provided, the gradient is accumulated into all the leaf Tensors
|
||||
/// that were used to compute the current tensor.
|
||||
/// When inputs are provided and a given input is not a leaf,
|
||||
/// the current implementation will call its grad_fn (even though it is not strictly needed to get this gradients).
|
||||
/// It is an implementation detail on which the user should not rely.
|
||||
/// See https://github.com/pytorch/pytorch/pull/60521#issuecomment-867061780 for more details.
|
||||
void backward(const Tensor & gradient={}, std::optional<bool> retain_graph=std::nullopt, bool create_graph=false, std::optional<TensorList> inputs=std::nullopt) const {
|
||||
// NB: Adding this wrapper to _backward here because we'd like our
|
||||
// 'backwards' api to accept the 'inputs' argument optionally. Since code gen
|
||||
// currently does not support optional of TensorList our approach is to replace
|
||||
// backward in native_functions.yaml with _backward and call it here instead.
|
||||
if (inputs.has_value()) {
|
||||
TORCH_CHECK(inputs.value().size() > 0, "'inputs' argument to backward cannot be empty")
|
||||
this->_backward(inputs.value(), gradient, retain_graph, create_graph);
|
||||
} else {
|
||||
this->_backward({}, gradient, retain_graph, create_graph);
|
||||
}
|
||||
}
|
||||
|
||||
/// \fn Tensor detach() const;
|
||||
///
|
||||
/// Returns a new Tensor, detached from the current graph.
|
||||
/// The result will never require gradient.
|
||||
|
||||
/// \fn Tensor & detach_() const;
|
||||
///
|
||||
/// Detaches the Tensor from the graph that created it, making it a leaf.
|
||||
/// Views cannot be detached in-place.
|
||||
|
||||
/// \fn void retain_grad() const;
|
||||
///
|
||||
/// Enables this Tensor to have their :attr:`grad` populated during
|
||||
/// :func:`backward`. This is a no-op for leaf tensors.
|
||||
|
||||
/// \fn bool retains_grad() const;
|
||||
///
|
||||
/// Is ``true`` if this Tensor is non-leaf and its :attr:`grad` is enabled to be
|
||||
/// populated during :func:`backward`, ``false`` otherwise.
|
||||
|
||||
const Tensor& set_requires_grad(bool requires_grad) const {
|
||||
TensorBase::set_requires_grad(requires_grad);
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Return a mutable reference to the gradient. This is conventionally
|
||||
/// used as `t.grad() = x` to set a gradient to a completely new tensor.
|
||||
/// Note that this function work with a non-const Tensor and is not
|
||||
/// thread safe.
|
||||
Tensor& mutable_grad() const {
|
||||
return impl_->mutable_grad();
|
||||
}
|
||||
|
||||
/// This function returns an undefined tensor by default and returns a defined tensor
|
||||
/// the first time a call to `backward()` computes gradients for this Tensor.
|
||||
/// The attribute will then contain the gradients computed and future calls
|
||||
/// to `backward()` will accumulate (add) gradients into it.
|
||||
const Tensor& grad() const {
|
||||
const Tensor& maybe_grad = impl_->grad();
|
||||
if (!is_leaf() && !retains_grad() && !maybe_grad.defined()) {
|
||||
TORCH_WARN(
|
||||
"The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad "
|
||||
"attribute won't be populated during autograd.backward(). If you indeed want the .grad "
|
||||
"field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. "
|
||||
"If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor "
|
||||
"instead. See github.com/pytorch/pytorch/pull/30531 for more information.");
|
||||
}
|
||||
return maybe_grad;
|
||||
}
|
||||
|
||||
// The Forward AD API functions below are low level and are not to be used by end
|
||||
// users who should use the API provided in torch/csrc/autograd.h
|
||||
|
||||
/// This function returns the forward gradient for this Tensor at the given level.
|
||||
const Tensor& _fw_grad(uint64_t level) const {
|
||||
return impl_->_fw_grad(level, *this);
|
||||
}
|
||||
|
||||
/// This function can be used to set the value of the forward grad.
|
||||
/// Note that the given new_grad might not be used directly if it has different
|
||||
/// metadata (size/stride/storage offset) compared to this Tensor. In that case,
|
||||
/// new_grad content will be copied into a new Tensor
|
||||
void _set_fw_grad(const TensorBase& new_grad, uint64_t level, bool is_inplace_op) const {
|
||||
impl_->_set_fw_grad(new_grad, *this, level, is_inplace_op);
|
||||
}
|
||||
|
||||
|
||||
// STOP. Thinking of adding a method here, which only makes use
|
||||
// of other ATen methods? Define it in native_functions.yaml.
|
||||
|
||||
//example
|
||||
//Tensor * add(Tensor & b);
|
||||
${tensor_method_declarations}
|
||||
|
||||
// Special C++ only overloads for std()-like functions (See gh-40287)
|
||||
// These are needed because int -> bool conversion takes precedence over int -> IntArrayRef
|
||||
// So, for example std(0) would select the std(unbiased=False) overload
|
||||
|
||||
Tensor var(int dim) const {
|
||||
return var(IntArrayRef{dim});
|
||||
}
|
||||
|
||||
Tensor std(int dim) const {
|
||||
return std(IntArrayRef{dim});
|
||||
}
|
||||
|
||||
// We changed .dtype() to return a TypeMeta in #12766. Ideally, we want the
|
||||
// at::kDouble and its friends to be TypeMeta's, but that hasn't happened yet.
|
||||
// Before that change, we make this method to maintain BC for C++ usage like
|
||||
// `x.to(y.dtype)`.
|
||||
// TODO: remove following two after at::kDouble and its friends are TypeMeta's.
|
||||
inline Tensor to(caffe2::TypeMeta type_meta, bool non_blocking=false, bool copy=false) const {
|
||||
return this->to(/*scalar_type=*/typeMetaToScalarType(type_meta), non_blocking, copy);
|
||||
}
|
||||
inline Tensor to(Device device, caffe2::TypeMeta type_meta, bool non_blocking=false, bool copy=false) const {
|
||||
return this->to(device, /*scalar_type=*/typeMetaToScalarType(type_meta), non_blocking, copy);
|
||||
}
|
||||
|
||||
template <typename F, typename... Args>
|
||||
decltype(auto) m(F func, Args&&... params) const {
|
||||
return func(*this, std::forward<Args>(params)...);
|
||||
}
|
||||
|
||||
/// NOTE: This is similar to the legacy `.data()` function on `Variable`, and is intended
|
||||
/// to be used from functions that need to access the `Variable`'s equivalent `Tensor`
|
||||
/// (i.e. `Tensor` that shares the same storage and tensor metadata with the `Variable`).
|
||||
///
|
||||
/// One notable difference with the legacy `.data()` function is that changes to the
|
||||
/// returned `Tensor`'s tensor metadata (e.g. sizes / strides / storage / storage_offset)
|
||||
/// will not update the original `Variable`, due to the fact that this function
|
||||
/// shallow-copies the `Variable`'s underlying TensorImpl.
|
||||
at::Tensor tensor_data() const {
|
||||
return TensorBase::tensor_data();
|
||||
}
|
||||
|
||||
/// NOTE: `var.variable_data()` in C++ has the same semantics as `tensor.data`
|
||||
/// in Python, which create a new `Variable` that shares the same storage and
|
||||
/// tensor metadata with the original `Variable`, but with a completely new
|
||||
/// autograd history.
|
||||
///
|
||||
/// NOTE: If we change the tensor metadata (e.g. sizes / strides /
|
||||
/// storage / storage_offset) of a variable created from `var.variable_data()`, those
|
||||
/// changes will not update the original variable `var`. In `.variable_data()`, we set
|
||||
/// `allow_tensor_metadata_change_` to false to make such changes explicitly illegal,
|
||||
/// in order to prevent users from changing metadata of `var.variable_data()`
|
||||
/// and expecting the original variable `var` to also be updated.
|
||||
at::Tensor variable_data() const {
|
||||
return TensorBase::variable_data();
|
||||
}
|
||||
|
||||
// Hooks
|
||||
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
template <typename T>
|
||||
using hook_return_void_t = std::enable_if_t<std::is_void<typename std::invoke_result_t<T&, Tensor>>::value, unsigned>;
|
||||
template <typename T>
|
||||
using hook_return_var_t = std::enable_if_t<std::is_same_v<typename std::invoke_result_t<T&, Tensor>, Tensor>, unsigned>;
|
||||
|
||||
/// Registers a backward hook.
|
||||
///
|
||||
/// The hook will be called every time a gradient with respect to the Tensor is computed.
|
||||
/// The hook should have one of the following signature:
|
||||
/// ```
|
||||
/// hook(Tensor grad) -> Tensor
|
||||
/// ```
|
||||
/// ```
|
||||
/// hook(Tensor grad) -> void
|
||||
/// ```
|
||||
/// The hook should not modify its argument, but it can optionally return a new gradient
|
||||
/// which will be used in place of `grad`.
|
||||
///
|
||||
/// This function returns the index of the hook in the list which can be used to remove hook.
|
||||
///
|
||||
/// Example:
|
||||
/// @code
|
||||
/// auto v = torch::tensor({0., 0., 0.}, torch::requires_grad());
|
||||
/// auto h = v.register_hook([](torch::Tensor grad){ return grad * 2; }); // double the gradient
|
||||
/// v.backward(torch::tensor({1., 2., 3.}));
|
||||
/// // This prints:
|
||||
/// // ```
|
||||
/// // 2
|
||||
/// // 4
|
||||
/// // 6
|
||||
/// // [ CPUFloatType{3} ]
|
||||
/// // ```
|
||||
/// std::cout << v.grad() << std::endl;
|
||||
/// v.remove_hook(h); // removes the hook
|
||||
/// @endcode
|
||||
template <typename T>
|
||||
hook_return_void_t<T> register_hook(T&& hook) const;
|
||||
template <typename T>
|
||||
hook_return_var_t<T> register_hook(T&& hook) const;
|
||||
|
||||
// Variable methods
|
||||
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Tensor data() const {
|
||||
return TensorBase::data();
|
||||
}
|
||||
|
||||
void _backward(TensorList inputs, const std::optional<Tensor>& gradient, std::optional<bool> keep_graph, bool create_graph) const;
|
||||
|
||||
const Tensor& requires_grad_(bool _requires_grad=true) const {
|
||||
TensorBase::requires_grad_(_requires_grad);
|
||||
return *this;
|
||||
}
|
||||
};
|
||||
|
||||
namespace detail {
|
||||
// Helper creator for Tensor class which doesn't requires the users to pass
|
||||
// in an intrusive_ptr instead it just converts the argument passed to
|
||||
// requested intrusive_ptr type.
|
||||
template <typename T, typename... Args>
|
||||
Tensor make_tensor(Args&&... args) {
|
||||
return Tensor(c10::make_intrusive<T>(std::forward<Args>(args)...));
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
} // namespace at
|
||||
|
||||
|
||||
namespace at {
|
||||
${tensor_method_definitions}
|
||||
} // namespace at
|
||||
|
||||
|
||||
namespace c10 {
|
||||
template <>
|
||||
struct MaybeOwnedTraits<at::Tensor> {
|
||||
using owned_type = at::Tensor;
|
||||
using borrow_type = at::Tensor;
|
||||
|
||||
static borrow_type createBorrow(const owned_type& from) {
|
||||
// NOTE: this can be implemented without the special
|
||||
// unsafe_borrow_t Tensor constructor as
|
||||
//
|
||||
// return borrow_type(c10::intrusive_ptr<at::TensorImpl, at::UndefinedTensorImpl>::reclaim(from.unsafeGetTensorImpl()));
|
||||
//
|
||||
// but that hurts inlining due to the nullptr check in the
|
||||
// Tensor(c10::intrusive_ptr<...>) constructor. We already know
|
||||
// that from.impl_ isn't null because from is a valid Tensor, so
|
||||
// we needn't do the check again. (using __builtin_assume can
|
||||
// avoid this, but wouldn't be portable to MSVC.)
|
||||
return borrow_type(borrow_type::unsafe_borrow_t{}, from);
|
||||
}
|
||||
|
||||
static void assignBorrow(borrow_type& lhs, const borrow_type& rhs) {
|
||||
lhs.unsafeReleaseTensorImpl();
|
||||
// See above note: this can be implemented with public API
|
||||
// similarly to createBorrow(), but that would hurt inlining.
|
||||
lhs = borrow_type(borrow_type::unsafe_borrow_t{}, rhs);
|
||||
}
|
||||
|
||||
static void destroyBorrow(borrow_type& toDestroy) {
|
||||
toDestroy.unsafeReleaseTensorImpl(); // "leak" it, but it was already +0.
|
||||
}
|
||||
|
||||
static const owned_type& referenceFromBorrow(const borrow_type& borrow) {
|
||||
return borrow;
|
||||
}
|
||||
|
||||
static const owned_type* pointerFromBorrow(const borrow_type& borrow) {
|
||||
return &borrow;
|
||||
}
|
||||
|
||||
static bool debugBorrowIsValid(const borrow_type& /*borrow*/) {
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct ExclusivelyOwnedTraits<at::Tensor> {
|
||||
using repr_type = at::Tensor;
|
||||
using pointer_type = at::Tensor*;
|
||||
using const_pointer_type = const at::Tensor*;
|
||||
|
||||
static repr_type nullRepr() {
|
||||
return at::Tensor();
|
||||
}
|
||||
|
||||
template <class... Args>
|
||||
static repr_type createInPlace(Args&&... args) {
|
||||
return at::Tensor(std::forward<Args>(args)...);
|
||||
}
|
||||
|
||||
static repr_type moveToRepr(at::Tensor&& x) {
|
||||
return std::move(x);
|
||||
}
|
||||
|
||||
static void destroyOwned(at::Tensor& x) {
|
||||
return ExclusivelyOwnedTraits<at::TensorBase>::destroyOwned(x);
|
||||
}
|
||||
|
||||
static at::Tensor take(at::Tensor& x) {
|
||||
return std::move(x);
|
||||
}
|
||||
|
||||
static pointer_type getImpl(repr_type& x) {
|
||||
return &x;
|
||||
}
|
||||
|
||||
static const_pointer_type getImpl(const repr_type& x) {
|
||||
return &x;
|
||||
}
|
||||
};
|
||||
} // namespace c10
|
||||
|
||||
namespace at {
|
||||
|
||||
inline c10::MaybeOwned<Tensor> borrow_from_optional_tensor(
|
||||
const std::optional<Tensor>& opt) {
|
||||
return opt.has_value()
|
||||
? c10::MaybeOwned<Tensor>::borrowed(*opt)
|
||||
: c10::MaybeOwned<Tensor>::owned(std::in_place);
|
||||
}
|
||||
|
||||
inline c10::MaybeOwned<Tensor> Tensor::expect_contiguous(MemoryFormat memory_format) const & {
|
||||
if (is_contiguous(memory_format)) {
|
||||
return c10::MaybeOwned<Tensor>::borrowed(*this);
|
||||
} else {
|
||||
return c10::MaybeOwned<Tensor>::owned(__dispatch_contiguous(memory_format));
|
||||
}
|
||||
}
|
||||
} // namespace at
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
#include <c10/core/Scalar.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
|
||||
#include <string_view>
|
||||
|
||||
namespace at {
|
||||
|
||||
namespace {
|
||||
|
||||
// Verifies the requested type is the same as the Tensor's type.
|
||||
void check_type(const TensorBase& tensor, ScalarType type) {
|
||||
TORCH_CHECK(
|
||||
tensor.scalar_type() == type
|
||||
|| (isQIntType(tensor.scalar_type())
|
||||
&& toUnderlying(tensor.scalar_type()) == type),
|
||||
"expected scalar type ", type, " but found ", tensor.scalar_type());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename T>
|
||||
const T* TensorBase::const_data_ptr() const {
|
||||
using NonConstT = std::remove_const_t<T>;
|
||||
check_type(*this, c10::CppTypeToScalarType<NonConstT>());
|
||||
return this->unsafeGetTensorImpl()->data_ptr_impl<NonConstT>();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* TensorBase::mutable_data_ptr() const {
|
||||
check_type(*this, c10::CppTypeToScalarType<T>());
|
||||
return this->unsafeGetTensorImpl()->mutable_data_ptr_impl<T>();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* TensorBase::data_ptr() const {
|
||||
return this->mutable_data_ptr<T>();
|
||||
}
|
||||
|
||||
#define DEFINE_CAST(T, name) \
|
||||
template TORCH_API const T* TensorBase::const_data_ptr<T>() const; \
|
||||
template TORCH_API const T* TensorBase::const_data_ptr<const T>() const; \
|
||||
template TORCH_API T* TensorBase::mutable_data_ptr() const; \
|
||||
template TORCH_API T* TensorBase::data_ptr() const;
|
||||
|
||||
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST)
|
||||
AT_FORALL_QINT_TYPES(DEFINE_CAST)
|
||||
DEFINE_CAST(uint16_t, UInt16)
|
||||
DEFINE_CAST(uint32_t, UInt32)
|
||||
DEFINE_CAST(uint64_t, UInt64)
|
||||
#undef DEFINE_CAST
|
||||
|
||||
#define DEFINE_ITEM(T, name) \
|
||||
template <> \
|
||||
TORCH_API T Tensor::item() const { \
|
||||
return item().to##name(); \
|
||||
}
|
||||
|
||||
AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_ITEM)
|
||||
#undef DEFINE_ITEM
|
||||
|
||||
} //namespace at
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
#define TORCH_ASSERT_NO_OPERATORS
|
||||
|
||||
#include <ATen/native/DispatchStub.h>
|
||||
#include <ATen/TensorIterator.h>
|
||||
#include <ATen/TensorMeta.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
// NB: this is explicitly copied here (via codegen) rather than
|
||||
// included via NativeFunctions.h to avoid recompiling this file when
|
||||
// NativeFunctions.h changes
|
||||
namespace meta {
|
||||
${meta_declaration}
|
||||
}
|
||||
|
||||
namespace native {
|
||||
${native_declaration}
|
||||
${native_definitions}
|
||||
}} // namespace at::native
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
#define TORCH_ASSERT_NO_OPERATORS
|
||||
|
||||
#include <ATen/native/ufunc/${name}.h>
|
||||
#include <ATen/native/DispatchStub.h>
|
||||
#include <ATen/TensorIterator.h>
|
||||
#include <ATen/native/cpu/Loops.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
|
||||
namespace at {
|
||||
namespace native {
|
||||
${native_definitions}
|
||||
}} // namespace at::native
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
#define TORCH_ASSERT_NO_OPERATORS
|
||||
|
||||
#include <ATen/native/ufunc/${name}.h>
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <ATen/native/DispatchStub.h>
|
||||
#include <c10/core/Scalar.h>
|
||||
${cuda_headers}
|
||||
|
||||
namespace at {
|
||||
|
||||
// NB: this is explicitly copied here (via codegen) rather than
|
||||
// included via NativeFunctions.h to avoid recompiling this file when
|
||||
// NativeFunctions.h changes
|
||||
namespace meta {
|
||||
${meta_declaration}
|
||||
}
|
||||
|
||||
namespace native {
|
||||
${native_declaration}
|
||||
${native_definitions}
|
||||
}} // namespace at::native
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
#include <ATen/UnboxingFunctions.h>
|
||||
#include <ATen/Functions.h>
|
||||
|
||||
#include <ATen/Tensor.h>
|
||||
#include <ATen/core/functional.h>
|
||||
#include <ATen/core/interned_strings.h>
|
||||
#include <ATen/core/ivalue.h>
|
||||
#include <ATen/core/stack.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <cstddef>
|
||||
#include <cstring>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <tuple>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
namespace at {
|
||||
namespace unboxing {
|
||||
|
||||
using ::c10::fmap;
|
||||
using ::c10::filter;
|
||||
using torch::jit::peek;
|
||||
using torch::jit::drop;
|
||||
using torch::jit::pack;
|
||||
using torch::jit::pop;
|
||||
|
||||
// Generated function declaration
|
||||
${definitions}
|
||||
|
||||
} // namespace unboxing
|
||||
} // namespace at
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
// ${generated_comment}
|
||||
|
||||
// Generated by tools/jit/gen_unboxing.py. This file declares code generated boxed C++ functions for operators,
|
||||
// base off of native_functions.yaml (or similar yaml file with the same syntax). The definition of such a boxed
|
||||
// function will pop out IValues from the stack then convert them into the correct C++ types based on given schema. This
|
||||
// unboxing logic is an alternative to template-based metaprogramming unboxing.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
namespace at {
|
||||
namespace unboxing {
|
||||
namespace {
|
||||
|
||||
template<typename T, size_t N>
|
||||
std::array<T, N> as_array(const c10::List<c10::IValue>& list) {
|
||||
std::array<T, N> res;
|
||||
AT_ASSERT(list.size() == N);
|
||||
std::vector<T> vec;
|
||||
for (c10::IValue elem : list) {
|
||||
vec.push_back(elem.to<T>());
|
||||
}
|
||||
std::copy(vec.begin(), vec.end(), res.begin());
|
||||
return res;
|
||||
}
|
||||
} // namespace <anonymous>
|
||||
using Stack = std::vector<c10::IValue>;
|
||||
// Generated function declaration
|
||||
${declarations}
|
||||
|
||||
} // namespace unboxing
|
||||
} // namespace at
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/FunctionalInverses.h>
|
||||
#include <ATen/ViewMetaClasses.h>
|
||||
|
||||
#ifndef AT_PER_OPERATOR_HEADERS
|
||||
#include <ATen/Operators.h>
|
||||
#include <ATen/NativeFunctions.h>
|
||||
#else
|
||||
${op_headers}
|
||||
#endif
|
||||
|
||||
namespace at {
|
||||
namespace functionalization {
|
||||
|
||||
${view_meta_implementations}
|
||||
|
||||
} // namespace functionalization
|
||||
} // namespace at
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
|
||||
// ${generated_comment}
|
||||
|
||||
#include <ATen/FunctionalStorageImpl.h>
|
||||
|
||||
namespace at {
|
||||
namespace functionalization {
|
||||
|
||||
${view_meta_declarations}
|
||||
|
||||
} // namespace functionalization
|
||||
} // namespace at
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
#include <ATen/ViewMetaClasses.h>
|
||||
#include <torch/csrc/functionalization/Module.h>
|
||||
|
||||
namespace torch::functionalization {
|
||||
|
||||
void initGenerated(PyObject* module) {
|
||||
auto functionalization = py::handle(module).cast<py::module>();
|
||||
$view_meta_bindings
|
||||
}
|
||||
|
||||
} // namespace torch::functionalization
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#if defined(TORCH_ASSERT_NO_OPERATORS) || defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS)
|
||||
#error This change adds a dependency on native_functions.yaml, \
|
||||
meaning the file will need to be re-compiled every time an operator \
|
||||
is changed or added. Consider if including <ATen/core/symbol.h> for \
|
||||
the c10::Symbol class would be sufficient, or if your change would be \
|
||||
better placed in another file.
|
||||
#endif
|
||||
|
||||
// ATen symbols correspond exactly to operators defined in ATen. Every
|
||||
// symbol here corresponds exactly to an ATen operation defined in
|
||||
// native_functions.yaml; attributes are in one-to-one correspondence
|
||||
// with their ATen name.
|
||||
|
||||
#define FORALL_ATEN_BASE_SYMBOLS(_) \
|
||||
${aten_symbols}
|
||||
|
||||
#define FORALL_ATTR_BASE_SYMBOLS(_) \
|
||||
${attr_symbols}
|
||||
@@ -0,0 +1,19 @@
|
||||
#pragma once
|
||||
|
||||
// ${generated_comment}
|
||||
|
||||
#include <torch/headeronly/macros/Macros.h>
|
||||
|
||||
HIDDEN_NAMESPACE_BEGIN(torch, headeronly)
|
||||
|
||||
// Enum of valid tags obtained from the entries in tags.yaml
|
||||
enum class Tag {
|
||||
${enum_of_valid_tags}
|
||||
};
|
||||
|
||||
HIDDEN_NAMESPACE_END(torch, headeronly)
|
||||
|
||||
// Re-expose in the at:: namespace for backward compatibility
|
||||
namespace at {
|
||||
using torch::headeronly::Tag;
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
load("//:tools/bazel.bzl", "rules")
|
||||
load(":build.bzl", "define_targets")
|
||||
|
||||
define_targets(rules = rules)
|
||||
@@ -0,0 +1,3 @@
|
||||
If you add a file to this directory, you **MUST** update
|
||||
`torch/CMakeLists.txt` and add the file as a dependency to
|
||||
the `add_custom_command` call.
|
||||
@@ -0,0 +1,20 @@
|
||||
def define_targets(rules):
|
||||
rules.py_library(
|
||||
name = "autograd",
|
||||
srcs = rules.glob(["*.py"]),
|
||||
data = rules.glob([
|
||||
"*.yaml",
|
||||
"templates/*",
|
||||
]),
|
||||
visibility = ["//:__subpackages__"],
|
||||
deps = [
|
||||
rules.requirement("PyYAML"),
|
||||
"//torchgen",
|
||||
],
|
||||
)
|
||||
|
||||
rules.filegroup(
|
||||
name = "deprecated_yaml",
|
||||
srcs = ["deprecated.yaml"],
|
||||
visibility = ["//:__subpackages__"],
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
|
||||
from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo as NFWDI
|
||||
from torchgen.context import native_function_manager
|
||||
from torchgen.utils import T
|
||||
|
||||
|
||||
# Like tools.api.context.with_native_function, but for
|
||||
# NativeFunctionWithDifferentiabilityInfo.
|
||||
def with_native_function_with_differentiability_info(
|
||||
func: Callable[[NFWDI], T],
|
||||
) -> Callable[[NFWDI], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(f: NFWDI) -> T:
|
||||
with native_function_manager(f.func):
|
||||
return func(f)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# Like the above but with an additional dispatch key string argument
|
||||
def with_native_function_with_differentiability_info_and_key(
|
||||
func: Callable[[NFWDI, str], T],
|
||||
) -> Callable[[NFWDI, str], T]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(f: NFWDI, key: str) -> T:
|
||||
with native_function_manager(f.func):
|
||||
return func(f, key)
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,134 @@
|
||||
# Deprecated function signatures. These are exposed in Python, but not included
|
||||
# in the error message suggestions.
|
||||
|
||||
- name: add(Tensor self, Scalar alpha, Tensor other) -> Tensor
|
||||
aten: add(self, other, alpha)
|
||||
|
||||
- name: add_(Tensor(a!) self, Scalar alpha, Tensor other) -> Tensor(a!)
|
||||
aten: add_(self, other, alpha)
|
||||
|
||||
- name: add(Tensor self, Scalar alpha, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: add_out(out, self, other, alpha)
|
||||
|
||||
- name: addbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor
|
||||
aten: addbmm(self, batch1, batch2, beta, alpha)
|
||||
|
||||
- name: addbmm_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor(a!)
|
||||
aten: addbmm_(self, batch1, batch2, beta, alpha)
|
||||
|
||||
- name: addbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addbmm_out(out, self, batch1, batch2, beta, alpha)
|
||||
|
||||
- name: addbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2) -> Tensor
|
||||
aten: addbmm(self, batch1, batch2, beta, 1)
|
||||
|
||||
- name: addbmm_(Scalar beta, Tensor(a!) self, Tensor batch1, Tensor batch2) -> Tensor(a!)
|
||||
aten: addbmm_(self, batch1, batch2, beta, 1)
|
||||
|
||||
- name: addbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addbmm_out(out, self, batch1, batch2, beta, 1)
|
||||
|
||||
- name: addcdiv(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor
|
||||
aten: addcdiv(self, tensor1, tensor2, value)
|
||||
|
||||
- name: addcdiv_(Tensor(a!) self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor(a!)
|
||||
aten: addcdiv_(self, tensor1, tensor2, value)
|
||||
|
||||
- name: addcdiv(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addcdiv_out(out, self, tensor1, tensor2, value)
|
||||
|
||||
- name: addcmul(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor
|
||||
aten: addcmul(self, tensor1, tensor2, value)
|
||||
|
||||
- name: addcmul_(Tensor(a!) self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor(a!)
|
||||
aten: addcmul_(self, tensor1, tensor2, value)
|
||||
|
||||
- name: addcmul(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addcmul_out(out, self, tensor1, tensor2, value)
|
||||
|
||||
- name: addmm(Scalar beta, Tensor self, Scalar alpha, Tensor mat1, Tensor mat2) -> Tensor
|
||||
aten: addmm(self, mat1, mat2, beta, alpha)
|
||||
|
||||
- name: addmm_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor mat1, Tensor mat2) -> Tensor(a!)
|
||||
aten: addmm_(self, mat1, mat2, beta, alpha)
|
||||
|
||||
- name: addmm(Scalar beta, Tensor self, Scalar alpha, Tensor mat1, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addmm_out(out, self, mat1, mat2, beta, alpha)
|
||||
|
||||
- name: addmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2) -> Tensor
|
||||
aten: addmm(self, mat1, mat2, beta, 1)
|
||||
|
||||
- name: addmm_(Scalar beta, Tensor(a!) self, Tensor mat1, Tensor mat2) -> Tensor(a!)
|
||||
aten: addmm_(self, mat1, mat2, beta, 1)
|
||||
|
||||
- name: addmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addmm_out(out, self, mat1, mat2, beta, 1)
|
||||
|
||||
- name: sspaddmm(Scalar beta, Tensor self, Scalar alpha, Tensor mat1, Tensor mat2) -> Tensor
|
||||
aten: sspaddmm(self, mat1, mat2, beta, alpha)
|
||||
|
||||
- name: sspaddmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2) -> Tensor
|
||||
aten: sspaddmm(self, mat1, mat2, beta, 1)
|
||||
|
||||
- name: addmv(Scalar beta, Tensor self, Scalar alpha, Tensor mat, Tensor vec) -> Tensor
|
||||
aten: addmv(self, mat, vec, beta, alpha)
|
||||
|
||||
- name: addmv_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor mat, Tensor vec) -> Tensor(a!)
|
||||
aten: addmv_(self, mat, vec, beta, alpha)
|
||||
|
||||
- name: addmv(Scalar beta, Tensor self, Scalar alpha, Tensor mat, Tensor vec, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addmv_out(out, self, mat, vec, beta, alpha)
|
||||
|
||||
- name: addmv(Scalar beta, Tensor self, Tensor mat, Tensor vec) -> Tensor
|
||||
aten: addmv(self, mat, vec, beta, 1)
|
||||
|
||||
- name: addmv_(Scalar beta, Tensor(a!) self, Tensor mat, Tensor vec) -> Tensor(a!)
|
||||
aten: addmv_(self, mat, vec, beta, 1)
|
||||
|
||||
- name: addmv(Scalar beta, Tensor self, Tensor mat, Tensor vec, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addmv_out(out, self, mat, vec, beta, 1)
|
||||
|
||||
- name: addr(Scalar beta, Tensor self, Scalar alpha, Tensor vec1, Tensor vec2) -> Tensor
|
||||
aten: addr(self, vec1, vec2, beta, alpha)
|
||||
|
||||
- name: addr_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor vec1, Tensor vec2) -> Tensor(a!)
|
||||
aten: addr_(self, vec1, vec2, beta, alpha)
|
||||
|
||||
- name: addr(Scalar beta, Tensor self, Scalar alpha, Tensor vec1, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addr_out(out, self, vec1, vec2, beta, alpha)
|
||||
|
||||
- name: addr(Scalar beta, Tensor self, Tensor vec1, Tensor vec2) -> Tensor
|
||||
aten: addr(self, vec1, vec2, beta, 1)
|
||||
|
||||
- name: addr_(Scalar beta, Tensor(a!) self, Tensor vec1, Tensor vec2) -> Tensor(a!)
|
||||
aten: addr_(self, vec1, vec2, beta, 1)
|
||||
|
||||
- name: addr(Scalar beta, Tensor self, Tensor vec1, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: addr_out(out, self, vec1, vec2, beta, 1)
|
||||
|
||||
- name: baddbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor
|
||||
aten: baddbmm(self, batch1, batch2, beta, alpha)
|
||||
|
||||
- name: baddbmm_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor(a!)
|
||||
aten: baddbmm_(self, batch1, batch2, beta, alpha)
|
||||
|
||||
- name: baddbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: baddbmm_out(out, self, batch1, batch2, beta, alpha)
|
||||
|
||||
- name: baddbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2) -> Tensor
|
||||
aten: baddbmm(self, batch1, batch2, beta, 1)
|
||||
|
||||
- name: baddbmm_(Scalar beta, Tensor(a!) self, Tensor batch1, Tensor batch2) -> Tensor(a!)
|
||||
aten: baddbmm_(self, batch1, batch2, beta, 1)
|
||||
|
||||
- name: baddbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: baddbmm_out(out, self, batch1, batch2, beta, 1)
|
||||
|
||||
- name: sub(Tensor self, Scalar alpha, Tensor other) -> Tensor
|
||||
aten: sub(self, other, alpha)
|
||||
|
||||
- name: sub_(Tensor(a!) self, Scalar alpha, Tensor other) -> Tensor(a!)
|
||||
aten: sub_(self, other, alpha)
|
||||
|
||||
- name: sub(Tensor self, Scalar alpha, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
|
||||
aten: sub_out(out, self, other, alpha)
|
||||
+3282
File diff suppressed because it is too large
Load Diff
+134
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
For procedural tests needed for __torch_function__, we use this function
|
||||
to export method names and signatures as needed by the tests in
|
||||
test/test_overrides.py.
|
||||
|
||||
python -m tools.autograd.gen_annotated_fn_args \
|
||||
aten/src/ATen/native/native_functions.yaml \
|
||||
aten/src/ATen/native/tags.yaml \
|
||||
$OUTPUT_DIR \
|
||||
tools/autograd
|
||||
|
||||
Where $OUTPUT_DIR is where you would like the files to be
|
||||
generated. In the full build system, OUTPUT_DIR is
|
||||
torch/testing/_internal/generated
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import textwrap
|
||||
from collections import defaultdict
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torchgen.api.python as python
|
||||
from torchgen.context import with_native_function
|
||||
from torchgen.gen import parse_native_yaml
|
||||
from torchgen.utils import FileManager
|
||||
|
||||
from .gen_python_functions import (
|
||||
is_py_fft_function,
|
||||
is_py_linalg_function,
|
||||
is_py_nn_function,
|
||||
is_py_special_function,
|
||||
is_py_torch_function,
|
||||
is_py_variable_method,
|
||||
should_generate_py_binding,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from torchgen.model import Argument, BaseOperatorName, NativeFunction
|
||||
|
||||
|
||||
def gen_annotated(
|
||||
native_yaml_path: str, tags_yaml_path: str, out: str, autograd_dir: str
|
||||
) -> None:
|
||||
native_functions = parse_native_yaml(
|
||||
native_yaml_path, tags_yaml_path
|
||||
).native_functions
|
||||
mappings = (
|
||||
(is_py_torch_function, "torch._C._VariableFunctions"),
|
||||
(is_py_nn_function, "torch._C._nn"),
|
||||
(is_py_linalg_function, "torch._C._linalg"),
|
||||
(is_py_special_function, "torch._C._special"),
|
||||
(is_py_fft_function, "torch._C._fft"),
|
||||
(is_py_variable_method, "torch.Tensor"),
|
||||
)
|
||||
annotated_args: list[str] = []
|
||||
for pred, namespace in mappings:
|
||||
groups: dict[BaseOperatorName, list[NativeFunction]] = defaultdict(list)
|
||||
for f in native_functions:
|
||||
if not should_generate_py_binding(f) or not pred(f):
|
||||
continue
|
||||
groups[f.func.name.name].append(f)
|
||||
for group in groups.values():
|
||||
for f in group:
|
||||
annotated_args.append(f"{namespace}.{gen_annotated_args(f)}")
|
||||
|
||||
template_path = os.path.join(autograd_dir, "templates")
|
||||
fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
|
||||
fm.write_with_template(
|
||||
"annotated_fn_args.py",
|
||||
"annotated_fn_args.py.in",
|
||||
lambda: {
|
||||
"annotated_args": textwrap.indent("\n".join(annotated_args), " "),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@with_native_function
|
||||
def gen_annotated_args(f: NativeFunction) -> str:
|
||||
def _get_kwargs_func_exclusion_list() -> list[str]:
|
||||
# functions that currently don't work with kwargs in test_overrides.py
|
||||
return [
|
||||
"diagonal",
|
||||
"round_",
|
||||
"round",
|
||||
"scatter_",
|
||||
]
|
||||
|
||||
def _add_out_arg(
|
||||
out_args: list[dict[str, Any]], args: Sequence[Argument], *, is_kwarg_only: bool
|
||||
) -> None:
|
||||
for arg in args:
|
||||
if arg.default is not None:
|
||||
continue
|
||||
out_arg: dict[str, Any] = {}
|
||||
out_arg["is_kwarg_only"] = str(is_kwarg_only)
|
||||
out_arg["name"] = arg.name
|
||||
out_arg["simple_type"] = python.argument_type_str(
|
||||
arg.type, simple_type=True
|
||||
)
|
||||
size_t = python.argument_type_size(arg.type)
|
||||
if size_t:
|
||||
out_arg["size"] = size_t
|
||||
out_args.append(out_arg)
|
||||
|
||||
out_args: list[dict[str, Any]] = []
|
||||
_add_out_arg(out_args, f.func.arguments.flat_positional, is_kwarg_only=False)
|
||||
if f"{f.func.name.name}" not in _get_kwargs_func_exclusion_list():
|
||||
_add_out_arg(out_args, f.func.arguments.flat_kwarg_only, is_kwarg_only=True)
|
||||
|
||||
return f"{f.func.name.name}: {repr(out_args)},"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Generate annotated_fn_args script")
|
||||
parser.add_argument(
|
||||
"native_functions", metavar="NATIVE", help="path to native_functions.yaml"
|
||||
)
|
||||
parser.add_argument("tags", metavar="TAGS", help="path to tags.yaml")
|
||||
parser.add_argument("out", metavar="OUT", help="path to output directory")
|
||||
parser.add_argument(
|
||||
"autograd", metavar="AUTOGRAD", help="path to template directory"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
gen_annotated(args.native_functions, args.tags, args.out, args.autograd)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
To run this file by hand from the root of the PyTorch
|
||||
repository, run:
|
||||
|
||||
python -m tools.autograd.gen_autograd \
|
||||
aten/src/ATen/native/native_functions.yaml \
|
||||
aten/src/ATen/native/tags.yaml \
|
||||
$OUTPUT_DIR \
|
||||
tools/autograd
|
||||
|
||||
Where $OUTPUT_DIR is where you would like the files to be
|
||||
generated. In the full build system, OUTPUT_DIR is
|
||||
torch/csrc/autograd/generated/
|
||||
"""
|
||||
|
||||
# gen_autograd.py generates C++ autograd functions and Python bindings.
|
||||
#
|
||||
# It delegates to the following scripts:
|
||||
#
|
||||
# gen_autograd_functions.py: generates subclasses of torch::autograd::Node
|
||||
# gen_variable_type.py: generates VariableType.h which contains all tensor methods
|
||||
# gen_python_functions.py: generates Python bindings to THPVariable
|
||||
#
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.autograd import (
|
||||
match_differentiability_info,
|
||||
NativeFunctionWithDifferentiabilityInfo,
|
||||
)
|
||||
from torchgen.gen import parse_native_yaml
|
||||
from torchgen.selective_build.selector import SelectiveBuilder
|
||||
|
||||
from . import gen_python_functions
|
||||
from .gen_autograd_functions import (
|
||||
gen_autograd_functions_lib,
|
||||
gen_autograd_functions_python,
|
||||
)
|
||||
from .gen_inplace_or_view_type import gen_inplace_or_view_type
|
||||
from .gen_trace_type import gen_trace_type
|
||||
from .gen_variable_factories import gen_variable_factories
|
||||
from .gen_variable_type import gen_variable_type
|
||||
from .gen_view_funcs import gen_view_funcs
|
||||
from .load_derivatives import load_derivatives
|
||||
|
||||
|
||||
def gen_autograd(
|
||||
native_functions_path: str,
|
||||
tags_path: str,
|
||||
out: str,
|
||||
autograd_dir: str,
|
||||
operator_selector: SelectiveBuilder,
|
||||
disable_autograd: bool = False,
|
||||
) -> None:
|
||||
# Parse and load derivatives.yaml
|
||||
differentiability_infos, used_dispatch_keys = load_derivatives(
|
||||
os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path, tags_path
|
||||
)
|
||||
|
||||
template_path = os.path.join(autograd_dir, "templates")
|
||||
|
||||
native_funcs = parse_native_yaml(native_functions_path, tags_path).native_functions
|
||||
fns = sorted(
|
||||
filter(
|
||||
operator_selector.is_native_function_selected_for_training, native_funcs
|
||||
),
|
||||
key=lambda f: cpp.name(f.func),
|
||||
)
|
||||
fns_with_diff_infos: list[NativeFunctionWithDifferentiabilityInfo] = (
|
||||
match_differentiability_info(fns, differentiability_infos)
|
||||
)
|
||||
|
||||
# Generate VariableType.h/cpp
|
||||
if not disable_autograd:
|
||||
gen_variable_type(
|
||||
out,
|
||||
native_functions_path,
|
||||
tags_path,
|
||||
fns_with_diff_infos,
|
||||
template_path,
|
||||
used_dispatch_keys,
|
||||
)
|
||||
|
||||
gen_inplace_or_view_type(
|
||||
out, native_functions_path, tags_path, fns_with_diff_infos, template_path
|
||||
)
|
||||
|
||||
# operator filter not applied as tracing sources are excluded in selective build
|
||||
gen_trace_type(out, native_funcs, template_path)
|
||||
# Generate Functions.h/cpp
|
||||
gen_autograd_functions_lib(out, differentiability_infos, template_path)
|
||||
|
||||
# Generate variable_factories.h
|
||||
gen_variable_factories(out, native_functions_path, tags_path, template_path)
|
||||
|
||||
# Generate ViewFuncs.h/cpp
|
||||
gen_view_funcs(out, fns_with_diff_infos, template_path)
|
||||
|
||||
|
||||
def gen_autograd_python(
|
||||
native_functions_path: str,
|
||||
tags_path: str,
|
||||
out: str,
|
||||
autograd_dir: str,
|
||||
) -> None:
|
||||
differentiability_infos, _ = load_derivatives(
|
||||
os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path, tags_path
|
||||
)
|
||||
|
||||
template_path = os.path.join(autograd_dir, "templates")
|
||||
|
||||
# Generate Functions.h/cpp
|
||||
gen_autograd_functions_python(out, differentiability_infos, template_path)
|
||||
|
||||
# Generate Python bindings
|
||||
deprecated_path = os.path.join(autograd_dir, "deprecated.yaml")
|
||||
gen_python_functions.gen(
|
||||
out, native_functions_path, tags_path, deprecated_path, template_path
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Generate autograd C++ files script")
|
||||
parser.add_argument(
|
||||
"native_functions", metavar="NATIVE", help="path to native_functions.yaml"
|
||||
)
|
||||
parser.add_argument("tags", metavar="NATIVE", help="path to tags.yaml")
|
||||
parser.add_argument("out", metavar="OUT", help="path to output directory")
|
||||
parser.add_argument(
|
||||
"autograd", metavar="AUTOGRAD", help="path to autograd directory"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
gen_autograd(
|
||||
args.native_functions,
|
||||
args.tags,
|
||||
args.out,
|
||||
args.autograd,
|
||||
SelectiveBuilder.get_nop_selector(),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+1082
File diff suppressed because it is too large
Load Diff
+676
@@ -0,0 +1,676 @@
|
||||
# Generates ADInplaceOrViewType.h/cpp
|
||||
#
|
||||
# NOTE: If any changes are being made to the ADInplaceOrView codegen please also check
|
||||
# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp
|
||||
# The fallback is expected to mimic this codegen, so we should keep the two in sync.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.autograd import (
|
||||
dispatch_strategy,
|
||||
gen_differentiable_outputs,
|
||||
NativeFunctionWithDifferentiabilityInfo,
|
||||
)
|
||||
from torchgen.api.types import (
|
||||
BaseCType,
|
||||
Binding,
|
||||
boolT,
|
||||
ConstRefCType,
|
||||
CType,
|
||||
DispatcherSignature,
|
||||
intArrayRefT,
|
||||
longT,
|
||||
OptionalCType,
|
||||
symIntArrayRefT,
|
||||
SymIntT,
|
||||
tensorT,
|
||||
)
|
||||
from torchgen.code_template import CodeTemplate
|
||||
from torchgen.context import with_native_function
|
||||
from torchgen.model import (
|
||||
NativeFunction,
|
||||
SchemaKind,
|
||||
SelfArgument,
|
||||
TensorOptionsArguments,
|
||||
Type,
|
||||
)
|
||||
from torchgen.utils import FileManager
|
||||
|
||||
from .context import with_native_function_with_differentiability_info
|
||||
from .gen_trace_type import (
|
||||
get_return_value,
|
||||
MANUAL_AUTOGRAD,
|
||||
tie_return_values,
|
||||
type_wrapper_name,
|
||||
)
|
||||
|
||||
|
||||
# See NOTE [ Autograd View Variables ] in variable.h for details.
|
||||
# If you update list VIEW_FUNCTIONS or RETURNS_VIEWS_OF_INPUT,
|
||||
# you **MUST** also update the public list of view ops accordingly in
|
||||
# docs/source/tensor_view.rst. Note not all ATen functions are exposed to public,
|
||||
# e.g alias & sparse_coo_tensor_with_dims_and_tensors.
|
||||
#
|
||||
# A map: function name => name of the argument that all outputs are view of
|
||||
|
||||
VIEW_FUNCTIONS_WITH_METADATA_CHANGE = [
|
||||
"view_as_complex",
|
||||
"view_as_real",
|
||||
"_conj",
|
||||
"_neg_view",
|
||||
"_nested_get_values",
|
||||
"_nested_view_from_buffer",
|
||||
"_nested_view_from_jagged",
|
||||
]
|
||||
|
||||
VIEW_FUNCTIONS = {
|
||||
"numpy_T": "self",
|
||||
"alias": "self",
|
||||
"as_strided": "self",
|
||||
"diagonal": "self",
|
||||
"expand": "self",
|
||||
"permute": "self",
|
||||
"select": "self",
|
||||
"slice": "self",
|
||||
"slice_inverse": "self",
|
||||
"split": "self",
|
||||
"split_with_sizes": "self",
|
||||
"squeeze": "self",
|
||||
"t": "self",
|
||||
"transpose": "self",
|
||||
"unfold": "self",
|
||||
"unsqueeze": "self",
|
||||
"flatten": "self",
|
||||
"view": "self",
|
||||
"unbind": "self",
|
||||
"_indices": "self",
|
||||
"_values": "self",
|
||||
"indices": "self",
|
||||
"values": "self",
|
||||
"crow_indices": "self",
|
||||
"col_indices": "self",
|
||||
"ccol_indices": "self",
|
||||
"row_indices": "self",
|
||||
# sparse_coo ctor output should really be views of both indices and values,
|
||||
# but we only supports making as view of a single variable, and indices is
|
||||
# discrete anyways.
|
||||
# FIXME: clone indices on construction.
|
||||
"sparse_coo_tensor_with_dims_and_tensors": "values",
|
||||
"_reshape_alias": "self",
|
||||
"_test_autograd_multiple_dispatch_view": "self",
|
||||
}
|
||||
|
||||
for key in VIEW_FUNCTIONS_WITH_METADATA_CHANGE:
|
||||
VIEW_FUNCTIONS[key] = "self"
|
||||
|
||||
# note: some VIEW_FUNCTIONS are just compositions of the view functions above
|
||||
# this list contains both the root view functions and any that are purely composed
|
||||
# of viewing functions, and is used by the JIT to determine when an operator
|
||||
# may return a view of its inputs; however they may sometimes return a copy.
|
||||
# (e.g. `contiguous`)
|
||||
RETURNS_VIEWS_OF_INPUT = set(VIEW_FUNCTIONS.keys()).union(
|
||||
{
|
||||
"chunk",
|
||||
"detach",
|
||||
"contiguous",
|
||||
"reshape",
|
||||
"reshape_as",
|
||||
"expand_as",
|
||||
"view_as",
|
||||
"real",
|
||||
"imag",
|
||||
"narrow",
|
||||
"movedim",
|
||||
"tensor_split",
|
||||
"swapdims",
|
||||
"swapaxes",
|
||||
"mT",
|
||||
"mH",
|
||||
"adjoint",
|
||||
"matrix_H",
|
||||
}
|
||||
)
|
||||
|
||||
# These are the functions we consider views for the purposes of validating
|
||||
# StorageImpl and TensorImpl in gen_variable_type.
|
||||
# `_unsafe_view` is not included in VIEW_FUNCTIONS above because it is not a
|
||||
# view for the purposes of ADInplaceOrView kernel, we do not want to call as_view
|
||||
# See NOTE [Unsafe View] for more info.
|
||||
ALL_VIEW_FUNCTIONS = {
|
||||
**VIEW_FUNCTIONS,
|
||||
"_unsafe_view": "self",
|
||||
}
|
||||
|
||||
ARRAYREF_TO_VEC = CodeTemplate(
|
||||
"""\
|
||||
auto ${vec} = ${arg}.vec();
|
||||
"""
|
||||
)
|
||||
|
||||
OPTIONAL_TO_VAL = CodeTemplate(
|
||||
"""\
|
||||
auto ${val} = ${arg}.value_or(${default});
|
||||
"""
|
||||
)
|
||||
|
||||
CALL_DISPATCH = CodeTemplate(
|
||||
"""\
|
||||
at::_ops::${unambiguous_name}::call(${unpacked_args})"""
|
||||
)
|
||||
|
||||
REVERSE_VIEW_DISPATCH = CodeTemplate(
|
||||
"""\
|
||||
${reverse_name}(${unpacked_args})"""
|
||||
)
|
||||
|
||||
MULTI_OUTPUT_VIEW_ITERATION = CodeTemplate(
|
||||
"""\
|
||||
for (auto ${view_idx} : c10::irange(${var}.size())) {
|
||||
${body}
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
SETUP_REPLAY_VIEW_IF_NOT_SUPPORT_AS_STRIDED_OR_VIEW_WITH_METADATA_CHANGE = CodeTemplate(
|
||||
"""\
|
||||
std::unique_ptr<torch::autograd::ViewFunc> func(nullptr);
|
||||
std::function<at::Tensor(const at::Tensor&)> rev_func=nullptr;
|
||||
if (${is_view_with_metadata_change} ||
|
||||
!self.unsafeGetTensorImpl()->support_as_strided() ||
|
||||
self.unsafeGetTensorImpl()->is_python_dispatch() ||
|
||||
c10::AutogradState::get_tls_state().get_view_replay_enabled()) {
|
||||
${replay_view_func}
|
||||
${reverse_replay_view_func}
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
REPLAY_VIEW_FUNC = CodeTemplate(
|
||||
"""\
|
||||
func = std::make_unique<${view_func_name}>(${view_func_args});
|
||||
"""
|
||||
)
|
||||
|
||||
REVERSE_REPLAY_VIEW_LAMBDA_FUNC = CodeTemplate(
|
||||
"""\
|
||||
rev_func = [=](const at::Tensor& ${input_view}) {
|
||||
return ${reverse_replay_view_call};
|
||||
};
|
||||
"""
|
||||
)
|
||||
|
||||
METHOD_DEFINITION = CodeTemplate(
|
||||
"""\
|
||||
${return_type} ${type_wrapper_name}(${formals}) {
|
||||
${type_definition_body}
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
WRAPPER_REGISTRATION = CodeTemplate(
|
||||
"""\
|
||||
m.impl("${unqual_operator_name_with_overload}",
|
||||
TORCH_FN(${class_type}::${type_wrapper_name})
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION = CodeTemplate(
|
||||
"""\
|
||||
m.impl("${unqual_operator_name_with_overload}", torch::autograd::autogradNotImplementedFallback());
|
||||
"""
|
||||
)
|
||||
|
||||
INPLACE_REDISPATCH = CodeTemplate(
|
||||
"""\
|
||||
{
|
||||
at::AutoDispatchBelowADInplaceOrView guard;
|
||||
at::_ops::${unambiguous_name}::redispatch(${unpacked_args});
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
ASSIGN_RETURN_VALUE = CodeTemplate(
|
||||
"""\
|
||||
${return_values} = ${rhs_value};
|
||||
"""
|
||||
)
|
||||
|
||||
VIEW_REDISPATCH = CodeTemplate(
|
||||
"""\
|
||||
${assign_return_values} ([&]() {
|
||||
at::AutoDispatchBelowADInplaceOrView guard;
|
||||
return at::_ops::${unambiguous_name}::redispatch(${unpacked_args});
|
||||
})();
|
||||
"""
|
||||
)
|
||||
|
||||
TMP_VAR = "_tmp"
|
||||
|
||||
|
||||
# FIXME: Ideally these functions should be methods on Type class, but we have a
|
||||
# comment in codegen/model.py there saying these concepts are not well defined.
|
||||
# Thus we put a version that commonly used by autograd codegen here.
|
||||
def is_tensor_type(t: Type) -> bool:
|
||||
# TODO: Should handle optional here?
|
||||
return t.is_tensor_like() and t.is_list_like() is None
|
||||
|
||||
|
||||
def is_tensor_list_type(t: Type) -> bool:
|
||||
# TODO: Should handle optional here?
|
||||
return t.is_tensor_like() and t.is_list_like() is not None
|
||||
|
||||
|
||||
UNPACK_TENSOR = CodeTemplate(
|
||||
"""\
|
||||
auto${ref} ${arg_name}_ = unpack${suffix}(${arg_name}, "${arg_name}", ${arg_pos});"""
|
||||
)
|
||||
|
||||
|
||||
def unpacked_name(arg_name: str) -> str:
|
||||
return arg_name + "_"
|
||||
|
||||
|
||||
# e.g. select.int -> select_copy_int_inverse()
|
||||
def inverse_view_name(f: NativeFunction) -> str:
|
||||
copy_variant = f"{f.root_name}_copy"
|
||||
overload = f"{f.func.name.overload_name}"
|
||||
if overload != "":
|
||||
overload = "_" + overload
|
||||
return f"{copy_variant}{overload}_inverse"
|
||||
|
||||
|
||||
def extract_bindings(f: NativeFunction) -> list[Binding]:
|
||||
return [
|
||||
r
|
||||
for a in f.func.schema_order_arguments()
|
||||
for r in cpp.argument(
|
||||
a,
|
||||
method=False,
|
||||
symint=True,
|
||||
cpp_no_default_args=set(),
|
||||
faithful=False,
|
||||
has_tensor_options=False,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@with_native_function
|
||||
def unpack_args(f: NativeFunction) -> tuple[list[str], list[Binding]]:
|
||||
body: list[str] = []
|
||||
unpacked_bindings: list[Binding] = []
|
||||
|
||||
for i, binding in enumerate(extract_bindings(f)):
|
||||
if isinstance(binding.argument, SelfArgument):
|
||||
raise AssertionError("Binding argument should not be SelfArgument")
|
||||
if isinstance(binding.argument, TensorOptionsArguments):
|
||||
raise RuntimeError("VariableKernel shouldn't take TensorOptions")
|
||||
|
||||
is_nullable = binding.argument.type.is_nullable()
|
||||
if not binding.argument.type.is_tensor_like() or is_nullable:
|
||||
unpacked_bindings.append(binding)
|
||||
continue
|
||||
|
||||
is_tensor_list = is_tensor_list_type(binding.argument.type)
|
||||
ref = (not is_nullable) and not is_tensor_list
|
||||
suffix = "_opt" if is_nullable and not is_tensor_list else ""
|
||||
body.append(
|
||||
UNPACK_TENSOR.substitute(
|
||||
arg_name=binding.name,
|
||||
arg_pos=i,
|
||||
suffix=suffix,
|
||||
ref="&" if ref else "",
|
||||
)
|
||||
)
|
||||
unpacked_bindings.append(
|
||||
Binding(
|
||||
name=unpacked_name(binding.name),
|
||||
nctype=binding.nctype,
|
||||
argument=binding.argument,
|
||||
default=binding.default,
|
||||
)
|
||||
)
|
||||
|
||||
return body, unpacked_bindings
|
||||
|
||||
|
||||
def get_base_name(f: NativeFunction) -> str:
|
||||
return f.func.name.name.base # TODO: should be str(f.func.name.name)?
|
||||
|
||||
|
||||
def get_view_info(f: NativeFunction) -> str | None:
|
||||
base_name = get_base_name(f)
|
||||
view_info = VIEW_FUNCTIONS.get(base_name)
|
||||
if view_info is None and base_name in RETURNS_VIEWS_OF_INPUT:
|
||||
view_info = "self"
|
||||
return view_info
|
||||
|
||||
|
||||
def emit_view_func(
|
||||
f: NativeFunction, bindings: list[Binding], view_idx: str | None = None
|
||||
) -> str:
|
||||
"""Generate an additional lambda function to recover views in backward when as_strided is not supported.
|
||||
See Note [View + Inplace update for base tensor] and [View + Inplace update for view tensor] for more details.
|
||||
"""
|
||||
# TODO: Clean this logic up if we get rid of reverse view funcs or reify them.
|
||||
input_base = "input_base"
|
||||
replay_view_func = ""
|
||||
updated_args: list[str] = []
|
||||
known_view_arg_simple_types: list[CType] = [
|
||||
BaseCType(longT),
|
||||
OptionalCType(BaseCType(longT)),
|
||||
BaseCType(SymIntT),
|
||||
OptionalCType(BaseCType(SymIntT)),
|
||||
BaseCType(boolT),
|
||||
BaseCType(intArrayRefT),
|
||||
BaseCType(symIntArrayRefT),
|
||||
ConstRefCType(BaseCType(tensorT)),
|
||||
ConstRefCType(OptionalCType(BaseCType(tensorT))),
|
||||
]
|
||||
for binding in bindings:
|
||||
arg, arg_type = binding.name, binding.nctype.type
|
||||
if arg == "self":
|
||||
updated_args.append(input_base)
|
||||
continue
|
||||
if arg_type not in known_view_arg_simple_types:
|
||||
known_types_str = ", ".join([str(t) for t in known_view_arg_simple_types])
|
||||
raise TypeError(
|
||||
f"You are adding an {arg_type} {arg} argument to op {cpp.name(f.func)} in addition to known types: "
|
||||
f"{known_types_str}. Please update the list or materialize it so that it can be closed "
|
||||
"over by value, also add a test in pytorch/xla/test/test_operations.py where this code "
|
||||
"is exercised."
|
||||
)
|
||||
if arg_type == BaseCType(intArrayRefT) or arg_type == BaseCType(
|
||||
symIntArrayRefT
|
||||
):
|
||||
# It's not safe to close over IntArrayRef by value, since this is a
|
||||
# reference type, so materialize a vector to close over by value
|
||||
arg_vec = arg + "_vec"
|
||||
replay_view_func += ARRAYREF_TO_VEC.substitute(arg=arg, vec=arg_vec)
|
||||
updated_args.append(arg_vec)
|
||||
elif arg_type == OptionalCType(BaseCType(longT)):
|
||||
# Materialize int64_t? to int64_t
|
||||
arg_value = arg + "_val"
|
||||
replay_view_func += OPTIONAL_TO_VAL.substitute(
|
||||
arg=arg, val=arg_value, default="0"
|
||||
)
|
||||
updated_args.append(arg_value)
|
||||
elif arg_type == ConstRefCType(BaseCType(tensorT)) or arg_type == ConstRefCType(
|
||||
OptionalCType(BaseCType(tensorT))
|
||||
):
|
||||
# NB: Closing over a tensor. If a user modifies this tensor, this will be silently
|
||||
# incorrect. The proper thing to do is to store the version counter and copy on write.
|
||||
updated_args.append(arg)
|
||||
else:
|
||||
updated_args.append(arg)
|
||||
|
||||
from .gen_view_funcs import view_func_name
|
||||
|
||||
view_func_args = [b.name for b in bindings if b.name != "self"]
|
||||
if view_idx is not None:
|
||||
view_func_args.append(f"{view_idx}")
|
||||
replay_view_func += REPLAY_VIEW_FUNC.substitute(
|
||||
view_func_name=view_func_name(f, include_namespace=True),
|
||||
view_func_args=view_func_args,
|
||||
)
|
||||
|
||||
input_view = "input_view"
|
||||
reverse_unpacked_args = [
|
||||
"self",
|
||||
f"{input_view}",
|
||||
# inverse_return_mode=
|
||||
"at::functionalization::InverseReturnMode::AlwaysView",
|
||||
*(() if view_idx is None else (f"{view_idx}",)),
|
||||
# skip input_base arg
|
||||
*updated_args[1:],
|
||||
]
|
||||
|
||||
from torchgen.api.functionalization import reverse_name
|
||||
|
||||
reverse_replay_view_call = REVERSE_VIEW_DISPATCH.substitute(
|
||||
reverse_name=reverse_name(f, include_namespace=True),
|
||||
unpacked_args=reverse_unpacked_args,
|
||||
)
|
||||
reverse_replay_view_func = REVERSE_REPLAY_VIEW_LAMBDA_FUNC.substitute(
|
||||
input_view=input_view, reverse_replay_view_call=reverse_replay_view_call
|
||||
)
|
||||
|
||||
is_view_with_metadata_change = (
|
||||
"true" if cpp.name(f.func) in VIEW_FUNCTIONS_WITH_METADATA_CHANGE else "false"
|
||||
)
|
||||
|
||||
return SETUP_REPLAY_VIEW_IF_NOT_SUPPORT_AS_STRIDED_OR_VIEW_WITH_METADATA_CHANGE.substitute(
|
||||
is_view_with_metadata_change=is_view_with_metadata_change,
|
||||
replay_view_func=replay_view_func,
|
||||
reverse_replay_view_func=reverse_replay_view_func,
|
||||
)
|
||||
|
||||
|
||||
def emit_view_body(
|
||||
fn: NativeFunctionWithDifferentiabilityInfo, var: str
|
||||
) -> tuple[str, str]:
|
||||
# See NOTE [ Autograd View Variables ] in variable.h for details.
|
||||
f = fn.func
|
||||
base_name = get_base_name(f)
|
||||
view_info = get_view_info(f)
|
||||
call = ""
|
||||
differentiable_outputs = gen_differentiable_outputs(fn)
|
||||
differentiable_output_vars = {r.name for r in differentiable_outputs}
|
||||
if not isinstance(view_info, str):
|
||||
raise TypeError(
|
||||
f"The view info should be a string for {base_name}, but it is: {view_info}"
|
||||
)
|
||||
if len(differentiable_output_vars) == 0:
|
||||
# no output is differentiable (.indices() for SparseTensors for example)
|
||||
rhs_value = (
|
||||
f"as_view({view_info}, {var}, "
|
||||
f"/* is_bw_differentiable */ false, /* is_fw_differentiable */ false)"
|
||||
)
|
||||
elif len(differentiable_output_vars) == 1:
|
||||
# Single differentiable output (Tensor or Tensor[])
|
||||
return_info = differentiable_outputs[0]
|
||||
# We only support simple Tensor or a TensorList for functions that return views
|
||||
if not is_tensor_type(return_info.type) and not is_tensor_list_type(
|
||||
return_info.type
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"{base_name} that return differentiable views can only return Tensor or Tensor[]"
|
||||
)
|
||||
|
||||
# See Note [ View + Inplace detection]
|
||||
def get_creation_meta_in_mode(original: str) -> str:
|
||||
creation_meta_with_grad_mode = f"(at::GradMode::is_enabled() ? {original} : CreationMeta::NO_GRAD_MODE)"
|
||||
return f"InferenceMode::is_enabled() ? CreationMeta::INFERENCE_MODE : {creation_meta_with_grad_mode}"
|
||||
|
||||
# Only allow rebasing of the history if we return a single Tensor
|
||||
# If we are in a no grad block, raise a warning
|
||||
# See NOTE [ View + Inplace detection ] for more details about this logic
|
||||
if is_tensor_list_type(return_info.type):
|
||||
creation_meta = get_creation_meta_in_mode("CreationMeta::MULTI_OUTPUT_NODE")
|
||||
view_idx = "view_idx"
|
||||
view_func = emit_view_func(
|
||||
f, extract_bindings(f), view_idx=view_idx
|
||||
).strip()
|
||||
as_view_call = (
|
||||
f"as_view(/* base */ {view_info}, /* output */ {var}[{view_idx}], "
|
||||
"/* is_bw_differentiable */ true, /* is_fw_differentiable */ true, "
|
||||
"/* view_func */ std::move(func), /* rev_view_func */ rev_func, "
|
||||
f"/* creation_meta */ {creation_meta});"
|
||||
)
|
||||
call += MULTI_OUTPUT_VIEW_ITERATION.substitute(
|
||||
var=var, view_idx=view_idx, body=f"{view_func}\n{as_view_call}"
|
||||
)
|
||||
rhs_value = f"std::move({var})"
|
||||
else:
|
||||
call += emit_view_func(f, extract_bindings(f), view_idx=None)
|
||||
creation_meta = get_creation_meta_in_mode("CreationMeta::DEFAULT")
|
||||
rhs_value = (
|
||||
f"as_view(/* base */ {view_info}, /* output */ {var}, /* is_bw_differentiable */ true, "
|
||||
"/* is_fw_differentiable */ true, "
|
||||
f"/* view_func */ std::move(func), /* rev_view_func */ rev_func, /* creation_meta */ {creation_meta})"
|
||||
)
|
||||
else:
|
||||
# This could be supported but we don't need it at the moment, so keeping things simple.
|
||||
raise RuntimeError(
|
||||
"Function that return multiple differentiable output "
|
||||
"when at least one of them is view is not supported."
|
||||
)
|
||||
return call, rhs_value
|
||||
|
||||
|
||||
def modifies_arguments(f: NativeFunction) -> bool:
|
||||
return f.func.kind() in [SchemaKind.inplace, SchemaKind.out]
|
||||
|
||||
|
||||
@with_native_function_with_differentiability_info
|
||||
def emit_inplace_or_view_body(fn: NativeFunctionWithDifferentiabilityInfo) -> list[str]:
|
||||
f = fn.func
|
||||
inplace_view_body: list[str] = []
|
||||
|
||||
dispatcher_sig = DispatcherSignature.from_schema(f.func)
|
||||
dispatcher_exprs = dispatcher_sig.exprs()
|
||||
|
||||
# code-generated ADInplaceOrView kernels plumb and recompute dispatch keys directly through the kernel for performance.
|
||||
# See Note [Plumbing Keys Through The Dispatcher] for details.
|
||||
dispatch_key_set = "ks & c10::after_ADInplaceOrView_keyset"
|
||||
redispatch_args = ", ".join([dispatch_key_set] + [a.expr for a in dispatcher_exprs])
|
||||
|
||||
# Note that this calls the slow, dispatching variants of manual_cpp_binding ops.
|
||||
# We could probably work harder to ensure that the fast variants are called instead, but the perf benefit would be minimal.
|
||||
if modifies_arguments(f): # inplace op
|
||||
inplace_view_body.append(
|
||||
INPLACE_REDISPATCH.substitute(
|
||||
unambiguous_name=f.func.name.unambiguous_name(),
|
||||
unpacked_args=redispatch_args,
|
||||
)
|
||||
)
|
||||
for r in cpp.return_names(f):
|
||||
inplace_view_body.append(f"increment_version({r});")
|
||||
else:
|
||||
if get_view_info(f) is None:
|
||||
raise AssertionError("Expected view info to be non-None")
|
||||
inplace_view_body.append(
|
||||
VIEW_REDISPATCH.substitute(
|
||||
assign_return_values="auto " + TMP_VAR + " = ",
|
||||
unambiguous_name=f.func.name.unambiguous_name(),
|
||||
unpacked_args=redispatch_args,
|
||||
)
|
||||
)
|
||||
call, rhs_value = emit_view_body(fn, TMP_VAR)
|
||||
inplace_view_body.append(call)
|
||||
if rhs_value is None:
|
||||
raise AssertionError("Expected rhs_value to be non-None")
|
||||
inplace_view_body.append(
|
||||
ASSIGN_RETURN_VALUE.substitute(
|
||||
return_values=tie_return_values(f), rhs_value=rhs_value
|
||||
)
|
||||
)
|
||||
if f.func.returns:
|
||||
inplace_view_body.append(f"return {get_return_value(f)};")
|
||||
return inplace_view_body
|
||||
|
||||
|
||||
@with_native_function
|
||||
def gen_formals(f: NativeFunction) -> str:
|
||||
return ", ".join(
|
||||
# code-generated autograd kernels plumb and recompute dispatch keys directly through the kernel for performance.
|
||||
# See Note [Plumbing Keys Through The Dispatcher] for details.
|
||||
["c10::DispatchKeySet ks"]
|
||||
+ [
|
||||
f"{cpp.argument_type(a, binds='__placeholder__', symint=True).cpp_type()} {a.name}"
|
||||
for a in f.func.schema_order_arguments()
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@with_native_function_with_differentiability_info
|
||||
def inplace_or_view_method_definition(
|
||||
fn: NativeFunctionWithDifferentiabilityInfo,
|
||||
) -> str | None:
|
||||
f = fn.func
|
||||
if get_view_info(f) is None and (
|
||||
# For functions that modify their inputs but don't return them,
|
||||
# we can't give them autograd support.
|
||||
# See https://github.com/pytorch/pytorch/issues/53796
|
||||
not modifies_arguments(f) or len(f.func.returns) == 0
|
||||
):
|
||||
return None
|
||||
return METHOD_DEFINITION.substitute(
|
||||
return_type=cpp.returns_type(f.func.returns, symint=True).cpp_type(),
|
||||
type_wrapper_name=type_wrapper_name(f),
|
||||
formals=gen_formals(f),
|
||||
type_definition_body=emit_inplace_or_view_body(fn),
|
||||
)
|
||||
|
||||
|
||||
@with_native_function_with_differentiability_info
|
||||
def inplace_or_view_method_registration(
|
||||
fn: NativeFunctionWithDifferentiabilityInfo,
|
||||
) -> str | None:
|
||||
f = fn.func
|
||||
if get_view_info(f) is None and (
|
||||
not modifies_arguments(f) or len(f.func.returns) == 0
|
||||
):
|
||||
return None
|
||||
return WRAPPER_REGISTRATION.substitute(
|
||||
unqual_operator_name_with_overload=f.func.name,
|
||||
type_wrapper_name=type_wrapper_name(f),
|
||||
class_type="ADInplaceOrView",
|
||||
)
|
||||
|
||||
|
||||
def use_derived(fn: NativeFunctionWithDifferentiabilityInfo) -> bool:
|
||||
f = fn.func
|
||||
name = cpp.name(f.func)
|
||||
return name not in MANUAL_AUTOGRAD and dispatch_strategy(fn) == "use_derived"
|
||||
|
||||
|
||||
def gen_inplace_or_view_type_env(
|
||||
fn: NativeFunctionWithDifferentiabilityInfo,
|
||||
) -> dict[str, list[str]]:
|
||||
definition = inplace_or_view_method_definition(fn)
|
||||
registration = inplace_or_view_method_registration(fn)
|
||||
|
||||
return {
|
||||
"ops_headers": (
|
||||
[f"#include <ATen/ops/{fn.func.root_name}_ops.h>"]
|
||||
if definition is not None
|
||||
else []
|
||||
),
|
||||
"inplace_or_view_method_definitions": [definition]
|
||||
if definition is not None
|
||||
else [],
|
||||
"inplace_or_view_wrapper_registrations": [registration]
|
||||
if registration is not None
|
||||
else [],
|
||||
}
|
||||
|
||||
|
||||
def gen_inplace_or_view_type(
|
||||
out: str,
|
||||
native_yaml_path: str,
|
||||
tags_yaml_path: str,
|
||||
fns_with_infos: list[NativeFunctionWithDifferentiabilityInfo],
|
||||
template_path: str,
|
||||
) -> None:
|
||||
# NOTE: see Note [Sharded File] at the top of the VariableType.cpp
|
||||
# template regarding sharding of the generated files.
|
||||
|
||||
fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
|
||||
fm.write_sharded(
|
||||
"ADInplaceOrViewType.cpp",
|
||||
[fn for fn in fns_with_infos if use_derived(fn)],
|
||||
key_fn=lambda fn: fn.func.root_name,
|
||||
base_env={
|
||||
"generated_comment": "@"
|
||||
+ f"generated from {fm.template_dir_for_comments()}/ADInplaceOrViewType.cpp",
|
||||
},
|
||||
env_callable=gen_inplace_or_view_type_env,
|
||||
num_shards=2,
|
||||
sharded_keys={
|
||||
"ops_headers",
|
||||
"inplace_or_view_method_definitions",
|
||||
"inplace_or_view_wrapper_registrations",
|
||||
},
|
||||
)
|
||||
+1413
File diff suppressed because it is too large
Load Diff
+542
@@ -0,0 +1,542 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import DispatcherSignature
|
||||
from torchgen.code_template import CodeTemplate
|
||||
from torchgen.context import with_native_function
|
||||
from torchgen.model import Argument, NativeFunction, SchemaKind, TensorOptionsArguments
|
||||
from torchgen.utils import FileManager
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# Note [Manual Backend kernels]
|
||||
# For these ops, we want to manually register to dispatch key Backend and
|
||||
# skip codegen-ed registration to all keys before Backend.
|
||||
# For codegen this means:
|
||||
# - op set below must match ops with manual_kernel_registration=True in native_functions.yaml
|
||||
# where we skip codegen backend kernels
|
||||
# - all ops below are part of MANUAL_AUTOGRAD to skip codegen Autograd kernel registration
|
||||
# - all ops below are part of MANUAL_TRACER to skip codegen Tracer kernel registration
|
||||
# Note: we still register to dispatch key Profiler for these ops, keeping it untouched for now.
|
||||
# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp
|
||||
MANUAL_BACKEND = {
|
||||
"options",
|
||||
"data",
|
||||
"set_data",
|
||||
"is_leaf",
|
||||
"output_nr",
|
||||
"_version",
|
||||
"retain_grad",
|
||||
"_backward",
|
||||
"requires_grad_",
|
||||
}
|
||||
|
||||
# For these ops we want to skip the codegen-ed registration to both Autograd and Tracer keys.
|
||||
# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp
|
||||
MANUAL_AUTOGRAD_AND_TRACER = {
|
||||
"resize_",
|
||||
"resize_as_",
|
||||
"detach",
|
||||
"detach_",
|
||||
"copy_",
|
||||
"_fw_primal",
|
||||
"_make_dual",
|
||||
}
|
||||
|
||||
# Currently MANUAL_AUTOGRAD and MANUAL_TRACER share the same set of ops:
|
||||
# union(MANUAL_BACKEND, MANUAL_AUTOGRAD_AND_TRACER)
|
||||
# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp
|
||||
MANUAL_AUTOGRAD = MANUAL_TRACER = MANUAL_BACKEND | MANUAL_AUTOGRAD_AND_TRACER
|
||||
|
||||
# These functions we don't want to record for tracing, because we always want
|
||||
# to trace their constituent parts. This is a temporary hack in lieue
|
||||
# of proper scopes, where subsequent compilation passes can ask for the unfolding
|
||||
# on demand. Only concrete ATen methods can be disabled this way; it will have
|
||||
# NO EFFECT otherwise.
|
||||
DONT_RECORD_TRACE = {
|
||||
"convolution",
|
||||
"conv1d",
|
||||
"conv2d",
|
||||
"conv3d",
|
||||
"conv_transpose1d",
|
||||
"conv_transpose2d",
|
||||
"conv_transpose3d",
|
||||
"lstm_cell",
|
||||
"gru_cell",
|
||||
"rnn_tanh_cell",
|
||||
"rnn_relu_cell",
|
||||
# FIXME: figure out a better way when we support sparse tensors in jit
|
||||
"_coalesced",
|
||||
}
|
||||
|
||||
|
||||
def should_trace(f: NativeFunction) -> bool:
|
||||
# Operations involving Storage or Type are not traceable at the moment
|
||||
if any(
|
||||
str(arg.type) in {"Storage", "Type"} for arg in f.func.schema_order_arguments()
|
||||
):
|
||||
return False
|
||||
# We can't trace functions which don't have any Tensor or TensorList returns
|
||||
if not any(r.type.is_tensor_like() for r in f.func.returns):
|
||||
return False
|
||||
return f.func.name.name.base not in DONT_RECORD_TRACE
|
||||
|
||||
|
||||
SELECT = CodeTemplate(
|
||||
"""\
|
||||
|
||||
if (${cond}) {
|
||||
${true}
|
||||
} else {
|
||||
${false}
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
OP_NAME = CodeTemplate(
|
||||
"""\
|
||||
op_name = c10::Symbol::fromQualString("aten::${trace_name}");
|
||||
"""
|
||||
)
|
||||
|
||||
# These functions have their names recorded under trace renamed,
|
||||
RENAME_TRACE = {
|
||||
"zero": "zeros_like", # replacing aten::zero_ with aten::zeros_like
|
||||
"fill": "full_like", # replacing aten::fill_ with aten::full_like
|
||||
}
|
||||
|
||||
|
||||
def format_trace_op_name(f: NativeFunction) -> str:
|
||||
# TODO: byte-for-byte compatible with old codegen behavior - should clean up
|
||||
if (
|
||||
f.func.kind() in (SchemaKind.functional, SchemaKind.out)
|
||||
or f.func.name.name.dunder_method
|
||||
):
|
||||
# special case for *_out functions: the in-place and out-of-place ops
|
||||
# are overloaded with the same name in the JIT
|
||||
trace_name = str(f.func.name.name)
|
||||
trace_name = RENAME_TRACE.get(trace_name, trace_name)
|
||||
return OP_NAME.substitute(trace_name=trace_name)
|
||||
|
||||
# otherwise, this is an in-place op and we need to emit both in- and
|
||||
# out-of-place versions
|
||||
outplace_trace_name = f.func.name.name.base
|
||||
inplace_trace_name = cpp.name(f.func)
|
||||
outplace_trace_name = RENAME_TRACE.get(outplace_trace_name, outplace_trace_name)
|
||||
inplace_trace_name = RENAME_TRACE.get(inplace_trace_name, inplace_trace_name)
|
||||
|
||||
return SELECT.substitute(
|
||||
cond="tracer_state->force_outplace",
|
||||
true=OP_NAME.substitute(trace_name=outplace_trace_name),
|
||||
false=OP_NAME.substitute(trace_name=inplace_trace_name),
|
||||
)
|
||||
|
||||
|
||||
ADD_TRACE_INPUT = CodeTemplate("""jit::tracer::addInputs(node, "${name}", ${input});""")
|
||||
|
||||
|
||||
def format_trace_inputs(f: NativeFunction) -> str:
|
||||
def dispatch_trace_input(arg: Argument | TensorOptionsArguments) -> Sequence[str]:
|
||||
if isinstance(arg, TensorOptionsArguments):
|
||||
name = "options"
|
||||
return [
|
||||
ADD_TRACE_INPUT.substitute(
|
||||
name=name, input="c10::optTypeMetaToScalarType(options.dtype_opt())"
|
||||
),
|
||||
ADD_TRACE_INPUT.substitute(name=name, input="options.layout()"),
|
||||
ADD_TRACE_INPUT.substitute(name=name, input="options.device()"),
|
||||
ADD_TRACE_INPUT.substitute(name=name, input="options.pinned_memory()"),
|
||||
]
|
||||
else:
|
||||
name = arg.name
|
||||
if str(arg.type) == "Tensor?[]":
|
||||
return [f'jit::tracer::addInputs(node, "{name}", {name});']
|
||||
else:
|
||||
return [ADD_TRACE_INPUT.substitute(name=name, input=name)]
|
||||
|
||||
args: list[Argument | TensorOptionsArguments] = list(
|
||||
f.func.schema_order_arguments()
|
||||
)
|
||||
|
||||
if f.func.is_out_fn():
|
||||
# *_out functions take the result as a separate argument, but we don't want to
|
||||
# trace that argument directly. Instead, we trace its TensorOptions.
|
||||
# So first, we need to remove the out argument from the list of arguments to trace.
|
||||
num_out_args = len(f.func.arguments.out)
|
||||
args = args[:-num_out_args]
|
||||
|
||||
trace_inputs = itertools.chain.from_iterable(
|
||||
dispatch_trace_input(arg) for arg in args
|
||||
)
|
||||
|
||||
if f.func.is_out_fn():
|
||||
# for *_out functions, handle the result argument differently for inplace/outplace.
|
||||
# For inplace: just add the input to the end to confirm with the JIT schema
|
||||
inplace = [
|
||||
ADD_TRACE_INPUT.substitute(
|
||||
name=f.func.arguments.out[i].name, input=f.func.arguments.out[i].name
|
||||
)
|
||||
# pyrefly: ignore [unbound-name]
|
||||
for i in range(num_out_args)
|
||||
]
|
||||
|
||||
# for outplace: do nothing, except if the function is a factory.
|
||||
# Factories are a bit special because their out-of-place overloads
|
||||
# take an extra TensorOptions argument, which is missing in the _out function
|
||||
has_tensor_return = any(r.type.is_tensor_like() for r in f.func.returns)
|
||||
has_tensor_input_arg = any(
|
||||
a.type.is_tensor_like() for a in f.func.arguments.flat_non_out
|
||||
)
|
||||
is_factory_method = f.category_override == "factory" or (
|
||||
has_tensor_return and not has_tensor_input_arg
|
||||
)
|
||||
|
||||
# HACK: preserve old codegen behavior - the old codegen set the `is_factory_method`
|
||||
# flag for the whole family of ops with the same basename if any of them is a
|
||||
# factory method. For most cases the whole family of ops are indeed all factory
|
||||
# method - 'normal' is the only exception. So we handle it specially here to avoid
|
||||
# cloning the old logic.
|
||||
if f.func.name.name.base == "normal":
|
||||
is_factory_method = True
|
||||
|
||||
if is_factory_method:
|
||||
outplace = [
|
||||
ADD_TRACE_INPUT.substitute(
|
||||
name="out",
|
||||
input="c10::optTypeMetaToScalarType(out.options().dtype_opt())",
|
||||
),
|
||||
ADD_TRACE_INPUT.substitute(name="out", input="out.options().layout()"),
|
||||
ADD_TRACE_INPUT.substitute(name="out", input="out.options().device()"),
|
||||
ADD_TRACE_INPUT.substitute(
|
||||
name="out", input="out.options().pinned_memory()"
|
||||
),
|
||||
]
|
||||
else:
|
||||
outplace = []
|
||||
|
||||
trace_inputs = itertools.chain(
|
||||
trace_inputs,
|
||||
[
|
||||
SELECT.substitute(
|
||||
cond="tracer_state->force_outplace",
|
||||
true="\n".join(outplace),
|
||||
false="\n".join(inplace),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
return "\n".join(trace_inputs)
|
||||
|
||||
|
||||
# `torch.jit.trace` have undocumented keyword argument `_force_outplace`,
|
||||
# which force jit to replace functions with outplace variants (for
|
||||
# example `aten::add_` becomes `aten::add`).
|
||||
#
|
||||
# This replacement implemented in-place with minimum modifications of
|
||||
# arguments stack (as it assumes that outplace call has the same arguments
|
||||
# as inplace version).
|
||||
#
|
||||
# However there are no such substitutions available for `aten::fill_`
|
||||
# and `aten::zero_` operators, as we never implemented `aten::fill`
|
||||
# and `aten::zero`. So jit tracing hack replacing `aten::zero_` with
|
||||
# `aten::zeros_like` and replacing `aten::fill_` with `aten::full_like`.
|
||||
#
|
||||
# But as they potentially can have different arguments, we also have
|
||||
# to hack into the stack and add missing ones.
|
||||
#
|
||||
# A possible alternative would be:
|
||||
#
|
||||
# - Add `aten::fill` and `aten::zero`
|
||||
#
|
||||
# - Or keep `aten::zeros_like` arguments aligned with `aten::zero_`
|
||||
# arguments (inside of the `native_functions.yaml`)
|
||||
RENAME_TRACE_ADD_ARGS = {
|
||||
"fill": """\
|
||||
jit::tracer::addInputs(node, "options", ::std::optional<ScalarType>());
|
||||
jit::tracer::addInputs(node, "options", layout_or_default(::std::nullopt));
|
||||
jit::tracer::addInputs(node, "options", device_or_default(::std::nullopt));
|
||||
jit::tracer::addInputs(node, "options", pinned_memory_or_default(::std::nullopt));
|
||||
::std::optional<MemoryFormat> memory_format = c10::MemoryFormat::Preserve;
|
||||
jit::tracer::addInputs(node, "memory_format", memory_format);
|
||||
""",
|
||||
"zero": """\
|
||||
jit::tracer::addInputs(node, "options", ::std::optional<ScalarType>());
|
||||
jit::tracer::addInputs(node, "options", layout_or_default(::std::nullopt));
|
||||
jit::tracer::addInputs(node, "options", device_or_default(::std::nullopt));
|
||||
jit::tracer::addInputs(node, "options", pinned_memory_or_default(::std::nullopt));
|
||||
::std::optional<MemoryFormat> memory_format = c10::MemoryFormat::Preserve;
|
||||
jit::tracer::addInputs(node, "memory_format", memory_format);
|
||||
""",
|
||||
}
|
||||
|
||||
INPLACE_GUARD = CodeTemplate(
|
||||
"""\
|
||||
jit::tracer::ensureUniqueIfOutOfPlaced("${name}", ${mutable_input});
|
||||
"""
|
||||
)
|
||||
|
||||
PRE_RECORD_TRACE = CodeTemplate(
|
||||
"""\
|
||||
torch::jit::Node* node = nullptr;
|
||||
std::shared_ptr<jit::tracer::TracingState> tracer_state;
|
||||
if (jit::tracer::isTracing()) {
|
||||
tracer_state = jit::tracer::getTracingState();
|
||||
at::Symbol op_name;
|
||||
${set_op_name}
|
||||
node = tracer_state->createNode(op_name, /*num_outputs=*/0);
|
||||
jit::tracer::recordSourceLocation(node);
|
||||
${add_trace_inputs}
|
||||
tracer_state->insertNode(node);
|
||||
${inplace_guard}
|
||||
jit::tracer::setTracingState(nullptr);
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def format_prerecord_trace(f: NativeFunction) -> str:
|
||||
if not should_trace(f):
|
||||
return ""
|
||||
|
||||
# TODO: clean up old codegen behavior
|
||||
is_inplace = (
|
||||
f.func.kind() in (SchemaKind.inplace, SchemaKind.out)
|
||||
and not f.func.name.name.dunder_method
|
||||
)
|
||||
add_args = (
|
||||
RENAME_TRACE_ADD_ARGS.get(f.func.name.name.base, "") if is_inplace else ""
|
||||
)
|
||||
additional_inputs = (
|
||||
SELECT.substitute(
|
||||
cond="tracer_state->force_outplace",
|
||||
true=add_args,
|
||||
false="",
|
||||
)
|
||||
if add_args
|
||||
else ""
|
||||
)
|
||||
|
||||
return PRE_RECORD_TRACE.substitute(
|
||||
set_op_name=format_trace_op_name(f),
|
||||
add_trace_inputs=format_trace_inputs(f) + additional_inputs,
|
||||
inplace_guard=INPLACE_GUARD.substitute(
|
||||
name=cpp.name(f.func),
|
||||
mutable_input=f.func.arguments.out[0].name
|
||||
if f.func.arguments.out
|
||||
else "self",
|
||||
)
|
||||
if is_inplace
|
||||
else "",
|
||||
)
|
||||
|
||||
|
||||
POST_RECORD_TRACE = CodeTemplate(
|
||||
"""\
|
||||
if (tracer_state) {
|
||||
jit::tracer::setTracingState(std::move(tracer_state));
|
||||
${add_trace_outputs}
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def format_postrecord_trace(f: NativeFunction) -> str:
|
||||
if not should_trace(f):
|
||||
return ""
|
||||
|
||||
# For outplacing ops, *_out overloads require special handling to move the
|
||||
# output *argument* to a return value
|
||||
if f.func.is_out_fn():
|
||||
output_names_outplace = [arg.name for arg in f.func.arguments.out]
|
||||
output_names_inplace = cpp.return_names(f)
|
||||
|
||||
# Code size optimization: the common case is that the return value is
|
||||
# the same for both variants
|
||||
if output_names_outplace == output_names_inplace:
|
||||
outputs = [
|
||||
f"jit::tracer::addOutput(node, {n});" for n in output_names_outplace
|
||||
]
|
||||
return POST_RECORD_TRACE.substitute(add_trace_outputs=outputs)
|
||||
|
||||
selection = SELECT.substitute(
|
||||
cond="force_outplace",
|
||||
true="\n".join(
|
||||
f"jit::tracer::addOutput(node, {n});" for n in output_names_outplace
|
||||
),
|
||||
false="\n".join(
|
||||
f"jit::tracer::addOutput(node, {n});" for n in output_names_inplace
|
||||
),
|
||||
)
|
||||
return POST_RECORD_TRACE.substitute(add_trace_outputs=selection)
|
||||
else:
|
||||
output_names = cpp.return_names(f)
|
||||
outputs = [f"jit::tracer::addOutput(node, {n});" for n in output_names]
|
||||
return POST_RECORD_TRACE.substitute(add_trace_outputs=outputs)
|
||||
|
||||
|
||||
def tie_return_values(f: NativeFunction) -> str:
|
||||
if len(f.func.returns) == 1:
|
||||
return f"auto {f.func.returns[0].name or 'result'}"
|
||||
names = cpp.return_names(f)
|
||||
return f"auto [{', '.join(names)}]"
|
||||
|
||||
|
||||
def get_return_value(f: NativeFunction) -> str:
|
||||
names = cpp.return_names(f)
|
||||
if len(f.func.returns) == 1:
|
||||
return names[0]
|
||||
if f.func.kind() == SchemaKind.out:
|
||||
return f"std::forward_as_tuple({', '.join(names)})"
|
||||
else:
|
||||
moved = ", ".join(f"std::move({name})" for name in names)
|
||||
return f"std::make_tuple({moved})"
|
||||
|
||||
|
||||
TRACE_DISPATCH = CodeTemplate(
|
||||
"""\
|
||||
${assign_return_values}at::_ops::${unambiguous_name}::redispatch(${unpacked_args});"""
|
||||
)
|
||||
|
||||
|
||||
def emit_trace_body(f: NativeFunction) -> list[str]:
|
||||
trace_body: list[str] = []
|
||||
|
||||
trace_body.append(format_prerecord_trace(f))
|
||||
|
||||
dispatcher_sig = DispatcherSignature.from_schema(f.func)
|
||||
dispatcher_exprs = dispatcher_sig.exprs()
|
||||
|
||||
# code-generated tracing kernels plumb and recompute dispatch keys directly through the kernel for performance.
|
||||
# See Note [Plumbing Keys Through The Dispatcher] for details.
|
||||
dispatch_key_set = "ks & c10::DispatchKeySet(c10::DispatchKeySet::FULL_AFTER, c10::DispatchKey::Tracer)"
|
||||
redispatch_args = ", ".join([dispatch_key_set] + [a.expr for a in dispatcher_exprs])
|
||||
|
||||
assign_return_values = (
|
||||
f"{tie_return_values(f)} = "
|
||||
if f.func.kind() in [SchemaKind.functional, SchemaKind.mutable]
|
||||
and f.func.returns
|
||||
else ""
|
||||
)
|
||||
|
||||
# Note that this calls the slow, dispatching variants of manual_cpp_binding ops.
|
||||
# We could probably work harder to ensure that the fast variants are
|
||||
# called instead, but the perf benefit would be minimal.
|
||||
trace_body.append(
|
||||
TRACE_DISPATCH.substitute(
|
||||
assign_return_values=assign_return_values,
|
||||
unambiguous_name=f.func.name.unambiguous_name(),
|
||||
unpacked_args=redispatch_args,
|
||||
)
|
||||
)
|
||||
|
||||
trace_body.append(format_postrecord_trace(f))
|
||||
if f.func.returns:
|
||||
trace_body.append(f"return {get_return_value(f)};")
|
||||
return trace_body
|
||||
|
||||
|
||||
METHOD_DEFINITION = CodeTemplate(
|
||||
"""\
|
||||
${return_type} ${type_wrapper_name}(${formals}) {
|
||||
${type_definition_body}
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def type_wrapper_name(f: NativeFunction, key: str = "Default") -> str:
|
||||
if f.func.name.overload_name:
|
||||
name = f"{cpp.name(f.func)}_{f.func.name.overload_name}"
|
||||
else:
|
||||
name = cpp.name(f.func)
|
||||
|
||||
# The key argument is only used in gen_variable_type where we need fns per autograd dispatch key.
|
||||
# In gen_trace_type and gen_inplace_view_type where only one fn per native_fn must be generated,
|
||||
# the key argument should not be passed.
|
||||
# We do not append key if it is Default so that generated functions from
|
||||
# before per-dispatch-key derivatives were added retain the same names.
|
||||
if key != "Default":
|
||||
name = name + f"_{key}"
|
||||
return name
|
||||
|
||||
|
||||
@with_native_function
|
||||
def method_definition(f: NativeFunction) -> str:
|
||||
if cpp.name(f.func) in MANUAL_TRACER:
|
||||
raise AssertionError(f"Function {cpp.name(f.func)} is in MANUAL_TRACER")
|
||||
|
||||
formals = ", ".join(
|
||||
# code-generated tracing kernels plumb and recompute dispatch keys directly through the kernel for performance.
|
||||
# See Note [Plumbing Keys Through The Dispatcher] for details.
|
||||
["c10::DispatchKeySet ks"]
|
||||
+ [
|
||||
f"{cpp.argument_type(a, binds='__placeholder__', symint=True).cpp_type()} {a.name}"
|
||||
for a in f.func.schema_order_arguments()
|
||||
]
|
||||
)
|
||||
|
||||
return METHOD_DEFINITION.substitute(
|
||||
return_type=cpp.returns_type(f.func.returns, symint=True).cpp_type(),
|
||||
type_wrapper_name=type_wrapper_name(f),
|
||||
formals=formals,
|
||||
type_definition_body=emit_trace_body(f),
|
||||
)
|
||||
|
||||
|
||||
WRAPPER_REGISTRATION = CodeTemplate(
|
||||
"""\
|
||||
m.impl("${name}",
|
||||
TORCH_FN(${class_type}::${type_wrapper_name})
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
@with_native_function
|
||||
def method_registration(f: NativeFunction) -> str:
|
||||
if cpp.name(f.func) in MANUAL_TRACER:
|
||||
raise AssertionError(f"Function {cpp.name(f.func)} is in MANUAL_TRACER")
|
||||
|
||||
return WRAPPER_REGISTRATION.substitute(
|
||||
name=f.func.name,
|
||||
type_wrapper_name=type_wrapper_name(f),
|
||||
class_type="TraceType",
|
||||
)
|
||||
|
||||
|
||||
def gen_trace_type_func(fn: NativeFunction) -> dict[str, list[str]]:
|
||||
return {
|
||||
"ops_headers": [f"#include <ATen/ops/{fn.root_name}_ops.h>"],
|
||||
"trace_method_definitions": [method_definition(fn)],
|
||||
"trace_wrapper_registrations": [method_registration(fn)],
|
||||
}
|
||||
|
||||
|
||||
def gen_trace_type(
|
||||
out: str, native_functions: list[NativeFunction], template_path: str
|
||||
) -> None:
|
||||
# NOTE: see Note [Sharded File] at the top of the VariableType.cpp
|
||||
# template regarding sharding of the generated files.
|
||||
fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
|
||||
fm.write_sharded(
|
||||
"TraceType.cpp",
|
||||
[fn for fn in native_functions if cpp.name(fn.func) not in MANUAL_TRACER],
|
||||
key_fn=lambda fn: fn.root_name,
|
||||
base_env={
|
||||
"generated_comment": "@"
|
||||
+ f"generated from {fm.template_dir_for_comments()}/TraceType.cpp",
|
||||
},
|
||||
env_callable=gen_trace_type_func,
|
||||
num_shards=5,
|
||||
sharded_keys={
|
||||
"ops_headers",
|
||||
"trace_method_definitions",
|
||||
"trace_wrapper_registrations",
|
||||
},
|
||||
)
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
# Generates C++ functions that wrap ATen tensor factory methods to turn them into Variables.
|
||||
#
|
||||
# This writes one file: variable_factories.h
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
import torchgen.api.python as python
|
||||
from torchgen.api import cpp
|
||||
from torchgen.api.types import CppSignatureGroup
|
||||
from torchgen.context import with_native_function
|
||||
from torchgen.gen import parse_native_yaml
|
||||
from torchgen.model import NativeFunction, TensorOptionsArguments, Variant
|
||||
from torchgen.utils import FileManager, mapMaybe
|
||||
|
||||
|
||||
OPTIONAL_TYPE_PATTERN = re.compile(r"std::optional<(.+)>")
|
||||
TYPE_PATTERN = re.compile(r"(?:const\s+)?([A-Z]\w+)")
|
||||
|
||||
|
||||
# Add 'at::' to types defined in ATen namespace, e.g. Tensor, TensorList, IntArrayRef and etc.
|
||||
# TODO: maybe update the cpp argument API to take optional namespace argument?
|
||||
def fully_qualified_type(argument_type: str) -> str:
|
||||
def maybe_optional_type(type: str, is_opt: bool) -> str:
|
||||
return f"std::optional<{type}>" if is_opt else type
|
||||
|
||||
opt_match = OPTIONAL_TYPE_PATTERN.match(argument_type)
|
||||
is_opt = opt_match is not None
|
||||
if opt_match:
|
||||
argument_type = argument_type[opt_match.start(1) : opt_match.end(1)]
|
||||
match = TYPE_PATTERN.match(argument_type)
|
||||
if match is None:
|
||||
return maybe_optional_type(argument_type, is_opt)
|
||||
index = match.start(1)
|
||||
qualified_type = f"{argument_type[:index]}at::{argument_type[index:]}"
|
||||
return maybe_optional_type(qualified_type, is_opt)
|
||||
|
||||
|
||||
def gen_variable_factories(
|
||||
out: str, native_yaml_path: str, tags_yaml_path: str, template_path: str
|
||||
) -> None:
|
||||
native_functions = parse_native_yaml(
|
||||
native_yaml_path, tags_yaml_path
|
||||
).native_functions
|
||||
factory_functions = [fn for fn in native_functions if is_factory_function(fn)]
|
||||
fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)
|
||||
fm.write_with_template(
|
||||
"variable_factories.h",
|
||||
"variable_factories.h",
|
||||
lambda: {
|
||||
"generated_comment": "@"
|
||||
+ f"generated from {fm.template_dir_for_comments()}/variable_factories.h",
|
||||
"ops_headers": [
|
||||
f"#include <ATen/ops/{fn.root_name}.h>" for fn in factory_functions
|
||||
],
|
||||
"function_definitions": list(mapMaybe(process_function, factory_functions)),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@with_native_function
|
||||
def is_factory_function(f: NativeFunction) -> bool:
|
||||
if Variant.function not in f.variants:
|
||||
return False
|
||||
|
||||
name = cpp.name(f.func)
|
||||
has_tensor_options = python.has_tensor_options(f)
|
||||
return has_tensor_options or name.endswith("_like")
|
||||
|
||||
|
||||
@with_native_function
|
||||
def process_function(f: NativeFunction) -> str | None:
|
||||
name = cpp.name(f.func)
|
||||
has_tensor_options = python.has_tensor_options(f)
|
||||
is_factory = has_tensor_options or name.endswith("_like")
|
||||
|
||||
if Variant.function not in f.variants or not is_factory:
|
||||
return None
|
||||
|
||||
cpp_sigs = CppSignatureGroup.from_native_function(f, method=False)
|
||||
sigs = [cpp_sigs.signature]
|
||||
if cpp_sigs.symint_signature is not None:
|
||||
sigs.append(cpp_sigs.symint_signature)
|
||||
r = ""
|
||||
for sig in sigs:
|
||||
formals: list[str] = []
|
||||
exprs: list[str] = []
|
||||
requires_grad = "false"
|
||||
for arg in sig.arguments():
|
||||
qualified_type = fully_qualified_type(arg.type)
|
||||
if arg.default:
|
||||
formals.append(f"{qualified_type} {arg.name} = {arg.default}")
|
||||
else:
|
||||
formals.append(f"{qualified_type} {arg.name}")
|
||||
|
||||
if isinstance(arg.argument, TensorOptionsArguments):
|
||||
# note: we remove the requires_grad setting from the TensorOptions because
|
||||
# it is ignored anyways (and we actually have an assertion that it isn't set
|
||||
# which would fail otherwise). We handle requires_grad explicitly here
|
||||
# instead of passing it through to the kernel.
|
||||
exprs.append(
|
||||
f"at::TensorOptions({arg.name}).requires_grad(::std::nullopt)"
|
||||
)
|
||||
# Manually set the requires_grad bit on the result tensor.
|
||||
requires_grad = f"{arg.name}.requires_grad()"
|
||||
else:
|
||||
exprs.append(arg.name)
|
||||
|
||||
r += f"""\
|
||||
inline at::Tensor {sig.name()}({", ".join(formals)}) {{
|
||||
at::AutoDispatchBelowADInplaceOrView guard;
|
||||
return autograd::make_variable(at::{sig.name()}({", ".join(exprs)}), /*requires_grad=*/{requires_grad});
|
||||
}}
|
||||
"""
|
||||
return r
|
||||
+2266
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user