Refactor autogen (#925)

.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

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
+from textwrap import dedent, indent
-
-import yaml
 
 # 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)) == "<class 'torchgen.api.types.OptionalCType'>"
 
-
+def reindent(text, prefix=""):
-def generate_native_functions(
+    return indent(dedent(text), prefix)
-    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
 
 
 @dataclass(frozen=True)
 class GenMlirLazyIr(torchgen.dest.GenLazyIR):
+    def isOptionalCType(self, arg):
+        return str(type(arg)) == "<class 'torchgen.api.types.OptionalCType'>"
 
-    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):
+                if self.isOptionalCType(arg.lazy_type):
-                    emplace_arguments.append(f"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr")
+                    emplace_arguments.append(
-                    continue
+                        f"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr"
-                emplace_arguments.append('loctx->GetOutputOp(operand(i++))')
+                    )
-                continue
+                else:
-            emplace_arguments.append(f'"{arg.name}", {arg.name}')
+                    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])
+            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_kwarg_values = [
-        emplace_kwarguments = "\n    ".join(
+            f'"{t.name}", loctx->GetOutputOp(operand(i++))'
-            [f"kwarguments.emplace_back({a});" for a in emplace_kwarg_values + emplace_kwarg_scalars])
+            for t in schema.keyword_values
-
-        return f"""
-  {signature} {{
-    PRINT_FUNCTION();
-    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::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")),
         ]
-    )
+        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
+        return reindent(
-    import re
+            f"""
+            {signature} {{
+                PRINT_FUNCTION();
+                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::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(
+                return {schema.aten_name}_out;
-        r"std::vector<torch::lazy::Shape>\s+(?P<name>\w+)\((?P<signature>[^\)]+)\)"
+            }}
-    )
+            """,
-    global_signatures = {}
+            "  ",
+        )
 
-    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(
+class GenTorchMlirLTC:
-        TORCH_DIR.joinpath("torch", "csrc", "lazy", "core", "shape_inference.h")
+    def __init__(self):
-    )
+        self.script_path = Path(__file__).resolve()
-    assert len(upstream_shape_inference_decls) > 0
+        self.config_path = (
-    shape_inference_decls = extract_signatures(
+            Path(__file__).resolve().parent.joinpath("autogen_ltc_backend.yaml")
-        backend_path.joinpath("LazyShapeInference.h")
+        )
-    )
+        self.torch_ops_file = TORCH_MLIR_DIR.joinpath(
-    assert len(shape_inference_decls) > 0
+            "include",
-    shape_inference_defs = extract_signatures(
+            "torch-mlir",
-        backend_path.joinpath("LazyShapeInference.cpp")
+            "Dialect",
-    )
+            "Torch",
-    assert len(shape_inference_decls) > len(shape_inference_defs)
+            "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 = (
+        self.torchgen_path = Path(torchgen.__path__[0]).resolve()
-        shape_inference_decls
+        assert self.torchgen_path.is_dir()
-        - upstream_shape_inference_decls
+
-        - shape_inference_defs
+        self.tensor_class = "torch::lazy::LazyTensor"
-    )
+
-    if missing_defs:
+        # Set the lazy value class
-        backend_path.joinpath("generated", "GenLazyShapeInference.cpp").write_text(
+        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<torch::lazy::Shape>\s+(?P<name>\w+)\((?P<signature>[^\)]+)\)"
+        )
+        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 <ATen/Tensor.h>
-                #include "../../utils/exception.h"
+                #include <c10/core/ScalarType.h>
+                #include <c10/util/Optional.h>
+                #include <torch/csrc/lazy/core/ir.h>
+                #include <torch/csrc/lazy/core/shape.h>
+                #include <torch/csrc/lazy/core/shape_inference.h>
+                #include <vector>
+
                 namespace torch {{
                 namespace lazy {{
+
                 {}
+
                 }}  // namespace lazy
                 }}  // namespace torch
                 """
-            ).format(
+            ).format(os.linesep.join(sorted(shape_inference_decls)))
-                "".join(
+        )
-                    dedent(
+
-                        f"""
+        shape_inference_decls = extract_signatures(
-                        std::vector<Shape> {name}({args}) {{
+            self.backend_path.joinpath("generated", "shape_inference.h").read_text()
-                            UNIMPLEMENTED_FUNCTION_ERROR();
+        )
-                        }}
+        assert len(shape_inference_decls) > 0
-                        """
+        upstream_shape_inference_decls = extract_signatures(
-                    )
+            TORCH_DIR.joinpath(
-                    for name, args in map(
+                "torch", "csrc", "lazy", "core", "shape_inference.h"
-                        global_signatures.get, sorted(missing_defs)
+            ).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<torch::lazy::Shape> {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
+        # Remove lazy_tensor_core imports
-    if unnecessary_defs:
+        subprocess.check_call(
-        unnecessary_defs = "\n\t".join(
+            [
-            f"{name}({args})"
+                "sed",
-            for name, args in map(global_signatures.get, unnecessary_defs)
+                "-i",
-        )
+                "/lazy_tensor_core/d",
-        warnings.warn(
+                str(self.backend_path.joinpath("generated", "LazyNativeFunctions.cpp")),
-            f"Unnecessary shape inference definitions found for:\n\t{unnecessary_defs}"
+            ]
         )
 
+    def __call__(self):
+        self.generate_native_functions()
+        self.generate_shape_inference()
+        self.generate_backend()
+
 
 def main(args):
-    script_path = Path(__file__).resolve()
+    generator = GenTorchMlirLTC()
-    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()
 
     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()
+    new_hash = generator.calculate_hash()
-
-    # 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()
 
     if args.force or new_hash != prev_hash:
-        parsed_yaml, grouped_native_functions = generate_native_functions(
+        generator()
-            config_path, torch_ops_file, native_functions
-        )
-
-        generate_backend(
-            native_functions,
-            backend_path,
-            parsed_yaml,
-            grouped_native_functions,
-        )
-
         hash_file.write_text(new_hash)
 
 

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

python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h

@@ -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 <ATen/Tensor.h>
-#include <c10/core/ScalarType.h>
-#include <c10/util/Optional.h>
-#include <torch/csrc/lazy/core/ir.h>
-#include <torch/csrc/lazy/core/shape.h>
-#include <torch/csrc/lazy/core/shape_inference.h>
-#include <vector>
-
-namespace torch {
-namespace lazy {
-
-// clang-format off
-
-TORCH_API std::vector<torch::lazy::Shape> compute_shape___and__(const at::Tensor & self, const at::Tensor & other);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape__reshape_alias(const at::Tensor & self, at::IntArrayRef size, at::IntArrayRef stride);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_abs(const at::Tensor & self);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_add(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_arange_out(const at::Scalar & start, const at::Scalar & end, const at::Scalar & step, at::Tensor & out);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_bernoulli(const at::Tensor & self, c10::optional<at::Generator> generator);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_bernoulli_functional(const at::Tensor & self, const at::Tensor & p, c10::optional<at::Generator> generator);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_bincount(const at::Tensor & self, const c10::optional<at::Tensor> & weights, int64_t minlength);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_bucketize(const at::Tensor & self, const at::Tensor & boundaries, bool out_int32, bool right);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_constant_pad_nd(const at::Tensor & self, at::IntArrayRef pad, const at::Scalar & value);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_convolution(const at::Tensor & input, const at::Tensor & weight, const c10::optional<at::Tensor> & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_convolution_overrideable(const at::Tensor & input, const at::Tensor & weight, const c10::optional<at::Tensor> & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_div(const at::Tensor & self, const at::Scalar & other);
-TORCH_API std::vector<torch::lazy::Shape> 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<torch::lazy::Shape> 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<torch::lazy::Shape> compute_shape_flip(const at::Tensor & self, at::IntArrayRef dims);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_fmod(const at::Tensor & self, const at::Scalar & other);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_hardswish(const at::Tensor & self);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_hardtanh(const at::Tensor & self, const at::Scalar & min_val, const at::Scalar & max_val);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_index_select(const at::Tensor & self, int64_t dim, const at::Tensor & index);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_logical_or(const at::Tensor & self, const at::Tensor & other);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_logsumexp(const at::Tensor & self, at::IntArrayRef dim, bool keepdim);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_masked_fill(const at::Tensor & self, const at::Tensor & mask, const at::Scalar & value);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_masked_select(const at::Tensor & self, const at::Tensor & mask);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_max(const at::Tensor & self);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_mean(const at::Tensor & self, c10::optional<at::ScalarType> dtype);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_mul(const at::Tensor & self, const at::Scalar & other);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_batch_norm(const at::Tensor & input, const c10::optional<at::Tensor> & weight, const c10::optional<at::Tensor> & bias, const c10::optional<at::Tensor> & running_mean, const c10::optional<at::Tensor> & running_var, bool training, double momentum, double eps);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_batch_norm_backward(const at::Tensor & grad_out, const at::Tensor & input, const c10::optional<at::Tensor> & weight, const c10::optional<at::Tensor> & running_mean, const c10::optional<at::Tensor> & running_var, const c10::optional<at::Tensor> & save_mean, const c10::optional<at::Tensor> & save_invstd, bool train, double eps, ::std::array<bool,3> output_mask);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_dropout(const at::Tensor & input, double p, c10::optional<bool> train);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_dropout_backward(const at::Tensor & grad_output, const at::Tensor & mask, double scale);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_native_layer_norm(const at::Tensor & input, at::IntArrayRef normalized_shape, const c10::optional<at::Tensor> & weight, const c10::optional<at::Tensor> & bias, double eps);
-TORCH_API std::vector<torch::lazy::Shape> 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<at::Tensor> & weight, const c10::optional<at::Tensor> & bias, ::std::array<bool,3> output_mask);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_new_empty(const at::Tensor & self, at::IntArrayRef size, c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout, c10::optional<at::Device> device, c10::optional<bool> pin_memory);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_relu(const at::Tensor & self);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_repeat(const at::Tensor & self, at::IntArrayRef repeats);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_resize_functional(const at::Tensor & self, at::IntArrayRef size, c10::optional<at::MemoryFormat> memory_format);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_sub(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_sum(const at::Tensor & self, c10::optional<at::ScalarType> dtype);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_uniform_functional(const at::Tensor & self, double from, double to, c10::optional<at::Generator> generator);
-TORCH_API std::vector<torch::lazy::Shape> compute_shape_zero_functional(const at::Tensor & self);
-
-
-// clang-format on
-
-} // namespace lazy
-} // namespace torch

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"
 

python/torch_mlir/csrc/base_lazy_backend/shape_inference.cpp

@@ -12,7 +12,7 @@
 #include <cmath>
 
 #include "../utils/exception.h"
-#include "LazyShapeInference.h"
+#include "generated/shape_inference.h"
 
 namespace torch {
 namespace lazy {