Build ARM64 support for lINUX

johnnynunez · johnnynunez · commit a1836e83c874 · 2025-02-11T10:54:24.000+01:00
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -156,6 +156,61 @@ jobs:
           destination_dir: built-wheels/${{ inputs.build_id }}/${{ !matrix.designated && 'matrix/' || '' }}
           endpoint: http://botmaster.tgr:9000
 
+  build_linux_arm64:
+    name: Build Linux
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        cuda: [ '', cuda ]
+        llvm: [ '', llvm ]
+        gl: [ '', gl ]
+        vk: [ '', vk ]
+        python: ${{ fromJSON(inputs.python) }}
+        include:
+          - { cuda: cuda, llvm: llvm, gl: gl, vk: vk, designated: designated }
+        exclude:
+          - { llvm: '', cuda: cuda }
+    runs-on: [ ubuntu24.04-arm ]
+    env:
+      TAICHI_CMAKE_ARGS: >-
+        -DTI_WITH_CUDA:BOOL=${{ matrix.cuda && 'ON' || 'OFF' }}
+        -DTI_WITH_LLVM:BOOL=${{ matrix.llvm && 'ON' || 'OFF' }}
+        -DTI_WITH_OPENGL:BOOL=${{ matrix.gl && 'ON' || 'OFF' }}
+        -DTI_WITH_VULKAN:BOOL=${{ matrix.vk && 'ON' || 'OFF' }}
+        -DTI_WITH_METAL:BOOL=OFF
+        -DTI_WITH_BACKTRACE:BOOL=ON
+        -DTI_BUILD_TESTS:BOOL=ON
+
+    steps:
+      - name: Workaround checkout Needed single revision issue
+        run: git submodule foreach 'git rev-parse HEAD > /dev/null 2>&1 || rm -rf $PWD' || true
+
+      - uses: actions/checkout@v3
+        with:
+          submodules: 'recursive'
+          fetch-depth: '0'
+
+      - name: Build
+        run: |
+          . .github/workflows/scripts/common-utils.sh
+          ci-docker-run-gpu \
+            -v $(pwd):/home/dev/taichi \
+            registry.botmaster.tgr/taichi-build-cuda:${{ env.CI_IMAGE_VERSION }} \
+            /home/dev/taichi/build.py ${NIGHTLY:+--nightly} --python=${{ matrix.python }} ${SHOULD_TAG_CONFIG:+--tag-config}
+        env:
+          SHOULD_TAG_CONFIG: ${{ !matrix.designated && 'yes' || '' }}
+
+      - name: Upload Built Wheel
+        uses: shallwefootball/s3-upload-action@v1.3.3
+        with:
+          aws_key_id: ${{ secrets.BOT_MINIO_ACCESS_KEY }}
+          aws_secret_access_key: ${{ secrets.BOT_MINIO_SECRET_KEY }}
+          aws_bucket: built-wheels
+          source_dir: dist
+          destination_dir: built-wheels/${{ inputs.build_id }}/${{ !matrix.designated && 'matrix/' || '' }}
+          endpoint: http://botmaster.tgr:9000
+
   build_amdgpu_linux:
     name: Build AMDGPU
     timeout-minutes: 30
@@ -350,4 +405,4 @@ jobs:
         with:
           name: libtaichi_c_api.iOS.a
           path: 'dist/C-API-iOS/*.a'
-          retention-days: 7
+          retention-days: 7
diff --git a/.github/workflows/initiator.yaml b/.github/workflows/initiator.yaml
@@ -36,7 +36,7 @@ jobs:
 
   preparation:
     name: Preparation
-    runs-on: [self-hosted, Linux]
+    runs-on: [self-hosted, Linux, ubuntu-24.04-arm]
     outputs:
       date: ${{ steps.gather.outputs.date }}
       short_sha: ${{ steps.gather.outputs.short_sha }}
@@ -58,7 +58,7 @@ jobs:
     with:
       build_id: ${{ needs.preparation.outputs.build_id }}
       nightly: false
