[Build] Update Docker distribution scripts for manylinux compatibility

- Changed base image from `tilelang-builder:18.04` to `tilelang-builder:manylinux` in both local and PyPI distribution scripts.
- Updated Dockerfile references to use `pypi.manylinux.Dockerfile`.
- Added `--gpus all` flag to the Docker run command to enable GPU support during execution.

Author: LeiWang1999

maint/scripts/docker_local_distribute.sh

@@ -1,9 +1,9 @@
 # Get the CUDA version from the command line
-IMAGE="tilelang-builder:18.04"
-docker build . -f "$(dirname "${BASH_SOURCE[0]}")/pypi.Dockerfile" --tag ${IMAGE}
+IMAGE="tilelang-builder:manylinux"
+docker build . -f "$(dirname "${BASH_SOURCE[0]}")/pypi.manylinux.Dockerfile" --tag ${IMAGE}
 
 install_pip="python3.8 -m pip install --upgrade pip && python3.8 -m pip install -r requirements-build.txt"
 
 tox_command="python3.8 -m tox -e py38,py39,py310,py311,py312"
 
-docker run --rm -v $(pwd):/tilelang ${IMAGE} /bin/bash -c "$install_pip && $tox_command"
+docker run --rm --gpus all -v $(pwd):/tilelang ${IMAGE} /bin/bash -c "$install_pip && $tox_command"

maint/scripts/docker_pypi_distribute.sh

@@ -1,9 +1,9 @@
 # Get the CUDA version from the command line
-IMAGE="tilelang-builder:18.04"
-docker build . -f "$(dirname "${BASH_SOURCE[0]}")/pypi.Dockerfile" --tag ${IMAGE}
+IMAGE="tilelang-builder:manylinux"
+docker build . -f "$(dirname "${BASH_SOURCE[0]}")/pypi.manylinux.Dockerfile" --tag ${IMAGE}
 
 install_pip="python3.8 -m pip install --upgrade pip && python3.8 -m pip install -r requirements-build.txt"
 
-tox_command="python3.8 -m tox -e py38-pypi,py39-pypi,py310-pypi,py311-pypi,py312-pypi,audit_2_27"
+tox_command="python3.8 -m tox -e py38-pypi,py39-pypi,py310-pypi,py311-pypi,py312-pypi"
 
-docker run --rm -v $(pwd):/tilelang ${IMAGE} /bin/bash -c "$install_pip && $tox_command"
+docker run --rm --gpus all -v $(pwd):/tilelang ${IMAGE} /bin/bash -c "$install_pip && $tox_command"

maint/scripts/pypi.manylinux.Dockerfile

@@ -0,0 +1,26 @@
+FROM pytorch/manylinux-builder:cuda12.1
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    TZ=Etc/UTC
+
+RUN set -eux; \
+    yum -y update && yum install -y \
+        zlib-devel openssl-devel \
+        libedit-devel libxml2-devel \
+        bzip2 bzip2-devel xz xz-devel \
+        epel-release
+
+RUN set -eux; \
+    conda create -n py38 python=3.8 -y && \
+    conda create -n py39 python=3.9 -y && \
+    conda create -n py310 python=3.10 -y && \
+    conda create -n py311 python=3.11 -y && \
+    conda create -n py312 python=3.12 -y && \
+    ln -sf /opt/conda/envs/py38/bin/python3.8 /usr/bin/python3.8 && \
+    ln -sf /opt/conda/envs/py39/bin/python3.9 /usr/bin/python3.9 && \
+    ln -sf /opt/conda/envs/py310/bin/python3.10 /usr/bin/python3.10 && \
+    ln -sf /opt/conda/envs/py311/bin/python3.11 /usr/bin/python3.11 && \
+    ln -sf /opt/conda/envs/py312/bin/python3.12 /usr/bin/python3.12 && \
+    conda install -y cmake patchelf
+
+WORKDIR /tilelang

pyproject.toml

