mirror of https://github.com/llvm/torch-mlir
Refactor autogen (#925)
Jae Hoon (Antonio) Kim
Henry Tu
parent
dfcc26556a
commit
a62d60829c
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)) == "<class 'torchgen.api.types.OptionalCType'>"
|
||||
|
||||
|
||||
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)) == "<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):
|
||||
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<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")),
|
||||
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<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(
|
||||
r"std::vector<torch::lazy::Shape>\s+(?P<name>\w+)\((?P<signature>[^\)]+)\)"
|
||||
)
|
||||
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<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 "../../utils/exception.h"
|
||||
#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 {{
|
||||
|
||||
{}
|
||||
|
||||
}} // namespace lazy
|
||||
}} // namespace torch
|
||||
"""
|
||||
).format(
|
||||
"".join(
|
||||
dedent(
|
||||
f"""
|
||||
std::vector<Shape> {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<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
|
||||
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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
#include <cmath>
|
||||
|
||||
#include "../utils/exception.h"
|
||||
#include "LazyShapeInference.h"
|
||||
#include "generated/shape_inference.h"
|
||||
|
||||
namespace torch {
|
||||
namespace lazy {
|
||||
Loading…
Reference in New Issue