wehub-resource-sync
1213 lines
52 KiB
Python
1213 lines
52 KiB
Python
# 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.
|
|
# ruff: noqa: F401
|
|
"""Relax to Python Function Converter.
|
|
|
|
This module provides functionality to convert Relax functions to Python functions
|
|
that can be executed directly in Python/PyTorch environment.
|
|
"""
|
|
|
|
import traceback
|
|
from typing import Any, Optional, Union
|
|
|
|
import numpy # pylint: disable=unused-import
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import tvm
|
|
from tvm import relax, runtime
|
|
from tvm.ir import IRModule, Op
|
|
|
|
|
|
class RelaxToPyFuncConverter:
|
|
"""Converter that works with IRModule to convert Relax functions to Python functions.
|
|
|
|
This converter transforms Relax functions into Python functions that can be executed
|
|
directly in Python/PyTorch environment. The conversion maps Relax operators to
|
|
corresponding PyTorch APIs and handles special cases like call_tir and call_dps_packed.
|
|
"""
|
|
|
|
def __init__(self, ir_module: IRModule):
|
|
"""Initialize the converter with an IRModule.
|
|
|
|
Args:
|
|
ir_module: The IRModule containing Relax functions to convert
|
|
"""
|
|
self.ir_module = ir_module
|
|
self.operator_map = self._get_op_map()
|
|
# Cache for RelaxExpressionConverter instances to avoid recreating them
|
|
self._converter_cache = {}
|
|
# Cache for operator mappings to avoid repeated lookups
|
|
self._op_cache = {}
|
|
|
|
def _create_fallback_tensor(
|
|
self, shape_hint: list[int] | None = None, dtype: str = "float32"
|
|
) -> torch.Tensor:
|
|
"""Create a fallback tensor with reasonable default shape."""
|
|
if shape_hint:
|
|
# Use the provided shape hint
|
|
return torch.zeros(shape_hint, dtype=getattr(torch, dtype))
|
|
else:
|
|
# Use a small default shape
|
|
return torch.zeros(1, dtype=getattr(torch, dtype))
|
|
|
|
def convert(self, relax_function_names: str | list[str]) -> IRModule:
|
|
"""Convert specified Relax functions to Python functions.
|
|
|
|
Args:
|
|
relax_function_names: Name(s) of Relax functions to convert
|
|
|
|
Returns:
|
|
Updated IRModule with converted Python functions stored in pyfuncs
|
|
|
|
Example:
|
|
>>> converter = RelaxToPyFuncConverter(ir_mod)
|
|
>>> # Convert a single function
|
|
>>> converted_ir_mod = converter.convert("my_relax_func")
|
|
>>> # Convert multiple functions
|
|
>>> converted_ir_mod = converter.convert(["func1", "func2"])
|
|
"""
|
|
if isinstance(relax_function_names, str):
|
|
relax_function_names = [relax_function_names]
|
|
|
|
# Create a copy of the current IRModule
|
|
new_ir_mod = self.ir_module.clone()
|
|
|
|
# Initialize pyfuncs if not exists
|
|
if not hasattr(new_ir_mod, "pyfuncs"):
|
|
new_ir_mod.pyfuncs = {}
|
|
|
|
# Get Relax function names from IRModule
|
|
relax_func_names = []
|
|
for global_var, func in self.ir_module.functions_items():
|
|
if isinstance(func, relax.Function):
|
|
relax_func_names.append(global_var.name_hint)
|
|
|
|
# Convert each Relax function
|
|
for func_name in relax_function_names:
|
|
if func_name not in relax_func_names:
|
|
raise ValueError(f"Relax function '{func_name}' not found in IRModule")
|
|
|
|
# Get the Relax function
|
|
relax_func = None
|
|
for global_var, func in self.ir_module.functions_items():
|
|
if global_var.name_hint == func_name and isinstance(func, relax.Function):
|
|
relax_func = func
|
|
break
|
|
|
|
if relax_func is None:
|
|
raise ValueError(f"Could not find Relax function '{func_name}'")
|
|
|
|
# Convert to Python function
|
|
py_func = self._convert_relax_func_to_python(relax_func, func_name)
|
|
|
|
# Store in pyfuncs
|
|
new_ir_mod.pyfuncs[func_name] = py_func
|
|
|
|
return new_ir_mod
|
|
|
|
def _convert_relax_func_to_python(self, relax_func: relax.Function, func_name: str) -> callable:
|
|
"""Convert a single Relax function to a Python function with caching."""
|
|
# Get function parameters
|
|
params = relax_func.params
|
|
|
|
# Create the Python function
|
|
def converted_function(*args, **_kwargs):
|
|
"""Converted Python function from Relax function."""
|
|
# Handle arguments
|
|
if len(args) != len(params):
|
|
raise ValueError(f"Expected {len(params)} arguments, got {len(args)}")
|
|
|
|
# Use cached converter or create new one
|
|
if func_name not in self._converter_cache:
|
|
self._converter_cache[func_name] = RelaxExpressionConverter(
|
|
self.operator_map, self.ir_module, self._op_cache
|
|
)
|
|
|
|
# Execute the converted function body
|
|
converter = self._converter_cache[func_name]
|
|
converter.current_params = params
|
|
return converter.convert_expr(relax_func.body, args)
|
|
|
|
# Set function metadata
|
|
converted_function.__name__ = func_name
|
|
converted_function.__doc__ = f"Converted Python function from Relax function: {func_name}"
|
|
|
|
return converted_function
|
|
|
|
@staticmethod
|
|
def _get_op_map() -> dict[str, str]:
|
|
"""Get the mapping from Relax operators to PyTorch operators."""
|
|
return {
|
|
# Binary operations
|
|
"relax.add": "torch.add",
|
|
"relax.subtract": "torch.sub",
|
|
"relax.multiply": "torch.mul",
|
|
"relax.divide": "torch.div",
|
|
"relax.power": "torch.pow",
|
|
"relax.maximum": "torch.maximum",
|
|
"relax.minimum": "torch.minimum",
|
|
"relax.floor_divide": "torch.floor_divide",
|
|
"relax.mod": "torch.fmod",
|
|
"relax.floor_mod": "torch.remainder",
|
|
"relax.log_add_exp": "torch.logaddexp",
|
|
# Bitwise operations
|
|
"relax.bitwise_and": "torch.bitwise_and",
|
|
"relax.bitwise_or": "torch.bitwise_or",
|
|
"relax.bitwise_xor": "torch.bitwise_xor",
|
|
"relax.left_shift": "torch.left_shift",
|
|
"relax.right_shift": "torch.right_shift",
|
|
# Unary operations
|
|
"relax.abs": "torch.abs",
|
|
"relax.negative": "torch.neg",
|
|
"relax.exp": "torch.exp",
|
|
"relax.log": "torch.log",
|
|
"relax.sqrt": "torch.sqrt",
|
|
"relax.rsqrt": "torch.rsqrt",
|
|
"relax.sin": "torch.sin",
|
|
"relax.cos": "torch.cos",
|
|
"relax.tanh": "torch.tanh",
|
|
"relax.sigmoid": "torch.sigmoid",
|
|
"relax.square": "torch.square",
|
|
"relax.sign": "torch.sign",
|
|
"relax.floor": "torch.floor",
|
|
"relax.ceil": "torch.ceil",
|
|
"relax.round": "torch.round",
|
|
"relax.trunc": "torch.trunc",
|
|
"relax.clip": "torch.clamp",
|
|
"relax.bitwise_not": "torch.bitwise_not",
|
|
# Trigonometric functions
|
|
"relax.acos": "torch.acos",
|
|
"relax.asin": "torch.asin",
|
|
"relax.atan": "torch.atan",
|
|
"relax.cosh": "torch.cosh",
|
|
"relax.sinh": "torch.sinh",
|
|
"relax.tan": "torch.tan",
|
|
"relax.acosh": "torch.acosh",
|
|
"relax.asinh": "torch.asinh",
|
|
"relax.atanh": "torch.atanh",
|
|
# Special functions
|
|
"relax.erf": "torch.erf",
|
|
"relax.isfinite": "torch.isfinite",
|
|
"relax.isinf": "torch.isinf",
|
|
"relax.isnan": "torch.isnan",
|
|
# Neural network operations
|
|
"relax.nn.relu": "F.relu",
|
|
"relax.nn.relu6": "F.relu6",
|
|
"relax.nn.gelu": "F.gelu",
|
|
"relax.nn.gelu_tanh": "F.gelu",
|
|
"relax.nn.softmax": "F.softmax",
|
|
"relax.nn.log_softmax": "F.log_softmax",
|
|
"relax.nn.dropout": "F.dropout",
|
|
"relax.nn.batch_norm": "F.batch_norm",
|
|
"relax.nn.layer_norm": "F.layer_norm",
|
|
"relax.nn.group_norm": "F.group_norm",
|
|
"relax.nn.instance_norm": "F.instance_norm",
|
|
"relax.nn.rms_norm": "F.layer_norm", # Approximate mapping
|
|
"relax.nn.linear": "F.linear",
|
|
"relax.nn.conv1d": "F.conv1d",
|
|
"relax.nn.conv2d": "F.conv2d",
|
|
"relax.nn.conv3d": "F.conv3d",
|
|
"relax.nn.conv1d_transpose": "F.conv_transpose1d",
|
|
"relax.nn.conv2d_transpose": "F.conv_transpose2d",
|
|
"relax.nn.conv3d_transpose": "F.conv_transpose3d",
|
|
"relax.nn.max_pool1d": "F.max_pool1d",
|
|
"relax.nn.max_pool2d": "F.max_pool2d",
|
|
"relax.nn.max_pool3d": "F.max_pool3d",
|
|
"relax.nn.avg_pool1d": "F.avg_pool1d",
|
|
"relax.nn.avg_pool2d": "F.avg_pool2d",
|
|
"relax.nn.avg_pool3d": "F.avg_pool3d",
|
|
"relax.nn.adaptive_avg_pool1d": "F.adaptive_avg_pool1d",
|
|
"relax.nn.adaptive_avg_pool2d": "F.adaptive_avg_pool2d",
|
|
"relax.nn.adaptive_avg_pool3d": "F.adaptive_avg_pool3d",
|
|
"relax.nn.leakyrelu": "F.leaky_relu",
|
|
"relax.nn.prelu": "F.prelu",
|
|
"relax.nn.selu": "F.selu",
|
|
"relax.nn.silu": "F.silu",
|
|
"relax.nn.softplus": "F.softplus",
|
|
"relax.nn.attention": "F.scaled_dot_product_attention", # Approximate mapping
|
|
"relax.nn.cross_entropy_with_logits": "F.cross_entropy",
|
|
"relax.nn.nll_loss": "F.nll_loss",
|
|
"relax.nn.pad": "F.pad",
|
|
"relax.nn.pixel_shuffle": "F.pixel_shuffle",
|
|
# Tensor operations
|
|
"relax.matmul": "torch.matmul",
|
|
"relax.linear": "F.linear",
|
|
"relax.einsum": "torch.einsum",
|
|
"relax.outer": "torch.outer",
|
|
"relax.reshape": "reshape", # Special handling needed
|
|
"relax.permute_dims": "permute_dims", # Special handling needed
|
|
"relax.expand_dims": "expand_dims", # Special handling needed
|
|
"relax.squeeze": "squeeze", # Special handling needed
|
|
"relax.concat": "concat", # Special handling needed
|
|
"relax.split": "split", # Special handling needed
|
|
"relax.stack": "stack", # Special handling needed
|
|
"relax.tile": "tile", # Special handling needed
|
|
"relax.repeat": "repeat", # Special handling needed
|
|
"relax.broadcast_to": "torch.broadcast_to",
|
|
"relax.flatten": "torch.flatten",
|
|
"relax.flip": "flip", # Special handling needed
|
|
"relax.roll": "torch.roll",
|
|
"relax.rot90": "torch.rot90",
|
|
"relax.meshgrid": "torch.meshgrid",
|
|
"relax.one_hot": "F.one_hot",
|
|
"relax.layout_transform": "torch.permute", # Approximate mapping
|
|
# Indexing operations
|
|
"relax.take": "take", # Special handling needed
|
|
"relax.gather_elements": "torch.gather",
|
|
"relax.gather_nd": "torch.gather",
|
|
"relax.scatter_elements": "torch.scatter",
|
|
"relax.scatter_nd": "torch.scatter",
|
|
"relax.index_put": "torch.index_put",
|
|
"relax.index_tensor": "torch.index_select",
|
|
"relax.strided_slice": "torch.slice",
|
|
"relax.dynamic_strided_slice": "torch.slice",
|
|
"relax.slice_scatter": "torch.scatter",
|
|
# Reduction operations
|
|
"relax.sum": "sum", # Special handling needed
|
|
"relax.mean": "mean", # Special handling needed
|
|
"relax.max": "max", # Special handling needed
|
|
"relax.min": "min", # Special handling needed
|
|
"relax.prod": "torch.prod",
|
|
"relax.std": "std", # Special handling needed
|
|
"relax.variance": "variance", # Special handling needed
|
|
"relax.cumsum": "torch.cumsum",
|
|
"relax.cumprod": "torch.cumprod",
|
|
"relax.argmax": "torch.argmax",
|
|
"relax.argmin": "torch.argmin",
|
|
# Comparison operations
|
|
"relax.equal": "torch.eq",
|
|
"relax.not_equal": "torch.ne",
|
|
"relax.greater": "torch.gt",
|
|
"relax.greater_equal": "torch.ge",
|
|
"relax.less": "torch.lt",
|
|
"relax.less_equal": "torch.le",
|
|
# Logical operations
|
|
"relax.logical_and": "torch.logical_and",
|
|
"relax.logical_or": "torch.logical_or",
|
|
"relax.logical_not": "torch.logical_not",
|
|
"relax.logical_xor": "torch.logical_xor",
|
|
# Creation operations
|
|
"relax.zeros": "torch.zeros",
|
|
"relax.ones": "torch.ones",
|
|
"relax.full": "torch.full",
|
|
"relax.full_like": "torch.full_like",
|
|
"relax.zeros_like": "torch.zeros_like",
|
|
"relax.ones_like": "torch.ones_like",
|
|
"relax.arange": "torch.arange",
|
|
"relax.eye": "torch.eye",
|
|
"relax.eye_like": "torch.eye",
|
|
"relax.tril": "torch.tril",
|
|
"relax.triu": "torch.triu",
|
|
"relax.hamming_window": "torch.hamming_window",
|
|
# Search operations
|
|
"relax.where": "torch.where",
|
|
"relax.bucketize": "torch.bucketize",
|
|
"relax.nonzero": "torch.nonzero",
|
|
"relax.unique": "torch.unique",
|
|
# Sorting operations
|
|
"relax.sort": "torch.sort",
|
|
"relax.argsort": "torch.argsort",
|
|
"relax.topk": "torch.topk",
|
|
# Sampling operations
|
|
"relax.multinomial_from_uniform": "torch.multinomial",
|
|
# Ternary operations
|
|
"relax.ewise_fma": "torch.fma", # Approximate mapping
|
|
# Data type operations
|
|
"relax.astype": "torch.to",
|
|
"relax.wrap_param": "torch.tensor",
|
|
# Mask operations
|
|
"relax.masked_fill": "torch.masked_fill",
|
|
# Quantization operations
|
|
"relax.quantize": "torch.quantize_per_tensor", # Approximate mapping
|
|
"relax.dequantize": "torch.dequantize", # Approximate mapping
|
|
# Special operations (handled separately)
|
|
"relax.call_tir": "call_tir",
|
|
"relax.call_tir_inplace": "call_tir_inplace",
|
|
"relax.call_dps_packed": "call_dps_packed",
|
|
"relax.call_pure_packed": "call_pure_packed",
|
|
"relax.call_tir_with_grad": "call_tir_with_grad",
|
|
"relax.call_builtin_with_ctx": "call_builtin_with_ctx",
|
|
"relax.call_inplace_packed": "call_inplace_packed",
|
|
"relax.invoke_closure": "invoke_closure",
|
|
"relax.invoke_pure_closure": "invoke_pure_closure",
|
|
"relax.make_closure": "make_closure",
|
|
"relax.null_value": "null_value",
|
|
"relax.print": "print",
|
|
"relax.shape_of": "shape_of",
|
|
"relax.shape_to_tensor": "shape_to_tensor",
|
|
"relax.tensor_to_shape": "tensor_to_shape",
|
|
"relax.to_vdevice": "to_vdevice",
|
|
"relax.hint_on_device": "hint_on_device",
|
|
"relax.assert_op": "assert_op",
|
|
}
|
|
|
|
|
|
class RelaxExpressionConverter:
|
|
"""Converter that transforms Relax expressions to Python/PyTorch code."""
|
|
|
|
def __init__(
|
|
self,
|
|
operator_map: dict[str, str],
|
|
ir_module: IRModule = None,
|
|
op_cache: dict[str, str] | None = None,
|
|
):
|
|
"""Initialize the expression converter.
|
|
|
|
Args:
|
|
operator_map: Mapping from Relax operators to PyTorch operators
|
|
ir_module: The IRModule containing TIR functions to compile
|
|
op_cache: Shared cache for operator mappings to avoid repeated lookups
|
|
"""
|
|
self.operator_map = operator_map
|
|
self.variable_map: dict[str, Any] = {}
|
|
self.current_params: list[relax.Var] = []
|
|
self.ir_module = ir_module
|
|
# Use shared operator cache or create new one
|
|
self._op_cache = op_cache if op_cache is not None else {}
|
|
|
|
def _create_fallback_tensor(
|
|
self, shape_hint: list[int] | None = None, dtype: str = "float32"
|
|
) -> torch.Tensor:
|
|
"""Create a fallback tensor with reasonable default shape."""
|
|
if shape_hint:
|
|
return torch.zeros(shape_hint, dtype=getattr(torch, dtype))
|
|
else:
|
|
return torch.zeros(1, dtype=getattr(torch, dtype))
|
|
|
|
def convert_expr(self, expr: relax.Expr, args: list[Any]) -> Any:
|
|
"""Convert a Relax expression to Python/PyTorch equivalent."""
|
|
if isinstance(expr, relax.Var):
|
|
return self._convert_var(expr, args)
|
|
elif isinstance(expr, relax.Call):
|
|
return self._convert_call(expr, args)
|
|
elif isinstance(expr, relax.Constant):
|
|
return self._convert_constant(expr)
|
|
elif isinstance(expr, relax.SeqExpr):
|
|
return self._convert_seq_expr(expr, args)
|
|
elif isinstance(expr, relax.Tuple):
|
|
return self._convert_tuple(expr, args)
|
|
elif isinstance(expr, relax.TupleGetItem):
|
|
return self._convert_tuple_get_item(expr, args)
|
|
elif isinstance(expr, relax.If):
|
|
return self._convert_if(expr, args)
|
|
elif isinstance(expr, relax.ShapeExpr):
|
|
return self._convert_shape_expr(expr)
|
|
else:
|
|
# Fallback for unknown expression types
|
|
return f"<unknown_expr: {type(expr).__name__}>"
|
|
|
|
def _convert_var(self, var: relax.Var, args: list[Any]) -> Any:
|
|
"""Convert a Relax variable to Python equivalent."""
|
|
if hasattr(var, "name_hint"):
|
|
var_name = var.name_hint
|
|
|
|
# Check if it's a function parameter
|
|
for i, param in enumerate(self.current_params):
|
|
if hasattr(param, "name_hint") and param.name_hint == var_name:
|
|
return args[i]
|
|
|
|
# Check if it's a bound variable
|
|
if var_name in self.variable_map:
|
|
return self.variable_map[var_name]
|
|
|
|
# Try to infer shape from var's type annotation
|
|
if hasattr(var, "ty") and hasattr(var.ty, "shape"):
|
|
shape = var.ty.shape
|
|
if shape and len(shape) > 0:
|
|
# Convert symbolic shapes to concrete values
|
|
concrete_shape = []
|
|
for dim in shape:
|
|
if isinstance(dim, int):
|
|
concrete_shape.append(dim)
|
|
else:
|
|
# For symbolic dimensions, use a reasonable default
|
|
concrete_shape.append(1)
|
|
return torch.zeros(concrete_shape, dtype=torch.float32)
|
|
|
|
if args and isinstance(args[0], torch.Tensor):
|
|
return torch.zeros_like(args[0])
|
|
# Use fallback tensor with shape inference
|
|
return self._create_fallback_tensor()
|
|
return self._create_fallback_tensor()
|
|
|
|
def _convert_call(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert a Relax call to Python/PyTorch equivalent."""
|
|
op = call.op
|
|
|
|
# Handle different types of calls
|
|
if isinstance(op, relax.GlobalVar):
|
|
# Function call
|
|
return self._convert_function_call(call, args)
|
|
elif isinstance(op, Op):
|
|
# Operator call
|
|
return self._convert_operator_call(call, args)
|
|
elif isinstance(op, relax.ExternFunc):
|
|
# External function call (like call_tir, call_dps_packed)
|
|
return self._convert_extern_func_call(call, args)
|
|
else:
|
|
return self._create_fallback_tensor()
|
|
|
|
def _convert_function_call(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert a Relax function call."""
|
|
func_name = call.op.name_hint
|
|
call_args = [self.convert_expr(arg, args) for arg in call.args]
|
|
|
|
# Handle special cases
|
|
if func_name in ["call_tir", "call_tir_inplace"]:
|
|
return self._convert_call_tir(call, args)
|
|
elif func_name in ["call_dps_packed", "call_pure_packed"]:
|
|
return self._convert_call_dps_packed(call, args)
|
|
else:
|
|
# Regular function call - return first argument as fallback
|
|
return call_args[0] if call_args else self._create_fallback_tensor()
|
|
|
|
def _convert_operator_call(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert a Relax operator call to PyTorch equivalent."""
|
|
op_name = call.op.name
|
|
call_args = [self.convert_expr(arg, args) for arg in call.args]
|
|
|
|
# Use cached operator mapping or look it up
|
|
if op_name not in self._op_cache:
|
|
self._op_cache[op_name] = self.operator_map.get(op_name)
|
|
pytorch_op = self._op_cache[op_name]
|
|
if pytorch_op:
|
|
try:
|
|
# Handle special operations
|
|
if pytorch_op == "call_tir":
|
|
return self._convert_call_tir(call, args)
|
|
elif pytorch_op == "call_tir_inplace":
|
|
return self._convert_call_tir(call, args)
|
|
elif pytorch_op == "call_dps_packed":
|
|
return self._convert_call_dps_packed(call, args)
|
|
elif pytorch_op == "call_pure_packed":
|
|
return self._convert_call_dps_packed(call, args)
|
|
elif pytorch_op == "expand_dims":
|
|
return self._convert_expand_dims(call, args)
|
|
elif pytorch_op in ["sum", "mean", "max", "min", "std", "variance"]:
|
|
return self._convert_reduction_op(call, args, pytorch_op)
|
|
elif pytorch_op == "squeeze":
|
|
return self._convert_squeeze(call, args)
|
|
elif pytorch_op in ["concat", "split", "stack"]:
|
|
return self._convert_tensor_ops(call, args, pytorch_op)
|
|
elif pytorch_op == "reshape":
|
|
return self._convert_reshape(call, args)
|
|
elif pytorch_op == "permute_dims":
|
|
return self._convert_permute_dims(call, args)
|
|
elif pytorch_op == "take":
|
|
return self._convert_take(call, args)
|
|
elif pytorch_op == "flip":
|
|
return self._convert_flip(call, args)
|
|
elif pytorch_op == "tile":
|
|
return self._convert_tile(call, args)
|
|
elif pytorch_op == "repeat":
|
|
return self._convert_repeat(call, args)
|
|
# Handle special cases for PyTorch operations
|
|
elif pytorch_op.startswith("F."):
|
|
return self._handle_functional_operation(pytorch_op, call, call_args)
|
|
elif pytorch_op.startswith("torch."):
|
|
# Regular PyTorch operation
|
|
func_name = pytorch_op[6:] # Remove "torch." prefix
|
|
func = getattr(torch, func_name)
|
|
return func(*call_args)
|
|
else:
|
|
# Direct function reference - use getattr for safer access
|
|
if pytorch_op.startswith("torch."):
|
|
module = torch
|
|
func_name = pytorch_op[6:] # Remove "torch." prefix
|
|
elif pytorch_op.startswith("F."):
|
|
module = F
|
|
func_name = pytorch_op[2:] # Remove "F." prefix
|
|
else:
|
|
return (
|
|
f"<exec_error: {pytorch_op}({', '.join(map(str, call_args))}) "
|
|
f"- unsupported operation>"
|
|
)
|
|
|
|
func = getattr(module, func_name, None)
|
|
if func is None:
|
|
return (
|
|
f"<exec_error: {pytorch_op}({', '.join(map(str, call_args))}) "
|
|
f"- function not found>"
|
|
)
|
|
return func(*call_args)
|
|
except (AttributeError, TypeError, ValueError) as error:
|
|
# This allows the test framework to catch and handle the errors appropriately
|
|
if pytorch_op.startswith("torch.") or pytorch_op.startswith("F."):
|
|
raise error
|
|
# Fallback to string representation for non-PyTorch operations
|
|
return f"<exec_error: {pytorch_op}({', '.join(map(str, call_args))}) - {error}>"
|
|
else:
|
|
# Unknown operator
|
|
return f"<unknown_op: {op_name}({', '.join(map(str, call_args))})>"
|
|
|
|
def _handle_functional_operation(
|
|
self, pytorch_op: str, call: relax.Call, call_args: list[Any]
|
|
) -> Any:
|
|
"""Handle PyTorch functional operations with special parameter handling."""
|
|
# Neural network function
|
|
func_name = pytorch_op[2:] # Remove "F." prefix
|
|
func = getattr(F, func_name)
|
|
|
|
# Special handling for functions that need dim parameter
|
|
if func_name in ["softmax", "log_softmax"]:
|
|
# Extract axis from call.attrs and convert to dim
|
|
axis = None
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
if axis is not None:
|
|
return func(call_args[0], dim=axis)
|
|
else:
|
|
# Default to last dimension if no axis specified
|
|
return func(call_args[0], dim=-1)
|
|
else:
|
|
return func(*call_args)
|
|
|
|
def _convert_extern_func_call(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert an external function call."""
|
|
func_name = call.op.global_symbol
|
|
call_args = [self.convert_expr(arg, args) for arg in call.args]
|
|
|
|
if func_name in ["call_tir", "call_tir_inplace"]:
|
|
return self._convert_call_tir(call, args)
|
|
elif func_name in ["call_dps_packed", "call_pure_packed"]:
|
|
return self._convert_call_dps_packed(call, args)
|
|
else:
|
|
return call_args[0] if call_args else self._create_fallback_tensor()
|
|
|
|
def _convert_call_tir(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert call_tir to Python equivalent with DLPack conversion."""
|
|
# Extract TIR function name and arguments
|
|
tir_func = call.args[0]
|
|
tir_args = call.args[1] if len(call.args) > 1 else []
|
|
out_ty = call.attrs.get("out_ty") if call.attrs else None
|
|
|
|
# Get function name
|
|
if isinstance(tir_func, relax.GlobalVar):
|
|
func_name = tir_func.name_hint
|
|
else:
|
|
# Convert the GlobalVar expression
|
|
func_name = self.convert_expr(tir_func, args)
|
|
if isinstance(func_name, str) and func_name.startswith("<"):
|
|
# If it's a placeholder, extract the name
|
|
func_name = str(tir_func)
|
|
|
|
# Convert arguments to PyTorch tensors
|
|
converted_args = [self.convert_expr(arg, args) for arg in tir_args]
|
|
|
|
try:
|
|
# First, try to get the TIR function from the current IRModule
|
|
tir_function = None
|
|
if self.ir_module:
|
|
# Look for the TIR function in the current IRModule
|
|
for global_var, func in self.ir_module.functions.items():
|
|
if global_var.name_hint == func_name and hasattr(func, "body"):
|
|
try:
|
|
# Compile the TIR function
|
|
target = tvm.target.Target("llvm")
|
|
with tvm.target.Target(target):
|
|
tir_function = tvm.compile(func, target=target)
|
|
break
|
|
except (RuntimeError, ValueError, TypeError) as compile_e:
|
|
print(
|
|
f"Warning: Failed to compile TIR function {func_name}: {compile_e}"
|
|
)
|
|
continue
|
|
|
|
# If not found in current module, try global registry
|
|
if tir_function is None:
|
|
tir_function = tvm.get_global_func(func_name)
|
|
|
|
if tir_function is None:
|
|
if len(converted_args) >= 2:
|
|
# Simple fallback: just add the tensors
|
|
return torch.add(converted_args[0], converted_args[1])
|
|
else:
|
|
return converted_args[0] if converted_args else torch.tensor([])
|
|
|
|
# Convert PyTorch tensors to TVM NDArrays via DLPack
|
|
tvm_args = []
|
|
for arg in converted_args:
|
|
try:
|
|
if isinstance(arg, torch.Tensor):
|
|
# Convert PyTorch tensor to TVM NDArray via DLPack
|
|
tvm_arg = runtime.from_dlpack(torch.to_dlpack(arg))
|
|
tvm_args.append(tvm_arg)
|
|
else:
|
|
tvm_args.append(arg)
|
|
except (AttributeError, TypeError, ValueError):
|
|
traceback.print_exc()
|
|
tvm_args.append(arg)
|
|
|
|
# For call_tir, we need to allocate output tensor
|
|
output_shape = None
|
|
if out_ty and hasattr(out_ty, "shape"):
|
|
output_shape = out_ty.shape
|
|
elif converted_args:
|
|
# Use the shape of the first input tensor
|
|
first_arg = converted_args[0]
|
|
if isinstance(first_arg, torch.Tensor):
|
|
output_shape = first_arg.shape
|
|
|
|
if output_shape is None:
|
|
if converted_args and isinstance(converted_args[0], torch.Tensor):
|
|
output_shape = converted_args[0].shape
|
|
else:
|
|
output_shape = (1,) # Default shape
|
|
|
|
# Allocate output tensor
|
|
output_tensor = runtime.empty(output_shape, dtype="float32")
|
|
tvm_args.append(output_tensor)
|
|
|
|
# Call the TIR function
|
|
try:
|
|
tir_function(*tvm_args)
|
|
# The result is in the output_tensor we allocated
|
|
# Convert result back to PyTorch tensor via DLPack
|
|
try:
|
|
result = torch.from_dlpack(output_tensor.to_dlpack())
|
|
return result
|
|
except AttributeError:
|
|
# Fallback: convert to numpy then to PyTorch
|
|
numpy_result = output_tensor.numpy()
|
|
result = torch.from_numpy(numpy_result)
|
|
return result
|
|
except (RuntimeError, ValueError, TypeError, AttributeError) as exc:
|
|
print(f"Warning: TIR function {func_name} execution failed: {exc}")
|
|
traceback.print_exc()
|
|
# Fallback to simple addition
|
|
if len(converted_args) >= 2:
|
|
return torch.add(converted_args[0], converted_args[1])
|
|
else:
|
|
return converted_args[0] if converted_args else torch.tensor([])
|
|
|
|
except (RuntimeError, ValueError, TypeError):
|
|
traceback.print_exc()
|
|
# Fallback implementation instead of error string
|
|
if len(converted_args) >= 2:
|
|
return torch.add(converted_args[0], converted_args[1])
|
|
else:
|
|
return converted_args[0] if converted_args else torch.tensor([])
|
|
|
|
def _convert_call_dps_packed(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert call_dps_packed to Python equivalent with DLPack conversion."""
|
|
# Extract packed function name and arguments
|
|
packed_func = call.args[0]
|
|
packed_args = call.args[1] if len(call.args) > 1 else []
|
|
_out_ty = call.attrs.get("out_ty") if call.attrs else None
|
|
|
|
# Get function name
|
|
if isinstance(packed_func, relax.GlobalVar):
|
|
func_name = packed_func.name_hint
|
|
elif isinstance(packed_func, relax.ExternFunc):
|
|
func_name = packed_func.global_symbol
|
|
else:
|
|
func_name = str(packed_func)
|
|
|
|
# Convert arguments to PyTorch tensors
|
|
converted_args = []
|
|
for arg in packed_args:
|
|
converted_arg = self.convert_expr(arg, args)
|
|
if isinstance(converted_arg, str) and converted_arg.startswith("<"):
|
|
# Handle Expr and other special cases
|
|
if "Expr" in converted_arg:
|
|
# Extract the value from Expr
|
|
try:
|
|
# Try to get the actual value from the Expr
|
|
if hasattr(arg, "value"):
|
|
converted_arg = arg.value
|
|
else:
|
|
converted_arg = 0.0 # Default value
|
|
except (AttributeError, ValueError, TypeError):
|
|
converted_arg = 0.0
|
|
else:
|
|
converted_arg = torch.tensor([]) # Fallback
|
|
converted_args.append(converted_arg)
|
|
|
|
try:
|
|
# Get the packed function from TVM
|
|
packed_function = tvm.get_global_func(func_name)
|
|
if packed_function is None:
|
|
return converted_args[0] if converted_args else torch.tensor([])
|
|
|
|
# Convert PyTorch tensors to TVM NDArrays via DLPack
|
|
tvm_args = []
|
|
for arg in converted_args:
|
|
if isinstance(arg, torch.Tensor):
|
|
# Convert PyTorch tensor to TVM NDArray via DLPack
|
|
tvm_arg = runtime.from_dlpack(torch.to_dlpack(arg))
|
|
tvm_args.append(tvm_arg)
|
|
else:
|
|
tvm_args.append(arg)
|
|
|
|
# Call the packed function
|
|
result = packed_function(*tvm_args)
|
|
|
|
# Convert result back to PyTorch tensor via DLPack
|
|
if isinstance(result, runtime.Tensor):
|
|
try:
|
|
pytorch_result = torch.from_dlpack(result.to_dlpack())
|
|
return pytorch_result
|
|
except AttributeError:
|
|
# Fallback: convert to numpy then to PyTorch
|
|
numpy_result = result.numpy()
|
|
pytorch_result = torch.from_numpy(numpy_result)
|
|
return pytorch_result
|
|
else:
|
|
return result
|
|
|
|
except (RuntimeError, ValueError, TypeError):
|
|
traceback.print_exc()
|
|
# Fallback: return the first argument
|
|
return converted_args[0] if converted_args else torch.tensor([])
|
|
|
|
def _convert_constant(self, const: relax.Constant) -> Any:
|
|
"""Convert a Relax constant to Python equivalent."""
|
|
if hasattr(const, "data"):
|
|
data = const.data
|
|
# Convert TVM NDArray to Python scalar if it's a scalar
|
|
if hasattr(data, "numpy"):
|
|
numpy_data = data.numpy()
|
|
if numpy_data.size == 1:
|
|
return float(numpy_data.item())
|
|
else:
|
|
# For multi-element arrays, convert to PyTorch tensor
|
|
return torch.from_numpy(numpy_data)
|
|
elif hasattr(data, "item"):
|
|
# Single element tensor
|
|
return data.item()
|
|
else:
|
|
return data
|
|
return self._create_fallback_tensor()
|
|
|
|
def _convert_seq_expr(self, seq: relax.SeqExpr, args: list[Any]) -> Any:
|
|
"""Convert a Relax sequence expression."""
|
|
# Convert blocks
|
|
for block in seq.blocks:
|
|
if hasattr(block, "bindings"):
|
|
for binding in block.bindings:
|
|
if isinstance(binding, relax.VarBinding):
|
|
var_name = binding.var.name_hint
|
|
value = self.convert_expr(binding.value, args)
|
|
self.variable_map[var_name] = value
|
|
|
|
# Convert body
|
|
return self.convert_expr(seq.body, args)
|
|
|
|
def _convert_tuple(self, tuple_expr: relax.Tuple, args: list[Any]) -> Any:
|
|
"""Convert a Relax tuple to Python tuple."""
|
|
elements = [self.convert_expr(elem, args) for elem in tuple_expr.fields]
|
|
return tuple(elements)
|
|
|
|
def _convert_tuple_get_item(self, get_item: relax.TupleGetItem, args: list[Any]) -> Any:
|
|
"""Convert a Relax tuple get item to Python equivalent."""
|
|
tuple_expr = self.convert_expr(get_item.tuple_value, args)
|
|
index = get_item.index
|
|
if isinstance(tuple_expr, torch.Tensor):
|
|
return tuple_expr[index] if index < len(tuple_expr) else self._create_fallback_tensor()
|
|
else:
|
|
return self._create_fallback_tensor()
|
|
|
|
def _convert_if(self, if_expr: relax.If, args: list[Any]) -> Any:
|
|
"""Convert a Relax if expression to Python equivalent."""
|
|
condition = self.convert_expr(if_expr.cond, args)
|
|
true_branch = self.convert_expr(if_expr.true_branch, args)
|
|
false_branch = self.convert_expr(if_expr.false_branch, args)
|
|
if isinstance(condition, torch.Tensor) and condition.item():
|
|
return (
|
|
true_branch
|
|
if isinstance(true_branch, torch.Tensor)
|
|
else self._create_fallback_tensor()
|
|
)
|
|
else:
|
|
return (
|
|
false_branch
|
|
if isinstance(false_branch, torch.Tensor)
|
|
else self._create_fallback_tensor()
|
|
)
|
|
|
|
def _convert_expand_dims(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert expand_dims to torch.unsqueeze with proper axis handling."""
|
|
if len(call.args) < 1:
|
|
return self._create_fallback_tensor()
|
|
|
|
# Convert the tensor argument
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Get the axis from call.attrs
|
|
axis = None
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
# Handle different types of axis
|
|
if hasattr(axis, "__iter__") and not isinstance(axis, str):
|
|
# It's an array/list, take the first element
|
|
axis = next(iter(axis)) if len(axis) > 0 else None
|
|
|
|
# Handle TVM types
|
|
if hasattr(axis, "value"):
|
|
# It's a TVM IntImm or similar, get the value
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
if axis is None:
|
|
return self._create_fallback_tensor()
|
|
|
|
# Use torch.unsqueeze with the correct axis
|
|
return torch.unsqueeze(tensor_arg, dim=axis)
|
|
|
|
def _convert_reduction_op(self, call: relax.Call, args: list[Any], op_name: str) -> Any:
|
|
"""Convert reduction operations with axis and keepdims parameters."""
|
|
if len(call.args) < 1:
|
|
return f"<{op_name}_error: insufficient arguments>"
|
|
|
|
# Convert the tensor argument
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Get axis and keepdims from call.attrs
|
|
axis = None
|
|
keepdims = False
|
|
|
|
if call.attrs:
|
|
if hasattr(call.attrs, "axis") and call.attrs.axis is not None:
|
|
axis = call.attrs.axis
|
|
# Handle different types of axis
|
|
if hasattr(axis, "__iter__") and not isinstance(axis, str):
|
|
# It's an array/list, convert to list of ints
|
|
axis = [
|
|
int(item.value) if hasattr(item, "value") else int(item) for item in axis
|
|
]
|
|
elif hasattr(axis, "value"):
|
|
# It's a TVM IntImm, get the value
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
if hasattr(call.attrs, "keepdims"):
|
|
keepdims = bool(call.attrs.keepdims)
|
|
|
|
# Get the PyTorch function
|
|
func = getattr(torch, op_name)
|
|
|
|
# Call with appropriate parameters
|
|
if axis is not None:
|
|
# For max and min, PyTorch returns (values, indices) tuple when dim is specified
|
|
if op_name in ["max", "min"]:
|
|
if isinstance(axis, list) and len(axis) == 1:
|
|
axis = axis[0]
|
|
elif isinstance(axis, list) and len(axis) > 1:
|
|
axis = axis[0]
|
|
result = func(tensor_arg, axis, keepdim=keepdims)
|
|
if isinstance(result, tuple):
|
|
return result[0]
|
|
else:
|
|
return result
|
|
else:
|
|
return func(tensor_arg, dim=axis, keepdim=keepdims)
|
|
else:
|
|
return func(tensor_arg)
|
|
|
|
def _convert_squeeze(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert squeeze to torch.squeeze with proper axis handling."""
|
|
if len(call.args) < 1:
|
|
return "<squeeze_error: insufficient arguments>"
|
|
|
|
# Convert the tensor argument
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Get axis from call.attrs
|
|
axis = None
|
|
if call.attrs and hasattr(call.attrs, "axis") and call.attrs.axis is not None:
|
|
axis = call.attrs.axis
|
|
# Handle different types of axis
|
|
if hasattr(axis, "__iter__") and not isinstance(axis, str):
|
|
axis = [int(item.value) if hasattr(item, "value") else int(item) for item in axis]
|
|
elif hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
# Call torch.squeeze with appropriate parameters
|
|
if axis is not None:
|
|
return torch.squeeze(tensor_arg, dim=axis)
|
|
else:
|
|
return torch.squeeze(tensor_arg)
|
|
|
|
def _convert_tensor_ops(self, call: relax.Call, args: list[Any], op_name: str) -> Any:
|
|
"""Convert tensor operations like concat, split, stack."""
|
|
if len(call.args) < 1:
|
|
return f"<{op_name}_error: insufficient arguments>"
|
|
|
|
# Convert arguments
|
|
converted_args = [self.convert_expr(arg, args) for arg in call.args]
|
|
|
|
if op_name == "concat":
|
|
# torch.cat(tensors, dim=0)
|
|
# In Relax, concat takes a tuple of tensors as first argument
|
|
if len(converted_args) == 1 and isinstance(converted_args[0], tuple):
|
|
# This is a tuple of tensors
|
|
tensors = converted_args[0]
|
|
else:
|
|
# Direct tensor arguments
|
|
tensors = converted_args
|
|
axis = 0
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
return torch.cat(tensors, dim=axis)
|
|
|
|
elif op_name == "split":
|
|
# torch.split(tensor, split_size_or_sections, dim=0)
|
|
tensor = converted_args[0]
|
|
split_size = converted_args[1] if len(converted_args) > 1 else 1
|
|
axis = 0
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
# Handle indices_or_sections parameter
|
|
if call.attrs and hasattr(call.attrs, "indices_or_sections"):
|
|
indices_or_sections = call.attrs.indices_or_sections
|
|
if hasattr(indices_or_sections, "value"):
|
|
indices_or_sections = int(indices_or_sections.value)
|
|
elif isinstance(indices_or_sections, int | float):
|
|
indices_or_sections = int(indices_or_sections)
|
|
|
|
# If indices_or_sections is an integer, it means split into N equal parts
|
|
if isinstance(indices_or_sections, int):
|
|
total_size = tensor.shape[axis]
|
|
split_size = total_size // indices_or_sections
|
|
result = torch.split(tensor, split_size, dim=axis)
|
|
return result
|
|
else:
|
|
result = torch.split(tensor, indices_or_sections, dim=axis)
|
|
return result
|
|
else:
|
|
result = torch.split(tensor, split_size, dim=axis)
|
|
return result
|
|
|
|
elif op_name == "stack":
|
|
# torch.stack(tensors, dim=0)
|
|
if len(converted_args) == 1 and isinstance(converted_args[0], tuple):
|
|
tensors = converted_args[0]
|
|
else:
|
|
tensors = converted_args
|
|
axis = 0
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
return torch.stack(tensors, dim=axis)
|
|
|
|
else:
|
|
return f"<{op_name}_error: unsupported operation>"
|
|
|
|
def _convert_reshape(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert reshape operation."""
|
|
if len(call.args) < 2:
|
|
return "<reshape_error: insufficient arguments>"
|
|
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
shape_arg = call.args[1]
|
|
|
|
# Convert shape argument to Python tuple
|
|
if isinstance(shape_arg, relax.ShapeExpr):
|
|
if hasattr(shape_arg, "values"):
|
|
shape = tuple(
|
|
int(v.value) if hasattr(v, "value") else int(v) for v in shape_arg.values
|
|
)
|
|
else:
|
|
shape = (int(shape_arg),)
|
|
elif isinstance(shape_arg, relax.Constant):
|
|
# Constant tensor case
|
|
shape_data = shape_arg.data.numpy()
|
|
shape = tuple(int(v) for v in shape_data)
|
|
else:
|
|
# Try to convert as expression
|
|
converted_shape = self.convert_expr(shape_arg, args)
|
|
if isinstance(converted_shape, list | tuple):
|
|
shape = tuple(int(v) for v in converted_shape)
|
|
else:
|
|
shape = (int(converted_shape),)
|
|
|
|
return torch.reshape(tensor_arg, shape)
|
|
|
|
def _convert_permute_dims(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert permute_dims operation."""
|
|
if len(call.args) < 1:
|
|
return "<permute_dims_error: insufficient arguments>"
|
|
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Extract axes from call.attrs
|
|
if call.attrs and hasattr(call.attrs, "axes"):
|
|
axes = call.attrs.axes
|
|
# Handle TVM Array type
|
|
if hasattr(axes, "__iter__") and not isinstance(axes, str):
|
|
# Convert TVM Array or Python list/tuple to tuple of ints
|
|
axes = tuple(int(v.value) if hasattr(v, "value") else int(v) for v in axes)
|
|
elif isinstance(axes, list | tuple):
|
|
axes = tuple(int(v) for v in axes)
|
|
else:
|
|
axes = (int(axes),)
|
|
else:
|
|
return "<permute_dims_error: no axes attribute>"
|
|
|
|
return torch.permute(tensor_arg, axes)
|
|
|
|
def _convert_take(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert take operation."""
|
|
if len(call.args) < 2:
|
|
return "<take_error: insufficient arguments>"
|
|
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
indices_arg = self.convert_expr(call.args[1], args)
|
|
|
|
# Extract axis from call.attrs
|
|
axis = None
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
if axis is not None:
|
|
# Use advanced indexing for specific axis
|
|
if axis == 0:
|
|
return tensor_arg[indices_arg]
|
|
else:
|
|
# For other axes, we need to use torch.index_select
|
|
return torch.index_select(tensor_arg, dim=axis, index=indices_arg)
|
|
else:
|
|
# No axis specified, use torch.take (flattens the tensor)
|
|
return torch.take(tensor_arg, indices_arg)
|
|
|
|
def _convert_flip(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert flip operation."""
|
|
if len(call.args) < 1:
|
|
return "<flip_error: insufficient arguments>"
|
|
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Extract axis from call.attrs
|
|
axis = None
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
if axis is not None:
|
|
# Convert single axis to list for torch.flip
|
|
dims = [axis]
|
|
else:
|
|
# Default: flip all dimensions
|
|
dims = list(range(tensor_arg.dim()))
|
|
|
|
return torch.flip(tensor_arg, dims=dims)
|
|
|
|
def _convert_tile(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert tile operation."""
|
|
if len(call.args) < 1:
|
|
return "<tile_error: insufficient arguments>"
|
|
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Extract repeats from call.attrs
|
|
if call.attrs and hasattr(call.attrs, "repeats"):
|
|
repeats = call.attrs.repeats
|
|
# Handle TVM Array type
|
|
if hasattr(repeats, "__iter__") and not isinstance(repeats, str):
|
|
repeats = tuple(int(v.value) if hasattr(v, "value") else int(v) for v in repeats)
|
|
elif isinstance(repeats, list | tuple):
|
|
repeats = tuple(int(v) for v in repeats)
|
|
else:
|
|
repeats = (int(repeats),)
|
|
else:
|
|
return "<tile_error: no repeats attribute>"
|
|
|
|
return torch.tile(tensor_arg, dims=repeats)
|
|
|
|
def _convert_repeat(self, call: relax.Call, args: list[Any]) -> Any:
|
|
"""Convert repeat operation."""
|
|
if len(call.args) < 1:
|
|
return "<repeat_error: insufficient arguments>"
|
|
|
|
tensor_arg = self.convert_expr(call.args[0], args)
|
|
|
|
# Extract repeats and axis from call.attrs
|
|
repeats = 1
|
|
axis = None
|
|
|
|
if call.attrs and hasattr(call.attrs, "repeats"):
|
|
repeats = call.attrs.repeats
|
|
if hasattr(repeats, "value"):
|
|
repeats = int(repeats.value)
|
|
elif isinstance(repeats, int | float):
|
|
repeats = int(repeats)
|
|
|
|
if call.attrs and hasattr(call.attrs, "axis"):
|
|
axis = call.attrs.axis
|
|
if hasattr(axis, "value"):
|
|
axis = int(axis.value)
|
|
elif isinstance(axis, int | float):
|
|
axis = int(axis)
|
|
|
|
if axis is not None:
|
|
return torch.repeat_interleave(tensor_arg, repeats=repeats, dim=axis)
|
|
else:
|
|
return torch.repeat_interleave(tensor_arg, repeats=repeats)
|
|
|
|
def _convert_shape_expr(self, shape_expr: relax.ShapeExpr) -> Any:
|
|
"""Convert a Relax shape expression to Python equivalent."""
|
|
if hasattr(shape_expr, "values"):
|
|
return f"<shape: ({', '.join(map(str, shape_expr.values))})>"
|
|
return f"<shape: {shape_expr}>"
|
|
|
|
|
|
def convert_relax_to_pyfunc(ir_module: IRModule, relax_function_names: str | list[str]) -> IRModule:
|
|
"""Convert Relax functions to Python functions.
|
|
|
|
Args:
|
|
ir_module: The IRModule containing Relax functions
|
|
relax_function_names: Name(s) of Relax functions to convert
|
|
|
|
Returns:
|
|
IRModule with converted Python functions stored in pyfuncs
|
|
|
|
Example:
|
|
>>> converted_ir_mod = convert_relax_to_pyfunc(ir_mod, "my_function")
|
|
>>> converted_ir_mod = convert_relax_to_pyfunc(ir_mod, ["func1", "func2"])
|
|
"""
|
|
converter = RelaxToPyFuncConverter(ir_module)
|
|
return converter.convert(relax_function_names)
|