From a62d60829c13076b5ff04d7d6402b69a886ed38f Mon Sep 17 00:00:00 2001 From: "Jae Hoon (Antonio) Kim" <17433012+antoniojkim@users.noreply.github.com> Date: Fri, 10 Jun 2022 10:25:10 -0400 Subject: [PATCH] Refactor autogen (#925) --- .gitignore | 2 - build_tools/autogen_ltc_backend.py | 665 ++++++++++-------- python/torch_mlir/csrc/CMakeLists.txt | 4 +- .../base_lazy_backend/LazyShapeInference.h | 71 -- .../mlir_native_functions.cpp | 2 +- ...ShapeInference.cpp => shape_inference.cpp} | 2 +- 6 files changed, 369 insertions(+), 377 deletions(-) delete mode 100644 python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h rename python/torch_mlir/csrc/base_lazy_backend/{LazyShapeInference.cpp => shape_inference.cpp} (98%) diff --git a/.gitignore b/.gitignore index 74ad3b2e9..3691deec1 100644 --- a/.gitignore +++ b/.gitignore @@ -25,8 +25,6 @@ __pycache__ bazel-* # Autogenerated files -/generated_native_functions.yaml -/generated_backend.hash /python/torch_mlir/csrc/base_lazy_backend/generated # Example backend diff --git a/build_tools/autogen_ltc_backend.py b/build_tools/autogen_ltc_backend.py index e1157dd48..4ee4e354e 100644 --- a/build_tools/autogen_ltc_backend.py +++ b/build_tools/autogen_ltc_backend.py @@ -2,153 +2,40 @@ import argparse import hashlib import importlib import os +import re import subprocess -import sys import warnings from collections import defaultdict from dataclasses import dataclass from pathlib import Path from shutil import which -from textwrap import dedent - -import yaml +from textwrap import dedent, indent # PyTorch's LTC backend autogen script import torchgen import torchgen.dest.lazy_ir import torchgen.gen_lazy_tensor -from torchgen.api.lazy import LazyIrSchema +import yaml +from torchgen.api.lazy import LazyIrSchema, setValueT +from torchgen.api.types import BaseCppType +from torchgen.dest import GenLazyShapeInferenceDefinition from torchgen.gen import get_grouped_native_functions, parse_native_yaml -from torchgen.model import NativeFunctionsGroup +from torchgen.gen_backend_stubs import parse_backend_yaml -TORCH_DIR = Path(importlib.util.find_spec('torch').origin).resolve().parent.parent +TORCH_DIR = Path(importlib.util.find_spec("torch").origin).resolve().parent.parent TORCH_MLIR_DIR = Path(__file__).resolve().parent.parent -def isOptionalCType(arg): - return str(type(arg)) == "" - -def generate_native_functions( - config_path: Path, torch_ops_file: Path, out_file: Path -): - print("Generating Native Functions Yaml") - - native_path = TORCH_DIR.joinpath("aten", "src", "ATen", "native") - native_yaml_path = native_path.joinpath("native_functions.yaml") - tags_yaml_path = native_path.joinpath("tags.yaml") - - parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path) - native_functions = parsed_yaml.native_functions - grouped_native_functions = get_grouped_native_functions(native_functions) - - def get_native_function_name(f): - func = f if hasattr(f, "func") else f.functional - return str(func.func.name), func - - def get_opnames(ops): - opnames = defaultdict(set) - for op in ops: - opname = op.split(".")[0] - opnames[opname].add(op) - return opnames - - native_functions = dict(map(get_native_function_name, native_functions)) - grouped_native_functions = dict(map(get_native_function_name, grouped_native_functions)) - aten_funcs = get_opnames(set(grouped_native_functions.keys())) - - with 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", []) - - if which("rg") is not None: # use ripgrep if available as its much faster - 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(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 native_functions: - continue - - func = 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 - - if func.has_composite_implicit_autograd_kernel and f"{op}_backward" not in torch_ops: - 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 - - # Additional ops to support that are not supported by Torch-MLIR explicitly - supported |= set(config.get("additional_ops", [])) - - with out_file.open("w") as f: - yaml.dump( - { - "backend": "Lazy", - "cpp_namespace": "torch::lazy", - "full_codegen": sorted(ops), - "supported": sorted(supported), - "non_native": non_native, - }, - 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)), - ) - ) - - return parsed_yaml, grouped_native_functions +def reindent(text, prefix=""): + return indent(dedent(text), prefix) @dataclass(frozen=True) class GenMlirLazyIr(torchgen.dest.GenLazyIR): + def isOptionalCType(self, arg): + return str(type(arg)) == "" - def lowering_function(self, schema): + def lowering_function(self, schema: LazyIrSchema): signature = "TorchMlirOpVector Lower(TorchMlirFunction function, TorchMlirLoweringContext* loctx) const override" if schema.properties.LowerDeclOnly: @@ -159,214 +46,392 @@ class GenMlirLazyIr(torchgen.dest.GenLazyIR): emplace_arguments = [] for arg in schema.positional_args: if arg.is_lazy_value: - if isOptionalCType(arg.lazy_type): - emplace_arguments.append(f"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr") - continue - emplace_arguments.append('loctx->GetOutputOp(operand(i++))') - continue - emplace_arguments.append(f'"{arg.name}", {arg.name}') + if self.isOptionalCType(arg.lazy_type): + emplace_arguments.append( + f"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr" + ) + else: + emplace_arguments.append("loctx->GetOutputOp(operand(i++))") + else: + 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'"{t.name}", loctx->GetOutputOp(operand(i++))' for t in schema.keyword_values] - emplace_kwarg_scalars = [f'"{t.name}", {t.name}' for t in schema.keyword_scalars] - emplace_kwarguments = "\n ".join( - [f"kwarguments.emplace_back({a});" for a in emplace_kwarg_values + emplace_kwarg_scalars]) - - return f""" - {signature} {{ - PRINT_FUNCTION(); - std::vector arguments; - std::vector 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::TorchMlirOpVector {schema.aten_name}_out = torch::lazy::LowerTorchMlirBuiltin(function, op().op, shapes(), arguments, kwarguments); - CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)}); - - return {schema.aten_name}_out; - }} - """.strip() - - -def generate_backend( - source_yaml: Path, - backend_path: Path, - parsed_yaml: dict, - grouped_native_functions: list, -): - print("Running Lazy Tensor Autogen") - - # No fallback code allowed - def gen_fallback_code(*args, **kwargs): - return "" - - torchgen.dest.lazy_ir.gen_fallback_code = gen_fallback_code - - torchgen.gen_lazy_tensor.run_gen_lazy_tensor( - backend_name="TorchMlir", - aten_path=str(TORCH_DIR.joinpath("aten", "src", "ATen")), - source_yaml=str(source_yaml), - output_dir=str(backend_path.joinpath("generated")), - dry_run=False, - impl_path=str(backend_path.joinpath("mlir_native_functions.cpp")), - node_base="torch::lazy::TorchMlirNode", - node_base_hdr=str(backend_path.joinpath("mlir_node.h")), - tensor_class="torch::lazy::LazyTensor", - tensor_class_hdr="torch/csrc/lazy/core/tensor.h", - shape_inference_hdr=str(backend_path.joinpath("LazyShapeInference.h")), - lazy_ir_generator=GenMlirLazyIr, - ) - - # Remove lazy_tensor_core imports - subprocess.check_call( - [ - "sed", - "-i", - "/lazy_tensor_core/d", - str(backend_path.joinpath("generated", "LazyNativeFunctions.cpp")), + f"arguments.emplace_back({a});" for a in emplace_arguments + ) + emplace_kwarg_values = [ + f'"{t.name}", loctx->GetOutputOp(operand(i++))' + for t in schema.keyword_values ] - ) + emplace_kwarg_scalars = [ + f'"{t.name}", {t.name}' for t in schema.keyword_scalars + ] + emplace_kwarguments = "\n ".join( + f"kwarguments.emplace_back({a});" + for a in emplace_kwarg_values + emplace_kwarg_scalars + ) - # programmatically check shape inference declarations - import re + return reindent( + f""" + {signature} {{ + PRINT_FUNCTION(); + std::vector arguments; + std::vector 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::TorchMlirOpVector {schema.aten_name}_out = torch::lazy::LowerTorchMlirBuiltin(function, op().op, shapes(), arguments, kwarguments); + CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)}); - sig_re = re.compile( - r"std::vector\s+(?P\w+)\((?P[^\)]+)\)" - ) - global_signatures = {} + return {schema.aten_name}_out; + }} + """, + " ", + ) - def extract_signatures(path): - signatures = set() - for name, args in sig_re.findall(path.read_text()): - signature = re.sub(r"\s+", "", f"{name}({args})") - global_signatures[signature] = (name, args) - signatures.add(signature) - return signatures - upstream_shape_inference_decls = extract_signatures( - TORCH_DIR.joinpath("torch", "csrc", "lazy", "core", "shape_inference.h") - ) - assert len(upstream_shape_inference_decls) > 0 - shape_inference_decls = extract_signatures( - backend_path.joinpath("LazyShapeInference.h") - ) - assert len(shape_inference_decls) > 0 - shape_inference_defs = extract_signatures( - backend_path.joinpath("LazyShapeInference.cpp") - ) - assert len(shape_inference_decls) > len(shape_inference_defs) +class GenTorchMlirLTC: + def __init__(self): + self.script_path = Path(__file__).resolve() + self.config_path = ( + Path(__file__).resolve().parent.joinpath("autogen_ltc_backend.yaml") + ) + self.torch_ops_file = TORCH_MLIR_DIR.joinpath( + "include", + "torch-mlir", + "Dialect", + "Torch", + "IR", + "GeneratedTorchOps.td", + ) + assert self.torch_ops_file.exists() + self.source_yaml = TORCH_MLIR_DIR.joinpath( + "build", "generated_native_functions.yaml" + ) + self.backend_path = TORCH_MLIR_DIR.joinpath( + "python", "torch_mlir", "csrc", "base_lazy_backend" + ) + assert self.backend_path.is_dir() + self.backend_path.joinpath("generated").mkdir(exist_ok=True) - missing_defs = ( - shape_inference_decls - - upstream_shape_inference_decls - - shape_inference_defs - ) - if missing_defs: - backend_path.joinpath("generated", "GenLazyShapeInference.cpp").write_text( + self.torchgen_path = Path(torchgen.__path__[0]).resolve() + assert self.torchgen_path.is_dir() + + self.tensor_class = "torch::lazy::LazyTensor" + + # Set the lazy value class + setValueT(BaseCppType("torch::lazy", "Value")) + + def calculate_hash(self): + m = hashlib.sha256() + + # Add file contents to hash + for path in ( + self.script_path, + self.config_path, + self.torch_ops_file, + self.source_yaml, + self.backend_path.joinpath("shape_inference.cpp"), + self.torchgen_path.joinpath("dest", "lazy_ir.py"), + self.torchgen_path.joinpath("api", "lazy.py"), + self.torchgen_path.joinpath("model.py"), + ): + if path.exists(): + m.update(path.read_bytes()) + + return m.hexdigest().strip() + + def generate_native_functions(self): + print("Generating Native Functions Yaml") + + native_path = self.torchgen_path.joinpath("packaged", "ATen", "native") + native_yaml_path = native_path.joinpath("native_functions.yaml") + tags_yaml_path = native_path.joinpath("tags.yaml") + + 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 + + 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 + + # 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)), + ) + ) + + def generate_shape_inference(self): + parsed_backend_yaml = parse_backend_yaml( + self.source_yaml, + self.grouped_native_functions, + self.backend_indices, + ) + backend_index = self.backend_indices[parsed_backend_yaml.backend_key] + + shape_gen = GenLazyShapeInferenceDefinition(backend_index, self.tensor_class) + + sig_re = re.compile( + r"std::vector\s+(?P\w+)\((?P[^\)]+)\)" + ) + global_signatures = {} + + def extract_signatures(text): + signatures = set() + for name, args in sig_re.findall(text): + signature = re.sub(r"\s+", "", f"{name}({args})") + global_signatures[signature] = (name, args) + signatures.add(signature) + return signatures + + shape_inference_decls = [] + for op in self.ops: + f = self.native_functions[op] + shape_sig = shape_gen(f) + shape_inference_decls.extend(shape_sig) + + self.backend_path.joinpath("generated", "shape_inference.h").write_text( dedent( """ - // This file contains autogenerated Lazy Shape Inference placeholders + // This file contains autogenerated Lazy Shape Inference declarations // for ops that dont have a corresponding structured kernel or shape definition - #include "../LazyShapeInference.h" - #include "../../utils/exception.h" + #include + #include + #include + #include + #include + #include + #include + namespace torch {{ namespace lazy {{ + {} + }} // namespace lazy }} // namespace torch """ - ).format( - "".join( - dedent( - f""" - std::vector {name}({args}) {{ - UNIMPLEMENTED_FUNCTION_ERROR(); - }} - """ - ) - for name, args in map( - global_signatures.get, sorted(missing_defs) + ).format(os.linesep.join(sorted(shape_inference_decls))) + ) + + shape_inference_decls = extract_signatures( + self.backend_path.joinpath("generated", "shape_inference.h").read_text() + ) + assert len(shape_inference_decls) > 0 + upstream_shape_inference_decls = extract_signatures( + TORCH_DIR.joinpath( + "torch", "csrc", "lazy", "core", "shape_inference.h" + ).read_text() + ) + assert len(upstream_shape_inference_decls) > 0 + shape_inference_defs = extract_signatures( + self.backend_path.joinpath("shape_inference.cpp").read_text() + ) + assert len(shape_inference_decls) > len(shape_inference_defs) + + missing_defs = ( + shape_inference_decls + - upstream_shape_inference_decls + - shape_inference_defs + ) + if missing_defs: + self.backend_path.joinpath("generated", "shape_inference.cpp").write_text( + dedent( + """ + // This file contains autogenerated Lazy Shape Inference placeholders + // for ops that dont have a corresponding structured kernel or shape definition + + #include "shape_inference.h" + #include "../../utils/exception.h" + namespace torch {{ + namespace lazy {{ + {} + }} // namespace lazy + }} // namespace torch + """ + ).format( + "".join( + dedent( + f""" + std::vector {name}({args}) {{ + UNIMPLEMENTED_FUNCTION_ERROR(); + }} + """ + ) + for name, args in map( + global_signatures.get, sorted(missing_defs) + ) ) ) ) + + unnecessary_defs = shape_inference_defs - shape_inference_decls + if unnecessary_defs: + unnecessary_defs = "\n\t".join( + f"{name}({args})" + for name, args in map(global_signatures.get, unnecessary_defs) + ) + warnings.warn( + f"Unnecessary shape inference definitions found for:\n\t{unnecessary_defs}" + ) + + def generate_backend(self): + print("Running Lazy Tensor Autogen") + + # No fallback code allowed + def gen_fallback_code(*args, **kwargs): + return "" + + torchgen.dest.lazy_ir.gen_fallback_code = gen_fallback_code + + torchgen.gen_lazy_tensor.run_gen_lazy_tensor( + backend_name="TorchMlir", + aten_path=str(TORCH_DIR.joinpath("aten", "src", "ATen")), + source_yaml=str(self.source_yaml), + output_dir=str(self.backend_path.joinpath("generated")), + dry_run=False, + impl_path=str(self.backend_path.joinpath("mlir_native_functions.cpp")), + node_base="torch::lazy::TorchMlirNode", + node_base_hdr=str(self.backend_path.joinpath("mlir_node.h")), + tensor_class=self.tensor_class, + tensor_class_hdr="torch/csrc/lazy/core/tensor.h", + shape_inference_hdr=str( + self.backend_path.joinpath("generated", "shape_inference.h") + ), + lazy_ir_generator=GenMlirLazyIr, ) - unnecessary_defs = shape_inference_defs - shape_inference_decls - if unnecessary_defs: - unnecessary_defs = "\n\t".join( - f"{name}({args})" - for name, args in map(global_signatures.get, unnecessary_defs) - ) - warnings.warn( - f"Unnecessary shape inference definitions found for:\n\t{unnecessary_defs}" + # Remove lazy_tensor_core imports + subprocess.check_call( + [ + "sed", + "-i", + "/lazy_tensor_core/d", + str(self.backend_path.joinpath("generated", "LazyNativeFunctions.cpp")), + ] ) + def __call__(self): + self.generate_native_functions() + self.generate_shape_inference() + self.generate_backend() + def main(args): - script_path = Path(__file__).resolve() - config_path = ( - Path(__file__).resolve().parent.joinpath("autogen_ltc_backend.yaml") - ) - torch_ops_file = TORCH_MLIR_DIR.joinpath( - "include", - "torch-mlir", - "Dialect", - "Torch", - "IR", - "GeneratedTorchOps.td", - ) - assert torch_ops_file.exists() - native_functions = TORCH_MLIR_DIR.joinpath( - "generated_native_functions.yaml" - ) - backend_path = TORCH_MLIR_DIR.joinpath( - "python", "torch_mlir", "csrc", "base_lazy_backend" - ) - assert backend_path.is_dir() - - torchgen_path = Path(torchgen.__path__[0]).resolve() - assert torchgen_path.is_dir() + generator = GenTorchMlirLTC() prev_hash = None - hash_file = TORCH_MLIR_DIR.joinpath("generated_backend.hash") + hash_file = TORCH_MLIR_DIR.joinpath("build", "generated_backend.hash") if hash_file.exists(): prev_hash = hash_file.read_text().strip() - m = hashlib.sha256() - - # Add file contents to hash - for path in ( - script_path, - config_path, - torch_ops_file, - native_functions, - backend_path.joinpath("LazyShapeInference.h"), - backend_path.joinpath("LazyShapeInference.cpp"), - torchgen_path.joinpath("dest", "lazy_ir.py"), - torchgen_path.joinpath("api", "lazy.py"), - torchgen_path.joinpath("model.py"), - ): - if path.exists(): - m.update(path.read_bytes()) - - new_hash = m.hexdigest().strip() + new_hash = generator.calculate_hash() if args.force or new_hash != prev_hash: - parsed_yaml, grouped_native_functions = generate_native_functions( - config_path, torch_ops_file, native_functions - ) - - generate_backend( - native_functions, - backend_path, - parsed_yaml, - grouped_native_functions, - ) - + generator() hash_file.write_text(new_hash) diff --git a/python/torch_mlir/csrc/CMakeLists.txt b/python/torch_mlir/csrc/CMakeLists.txt index 781a87731..14fd9d1fd 100644 --- a/python/torch_mlir/csrc/CMakeLists.txt +++ b/python/torch_mlir/csrc/CMakeLists.txt @@ -33,9 +33,9 @@ add_custom_target( add_library(torch_mlir_ltc_backend SHARED base_lazy_backend/backend_impl.cpp base_lazy_backend/generated/LazyNativeFunctions.cpp - base_lazy_backend/generated/GenLazyShapeInference.cpp base_lazy_backend/generated/RegisterLazy.cpp - base_lazy_backend/LazyShapeInference.cpp + base_lazy_backend/generated/shape_inference.cpp + base_lazy_backend/shape_inference.cpp base_lazy_backend/mlir_lowering_context.cpp base_lazy_backend/mlir_native_functions.cpp base_lazy_backend/mlir_node.cpp diff --git a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h b/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h deleted file mode 100644 index 956f9fc46..000000000 --- a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h +++ /dev/null @@ -1,71 +0,0 @@ -//===- LazyShapeInference.h -----------------------------------------------===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Also available under a BSD-style license. See LICENSE. -// -//===----------------------------------------------------------------------===// - -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -namespace torch { -namespace lazy { - -// clang-format off - -TORCH_API std::vector compute_shape___and__(const at::Tensor & self, const at::Tensor & other); -TORCH_API std::vector compute_shape__reshape_alias(const at::Tensor & self, at::IntArrayRef size, at::IntArrayRef stride); -TORCH_API std::vector compute_shape_abs(const at::Tensor & self); -TORCH_API std::vector compute_shape_add(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha); -TORCH_API std::vector compute_shape_arange_out(const at::Scalar & start, const at::Scalar & end, const at::Scalar & step, at::Tensor & out); -TORCH_API std::vector compute_shape_bernoulli(const at::Tensor & self, c10::optional generator); -TORCH_API std::vector compute_shape_bernoulli_functional(const at::Tensor & self, const at::Tensor & p, c10::optional generator); -TORCH_API std::vector compute_shape_bincount(const at::Tensor & self, const c10::optional & weights, int64_t minlength); -TORCH_API std::vector compute_shape_bucketize(const at::Tensor & self, const at::Tensor & boundaries, bool out_int32, bool right); -TORCH_API std::vector compute_shape_constant_pad_nd(const at::Tensor & self, at::IntArrayRef pad, const at::Scalar & value); -TORCH_API std::vector compute_shape_convolution(const at::Tensor & input, const at::Tensor & weight, const c10::optional & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups); -TORCH_API std::vector compute_shape_convolution_overrideable(const at::Tensor & input, const at::Tensor & weight, const c10::optional & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups); -TORCH_API std::vector compute_shape_div(const at::Tensor & self, const at::Scalar & other); -TORCH_API std::vector compute_shape_embedding(const at::Tensor & weight, const at::Tensor & indices, int64_t padding_idx, bool scale_grad_by_freq, bool sparse); -TORCH_API std::vector compute_shape_embedding_dense_backward(const at::Tensor & grad_output, const at::Tensor & indices, int64_t num_weights, int64_t padding_idx, bool scale_grad_by_freq); -TORCH_API std::vector compute_shape_flip(const at::Tensor & self, at::IntArrayRef dims); -TORCH_API std::vector compute_shape_fmod(const at::Tensor & self, const at::Scalar & other); -TORCH_API std::vector compute_shape_hardswish(const at::Tensor & self); -TORCH_API std::vector compute_shape_hardtanh(const at::Tensor & self, const at::Scalar & min_val, const at::Scalar & max_val); -TORCH_API std::vector compute_shape_index_select(const at::Tensor & self, int64_t dim, const at::Tensor & index); -TORCH_API std::vector compute_shape_logical_or(const at::Tensor & self, const at::Tensor & other); -TORCH_API std::vector compute_shape_logsumexp(const at::Tensor & self, at::IntArrayRef dim, bool keepdim); -TORCH_API std::vector compute_shape_masked_fill(const at::Tensor & self, const at::Tensor & mask, const at::Scalar & value); -TORCH_API std::vector compute_shape_masked_select(const at::Tensor & self, const at::Tensor & mask); -TORCH_API std::vector compute_shape_max(const at::Tensor & self); -TORCH_API std::vector compute_shape_mean(const at::Tensor & self, c10::optional dtype); -TORCH_API std::vector compute_shape_mul(const at::Tensor & self, const at::Scalar & other); -TORCH_API std::vector compute_shape_native_batch_norm(const at::Tensor & input, const c10::optional & weight, const c10::optional & bias, const c10::optional & running_mean, const c10::optional & running_var, bool training, double momentum, double eps); -TORCH_API std::vector compute_shape_native_batch_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, const c10::optional & weight, const c10::optional & running_mean, const c10::optional & running_var, const c10::optional & save_mean, const c10::optional & save_invstd, bool train, double eps, ::std::array output_mask); -TORCH_API std::vector compute_shape_native_dropout(const at::Tensor & input, double p, c10::optional train); -TORCH_API std::vector compute_shape_native_dropout_backward(const at::Tensor & grad_output, const at::Tensor & mask, double scale); -TORCH_API std::vector compute_shape_native_layer_norm(const at::Tensor & input, at::IntArrayRef normalized_shape, const c10::optional & weight, const c10::optional & bias, double eps); -TORCH_API std::vector compute_shape_native_layer_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, at::IntArrayRef normalized_shape, const at::Tensor & mean, const at::Tensor & rstd, const c10::optional & weight, const c10::optional & bias, ::std::array output_mask); -TORCH_API std::vector compute_shape_new_empty(const at::Tensor & self, at::IntArrayRef size, c10::optional dtype, c10::optional layout, c10::optional device, c10::optional pin_memory); -TORCH_API std::vector compute_shape_relu(const at::Tensor & self); -TORCH_API std::vector compute_shape_repeat(const at::Tensor & self, at::IntArrayRef repeats); -TORCH_API std::vector compute_shape_resize_functional(const at::Tensor & self, at::IntArrayRef size, c10::optional memory_format); -TORCH_API std::vector compute_shape_sub(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha); -TORCH_API std::vector compute_shape_sum(const at::Tensor & self, c10::optional dtype); -TORCH_API std::vector compute_shape_uniform_functional(const at::Tensor & self, double from, double to, c10::optional generator); -TORCH_API std::vector compute_shape_zero_functional(const at::Tensor & self); - - -// clang-format on - -} // namespace lazy -} // namespace torch diff --git a/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp b/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp index 83f6af571..7b743613e 100644 --- a/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp +++ b/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp @@ -29,7 +29,7 @@ #include "../utils/exception.h" #include "../utils/sys_utils.h" -#include "LazyShapeInference.h" +#include "generated/shape_inference.h" #include "generated/LazyNativeFunctions.h" #include "ops/to_copy.h" diff --git a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp b/python/torch_mlir/csrc/base_lazy_backend/shape_inference.cpp similarity index 98% rename from python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp rename to python/torch_mlir/csrc/base_lazy_backend/shape_inference.cpp index b475d576e..4ef9029aa 100644 --- a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp +++ b/python/torch_mlir/csrc/base_lazy_backend/shape_inference.cpp @@ -12,7 +12,7 @@ #include #include "../utils/exception.h" -#include "LazyShapeInference.h" +#include "generated/shape_inference.h" namespace torch { namespace lazy {