[Dev] Separate tvm and tilelang and utilize the upstream tilelang release

.gitmodules

@@ -1,8 +1,12 @@
 [submodule "3rdparty/tvm"]
 	path = 3rdparty/tvm
-	url = https://github.com/TileLang/tvm.git
-	branch = upstream
+	url = https://github.com/tile-ai/tvm.git
+	branch = tilelang_codebase
+[submodule "3rdparty/tilelang"]
+	path = 3rdparty/tilelang
+	url = https://github.com/tile-ai/tilelang
+	branch = bitblas
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
-	url = https://github.com/TileLang/cutlass
+	url = https://github.com/tile-ai/cutlass
 	branch = tldev

bitblas/__init__.py

@@ -90,7 +90,7 @@ def new_func(*args, **kwargs):
 
 if TVM_IMPORT_PYTHON_PATH is not None:
     os.environ["PYTHONPATH"] = (TVM_IMPORT_PYTHON_PATH + ":" + os.environ.get("PYTHONPATH", ""))
-    sys.path.insert(0, TVM_IMPORT_PYTHON_PATH + "/python")
+    sys.path.insert(0, TVM_IMPORT_PYTHON_PATH)
 else:
     # remove the existing tvm path in PYTHONPATH
     def remove_tvm_path(path):
@@ -107,6 +107,7 @@ def remove_tvm_path(path):
         os.environ["PYTHONPATH"] = (
             install_tvm_path + "/python:" + os.environ.get("PYTHONPATH", ""))
         sys.path.insert(0, install_tvm_path + "/python")
+        os.environ["TVM_IMPORT_PYTHON_PATH"] = install_tvm_path + "/python"
 
     # developed 3rdparty tvm
     develop_tvm_path = os.path.join(
@@ -119,6 +120,22 @@ def remove_tvm_path(path):
         os.environ["PYTHONPATH"] = (
             develop_tvm_path + "/python:" + os.environ.get("PYTHONPATH", ""))
         sys.path.insert(0, develop_tvm_path + "/python")
+        os.environ["TVM_IMPORT_PYTHON_PATH"] = develop_tvm_path + "/python"
+
+# TILELANG PATH
+if os.environ.get("TILELANG_IMPORT_PATH", None) is None:
+    install_tilelang_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "3rdparty", "tilelang")
+    develop_tilelang_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "..", "3rdparty", "tilelang")
+    if os.path.exists(install_tilelang_path):
+        os.environ["PYTHONPATH"] = install_tilelang_path + ":" + os.environ.get("PYTHONPATH", "")
+        sys.path.insert(0, install_tilelang_path)
+    elif (os.path.exists(develop_tilelang_path) and develop_tilelang_path not in sys.path):
+        os.environ["PYTHONPATH"] = develop_tilelang_path + ":" + os.environ.get("PYTHONPATH", "")
+        sys.path.insert(0, develop_tilelang_path)
+    else:
+        logger.warning(TL_TEMPLATE_NOT_FOUND_MESSAGE)
 
 if os.environ.get("TL_CUTLASS_PATH", None) is None:
     install_cutlass_path = os.path.join(
@@ -133,6 +150,8 @@ def remove_tvm_path(path):
         logger.warning(CUTLASS_NOT_FOUND_MESSAGE)
 
 import tvm as tvm  # noqa: E402
+import tilelang as tilelang  # noqa: E402
+
 from .base import (
     TileDevice,  # noqa: F401
     fast_tune,  # noqa: F401

bitblas/base/base_scheduler.py

@@ -1,9 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from tvm import te
 from tvm import IRModule
 from tvm.tir import PrimFunc
 from typing import Optional, Union, Callable, List, Dict
 from dataclasses import dataclass, field
-from tvm.tl.transform import Simplify
+from tilelang.transform import Simplify
 from abc import ABC, abstractmethod
 from bitblas.base.arch import TileDevice, is_volta_arch, is_ampere_arch, is_cdna_arch, auto_infer_current_arch
 from bitblas.base.roller.hint import Hint

bitblas/builder/lib_generator/__init__.py

@@ -67,24 +67,24 @@ def compile_lib(self, timeout: float = None, with_tl: bool = False):
             raise ValueError(f"Unsupported platform: {platform}")
 
         if with_tl:
-            install_tvm_path = os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "../..", "3rdparty", "tvm")
-            develop_tvm_path = os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "../../..", "3rdparty", "tvm")
+            install_tilelang_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "../..", "3rdparty", "tilelang")
+            develop_tilelang_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "../../..", "3rdparty", "tilelang")
 
-            tvm_root = next((path for path in [install_tvm_path, develop_tvm_path]
-                             if os.path.exists(path) and path not in sys.path), None)
+            tilelang_root = next((path for path in [install_tilelang_path, develop_tilelang_path]
+                                  if os.path.exists(path) and path not in sys.path), None)
 
             if "TL_TEMPLATE_PATH " in os.environ:
                 tl_template_path = os.environ["TL_TEMPLATE_PATH"]
             else:
-                tl_template_path = osp.abspath(osp.join(tvm_root, "src/tl"))
+                tl_template_path = osp.abspath(osp.join(tilelang_root, "src"))
 
-            tl_template_path = osp.abspath(osp.join(tvm_root, "src/tl"))
+            tl_template_path = osp.abspath(osp.join(tilelang_root, "src"))
             if "TL_CUTLASS_PATH" in os.environ:
                 cutlass_path = os.environ["TL_CUTLASS_PATH"]
             else:
-                cutlass_path = osp.abspath(osp.join(tvm_root, "3rdparty/cutlass/include"))
+                cutlass_path = osp.abspath(osp.join(tilelang_root, "3rdparty/cutlass/include"))
 
             command += [
                 "-I" + tl_template_path,

bitblas/gpu/intrin/hip.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from tvm.runtime import convert
 from tvm.tir.expr import Cast, IntImm
 from tvm.tir.function import TensorIntrin

bitblas/ops/general_flashatten/tilelang/flashatten.py

@@ -3,7 +3,8 @@
 
 from bitblas import tvm as tvm
 from bitblas.base.base_scheduler import BaseScheduler
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from dataclasses import dataclass
 from typing import Optional
 import logging

bitblas/ops/general_matmul/tilelang/dense/gemv_simt.py

@@ -1,9 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from functools import reduce
 from typing import Optional, List
-import tvm.tl.language as T
+import tilelang.language as T
 from tvm import DataType
 from tvm.tir import PrimFunc
 

bitblas/ops/general_matmul/tilelang/dense/matmul_mma.py

@@ -3,8 +3,9 @@
 # tile represents tile library
 
 from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from tvm import DataType
-import tvm.tl.language as T
+import tilelang.language as T
 from typing import Optional, List
 from bitblas.tl.utils import (
     get_mma_micro_size,

bitblas/ops/general_matmul/tilelang/dense/matmul_simt.py

@@ -1,8 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from typing import Optional, List
-import tvm.tl.language as T
+import tilelang.language as T
 from tvm import DataType
 from tvm.tir import PrimFunc
 

bitblas/ops/general_matmul/tilelang/dense/matmul_tile.py

@@ -3,7 +3,8 @@
 # tile represents tile library
 
 from bitblas import tvm as tvm
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Optional, List
 from bitblas.tl.utils import (
     get_mma_micro_size,

bitblas/ops/general_matmul/tilelang/dequantize/gemv_dequantize_simt.py

@@ -1,9 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from functools import reduce
 from typing import Optional, List
-import tvm.tl.language as T
+import tilelang.language as T
 from tvm import DataType
 from tvm.tir import PrimFunc
 

bitblas/ops/general_matmul/tilelang/dequantize/matmul_dequantize_mma.py

@@ -2,7 +2,8 @@
 # Licensed under the MIT License.
 from bitblas import tvm as tvm
 from tvm import DataType
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Optional, List
 from bitblas.tl.utils import (
     get_mma_micro_size,  # noqa: F401

bitblas/ops/general_matmul/tilelang/dequantize/matmul_dequantize_mma_weight_transform.py

@@ -1,8 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from tvm import DataType
-import tvm.tl.language as T
+import tilelang.language as T
 from typing import Optional, List
 from bitblas.tl.utils import (
     get_mma_micro_size,  # noqa: F401
@@ -680,7 +681,8 @@ def is_b_smooth(self):
 
 
 @dataclass
-class MatmulINT4DequantizeMMAWeightPropagationScheduler(MatmulDequantizeMMAWeightPropagationScheduler):
+class MatmulINT4DequantizeMMAWeightPropagationScheduler(
+        MatmulDequantizeMMAWeightPropagationScheduler):
 
     class TLHint(MatmulDequantizeMMAWeightPropagationScheduler.TLHint):
         hint_type: str = "MatmulINT4DequantizeMMAWeightPropagationScheduler"

bitblas/ops/general_matmul/tilelang/dequantize/matmul_dequantize_simt.py

@@ -3,7 +3,8 @@
 from bitblas import tvm as tvm
 from tvm import DataType
 from tvm.tir import PrimFunc
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Optional, List
 from bitblas.base.arch import TileDevice
 from bitblas.base.roller.hint import Hint

bitblas/ops/general_matmul/tilelang/dequantize/matmul_dequantize_tile.py

@@ -4,7 +4,8 @@
 
 from bitblas import tvm as tvm
 from tvm import DataType
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Optional, List
 from bitblas.base.arch import TileDevice
 from bitblas.base.roller.hint import Hint

bitblas/ops/operator.py

@@ -1,8 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from abc import ABC, abstractmethod
-from bitblas import tvm
-from tvm import tl
+from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from tvm import IRModule
 from tvm.runtime.module import Module
 from tvm.target import Target
@@ -192,7 +192,7 @@ def tvm_callback_hip_postproc(code, _):
                     if self.is_tir_backend():
                         rt_mod = tvm.build(self.scheduled_ir_module, target=target)
                     elif self.is_tilelang_backend():
-                        rt_mod = tl.lower(
+                        rt_mod = tilelang.lower(
                             self.scheduled_ir_module, target=target, runtime_only=True)
                     else:
                         raise ValueError(f"Unsupported backend: {self.backend}")

bitblas/tl/mfma_layout.py

@@ -1,8 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-
+from bitblas import tilelang as tilelang
 from tvm import DataType
-import tvm.tl.language as T
+import tilelang.language as T
 from tvm.runtime import convert
 
 

bitblas/tl/mfma_macro_generator.py

@@ -1,7 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Tuple
 from tvm import DataType
 from tvm.tir import PrimExpr

bitblas/tl/mma_layout.py

@@ -3,7 +3,8 @@
 
 from typing import Union
 from tvm import arith, DataType
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 
 
 def ldmatrix_32x8_to_shared_16x16_layout(thread_id, local_id):

bitblas/tl/mma_macro_generator.py

@@ -1,7 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Union, Tuple, Optional
 from bitblas.base.operator_common import TransformKind
 from tvm import DataType

bitblas/tl/tuner.py

@@ -1,7 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from bitblas import tvm
+from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 import os
 import logging
 import tempfile
@@ -10,7 +11,6 @@
 from tvm import IRModule
 from tvm.runtime import Module
 from tvm.tir import Schedule
-import tvm.tl as tl
 from bitblas.tl.base_hint import BaseTLHint
 from bitblas.base.arch import TileDevice
 from bitblas.base.utils import get_dummy_input_arrays
@@ -122,7 +122,7 @@ def tvm_callback_cuda_postproc(code, _):
                     "tir.disable_cse_tir": True,
                     **(config.pass_context if config.pass_context else {})
                 }):
-            rt_mod = tl.lower(tl_prim_func, arch.target, runtime_only=True)
+            rt_mod = tilelang.lower(tl_prim_func, arch.target, runtime_only=True)
 
         from tvm.contrib.tar import tar  # Import the tar module
 

bitblas/tl/wmma_macro_generator.py

@@ -1,7 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import tvm.tl.language as T
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 from typing import Tuple, Optional
 from tvm import DataType
 from tvm.tir import PrimExpr

bitblas/utils/rtmod_analysis.py

@@ -7,27 +7,27 @@
 from tvm.target import Target
 from typing import Tuple, List
 from tvm import tir
-from tvm import tl
-from tvm.tl.engine import is_device_call
+from bitblas import tilelang as tilelang
+from tilelang.engine import is_device_call
 
 
 def get_annotated_device_mod_from_tl(mod: IRModule, target: Target) -> "IRModule":
     target_host = tvm.target.Target("llvm -keys=cpu")
     target = tvm.target.Target(target, target_host)
     mod = tir.transform.BindTarget(target)(mod)
 
-    mod = tl.transform.FrontendLegalize()(mod)
+    mod = tilelang.transform.FrontendLegalize()(mod)
     mod = tir.transform.Simplify()(mod)
-    mod = tl.transform.LayoutInference()(mod)
-    mod = tl.transform.LowerTileOp()(mod)
+    mod = tilelang.transform.LayoutInference()(mod)
+    mod = tilelang.transform.LowerTileOp()(mod)
     mod = tir.transform.Simplify()(mod)
 
     if target.arch == "sm_90":
-        mod = tl.transform.WarpSpecializedPipeline()(mod)
+        mod = tilelang.transform.WarpSpecializedPipeline()(mod)
     else:
         mod = tir.transform.PlanAndUpdateBufferAllocationLocation()(mod)
-        mod = tl.transform.PipelinePlanning()(mod)
-        mod = tl.transform.InjectSoftwarePipeline()(mod)
+        mod = tilelang.transform.PipelinePlanning()(mod)
+        mod = tilelang.transform.InjectSoftwarePipeline()(mod)
 
     mod = tir.transform.LowerOpaqueBlock()(mod)
     mod = tir.transform.FlattenBuffer()(mod)
@@ -57,7 +57,7 @@ def get_annotated_device_mod_from_tl(mod: IRModule, target: Target) -> "IRModule
     # the Legalization.
     mod = tir.transform.LowerThreadAllreduce()(mod)
     mod = tir.transform.ThreadSync("shared.dyn")(mod)
-    mod = tl.transform.LowerHopperIntrin()(mod)
+    mod = tilelang.transform.LowerHopperIntrin()(mod)
     mod = tir.transform.InjectPTXAsyncCopy()(mod)
 
     mod = tir.transform.AnnotateDeviceRegions()(mod)