# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name, missing-function-docstring, not-callable # pylint: disable=import-outside-toplevel, unused-argument, use-list-literal # mypy: ignore-errors """PyTorch Dynamo backend of Relax.""" import functools import tvm_ffi import tvm from tvm.relax import build as relax_build from .fx_translator import from_fx def device_from_inputs(example_inputs): for x in example_inputs: if hasattr(x, "device"): return x.device return None def relax_dynamo(pipeline: tvm.transform.Pass | None = None): """A helper function to create a relax backend. Parameters ---------- pipeline : Optional[tvm.transform.Pass] The pipeline to be applied to the relax module before sent to build. Returns ------- backend : Callable[[torch.fx.GraphModule, List[torch.Tensor]], Callable] The relax dynamo backend. """ def _relax_backend(graph_module, example_inputs): import torch # type: ignore[import] assert isinstance(graph_module, torch.fx.GraphModule) def to_torch_tensor(nd_tensor): """A helper function to transfer a Tensor to torch.tensor.""" if isinstance(nd_tensor, torch.Tensor): # tvm-ffi #517 (Recursive DLPack container conversion) auto-converts # ffi::Tensor items returned in containers back to torch.Tensor when # the call site passed torch.Tensor inputs. return nd_tensor if isinstance(nd_tensor, tvm.runtime.Tensor): return torch.from_numpy(nd_tensor.numpy()) elif isinstance(nd_tensor, tvm_ffi.Array): return tuple(to_torch_tensor(x) for x in nd_tensor) else: raise ValueError(f"Unsupported type {type(nd_tensor)}") graph_module.graph.eliminate_dead_code() device = device_from_inputs(example_inputs) assert len(example_inputs) fake_inputs = [] if isinstance(example_inputs[0], torch._subclasses.fake_tensor.FakeTensor): # Fake tensors fake_inputs = example_inputs else: # Real tensors for node in graph_module.graph.nodes: if node.op != "placeholder": continue if "grapharg" not in node.meta: continue fake_tensor = node.meta["grapharg"].fake_tensor if fake_tensor is None: continue fake_inputs.append(fake_tensor) input_info = [] shape_vars = {} for tensor in fake_inputs: shape = [] for s in tensor.shape: if isinstance(s, torch.SymInt): if str(s) not in shape_vars: shape_vars[str(s)] = tvm.tirx.Var(str(s), "int64") shape.append(shape_vars[str(s)]) else: shape.append(s) input_info.append((shape, tensor.dtype)) mod = from_fx(graph_module, input_info) if device.type == "cuda": dev = tvm.cuda(device.index) target = tvm.target.Target("cuda") else: dev = tvm.cpu(0) target = tvm.target.Target(llvm_target()) # invoke optimization pipeline. if pipeline is None: # get default pipeline seq = tvm.relax.get_pipeline() elif isinstance(pipeline, str): # lookup by name seq = tvm.relax.get_pipeline(pipeline) else: seq = pipeline mod = mod.with_attr("target", target) mod = seq(mod) ex = relax_build(mod, target=target) vm = tvm.relax.VirtualMachine(ex.mod, device=dev) def exec_tvm(*i_args): args = [a.contiguous() for a in i_args if isinstance(a, torch.Tensor)] vm_args = list() for arg in args: if arg.requires_grad: arg = arg.detach() if isinstance(arg, torch._subclasses.fake_tensor.FakeTensor): # Materialize a real (eager) Tensor arg = torch.randn(arg.shape, dtype=arg.dtype, device=device) vm_args.append(arg) outputs = vm["main"](*vm_args) return to_torch_tensor(outputs) return exec_tvm return _relax_backend def dynamo_capture_subgraphs(model, *params, **kwargs) -> tvm.IRModule: """Capture subgraphs of the PyTorch model using torch.compile into an IRModule. Parameters ---------- model : torch.nn.Module The PyTorch model to be captured. params : List[torch.Tensor] The parameters of the PyTorch model. keep_params_as_input : bool Whether to keep model parameters as input variables of the captured Relax functions. Returns ------- output : ImporterOutput The output of translation, including the translated IRModule. If `keep_params_as_input` is true, the functions in the IRModule have an attribute "params" that contains the weights of the input model. The weights can be detached by `relax.frontend.detach_params`. """ import torch # type: ignore[import] from torch import _dynamo as dynamo # type: ignore[import] from torch import fx # type: ignore[import] keep_params_as_input = "keep_params_as_input" in kwargs and kwargs["keep_params_as_input"] kwargs.pop("keep_params_as_input", None) mod = tvm.IRModule() def _capture(graph_module: fx.GraphModule, example_inputs): assert isinstance(graph_module, torch.fx.GraphModule) input_info = [(tuple(tensor.shape), str(tensor.dtype)) for tensor in example_inputs] mod_ = from_fx( graph_module, input_info, keep_params_as_input=keep_params_as_input, unwrap_unit_return_tuple=True, ) new_name = f"subgraph_{len(mod.get_global_vars())}" mod[new_name] = mod_["main"].with_attr("global_symbol", new_name) return graph_module.forward dynamo.reset() compiled_model = torch.compile(model, backend=_capture) with torch.no_grad(): compiled_model(*params, **kwargs) return mod @functools.lru_cache(None) def llvm_target(): import platform import subprocess AVX512_TARGET = {"kind": "llvm", "mcpu": "skylake-avx512"} AVX2_TARGET = {"kind": "llvm", "mcpu": "core-avx2"} DEFAULT_TARGET = "llvm" system = platform.system() if system == "Linux": try: with open("/proc/cpuinfo") as f: cpuinfo = f.read() if "avx512" in cpuinfo: return AVX512_TARGET return AVX2_TARGET except FileNotFoundError: pass elif system == "Darwin": try: result = subprocess.run( ["sysctl", "-n", "machdep.cpu.features"], capture_output=True, text=True, check=False, ) if result.returncode == 0: cpu_features = result.stdout.lower() if "avx512" in cpu_features: return AVX512_TARGET if "avx2" in cpu_features: return AVX2_TARGET except (FileNotFoundError, subprocess.SubprocessError): pass if platform.machine() == "arm64": return DEFAULT_TARGET # Default fallback return DEFAULT_TARGET