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,64 @@
+"""A compiler pass that dispatch low-batch-gemm to gemv schedule."""
+
+import tvm
+import tvm_ffi
+from tvm import tirx
+from tvm.ir.module import IRModule
+from tvm.s_tir import dlight as dl
+
+
+@tvm.transform.module_pass(opt_level=0, name="LowBatchGemvSpecialize")
+class LowBatchGemvSpecialize:
+    """A compiler pass that dispatch low-batch-gemm to gemv schedule."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, tirx.PrimFunc):
+                low_batch_range = [2, 8]
+                buckets = [2, 4]
+                low_batch_funcs = []
+                for bucket in buckets:
+                    low_batch_mod = IRModule({})
+                    low_batch_mod["main"] = func
+                    low_batch_mod = dl.ApplyDefaultSchedule(
+                        dl.gpu.LowBatchGEMV(bucket),
+                    )(low_batch_mod)
+                    low_batch_funcs.append(low_batch_mod["main"])
+                if any(
+                    tvm_ffi.structural_equal(low_batch_func, func)
+                    for low_batch_func in low_batch_funcs
+                ):
+                    continue
+                buffers = func.buffer_map.values()
+                shapes = [buffer.shape for buffer in buffers]
+                symbolic_vars = set(
+                    expr for shape in shapes for expr in shape if isinstance(expr, tirx.Var)
+                )
+                if len(symbolic_vars) != 1:
+                    continue
+                gemm_mod = IRModule({})
+                gemm_mod["main"] = func
+                gemm_mod = dl.ApplyDefaultSchedule(
+                    dl.gpu.Matmul(),
+                )(gemm_mod)
+                gemm_func = gemm_mod["main"]
+                sym_var = next(iter(symbolic_vars))
+                body = gemm_func.body
+                for i, range_limit in reversed(list(enumerate(low_batch_range))):
+                    body = tirx.IfThenElse(
+                        tirx.op.tvm_thread_invariant(sym_var <= range_limit),
+                        low_batch_funcs[i].body,
+                        body,
+                    )
+                body = tirx.SBlock([], [], [], "root", body)
+                body = tirx.SBlockRealize([], True, body)
+                new_func = func.with_body(body)
+                new_func = new_func.with_attr("tirx.is_scheduled", 1)
+                new_func = new_func.with_attr("tirx.HoistIfThenElseExprWithBlock", 1)
+                mod.update_func(g_var, new_func)
+        return mod