Refactor: Decoupling NPU-related ttshared Pass Patches and Integrating Non-bisheng-compiler Stages into Pybind11 and Stage Consolidation (#125)

* Refactor: Integrate Pybind11 and Refactor NPU Passes for Compile Speedup and Size Reduction

This commit introduces significant improvements to the NPU backend focusing on code architecture and build performance:

1. NPU Code Refactoring: NPU-specific logic was successfully decoupled from the monolithic 'triton-shared' component and re-implemented as dedicated, cleaner passes and patterns. This greatly enhances code clarity and maintainability.

2. Performance Boost (Pybind11): Integrated Pybind11, effectively eliminating Python-C++ interface overhead. Compilation benchmarks show intermediate stages (e.g., ttir_post, ttshared, linkedir) are now over 60% faster on average.

3. Binary Size Optimization: The implementation allows for the removal of the external 'opt' dependency, resulting in a substantial reduction of the release package size: Before: ~488MB After: ~330MB

Note for Publish Build:

To generate the optimized release version, ensure the environment variable is set: export IS_NOT_PUBLISH=0

* fix  Improve error handling for bishengir-compile failures

Author: Truth-Ke

.gitignore

@@ -96,6 +96,7 @@ tags
 kernel_meta/
 fusion_result.json
 *.log
+launcher_cxx11abi*
 
 # package
 backend/triton-shared-opt-v3*

CMakeLists.txt

@@ -1,5 +1,6 @@
 set(DC_TRITON_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
 set(DC_TRITON_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR})
+set(TRTION_SHARED_NPU_SPECIFIC_SOURCES  "${DC_TRITON_SOURCE_DIR}/compiler/lib/Conversion/TritonToLinalgNPU")
 
 set(DC_TRITON_INCLUDE_DIR "")
 set(DC_TRITON_LINK_DIR "")
@@ -39,5 +40,24 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/tools)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/compiler)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/triton_shared)
 
-add_triton_plugin(DICPTriton ${CMAKE_CURRENT_SOURCE_DIR}/dicp_triton.cc LINK_LIBS ${LIBS})
-target_include_directories(DICPTriton PUBLIC "${pybind11_INCLUDE_DIRS}")
+if (TRITON_BUILD_PYTHON_MODULE)
+    add_triton_plugin(tritonDicpTriton ${CMAKE_CURRENT_SOURCE_DIR}/triton_dicp_triton.cc 
+    LINK_LIBS 
+
+    MLIRAffineToStandard
+    MLIRIR
+    MLIRPass
+    MLIRTransforms
+    MLIRSupport
+    MLIRBytecodeWriter
+    
+    TritonToLinalgNPUCoversion
+
+    LinalgExtTransforms
+    TritonExtTransforms
+
+    LinalgToLinked
+    LinkedToHIVM
+    )
+    target_link_libraries(tritonDicpTriton PRIVATE Python3::Module pybind11::headers)
+endif()

backend/compiler.py

@@ -235,10 +235,9 @@ def add_stages(self, stages, options, language=None):
             from triton.backends.dicp_triton.npu import (
                 make_ttir,
                 ttir_to_linalg,
-                ttir_to_ttsharedir,
+                ttir_to_ttsharedir_ascend,
                 ttsharedir_to_linkedir,
                 linalg_to_bin_enable_npu_compile,
-                ttir_post,
             )
 
             stages["ttir"] = lambda src, metadata: make_ttir(src, metadata, options)
@@ -255,11 +254,10 @@ def add_stages(self, stages, options, language=None):
                     )
             else:
                 if options.enable_npu_compile:
-                    stages["ttir_post"] = lambda src, metadata: ttir_post(
-                        src, metadata, options
-                    )
-                    stages["ttshared"] = lambda src, metadata: ttir_to_ttsharedir(
-                        src, metadata, options, named_ops=True
+                    stages["ttshared"] = (
+                        lambda src, metadata: ttir_to_ttsharedir_ascend(
+                            src, metadata, options, named_ops=True
+                        )
                     )
                     stages["linkedir"] = lambda src, metadata: ttsharedir_to_linkedir(
                         src, metadata, options, named_ops=True

backend/driver.py

@@ -381,3 +381,7 @@ def map_python_to_cpp_type(self, ty: str) -> str:
             "f32": "float",
             "fp64": "double",
         }[ty]
+
+    @classmethod
+    def clear_cache(self, cache):
+        cache.zero_()

backend/npu.py

@@ -7,8 +7,9 @@
 import hashlib
 from triton.runtime.cache import get_cache_manager, get_dump_manager
 from triton.backends.compiler import GPUTarget
-from triton._C.libtriton import ir, passes
+from triton._C.libtriton import ir, passes, dicp_triton
 from triton.runtime.cache import get_dump_manager
+import triton.backends.dicp_triton.utils as dicp_utils
 from dataclasses import dataclass
 from typing import Any, Union, Tuple, Dict
 import ctypes
@@ -25,8 +26,11 @@
 replace_linked_ir = os.environ.get("DLC_REPLACE_LINKED_IR_FILE", None)
 if dump_ir or (replace_ttshared_ir is not None) or (replace_linked_ir is not None):
     os.environ["TRITON_ALWAYS_COMPILE"] = "1"
-    if not os.path.exists("./tmp"):
-        os.makedirs("./tmp")
+    dump_dir = "./tmp"
+    os.environ["TRITON_DUMP_DIR"] = os.environ.get("TRITON_DUMP_DIR", dump_dir)
+    if os.path.exists(dump_dir):
+        print(f"Directory **{dump_dir}** exists. Deleting the entire directory...")
+        shutil.rmtree(dump_dir)
 
 local_bishengir_path = os.path.join(os.path.dirname(__file__), "../../_C/bishengir")
 bisheng_install_path = os.environ.get("BISHENG_INSTALL_PATH", None)
@@ -49,6 +53,19 @@ def downgrade_llir(llir):
     return llir
 
 
+def _replace_mod_ir_with_file(mod, filepath: str, stage_name: str):
+    p = Path(filepath)
+    if not p.exists():
+        raise FileNotFoundError(f"Replacement MLIR file not found: {filepath}")
+    print(f"[DEBUG] replacing '{stage_name}' IR with file '{filepath}'")
+    try:
+        new_mod = ir.parse_mlir_module(str(p), mod.context)
+        new_mod.context = mod.context
+        return new_mod
+    except Exception as e:
+        raise RuntimeError(f"Failed to parse replacement MLIR file '{filepath}': {e}")
+
+
 def _downgrade_mem_attrs(llir: str):
     memory_pattern = r"memory\([^()]*\)"
 
@@ -330,6 +347,7 @@ def min_dot_size(target: GPUTarget):
 def make_ttir(mod, metadata, opt):
     if "hash" not in metadata:
         metadata["hash"] = hashlib.md5(f"{mod}-{metadata}".encode()).hexdigest()
+    mod.set_attr("dicp.backend", ir.builder(mod.context).get_string_attr("ascend"))
     # the same optimize pass for triton-ir as all other backends
     pm = ir.pass_manager(mod.context)
     pm.enable_debug()
@@ -341,44 +359,17 @@ def make_ttir(mod, metadata, opt):
     passes.common.add_licm(pm)
     passes.common.add_symbol_dce(pm)
     pm.run(mod)
-    if opt.debug:
-        dump_manager = get_dump_manager(metadata["hash"])
-        print(f"Dumping intermediate results to {dump_manager.cache_dir}")
-        dump_manager.put(str(mod), "kernel.ttir.mlir", binary=False)
-
+    if opt.debug or dump_ir:
+        dicp_utils._dump_stage_ir(str(mod), metadata["hash"], "kernel.ttir.mlir")
     return mod
 
 
-def ttir_post(mod, metadata, opt):
-    ttir_code = str(mod)
-    with tempfile.TemporaryDirectory() as tmpdir:
-        src_path = os.path.join(tmpdir, "kernel.ttir.mlir")
-        dst_path = os.path.join(tmpdir, "kernel.ttir_post.mlir")
-        Path(src_path).write_text(ttir_code)
-        dicp_opt_path = _get_dicp_opt_path()
-        dicp_cmd_list = [
-            dicp_opt_path,
-            src_path,
-            "--canonicalize-cmpi",
-            "--canonicalize-triton-ir-ascend",
-            "-o",
-            dst_path,
-        ]
-        if dump_ir:
-            shutil.copy(src_path, "./tmp/kernel.ttir.mlir")
-            print(f"DEBUG dump ir[ttir_post] command: {dicp_cmd_list}")
-        ret = subprocess.run(dicp_cmd_list, capture_output=True, check=True)
-        if dump_ir:
-            shutil.copy(dst_path, "./tmp/kernel.ttir_post.mlir")
-        return Path(dst_path).read_text()
-
-
 def ttir_to_linalg(mod, metadata, opt, *, named_ops=True):
     # use triton_adapter to lower Triton-MLIR to linalg
     # Get Triton-MLIR as string
     ttir_code = str(mod)
     with tempfile.TemporaryDirectory() as tmpdir:
-        src_path = os.path.join(tmpdir, "kernel.ttir_post.mlir")
+        src_path = os.path.join(tmpdir, "kernel.ttir.mlir")
         dst_path = os.path.join(tmpdir, "kernel.ttadapter.mlir")
         Path(src_path).write_text(ttir_code)
         triton_adapter_opt_path = _get_triton_adapter_opt_path()
@@ -403,90 +394,72 @@ def ttir_to_linalg(mod, metadata, opt, *, named_ops=True):
         return Path(dst_path).read_text()
 
 
-def ttir_to_ttsharedir(mod, metadata, opt, *, named_ops=False):
-    ttir_code = str(mod)
-    # 注释掉gpu.barrier
-    ttir_code = re.sub(r"gpu\.barrier", r"// gpu.barrier", ttir_code)
-    with tempfile.TemporaryDirectory() as tmpdir:
-        src_path = os.path.join(tmpdir, "kernel.ttir_post.mlir")
-        dst_ttshared_path = os.path.join(tmpdir, "kernel.ttshared.mlir")
-        Path(src_path).write_text(ttir_code)
-        triton_shared_opt_path = _get_triton_shared_opt_path()
-        ttshared_cmd = (
-            "--triton-to-linalg-experimental"
-            if "v3_2" not in triton_shared_opt_path
-            else "--triton-to-linalg"
-        )
-
-        cmd_shared_list = [
-            triton_shared_opt_path,
-            src_path,
-            ttshared_cmd,
-            "-o",
-            dst_ttshared_path,
+def ttir_to_ttsharedir_ascend(mod, metadata, opt, *, named_ops=False):
+    pm = ir.pass_manager(mod.context)
+    dicp_triton.passes.triton_shared_ascend.add_canonicalize_cmpi(pm)
+    dicp_triton.passes.triton_shared_ascend.add_canonicalize_triton_ir_ascend(pm)
+    dicp_triton.passes.triton_shared_ascend.add_triton_to_linalg_npu(pm)
+    pm.run(mod)
+    if opt.debug or dump_ir:
+        cmd_list = [
+            _get_dicp_opt_path(),
+            "kernel.ttir.mlir",
+            "--canonicalize-cmpi",
+            "--canonicalize-triton-ir-ascend",
+            "--triton-to-linalg-npu-conversion",
         ]
-        if dump_ir:
-            print(f"DEBUG dump ir[ttir_to_ttsharedir] command: {cmd_shared_list}")
-        ret = subprocess.run(cmd_shared_list, capture_output=True, check=True)
-
-        if dump_ir:
-            shutil.copy(dst_ttshared_path, "./tmp/kernel.ttshared.mlir")
-        if replace_ttshared_ir is not None:
-            print(f"[DEBUG] Replace ttsharedir with {replace_ttshared_ir}")
-            return Path(replace_ttshared_ir).read_text()
-        return Path(dst_ttshared_path).read_text()
+        dicp_utils._dump_stage_ir(
+            str(mod), metadata["hash"], "kernel.ttshared.mlir", cmd_list
+        )
+    if replace_ttshared_ir is not None:
+        return _replace_mod_ir_with_file(
+            mod, replace_ttshared_ir, "ttir_to_ttsharedir_ascend"
+        )
+    return mod
 
 
 def ttsharedir_to_linkedir(mod, metadata, opt, *, named_ops=False):
-    ttsharedir_code = str(mod)
-    with tempfile.TemporaryDirectory() as tmpdir:
-        src_path = os.path.join(tmpdir, "kernel.ttshared.mlir")
-        dst_path = os.path.join(tmpdir, "kernel.linkedir.mlir")
-        Path(src_path).write_text(ttsharedir_code)
-        dicp_opt_path = _get_dicp_opt_path()
-        dicp_cmd_list = [
-            dicp_opt_path,
-            src_path,
+    pm = ir.pass_manager(mod.context)
+    dicp_triton.passes.linked_npu.add_lower_affine(pm)
+    dicp_triton.passes.linked_npu.add_normalize_slice_ops(pm)
+    dicp_triton.passes.linked_npu.add_linalg_if_to_select(pm)
+    dicp_triton.passes.linked_npu.add_linalg_generic_to_scf(pm)
+    dicp_triton.passes.linked_npu.add_scalar_to_1d_tensor(pm)
+    dicp_triton.passes.linked_npu.add_linalg_to_linked(pm, False, True)
+    dicp_triton.passes.linked_npu.add_linked_to_hivm(pm)
+    pm.run(mod)
+
+    # TODO(zmz): 修改test_path 中内容，暂时在python中处理，bishengir-compile后续会支持，去掉这里逻辑。
+    content = str(mod)
+    # 将"*xfxxx"替换成"?xfxxx"
+    content = content.replace("*xf", "?xf")
+    content = content.replace("*xi", "?xi")
+    content = content.replace("*xbf", "?xbf")
+    # 匹配形如 "memref<...> to tensor<...>" 的模式
+    pattern = r"(memref\<.*?\>)\s+to\s+(tensor\<.*?\>)"
+    # 使用正则替换，保留memref和tensor类型，中间插入注释
+    content = re.sub(pattern, r"\1 // to \2", content)
+
+    if opt.debug or dump_ir:
+        cmd_list = [
+            _get_dicp_opt_path(),
+            "kernel.ttshared.mlir",
             "--lower-affine",
             "--normalize-slice-ops",
             "--linalg-if-to-select",
             "--linalg-generic-to-scf",
             "--scalar-to-1d-tensor",
             f"--linalg-to-linked=global-kernel=false named-ops=true",
             "--linked-to-hivm",
-            "-o",
-            dst_path,
         ]
-        if dump_ir:
-            print(f"DEBUG dump ir[ttsharedir_to_linkedir] command: {dicp_cmd_list}")
-        ret = subprocess.run(dicp_cmd_list, capture_output=True, check=True)
-        # TODO(zmz): 修改test_path 中内容，暂时在python中处理，bishengir-compile后续会支持，去掉这里逻辑。
-        with open(dst_path, "r") as f:
-            content = f.read()
-            # 将"*xfxxx"替换成"?xfxxx"
-            content = content.replace("*xf", "?xf")
-            content = content.replace("*xi", "?xi")
-            content = content.replace("*xbf", "?xbf")
-            with open(dst_path, "w") as f:
-                f.write(content)
-
-        # 匹配形如 "memref<...> to tensor<...>" 的模式
-        pattern = r"(memref\<.*?\>)\s+to\s+(tensor\<.*?\>)"
-        with open(dst_path, "r") as f:
-            lines = f.readlines()
-        modified = []
-        for line in lines:
-            # 使用正则替换，保留memref和tensor类型，中间插入注释
-            new_line = re.sub(pattern, r"\1 // to \2", line)
-            modified.append(new_line)
-        with open(dst_path, "w") as f:
-            f.writelines(modified)
-        if dump_ir:
-            shutil.copy(dst_path, "./tmp/kernel.linkedir.mlir")
-        if replace_linked_ir is not None:
-            print(f"[DEBUG] Replace Linkedir with {replace_linked_ir}")
-            return Path(replace_linked_ir).read_text()
-        return Path(dst_path).read_text()
+        dicp_utils._dump_stage_ir(
+            content, metadata["hash"], "kernel.linkedir.mlir", cmd_list
+        )
+
+    if replace_linked_ir is not None:
+        print(f"[DEBUG] Replace Linkedir with {replace_linked_ir}")
+        return Path(replace_linked_ir).read_text()
+    return content
 
 
 def linalg_to_llir(linalg: str, metadata, opt):
@@ -721,7 +694,14 @@ def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt):
         )
         if dump_ir:
             print(f"DEBUG dump ir[bishengir-compile] command: {cmd_list}")
-        ret = subprocess.run(cmd_list, capture_output=True, check=True)
+        try:
+            ret = subprocess.run(cmd_list, capture_output=True, check=True, text=True)
+        except subprocess.CalledProcessError as e:
+            # Print compilation error details
+            print(f"bishengir-compile compilation failed with exit code {e.returncode}")
+            print(f"Stderr:\n{e.stderr}")
+            raise RuntimeError("bishengir-compile compilation failed") from e
+
         if not Path(bin_path).is_file():
             print(ret.stderr.decode("utf-8"))
         if Path(callback_path).is_file():