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,51 @@
+"""A compiler pass that dispatches patterns to CUBLAS."""
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.backend import get_patterns_with_prefix
+
+try:
+    import tvm.relax.backend.cuda.cublas as _cublas  # noqa: F401
+    import tvm.relax.backend.rocm.hipblas as _hipblas  # noqa: F401
+except ImportError:
+    # Note: legacy path of cublas/hipblas for backward compatibility
+    pass
+
+
+@tvm.transform.module_pass(opt_level=0, name="BLASDispatch")
+class BLASDispatch:
+    """A compiler pass that dispatches patterns to cuBLAS/hipBLAS."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        if target.kind.name == "cuda":
+            self.has_blas = tvm.get_global_func("relax.ext.cublas", True)
+            if not self.has_blas:
+                raise Exception("cuBLAS is not enabled.")
+            self.patterns = get_patterns_with_prefix("cublas")
+        elif target.kind.name == "rocm":
+            self.has_blas = tvm.get_global_func("relax.ext.hipblas", True)
+            if not self.has_blas:
+                raise Exception("hipBLAS is not enabled.")
+            self.patterns = get_patterns_with_prefix("hipblas")
+        else:
+            raise Exception(f"Unsupported target {target.kind.name} for BLAS dispatch.")
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        model_names = [
+            gv.name_hint for gv, func in mod.functions.items() if isinstance(func, relax.Function)
+        ]
+        # exclude single batch decode
+        model_names = [name for name in model_names if "batch" in name or "decode" not in name]
+        mod = tvm.transform.Sequential(
+            [
+                relax.transform.FuseOpsByPattern(
+                    self.patterns,
+                    bind_constants=False,
+                    annotate_codegen=True,
+                    entry_functions=model_names,
+                ),
+                relax.transform.RunCodegen({}, entry_functions=model_names),
+            ]
+        )(mod)
+        return mod