chore: import upstream snapshot with attribution

python/tvm/relax/frontend/stablehlo/__init__.py

@@ -0,0 +1,22 @@
+# isort: skip_file
+# 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.
+"""
+StableHLO Frontends for constructing Relax programs, with the model importers
+"""
+
+from .stablehlo_translator import from_stablehlo

python/tvm/relax/frontend/stablehlo/stablehlo_translator.py

@@ -0,0 +1,445 @@
+# 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=import-outside-toplevel, unused-argument
+
+"""StableHLO frontend of Relax."""
+
+from collections.abc import Callable
+from typing import Any
+
+import tvm
+from tvm import relax, tirx
+
+
+class StableHLOImporter:
+    """An importer from StableHLO to Relax."""
+
+    from jaxlib import mlir
+    from jaxlib.mlir.dialects import stablehlo
+
+    def __init__(self) -> None:
+        from jaxlib import mlir
+
+        self._nodes: dict[str | mlir.ir.Operation, relax.Expr] = {}
+        self.block_builder: relax.BlockBuilder = None
+        self.create_convert_map()
+
+    @staticmethod
+    def _convert_data_type(input_type):
+        """converts the data type from mlir to tvm."""
+        from jaxlib import mlir
+
+        if mlir.ir.ShapedType.isinstance(input_type):
+            input_type = mlir.ir.ShapedType(input_type).element_type
+
+        input_type = str(input_type)
+        if input_type == "f16":
+            return "float16"
+        elif input_type in ["f32", "F32Type"]:
+            return "float32"
+        elif input_type in ["f64", "F64Type"]:
+            return "float64"
+        elif input_type == "i1":
+            return "bool"
+        elif input_type == "i8":
+            return "int8"
+        elif input_type == "i16":
+            return "int16"
+        elif input_type == "i32":
+            return "int32"
+        elif input_type == "i64":
+            return "int64"
+        elif input_type == "ui8":
+            return "uint8"
+        elif input_type == "ui16":
+            return "uint16"
+        elif input_type == "ui32":
+            return "uint32"
+        elif input_type == "ui64":
+            return "uint64"
+        else:
+            raise NotImplementedError(f"input_type {input_type} is not handled yet")
+
+    def _attr2value(self, node) -> Any | list[Any]:
+        import numpy as np
+        from jaxlib import mlir
+
+        if mlir.ir.IntegerAttr.isinstance(node):
+            int_attr = mlir.ir.IntegerAttr(node)
+            return int_attr.value
+        if mlir.ir.FloatAttr.isinstance(node):
+            float_attr = mlir.ir.FloatAttr(node)
+            return float_attr.value
+        if mlir.ir.DenseIntElementsAttr.isinstance(node):
+            dense_attr = mlir.ir.DenseIntElementsAttr(node)
+        elif mlir.ir.DenseFPElementsAttr.isinstance(node):
+            dense_attr = mlir.ir.DenseFPElementsAttr(node)
+        else:
+            raise ValueError("Unsupported Attribute type: " + str(type(node)))
+        ret = []
+        for val in dense_attr:
+            ret.append(val)
+        shape = self.get_shape(node.type)
+        dtype = self._convert_data_type(node.type)
+        return np.asarray(ret, dtype).reshape(shape).tolist()
+
+    def retrieve_operands(self, node):
+        return self._retrieve_operands(node.operands)
+
+    def _retrieve_operands(self, node):
+        from jaxlib import mlir
+
+        # the operand is one of the inputs of FuncOp
+        if isinstance(node, mlir.ir.Operation):
+            return self._nodes[node]
+        if isinstance(node, tuple):
+            return tuple(self._retrieve_operands(x) for x in node)
+        if isinstance(node, list | mlir.ir.OpOperandList):
+            return [self._retrieve_operands(x) for x in node]
+        if isinstance(node, dict):
+            return {self._retrieve_operands(k): self._retrieve_operands(v) for k, v in node.items()}
+        if isinstance(node, mlir.ir.Value):
+            if isinstance(node.owner, mlir.ir.Block):
+                block_arg = mlir.ir.BlockArgument(node)
+                return self._nodes["arg" + str(block_arg.arg_number)]
+            return self._retrieve_operands(node.owner)
+        return node
+
+    def get_shape(self, inpt_type) -> list[Any]:
+        """Get the shape from Type like tensor<?x?xf32>"""
+        from jaxlib import mlir
+
+        shape_type = inpt_type
+        if isinstance(shape_type, mlir.ir.Type):
+            shape_type = mlir.ir.ShapedType(shape_type)
+        ret = []
+        for i in range(shape_type.rank):
+            # get_dim_size
+            if shape_type.is_dynamic_dim(i):
+                n = tirx.Var("n", "int64")
+                ret.append(n)
+            else:
+                ret.append(shape_type.get_dim_size(i))
+
+        return ret
+
+    @staticmethod
+    def _promote_binary_op_args(lhs, rhs):
+        if not isinstance(lhs, relax.Expr) and not isinstance(rhs, relax.Expr):
+            msg = "Both the lhs and the rhs are not expressions."
+            raise AssertionError(msg)
+        if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
+            return lhs, rhs
+        if isinstance(lhs, relax.Expr):
+            assert isinstance(lhs.ty, relax.TensorType)
+            return lhs, relax.const(rhs, lhs.ty.dtype)
+        assert isinstance(rhs.ty, relax.TensorType)
+        return relax.const(lhs, rhs.ty.dtype), rhs
+
+    def _call_binary_op(self, op, lhs, rhs):
+        lhs, rhs = StableHLOImporter._promote_binary_op_args(lhs, rhs)
+        return self.block_builder.emit(op(lhs, rhs))
+
+    def _add(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.add, lhs, rhs)
+        return lhs + rhs
+
+    def _maximum(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.maximum(lhs, rhs))
+
+    def _minimum(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.minimum(lhs, rhs))
+
+    def _divide(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.divide, lhs, rhs)
+        return lhs / rhs
+
+    def _multiply(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.multiply, lhs, rhs)
+        return lhs * rhs
+
+    def _subtract(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.subtract, lhs, rhs)
+        return lhs - rhs
+
+    def _broadcast_in_dim(self, node: mlir.ir.Operation) -> relax.Expr:
+        operands = self.retrieve_operands(node)
+        data = operands[0]
+        # broadcast_dims = self._attr2value(node.attributes["broadcast_dimensions"])
+        shape = self.get_shape(node.result.type)
+        # scalar
+        if len(shape) == 0:
+            return data
+        return self.block_builder.emit(relax.op.broadcast_to(data, shape))
+
+    def _const(self, node: mlir.ir.Operation) -> relax.Expr:
+        const_value = self._attr2value(node.attributes["value"])
+        dtype = self._convert_data_type(node.result.type)
+        return relax.const(const_value, dtype)
+
+    def _dot_general(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.matmul(lhs, rhs))
+
+    def _convolution(self, node) -> relax.Expr:
+        from jaxlib import mlir
+
+        x, weight = self.retrieve_operands(node)
+        shaped_type = mlir.ir.ShapedType(node.result.type)
+        out_dtype = self._convert_data_type(shaped_type.element_type)
+        strides = self._attr2value(node.attributes["window_strides"])
+        padding = self._attr2value(node.attributes["padding"])
+        lhs_dilation = self._attr2value(node.attributes["lhs_dilation"])
+        rhs_dilation = self._attr2value(node.attributes["rhs_dilation"])
+        if len(lhs_dilation) > 0:
+            lhs_dilation = lhs_dilation[0]
+        if len(rhs_dilation) > 0:
+            rhs_dilation = rhs_dilation[0]
+        dilation = (lhs_dilation, rhs_dilation)
+        groups = self._attr2value(node.attributes["batch_group_count"])
+        conv2d = relax.op.nn.conv2d(
+            x,
+            weight,
+            strides=strides,
+            padding=padding[0],
+            dilation=dilation,
+            groups=groups,
+            data_layout="NHWC",
+            kernel_layout="HWIO",
+            out_dtype=out_dtype,
+        )
+
+        return self.block_builder.emit(conv2d)
+
+    def _reshape(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        if isinstance(data, list):
+            assert len(data) == 1
+            data = data[0]
+        new_shape = self.get_shape(node.result.type)
+        return self.block_builder.emit(relax.op.reshape(data, new_shape))
+
+    def _reduce(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        dimensions = self._attr2value(node.attributes["dimensions"])
+        if node.body is not None:
+            reducer_op = node.body.blocks[0].operations[0].OPERATION_NAME
+            assert reducer_op == "stablehlo.add", f"reducer {reducer_op} in reduce is not supported"
+        return self.block_builder.emit(relax.op.sum(data[0], axis=dimensions))
+
+    def _reduce_window(self, node: mlir.ir.Operation) -> relax.Expr:
+        operands = self.retrieve_operands(node)
+        window_dimensions = self._attr2value(node.attributes["window_dimensions"])
+        window_dilations = self._attr2value(node.attributes["window_dilations"])
+
+        if node.body is not None:
+            reducer_op = node.body.blocks[0].operations[0].OPERATION_NAME
+            assert reducer_op == "stablehlo.maximum", (
+                f"the reducer {reducer_op} in reduce_window is not supported"
+            )
+
+        pool_size = []
+        for i, window_dim in enumerate(window_dimensions):
+            if window_dim == 0:
+                pool_size.append(0)
+            else:
+                dilated_window_size = (window_dim - 1) * window_dilations[i] + 1
+                pool_size.append(dilated_window_size)
+        strides = self._attr2value(node.attributes["window_strides"])
+        # padding = self._attr2value(node.attributes["padding"])
+
+        # TODO (yongwww): Infer the layout automatically
+        layout = "NHWC"
+
+        ret = self.block_builder.emit(
+            relax.op.nn.max_pool2d(
+                operands[0],
+                pool_size=pool_size[1:3],  # HW
+                strides=strides[1:3],
+                padding=[1, 1],
+                dilation=window_dilations[1:3],
+                layout=layout,
+            )
+        )
+        return ret
+
+    def _rsqrt(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.rsqrt(data[0]))
+
+    def _sin(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sin(data[0]))
+
+    def _sinh(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sinh(data[0]))
+
+    def _cos(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.cos(data[0]))
+
+    def _cosh(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.cosh(data[0]))
+
+    def _sqrt(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sqrt(data[0]))
+
+    def _round(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.round(data[0]))
+
+    def _exp(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.exp(data[0]))
+
+    def _return(self, node: mlir.ir.Operation) -> relax.Expr:
+        outputs = self.retrieve_operands(node)
+        return self.block_builder.emit_output(self.nodes[outputs])
+
+    def create_convert_map(self):
+        from jaxlib import mlir
+
+        self.convert_map: dict[str, Callable[[mlir.ir.Operation], relax.Var]] = {
+            "stablehlo.add": self._add,
+            "stablehlo.broadcast_in_dim": self._broadcast_in_dim,
+            "stablehlo.constant": self._const,
+            "stablehlo.convolution": self._convolution,
+            "stablehlo.cosine": self._cos,
+            "stablehlo.cosh": self._cosh,
+            "stablehlo.divide": self._divide,
+            "stablehlo.dot_general": self._dot_general,
+            "stablehlo.exponential": self._exp,
+            "stablehlo.maximum": self._maximum,
+            "stablehlo.minimum": self._minimum,
+            "stablehlo.multiply": self._multiply,
+            "stablehlo.reshape": self._reshape,
+            "stablehlo.reduce": self._reduce,
+            "stablehlo.reduce_window": self._reduce_window,
+            "stablehlo.round_nearest_afz": self._round,
+            "stablehlo.rsqrt": self._rsqrt,
+            "stablehlo.sine": self._sin,
+            "chlo.sinh": self._sinh,
+            "stablehlo.sqrt": self._sqrt,
+            "stablehlo.subtract": self._subtract,
+            "func.return": self._return,
+            "stablehlo.return": self._return,
+        }
+
+    def from_stablehlo(self, model, input_info: list[tuple[tuple[int], str]]) -> tvm.IRModule:
+        """Convert a StableHLO Module to a Relax program.
+
+        Parameters
+        ----------
+        model : mlir.ir.Module
+            The StableHLO Module to convert.
+
+        input_info : List[Tuple[Tuple[int], str]]
+            A list of shapes and data types of input tensors.
+
+        Returns
+        -------
+        output : tvm.IRModule
+            The result IRModule with entry function "main"
+        """
+        from jaxlib import mlir
+        from jaxlib.mlir.dialects import stablehlo
+
+        assert isinstance(model, mlir.ir.Module)
+        block: mlir.ir.Block = model.body.operations[0].regions[0].blocks[0]
+
+        # inputs of the function
+        inputs = []
+        for idx, arg in enumerate(block.arguments.types):
+            arg_shape = mlir.ir.ShapedType(arg)
+            ipt_shape = self.get_shape(arg_shape)
+            ipt_dtype = self._convert_data_type(arg_shape.element_type)
+            ipt_name = "arg" + str(idx)
+            ipt_var = relax.Var(f"arg{idx}", relax.TensorType(ipt_shape, ipt_dtype))
+            self._nodes[ipt_name] = ipt_var
+            inputs.append(ipt_var)
+
+        # TODO (yongwww): Handle mlir.ir.Module with multiple functions
+        # Initialize the block builder with a function and a dataflow block.
+        # Raise error if the input stablehlo op is impure
+        func_name = "main"
+        self.block_builder = relax.BlockBuilder()
+
+        with self.block_builder.function(name=func_name, params=inputs.copy()):
+            output = None
+            with self.block_builder.dataflow():
+                block = model.body.operations[0].regions[0].blocks[0]
+                for operation in block.operations:
+                    if isinstance(operation, mlir.dialects.func.ReturnOp | stablehlo.ReturnOp):
+                        operation = operation.operands[0].owner
+                        # TODO (yongwww): handle multiple outputs
+                        output = self.block_builder.emit_output(self._nodes[operation])
+                        break
+
+                    if isinstance(operation, mlir.ir.OpView):
+                        op_name = operation.operation.name
+                        assert op_name in self.convert_map, f"Unsupported operation {op_name}"
+                        self._nodes[operation] = self.convert_map[op_name](operation)
+                    else:
+                        raise ValueError(f"Unsupported op {operation}")
+            assert output is not None
+            self.block_builder.emit_func_output(output)
+
+        mod = self.block_builder.get()
+        return mod
+
+
+def from_stablehlo(
+    stablehlo_module,
+    input_info: list[tuple[tuple[int], str]] | None = None,
+) -> tvm.IRModule:
+    """Convert a StableHLO Module to a Relax program
+
+    Parameters
+    ----------
+    stablehlo_module : Union[str, mlir.ir.Module]
+        The StableHLO Module to convert.
+
+    input_info : List[Tuple[Tuple[int], str]]
+        A list of shapes and data types of input tensors.
+
+    Returns
+    -------
+    output : tvm.IRModule
+        The result IRModule with entry function "main"
+    """
+    from jax._src.interpreters import mlir as jax_mlir
+
+    if isinstance(stablehlo_module, str):
+        # TODO (yongwww): support the serialized bytecode format of StableHLO
+        # model using stablehlo.deserialize_portable_artifact(ir) if the python
+        # binding is ready
+        context = jax_mlir.make_ir_context()
+        stablehlo_module = jax_mlir.ir.Module.parse(stablehlo_module, context)
+    return StableHLOImporter().from_stablehlo(stablehlo_module, input_info)