@@ -1,7 +1,6 @@
 [build-system]
 requires = [
     "build",
-    "cmake>=3.26",
     "packaging",
     "setuptools>=61",
     "wheel",

tilelang/transform/add_bufstore_wrapper.py

@@ -1,7 +1,7 @@
 from tvm.tir import (BufferStore, For, AttrStmt, ForKind, Var, PrimFunc, BufferLoad, Buffer, IntImm)
 from tvm.tir.stmt_functor import ir_transform, post_order_visit
 from tvm.tir.transform import prim_func_pass
-
+from typing import Tuple, List, Dict
 
 def AddWrapperForSingleBufStore():
     """
@@ -41,7 +41,7 @@ def visit_variable(node):
             post_order_visit(operation, visit_variable)
             return used_variables
 
-        def collect_buffer_accesses(statement) -> tuple[list[Buffer], list[Buffer]]:
+        def collect_buffer_accesses(statement) -> Tuple[List[Buffer], List[Buffer]]:
             """
             Categorizes buffers accessed in the statement by their scope.
 
@@ -68,7 +68,7 @@ def visit_buffer_access(node):
                     local_buffers.append(buffer)
             return local_buffers, fragment_buffers
 
-        def collect_buffer_indices(statement) -> dict[Buffer, list[int]]:
+        def collect_buffer_indices(statement) -> Dict[Buffer, List[int]]:
             """
             Maps each buffer to its access indices.
 

tilelang/utils/sparse.py

@@ -1,7 +1,7 @@
 import os
 import torch
 import warnings
-from typing import Optional
+from typing import Optional, Tuple
 from tilelang.contrib import nvcc
 from torch.utils.cpp_extension import load, _import_module_from_library
 from tilelang import env
@@ -44,7 +44,7 @@ def _get_cached_lib():
 
 
 def compress_sm90(A: torch.Tensor, block_k: int,
-                  transposed: bool) -> tuple[torch.Tensor, torch.Tensor]:
+                  transposed: bool) -> Tuple[torch.Tensor, torch.Tensor]:
     if block_k > 128:
         block_k = 128
         # Ref: https://github.com/NVIDIA/cutlass/blob/c2ad7c5b20f131c4ba33601860f1da3f9c9df0f3/include/cutlass/gemm/collective/builders/sm90_sparse_gmma_builder.inl#L145-L146
@@ -56,7 +56,7 @@ def compress_sm90(A: torch.Tensor, block_k: int,
     return compress_lib.compress_sm90(A, block_k, transposed)
 
 
-def compress_sm80(A: torch.Tensor, transposed: bool) -> tuple[torch.Tensor, torch.Tensor]:
+def compress_sm80(A: torch.Tensor, transposed: bool) -> Tuple[torch.Tensor, torch.Tensor]:
     try:
         from torch.sparse import to_sparse_semi_structured, SparseSemiStructuredTensor
     except ImportError as err:
@@ -76,7 +76,7 @@ def compress_sm80(A: torch.Tensor, transposed: bool) -> tuple[torch.Tensor, torc
 def compress(A: torch.Tensor,
              transposed: bool,
              arch: Optional[str] = None,
-             **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+             **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Compress a tensor using the appropriate method based on the CUDA architecture.
     """

tox.ini

@@ -1,30 +1,36 @@
 [tox]
 envlist = py38,py39,py310,py311,py312
-isolated_build = True
+isolated_build = False
 
 [testenv:py{38,39,310,311,312}]
+skip_install = false
 deps =
     wheel
     build
+setenv =
+    PYTHON_EXECUTABLE = {envpython}
+    Python3_EXECUTABLE = {envpython}
 commands =
     python -m build --wheel -o {toxinidir}/dist
 
-
 [testenv:py{38,39,310,311,312}-pypi]
+skip_install = false
 setenv =
     PYPI_BUILD = TRUE
+    PYTHON_EXECUTABLE = {envpython}
+    Python3_EXECUTABLE = {envpython}
 commands =
-    python setup.py bdist_wheel
+    python setup.py bdist_wheel --plat-name=manylinux2014_x86_64
 
-[testenv:audit_2_27]
+[testenv:audit_manylinux2014]
 skip_install = true
 allowlist_externals =
     bash
 deps =
     auditwheel
     patchelf
 commands =
-    bash -c 'auditwheel repair -L=/lib --exclude=/usr/local/cuda* --exclude=libcuda.so.1 --plat=manylinux_2_27_x86_64 dist/*'
+    bash -c 'auditwheel repair -L=/lib --exclude=/usr/local/cuda* --exclude=libcuda.so.1 --plat=manylinux2014_x86_64 dist/*'
 
 [testenv:py38]
 basepython = python3.8