[CI] Use auditwheel to generate manylinux wheels (#251)

* use auditwheel to get correct manylinux wheels

* fix

* make py3.8 happy

* trivial updates

* Add typing.Tuple import and update annotations

* fmt

* Remove unused import and update type hints

* lint fix

---------
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>

.gitignore

@@ -19,6 +19,7 @@
 debug/
 build/
 dist/
+wheelhouse/
 __pycache__
 nnfusion.tar.gz
 

maint/scripts/docker_local_distribute.sh

 # Get the CUDA version from the command line
-IMAGE=nvidia/cuda:12.1.0-devel-ubuntu18.04
+IMAGE="tilelang-builder:18.04"
+docker build . -f "$(dirname "${BASH_SOURCE[0]}")/pypi.Dockerfile" --tag ${IMAGE}
 
-docker pull ${IMAGE}
-
-apt_command="apt update && apt install -y software-properties-common && add-apt-repository -y ppa:deadsnakes/ppa && apt update && apt install -y wget curl libtinfo-dev zlib1g-dev libssl-dev build-essential libedit-dev libxml2-dev"
-
-install_python_env="curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b -p ~/miniconda3 && export PATH=~/miniconda3/bin/:$PATH && 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 -s ~/miniconda3/envs/py38/bin/python3.8 /usr/bin/python3.8 && ln -s ~/miniconda3/envs/py39/bin/python3.9 /usr/bin/python3.9 && ln -s ~/miniconda3/envs/py310/bin/python3.10 /usr/bin/python3.10 && ln -s ~/miniconda3/envs/py311/bin/python3.11 /usr/bin/python3.11 && ln -s ~/miniconda3/envs/py312/bin/python3.12 /usr/bin/python3.12"
-
-install_cmake="wget https://github.com/Kitware/CMake/releases/download/v3.28.4/cmake-3.28.4-linux-x86_64.tar.gz && tar -xvzf cmake-*.tar.gz && rm cmake-*.tar.gz && cd cmake-* && cp bin/* /usr/local/bin/ && mv share/* /usr/local/share/ && mv man/* /usr/local/man/ && hash -r && cd /tilelang && export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH"
-
-install_pip="python3.8 -m pip install  --upgrade pip && python3.8 -m pip install -r requirements-build.txt"
+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"
 

maint/scripts/docker_pypi_distribute.sh

 # Get the CUDA version from the command line
-IMAGE=nvidia/cuda:12.1.0-devel-ubuntu18.04
-
-docker pull ${IMAGE}
-
-apt_command="apt update && apt install -y software-properties-common && add-apt-repository -y ppa:deadsnakes/ppa && apt update && apt install -y wget curl libtinfo-dev zlib1g-dev libssl-dev build-essential libedit-dev libxml2-dev"
-
-install_python_env="curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b -p ~/miniconda3 && export PATH=~/miniconda3/bin/:$PATH && 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 -s ~/miniconda3/envs/py38/bin/python3.8 /usr/bin/python3.8 && ln -s ~/miniconda3/envs/py39/bin/python3.9 /usr/bin/python3.9 && ln -s ~/miniconda3/envs/py310/bin/python3.10 /usr/bin/python3.10 && ln -s ~/miniconda3/envs/py311/bin/python3.11 /usr/bin/python3.11 && ln -s ~/miniconda3/envs/py312/bin/python3.12 /usr/bin/python3.12"
-
-install_cmake="wget https://github.com/Kitware/CMake/releases/download/v3.28.4/cmake-3.28.4-linux-x86_64.tar.gz && tar -xvzf cmake-*.tar.gz && rm cmake-*.tar.gz && cd cmake-* && cp bin/* /usr/local/bin/ && mv share/* /usr/local/share/ && mv man/* /usr/local/man/ && hash -r && cd /tilelang && export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH"
+IMAGE="tilelang-builder:18.04"
+docker build . -f "$(dirname "${BASH_SOURCE[0]}")/pypi.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"
+tox_command="python3.8 -m tox -e py38-pypi,py39-pypi,py310-pypi,py311-pypi,py312-pypi,audit_2_27"
 
-docker run --rm -v $(pwd):/tilelang ${IMAGE} /bin/bash -c "$apt_command && $install_python_env && $install_cmake && $install_pip && $tox_command"
+docker run --rm -v $(pwd):/tilelang ${IMAGE} /bin/bash -c "$install_pip && $tox_command"

maint/scripts/pypi.Dockerfile

+FROM nvidia/cuda:12.1.0-devel-ubuntu18.04
+
+RUN set -eux; \
+    apt-get update; \
+    apt-get install -y wget curl libtinfo-dev zlib1g-dev libssl-dev build-essential libedit-dev libxml2-dev; \
+    curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh; \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /miniconda3; \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/miniconda3/bin/:$PATH
+
+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 -s /miniconda3/envs/py38/bin/python3.8 /usr/bin/python3.8; \
+    ln -s /miniconda3/envs/py39/bin/python3.9 /usr/bin/python3.9; \
+    ln -s /miniconda3/envs/py310/bin/python3.10 /usr/bin/python3.10; \
+    ln -s /miniconda3/envs/py311/bin/python3.11 /usr/bin/python3.11; \
+    ln -s /miniconda3/envs/py312/bin/python3.12 /usr/bin/python3.12; \
+    conda install -y cmake patchelf
+
+WORKDIR /tilelang
\ No newline at end of file

requirements-build.txt

@@ -6,3 +6,5 @@ setuptools>=61
 torch
 wheel
 tox
+auditwheel
+patchelf

setup.py

@@ -226,6 +226,15 @@ def setup_llvm_for_tvm():
     return extract_path, llvm_config_path
 
 
+def patch_libs(libpath):
+    """
+    tvm and tilelang libs are copied from elsewhere into wheels
+    and have a hard-coded rpath.
+    Set rpath to the directory of libs so auditwheel works well.
+    """
+    subprocess.run(['patchelf', '--set-rpath', '$ORIGIN', libpath])
+
+
 class TileLangBuilPydCommand(build_py):
     """Customized setuptools install command - builds TVM after setting up LLVM."""
 
@@ -302,6 +311,7 @@ class TileLangBuilPydCommand(build_py):
                     break
 
             if source_lib_file:
+                patch_libs(source_lib_file)
                 target_dir_release = os.path.join(self.build_lib, PACKAGE_NAME, "lib")
                 target_dir_develop = os.path.join(PACKAGE_NAME, "lib")
                 os.makedirs(target_dir_release, exist_ok=True)
@@ -453,6 +463,7 @@ class TileLangDevelopCommand(develop):
             if not os.path.exists(target_dir):
                 os.makedirs(target_dir)
             if os.path.exists(source_lib_file):
+                patch_libs(source_lib_file)
                 shutil.copy2(source_lib_file, target_dir)
                 # remove the original file
                 os.remove(source_lib_file)

tilelang/jit/adapter/utils.py

+from __future__ import annotations
+
 import re
 from typing import Union, Optional, Literal
 from tilelang import tvm as tvm

tilelang/utils/tensor.py

+from __future__ import annotations
 """The profiler and convert to torch utils"""
 from enum import Enum
 import torch

tox.ini

@@ -14,7 +14,17 @@ commands =
 setenv =
     PYPI_BUILD = TRUE
 commands =
-    python setup.py bdist_wheel --plat-name=manylinux1_x86_64
+    python setup.py bdist_wheel
+
+[testenv:audit_2_27]
+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/*'
 
 [testenv:py38]
 basepython = python3.8