Add a pass management interface for Ascend

CMakeLists.txt

@@ -135,6 +135,10 @@ function(add_triton_object name)
   )
 
   if (FLAGTREE_BACKEND STREQUAL "ascend")
+    set(ASCENDNPU_IR_SRC_DIR ${PROJECT_SOURCE_DIR}/third_party/ascendnpu-ir)
+    set(ASCENDNPU_IR_BINARY_DIR ${PROJECT_BINARY_DIR}/third_party/ascendnpu-ir)
+    include_directories(${ASCENDNPU_IR_SRC_DIR}/bishengir/include)
+    include_directories(${ASCENDNPU_IR_BINARY_DIR}/bishengir/include)
     set(patched_depends "")
     foreach(dep ${ARG_DEPENDS})
       list(FIND PATCHED_TRITON_DEPENDS "${dep}" index)

third_party/ascend/CMakeLists.txt

@@ -2,6 +2,10 @@
 #add_subdirectory(test)
 
 add_triton_plugin(TritonAscend ${CMAKE_CURRENT_SOURCE_DIR}/triton_ascend.cpp)
-target_include_directories(TritonAscend PRIVATE ${CMAKE_SOURCE_DIR}/third_party/flir/include)
+target_include_directories(TritonAscend PRIVATE 
+	${CMAKE_SOURCE_DIR}/third_party/flir/include
+		${CMAKE_BINARY_DIR}/third_party/flir/include)
+
+add_dependencies(TritonAscend TritonToLinalgIncubatedPassIncGen)
 
 add_triton_library(Registrar Registrar.cc)

third_party/ascend/backend/compiler.py

@@ -68,14 +68,56 @@ def make_ttir(mod, metadata, opt):
 def ttir_to_linalg(mod, metadata, opt, *, named_ops=False):
     pm = ir.pass_manager(mod.context)
     pm.enable_debug()
+    enable_nd2nz_on_vector = metadata["enable_nd2nz_on_vector"]
     # Add pass here.
-    ascend.passes.convert.add_triton_to_linalg_pipeline(pm)
+    #ascend.passes.convert.add_triton_to_linalg_pipeline(pm)
+    ascend.passes.convert.add_triton_linearize(pm)
+    ascend.passes.convert.add_triton_to_hivm(pm)
+    ascend.passes.convert.add_triton_to_hfusion(pm)
+    ascend.passes.convert.add_triton_to_llvm(pm)
+    ascend.passes.convert.add_triton_to_linalg_incubated(
+    pm,
+    global_kernel=False,
+    named_ops=True,
+    enable_nd2nz_on_vector=enable_nd2nz_on_vector
+    )
     pm.run(mod)
     return str(mod)
+    # use triton_adapter to lower Triton-MLIR to linalg
+    # Get Triton-MLIR as string
     '''
-    with open('/home/zhengyang/FlagTree/triton-op-ir/ops/01_vector_add/cache_ascend/add_kernel.ttadapter', 'r', encoding='utf-8') as f:
-        content = f.read()
-    return content
+    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.ttadapter.mlir")
+        Path(src_path).write_text(ttir_code)
+        triton_adapter_opt_path = _get_triton_adapter_opt_path()
+
+        enable_nd2nz_on_vector = metadata["enable_nd2nz_on_vector"]
+        cmd_list = [
+            triton_adapter_opt_path,
+            src_path,
+            "--discrete-mask-access-conversion",
+            "--triton-to-annotation",
+            "--triton-to-unstructure",
+            "--triton-to-hivm",
+            "--bubble-up-operation",
+            f"--triton-to-linalg=global-kernel=false named-ops={named_ops} "\
+            f"enable-nd2nz-on-vector={enable_nd2nz_on_vector}",
+            "-o",
+            dst_path,
+        ]
+        if _is_ascend_sanitizer_enabled() or not _is_debug_line_info_disabled():
+            cmd_list += ["--mlir-print-debuginfo"]  # pass debug info
+
+        ret = subprocess.run(cmd_list, capture_output=True, check=True)
+        if opt.debug:
+            dump_manager = get_dump_manager(metadata["hash"])
+            dump_manager.put(
+                Path(dst_path).read_text(), "kernel.ttadapter.mlir", binary=False
+            )
+
+        return Path(dst_path).read_text()
     '''
 
 

third_party/ascend/triton_ascend.cpp

@@ -4,6 +4,11 @@
 #include "mlir/Pass/PassManager.h"
 #include "passes.h"
 #include "triton-shared/Conversion/TritonToLinalgExperimental/TritonToLinalgExperimental.h"
+#include "triton-shared/TritonToHFusion/TritonToHFusion.h"
+#include "triton-shared/TritonLinearize/TritonLinearize.h"
+#include "triton-shared/TritonToLinalgIncubated/TritonToLinalgIncubatedPass.h"
+#include "triton-shared/TritonToHIVM/TritonToHIVM.h"
+#include "triton-shared/TritonToLLVM/TritonToLLVM.h"
 
 #define PY_SSIZE_T_CLEAN
 #include <pybind11/pybind11.h>
@@ -12,10 +17,32 @@ namespace py = pybind11;
 void init_triton_ascend_passes_convert(py::module &&m) {
   ADD_PASS_WRAPPER_0("add_triton_to_linalg_pipeline",
                      mlir::triton::createTritonToLinalgExperimentalPass);
+  ADD_PASS_WRAPPER_0("add_triton_linearize",
+                     mlir::triton::createTritonLinearizePass);
+  ADD_PASS_WRAPPER_0("add_triton_to_hivm",
+                     mlir::triton::createTritonToHIVMPass);
+  ADD_PASS_WRAPPER_0("add_triton_to_hfusion",
+                     mlir::triton::createTritonToHFusionPass);
+  ADD_PASS_WRAPPER_0("add_triton_to_llvm",
+                     mlir::triton::createTritonToLLVMPass);
+  m.def(
+      "add_triton_to_linalg_incubated",
+      [](mlir::PassManager &pm,
+         bool global_kernel,
+         bool named_ops,
+         bool enable_nd2nz_on_vector) {
+        pm.addPass(mlir::triton::Incubated::createTritonToLinalgIncubatedPass(
+            global_kernel, named_ops, enable_nd2nz_on_vector));
+      },
+      py::arg("pm"),
+      py::arg("global_kernel"),
+      py::arg("named_ops"),
+      py::arg("enable_nd2nz_on_vector"));
 }
 
+
 // register ascend passes to triton
 void init_triton_ascend(py::module &&m) {
   auto passes = m.def_submodule("passes");
   init_triton_ascend_passes_convert(passes.def_submodule("convert"));
-}
+}