[build] Blackwell Support LLVM20 + ARM Support

.github/workflows/scripts/check_clang_tidy.sh

@@ -1,14 +1,46 @@
 #!/bin/bash
 
-# TODO: move inside docker
-wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm15/taichi-llvm-15-linux.zip
-unzip taichi-llvm-15-linux.zip && rm taichi-llvm-15-linux.zip
-export LLVM_DIR="$PWD/taichi-llvm-15"
-
-pushd taichi
-python3 -m pip install -r requirements_dev.txt
-
-CI_SETUP_CMAKE_ARGS=$1
-export CI_SETUP_CMAKE_ARGS
-python3 ./scripts/run_clang_tidy.py $PWD/taichi -clang-tidy-binary clang-tidy-10 -header-filter=$PWD/taichi -j2
-popd
+# we need to generate a special file "compile_commands.json"
+# to do this, we need to run cmake with the following options:
+# -DCMAKE_EXPORT_COMPILE_COMMANDS=ON
+# I'm fairly sure that we don't need to run the actual build,
+# but it's not obvious to me how to do this. So, I'm just going
+# to run a full build for now, and we can FIXME this later.
+
+set -ex
+
+sudo apt-get update
+sudo apt-get install -y clang-tidy-20
+git submodule update --init --recursive
+
+sudo apt install -y \
+    freeglut3-dev \
+    libglfw3-dev \
+    libglm-dev \
+    libglu1-mesa-dev \
+    libwayland-dev \
+    libx11-xcb-dev \
+    libxcb-dri3-dev \
+    libxcb-ewmh-dev \
+    libxcb-keysyms1-dev \
+    libxcb-randr0-dev \
+    libxcursor-dev \
+    libxi-dev \
+    libxinerama-dev \
+    libxrandr-dev \
+    pybind11-dev \
+    libc++-15-dev \
+    libc++abi-15-dev \
+    clang-15 \
+    libclang-common-15-dev \
+    libclang-cpp15 \
+    libclang1-15 \
+    cmake \
+    ninja-build \
+    python3-dev \
+    python3-pip
+
+pip install scikit-build
+export TAICHI_CMAKE_ARGS="-DCMAKE_EXPORT_COMPILE_COMMANDS=ON"
+./build.py wheel
+python ./scripts/run_clang_tidy.py $PWD/taichi -clang-tidy-binary clang-tidy-14 -header-filter=$PWD/taichi -j2

.github/workflows/scripts/ti_build/compiler.py

@@ -33,7 +33,7 @@ def setup_clang(as_compiler=True) -> None:
     """
     u = platform.uname()
     if u.system == "Linux":
-        for v in ("", "-14", "-13", "-12", "-11", "-10"):
+        for v in ("", "-20", "-19", "-18", "-17", "-16", "-15", "-14", "-13", "-12", "-11", "-10"):
             clang = shutil.which(f"clang{v}")
             if clang is not None:
                 clangpp = shutil.which(f"clang++{v}")

.github/workflows/scripts/ti_build/entry.py

@@ -52,11 +52,19 @@ def build_wheel(python: Command, pip: Command) -> None:
     elif wheel_tag:
         proj_tags.extend(["egg_info", f"--tag-build={wheel_tag}"])
 
-    if platform.system() == "Linux":
+    u = platform.uname()
+    if (u.system, u.machine) == ("Linux", "x86_64"):
         if is_manylinux2014():
             extra.extend(["-p", "manylinux2014_x86_64"])
         else:
             extra.extend(["-p", "manylinux_2_27_x86_64"])
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        if is_manylinux2014():
+            extra.extend(["-p", "manylinux2014_aarch64"])
+        else:
+            extra.extend(["-p", "manylinux_2_27_aarch64"])
+    else:
+        extra.extend(["-p", "manylinux2014_x86_64"])
 
     python("setup.py", "clean")
     python("misc/make_changelog.py", "--ver", "origin/master", "--repo_dir", "./", "--save")

.github/workflows/scripts/ti_build/llvm.py

@@ -19,7 +19,7 @@ def setup_llvm() -> None:
     Download and install LLVM.
     """
     u = platform.uname()
-    if u.system == "Linux":
+    if (u.system, u.machine) == ("Linux", "x86_64"):
         if cmake_args.get_effective("TI_WITH_AMDGPU"):
             out = get_cache_home() / "llvm15-amdgpu-005"
             url = "https://github.com/GaleSeLee/assets/releases/download/v0.0.5/taichi-llvm-15.0.0-linux.zip"
@@ -31,6 +31,11 @@ def setup_llvm() -> None:
             out = get_cache_home() / "llvm15"
             url = "https://github.com/taichi-dev/taichi_assets/releases/download/llvm15/taichi-llvm-15-linux.zip"
         download_dep(url, out, strip=1)
