Add debugger and CANN profiler support for the Ascend backend

.gitignore

@@ -11,6 +11,7 @@ python/build/
 python/dist/
 python/triton*.egg-info/
 python/triton_kernels/triton*.egg-info/
+*.egg-info/
 
 python/triton/_C/*.pyd
 python/triton/_C/*.so
@@ -91,5 +92,15 @@ docs/sg_execution_times.rst
 # macOS
 .DS_Store
 
+# Local source downloads and runtime artifacts
+/.external/
+/*.lock
+/extra-info/
+/kernel_meta/
+/third_party/ascend/AscendNPU-IR/
+/third_party/flir/
+/third_party/Debugger/samples/**/source/generated_kernels/
+/third_party/Debugger/samples/**/source/kernel_sources/
+
 # Agent
 AGENTS.md

CMakeLists.txt

@@ -99,6 +99,12 @@ list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
 
 # Options
 option(TRITON_BUILD_PYTHON_MODULE "Build Python Triton bindings" OFF)
+option(FLAGTREE_ENABLE_DEBUGGER "Build FlagTree debugger support" ON)
+if(FLAGTREE_ENABLE_DEBUGGER)
+  add_compile_definitions(FLAGTREE_ENABLE_DEBUGGER=1)
+else()
+  add_compile_definitions(FLAGTREE_ENABLE_DEBUGGER=0)
+endif()
 if(FLAGTREE_BACKEND)
   option(TRITON_BUILD_PROTON "Build the Triton Proton profiler" OFF)
   option(TRITON_BUILD_UT "Build C++ Triton Unit Tests" OFF)
@@ -280,6 +286,10 @@ else()
 endif()
 include_directories(${PROJECT_SOURCE_DIR}/third_party)
 include_directories(${PROJECT_BINARY_DIR}/third_party) # Tablegen'd files
+if(FLAGTREE_ENABLE_DEBUGGER)
+  include_directories(${PROJECT_SOURCE_DIR}/third_party/Debugger/include)
+  include_directories(${PROJECT_BINARY_DIR}/third_party/Debugger/include)
+endif()
 
 # link_directories(${LLVM_LIBRARY_DIR})
 if (FLAGTREE_BACKEND MATCHES "^(cambricon|aipu|tsingmicro|enflame)$")
@@ -292,6 +302,10 @@ elseif(NOT FLAGTREE_BACKEND)
   add_subdirectory(lib)
 endif()
 
+if(FLAGTREE_ENABLE_DEBUGGER)
+  add_subdirectory(third_party/Debugger)
+endif()
+
 # TODO: Figure out which target is sufficient to fix errors; triton is
 # apparently not enough. Currently set linking libstdc++fs for all targets
 # to support some old version GCC compilers like 8.3.0.
@@ -500,6 +514,11 @@ if(TRITON_BUILD_PYTHON_MODULE)
 
   set(TRITON_BACKENDS_TUPLE "(${TRITON_BACKENDS_TUPLE})")
   add_compile_definitions(TRITON_BACKENDS_TUPLE=${TRITON_BACKENDS_TUPLE})
+  set(TRITON_DEBUGGER_BINDING_SOURCE)
+  if(FLAGTREE_ENABLE_DEBUGGER)
+    list(APPEND TRITON_DEBUGGER_BINDING_SOURCE
+      ${CMAKE_CURRENT_SOURCE_DIR}/third_party/Debugger/python/Bindings.cpp)
+  endif()
   if(FLAGTREE_BACKEND STREQUAL "cambricon")
     add_library(triton SHARED)
   elseif(FLAGTREE_BACKEND STREQUAL "ascend")
