def build_mlir_module(op: OpOverload, kwargs: Dict[str, Any]) -> ir.Module: """Translate input function into an MLIR module in the `torch` dialect. Parameters ---------- op: OpOverload Callable from the torch.ops.aten module/namespace that has a _schema field. kwargs: Dict A dictionary with all arguments passed in through __torch_dispatch__ (including int/float,bool params). Returns ------- ir.Module Translation of the input module into an MLIR module. """ # The assert here is to catch tensor shapes that have size 0 dimensions, such as those produced in # the course of evaluating SliceEndSleStartModule_basic and SliceOutOfLowerBoundEndIndexModule_basic. # Such 0 size dimensions fail the assert at mlir/lib/IR/BuiltinTypes.cpp, line 887 annotations = [] for arg_name, arg in kwargs.items(): if isinstance(arg, torch.Tensor): assert np.prod( arg.shape) != 0, f"{arg_name} has invalid shape {arg.shape}" annotations.append( TorchTensorType(shape=tuple(arg.shape), dtype=arg.dtype)) annotations = tuple(annotations) script_fun = build_ts_script_function(op._schema, kwargs) assert len(annotations) == len(list(script_fun.graph.inputs( ))), "Number of annotations and number of graph inputs differs." mb = ModuleBuilder() mb.import_function(script_fun) func_op = get_func_op_with_name(mb.module, script_fun.name) assert ( func_op is not None ), "Unable to find FuncOp in new module. Make sure function was imported correctly into ModuleBuilder" func_annotation = Annotation(annotations) arg_attrs = AnnotationConverter.to_mlir_array_attr(func_annotation, mb.context) func_op.attributes["arg_attrs"] = arg_attrs return mb.module
def convert_torchscript_module_to_torch_backend_contract_mlir(program: torch.nn.Module): """Perform common lowering from TorchScript to Torch MLIR Returns an MLIR module that satisfies the Torch backend contract. """ mb = ModuleBuilder() scripted = torch.jit.script(program) class_annotator = ClassAnnotator() extract_annotations(program, scripted, class_annotator) # TODO: Find a way to make each of these calls own its own # "debuggable error report" situation. try: original_stderr = sys.stderr sys.stderr = StringIO() # Import the TorchScript module to MLIR mb.import_module(scripted._c, class_annotator) except Exception as e: raise Exception(f""" PyTorch TorchScript module -> torch-mlir Object Graph IR import failed with: Exception: {e} Diagnostics: {sys.stderr.getvalue()} """) from None finally: sys.stderr = original_stderr run_pipeline_with_repro_report( mb.module, "torchscript-module-to-torch-backend-pipeline", "Lowering TorchScript Object Graph IR -> Torch Backend IR") return mb.module
# -*- Python -*- # This file is licensed under a pytorch-style license # See LICENSE.pytorch for license information. import torch from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder # RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s mb = ModuleBuilder() # Note: The "if without else" case is handled by yielding None from the # else branch and making all defined values optional, so no special handling # is needed. # CHECK-LABEL: @__torch__.prim_If( # CHECK-SAME: %[[B:.*]]: !torch.bool, # CHECK-SAME: %[[I:.*]]: !torch.int) -> !torch.int { @mb.import_function @torch.jit.script def prim_If(b: bool, i: int): # CHECK: %[[RES:.*]] = torch.prim.If %[[B]] -> (!torch.int) { # CHECK: %[[ADD:.*]] = torch.aten.add.int %[[I]], %[[I]] # CHECK: torch.prim.If.yield %[[ADD]] : !torch.int # CHECK: } else { # CHECK: %[[MUL:.*]] = torch.aten.mul.int %[[I]], %[[I]] # CHECK: torch.prim.If.yield %[[MUL]] : !torch.int # CHECK: } # CHECK: return %[[RES:.*]] : !torch.int if b:
# -*- Python -*- # This file is licensed under a pytorch-style license # See LICENSE.pytorch for license information. import typing import torch from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder # RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s mb = ModuleBuilder() # CHECK-LABEL: func.func private @__torch__.TestModule.forward # CHECK-SAME: (%[[ARG0:.*]]: !torch.nn.Module<"__torch__.TestModule">, %[[ARG1:.*]]: !torch.tensor) -> !torch.tensor { # CHECK: %[[VAL_2:.*]] = constant @__torch__.identity : (!torch.tensor) -> !torch.tensor # CHECK: %[[VAL_3:.*]] = call_indirect %[[VAL_2]](%[[ARG1]]) : (!torch.tensor) -> !torch.tensor # CHECK: return %[[VAL_3]] : !torch.tensor # CHECK: } # CHECK-LABEL: func.func private @__torch__.identity # CHECK-SAME: (%[[ARG:.*]]: !torch.tensor) -> !torch.tensor { # CHECK: return %[[ARG]] : !torch.tensor # CHECK: } # CHECK-LABEL: torch.class_type @__torch__.TestModule { # CHECK: torch.method "forward", @__torch__.TestModule.forward # CHECK: } def identity(x): return x
# -*- Python -*- # This file is licensed under a pytorch-style license # See LICENSE.pytorch for license information. import torch from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder from typing import Tuple, Optional, NamedTuple from utils import create_script_function # RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s mb = ModuleBuilder() NT = NamedTuple('NT', [('f1', Optional[torch.Tensor]), ('f2', Optional[torch.Tensor])]) # CHECK-LABEL: func.func @__torch__.tuple( # CHECK-SAME: %[[T0:.*]]: !torch.tensor, # CHECK-SAME: %[[T1:.*]]: !torch.tensor) -> # CHECK-SAME: !torch.tuple<tensor, tensor> { # CHECK: %[[RET:.*]] = torch.prim.TupleConstruct %[[T0]], %[[T1]] : # CHECK-SAME: !torch.tensor, !torch.tensor -> !torch.tuple<tensor, tensor> # CHECK: return %[[RET]] : !torch.tuple<tensor, tensor> @mb.import_function @torch.jit.script def tuple(t0, t1): return t0, t1 # CHECK-LABEL: func.func @__torch__.tuple_optional(
# -*- Python -*- # This file is licensed under a pytorch-style license # See LICENSE.pytorch for license information. import typing import torch from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder # RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s mb = ModuleBuilder() class TestModule(torch.nn.Module): def __init__(self): super().__init__() self.exported = 1 self.not_exported = 2 def forward(self): return self.not_exported_method() def not_exported_method(self): return test_module = TestModule() recursivescriptmodule = torch.jit.script(test_module) annotator = ClassAnnotator() class_type = recursivescriptmodule._c._type()
# -*- Python -*- # This file is licensed under a pytorch-style license # See LICENSE.pytorch for license information. import typing import torch from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder from utils import create_script_function import typing # RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s mb = ModuleBuilder() # CHECK-LABEL: func.func @__torch__.prim_NumToTensor( # CHECK-SAME: %[[ARG:.*]]: !torch.int) -> !torch.tensor { # CHECK: %[[RET:.*]] = torch.prim.NumToTensor.Scalar %[[ARG]] : !torch.int -> !torch.tensor # CHECK: return %[[RET]] : !torch.tensor # CHECK: } @mb.import_function @torch.jit.script def prim_NumToTensor(i: int): return _to_tensor(i) # CHECK-LABEL: func.func @__torch__.prim_Print( # CHECK-SAME: %[[ARG:.*]]: !torch.tensor) -> !torch.none { # CHECK: %[[STR:.*]] = torch.constant.str "x" # CHECK: torch.prim.Print(%[[STR]], %[[ARG]]) : !torch.str, !torch.tensor
def compile(model: torch.nn.Module, example_args: Union[_example_arg, Sequence[_example_arg]], output_type: Union[str, "OutputType"] = OutputType.TORCH, use_tracing: bool = False, verbose: bool = False): """Convert a PyTorch model to MLIR. Args: model: The PyTorch model to convert. example_args: A list of example arguments to use when inferring the shapes of the arguments to `forward` method of the model. A single tensor is treated as a list of a single tensor. A TensorPlaceholder object is also allowed in the place of any Tensor. output_type: The kind of output to produce. See `OutputType` for more details. use_tracing: If True, use `torch.jit.trace` to convert the model to JIT IR rather than `torch.jit.script`. Returns: An MLIR module that contains the converted model in the specified output type. """ output_type = OutputType.get(output_type) # Special case -- many models have just one input, so canonicalize a single # tensor to a list of a single tensor to make the API more ergonomic. if isinstance(example_args, (torch.Tensor, TensorPlaceholder)): example_args = (example_args, ) # TODO: Don't hardcode "forward". See `torch.onnx.export` and # `torch.jit.trace_module` for API inspiration. if use_tracing: scripted = torch.jit.trace(model, tuple(example_args)) else: scripted = torch.jit.script(model) # Convert all concrete inputs to TensorPlaceholder's, for consistency. arg_placeholders = [] for arg in example_args: if isinstance(arg, TensorPlaceholder): arg_placeholders.append(arg) else: assert isinstance(arg, torch.Tensor) arg_placeholders.append(TensorPlaceholder.like(arg)) class_annotator = ClassAnnotator() forward_annotation = [None] for arg in arg_placeholders: # Assume that all tensors have value semantics for now. forward_annotation.append((arg.shape, arg.dtype, True)) class_annotator.exportNone(scripted._c._type()) class_annotator.exportPath(scripted._c._type(), ["forward"]) class_annotator.annotateArgs(scripted._c._type(), ["forward"], forward_annotation) mb = ModuleBuilder() mb.import_module(scripted._c, class_annotator) if output_type == OutputType.RAW: return mb.module run_pipeline_with_repro_report( mb.module, "torchscript-module-to-torch-backend-pipeline", "Lowering TorchScript IR -> Torch Backend IR") if verbose: print("\n====================") print("Torch Backend IR") print(mb.module) if output_type == OutputType.TORCH: return mb.module if output_type == OutputType.TOSA: run_pipeline_with_repro_report( mb.module, "torch-backend-to-tosa-backend-pipeline", "Lowering Torch Backend IR -> TOSA Backend IR") if verbose: print("\n====================") print("TOSA Backend IR") print(mb.module) return mb.module if output_type == OutputType.LINALG_ON_TENSORS: run_pipeline_with_repro_report( mb.module, "torch-backend-to-linalg-on-tensors-backend-pipeline", "Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR") if verbose: print("\n====================") print("LINALG Backend IR") print(mb.module) return mb.module elif output_type == OutputType.MHLO: run_pipeline_with_repro_report( mb.module, "torch-backend-to-mhlo-backend-pipeline", "Lowering Torch Backend IR -> MHLO Backend IR") if verbose: print("\n====================") print("MHLO Backend IR") print(mb.module) return mb.module raise Exception(f"Unknown OutputType: {output_type}")
def compile(model: torch.nn.Module, example_args: Union[_example_arg, Sequence[_example_arg]], output_type: Union[str, "OutputType"] = OutputType.TORCH, use_tracing: bool = False, ignore_traced_shapes=False, verbose: bool = False): """Convert a PyTorch model to MLIR. Args: model: The PyTorch model to convert. example_args: A list of example arguments to use when inferring the shapes of the arguments to `forward` method of the model. A single tensor is treated as a list of a single tensor. A TensorPlaceholder object is also allowed in the place of any Tensor. output_type: The kind of output to produce. See `OutputType` for more details. use_tracing: If True, use `torch.jit.trace` to convert the model to JIT IR rather than `torch.jit.script`. ignore_traced_shapes: If True, ignore the shapes that were observed during tracing. This should only be used if one knows that the original traced program would result in the same trace (modulo shapes) for all shape combinations implied by any `TensorPlaceholder`'s used as `example_args`. Also, strictly-speaking, this option covers dtypes too, but we just say "shapes" to be succinct. verbose: If true, print extra information about the conversion. Returns: An MLIR module that contains the converted model in the specified output type. """ output_type = OutputType.get(output_type) if ignore_traced_shapes and not use_tracing: raise Exception("`ignore_traced_shapes` requires `use_tracing`") # Special case -- many models have just one input, so canonicalize a single # tensor to a list of a single tensor to make the API more ergonomic. if isinstance(example_args, (torch.Tensor, TensorPlaceholder)): example_args = (example_args, ) # TODO: Don't hardcode "forward". See `torch.onnx.export` and # `torch.jit.trace_module` for API inspiration. if use_tracing: example_args_for_trace = [] for arg in example_args: if isinstance(arg, TensorPlaceholder): if not ignore_traced_shapes: # To avoid accidental footguns, we require # `ignore_traced_shapes` to be true if we're using # TensorPlaceholder's, as it falls into the same # "hopefully the trace works for different inputs" bucket # of concerns. raise Exception( "TensorPlaceholder can only be used with tracing when `ignore_traced_shapes=True`" ) # For any dynamic dimensions, replace them with "7" arbitrarily. # If a user is using dynamic dimensions with tracing, they are # walking on thin ice already -- assume they know what they are # doing. shape = [s if s != -1 else 7 for s in arg.shape] example_args_for_trace.append( torch.ones(*shape, dtype=arg.dtype)) else: example_args_for_trace.append(arg) scripted = torch.jit.trace(model, tuple(example_args_for_trace)) else: scripted = torch.jit.script(model) # Convert all concrete inputs to TensorPlaceholder's, for consistency. arg_placeholders = [] for arg in example_args: if isinstance(arg, TensorPlaceholder): arg_placeholders.append(arg) else: assert isinstance(arg, torch.Tensor) arg_placeholders.append(TensorPlaceholder.like(arg)) class_annotator = ClassAnnotator() forward_annotation = [None] for arg in arg_placeholders: # Assume that all tensors have value semantics for now. forward_annotation.append((arg.shape, arg.dtype, True)) class_annotator.exportNone(scripted._c._type()) class_annotator.exportPath(scripted._c._type(), ["forward"]) class_annotator.annotateArgs(scripted._c._type(), ["forward"], forward_annotation) mb = ModuleBuilder() import_options = ImportOptions() import_options.ignoreExistingTensorShapesAndDtypes = ignore_traced_shapes mb.import_module(scripted._c, class_annotator, import_options) if output_type == OutputType.RAW: return mb.module run_pipeline_with_repro_report( mb.module, "torchscript-module-to-torch-backend-pipeline", "Lowering TorchScript IR -> Torch Backend IR") if verbose: print("\n====================") print("Torch Backend IR") print(mb.module) if output_type == OutputType.TORCH: return mb.module if output_type == OutputType.TOSA: run_pipeline_with_repro_report( mb.module, "torch-backend-to-tosa-backend-pipeline", "Lowering Torch Backend IR -> TOSA Backend IR") if verbose: print("\n====================") print("TOSA Backend IR") print(mb.module) return mb.module if output_type == OutputType.LINALG_ON_TENSORS: run_pipeline_with_repro_report( mb.module, "torch-backend-to-linalg-on-tensors-backend-pipeline", "Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR") if verbose: print("\n====================") print("LINALG Backend IR") print(mb.module) return mb.module elif output_type == OutputType.MHLO: run_pipeline_with_repro_report( mb.module, "torch-backend-to-mhlo-backend-pipeline", "Lowering Torch Backend IR -> MHLO Backend IR") if verbose: print("\n====================") print("MHLO Backend IR") print(mb.module) return mb.module raise Exception(f"Unknown OutputType: {output_type}")