+
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        out = get_cache_home() / "llvm15-manylinux2014"
+        # FIXME: ARM LLVM!
+        pass
     elif (u.system, u.machine) == ("Darwin", "arm64"):
         out = get_cache_home() / "llvm15-m1-nozstd"
         url = "https://github.com/taichi-dev/taichi_assets/releases/download/llvm15/taichi-llvm-15-m1-nozstd.zip"
@@ -48,4 +53,8 @@ def setup_llvm() -> None:
 
     # We should use LLVM toolchains shipped with OS.
     # path_prepend('PATH', out / 'bin')
-    os.environ["LLVM_DIR"] = str(out)
+    if (u.system, u.machine) not in (("Linux", "arm64"), ("Linux", "aarch64")):
+        os.environ["LLVM_DIR"] = str(out)
+        # os.environ["LLVM_DIR"] = "/usr/lib/llvm-20/cmake"
+        # os.environ["CUDA_HOME"] = "/usr/local/cuda"
+        # os.environ["CPATH"] = "/usr/local/cuda/include"

.github/workflows/scripts/ti_build/ospkg.py

@@ -24,7 +24,7 @@
     "liblz4-dev",
     "libpng-dev",
     "libssl-dev",
-    "libtinfo-dev",
+    "libncurses-dev",
     "libwayland-dev",
     "libx11-xcb-dev",
     "libxcb-dri3-dev",

.github/workflows/scripts/ti_build/python.py

@@ -17,20 +17,23 @@
 
 
 # -- code --
-def setup_mambaforge(prefix):
+def setup_miniforge(prefix):
     u = platform.uname()
-    if u.system == "Linux":
-        url = "https://github.com/conda-forge/miniforge/releases/download/23.1.0-1/Mambaforge-23.1.0-1-Linux-x86_64.sh"
+    if (u.system, u.machine) == ("Linux", "x86_64"):
+        url = "https://github.com/conda-forge/miniforge/releases/download/25.3.0-3/Miniforge3-25.3.0-3-Linux-x86_64.sh"
+        download_dep(url, prefix, args=["-bfp", str(prefix)])
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        url = "https://github.com/conda-forge/miniforge/releases/download/25.3.0-3/Miniforge3-25.3.0-3-Linux-aarch64.sh"
         download_dep(url, prefix, args=["-bfp", str(prefix)])
     elif (u.system, u.machine) == ("Darwin", "arm64"):
-        url = "https://github.com/conda-forge/miniforge/releases/download/23.1.0-1/Mambaforge-23.1.0-1-MacOSX-arm64.sh"
+        url = "https://github.com/conda-forge/miniforge/releases/download/25.3.0-3/Miniforge3-25.3.0-3-MacOSX-arm64.sh"
         download_dep(url, prefix, args=["-bfp", str(prefix)])
     elif (u.system, u.machine) == ("Darwin", "x86_64"):
