[Dev][TL] Decouple 3rdparty TileLang Backend with TVM

.gitmodules

@@ -6,3 +6,7 @@
 	path = 3rdparty/cutlass
 	url = https://github.com/TileLang/cutlass
 	branch = tldev
+[submodule "3rdparty/tile-lang"]
+	path = 3rdparty/tile-lang
+	url = https://github.com/TileLang/tile-lang
+	branch = dev

bitblas/__init__.py

@@ -3,62 +3,21 @@
 import sys
 import os
 
-# installing tvm
-install_tvm_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "3rdparty", "tvm")
-install_cutlass_path = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), "3rdparty", "cutlass")
-if os.path.exists(install_tvm_path) and install_tvm_path not in sys.path:
-    os.environ["PYTHONPATH"] = install_tvm_path + "/python:" + os.environ.get("PYTHONPATH", "")
-    os.environ["TL_CUTLASS_PATH"] = install_cutlass_path + "/include"
-    os.environ["TL_TEMPLATE_PATH"] = os.path.join(install_tvm_path, "src/tl")
-    sys.path.insert(0, install_tvm_path + "/python")
-
-develop_tvm_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "3rdparty", "tvm")
-develop_cutlass_path = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), "..", "3rdparty", "cutlass")
-if os.path.exists(develop_tvm_path) and develop_tvm_path not in sys.path:
-    os.environ["PYTHONPATH"] = develop_tvm_path + "/python:" + os.environ.get("PYTHONPATH", "")
-    os.environ["TL_CUTLASS_PATH"] = develop_cutlass_path + "/include"
-    os.environ["TL_TEMPLATE_PATH"] = os.path.join(install_tvm_path, "src/tl")
-    sys.path.insert(0, develop_tvm_path + "/python")
-
-import tvm as tvm  # noqa: E402
-from . import gpu  # noqa: F401
-from .base import (
-    TileDevice,  # noqa: F401
-    fast_tune,  # noqa: F401
-    ApplyDefaultSchedule,  # noqa: F401
-    ApplyFastTuning,  # noqa: F401
-    BlockInfo,  # noqa: F401
-    IterInfo,  # noqa: F401
-    ScheduleRule,  # noqa: F401
-    normalize_prim_func,  # noqa: F401
-    try_inline,  # noqa: F401
-    try_inline_contiguous_spatial,  # noqa: F401
-)
-
-from . import testing  # noqa: F401
-from .utils import auto_detect_nvidia_target, apply_transform_on_input  # noqa: F401
-from .ops.general_matmul import MatmulConfig, Matmul  # noqa: F401
-from .ops.general_matmul_splitk import MatmulConfigWithSplitK, MatmulWithSplitK  # noqa: F401
-from .ops.general_flashatten import FlashAttenConfig, FlashAtten  # noqa: F401
-from .module import Linear  # noqa: F401
-
 import warnings
 import functools
 import logging
 from tqdm import tqdm
 
 
 class TqdmLoggingHandler(logging.Handler):
-    """ Custom logging handler that directs log output to tqdm progress bar to avoid interference. """
+    """Custom logging handler that directs log output to tqdm progress bar to avoid interference."""
 
     def __init__(self, level=logging.NOTSET):
-        """ Initialize the handler with an optional log level. """
+        """Initialize the handler with an optional log level."""
         super().__init__(level)
 
     def emit(self, record):
-        """ Emit a log record. Messages are written to tqdm to ensure output in progress bars isn't corrupted. """
+        """Emit a log record. Messages are written to tqdm to ensure output in progress bars isn't corrupted."""
         try:
             msg = self.format(record)
             tqdm.write(msg)
@@ -67,8 +26,8 @@ def emit(self, record):
 
 
 def set_log_level(level):
-    """ Set the logging level for the module's logger.
-    
+    """Set the logging level for the module's logger.
+
     Args:
         level (str or int): Can be the string name of the level (e.g., 'INFO') or the actual level (e.g., logging.INFO).
         OPTIONS: 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'
@@ -80,15 +39,17 @@ def set_log_level(level):
 
 
 def _init_logger():
-    """ Initialize the logger specific for this module with custom settings and a Tqdm-based handler. """
+    """Initialize the logger specific for this module with custom settings and a Tqdm-based handler."""
     logger = logging.getLogger(__name__)
     handler = TqdmLoggingHandler()
     formatter = logging.Formatter(
-        fmt="%(asctime)s [BitBLAS:%(levelname)s]: %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
+        fmt="%(asctime)s [BitBLAS:%(levelname)s]: %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
     handler.setFormatter(formatter)
     logger.addHandler(handler)
     logger.propagate = False
-    set_log_level('WARNING')
+    set_log_level("WARNING")
 
 
 _init_logger()
@@ -107,12 +68,108 @@ def new_func(*args, **kwargs):
             warnings.warn(
                 f"Call to deprecated function {func.__name__} ({reason}).",
                 category=DeprecationWarning,
-                stacklevel=2)
+                stacklevel=2,
+            )
             return func(*args, **kwargs)
 
         return new_func
 
     return decorator
 
 
