mirror of https://github.com/llvm/torch-mlir
505 lines
18 KiB
Python
505 lines
18 KiB
Python
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import logging
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import warnings
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from shutil import which
|
|
from textwrap import dedent, indent
|
|
|
|
# PyTorch's LTC backend autogen script
|
|
import torchgen
|
|
import torchgen.dest.lazy_ir
|
|
import torchgen.gen_lazy_tensor
|
|
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.gen_backend_stubs import parse_backend_yaml
|
|
|
|
TORCH_DIR = Path(importlib.util.find_spec("torch").origin).resolve().parent.parent
|
|
if TORCH_DIR.joinpath("torch", "include").is_dir():
|
|
TORCH_DIR = TORCH_DIR.joinpath("torch", "include")
|
|
TORCHGEN_DIR = Path(torchgen.__path__[0]).resolve()
|
|
TORCH_MLIR_DIR = Path(__file__).resolve().parent.parent
|
|
|
|
|
|
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: LazyIrSchema):
|
|
signature = "TorchMlirOpVector Lower(TorchMlirFunction function, TorchMlirLoweringContext* loctx) const override"
|
|
|
|
if schema.properties.LowerDeclOnly:
|
|
return f"{signature};"
|
|
elif not schema.properties.Lower:
|
|
return ""
|
|
|
|
emplace_arguments = []
|
|
for arg in schema.positional_args:
|
|
if arg.is_lazy_value:
|
|
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 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);
|
|
TORCH_CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)});
|
|
|
|
return {schema.aten_name}_out;
|
|
}}
|
|
""",
|
|
" ",
|
|
)
|
|
|
|
|
|
class GenTorchMlirLTC:
|
|
def __init__(self, binary_dir):
|
|
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(
|
|
# fmt: off
|
|
"include", "torch-mlir", "Dialect", "Torch", "IR", "GeneratedTorchOps.td",
|
|
# fmt: on
|
|
)
|
|
assert self.torch_ops_file.exists()
|
|
self.binary_dir = Path(binary_dir)
|
|
assert self.binary_dir.is_dir(), f"Binary directory not found: {self.binary_dir}"
|
|
self.source_yaml = self.binary_dir.joinpath("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.generated_path = self.binary_dir.joinpath(
|
|
"python", "torch_mlir", "csrc", "base_lazy_backend", "generated"
|
|
)
|
|
self.generated_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Create symlink to match doc structure
|
|
generated_path = self.backend_path.joinpath("generated").resolve()
|
|
if not generated_path.exists():
|
|
generated_path.symlink_to(
|
|
os.path.relpath(self.generated_path, generated_path.parent),
|
|
target_is_directory=True,
|
|
)
|
|
|
|
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"),
|
|
TORCHGEN_DIR.joinpath("dest", "lazy_ir.py"),
|
|
TORCHGEN_DIR.joinpath("api", "lazy.py"),
|
|
TORCHGEN_DIR.joinpath("model.py"),
|
|
):
|
|
if path.exists():
|
|
m.update(path.read_bytes())
|
|
|
|
return m.hexdigest().strip()
|
|
|
|
def generate_native_functions(self):
|
|
logging.info("Generating Native Functions Yaml")
|
|
|
|
native_path = TORCHGEN_DIR.joinpath("packaged", "ATen", "native")
|
|
native_yaml_path = native_path.joinpath("native_functions.yaml")
|
|
tags_yaml_path = native_path.joinpath("tags.yaml")
|
|
|
|
ts_native_yaml_path = TORCH_DIR.joinpath(
|
|
"aten", "src", "ATen", "native", "ts_native_functions.yaml"
|
|
)
|
|
ts_native_yaml = None
|
|
if ts_native_yaml_path.exists():
|
|
ts_native_yaml = yaml.load(ts_native_yaml_path.read_text(), yaml.CLoader)
|
|
|
|
parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
|
|
self.native_functions = parsed_yaml.native_functions
|
|
self.backend_indices = parsed_yaml.backend_indices
|
|
self.grouped_native_functions = get_grouped_native_functions(
|
|
self.native_functions
|
|
)
|
|
|
|
def get_native_function_name(f):
|
|
func = f if hasattr(f, "func") else f.functional
|
|
return str(func.func.name)
|
|
|
|
self.native_functions = {
|
|
get_native_function_name(f): f for f in self.native_functions
|
|
}
|
|
|
|
def get_opnames(ops):
|
|
opnames = defaultdict(set)
|
|
for op in ops:
|
|
opname = op.split(".")[0]
|
|
opnames[opname].add(op)
|
|
return opnames
|
|
|
|
aten_funcs = get_opnames(
|
|
map(get_native_function_name, self.grouped_native_functions)
|
|
)
|
|
|
|
with self.config_path.open() as f:
|
|
config = yaml.load(f, yaml.CLoader)
|
|
|
|
# List of unsupported ops in LTC autogen because of some error
|
|
blacklist = set(config.get("blacklist", []))
|
|
|
|
# List of supported ops that we don't want to do the full codegen for
|
|
# primarily view ops
|
|
supported = set(config.get("supported", []))
|
|
|
|
# List of non-native ops to do IR codegen for
|
|
non_native = config.get("non_native", [])
|
|
|
|
# use ripgrep if available as its much faster
|
|
if which("rg") is not None:
|
|
cmd = ["rg", "-o", "-N", r"aten::[0-9a-zA-Z_\.]+"]
|
|
else:
|
|
cmd = ["grep", "-o", r"aten::[0-9a-zA-Z_\.]\+"]
|
|
|
|
torch_ops = set(
|
|
op[6:]
|
|
for op in subprocess.check_output(
|
|
cmd + [str(self.torch_ops_file)],
|
|
encoding="utf-8",
|
|
)
|
|
.strip()
|
|
.split(os.linesep)
|
|
)
|
|
torch_opnames = get_opnames(torch_ops)
|
|
|
|
# process ops list
|
|
ops = set()
|
|
composite_implicit = set()
|
|
|
|
for op in torch_ops:
|
|
if op not in self.native_functions:
|
|
continue
|
|
|
|
func = self.native_functions[op]
|
|
base = func.func.name.name.base
|
|
|
|
if base in blacklist or op in blacklist:
|
|
continue
|
|
if base in supported or op in supported:
|
|
continue
|
|
# Blacklist new_/_like ops since they are non-differentiable.
|
|
if any(o.startswith("new_") or o.endswith("_like") for o in (base, op)):
|
|
continue
|
|
|
|
if func.has_composite_implicit_autograd_kernel:
|
|
composite_implicit.add(op)
|
|
elif func.func.name.name.inplace:
|
|
for autogen in func.autogen:
|
|
if "functional" in autogen.overload_name:
|
|
ops.add(str(autogen))
|
|
else:
|
|
ops.add(op)
|
|
|
|
skipped = set(torch_ops) - ops - supported - composite_implicit
|
|
|
|
# List of ops autogen even if not explicitly supported by Torch-MLIR explicitly
|
|
ops |= set(config.get("whitelist", []))
|
|
|
|
# Additional ops to support that are not supported by Torch-MLIR explicitly
|
|
supported |= set(config.get("additional_ops", []))
|
|
|
|
self.ops = sorted(ops)
|
|
|
|
with self.source_yaml.open("w") as f:
|
|
source_yaml = {
|
|
"backend": "Lazy",
|
|
"cpp_namespace": "torch::lazy",
|
|
"full_codegen": self.ops,
|
|
"supported": sorted(supported),
|
|
"non_native": non_native,
|
|
}
|
|
yaml.dump(source_yaml, f, default_flow_style=False)
|
|
f.write(
|
|
dedent(
|
|
"""
|
|
|
|
# Composite implicit ops (supported by Torch-MLIR but not differentiable)
|
|
{composite_implicit}
|
|
# Skipped ops (supported by Torch-MLIR but no equivalent native function)
|
|
{skipped}
|
|
"""
|
|
).format(
|
|
composite_implicit=os.linesep.join(
|
|
f"# - {op}" for op in sorted(composite_implicit)
|
|
),
|
|
skipped=os.linesep.join(f"# - {op}" for op in sorted(skipped)),
|
|
)
|
|
)
|
|
|
|
if ts_native_yaml:
|
|
ts_full_codegen = set(ts_native_yaml["full_codegen"])
|
|
mlir_full_codegen = set(self.ops)
|
|
|
|
if ts_full_codegen - mlir_full_codegen:
|
|
logging.debug(
|
|
"Full Codegen ops supported by the TorchScript backend "
|
|
"but not by the Torch-MLIR backend:\n {}".format(
|
|
"\n ".join(sorted(ts_full_codegen - mlir_full_codegen))
|
|
)
|
|
)
|
|
|
|
if mlir_full_codegen - ts_full_codegen:
|
|
logging.debug(
|
|
"Full Codegen ops supported by the Torch-MLIR backend "
|
|
"but not by the TorchScript backend:\n {}".format(
|
|
"\n ".join(sorted(mlir_full_codegen - ts_full_codegen))
|
|
)
|
|
)
|
|
|
|
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.generated_path.joinpath("shape_inference.h").write_text(
|
|
dedent(
|
|
"""
|
|
// This file contains autogenerated Lazy Shape Inference declarations
|
|
// for ops that dont have a corresponding structured kernel or shape definition
|
|
|
|
#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(os.linesep.join(sorted(shape_inference_decls)))
|
|
)
|
|
|
|
shape_inference_decls = extract_signatures(
|
|
self.generated_path.joinpath("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.generated_path.joinpath("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 "torch_mlir/csrc/base_lazy_backend/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):
|
|
logging.info("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(TORCHGEN_DIR.joinpath("packaged", "ATen")),
|
|
source_yaml=str(self.source_yaml),
|
|
output_dir=str(self.generated_path),
|
|
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.generated_path.joinpath("shape_inference.h")),
|
|
lazy_ir_generator=GenMlirLazyIr,
|
|
)
|
|
|
|
def __call__(self):
|
|
self.generate_native_functions()
|
|
self.generate_shape_inference()
|
|
self.generate_backend()
|
|
|
|
|
|
def main(args):
|
|
generator = GenTorchMlirLTC(args.binary_dir)
|
|
|
|
hash_file = generator.binary_dir.joinpath("generated_backend.hash")
|
|
|
|
prev_hash = None
|
|
if hash_file.exists():
|
|
prev_hash = hash_file.read_text().strip()
|
|
|
|
new_hash = generator.calculate_hash()
|
|
|
|
if args.force or new_hash != prev_hash:
|
|
generator()
|
|
hash_file.write_text(new_hash)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"-b",
|
|
"--binary_dir",
|
|
type=str,
|
|
default=os.getenv(
|
|
"TORCH_MLIR_BINARY_DIR",
|
|
TORCH_MLIR_DIR.joinpath("build"),
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"-f",
|
|
"--force",
|
|
action="store_true",
|
|
)
|
|
parser.add_argument(
|
|
"-d",
|
|
"--debug",
|
|
help="Print lots of debugging statements",
|
|
action="store_const",
|
|
dest="loglevel",
|
|
const=logging.DEBUG,
|
|
default=logging.WARNING,
|
|
)
|
|
parser.add_argument(
|
|
"-v",
|
|
"--verbose",
|
|
help="Be verbose",
|
|
action="store_const",
|
|
dest="loglevel",
|
|
const=logging.INFO,
|
|
)
|
|
args = parser.parse_args()
|
|
logging.basicConfig(level=args.loglevel)
|
|
main(args)
|