-        url = "https://github.com/conda-forge/miniforge/releases/download/23.1.0-1/Mambaforge-23.1.0-1-MacOSX-x86_64.sh"
+        url = "https://github.com/conda-forge/miniforge/releases/download/25.3.0-3/Miniforge3-25.3.0-3-MacOSX-x86_64.sh"
         download_dep(url, prefix, args=["-bfp", str(prefix)])
     elif u.system == "Windows":
         url = (
-            "https://github.com/conda-forge/miniforge/releases/download/23.1.0-1/Mambaforge-23.1.0-1-Windows-x86_64.exe"
+            "https://github.com/conda-forge/miniforge/releases/download/25.3.0-3/Miniforge3-25.3.0-3-Windows-x86_64.exe"
         )
         download_dep(
             url,
@@ -82,8 +85,8 @@ def setup_python(version: str) -> Tuple[Command, Command]:
 
     windows = platform.system() == "Windows"
 
-    prefix = get_cache_home() / "mambaforge"
-    setup_mambaforge(prefix)
+    prefix = get_cache_home() / "miniforge"
+    setup_miniforge(prefix)
 
     if windows:
         conda_path = prefix / "Scripts" / "conda.exe"
@@ -92,9 +95,9 @@ def setup_python(version: str) -> Tuple[Command, Command]:
 
     if not conda_path.exists():
         shutil.rmtree(prefix, ignore_errors=True)
-        setup_mambaforge(prefix)
+        setup_miniforge(prefix)
         if not conda_path.exists():
-            raise RuntimeError(f"Failed to setup mambaforge at {prefix}")
+            raise RuntimeError(f"Failed to setup miniforge at {prefix}")
 
     conda = sh.bake(str(conda_path))
 

.github/workflows/scripts/ti_build/sccache.py

@@ -31,18 +31,16 @@ def setup_sccache() -> Command:
         raise RuntimeError(f"Unsupported platform: {u.system} {u.machine}")
 
     if not exe.exists():
-        if u.system == "Linux":
-            url = "https://github.com/mozilla/sccache/releases/download/v0.4.1/sccache-v0.4.1-x86_64-unknown-linux-musl.tar.gz"
+        if (u.system, u.machine) == ("Linux", "x86_64"):
+            url = "https://github.com/mozilla/sccache/releases/download/v0.10.0/sccache-v0.10.0-x86_64-unknown-linux-musl.tar.gz"
+        elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+            url = "https://github.com/mozilla/sccache/releases/download/v0.10.0/sccache-v0.10.0-aarch64-unknown-linux-musl.tar.gz"
         elif (u.system, u.machine) == ("Darwin", "arm64"):
-            url = (
-                "https://github.com/mozilla/sccache/releases/download/v0.4.1/sccache-v0.4.1-aarch64-apple-darwin.tar.gz"
-            )
+            url = "https://github.com/mozilla/sccache/releases/download/v0.10.0/sccache-v0.10.0-aarch64-apple-darwin.tar.gz"
         elif (u.system, u.machine) == ("Darwin", "x86_64"):
-            url = (
-                "https://github.com/mozilla/sccache/releases/download/v0.4.1/sccache-v0.4.1-x86_64-apple-darwin.tar.gz"
-            )
+            url = "https://github.com/mozilla/sccache/releases/download/v0.10.0/sccache-v0.10.0-x86_64-apple-darwin.tar.gz"
         elif u.system == "Windows":
-            url = "https://github.com/mozilla/sccache/releases/download/v0.4.1/sccache-v0.4.1-x86_64-pc-windows-msvc.tar.gz"
+            url = "https://github.com/mozilla/sccache/releases/download/v0.10.0/sccache-v0.10.0-x86_64-pc-windows-msvc.tar.gz"
         else:
             raise RuntimeError(f"Unsupported platform: {u.system} {u.machine}")
 

.github/workflows/scripts/ti_build/vulkan.py

@@ -12,23 +12,32 @@
 
 
 # -- code --
-@banner("Setup Vulkan 1.3.236.0")
+@banner("Setup Vulkan 1.4.313.0")
 def setup_vulkan():
     u = platform.uname()
-    if u.system == "Linux":
-        url = "https://sdk.lunarg.com/sdk/download/1.3.236.0/linux/vulkansdk-linux-x86_64-1.3.236.0.tar.gz"
-        prefix = get_cache_home() / "vulkan-1.3.236.0"
+    if (u.system, u.machine) == ("Linux", "x86_64"):
+        url = "https://sdk.lunarg.com/sdk/download/1.4.313.0/linux/vulkansdk-linux-x86_64-1.4.313.0.tar.xz"
+        prefix = get_cache_home() / "vulkan-1.4.313.0"
         download_dep(url, prefix, strip=1)
         sdk = prefix / "x86_64"
         os.environ["VULKAN_SDK"] = str(sdk)
         path_prepend("PATH", sdk / "bin")
         path_prepend("LD_LIBRARY_PATH", sdk / "lib")
         os.environ["VK_LAYER_PATH"] = str(sdk / "etc" / "vulkan" / "explicit_layer.d")
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        url = "https://github.com/johnnynunez/vulkan-sdk-arm/releases/download/1.4.313.0/vulkansdk-linux-arm64-ubuntu-24.04-arm-1.4.313.0.zip"
+        prefix = get_cache_home() / "vulkan-1.4.313.0"
+        download_dep(url, prefix, strip=1)
+        sdk = prefix / "arm64"
+        os.environ["VULKAN_SDK"] = str(sdk)
+        path_prepend("PATH", sdk / "bin")
+        path_prepend("LD_LIBRARY_PATH", sdk / "lib")
+        os.environ["VK_LAYER_PATH"] = str(sdk / "etc" / "vulkan" / "explicit_layer.d")
     # elif (u.system, u.machine) == ("Darwin", "arm64"):
     # elif (u.system, u.machine) == ("Darwin", "x86_64"):
     elif (u.system, u.machine) == ("Windows", "AMD64"):
-        url = "https://sdk.lunarg.com/sdk/download/1.3.236.0/windows/VulkanSDK-1.3.236.0-Installer.exe"
-        prefix = get_cache_home() / "vulkan-1.3.236.0"
+        url = "https://sdk.lunarg.com/sdk/download/1.4.313.0/windows/VulkanSDK-1.4.313.0-Installer.exe"
+        prefix = get_cache_home() / "vulkan-1.4.313.0"
         download_dep(
             url,
             prefix,