+logger = logging.getLogger(__name__)
+
+# SETUP ENVIRONMENT VARIABLES
+CUTLASS_NOT_FOUND_MESSAGE = ("CUTLASS is not installed or found in the expected path")
+", which may lead to compilation bugs when utilize tilelang backend."
+TL_TEMPLATE_NOT_FOUND_MESSAGE = ("TileLang is not installed or found in the expected path")
+", which may lead to compilation bugs when utilize tilelang backend."
+
+# Handle TVM_IMPORT_PYTHON_PATH to import tvm from the specified path
+TVM_IMPORT_PYTHON_PATH = os.environ.get("TVM_IMPORT_PYTHON_PATH", None)
+
+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")
+else:
+    # installed 3rdparty tvm
+    install_tvm_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "3rdparty", "tvm")
+    if os.path.exists(install_tvm_path) and install_tvm_path not in sys.path:
+        os.environ["PYTHONPATH"] = (
+            install_tvm_path + "/python:" + os.environ.get("PYTHONPATH", ""))
+        sys.path.insert(0, install_tvm_path + "/python")
+
+    # developed 3rdparty tvm
+    develop_tvm_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "..", "3rdparty", "tvm")
+    if os.path.exists(develop_tvm_path) and develop_tvm_path not in sys.path:
+        os.environ["PYTHONPATH"] = (
+            develop_tvm_path + "/python:" + os.environ.get("PYTHONPATH", ""))
+        sys.path.insert(0, develop_tvm_path + "/python")
+
+if os.environ.get("TL_CUTLASS_PATH", None) is None:
+    install_cutlass_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "3rdparty", "cutlass")
+    develop_cutlass_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "..", "3rdparty", "cutlass")
+    if os.path.exists(install_cutlass_path):
+        os.environ["TL_CUTLASS_PATH"] = install_cutlass_path + "/include"
+    elif (os.path.exists(develop_cutlass_path) and develop_cutlass_path not in sys.path):
+        os.environ["TL_CUTLASS_PATH"] = develop_cutlass_path + "/include"
+    else:
+        logger.warning(CUTLASS_NOT_FOUND_MESSAGE)
+
+install_tilelang_path = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "3rdparty", "tile-lang")
+develop_tilelang_path = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "..", "3rdparty", "tile-lang")
+
+if os.environ.get("TL_TEMPLATE_PATH", None) is None:
+    sys.path.insert(0, install_tilelang_path + "/python")
+    if os.path.exists(install_tilelang_path):
+        os.environ["TL_TEMPLATE_PATH"] = os.path.join(install_tilelang_path, "src")
+    elif os.path.exists(develop_tilelang_path):
+        os.environ["TL_TEMPLATE_PATH"] = os.path.join(develop_tilelang_path, "src")
+    else:
+        logger.warning(TL_TEMPLATE_NOT_FOUND_MESSAGE)
+
+if (os.path.exists(install_tilelang_path) and install_tilelang_path not in sys.path):
+    os.environ["PYTHONPATH"] = (
+        install_tilelang_path + "/python:" + os.environ.get("PYTHONPATH", ""))
+    sys.path.insert(0, install_tilelang_path + "/python")
+
+if (os.path.exists(develop_tilelang_path) and develop_tilelang_path not in sys.path):
+    os.environ["PYTHONPATH"] = (
+        develop_tilelang_path + "/python:" + os.environ.get("PYTHONPATH", ""))
+    sys.path.insert(0, develop_tilelang_path + "/python")
+
+import tvm as tvm  # noqa: E402
+import tilelang as tilelang  # noqa: E402
+from . import gpu  # noqa: F401
+from .base import (
+    TileDevice,  # noqa: F401
+    fast_tune,  # noqa: F401
+    ApplyDefaultSchedule,  # noqa: F401
+    ApplyFastTuning,  # noqa: F401
+    BlockInfo,  # noqa: F401
+    IterInfo,  # noqa: F401
+    ScheduleRule,  # noqa: F401
+    normalize_prim_func,  # noqa: F401
+    try_inline,  # noqa: F401
+    try_inline_contiguous_spatial,  # noqa: F401
+)
+
+from . import testing  # noqa: F401
+from .utils import (
+    auto_detect_nvidia_target,  # noqa: F401
+    apply_transform_on_input,  # noqa: F401
+)
+from .ops.general_matmul import MatmulConfig, Matmul  # noqa: F401
+from .ops.general_matmul_splitk import (
+    MatmulConfigWithSplitK,  # noqa: F401
+    MatmulWithSplitK,  # noqa: F401
+)
+from .ops.general_flashatten import FlashAttenConfig, FlashAtten  # noqa: F401
+from .module import Linear  # noqa: F401
+
 __version__ = "0.0.1.dev15"

