[AscendNPU-IR] Support CANN-9.0.0

.github/workflows/build-npuir.yml

@@ -9,6 +9,10 @@ on:
         description: "Git ref to checkout (leave empty for default event ref)"
         type: string
         default: ''
+      ascend-npuir-commit:
+        description: "Override AscendNPU-IR commit hash for cache/build selection"
+        type: string
+        default: ''
     outputs:
       npuir-commit:
         description: "AscendNPU-IR commit hash"
@@ -42,7 +46,11 @@ jobs:
     - name: Get AscendNPU-IR commit ID
       id: npuir-version
       run: |
-        NPUIR_COMMIT=$(git rev-parse HEAD:3rdparty/AscendNPU-IR)
+        if [ -n "${{ inputs.ascend-npuir-commit }}" ]; then
+          NPUIR_COMMIT="${{ inputs.ascend-npuir-commit }}"
+        else
+          NPUIR_COMMIT=$(git rev-parse HEAD:3rdparty/AscendNPU-IR)
+        fi
         echo "commit=${NPUIR_COMMIT}" >> "$GITHUB_OUTPUT"
         echo "AscendNPU-IR commit: ${NPUIR_COMMIT}"
 
@@ -62,6 +70,10 @@ jobs:
       run: |
         git submodule update --init 3rdparty/AscendNPU-IR
         cd 3rdparty/AscendNPU-IR
+        if [ -n "${{ inputs.ascend-npuir-commit }}" ]; then
+          git fetch --depth=1 origin "${{ inputs.ascend-npuir-commit }}"
+          git checkout --detach "${{ inputs.ascend-npuir-commit }}"
+        fi
         git submodule update --init
 
     - name: Set up Python 3.11

.github/workflows/build-tilelang-npuir.yml

@@ -7,12 +7,23 @@ on:
         description: "Git ref to checkout"
         type: string
         required: true
+      cann-version:
+        description: "CANN compatibility version"
+        type: string
+        required: true
   workflow_dispatch:
     inputs:
       checkout-ref:
         description: "Git ref to checkout (leave empty for default)"
         type: string
         default: ''
+      cann-version:
+        description: "CANN compatibility version"
+        type: choice
+        options:
+          - 8.5.0
+          - 9.0.0.beta2
+        default: 8.5.0
 
 permissions:
   contents: read
@@ -23,18 +34,47 @@ env:
   TVM_INSTALL_DIR: tvm-prebuild
 
 jobs:
+  select-npuir-version:
+    runs-on: ubuntu-22.04
+    outputs:
+      npuir-commit: ${{ steps.resolve.outputs.npuir-commit }}
+      cann-version-slug: ${{ steps.resolve.outputs.cann-version-slug }}
+      wheel-build-number: ${{ steps.resolve.outputs.wheel-build-number }}
+    steps:
+    - name: Resolve AscendNPU-IR cache version
+      id: resolve
+      run: |
+        case "${{ inputs.cann-version }}" in
+          8.5.0)
+            echo "npuir-commit=31f690369d1247fbd5529a3f88b758f7d470ae4f" >> "$GITHUB_OUTPUT"
+            echo "cann-version-slug=cann-8.5.0" >> "$GITHUB_OUTPUT"
+            echo "wheel-build-number=850cann850" >> "$GITHUB_OUTPUT"
+            ;;
+          9.0.0.beta2)
+            echo "npuir-commit=139bb6c7aef6cf923babb88318aaec3697f2091f" >> "$GITHUB_OUTPUT"
+            echo "cann-version-slug=cann-9.0.0.beta2" >> "$GITHUB_OUTPUT"
+            echo "wheel-build-number=9000cann9000beta2" >> "$GITHUB_OUTPUT"
+            ;;
+          *)
+            echo "::error::Unsupported CANN version: ${{ inputs.cann-version }}"
+            exit 1
+            ;;
+        esac
+
   prebuild-npuir:
+    needs: select-npuir-version
     uses: ./.github/workflows/build-npuir.yml
     with:
       checkout-ref: ${{ inputs.checkout-ref }}
+      ascend-npuir-commit: ${{ needs.select-npuir-version.outputs.npuir-commit }}
 
   prebuild-tvm:
     uses: ./.github/workflows/build-tvm.yml
     with:
       checkout-ref: ${{ inputs.checkout-ref }}
 
   build-wheel:
-    needs: [prebuild-npuir, prebuild-tvm]
+    needs: [select-npuir-version, prebuild-npuir, prebuild-tvm]
     runs-on: ${{ matrix.runner }}
     timeout-minutes: 120
     strategy:
@@ -63,7 +103,7 @@ jobs:
       uses: actions/cache/restore@v4
       with:
         path: ${{ env.NPUIR_INSTALL_DIR }}
-        key: npuir-prebuild-rtti-${{ matrix.arch }}-${{ needs.prebuild-npuir.outputs.npuir-commit }}
+        key: npuir-prebuild-rtti-${{ matrix.arch }}-${{ needs.select-npuir-version.outputs.npuir-commit }}
         fail-on-cache-miss: true
 
     - name: Restore TVM prebuild from cache
@@ -94,6 +134,7 @@ jobs:
         cp 3rdparty/tvm/cmake/config.cmake build/
         echo "set(USE_NPUIR ON)" >> build/config.cmake
         echo "set(BISHENGIR_ROOT_PATH ${BISHENGIR_PATH})" >> build/config.cmake
+        echo "set(TILELANG_ASCEND_CANN_VERSION ${{ inputs.cann-version }})" >> build/config.cmake
         echo "set(USE_CUDA OFF)" >> build/config.cmake
         echo "set(USE_ROCM OFF)" >> build/config.cmake
         echo "set(USE_LLVM OFF)" >> build/config.cmake
@@ -112,11 +153,12 @@ jobs:
         USE_NPUIR: "true"
         BISHENGIR_PATH: ${{ github.workspace }}/npuir-prebuild
         TVM_PREBUILD_PATH: ${{ github.workspace }}/tvm-prebuild
-      run: python setup.py bdist_wheel
+      run: |
+        python setup.py bdist_wheel --build-number "${{ needs.select-npuir-version.outputs.wheel-build-number }}"
 
     - name: Upload wheel artifact
       uses: actions/upload-artifact@v4
       with:
-        name: tilelang-npuir-py${{ env.PYTHON_VERSION }}-${{ matrix.arch }}
+        name: tilelang-npuir-${{ needs.select-npuir-version.outputs.cann-version-slug }}-py${{ env.PYTHON_VERSION }}-${{ matrix.arch }}
         path: dist/*.whl
         retention-days: 30

.github/workflows/ci_npuir.yml

@@ -53,15 +53,36 @@ jobs:
 
   build:
     needs: check-trigger
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cann-version: 8.5.0
+            cann-version-slug: cann-8.5.0
+          - cann-version: 9.0.0.beta2
+            cann-version-slug: cann-9.0.0.beta2
     uses: ./.github/workflows/build-tilelang-npuir.yml
     with:
       checkout-ref: refs/pull/${{ needs.check-trigger.outputs.pr-number }}/merge
+      cann-version: ${{ matrix.cann-version }}
 
   test-npuir:
     needs: [check-trigger, build]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cann-version: 8.5.0
+            cann-version-slug: cann-8.5.0
+            container-image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11
+          - cann-version: 9.0.0.beta2
+            cann-version-slug: cann-9.0.0.beta2
+            container-image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:9.0.0.beta.2-a3-ubuntu22.04-py3.11
     runs-on: linux-aarch64-a3-2
     container:
-      image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.5.0-a3-ubuntu22.04-py3.11
+      image: ${{ matrix.container-image }}
+    env:
+      TILELANG_ASCEND_CANN_VERSION: ${{ matrix.cann-version }}
     timeout-minutes: 120
     steps:
 
@@ -74,7 +95,7 @@ jobs:
     - name: Download arm64 wheel
       uses: actions/download-artifact@v4
       with:
-        name: tilelang-npuir-py3.11-arm64
+        name: tilelang-npuir-${{ matrix.cann-version-slug }}-py3.11-arm64
         path: dist/
 
     - name: Install tilelang and test dependencies
@@ -90,7 +111,7 @@ jobs:
         pip install torch-npu==2.7.1
         pip install dist/*.whl
         pip install pytest pytest-xdist pytest-html numpy
-    
+
     - name: Install clang
       run: |
         apt-get update && apt-get install -y clang
@@ -165,14 +186,21 @@ jobs:
       if: always()
       uses: actions/upload-artifact@v4
       with:
-        name: npuir-test-report
+        name: npuir-test-report-${{ matrix.cann-version-slug }}
         path: testing/npuir/output/
 
   build-release:
     if: github.event_name == 'release'
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cann-version: 8.5.0
+          - cann-version: 9.0.0.beta2
     uses: ./.github/workflows/build-tilelang-npuir.yml
     with:
       checkout-ref: refs/tags/${{ github.event.release.tag_name }}
+      cann-version: ${{ matrix.cann-version }}
     secrets: inherit
 
   upload-release-assets:

CMakeLists.txt

@@ -42,6 +42,25 @@ endif()
 
 # Setting up MLIR infrastructure
 set(NPUIR_LIBS "")
+set(_tilelang_ascend_cann_version "${TILELANG_ASCEND_CANN_VERSION}")
+if("${_tilelang_ascend_cann_version}" STREQUAL "")
+  if(DEFINED ENV{ASCEND_HOME_PATH} AND NOT "$ENV{ASCEND_HOME_PATH}" STREQUAL "")
+    set(_tilelang_ascend_home "$ENV{ASCEND_HOME_PATH}")
+  elseif(DEFINED ENV{ASCEND_HOME} AND NOT "$ENV{ASCEND_HOME}" STREQUAL "")
+    set(_tilelang_ascend_home "$ENV{ASCEND_HOME}")
+  endif()
+
+  if(DEFINED _tilelang_ascend_home)
+    if(_tilelang_ascend_home MATCHES "cann-8\\.5\\.0(/|$)")
+      set(_tilelang_ascend_cann_version "8.5.0")
+    elseif(_tilelang_ascend_home MATCHES "cann-9\\.0\\.0-beta\\.2(/|$)")
+      set(_tilelang_ascend_cann_version "9.0.0.beta2")
+    endif()
+  endif()
+endif()
+set(TILELANG_ASCEND_CANN_VERSION "${_tilelang_ascend_cann_version}" CACHE STRING "CANN compatibility version for TileLang NPUIR build" FORCE)
+unset(_tilelang_ascend_cann_version)
+unset(_tilelang_ascend_home)
 if(USE_NPUIR)
   set(BISHENGIR_INCLUDE_DIRS "")
   set(BISHENGIR_LIBRARY_DIRS "")
@@ -227,6 +246,19 @@ set(TILE_LANG_COMPILE_DEFS
   PICOJSON_USE_INT64
 )
 
+if(USE_NPUIR)
+  if(TILELANG_ASCEND_CANN_VERSION STREQUAL "")
+    message(FATAL_ERROR "TILELANG_ASCEND_CANN_VERSION must be set when USE_NPUIR=ON")
+  endif()
+  if(TILELANG_ASCEND_CANN_VERSION STREQUAL "8.5.0")
+    list(APPEND TILE_LANG_COMPILE_DEFS TILELANG_ASCEND_CANN_VERSION_8_5_0=1)
+  elseif(TILELANG_ASCEND_CANN_VERSION STREQUAL "9.0.0.beta2")
+    list(APPEND TILE_LANG_COMPILE_DEFS TILELANG_ASCEND_CANN_VERSION_9_0_0_BETA2=1)
+  else()
+    message(FATAL_ERROR "Unsupported TILELANG_ASCEND_CANN_VERSION=${TILELANG_ASCEND_CANN_VERSION}")
+  endif()
+endif()
+
 # Set target properties for object library
 target_include_directories(tilelang_objs PRIVATE ${TILE_LANG_INCLUDES})
 target_compile_definitions(tilelang_objs PRIVATE ${TILE_LANG_COMPILE_DEFS})

README.md

@@ -15,7 +15,6 @@ Within the TileLang ecosystem, we have developed an NPU Intermediate Representat
 <img src=./images/npuir_architecture.png style="width: 50%";/>
 </div>
 
-
 ## Latest News
 - 3/28/2026 🚀: We provide a free environment to facilitate user experience and development for TileLang [Pull Request#708](https://github.com/tile-ai/tilelang-ascend/pull/708).
 
@@ -64,7 +63,6 @@ If you need to access Ascend NPU computing resources for development or testing,
 
 Within the `examples` directory, you will also find additional complex kernels—such as convolutions, forward/backward passes for FlashAttention, more operators will continuously be added.
 
-
 ## Installation
 ### Environment Setup
 
@@ -85,15 +83,12 @@ source /path/to/install/Ascend/ascend-toolkit/set_env.sh
 
 Prepare a Python environment with Python version between 3.7.*x* and 3.11.4 (inclusive) and ensure that `pip3` is available.
 
-
    Ascend Toolkit Installation Requirements
 
    ```shell
    pip3 install attrs cython 'numpy>=1.19.2,<=1.24.0' decorator sympy cffi pyyaml pathlib2 psutil protobuf==3.20.0 scipy requests absl-py
    ```
 
-
-
 <!-- 补充环境变量设置 -->
 Set Environment Variables
 
@@ -102,11 +97,8 @@ export ACL_OP_INIT_MODE=1
 ```
   <!-- 注意：如果用户需要新的编译器安装包，请联系社区管理员zhaojiqiao@huawei.com,yangsichan@huawei.com TEL:15901269653 -->
 
-  Note: If you require a new compiler installation package, please contact the community administrators:  
-**zhaojiqiao@huawei.com**, **yangsichan@huawei.com**  
-
-
-
+  Note: If you require a new compiler installation package, please contact the community administrators:
+**zhaojiqiao@huawei.com**, **yangsichan@huawei.com**
 
 #### Build
 
@@ -124,6 +116,8 @@ Run the installation script
 cd tilelang-ascend
 # build AscendNPU-IR in 3rdparty
 bash install_npuir.sh
+# build with an explicit CANN compatibility version
+bash install_npuir.sh --cann-version=8.5.0
 # Alternative way of building with local AscendNPU-IR
 bash install_npuir.sh --bishengir-path=/path/to/bishengir-compile
 ```
@@ -160,12 +154,12 @@ seq_len = 4096  # Length of the vectors to be added
 def vec_add(N, block_N, dtype="float32"):
     """
     Define a vector addition kernel using TileLang.
-    
+
     Parameters:
     - N: Total length of the vectors.
     - block_N: Number of elements processed per kernel thread/block.
     - dtype: Data type of the tensors (default: "float32").
-    
+
     Returns:
     - A TileLang prim_func representing the vector addition kernel.
     """
@@ -254,14 +248,14 @@ def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="flo
         C: T.Tensor((M, N), dtype), # Output matrix C
     ):
         with T.Kernel(T.ceildiv(N, block_N) * T.ceildiv(M, block_M), is_npu=True) as (cid, _):
-          
+
             by = cid // T.ceildiv(N, block_N) # Block row index
             bx = cid % T.ceildiv(N, block_N)  # Block column index
 
             # Alloc shared memory for inputs
             A_shared = T.alloc_shared((block_M, block_K), dtype)
             B_shared = T.alloc_shared((block_K, block_N), dtype)
-          
+
             # Alloc local fragment for accumulation
             C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
 
@@ -270,7 +264,7 @@ def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="flo
                 # Copy the data from global memory to shared memory
                 T.copy(A[by * block_M, k * block_K], A_shared)
                 T.copy(B[k * block_K, bx * block_N], B_shared)
-              
+
                 # Perform matrix multiplication with accumulation
                 # If 'initC' is true, the result matrix will be initialized to zero before accumulation
                 T.gemm(A_shared, B_shared, C_local, initC=(k == 0))
@@ -287,4 +281,3 @@ def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="flo
 ## Acknowledgements
 
 Peking University Kunpeng & Ascend Center for Excellence in Science, Education, Innovation
-

docs/developer/EnvironmentVariables.md

@@ -28,6 +28,8 @@ export TILELANG_DUMP_IR=TRUE
 | `TILELANG_ASCEND_MODE` | `Expert` | Set the TileLang Mode; currently, Expert mode and Developer mode are supported | `Expert`: Expert Mode<br>`Developer`: Developer Mode |
 | `TILELANG_ASCEND_DEVICE_NAME` | `Ascend910B` | Override the target device name for compilation (e.g. for cross-compilation). If not set, runtime hardware detection is used. | String, e.g., `Ascend910B`|
 
+For NPUIR builds and JIT compilation, TileLang resolves the CANN compatibility version from `ASCEND_HOME_PATH` by default. `install_npuir.sh --cann-version=<version>` uses the same version source to build both `AscendNPU-IR` and TileLang, and JIT compilation forwards the matching `hivmc-version` to `adapt_triton_kernel` and `bishengir-compile`.
+
 ## Autotuner
 
 | Variable | Default | Description | Valid Values |