def main() -> None:
parser = argparse.ArgumentParser(description="Generate ATen source files")
parser.add_argument(
"-s",
"--source-path",
help="path to source directory for ATen",
default="caffe2/aten/src/ATen",
)
parser.add_argument(
"-p",
"--generated-ops-cpp-path",
help="path to directory to generate op dispatcher .cpp file",
default="caffe2/torch/csrc/jit/runtime/static/generated_ops.cpp",
)
parser.add_argument(
"-t",
"--generated-ops-test-cpp-path",
help="path to directory to generate op dispatcher .cpp file",
default="caffe2/benchmarks/static_runtime/test_generated_ops.cc",
)
options = parser.parse_args()
native_yaml_path = os.path.join(options.source_path,
"native/native_functions.yaml")
tags_yaml_path = os.path.join(options.source_path, "native/tags.yaml")
parsed_yaml = gen.parse_native_yaml(native_yaml_path, tags_yaml_path)
native_functions, backend_indices = (
parsed_yaml.native_functions,
parsed_yaml.backend_indices,
)
grouped_native_functions = gen.get_grouped_native_functions(
native_functions)
structured_native_functions = [
g for g in grouped_native_functions
if isinstance(g, NativeFunctionsGroup)
]
supported_function_groups = group_functions_by_op_name(
structured_native_functions)
gen_out_variant_dispatcher = generator.GenOutVariantDispatcher()
result = [
gen_out_variant_dispatcher(groups, backend_indices[DispatchKey.CPU])
for groups in supported_function_groups
]
gen_out_variant_dispatcher_test_case = generator.GenOutVariantDispatcherTestCase(
)
test_result = [
gen_out_variant_dispatcher_test_case(groups)
for groups in supported_function_groups
]
write_cpp(result, options.generated_ops_cpp_path)
write_test_cpp(test_result, options.generated_ops_test_cpp_path)
print("total grouped native ops: %d" % len(grouped_native_functions))
print("structured grouped native ops: %d" %
len(structured_native_functions))
supported_grouped_functions = sum(
[len(groups) for groups in supported_function_groups])
print("generated grouped native ops: %d" % supported_grouped_functions)
def main() -> None:
parser = argparse.ArgumentParser(description="Generate unboxing source files")
parser.add_argument(
"-s",
"--source-path",
help="path to source directory for ATen",
default="aten/src/ATen",
)
parser.add_argument(
"-d", "--install_dir", help="output directory", default="build/aten/src/ATen"
)
parser.add_argument(
"-o",
"--output-dependencies",
help="output a list of dependencies into the given file and exit",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="run without writing any files (still updates outputs)",
)
parser.add_argument(
"--op_selection_yaml_path",
help="Provide a path to the operator selection (for custom build) YAML "
"that contains the information about the set of selected operators "
"and their categories (training, ...). Each operator is either a "
"full operator name with overload or just a bare operator name. "
"The operator names also contain the namespace prefix (e.g. aten::)",
)
options = parser.parse_args()
if options.op_selection_yaml_path is not None:
selector = SelectiveBuilder.from_yaml_path(options.op_selection_yaml_path)
else:
selector = SelectiveBuilder.get_nop_selector()
native_yaml_path = os.path.join(options.source_path, "native/native_functions.yaml")
tags_yaml_path = os.path.join(options.source_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,
)
cpu_fm = make_file_manager(options=options)
gen_unboxing(native_functions=native_functions, cpu_fm=cpu_fm, selector=selector)
if options.output_dependencies:
depfile_path = pathlib.Path(options.output_dependencies).resolve()
depfile_name = depfile_path.name
depfile_stem = depfile_path.stem
path = depfile_path.parent / depfile_name
cpu_fm.write_outputs(depfile_stem, str(path))
def load_derivatives(derivatives_yaml_path: str, native_yaml_path: str,
tags_yaml_path: str) -> Sequence[DifferentiabilityInfo]:
# Do some caching as this is a deterministic function
global _GLOBAL_LOAD_DERIVATIVE_CACHE
key = (derivatives_yaml_path, native_yaml_path)
if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE:
with open(derivatives_yaml_path, "r") as f:
definitions = yaml.load(f, Loader=YamlLoader)
funcs = parse_native_yaml(native_yaml_path,
tags_yaml_path).native_functions
# From the parsed native functions, separate out the (generated) view_copy functions,
# so we can generate derivatives for them separately.
native_functions_with_view_groups = get_grouped_by_view_native_functions(
funcs)
native_functions_without_view_copies = concatMap(
# We need to pull out the view_inplace ops too, since they might have their own derivative entries.
lambda g: [g] if isinstance(g, NativeFunction) else list(
g.functions(include_copy=False)),
native_functions_with_view_groups,
)
view_groups = [
g for g in native_functions_with_view_groups
if isinstance(g, NativeFunctionsViewGroup)
]
# What's the difference between function schema v.s. signature?
# function schema is the complete declaration including mutability annotation / default value and etc.
# signature is the canonical schema for a group of functions (in-place/out/functional variants)
# that are semantically related.
functions_by_signature: Dict[FunctionSchema,
List[NativeFunction]] = defaultdict(list)
functions_by_schema: Dict[str, NativeFunction] = dict()
for function in native_functions_without_view_copies:
functions_by_signature[function.func.signature()].append(function)
assert str(function.func) not in functions_by_schema
functions_by_schema[str(function.func)] = function
# Keep track of how many of which ops we've seen so we can
# disambiguate them with a numeric suffix.
op_counter = Counter[str]()
infos = [
create_differentiability_info(defn, functions_by_signature,
functions_by_schema, op_counter)
for defn in definitions
]
infos += add_view_copy_derivatives(infos, view_groups)
_GLOBAL_LOAD_DERIVATIVE_CACHE[key] = infos
return _GLOBAL_LOAD_DERIVATIVE_CACHE[key]
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 = list(
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)
def generate_code(install_dir='functorch/csrc', source_path=None):
if source_path is None:
# infer the source path via torchgen
source_path = os.path.join(get_torchgen_root(), "packaged/ATen")
native_yaml_path = os.path.join(source_path, 'native/native_functions.yaml')
tags_path = os.path.join(source_path, 'native/tags.yaml')
parsed_yaml = parse_native_yaml(native_yaml_path, tags_path)
native_functions, _ = parsed_yaml.native_functions, parsed_yaml.backend_indices
template_dir = os.path.join(source_path, "templates")
def make_file_manager(install_dir: str) -> FileManager:
return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=False)
cpu_fm = make_file_manager(install_dir)
cpu_fm.write('VmapGeneratedPlumbing.h', lambda: gen_all_vmap_plumbing(native_functions))
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}/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)),
},
)
def main() -> None:
parser = argparse.ArgumentParser(description='functorch codegen')
parser.add_argument(
'-s',
'--source-path',
help='path to source directory for ATen',
default='/scratch/rzou/pt/debug-cpu/aten/src/ATen')
parser.add_argument(
'-o',
'--output-dependencies',
help='output a list of dependencies into the given file and exit')
parser.add_argument(
'--dry-run', action='store_true',
help='run without writing any files (still updates outputs)')
parser.add_argument(
'-d', '--install_dir', help='output directory',
default='functorch/csrc')
options = parser.parse_args()
native_yaml_path = os.path.join(options.source_path, 'native/native_functions.yaml')
parsed_yaml = parse_native_yaml(native_yaml_path)
native_functions, _ = parsed_yaml.native_functions, parsed_yaml.backend_indices
template_dir = os.path.join(options.source_path, "templates")
def make_file_manager(install_dir: str) -> FileManager:
return FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=options.dry_run)
cpu_fm = make_file_manager(options.install_dir)
cpu_fm.write('VmapGeneratedPlumbing.h', lambda: gen_all_vmap_plumbing(native_functions))
if options.output_dependencies:
depfile_path = pathlib.Path(options.output_dependencies).resolve()
depfile_name = depfile_path.name
depfile_stem = depfile_path.stem
for fm, prefix in [
(cpu_fm, ""),
]:
varname = prefix + depfile_stem
path = depfile_path.parent / (prefix + depfile_name)
fm.write_outputs(varname, str(path))
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), " "
),
},
)
def gen_pyi(
native_yaml_path: str,
tags_yaml_path: str,
deprecated_yaml_path: str,
fm: FileManager,
) -> None:
"""gen_pyi()
This function generates a pyi file for torch.
"""
# Some of this logic overlaps with generate_python_signature in
# tools/autograd/gen_python_functions.py; however, this
# function is all about generating mypy type signatures, whereas
# the other function generates are custom format for argument
# checking. If you are update this, consider if your change
# also needs to update the other file.
# Dictionary for NamedTuple definitions
namedtuples: Dict[str, str] = {}
# Generate type signatures for top-level functions
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
unsorted_function_hints: Dict[str,
List[str]] = collections.defaultdict(list)
for n, n1, n2 in [
("csr", "crow", "col"),
("csc", "ccol", "row"),
("bsr", "crow", "col"),
("bsc", "ccol", "row"),
]:
unsorted_function_hints.update({
f"sparse_{n}_tensor": [
f"def sparse_{n}_tensor({n1}_indices: Union[Tensor, List],"
f"{n2}_indices: Union[Tensor, List],"
" values: Union[Tensor, List], size: Optional[_size]=None,"
" *, dtype: Optional[_dtype]=None,"
" device: Union[_device, str, None]=None, requires_grad:_bool=False) -> Tensor: ..."
],
f"_sparse_{n}_tensor_unsafe": [
f"def _sparse_{n}_tensor_unsafe({n1}_indices: Union[Tensor, List],"
f"{n2}_indices: Union[Tensor, List],"
" values: Union[Tensor, List], size: List[int],"
" dtype: Optional[_dtype] = None, device: Optional[_device] = None,"
" requires_grad: bool = False) -> Tensor: ..."
],
})
unsorted_function_hints.update({
"set_flush_denormal":
["def set_flush_denormal(mode: _bool) -> _bool: ..."],
"get_default_dtype": ["def get_default_dtype() -> _dtype: ..."],
"asarray": [
"def asarray(obj: Any, *, dtype: Optional[_dtype]=None, "
"device: Union[_device, str, None]=None, copy: Optional[_bool]=None, "
"requires_grad: _bool=False) -> Tensor: ..."
],
"from_numpy": ["def from_numpy(ndarray) -> Tensor: ..."],
"frombuffer": [
"def frombuffer(buffer: Any, *, dtype: _dtype, count: int=-1, "
"offset: int=0, device: Union[_device, str, None]=None, "
"requires_grad: _bool=False) -> Tensor: ..."
],
"numel": ["def numel(self: Tensor) -> _int: ..."],
"as_tensor": [
"def as_tensor(data: Any, dtype: _dtype=None, device: Optional[_device]=None) -> Tensor: ..."
],
"get_num_threads": ["def get_num_threads() -> _int: ..."],
"set_num_threads": ["def set_num_threads(num: _int) -> None: ..."],
"init_num_threads": ["def init_num_threads() -> None: ..."],
"get_num_interop_threads":
["def get_num_interop_threads() -> _int: ..."],
"set_num_interop_threads":
["def set_num_interop_threads(num: _int) -> None: ..."],
# These functions are explicitly disabled by
# SKIP_PYTHON_BINDINGS because they are hand bound.
# Correspondingly, we must hand-write their signatures.
"tensor":
["def tensor(data: Any, {}) -> Tensor: ...".format(FACTORY_PARAMS)],
"sparse_coo_tensor": [
"def sparse_coo_tensor(indices: Tensor, values: Union[Tensor,List],"
" size: Optional[_size]=None, *, dtype: Optional[_dtype]=None,"
" device: Union[_device, str, None]=None, requires_grad:_bool=False) -> Tensor: ..."
],
"_sparse_coo_tensor_unsafe": [
"def _sparse_coo_tensor_unsafe(indices: Tensor, values: Tensor, size: List[int],"
" dtype: Optional[_dtype] = None, device: Optional[_device] = None,"
" requires_grad: bool = False) -> Tensor: ..."
],
"sparse_compressed_tensor": [
"def sparse_compressed_tensor(compressed_indices: Union[Tensor, List],"
"plain_indices: Union[Tensor, List],"
" values: Union[Tensor, List], size: Optional[_size]=None,"
" *, dtype: Optional[_dtype]=None, layout: Optional[_layout] = None,"
" device: Union[_device, str, None]=None, requires_grad:_bool=False) -> Tensor: ..."
],
"_sparse_compressed_tensor_unsafe": [
"def _sparse_compressed_tensor_unsafe(comp_indices: Union[Tensor, List],"
"plain_indices: Union[Tensor, List],"
" values: Union[Tensor, List], size: List[int],"
" dtype: Optional[_dtype] = None, layout: Optional[_layout] = None,"
" device: Optional[_device] = None,"
" requires_grad: bool = False) -> Tensor: ..."
],
"_is_functional_tensor":
["def _is_functional_tensor(t: Tensor) -> _bool: ..."],
"range": [
"def range(start: Number, end: Number,"
" step: Number=1, *, out: Optional[Tensor]=None, {}) -> Tensor: ..."
.format(FACTORY_PARAMS)
],
"arange": [
"def arange(start: Number, end: Number, step: Number, *,"
" out: Optional[Tensor]=None, {}) -> Tensor: ...".format(
FACTORY_PARAMS),
"def arange(start: Number, end: Number, *, out: Optional[Tensor]=None, {}) -> Tensor: ..."
.format(FACTORY_PARAMS),
"def arange(end: Number, *, out: Optional[Tensor]=None, {}) -> Tensor: ..."
.format(FACTORY_PARAMS),
],
"linspace": [
"def linspace(start: Number, end: Number, steps: Optional[_int]=None, *,"
" out: Optional[Tensor]=None, {}) -> Tensor: ...".format(
FACTORY_PARAMS)
],
"logspace": [
"def logspace(start: Number, end: Number, steps: Optional[_int]=None, base: _float=10.0, *,"
" out: Optional[Tensor]=None, {}) -> Tensor: ...".format(
FACTORY_PARAMS)
],
"randint": [
"def randint(low: _int, high: _int, size: _size, *,"
" generator: Optional[Generator]=None, {}) -> Tensor: ...".format(
FACTORY_PARAMS),
"def randint(high: _int, size: _size, *,"
" generator: Optional[Generator]=None, {}) -> Tensor: ...".format(
FACTORY_PARAMS),
],
"full": [
"def full(size: _size, fill_value: Number, *,"
" out: Optional[Tensor]=None,"
" layout: _layout=strided, {}) -> Tensor: ...".format(
FACTORY_PARAMS),
"def full(size: _size, fill_value: Number, *,"
" names: List[Union[str, None]],"
" layout: _layout=strided, {}) -> Tensor: ...".format(
FACTORY_PARAMS),
],
"is_grad_enabled": ["def is_grad_enabled() -> _bool: ..."],
"is_inference_mode_enabled":
["def is_inference_mode_enabled() -> _bool: ..."],
"nonzero": [
"def nonzero(input: Tensor, *, as_tuple: Literal[False]=False, out: Optional[Tensor]=None) -> Tensor: ...",
"def nonzero(input: Tensor, *, as_tuple: Literal[True]) -> Tuple[Tensor, ...]: ...",
],
"binary_cross_entropy_with_logits": [
"def binary_cross_entropy_with_logits(input: Tensor, target: Tensor, "
"weight: Optional[Tensor] = None, size_average: Optional[bool] = None, "
"reduce: Optional[bool] = None, reduction: str = ..., "
"pos_weight: Optional[Tensor] = None) -> Tensor: ..."
],
"cosine_embedding_loss": [
"def cosine_embedding_loss(input1: Tensor, input2: Tensor, "
"target: Tensor, margin: float = ..., size_average: Optional[bool] = ..., "
"reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ..."
],
"ctc_loss": [
"def ctc_loss(log_probs: Tensor, targets: Tensor, input_lengths: Tensor, target_lengths: Tensor,"
" blank: int = ..., reduction: str = ..., zero_infinity: bool = ...) -> Tensor: ..."
],
"hinge_embedding_loss": [
"def hinge_embedding_loss(input: Tensor, target: Tensor, margin: float = ...,"
" size_average: Optional[bool] = ..., reduce: Optional[bool] = ..., "
"reduction: str = ...) -> Tensor: ..."
],
"kl_div": [
"def kl_div(input: Tensor, target: Tensor, size_average: Optional[bool] = ..., "
"reduce: Optional[bool] = ..., reduction: str = ..., log_target: bool = ...) -> Tensor: ..."
],
"margin_ranking_loss": [
"def margin_ranking_loss(input1: Tensor, input2: Tensor, target: Tensor,"
" margin: float = ..., size_average: Optional[bool] = ..., "
" reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ..."
],
"triplet_margin_loss": [
"def triplet_margin_loss(anchor: Tensor, positive: Tensor, negative: Tensor, "
"margin: float = ..., p: float = ..., eps: float = ..., swap: bool = ..., "
"size_average: Optional[bool] = ..., "
"reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ..."
],
"dsmm": ["def dsmm(input: Tensor, mat2: Tensor) -> Tensor: ..."],
"hsmm": ["def hsmm(input: Tensor, mat2: Tensor) -> Tensor: ..."],
"saddmm": [
"def saddmm(input: Tensor, mat1: Tensor, mat2: Tensor, *, beta: Number=1, "
"alpha: Number=1, out: Optional[Tensor]=None) -> Tensor: ..."
],
"spmm": ["def spmm(input: Tensor, mat2: Tensor) -> Tensor: ..."],
"div": [
"def div(input: Union[Tensor, Number], other: Union[Tensor, Number], *, "
"rounding_mode: Optional[str] = None, out: Optional[Tensor]=None) -> Tensor: ..."
],
})
for binop in ["mul", "true_divide", "floor_divide"]:
unsorted_function_hints[binop].append(
"def {}(input: Union[Tensor, Number],"
" other: Union[Tensor, Number],"
" *, out: Optional[Tensor]=None) -> Tensor: ...".format(binop))
for binop in ["add", "sub"]:
unsorted_function_hints[binop].append(
"def {}(input: Union[Tensor, Number],"
" other: Union[Tensor, Number],"
" *, alpha: Optional[Number]=1, out: Optional[Tensor]=None) -> Tensor: ..."
.format(binop))
native_functions = parse_native_yaml(native_yaml_path,
tags_yaml_path).native_functions
native_functions = list(
filter(should_generate_py_binding, native_functions))
function_signatures = load_signatures(native_functions,
deprecated_yaml_path,
method=False,
pyi=True)
sig_groups = get_py_torch_functions(function_signatures)
for group in sorted(sig_groups, key=lambda g: g.signature.name):
name = group.signature.name
unsorted_function_hints[name] += generate_type_hints(group)
named_tuple = returns_named_tuple_pyi(group.signature)
if named_tuple is not None and not group.signature.deprecated:
# deprecated namedtuples are currently not included for torch functions
tuple_name, tuple_def = named_tuple
if tuple_name in namedtuples:
assert namedtuples[tuple_name] == tuple_def
else:
namedtuples[tuple_name] = tuple_def
function_hints = []
for name, hints in sorted(unsorted_function_hints.items()):
if len(hints) > 1:
hints = ["@overload\n" + h for h in hints]
function_hints += hints
# Generate type signatures for Tensor methods
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
unsorted_tensor_method_hints: Dict[
str, List[str]] = collections.defaultdict(list)
unsorted_tensor_method_hints.update({
"size": [
"def size(self) -> Size: ...",
"def size(self, dim: _int) -> _int: ...",
],
"stride": [
"def stride(self) -> Tuple[_int]: ...",
"def stride(self, _int) -> _int: ...",
],
"new_ones": [
"def new_ones(self, size: _size, {}) -> Tensor: ...".format(
FACTORY_PARAMS)
],
"new_tensor": [
"def new_tensor(self, data: Any, {}) -> Tensor: ...".format(
FACTORY_PARAMS)
],
# new and __init__ have the same signatures differ only in return type
# Adapted from legacy_tensor_ctor and legacy_tensor_new
"new": [
"def new(self, *args: Any, {}) ->Tensor: ...".format(DEVICE_PARAM),
"def new(self, storage: Storage) -> Tensor: ...",
"def new(self, other: Tensor) -> Tensor: ...",
"def new(self, size: _size, *, {}) -> Tensor: ...".format(
DEVICE_PARAM),
],
"__init__": [
"def __init__(self, *args: Any, {}) -> None: ...".format(
DEVICE_PARAM),
"def __init__(self, storage: Storage) -> None: ...",
"def __init__(self, other: Tensor) -> None: ...",
"def __init__(self, size: _size, *, {}) -> None: ...".format(
DEVICE_PARAM),
],
"as_subclass": ["def as_subclass(self, cls: Tensor) -> Tensor: ..."],
"_make_subclass": [
"def _make_subclass(cls, data: Tensor, require_grad: _bool = False, dispatch_strides: _bool=False,"
" dispatch_device: _bool=False) -> Tensor: ..."
],
"__getitem__":
["def __getitem__(self, {}) -> Tensor: ...".format(INDICES)],
"__setitem__": [
"def __setitem__(self, {}, val: Union[Tensor, Number])"
" -> None: ...".format(INDICES)
],
"tolist": ["def tolist(self) -> List: ..."],
"requires_grad_":
["def requires_grad_(self, mode: _bool=True) -> Tensor: ..."],
"element_size": ["def element_size(self) -> _int: ..."],
"data_ptr": ["def data_ptr(self) -> _int: ..."],
"dim": ["def dim(self) -> _int: ..."],
"nonzero": [
"def nonzero(self, *, as_tuple: Literal[False]=False) -> Tensor: ...",
"def nonzero(self, *, as_tuple: Literal[True]) -> Tuple[Tensor, ...]: ...",
],
"numel": ["def numel(self) -> _int: ..."],
"ndimension": ["def ndimension(self) -> _int: ..."],
"nelement": ["def nelement(self) -> _int: ..."],
"cuda": [
"def cuda(self, device: Optional[Union[_device, _int, str]]=None, non_blocking: _bool=False) -> Tensor: ..."
],
"numpy": ["def numpy(self) -> Any: ..."],
"apply_": ["def apply_(self, callable: Callable) -> Tensor: ..."],
"map_":
["def map_(self, tensor: Tensor, callable: Callable) -> Tensor: ..."],
"map2_": [
"def map2_(self, x: Tensor, y: Tensor, callable: Callable) -> Tensor: ..."
],
"storage": ["def _storage(self) -> Storage: ..."],
"storage_type": ["def storage_type(self) -> Storage: ..."],
"type": [
"def type(self, dtype: None=None, non_blocking: _bool=False) -> str: ...",
"def type(self, dtype: Union[str, _dtype], non_blocking: _bool=False) -> Tensor: ...",
],
"get_device": ["def get_device(self) -> _int: ..."],
"contiguous": [
"def contiguous(self, memory_format=torch.contiguous_format) -> Tensor: ..."
],
"has_names": ["def has_names(self) -> _bool: ..."],
"is_contiguous": [
"def is_contiguous(self, memory_format=torch.contiguous_format) -> _bool: ..."
],
"_is_view": ["def _is_view(self) -> _bool: ..."],
"is_cuda": ["is_cuda: _bool"],
"is_leaf": ["is_leaf: _bool"],
"is_nested": ["is_nested: _bool"],
"is_sparse": ["is_sparse: _bool"],
"is_sparse_csr": ["is_sparse_csr: _bool"],
"is_quantized": ["is_quantized: _bool"],
"is_meta": ["is_meta: _bool"],
"is_mps": ["is_mps: _bool"],
"is_ort": ["is_ort: _bool"],
"is_mkldnn": ["is_mkldnn: _bool"],
"is_vulkan": ["is_vulkan: _bool"],
"is_ipu": ["is_ipu: _bool"],
"storage_offset": ["def storage_offset(self) -> _int: ..."],
"to": [
"def to(self, dtype: _dtype, non_blocking: _bool=False, copy: _bool=False) -> Tensor: ...",
"def to(self, device: Optional[Union[_device, str]]=None, dtype: Optional[_dtype]=None, "
"non_blocking: _bool=False, copy: _bool=False) -> Tensor: ...",
"def to(self, other: Tensor, non_blocking: _bool=False, copy: _bool=False) -> Tensor: ...",
],
"item": ["def item(self) -> Number: ..."],
"copy_": [
"def copy_(self, src: Tensor, non_blocking: _bool=False) -> Tensor: ..."
],
"set_": [
"def set_(self, storage: Union[Storage, _TypedStorage], offset: _int, size: _size, stride: _size) -> Tensor: ...",
"def set_(self, storage: Union[Storage, _TypedStorage]) -> Tensor: ...",
],
"split": [
"def split(self, split_size: _int, dim: _int=0) -> Sequence[Tensor]: ...",
"def split(self, split_size: Tuple[_int, ...], dim: _int=0) -> Sequence[Tensor]: ...",
],
"div": [
"def div(self, other: Union[Tensor, Number], *, rounding_mode: Optional[str] = None) -> Tensor: ..."
],
"div_": [
"def div_(self, other: Union[Tensor, Number], *, rounding_mode: Optional[str] = None) -> Tensor: ..."
],
})
for binop in ["mul", "true_divide", "floor_divide"]:
for inplace in [False, True]:
out_suffix = ", *, out: Optional[Tensor]=None"
if inplace:
binop += "_"
out_suffix = ""
unsorted_tensor_method_hints[binop].append(
"def {}(self, other: Union[Tensor, Number]{})"
" -> Tensor: ...".format(binop, out_suffix))
for binop in ["add", "sub"]:
for inplace in [False, True]:
out_suffix = ", out: Optional[Tensor]=None"
if inplace:
binop += "_"
out_suffix = ""
unsorted_tensor_method_hints[binop].append(
"def {}(self, other: Union[Tensor, Number], "
"*, alpha: Optional[Number]=1{})"
" -> Tensor: ...".format(binop, out_suffix))
simple_conversions = [
"byte",
"char",
"cpu",
"double",
"float",
"half",
"int",
"long",
"short",
"bool",
"bfloat16",
]
for name in simple_conversions:
unsorted_tensor_method_hints[name].append(
"def {}(self) -> Tensor: ...".format(name))
# pyi tensor methods don't currently include deprecated signatures for some reason
# TODO: we should probably add them in
tensor_method_signatures = load_signatures(
native_functions,
deprecated_yaml_path,
method=True,
skip_deprecated=True,
pyi=True,
)
tensor_method_sig_groups = get_py_torch_functions(tensor_method_signatures,
method=True)
for group in sorted(tensor_method_sig_groups,
key=lambda g: g.signature.name):
name = group.signature.name
unsorted_tensor_method_hints[name] += generate_type_hints(group)
named_tuple = returns_named_tuple_pyi(group.signature)
if named_tuple is not None and not group.signature.deprecated:
# deprecated namedtuples are currently not included for torch functions
tuple_name, tuple_def = named_tuple
if tuple_name in namedtuples:
assert namedtuples[tuple_name] == tuple_def
else:
namedtuples[tuple_name] = tuple_def
for op in all_ops:
name = "__{}__".format(op)
unsorted_tensor_method_hints[name] += sig_for_ops(name)
tensor_method_hints = []
for name, hints in sorted(unsorted_tensor_method_hints.items()):
if len(hints) > 1:
hints = ["@overload\n" + h for h in hints]
tensor_method_hints += hints
# TODO: Missing type hints for nn
# Generate namedtuple definitions
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
namedtuple_defs = [
"{} = {}".format(name, defn) for name, defn in namedtuples.items()
]
# Generate type signatures for legacy classes
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
legacy_storage_base_hints = ["class StorageBase(object): ..."]
legacy_class_hints = []
for c in (
"DoubleTensor",
"FloatTensor",
"LongTensor",
"IntTensor",
"ShortTensor",
"HalfTensor",
"CharTensor",
"ByteTensor",
"BoolTensor",
):
legacy_class_hints.append("class {}(Tensor): ...".format(c))
# Generate type signatures for dtype classes
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# TODO: don't explicitly list dtypes here; get it from canonical
# source
dtype_class_hints = [
"{}: dtype = ...".format(n) for n in [
"float32",
"float",
"float64",
"double",
"float16",
"bfloat16",
"half",
"uint8",
"int8",
"int16",
"short",
"int32",
"int",
"int64",
"long",
"complex32",
"complex64",
"cfloat",
"complex128",
"cdouble",
"quint8",
"qint8",
"qint32",
"bool",
"quint4x2",
"quint2x4",
]
]
# Generate __all__ directive
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Include only the functions that contain hints, to prevent undefined
# symbols to be included in the `__all__` directive.
hinted_function_names = [
name for name, hint in unsorted_function_hints.items() if hint
]
all_symbols = sorted(list(namedtuples.keys()) + hinted_function_names)
all_directive = pformat(all_symbols, width=100, compact=True).split("\n")
all_directive[0] = "__all__ = {}".format(all_directive[0])
# Write out the stub
# ~~~~~~~~~~~~~~~~~~
env = {
"namedtuple_defs": namedtuple_defs,
"function_hints": function_hints,
"tensor_method_hints": tensor_method_hints,
"legacy_class_hints": legacy_class_hints,
"legacy_storage_base_hints": legacy_storage_base_hints,
"dtype_class_hints": dtype_class_hints,
"all_directive": all_directive,
}
fm.write_with_template(
"torch/_C/__init__.pyi",
"torch/_C/__init__.pyi.in",
lambda: {
"generated_comment": "@" +
"generated from torch/_C/__init__.pyi.in",
**env,
},
)
fm.write_with_template(
"torch/_C/_VariableFunctions.pyi",
"torch/_C/_VariableFunctions.pyi.in",
lambda: {
"generated_comment": "@" +
"generated from torch/_C/_VariableFunctions.pyi.in",
**env,
},
)
fm.write_with_template(
"torch/_VF.pyi",
"torch/_C/_VariableFunctions.pyi.in",
lambda: {
"generated_comment": "@" +
"generated from torch/_C/_VariableFunctions.pyi.in",
**env,
},
)
fm.write_with_template(
"torch/return_types.pyi",
"torch/_C/return_types.pyi.in",
lambda: {
"generated_comment": "@" +
"generated from torch/_C/return_types.pyi",
**env,
},
)
gen_nn_functional(fm)
def gen(
out: str,
native_yaml_path: str,
tags_yaml_path: str,
deprecated_yaml_path: str,
template_path: str,
) -> None:
fm = FileManager(install_dir=out,
template_dir=template_path,
dry_run=False)
native_functions = parse_native_yaml(native_yaml_path,
tags_yaml_path).native_functions
native_functions = list(
filter(should_generate_py_binding, native_functions))
methods = load_signatures(native_functions,
deprecated_yaml_path,
method=True)
create_python_bindings(
fm,
methods,
is_py_variable_method,
None,
"python_variable_methods.cpp",
method=True,
)
# NOTE: num_shards here must be synced with gatherTorchFunctions in
# torch/csrc/autograd/python_torch_functions_manual.cpp
functions = load_signatures(native_functions,
deprecated_yaml_path,
method=False)
create_python_bindings_sharded(
fm,
functions,
is_py_torch_function,
"torch",
"python_torch_functions.cpp",
method=False,
num_shards=3,
)
create_python_bindings(
fm,
functions,
is_py_nn_function,
"torch.nn",
"python_nn_functions.cpp",
method=False,
)
create_python_bindings(
fm,
functions,
is_py_fft_function,
"torch.fft",
"python_fft_functions.cpp",
method=False,
)
create_python_bindings(
fm,
functions,
is_py_linalg_function,
"torch.linalg",
"python_linalg_functions.cpp",
method=False,
)
create_python_bindings(
fm,
functions,
is_py_sparse_function,
"torch.sparse",
"python_sparse_functions.cpp",
method=False,
)
create_python_bindings(
fm,
functions,
is_py_special_function,
"torch.special",
"python_special_functions.cpp",
method=False,
)
# Currently, we only use `functions` to generate `return_types` bindings.
# All methods which return namedtuple have function variant at this point.
# If any method only operator with namedtuple is added in the future,
# we will have to address that.
create_python_return_type_bindings(fm, functions, lambda fn: True,
"python_return_types.cpp")
valid_tags = parse_tags_yaml(tags_yaml_path)
def gen_tags_enum() -> Dict[str, str]:
return {
"enum_of_valid_tags": ("".join(
[f'\n.value("{tag}", at::Tag::{tag})' for tag in valid_tags]))
}
fm.write("python_enum_tag.cpp", gen_tags_enum)
def gen_external(native_functions_path, tags_path, external_path):
native_functions = parse_native_yaml(native_functions_path, tags_path)
func_decls = []
func_registrations = []
for func in native_functions:
schema = func.func
name = schema.name.name.base
args = schema.arguments
# Only supports extern calls for functions with out variants
if not schema.is_out_fn():
continue
# Doesn't currently support functions with more than one out parameter
if len(args.out) > 1:
continue
# Doesn't currently support kwarg arguments
if len(args.pre_tensor_options_kwarg_only) > 0 or len(
args.post_tensor_options_kwarg_only) > 0:
continue
self_arg = [args.self_arg.argument
] if args.self_arg is not None else []
args = list(args.pre_self_positional) + self_arg + list(
args.post_self_positional)
tensor_args = [
arg for arg in args if isinstance(arg.type, model.BaseType)
and arg.type.name == model.BaseTy.Tensor
]
if len(tensor_args) != len(args):
continue
arg_names = [None] * len(args)
tensor_decls = []
for idx, arg in enumerate(tensor_args):
s = f"const at::Tensor& {arg.name} = tensors[{idx + 1}];"
tensor_decls.append(s)
arg_names[idx] = arg.name
nl = '\n'
# print(tensor_decls, name, arg_names)
func_decl = f"""\
void nnc_aten_{name}(
int64_t bufs_num,
void** buf_data,
int64_t* buf_ranks,
int64_t* buf_dims,
int64_t* buf_strides,
int8_t* buf_dtypes,
int64_t args_num,
int64_t* extra_args) {{
std::vector<at::Tensor> tensors =
constructTensors(bufs_num, buf_data, buf_ranks, buf_dims, buf_strides, buf_dtypes);
at::Tensor& r = tensors[0];
{nl.join(tensor_decls)}
try {{
at::{name}_out({', '.join(['r'] + arg_names)});
}} catch (...) {{
}}
}}"""
func_registration = f"""\
const static RegisterNNCExternalFunction nnc_{name}(
"nnc_aten_{name}",
nnc_aten_{name});"""
func_decls.append(func_decl)
func_registrations.append(func_registration)
fm = FileManager(install_dir='.', template_dir='.', dry_run=False)
fm.write_with_template(
'external_functions_codegen.cpp', external_path, lambda: {
'external_registrations': func_registrations,
'external_functions': func_decls
})
def load_derivatives(
derivatives_yaml_path: str, native_yaml_path: str, tags_yaml_path: str
) -> Tuple[Dict[FunctionSchema, Dict[str, DifferentiabilityInfo]], Set[str]]:
# Do some caching as this is a deterministic function
global _GLOBAL_LOAD_DERIVATIVE_CACHE
key = (derivatives_yaml_path, native_yaml_path)
if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE:
with open(derivatives_yaml_path, "r") as f:
definitions = yaml.load(f, Loader=YamlLoader)
funcs = parse_native_yaml(native_yaml_path, tags_yaml_path).native_functions
# From the parsed native functions, separate out the (generated) view_copy functions,
# so we can generate derivatives for them separately.
native_functions_with_view_groups = get_grouped_by_view_native_functions(funcs)
native_functions_without_view_copies = concatMap(
# We need to pull out the view_inplace ops too, since they might have their own derivative entries.
lambda g: [g]
if isinstance(g, NativeFunction)
else list(g.functions(include_copy=False)),
native_functions_with_view_groups,
)
view_groups = [
g
for g in native_functions_with_view_groups
if isinstance(g, NativeFunctionsViewGroup)
]
# What's the difference between function schema v.s. signature?
# function schema is the complete declaration including mutability annotation / default value and etc.
# signature is the canonical schema for a group of functions (in-place/out/functional variants)
# that are semantically related.
functions_by_signature: Dict[
FunctionSchema, List[NativeFunction]
] = defaultdict(list)
functions_by_schema: Dict[str, NativeFunction] = dict()
for function in native_functions_without_view_copies:
functions_by_signature[function.func.signature()].append(function)
assert str(function.func) not in functions_by_schema
functions_by_schema[str(function.func)] = function
# Keep track of how many of which ops we've seen so we can
# disambiguate them with a numeric suffix.
op_counter = Counter[str]()
# infos is a dict that maps FunctionSchema -> a dict of per dispatch key DifferentiabilityInfos
# this is useful because in tools/autograd/gen_autograd.py:match_differentiability_info
# we ultimately need to categorize the DifferentiabilityInfos by FunctionSchema
infos: Dict[FunctionSchema, Dict[str, DifferentiabilityInfo]] = dict()
used_dispatch_keys: Set[str] = set()
for defn_dict in definitions:
# Ensure that the old derivatives.yaml schema with no dispatch key can be loaded.
if "dispatch" not in defn_dict:
specification = defn_dict.pop("name")
output_differentiability = defn_dict.pop(
"output_differentiability", None
)
defn_dict = {"name": specification, "dispatch": {"Default": defn_dict}}
if output_differentiability:
defn_dict["output_differentiability"] = output_differentiability
name, per_dispatch_diffinfos = create_differentiability_info(
defn_dict,
functions_by_signature,
functions_by_schema,
op_counter,
used_dispatch_keys,
)
infos[name] = per_dispatch_diffinfos
add_view_copy_derivatives(infos, view_groups)
# cache both loaded infos as well a a set of all the dispatch_keys/aliases
# that appear in derivatives.yaml. used_dispatch_keys is useful for generating
# VariableType.cpp where we need a TORCH_LIBRARY_IMPL for every autograd dispatch key used
_GLOBAL_LOAD_DERIVATIVE_CACHE[key] = infos, used_dispatch_keys
return _GLOBAL_LOAD_DERIVATIVE_CACHE[key]
def run_gen_lazy_tensor(
aten_path: str,
source_yaml: str,
output_dir: str,
dry_run: bool,
impl_path: Optional[str],
node_base: str = default_args.node_base,
node_base_hdr: Optional[str] = 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,
# 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,
# 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: Union[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
full_codegen = parse_full_codegen_ops(source_yaml, grouped_native_functions)
def concat_map_codegen(
func: Callable[[NativeFunction], Sequence[str]],
xs: Iterable[Union[NativeFunctionsGroup, NativeFunction]],
*,
codegenInplaceVariant: bool = False,
) -> 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.
Note: If xs is not sorted, there may be an edge case when generating IR classes. Considering relu and relu_, if
we encounter relu_ before relu. we will then generate an IR class with op = at::aten::relu_ for both relu and
relu_ which will cause problems for relu.
TODO(alanwaketan): Once all ops are grouped properly, we should no longer need this hack.
"""
generated = set()
def gen_key(func: FunctionSchema) -> Tuple[str, str]:
# we want to generate unique entries for overloads of functional variants,
# but not for inplace variants unless explicitly told `codegenInplaceVariant`
return (func.name.name.base, func.name.overload_name)
for x in xs:
f = x.functional if isinstance(x, NativeFunctionsGroup) else x
# For the 'or'd terms:
# 1. codegenInplaceVariant means we can generate the in-place variant corresponding items.
# 2. not f.func.name.name.inplace means the op is not a in-place variant, so we can generate the item.
# 3. f.func.name.name.base not in generated means even for in-place ops we still need to generate the item
# as if they were the functional variants for one time.
if f.func.name in full_codegen and (
codegenInplaceVariant
or not f.func.name.name.inplace
or gen_key(f.func) not in generated
):
generated.add(gen_key(f.func))
for r in func(f):
yield r
selector = SelectiveBuilder.get_nop_selector()
assert backend_key is not 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} becuase 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,
codegenInplaceVariant=True,
)
)
validate_shape_inference_header(shape_inference_hdr, expected_shape_infr_decls)
assert class_name is not 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,
cpp_namespace,
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/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 []
)
],
"native_functions_include": "",
"namespace_prologue": ns_helper.prologue,
"namespace_epilogue": ns_helper.epilogue,
"native_function_definitions": list(
concat_map_codegen(
dest.GenLazyNativeFuncDefinition(
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,
codegenInplaceVariant=True,
)
),
},
)
# Generate IR node classes
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",
"c10/util/Optional.h",
"torch/csrc/lazy/core/hash.h",
"torch/csrc/lazy/core/ir.h",
"torch/csrc/lazy/core/shape.h",
"vector",
]
],
"lazy_ir_inc": [
f'#include "{path}"'
for path in [node_base_hdr if node_base_hdr is not None else None]
if path is not None
],
"ir_declarations": list(
concat_map_codegen(
lazy_ir_generator(backend_indices[backend_key], node_base),
grouped_native_functions,
)
),
"namespace_prologue": ns_helper.prologue,
"namespace_epilogue": ns_helper.epilogue,
},
)
def main() -> None:
parser = argparse.ArgumentParser(description="Generate ATen source files")
parser.add_argument(
"-s",
"--source-path",
help="path to source directory for ATen",
default="caffe2/aten/src/ATen",
)
parser.add_argument(
"-p",
"--generated-ops-cpp-path",
help="path to directory to generate op dispatcher .cpp file",
default="caffe2/torch/csrc/jit/runtime/static/generated_ops.cpp",
)
parser.add_argument(
"-t",
"--generated-ops-test-cpp-path",
help="path to directory to generate op dispatcher .cpp file",
default="caffe2/benchmarks/static_runtime/test_generated_ops.cc",
)
options = parser.parse_args()
native_yaml_path = os.path.join(options.source_path,
"native/native_functions.yaml")
tags_yaml_path = os.path.join(options.source_path, "native/tags.yaml")
parsed_yaml = gen.parse_native_yaml(native_yaml_path, tags_yaml_path)
native_functions, backend_indices = (
parsed_yaml.native_functions,
parsed_yaml.backend_indices,
)
op_generator = generator.GenOpDispatcher()
test_case_generator = generator.GenOpTestCase()
native_functions_groups = [
g for g in gen.get_grouped_native_functions(native_functions)
if isinstance(g, NativeFunctionsGroup)
]
supported_functions_groups = group_functions_by_op_name(
native_functions_groups)
out_variant_op_result = [
op_generator.out_variant(groups, backend_indices[DispatchKey.CPU])
for groups in supported_functions_groups
]
out_variant_test_result = [
test_case_generator.out_variant(groups)
for groups in supported_functions_groups
]
native_functions_view_groups = [
g for g in gen.get_grouped_by_view_native_functions(native_functions)
if isinstance(g, NativeFunctionsViewGroup)
]
supported_functions_view_groups = group_functions_by_op_name(
native_functions_view_groups)
view_op_result = [
op_generator.view(groups, backend_indices[DispatchKey.CPU])
for groups in supported_functions_view_groups
]
view_test_result = [
test_case_generator.view(groups)
for groups in supported_functions_view_groups
]
op_result = out_variant_op_result + ["\n\n"] + view_op_result
test_result = out_variant_test_result + ["\n\n"] + view_test_result
write_cpp(op_result, options.generated_ops_cpp_path)
write_test_cpp(test_result, options.generated_ops_test_cpp_path)
print("\ntotal grouped native ops: %d" %
len(gen.get_grouped_native_functions(native_functions)))
print("grouped native ops with out variant: %d" %
len(native_functions_groups))
supported_functions_num = sum(
[len(groups) for groups in supported_functions_groups])
print("generated functions groups with out variant: %d" %
supported_functions_num)
print("\nview grouped native ops: %d" % len(native_functions_view_groups))
supported_view_functions_num = sum(
[len(groups) for groups in supported_functions_view_groups])
print("generated functions view groups: %d" % supported_view_functions_num)
print("\noverall generated : %d" %
(supported_functions_num + supported_view_functions_num))
def run(source_yaml: str,
output_dir: str,
dry_run: bool,
impl_path: Optional[str] = None) -> None:
# Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py
pytorch_root = pathlib.Path(__file__).parent.parent.absolute()
template_dir = os.path.join(pytorch_root, "aten/src/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(
pytorch_root, "aten/src/ATen/native/native_functions.yaml")
tags_yaml_path = os.path.join(pytorch_root,
"aten/src/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()
assert class_name is not 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,
)
def main(args: List[str]) -> None:
parser = argparse.ArgumentParser(
description="Generate unboxing source files")
parser.add_argument(
"-s",
"--source-path",
help="path to source directory for ATen",
default="aten/src/ATen",
)
parser.add_argument("-d",
"--install_dir",
help="output directory",
default="build/aten/src/ATen")
parser.add_argument(
"-o",
"--output-dependencies",
help="output a list of dependencies into the given file and exit",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="run without writing any files (still updates outputs)",
)
parser.add_argument(
"--op_selection_yaml_path",
help="Provide a path to the operator selection (for custom build) YAML "
"that contains the information about the set of selected operators "
"and their categories (training, ...). Each operator is either a "
"full operator name with overload or just a bare operator name. "
"The operator names also contain the namespace prefix (e.g. aten::)",
)
parser.add_argument(
"--op_registration_allowlist",
nargs="*",
help="filter op registrations by the allowlist (if set); "
"each item is `namespace`::`operator name` without overload name; "
"e.g.: aten::empty aten::conv2d ...",
)
parser.add_argument(
"--TEST_ONLY_op_registration_allowlist_yaml_path",
help="Provide a path to the operator selection (for custom build) YAML "
"which contains a list of operators. It is to serve testing purpose and "
"each item is `namespace`::`operator name` without overload name; "
"e.g.: aten::empty aten::conv2d ...",
)
options = parser.parse_args(args)
if options.op_registration_allowlist:
op_registration_allowlist = options.op_registration_allowlist
elif options.TEST_ONLY_op_registration_allowlist_yaml_path:
with open(options.TEST_ONLY_op_registration_allowlist_yaml_path,
"r") as f:
op_registration_allowlist = yaml.safe_load(f)
else:
op_registration_allowlist = None
selector = get_custom_build_selector(
options.op_registration_allowlist,
options.op_selection_yaml_path,
)
native_yaml_path = os.path.join(options.source_path,
"native/native_functions.yaml")
tags_yaml_path = os.path.join(options.source_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,
)
cpu_fm = make_file_manager(options=options)
gen_unboxing(native_functions=native_functions,
cpu_fm=cpu_fm,
selector=selector)
if options.output_dependencies:
depfile_path = pathlib.Path(options.output_dependencies).resolve()
depfile_name = depfile_path.name
depfile_stem = depfile_path.stem
path = depfile_path.parent / depfile_name
cpu_fm.write_outputs(depfile_stem, str(path))
def generate_native_functions(self):
logging.info("Generating Native Functions Yaml")
native_path = TORCHGEN_DIR.joinpath("packaged", "ATen", "native")
native_yaml_path = native_path.joinpath("native_functions.yaml")
tags_yaml_path = native_path.joinpath("tags.yaml")
ts_native_yaml_path = TORCH_DIR.joinpath(
"aten", "src", "ATen", "native", "ts_native_functions.yaml"
)
ts_native_yaml = None
if ts_native_yaml_path.exists():
ts_native_yaml = yaml.load(ts_native_yaml_path.read_text(), yaml.CLoader)
parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
self.native_functions = parsed_yaml.native_functions
self.backend_indices = parsed_yaml.backend_indices
self.grouped_native_functions = get_grouped_native_functions(
self.native_functions
)
def get_native_function_name(f):
func = f if hasattr(f, "func") else f.functional
return str(func.func.name)
self.native_functions = {
get_native_function_name(f): f for f in self.native_functions
}
def get_opnames(ops):
opnames = defaultdict(set)
for op in ops:
opname = op.split(".")[0]
opnames[opname].add(op)
return opnames
aten_funcs = get_opnames(
map(get_native_function_name, self.grouped_native_functions)
)
with self.config_path.open() as f:
config = yaml.load(f, yaml.CLoader)
# List of unsupported ops in LTC autogen because of some error
blacklist = set(config.get("blacklist", []))
# List of supported ops that we don't want to do the full codegen for
# primarily view ops
supported = set(config.get("supported", []))
# List of non-native ops to do IR codegen for
non_native = config.get("non_native", [])
# use ripgrep if available as its much faster
if which("rg") is not None:
cmd = ["rg", "-o", "-N", r"aten::[0-9a-zA-Z_\.]+"]
else:
cmd = ["grep", "-o", r"aten::[0-9a-zA-Z_\.]\+"]
torch_ops = set(
op[6:]
for op in subprocess.check_output(
cmd + [str(self.torch_ops_file)],
encoding="utf-8",
)
.strip()
.split(os.linesep)
)
torch_opnames = get_opnames(torch_ops)
# process ops list
ops = set()
composite_implicit = set()
for op in torch_ops:
if op not in self.native_functions:
continue
func = self.native_functions[op]
base = func.func.name.name.base
if base in blacklist or op in blacklist:
continue
if base in supported or op in supported:
continue
# Blacklist new_/_like ops since they are non-differentiable.
if any(o.startswith("new_") or o.endswith("_like") for o in (base, op)):
continue
if func.has_composite_implicit_autograd_kernel:
composite_implicit.add(op)
elif func.func.name.name.inplace:
for autogen in func.autogen:
if "functional" in autogen.overload_name:
ops.add(str(autogen))
else:
ops.add(op)
skipped = set(torch_ops) - ops - supported - composite_implicit
# List of ops autogen even if not explicitly supported by Torch-MLIR explicitly
ops |= set(config.get("whitelist", []))
# Additional ops to support that are not supported by Torch-MLIR explicitly
supported |= set(config.get("additional_ops", []))
self.ops = sorted(ops)
with self.source_yaml.open("w") as f:
source_yaml = {
"backend": "Lazy",
"cpp_namespace": "torch::lazy",
"full_codegen": self.ops,
"supported": sorted(supported),
"non_native": non_native,
}
yaml.dump(source_yaml, f, default_flow_style=False)
f.write(
dedent(
"""
# Composite implicit ops (supported by Torch-MLIR but not differentiable)
{composite_implicit}
# Skipped ops (supported by Torch-MLIR but no equivalent native function)
{skipped}
"""
).format(
composite_implicit=os.linesep.join(
f"# - {op}" for op in sorted(composite_implicit)
),
skipped=os.linesep.join(f"# - {op}" for op in sorted(skipped)),
)
)
if ts_native_yaml:
ts_full_codegen = set(ts_native_yaml["full_codegen"])
mlir_full_codegen = set(self.ops)
if ts_full_codegen - mlir_full_codegen:
logging.debug(
"Full Codegen ops supported by the TorchScript backend "
"but not by the Torch-MLIR backend:\n {}".format(
"\n ".join(sorted(ts_full_codegen - mlir_full_codegen))
)
)
if mlir_full_codegen - ts_full_codegen:
logging.debug(
"Full Codegen ops supported by the Torch-MLIR backend "
"but not by the TorchScript backend:\n {}".format(
"\n ".join(sorted(mlir_full_codegen - ts_full_codegen))
)
)