@@ -509,13 +528,15 @@ if(TRITON_BUILD_PYTHON_MODULE)
     add_library(triton SHARED ${PYTHON_SRC_PATH}/main.cc
                 ${PYTHON_SRC_PATH}/ir.cc
                 ${BUFFER_IR_SRC_PATH}/buffer_ir.cc
+                ${TRITON_DEBUGGER_BINDING_SOURCE}
                 ${PYTHON_ROOT_SRC_PATH}/gluon_ir.cc
                 ${PYTHON_ROOT_SRC_PATH}/passes.cc
                 ${PYTHON_ROOT_SRC_PATH}/interpreter.cc
                 ${PYTHON_ROOT_SRC_PATH}/llvm.cc)
   else()
     add_library(triton SHARED ${PYTHON_SRC_PATH}/main.cc
                 ${PYTHON_SRC_PATH}/ir.cc
+                ${TRITON_DEBUGGER_BINDING_SOURCE}
                 ${PYTHON_SRC_PATH}/gluon_ir.cc
                 ${PYTHON_SRC_PATH}/passes.cc
                 ${PYTHON_SRC_PATH}/interpreter.cc
@@ -586,8 +607,15 @@ if(FLAGTREE_BACKEND MATCHES "^(ascend)$")
   add_subdirectory(test)
 endif()
 
+if(FLAGTREE_ENABLE_DEBUGGER)
+  add_subdirectory(third_party/Debugger/test/lit)
+endif()
+
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)
+  if(FLAGTREE_ENABLE_DEBUGGER)
+    add_subdirectory(third_party/Debugger/test/unittest)
+  endif()
   # This target runs all the unit tests.
   add_custom_target(check-triton-unit-tests
     COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure

bin/CMakeLists.txt

@@ -1,11 +1,16 @@
 get_property(triton_libs GLOBAL PROPERTY TRITON_LIBS)
+set(flagtree_debugger_libs)
+if(FLAGTREE_ENABLE_DEBUGGER)
+  list(APPEND flagtree_debugger_libs FlagTreeDebugger)
+endif()
 
 add_llvm_executable(triton-opt triton-opt.cpp PARTIAL_SOURCES_INTENDED)
 
 # TODO: what's this?
 llvm_update_compile_flags(triton-opt)
 target_link_libraries(triton-opt PRIVATE
   ${triton_libs}
+  ${flagtree_debugger_libs}
   # tests
   TritonTestAnalysis
   TritonTestDialect
@@ -27,6 +32,7 @@ mlir_check_all_link_libraries(triton-reduce)
 llvm_update_compile_flags(triton-reduce)
 target_link_libraries(triton-reduce PRIVATE
   ${triton_libs}
+  ${flagtree_debugger_libs}
   # tests
   TritonTestAnalysis
   TritonTestDialect
@@ -47,6 +53,7 @@ add_llvm_executable(triton-lsp triton-lsp.cpp PARTIAL_SOURCES_INTENDED)
 llvm_update_compile_flags(triton-lsp)
 target_link_libraries(triton-lsp PRIVATE
   ${triton_libs}
+  ${flagtree_debugger_libs}
   # tests
   TritonTestAnalysis
   TritonTestDialect
@@ -86,6 +93,7 @@ export_executable_symbols_for_plugins(triton-llvm-opt)
 add_llvm_executable(triton-tensor-layout triton-tensor-layout.cpp PARTIAL_SOURCES_INTENDED)
 target_link_libraries(triton-tensor-layout PRIVATE
   ${triton_libs}
+  ${flagtree_debugger_libs}
   TritonTestAnalysis
   TritonTestDialect
   TritonTestProton

bin/RegisterTritonDialects.h

@@ -1,4 +1,9 @@
 #pragma once
+#if FLAGTREE_ENABLE_DEBUGGER
+#include "Debugger/IR/Dialect.h"
+#include "Debugger/Instrumentation/Passes.h"
+#include "Debugger/Metadata/Passes.h"
+#endif
 #include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
 #include "amd/include/TritonAMDGPUTransforms/Passes.h"
 #include "nvidia/include/Dialect/NVGPU/IR/Dialect.h"
@@ -60,6 +65,10 @@ void registerTestScopeIdAllocationPass();
 
 inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerAllPasses();
+#if FLAGTREE_ENABLE_DEBUGGER
+  mlir::flagtree::debugger::registerFlagTreeDebuggerInstrumentationPasses();
+  mlir::flagtree::debugger::registerFlagTreeDebuggerMetadataPasses();
+#endif
   mlir::triton::registerTritonPasses();
   mlir::triton::gpu::registerTritonGPUPasses();
   mlir::triton::nvidia_gpu::registerTritonNvidiaGPUPasses();
@@ -136,6 +145,9 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
       mlir::triton::amdgpu::TritonAMDGPUDialect,
       mlir::triton::proton::ProtonDialect,
       mlir::triton::proton::gpu::ProtonGPUDialect, mlir::ROCDL::ROCDLDialect,
+#if FLAGTREE_ENABLE_DEBUGGER
+      mlir::flagtree::debugger::FlagTreeDebugDialect,
+#endif
 #ifdef __TLE__
       mlir::triton::gluon::GluonDialect,
       mlir::triton::tle::TleDialect // flagtree tle raw

pyproject.toml

@@ -38,3 +38,5 @@ ignore = ["E501", "E701", "E731", "E741"]
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401", "E402"]
+"third_party/Debugger/test/python/language/test_module_a_*.py" = ["E402"]
+"third_party/Debugger/test/python/runtime/test_module_a_*.py" = ["E402"]

python/src/ir.cc

@@ -38,6 +38,9 @@
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h"
 #include "triton/Tools/Sys/GetEnv.hpp"
+#if FLAGTREE_ENABLE_DEBUGGER
+#include "Debugger/IR/Dialect.h"
+#endif
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/SourceMgr.h"
@@ -365,6 +368,9 @@ void init_triton_ir(py::module &&m) {
                     ::mlir::gpu::GPUDialect, cf::ControlFlowDialect,
                     LLVM::LLVMDialect, mlir::ub::UBDialect,
                     mlir::triton::gluon::GluonDialect,
+#if FLAGTREE_ENABLE_DEBUGGER
+                    mlir::flagtree::debugger::FlagTreeDebugDialect,
+#endif
                     DLTIDialect,                  // flagtree tle raw
                     mlir::triton::tle::TleDialect // flagtree tle raw
                     >();
@@ -1835,6 +1841,29 @@ void init_triton_ir(py::module &&m) {
       // Force GPU barrier
       .def("create_barrier",
            [](TritonOpBuilder &self) { self.create<mlir::gpu::BarrierOp>(); })
+#if FLAGTREE_ENABLE_DEBUGGER
+      .def("create_debug_collect_begin",
+           [](TritonOpBuilder &self, int32_t level) {
+             auto attr = self.getBuilder().getI32IntegerAttr(level);
+             mlir::IntegerAttr addrAttr;
+             self.create<mlir::flagtree::debugger::CollectBeginOp>(
+                 attr, addrAttr, /*scope_id=*/nullptr);
+           })
+      .def("create_debug_collect_begin",
+           [](TritonOpBuilder &self, int32_t level, int32_t addrLevel) {
+             auto attr = self.getBuilder().getI32IntegerAttr(level);
+             mlir::IntegerAttr addrAttr;
+             if (addrLevel >= 0)
+               addrAttr = self.getBuilder().getI32IntegerAttr(addrLevel);
+             self.create<mlir::flagtree::debugger::CollectBeginOp>(
+                 attr, addrAttr, /*scope_id=*/nullptr);
+           })
+      .def("create_debug_collect_end",
+           [](TritonOpBuilder &self) {
+             self.create<mlir::flagtree::debugger::CollectEndOp>(
+                 /*scope_id=*/nullptr);
+           })
+#endif
       // Make a block pointer (tensor pointer in Triton IR)
       .def("create_make_block_ptr",
            [](TritonOpBuilder &self, Value &base, std::vector<Value> &shape,

python/src/main.cc

@@ -44,6 +44,9 @@ void init_triton_interpreter(pybind11::module &&m);
 void init_triton_passes(pybind11::module &&m);
 void init_triton_stacktrace_hook(pybind11::module &m);
 void init_gluon_ir(pybind11::module &&m);
+#if FLAGTREE_ENABLE_DEBUGGER
+void init_triton_debugger(pybind11::module &&m);
+#endif
 FOR_EACH_P(DECLARE_BACKEND, TRITON_BACKENDS_TUPLE)
 
 PYBIND11_MODULE(libtriton, m) {
@@ -55,5 +58,8 @@ PYBIND11_MODULE(libtriton, m) {
   init_triton_interpreter(m.def_submodule("interpreter"));
   init_triton_llvm(m.def_submodule("llvm"));
   init_gluon_ir(m.def_submodule("gluon_ir"));
+#if FLAGTREE_ENABLE_DEBUGGER
+  init_triton_debugger(m.def_submodule("debugger"));
+#endif
   FOR_EACH_P(INIT_BACKEND, TRITON_BACKENDS_TUPLE)
 }

python/src/passes.cc

@@ -15,6 +15,13 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#if FLAGTREE_ENABLE_DEBUGGER
+#include "Debugger/Instrumentation/Passes.h"
+#include "Debugger/Metadata/Passes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include <cstdint>
+#endif
+
 namespace py = pybind11;
 
 void init_triton_analysis(py::module &&m) {
@@ -119,6 +126,69 @@ void init_gluon_passes(py::module &&m) {
   ADD_PASS_WRAPPER_0("add_inliner", gluon::createGluonInline);
 }
 
+#if FLAGTREE_ENABLE_DEBUGGER
+void init_flagtree_debug_passes(py::module &&m) {
+  using namespace mlir;
+  using namespace mlir::flagtree::debugger;
+  m.def("has_debug_collect_markers",
+        [](ModuleOp mod) { return hasDebugCollectMarkers(mod); });
+  m.def("insert_default_debug_collect_markers", [](ModuleOp mod, int32_t level,
+                                                   int32_t addrLevel) {
+    return succeeded(insertDefaultDebugCollectMarkers(mod, level, addrLevel));
+  });
+  m.def("get_debug_tracked_op_table_json",
+        [](ModuleOp mod) { return getDebugTrackedOpTableJson(mod); });
+  m.def("get_debug_kernel_metadata_json",
+        [](ModuleOp mod) { return getDebugKernelMetadataJson(mod); });
+  m.def("get_debug_kernel_id",
+        [](ModuleOp mod) { return getDebugKernelId(mod); });
+  m.def("get_debug_records_per_instance",
+        [](ModuleOp mod) { return getDebugRecordsPerInstance(mod); });
+  m.def("get_debug_record_size",
+        [](ModuleOp mod) { return getDebugRecordSize(mod); });
+  m.def("get_debug_record_layout",
+        [](ModuleOp mod) { return getDebugRecordLayout(mod); });
+  m.def("get_debug_record_plan_json",
+        [](ModuleOp mod) { return getDebugRecordPlanJson(mod); });
+  m.def("get_debug_full_dump_payload_bytes_per_instance", [](ModuleOp mod) {
+    return getDebugFullDumpPayloadBytesPerInstance(mod);
+  });
+  m.def("get_debug_full_dump_plan_json",
+        [](ModuleOp mod) { return getDebugFullDumpPlanJson(mod); });
+  m.def("set_debug_kernel_id_seed", [](ModuleOp mod, const std::string &seed) {
+    setDebugKernelIdSeed(mod, seed);
+  });
+  m.def("set_debug_hidden_arg_abi_enabled", [](ModuleOp mod, bool enabled) {
+    setDebugHiddenArgAbiEnabled(mod, enabled);
+  });
+  m.def("set_debug_addr_level", [](ModuleOp mod, int32_t addrLevel) {
+    setDebugAddrLevel(mod, addrLevel);
+  });
+  m.def("set_debug_timeline_enabled", [](ModuleOp mod, bool enabled) {
+    setDebugTimelineEnabled(mod, enabled);
+  });
+  m.def("set_debug_timeline_only",
+        [](ModuleOp mod, bool enabled) { setDebugTimelineOnly(mod, enabled); });
+  m.def("assign_debug_collect_scope_ids_without_erase", [](ModuleOp mod) {
+    return succeeded(assignDebugCollectScopeIdsWithoutErase(mod));
+  });
+  m.def("assign_debug_op_ids_and_metadata_without_pass_manager",
+        [](ModuleOp mod) {
+          return succeeded(assignDebugOpIdsAndMetadataWithoutPassManager(mod));
+        });
+  m.def("erase_debug_collect_markers",
+        [](ModuleOp mod) { eraseDebugCollectMarkers(mod); });
+  m.def("has_triton_tensor_pointer_types",
+        [](ModuleOp mod) { return hasTritonTensorPointerTypes(mod); });
+  ADD_PASS_WRAPPER_0("add_resolve_debug_scope", createResolveDebugScopePass);
+  ADD_PASS_WRAPPER_0("add_assign_debug_op_id", createAssignOpIdPass);
+  ADD_PASS_WRAPPER_0("add_insert_instrumentation",
+                     createInsertInstrumentationPass);
+  ADD_PASS_WRAPPER_0("add_simplify_record_memref_writes",
+                     createSimplifyRecordMemrefWritesPass);
+}
+#endif
+
 void init_triton_passes(py::module &&m) {
   init_triton_analysis(m.def_submodule("analysis"));
   init_triton_passes_common(m.def_submodule("common"));
@@ -127,4 +197,7 @@ void init_triton_passes(py::module &&m) {
   init_triton_passes_ttgpuir(m.def_submodule("ttgpuir"));
   init_triton_passes_llvmir(m.def_submodule("llvmir"));
   init_gluon_passes(m.def_submodule("gluon"));
+#if FLAGTREE_ENABLE_DEBUGGER
+  init_flagtree_debug_passes(m.def_submodule("flagtree_debug"));
+#endif
 }

python/triton/__init__.py

@@ -36,6 +36,8 @@
     "CompilationError",
     "compile",
     "Config",
+    "disable_debug",
+    "enable_debug",
     "constexpr_function",
     "FutureKernel",
     "heuristics",
@@ -80,3 +82,17 @@ def next_power_of_2(n: int):
     n |= n >> 32
     n += 1
     return n
+
+
+def enable_debug(level: int = 1, addr_level: int = 0):
+    """Enable FlagTree debugger instrumentation for subsequently compiled kernels."""
+    from triton.runtime import debugger
+
+    return debugger.activate(level=level, addr_level=addr_level)
+
+
+def disable_debug():
+    """Disable FlagTree debugger instrumentation."""
+    from triton.runtime import debugger
+
+    return debugger.deactivate()