chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,107 @@
+"""A compiler pass that fuses transpose + dequantize."""
+
+import tvm
+from tvm import relax, s_tir, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTranspose")
+class FuseDequantizeTranspose:
+    """A compiler pass that fuses transpose + dequantize."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _DequantizeTransposeFuser(mod).transform()
+
+
+@mutator
+class _DequantizeTransposeFuser(PyExprMutator):
+    def __init__(
+        self,
+        mod: IRModule,
+    ):
+        super().__init__(mod)
+        self.mod = mod
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                updated_func = remove_all_unused(updated_func)
+                self.builder_.update_func(g_var, updated_func)
+        return self.builder_.get()
+
+    def visit_call_(
+        self,
+        call: relax.Call,
+    ) -> relax.Expr:
+        call = self.visit_expr_post_order(call)
+        if call.op != tvm.ir.Op.get("relax.matmul"):
+            return call
+        # Do not fuse dequantize-transpose for GeMM
+        if (
+            call.args[0].ty.ndim < 2
+            or not isinstance(call.args[0].ty.shape[-2], tirx.IntImm)
+            or call.args[0].ty.shape[-2].value != 1
+        ):
+            return call
+
+        matmul_rhs = self.lookup_binding(call.args[1])
+        if (
+            not isinstance(matmul_rhs, relax.Call)
+            or matmul_rhs.op != tvm.ir.Op.get("relax.permute_dims")
+            or matmul_rhs.args[0].ty.ndim != 2
+            or matmul_rhs.attrs.axes is not None
+        ):
+            return call
+
+        transpose_input = self.lookup_binding(matmul_rhs.args[0])
+        if (
+            not isinstance(transpose_input, relax.Call)
+            or transpose_input.op != tvm.ir.Op.get("relax.call_tir")
+            or not transpose_input.args[0].name_hint.startswith("dequantize")
+            or not isinstance(transpose_input.ty, relax.TensorType)
+        ):
+            return call
+
+        dequantize_tir_func = self.mod[transpose_input.args[0]]
+        assert isinstance(dequantize_tir_func, tirx.PrimFunc)
+        if (
+            len(dequantize_tir_func.body.block.alloc_buffers) != 1
+            or not isinstance(dequantize_tir_func.body.block.body, tirx.SeqStmt)
+            or len(dequantize_tir_func.body.block.body) != 2
+            or not isinstance(dequantize_tir_func.body.block.body[1], tirx.For)
+            or not isinstance(dequantize_tir_func.body.block.body[1].body.body, tirx.SBlockRealize)
+            or dequantize_tir_func.body.block.body[1].body.body.block.name_hint != "T_transpose"
+        ):
+            return call
+
+        new_func_buffers = [
+            dequantize_tir_func.buffer_map[var] for var in dequantize_tir_func.params
+        ]
+        new_func_buffers[-1] = dequantize_tir_func.body.block.alloc_buffers[0]
+        new_func = tirx.PrimFunc(
+            params=new_func_buffers,
+            body=tirx.SBlockRealize(
+                iter_values=[],
+                predicate=True,
+                block=tirx.SBlock(
+                    iter_vars=[],
+                    reads=[],
+                    writes=[],
+                    name_hint="root",
+                    body=dequantize_tir_func.body.block.body[0],
+                ),
+            ),
+        )
+        # Call `renew_defs` for deep-copy to avoid IR node duplication in
+        # different PrimFuncs of an IRModule.
+        new_func = s_tir.renew_defs(new_func)
+        g_var = self.builder_.add_func(new_func, func_name="dequantize")
+        dequantize_matmul_rhs = self.builder_.emit(
+            relax.call_tir(g_var, transpose_input.args[1], out_ty=matmul_rhs.ty)
+        )
+        return relax.op.matmul(call.args[0], dequantize_matmul_rhs, out_dtype=call.attrs.out_dtype)