add custom op in ascend

.github/workflows/main.yml

@@ -56,101 +56,130 @@ jobs:
       - name: Build and install DLCompiler
         run: |
           set -ex
-          source /home/dlc_ci/.bashrc
-          which conda
-          echo "which conda? $(which conda)"
-          conda activate dlcompiler
-          source /usr/local/Ascend/cann-8.5.0/set_env.sh
-          cd ${{ env.CI_PATH }}
-          export JSON_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/include.zip
-          export GOOGLETEST_DIR34=${{ vars.CI_BASE_PATH }}/data/v34/googletest
-          export LLVM_TGZ_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/llvm-064f02da-ubuntu-arm64.tar.gz
-
-          rm -rf ./third_party/*
-          git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/ascendnpu-ir ./third_party/ascendnpu-ir
-          git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/json ./third_party/json
-          git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton_shared ./third_party/triton_shared
-          git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton ./third_party/triton
-          echo "whoami? $(whoami)"
-          echo "which python? $(which python)"
-          git submodule update --init
-          echo "git submodule update done."
-          pip install nanobind -i https://mirrors.huaweicloud.com/repository/pypi/simple
-          bash compile_shared.sh apply_patch=true
+          echo "=== 切换前用户信息 ==="
+          echo "whoami: $(whoami)"
+          echo "user id: $(id)"
+          echo "PATH: $PATH"
+          echo "======================"
+
+          # 切换到 root 用户执行，确保有权限访问 NPU 设备
+          sudo -E bash -c '
+            source /home/dlc_ci/.bashrc
+            which conda
+            echo "which conda? $(which conda)"
+            conda activate dlcompiler
+            source /usr/local/Ascend/cann-8.5.0/set_env.sh
+            cd ${{ env.CI_PATH }}
+            export JSON_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/include.zip
+            export GOOGLETEST_DIR34=${{ vars.CI_BASE_PATH }}/data/v34/googletest
+            export LLVM_TGZ_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/llvm-064f02da-ubuntu-arm64.tar.gz
+
+            rm -rf ./third_party/*
+            git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/ascendnpu-ir ./third_party/ascendnpu-ir
+            git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/json ./third_party/json
+            git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton_shared ./third_party/triton_shared
+            git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton ./third_party/triton
+            echo "whoami before compile: $(whoami)"
+            echo "which python? $(which python)"
+            echo "npu-smi path: $(which npu-smi 2>/dev/null || echo "not found")"
+            git submodule update --init
+            echo "git submodule update done."
+            pip install nanobind -i https://mirrors.huaweicloud.com/repository/pypi/simple
+
+            bash compile_shared.sh apply_patch=true
+          '
 
       - name: Build and install tilelang-dlc
         run: |
           set -ex
-          source /home/dlc_ci/.bashrc
-          conda activate dlcompiler
-          source /usr/local/Ascend/cann-8.5.0/set_env.sh
-          cd ${{ env.CI_PATH }}
-          export TILELANG_DLC_PATH=${{ vars.CI_BASE_PATH }}/data/tilelang-dlc
-          export DLCOMPILER_SOURCE=${{ env.CI_PATH }}
-          export TILELANG_USE_DLCOMPILER=1
-          echo "whoami? $(whoami)"
-          echo "which python? $(which python)"
-          bash scripts/install_tilelang-dlc.sh
+          # 切换到 root 用户执行
+          sudo -E bash -c '
+            source /home/dlc_ci/.bashrc
+            conda activate dlcompiler
+            source /usr/local/Ascend/cann-8.5.0/set_env.sh
+            cd ${{ env.CI_PATH }}
+            export TILELANG_DLC_PATH=${{ vars.CI_BASE_PATH }}/data/tilelang-dlc
+            export DLCOMPILER_SOURCE=${{ env.CI_PATH }}
+            export TILELANG_USE_DLCOMPILER=1
+            echo "whoami? $(whoami)"
+            echo "which python? $(which python)"
+            bash scripts/install_tilelang-dlc.sh
+          '
 
       - name: Run tilelang-dlc tests on ascend
         run: |
           set -ex
-          source /home/dlc_ci/.bashrc
-          conda activate dlcompiler
-          source /usr/local/Ascend/cann-8.5.0/set_env.sh
-          cd ${{env.CI_PATH }}
-          export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
-          export ASCEND_RT_VISIBLE_DEVICES=7
-          export TILELANG_USE_DLCOMPILER=1
-          bash test/commonir/run_tests.sh
+          # 切换到 root 用户执行
+          sudo -E bash -c '
+            source /home/dlc_ci/.bashrc
+            conda activate dlcompiler
+            source /usr/local/Ascend/cann-8.5.0/set_env.sh
+            cd ${{env.CI_PATH }}
+            export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
+            export ASCEND_RT_VISIBLE_DEVICES=7
+            export TILELANG_USE_DLCOMPILER=1
+            bash test/commonir/run_tests.sh
+          '
 
       - name: Run triton tests on ascend
         run: |
           set -ex
-          source /home/dlc_ci/.bashrc
-          conda activate dlcompiler
-          source /usr/local/Ascend/cann-8.5.0/set_env.sh
-          cd ${{env.CI_PATH }}
-          echo "whoami? $(whoami)"
-          echo "which python? $(which python)"
-          export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
-          export ASCEND_RT_VISIBLE_DEVICES=7
-          bash test/ascend/run_tests.sh
+          # 切换到 root 用户执行
+          sudo -E bash -c '
+            source /home/dlc_ci/.bashrc
+            conda activate dlcompiler
+            source /usr/local/Ascend/cann-8.5.0/set_env.sh
+            cd ${{env.CI_PATH }}
+            echo "whoami? $(whoami)"
+            echo "which python? $(which python)"
+            export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
+            export ASCEND_RT_VISIBLE_DEVICES=7
+            bash test/ascend/run_tests.sh
+          '
 
       - name: Run MLIR tests
         run: |
           set -ex
-          source /home/dlc_ci/.bashrc
-          conda activate dlcompiler
-          source /usr/local/Ascend/cann-8.5.0/set_env.sh
-          cd ${{env.CI_PATH }}
-          echo "whoami? $(whoami)"
-          echo "which python? $(which python)"
-          export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
-          export ASCEND_RT_VISIBLE_DEVICES=7
-          bash test/ascend/test_mlir.sh
+          # 切换到 root 用户执行
+          sudo -E bash -c '
+            source /home/dlc_ci/.bashrc
+            conda activate dlcompiler
+            source /usr/local/Ascend/cann-8.5.0/set_env.sh
+            cd ${{env.CI_PATH }}
+            echo "whoami? $(whoami)"
+            echo "which python? $(which python)"
+            export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
+            export ASCEND_RT_VISIBLE_DEVICES=7
+            bash test/ascend/test_mlir.sh
+          '
 
       - name: Run DSL tests
         run: |
           set -ex
-          source /home/dlc_ci/.bashrc
-          conda activate dlcompiler
-          source /usr/local/Ascend/cann-8.5.0/set_env.sh
-          cd ${{env.CI_PATH }}
-          echo "whoami? $(whoami)"
-          echo "which python? $(which python)"
-          export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
-          export ASCEND_RT_VISIBLE_DEVICES=7
-          bash test/dsl/run_tests.sh
+          # 切换到 root 用户执行
+          sudo -E bash -c '
+            source /home/dlc_ci/.bashrc
+            conda activate dlcompiler
+            source /usr/local/Ascend/cann-8.5.0/set_env.sh
+            cd ${{env.CI_PATH }}
+            echo "whoami? $(whoami)"
+            echo "which python? $(which python)"
+            export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH
+            export ASCEND_RT_VISIBLE_DEVICES=7
+            bash test/dsl/run_tests.sh
+          '
 
       - name: Clear workfile
         if: always()
         run: |
-          export workdir=$(pwd)
-          cd ..
-          rm -rf $workdir
-          mkdir $workdir
-          chmod -R 777 $workdir  
-          if [ -d "${{ env.CI_PATH }}" ]; then
-            rm -rf ${{ env.CI_PATH }}
-          fi
+          # 使用 root 权限清理，确保能删除之前以 root 创建的文件
+          sudo bash -c '
+            export workdir=$(pwd)
+            cd ..
+            rm -rf $workdir
+            mkdir $workdir
+            chmod -R 777 $workdir  
+            if [ -d "${{ env.CI_PATH }}" ]; then
+              rm -rf ${{ env.CI_PATH }}
+            fi
+          '

backend/npu.py

@@ -521,6 +521,15 @@ def ttsharedir_to_linkedir(mod, metadata, opt, *, named_ops=False, cpu_verify=Fa
     pattern = r"(memref\<.*?\>)\s+to\s+(tensor\<.*?\>)"
     # 使用正则替换，保留memref和tensor类型，中间插入注释
     content = re.sub(pattern, r"\1 // to \2", content)
+    # 处理customop的attr
+    if len(re.findall("hivm\.hir\.custom", content)) > 0:
+        content = re.sub(r'"#hivm\.pipe<([A-Za-z0-9_]*)>"', r"#hivm.pipe<\1>", content)
+        content = re.sub(
+            r'"#hivm\.tcore_type<([A-Za-z0-9_]*)>"', r"#hivm.tcore_type<\1>", content
+        )
+        content = re.sub(
+            r'"#hivm\.vf_mode<([A-Za-z0-9_]*)>"', r"#hivm.vf_mode<\1>", content
+        )
 
     if opt.debug or dump_ir:
         cmd_list = [
@@ -681,6 +690,10 @@ def _parse_linalg_metadata(linalg: str, metadata: dict):
     TENSOR_KIND_REGEX = (
         r"%arg(\d+):[^,)]*?\{[^}]*?tt\.tensor_kind\s*=\s*([^:\s}]+)\s*:[^}]*?\}"
     )
+
+    # Example: bitcode = "a.bc"
+    BITCODES_REGEX = r'bitcode\s*=\s*(?:"([^"]+)"|\'([^\']+)\'|(\w+))'
+
     # Example removal:   ', mix_mode = "aiv"' → ''
     REMOVE_MIX_MODE_REGEX = r', mix_mode\s*=\s*"[^"]*"'
     # Note: Compiled Kernel requires to estimate size of shared memory to occupy
@@ -696,6 +709,11 @@ def _parse_linalg_metadata(linalg: str, metadata: dict):
     metadata["tensor_kinds"] = [
         int(kind) for _, kind in re.findall(TENSOR_KIND_REGEX, linalg)
     ]
+
+    # Parse all bitcode paths
+    bitcodes = re.findall(BITCODES_REGEX, linalg)
+    metadata["bitcodes"] = [val for group in bitcodes for val in group if val]
+
     # remove the mix_mode attribute
     linalg = re.sub(REMOVE_MIX_MODE_REGEX, "", linalg)
     return linalg, metadata
@@ -729,6 +747,12 @@ def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt):
             _compile_option_list += ["--enable-sanitizer=true"]
         if _is_auto_map_parallel_blocks_enabled():
             _compile_option_list += ["--enable-auto-blockify-loop"]
+
+        bitcodes = metadata["bitcodes"]
+        if bitcodes is not None:
+            for bitcodes in bitcodes:
+                _compile_option_list += [f"--link-aicore-bitcode={bitcodes}"]
+
         npu_compiler_path = _get_npucompiler_path()
 
         # support bishengir-compile more version

compiler/lib/Conversion/LinkedToHIVM/LinkedToHIVM.cpp

@@ -123,13 +123,32 @@ struct TritonCustomSyncOpToHIVMSyncOpConversion
   }
 };
 
+// Convert CustomOp after operand type changed,
+// for example tt.ptr changed to memref.
+class TritonCustomOpToHIVMCustomOpConversion
+    : public OpConversionPattern<triton::CustomOp> {
+public:
+  using OpConversionPattern<triton::CustomOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(triton::CustomOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto res_types = adaptor.getOutputs().getTypes();
+    auto new_op = rewriter.create<hivm::CustomOp>(
+        op->getLoc(), res_types, adaptor.getOperands(), op->getAttrs());
+    rewriter.replaceOp(op, new_op);
+    return success();
+  }
+};
+
 void LinkedToHIVMPass::runOnOperation() {
   auto module = getOperation();
   ConversionTarget target(getContext());
   target.addLegalDialect<hivm::HIVMDialect>();
 
   RewritePatternSet patterns(&getContext());
   patterns.add<TritonCustomSyncOpToHIVMSyncOpConversion>(patterns.getContext());
+  patterns.add<TritonCustomOpToHIVMCustomOpConversion>(patterns.getContext());
   if (failed(applyPartialConversion(module, target, std::move(patterns)))) {
     signalPassFailure();
   }

language/deeplink/__init__.py

@@ -21,6 +21,14 @@
     L0C,
     SyncFlag,
 )
+from .custom_op import (
+    custom,
+    custom_semantic,
+    register_custom_op,
+    CORE,
+    PIPE,
+    MODE,
+)
 
 __all__ = [
     "libdevice",
@@ -43,6 +51,12 @@
     "L0C",
     "SyncFlag",
     "async_task",
+    "custom",
+    "custom_semantic",
+    "register_custom_op",
+    "CORE",
+    "PIPE",
+    "MODE",
 ]
 
 init_dicp_driver()