bitblas/builder/lib_generator/__init__.py

@@ -50,24 +50,24 @@ def compile_lib(self, timeout: float = None, with_tl: bool = False):
         ]
 
         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", "tile-lang")
+            develop_tilelang_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "../../..", "3rdparty", "tile-lang")
 
-            tvm_root = next((path for path in [install_tvm_path, develop_tvm_path]
+            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(tvm_root, "src/tl"))
-            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"))
+            
+            install_cutlass_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "../..", "3rdparty", "cutlass")
+            develop_cutlass_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "../../..", "3rdparty", "cutlass")
+            cutlass_root = next((path for path in [install_cutlass_path, develop_cutlass_path]
+                                if os.path.exists(path) and path not in sys.path), None)
+
+            tl_template_path = tl_template_path = os.environ["TL_TEMPLATE_PATH"] if "TL_TEMPLATE_PATH" in os.environ else osp.abspath(osp.join(tilelang_root, "src"))
+
+            cutlass_path = os.environ["TL_CUTLASS_PATH"] if "TL_CUTLASS_PATH" in os.environ else osp.abspath(osp.join(cutlass_root, "include"))
 
             command += [
                 "-I" + tl_template_path,

bitblas/ops/base_scheduler.py

@@ -1,8 +1,12 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from tvm import IRModule
 from tvm.tir import PrimFunc
 from typing import Union, Callable
 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
 

bitblas/ops/general_flashatten/tilelang/flashatten.py

@@ -2,8 +2,9 @@
 # Licensed under the MIT License.
 
 from bitblas import tvm as tvm
+from bitblas import tilelang as tilelang
 from bitblas.ops.base_scheduler import BaseScheduler
-import tvm.tl.language as T
+import tilelang.language as T
 from dataclasses import dataclass
 from typing import Optional
 import logging

bitblas/ops/general_matmul/tilelang/dense/matmul_tensorcore.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,

bitblas/ops/general_matmul/tilelang/dequantize/block_primitive_tensorcore.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, Literal
 from bitblas.ops.base_scheduler import BaseScheduler
 from bitblas.base.arch import TileDevice

bitblas/ops/general_matmul/tilelang/dequantize/finegrained_primitive_tensorcore.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, Literal
 from bitblas.tl.utils import (
     get_mma_micro_size,  # noqa: F401

bitblas/ops/operator.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 from abc import ABC, abstractmethod
 from bitblas import tvm
-from tvm import tl
+from bitblas import tilelang
 from tvm import IRModule
 from tvm.runtime.module import Module
 from tvm.target import Target
@@ -196,7 +196,7 @@ def tvm_callback_cuda_postproc(code, _):
                                     "tir.disable_cse_tir": True,
                                     **(self.pass_context if self.pass_context else {})
                                 }):
-                            rt_mod = tl.lower(tl_prim_func, target=target, runtime_only=True)
+                            rt_mod = tilelang.lower(tl_prim_func, target=target, runtime_only=True)
                     else:
                         raise ValueError(f"Unsupported backend: {self.backend}")
             except Exception as build_runtime_error:  # noqa: F841

bitblas/tl/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 tvm as tvm
+from bitblas import tilelang as tilelang
+import tilelang.language as T
 
 from typing import Union
 from bitblas.ops.common import TransformKind

bitblas/tl/tuner.py

@@ -1,14 +1,14 @@
 # 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
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import List, Tuple, Optional, Literal
 from tvm import tir, IRModule
 from tvm.runtime import Module
 from tvm.tir import Schedule
-import tvm.tl as tl
 from bitblas.ops.base_scheduler import BaseScheduler
 from bitblas.base.arch import CUDA
 from bitblas.base.utils import get_dummy_input_arrays
@@ -133,7 +133,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/utils.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 tvm import arith
 from tvm import DataType
-import tvm.tl.language as T
+import tilelang.language as T
 from typing import Union, Literal
 from .mma_layout import (
     ldmatrix_32x8_to_shared_16x16_layout,