-      python: '["3.9", "3.10", "3.11"]'
+      python: '["3.10", "3.11", "3.12", "3.13"]'
     secrets:
       BOT_MINIO_ACCESS_KEY: ${{ secrets.BOT_MINIO_ACCESS_KEY }}
       BOT_MINIO_SECRET_KEY: ${{ secrets.BOT_MINIO_SECRET_KEY }}
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -34,7 +34,7 @@ jobs:
     if: github.repository_owner == 'taichi-dev'
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
 
       - name: Save new version
         run: |
@@ -85,7 +85,59 @@ jobs:
         -DTI_BUILD_TESTS:BOOL=ON
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
+        with:
+          submodules: 'recursive'
+          fetch-depth: '0'
+
+      - name: Prepare Environment
+        run: |
+          . .github/workflows/scripts/common-utils.sh
+          prepare-build-cache
+          echo CI_DOCKER_RUN_EXTRA_ARGS="-v $(pwd):/home/dev/taichi" >> $GITHUB_ENV
+
+      - name: Build
+        run: |
+          . .github/workflows/scripts/common-utils.sh
+
+          ci-docker-run-gpu --name taichi-build \
+            registry.botmaster.tgr/taichi-build-cuda:${{ env.CI_IMAGE_VERSION }} \
+            /home/dev/taichi/.github/workflows/scripts/build.py $NIGHTLY
+
+      - name: Archive Wheel Artifacts
+        uses: actions/upload-artifact@v3
+        with:
+          name: dist
+          path: dist/*.whl
+          retention-days: 20
+
+      - name: Test
+        run: |
+          . .github/workflows/scripts/common-utils.sh
+
+          ci-docker-run-gpu --name taichi-test \
+             registry.botmaster.tgr/taichi-test-cuda:${{ env.CI_IMAGE_VERSION }} \
+             /home/dev/taichi/.github/workflows/scripts/unix_test.sh
+        env:
+          PY: ${{ matrix.python }}
+
+  build_and_test_linux_arm:
+    name: Build and Upload (linux Arm only)
+    needs: matrix_prep
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJson(needs.matrix_prep.outputs.matrix) }}
+    runs-on: [ ubuntu24.04-arm, cuda, vulkan, cn, release, sm70 ]
+    env:
+      PY: ${{ matrix.python }}
+      PROJECT_NAME: ${{ matrix.name }}
+      TAICHI_CMAKE_ARGS: >-
+        -DTI_WITH_OPENGL:BOOL=ON
+        -DTI_WITH_VULKAN:BOOL=ON
+        -DTI_BUILD_TESTS:BOOL=ON
+
+    steps:
+      - uses: actions/checkout@v4
         with:
           submodules: 'recursive'
           fetch-depth: '0'
@@ -136,7 +188,7 @@ jobs:
   #       -DTI_WITH_OPENGL:BOOL=OFF
   #       -DTI_BUILD_TESTS:BOOL=ON
   #   steps:
-  #     - uses: actions/checkout@v3
+  #     - uses: actions/checkout@v4
   #       with:
   #         submodules: 'recursive'
   #         fetch-depth: '0'
@@ -188,7 +240,7 @@ jobs:
         -DTI_BUILD_TESTS:BOOL=ON
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           submodules: 'recursive'
           fetch-depth: '0'
@@ -243,12 +295,12 @@ jobs:
           git config --system core.longpaths true
           git submodule foreach 'git rev-parse HEAD > /dev/null 2>&1 || rm -rf $PWD'
 
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           submodules: 'recursive'
           fetch-depth: '0'
 
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python }}
 
@@ -289,11 +341,11 @@ jobs:
 
   upload_to_pypi:
     name: Upload release to PyPI
-    needs: [build_and_test_linux, build_and_test_m1, build_and_test_windows]
+    needs: [build_and_test_linux, build_and_test_m1, build_and_test_windows, build_and_test_linux_arm]
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
         with:
           python-version: 3.9
 
@@ -331,11 +383,11 @@ jobs:
     runs-on: ubuntu-latest
     if: ${{ github.event.inputs.version }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v5
         with:
           python-version: 3.9
 
diff --git a/.github/workflows/scripts/ti_build/dep.py b/.github/workflows/scripts/ti_build/dep.py
@@ -153,6 +153,9 @@ def download_dep(url, outdir, *, strip=0, force=False, args=None, plain=False, e
     elif name.endswith(".tar.gz") or name.endswith(".tgz"):
         outdir.mkdir(parents=True, exist_ok=True)
         tar("-xzf", local_cached, "-C", outdir, f"--strip-components={strip}")
+    elif name.endswith(".tar.xz"):
+        outdir.mkdir(parents=True, exist_ok=True)
+        tar("-xJf", local_cached, "-C", outdir, f"--strip-components={strip}")
     elif name.endswith(".sh"):
         bash(local_cached, *args)
     elif "." not in name and args is not None:
diff --git a/.github/workflows/scripts/ti_build/entry.py b/.github/workflows/scripts/ti_build/entry.py
@@ -51,11 +51,17 @@ 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"])
+            extra.extend(["-p", "manylinux_2_28_x86_64"])
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        extra.extend(["-p", "manylinux_2_28_aarch64"])
+    else:
+        extra.extend(["-p", "manylinux_2_28_x86_64"])
+
 
     python("setup.py", "clean")
     python("misc/make_changelog.py", "--ver", "origin/master", "--repo_dir", "./", "--save")
diff --git a/.github/workflows/scripts/ti_build/llvm.py b/.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"
diff --git a/.github/workflows/scripts/ti_build/python.py b/.github/workflows/scripts/ti_build/python.py
@@ -15,22 +15,24 @@
 from .misc import banner, get_cache_home, path_prepend, info
 from .tinysh import Command, sh
 
-
 # -- 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/24.11.3-0/Miniforge3-24.11.3-0-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/24.11.3-0/Miniforge3-24.11.3-0-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/24.11.3-0/Miniforge3-24.11.3-0-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/24.11.3-0/Miniforge3-24.11.3-0-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/24.11.3-0/Miniforge3-Windows-x86_64.exe"
         )
         download_dep(
             url,
@@ -82,8 +84,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 +94,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))
 
diff --git a/.github/workflows/scripts/ti_build/sccache.py b/.github/workflows/scripts/ti_build/sccache.py
@@ -31,18 +31,22 @@ 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.9.1/sccache-v0.9.1-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.9.1/sccache-v0.9.1-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"
+                "https://github.com/mozilla/sccache/releases/download/v0.9.1/sccache-v0.9.1-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"
+                "https://github.com/mozilla/sccache/releases/download/v0.9.1/sccache-v0.9.1-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.9.1/sccache-v0.9.1-x86_64-pc-windows-msvc.tar.gz"
         else:
             raise RuntimeError(f"Unsupported platform: {u.system} {u.machine}")
 
diff --git a/.github/workflows/scripts/ti_build/vulkan.py b/.github/workflows/scripts/ti_build/vulkan.py
@@ -12,23 +12,34 @@
 
 
 # -- code --
-@banner("Setup Vulkan 1.3.236.0")
+@banner("Setup Vulkan 1.4.304.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.304.0/linux/vulkansdk-linux-x86_64-1.4.304.0.tar.xz"
+        prefix = get_cache_home() / "vulkan-1.4.304.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) == ("Darwin", "arm64"):
+    elif (u.system, u.machine) in (("Linux", "arm64"), ("Linux", "aarch64")):
+        url = (
+            "https://github.com/johnnynunez/vulkan-sdk-arm/releases/download/1.4.304.0/"
+            "vulkansdk-linux-arm64-ubuntu-22.04-arm-1.4.304.0.zip"
+        )
+        prefix = get_cache_home() / "vulkan-1.4.304.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", "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.304.0/windows/VulkanSDK-1.4.304.0-Installer.exe"
+        prefix = get_cache_home() / "vulkan-1.4.304.0"
         download_dep(
             url,
             prefix,
diff --git a/docs/lang/articles/contribution/dev_install.md b/docs/lang/articles/contribution/dev_install.md
diff --git a/python/taichi/_version_check.py b/python/taichi/_version_check.py
