diff --git a/.gitignore b/.gitignore index 7a9396476..ae77de517 100644 --- a/.gitignore +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 3b39b5421..7294ce2de 100644 --- a/CMakeLists.txt +++ b/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,6 +528,7 @@ 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 @@ -516,6 +536,7 @@ if(TRITON_BUILD_PYTHON_MODULE) 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 diff --git a/bin/CMakeLists.txt b/bin/CMakeLists.txt index cc0b4fa15..cbc9516a0 100644 --- a/bin/CMakeLists.txt +++ b/bin/CMakeLists.txt @@ -1,4 +1,8 @@ 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) @@ -6,6 +10,7 @@ add_llvm_executable(triton-opt triton-opt.cpp PARTIAL_SOURCES_INTENDED) 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 diff --git a/bin/RegisterTritonDialects.h b/bin/RegisterTritonDialects.h index 39bd4a3a9..1360cc9d3 100644 --- a/bin/RegisterTritonDialects.h +++ b/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 ®istry) { 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 ®istry) { 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 diff --git a/pyproject.toml b/pyproject.toml index b7d791b47..ff44e77e2 100644 --- a/pyproject.toml +++ b/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"] diff --git a/python/src/ir.cc b/python/src/ir.cc index d2529340e..f245923d5 100644 --- a/python/src/ir.cc +++ b/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(); }) +#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( + 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( + attr, addrAttr, /*scope_id=*/nullptr); + }) + .def("create_debug_collect_end", + [](TritonOpBuilder &self) { + self.create( + /*scope_id=*/nullptr); + }) +#endif // Make a block pointer (tensor pointer in Triton IR) .def("create_make_block_ptr", [](TritonOpBuilder &self, Value &base, std::vector &shape, diff --git a/python/src/main.cc b/python/src/main.cc index 9bbd5e636..68388a286 100644 --- a/python/src/main.cc +++ b/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) } diff --git a/python/src/passes.cc b/python/src/passes.cc index c328137d4..ca78e57a1 100644 --- a/python/src/passes.cc +++ b/python/src/passes.cc @@ -15,6 +15,13 @@ #include #include +#if FLAGTREE_ENABLE_DEBUGGER +#include "Debugger/Instrumentation/Passes.h" +#include "Debugger/Metadata/Passes.h" +#include "mlir/IR/BuiltinOps.h" +#include +#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 } diff --git a/python/triton/__init__.py b/python/triton/__init__.py index dae98edc7..d9ac67896 100644 --- a/python/triton/__init__.py +++ b/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() diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py index 66b3ef4e3..b908c183d 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py @@ -505,6 +505,53 @@ def set_value(self, name: str, value: Union[base_value, constexpr]) -> None: self.lscope[name] = value self.local_defs[name] = value + def _statement_source(self, node): + source = None + if hasattr(node, "lineno"): + try: + source = ast.get_source_segment(self.jit_fn.src, node) + except Exception: + source = None + if source is None and hasattr(ast, "unparse"): + try: + source = ast.unparse(node) + except Exception: + source = None + return "" if source is None else " ".join(source.strip().split()) + + def _statement_id(self, node): + line = int(self.begin_line + getattr(node, "lineno", 0)) + col = int(getattr(node, "col_offset", 0)) + return max(0, min(line * 1000 + min(col, 999), (1 << 31) - 1)) + + def _annotate_assignment_value(self, node, name, value, statement_source): + if not _is_triton_tensor(value): + return + handle = getattr(value, "handle", None) + if handle is None or not hasattr(handle, "set_attr"): + return + if statement_source: + handle.set_attr("flagtree.debug.triton_statement", self.builder.get_string_attr(statement_source)) + handle.set_attr("flagtree.debug.statement_result_name", self.builder.get_string_attr(str(name))) + handle.set_attr("flagtree.debug.statement_id", self.builder.get_int32_attr(self._statement_id(node))) + + def _annotate_assignment_target(self, node, target, value, statement_source): + if isinstance(target, ast.Name): + self._annotate_assignment_value(node, target.id, value, statement_source) + elif isinstance(target, ast.Tuple): + for child, child_value in zip(target.elts, value.values): + self._annotate_assignment_target(node, child, child_value, statement_source) + + def _annotate_expression_value(self, node, value, statement_source): + if not _is_triton_tensor(value): + return + handle = getattr(value, "handle", None) + if handle is None or not hasattr(handle, "set_attr"): + return + if statement_source: + handle.set_attr("flagtree.debug.triton_statement", self.builder.get_string_attr(statement_source)) + handle.set_attr("flagtree.debug.statement_id", self.builder.get_int32_attr(self._statement_id(node))) + def _get_insertion_point_and_loc(self): # XXX: this is a hack to get the location of the insertion point. # The insertion point's location could be invalid sometimes, @@ -745,6 +792,7 @@ def _sanitize_value(value): values = _sanitize_value(self.visit(node.value)) else: values = _sanitize_value(self.visit(node.value)) + self._annotate_assignment_target(node, target, values, self._statement_source(node)) self.assignTarget(target, values) def visit_AugAssign(self, node): @@ -1457,6 +1505,10 @@ def call_Function(self, node, fn, args, kws): extra_kwargs["_semantic"] = self.semantic if '_generator' in sig.parameters: extra_kwargs['_generator'] = self + if '_debug_statement_source' in sig.parameters: + extra_kwargs['_debug_statement_source'] = self._statement_source(node) + if '_debug_statement_id' in sig.parameters: + extra_kwargs['_debug_statement_id'] = self._statement_id(node) try: # Special handling for tl.load with hints hint_trigger("inject_kwargs_with_hints", fn, flagtree_hints, node.lineno, kws) @@ -1599,7 +1651,9 @@ def visit_Attribute(self, node): def visit_Expr(self, node): node.value._is_unused = True - ast.NodeVisitor.generic_visit(self, node) + value = self.visit(node.value) + self._annotate_expression_value(node, value, self._statement_source(node)) + return value def visit_NoneType(self, node): return None diff --git a/python/triton/compiler/compiler.py b/python/triton/compiler/compiler.py index 1f38e9ddf..7c345010e 100644 --- a/python/triton/compiler/compiler.py +++ b/python/triton/compiler/compiler.py @@ -11,6 +11,7 @@ from ..runtime.driver import driver from ..tools.disasm import get_sass from pathlib import Path +import importlib import re import functools import os @@ -278,6 +279,13 @@ def compile(src, target=None, options=None, _env_vars=None): **options.__dict__, **env_vars, } + # Debug instrumentation consumes these fields during the TTIR stage. They + # must come from GPUTarget: backend options do not define backend_name on + # every target (notably Ascend), and filling them after lowering is too late + # for both kernel metadata and transfer-backend selection. + metadata["backend_name"] = str(target.backend) + metadata["debug_backend_name"] = str(target.backend) + metadata["debug_target_name"] = str(target.arch) metadata["triton_version"] = __version__ # run compilation pipeline and populate metadata stages = dict() @@ -343,6 +351,19 @@ def compile(src, target=None, options=None, _env_vars=None): module = next_module if compilation_listener: timer.stage_finished(ext) + if str(metadata.get("instrumentation_mode", "")).startswith("debugger"): + try: + debugger_config = importlib.import_module("triton.runtime.debugger").current_compile_config() + except Exception: + debugger_config = {} + + metadata.update(debugger_config) + metadata.setdefault("debug_enabled", True) + metadata.setdefault("debug_protocol_version", 2) + metadata.setdefault("debug_kernel_id", int(hash[:8], 16) or 1) + metadata.setdefault("debug_backend_name", metadata.get("backend_name", "")) + metadata.setdefault("debug_target_name", str(getattr(target, "arch", ""))) + metadata.setdefault("debug_kernel_name", metadata.get("name", "")) # write-back metadata metadata_group[metadata_filename] = fn_cache_manager.put(json.dumps(metadata, default=vars), metadata_filename, binary=False) @@ -485,10 +506,27 @@ def run(self): return self._run def launch_metadata(self, grid, stream, *args): - if knobs.runtime.launch_enter_hook is None: + debugger_active = False + try: + debugger_active = importlib.import_module("triton.runtime.debugger").is_active() + except Exception: + pass + if (knobs.runtime.launch_enter_hook is None and not getattr(self.metadata, "debug_enabled", False) + and not debugger_active): return None self._init_handles() - ret = LazyDict({"name": self.name, "function": self.function, "stream": stream}) + grid_size = len(grid) + normalized_grid = ( + int(grid[0]), + int(grid[1]) if grid_size > 1 else 1, + int(grid[2]) if grid_size > 2 else 1, + ) + ret = LazyDict({ + "name": self.name, + "function": self.function, + "stream": stream, + "grid": normalized_grid, + }) if not isinstance(self.src, ASTSource) or self.src.fn.launch_metadata is None: return ret arg_dict = {name: arg for name, arg in zip(self.src.fn.arg_names, args)} @@ -502,8 +540,44 @@ def runner(*args, stream=None): if stream is None: device = driver.active.get_current_device() stream = driver.active.get_current_stream(device) + debug_ctx = None + launcher_manages_debug = bool(getattr(self._run, "manages_debug_launch", False)) + records_per_instance = int(getattr(self.metadata, "debug_records_per_instance", 0)) + has_hidden_arg = (bool(getattr(self.metadata, "debug_launch_hidden_arg", False)) + and records_per_instance > 0) + if (getattr(self.metadata, "debug_launch_hidden_arg", False) and records_per_instance <= 0): + self._run.debug_launch_hidden_arg = False + self._run.debug_ctrl_ptr = 0 + if (getattr(self.metadata, "debug_enabled", False) and has_hidden_arg and not launcher_manages_debug): + from triton.runtime.debug_collect_runtime import default_debug_collect_runtime + + grid_size = len(grid) + runtime_metadata = { + "grid": ( + int(grid[0]), + int(grid[1]) if grid_size > 1 else 1, + int(grid[2]) if grid_size > 2 else 1, + ), + "records_per_instance": records_per_instance, + } + debug_ctx = default_debug_collect_runtime.prepare(self.metadata, stream, runtime_metadata) + self._debug_ctrl_ptr = default_debug_collect_runtime.hidden_arg(debug_ctx) + if has_hidden_arg and not launcher_manages_debug: + from triton.compiler.flagtree_debug import prepare_launch_debug_ctrl + + prepare_launch_debug_ctrl(self, stream) launch_metadata = self.launch_metadata(grid, stream, *args) - self.run(grid[0], grid[1], grid[2], stream, self.function, self.packed_metadata, launch_metadata, - knobs.runtime.launch_enter_hook, knobs.runtime.launch_exit_hook, *args) + try: + self.run(grid[0], grid[1], grid[2], stream, self.function, self.packed_metadata, launch_metadata, + knobs.runtime.launch_enter_hook, knobs.runtime.launch_exit_hook, *args) + if debug_ctx is not None: + from triton.runtime.debug_collect_runtime import default_debug_collect_runtime + + default_debug_collect_runtime.export(debug_ctx, stream) + finally: + if debug_ctx is not None: + from triton.runtime.debug_collect_runtime import default_debug_collect_runtime + + default_debug_collect_runtime.release(debug_ctx) return runner diff --git a/python/triton/compiler/flagtree_debug.py b/python/triton/compiler/flagtree_debug.py new file mode 100644 index 000000000..c8f33fb3e --- /dev/null +++ b/python/triton/compiler/flagtree_debug.py @@ -0,0 +1,200 @@ +"""FlagTree debugger: TTIR pass orchestration and compile-time metadata keys.""" +from __future__ import annotations + +import json +import os + +from triton._C.libtriton import ir, passes + +_instrumentation_mode = "" +_DISABLED_BUILD_MESSAGE = ("FlagTree debugger support is not available in this build; rebuild with " + "-DFLAGTREE_ENABLE_DEBUGGER=ON") + + +def set_instrumentation_mode(mode: str) -> None: + global _instrumentation_mode + _instrumentation_mode = str(mode or "") + + +def get_instrumentation_mode() -> str: + return _instrumentation_mode + + +def _get_debug_passes(): + return getattr(passes, "flagtree_debug", None) + + +def _run_pass_manager(pm, mod, description: str) -> None: + try: + pm.run(mod, description) + except TypeError: + pm.run(mod) + + +def _debug_launch_hidden_arg_enabled() -> bool: + # Keep the environment variable as a compatibility hook for subprocesses + # and older scripts, but make the Python debugger API the normal path. + if os.environ.get("TRITON_FLAGTREE_DEBUG_LAUNCH_PTR", "") == "1": + return True + try: + from triton.runtime import debugger + + return debugger.is_active() + except Exception: + return False + + +def _kernel_internal_timeline_supported() -> bool: + try: + import triton + + backend = str(triton.runtime.driver.active.get_current_target().backend).lower() + except Exception: + return False + return backend in {"ascend", "npu", "cann"} + + +def _finish_metadata_only_tensor_pointer_debug(fd, mod, metadata: dict) -> bool: + has_tensor_pointer = getattr(fd, "has_triton_tensor_pointer_types", None) + if has_tensor_pointer is None or not has_tensor_pointer(mod): + return False + + if not fd.assign_debug_collect_scope_ids_without_erase(mod): + raise RuntimeError("failed to resolve debug collect scopes") + assign_metadata = getattr(fd, "assign_debug_op_ids_and_metadata_without_pass_manager", None) + erase_markers = getattr(fd, "erase_debug_collect_markers", None) + if assign_metadata is None or erase_markers is None: + raise RuntimeError("debug tensor-pointer metadata fallback is unavailable") + if not assign_metadata(mod): + raise RuntimeError("failed to assign debug op ids") + erase_markers(mod) + + tracked_table_json = fd.get_debug_tracked_op_table_json(mod) + kernel_metadata_json = fd.get_debug_kernel_metadata_json(mod) + metadata["debug_kernel_id"] = fd.get_debug_kernel_id(mod) + metadata["debug_metadata_json"] = kernel_metadata_json + metadata["debug_tracked_table"] = json.loads(tracked_table_json) + metadata["debug_records_per_instance"] = 0 + metadata["debug_record_size"] = 64 + metadata["debug_record_layout"] = "" + metadata["debug_record_plan"] = [] + metadata["debug_launch_hidden_arg"] = False + metadata["debug_metadata_only_reason"] = "triton_tensor_pointer" + return True + + +def run_ttir_debug_passes_if_needed(mod, metadata: dict) -> None: + """If the module contains debug collect markers, run debug passes and set metadata.""" + fd = _get_debug_passes() + if fd is None: + if get_instrumentation_mode().startswith("debugger"): + raise RuntimeError(_DISABLED_BUILD_MESSAGE) + return + has_markers = fd.has_debug_collect_markers(mod) + auto_collect = get_instrumentation_mode() == "debugger_auto" + + if auto_collect: + try: + from triton.runtime import debugger + + debug_config = debugger.current_compile_config() + except Exception: + debug_config = {} + level = int(debug_config.get("debug_record_level", 1)) + addr_level = int(debug_config.get("debug_addr_level", 0)) + has_markers = bool(fd.insert_default_debug_collect_markers(mod, level, addr_level)) + + if not has_markers: + metadata["debug_enabled"] = False + metadata["debug_launch_hidden_arg"] = False + return + try: + from triton.runtime import debugger + + debug_config = debugger.current_compile_config() + except Exception: + debug_config = {} + + metadata["debug_enabled"] = True + metadata["debug_protocol_version"] = 2 + metadata["debug_record_level"] = int(debug_config.get("debug_record_level", 1)) + metadata["debug_addr_level"] = int(debug_config.get("debug_addr_level", 0)) + metadata["debug_export_mode"] = debug_config.get("debug_export_mode", "POST_KERNEL_EXPORT") + if "debug_record_capacity" in debug_config: + metadata["debug_record_capacity"] = int(debug_config["debug_record_capacity"]) + metadata["debug_launch_hidden_arg"] = _debug_launch_hidden_arg_enabled() + fd.set_debug_kernel_id_seed(mod, str(metadata.get("hash") or "")) + fd.set_debug_hidden_arg_abi_enabled(mod, bool(metadata["debug_launch_hidden_arg"])) + fd.set_debug_addr_level(mod, int(metadata["debug_addr_level"])) + fd.set_debug_timeline_enabled(mod, bool(auto_collect and _kernel_internal_timeline_supported())) + fd.set_debug_timeline_only(mod, bool(auto_collect)) + + if _finish_metadata_only_tensor_pointer_debug(fd, mod, metadata): + return + + metadata_pm = ir.pass_manager(mod.context) + metadata_pm.enable_debug() + fd.add_resolve_debug_scope(metadata_pm) + fd.add_assign_debug_op_id(metadata_pm) + # Debug collect markers in called helper functions are ignored by the + # metadata pass to avoid call-graph id propagation. After those markers are + # erased, run the regular inliner so helper calls such as FlagGems + # triton_lang_extension.program_id return to the call-free TTIR shape that + # Ascend's ttadapter path expects. + passes.common.add_inliner(metadata_pm) + passes.common.add_symbol_dce(metadata_pm) + _run_pass_manager(metadata_pm, mod, "flagtree_debug_collect_metadata") + + tracked_table_json = fd.get_debug_tracked_op_table_json(mod) + kernel_metadata_json = fd.get_debug_kernel_metadata_json(mod) + metadata["debug_kernel_id"] = fd.get_debug_kernel_id(mod) + metadata["debug_metadata_json"] = kernel_metadata_json + metadata["debug_tracked_table"] = json.loads(tracked_table_json) + + instrumentation_pm = ir.pass_manager(mod.context) + instrumentation_pm.enable_debug() + fd.add_insert_instrumentation(instrumentation_pm) + passes.common.add_cse(instrumentation_pm) + passes.common.add_canonicalizer(instrumentation_pm) + _run_pass_manager(instrumentation_pm, mod, "flagtree_debug_collect") + metadata["debug_records_per_instance"] = int(fd.get_debug_records_per_instance(mod)) + metadata["debug_record_size"] = int(fd.get_debug_record_size(mod)) + metadata["debug_record_layout"] = fd.get_debug_record_layout(mod) + metadata["debug_record_plan"] = json.loads(fd.get_debug_record_plan_json(mod)) + metadata["debug_full_dump_payload_bytes_per_instance"] = int(fd.get_debug_full_dump_payload_bytes_per_instance(mod)) + metadata["debug_full_dump_plan"] = json.loads(fd.get_debug_full_dump_plan_json(mod)) + if metadata["debug_records_per_instance"] <= 0: + # The user may request dynamic debugger collection, but the IR pass is + # the source of truth for whether a hidden-arg ABI was actually added. + # Metadata-only regions must not make the launcher append + # __debug_ctrl_ptr, otherwise Ascend kernel arguments become misaligned. + metadata["debug_launch_hidden_arg"] = False + + +def run_ttadapter_debug_passes_if_needed(mod, metadata: dict) -> None: + """Simplify debugger-only TTAdapter writes after hidden arg becomes memref.""" + if not metadata.get("debug_enabled", False): + return + if not metadata.get("debug_launch_hidden_arg", False): + return + if int(metadata.get("debug_records_per_instance", 0)) <= 0: + return + + fd = _get_debug_passes() + if fd is None: + raise RuntimeError(_DISABLED_BUILD_MESSAGE) + + pm = ir.pass_manager(mod.context) + pm.enable_debug() + fd.add_simplify_record_memref_writes(pm) + _run_pass_manager(pm, mod, "flagtree_debug_ttadapter") + + +def prepare_launch_debug_ctrl(compiled_kernel, stream) -> None: + """Set CudaLauncher.debug_ctrl_ptr before launch through the transfer engine.""" + del stream # reserved for async engine / stream-ordered alloc + if compiled_kernel._run is None: + compiled_kernel._init_handles() + launcher = compiled_kernel._run + if getattr(launcher, "debug_launch_hidden_arg", False): + launcher.debug_ctrl_ptr = int(getattr(compiled_kernel, "_debug_ctrl_ptr", 0)) diff --git a/python/triton/knobs.py b/python/triton/knobs.py index bc5752bad..353a6b8cf 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -363,6 +363,9 @@ class compilation_knobs(base_knobs): front_end_debugging: env_bool = env_bool("TRITON_FRONT_END_DEBUGGING") allow_non_constexpr_globals: env_bool = env_bool("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS") enable_experimental_consan: env_bool = env_bool("TRITON_ENABLE_EXPERIMENTAL_CONSAN") + # Backend options include this value so instrumentation mode changes use a + # distinct in-memory and on-disk kernel cache entry. + instrumentation_mode: str = "" listener: Union[CompilationListener, None] = None diff --git a/python/triton/language/__init__.py b/python/triton/language/__init__.py index de88cee95..77dcc83de 100644 --- a/python/triton/language/__init__.py +++ b/python/triton/language/__init__.py @@ -72,6 +72,8 @@ constexpr, constexpr_type, debug_barrier, + debug_collect_end, + debug_collect_start, device_assert, device_print, dot, @@ -191,6 +193,8 @@ "cumprod", "cumsum", "debug_barrier", + "debug_collect_end", + "debug_collect_start", "device_assert", "device_print", "div_rn", diff --git a/python/triton/language/core.py b/python/triton/language/core.py index 5b16c4597..f04b2d7ee 100644 --- a/python/triton/language/core.py +++ b/python/triton/language/core.py @@ -2146,7 +2146,8 @@ def store_tensor_descriptor(desc: tensor_descriptor_base, offsets: Sequence[cons @_tensor_member_fn @builtin -def store(pointer, value, mask=None, boundary_check=(), cache_modifier="", eviction_policy="", _semantic=None): +def store(pointer, value, mask=None, boundary_check=(), cache_modifier="", eviction_policy="", _semantic=None, + _debug_statement_source=None, _debug_statement_id=None): """ Store a tensor of data into memory locations defined by `pointer`. @@ -2192,7 +2193,8 @@ def store(pointer, value, mask=None, boundary_check=(), cache_modifier="", evict mask = _semantic.to_tensor(mask) cache_modifier = _unwrap_if_constexpr(cache_modifier) eviction_policy = _unwrap_if_constexpr(eviction_policy) - return _semantic.store(pointer, value, mask, boundary_check, cache_modifier, eviction_policy) + return _semantic.store(pointer, value, mask, boundary_check, cache_modifier, eviction_policy, + _debug_statement_source=_debug_statement_source, _debug_statement_id=_debug_statement_id) @builtin @@ -2852,6 +2854,27 @@ def debug_barrier(_semantic=None): return _semantic.debug_barrier() +@builtin +def debug_collect_start(level=1, addr_level=None, _semantic=None): + '''Begin a FlagTree debug collect region.''' + level_v = level.value if isinstance(level, constexpr) else level + if not isinstance(level_v, int): + raise TypeError("tl.debug_collect_start: level must be an integer") + + addr_level_v = addr_level.value if isinstance(addr_level, constexpr) else addr_level + if addr_level_v is not None and not isinstance(addr_level_v, int): + raise TypeError("tl.debug_collect_start: addr_level must be an integer") + if addr_level_v is not None and not 0 <= addr_level_v <= 2: + raise ValueError("tl.debug_collect_start: addr_level must be 0, 1, or 2") + return _semantic.debug_collect_start(int(level_v), addr_level_v) + + +@builtin +def debug_collect_end(_semantic=None): + '''End a FlagTree debug collect region.''' + return _semantic.debug_collect_end() + + @builtin def multiple_of(input, values, _semantic=None): """ diff --git a/python/triton/language/semantic.py b/python/triton/language/semantic.py index a0e780ebc..a1fe5073f 100644 --- a/python/triton/language/semantic.py +++ b/python/triton/language/semantic.py @@ -1213,7 +1213,17 @@ def descriptor_scatter(self, desc, value: TensorTy, x_offsets, y_offset) -> Tens self.builder.create_descriptor_scatter(desc.handle, value.handle, x_offsets.handle, y_offset) return self.tensor(None, tl.void) - def _store_block_pointer(self, ptr, val, mask, boundary_check, cache, eviction): + def _annotate_debug_statement(self, op, source, statement_id): + if op is None or not hasattr(op, "set_attr"): + return op + if source: + op.set_attr("flagtree.debug.triton_statement", self.builder.get_string_attr(source)) + if statement_id is not None: + op.set_attr("flagtree.debug.statement_id", self.builder.get_int32_attr(statement_id)) + return op + + def _store_block_pointer(self, ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source=None, + _debug_statement_id=None): # Store by a block pointer: `pointer_type>` # Block pointers can not have the `mask` argument if mask is not None: @@ -1238,10 +1248,11 @@ def _store_block_pointer(self, ptr, val, mask, boundary_check, cache, eviction): val = self.cast(val, elt_ty) # Build IR - return self.tensor( - self.builder.create_tensor_pointer_store(ptr.handle, val.handle, boundary_check, cache, eviction), tl.void) + op = self.builder.create_tensor_pointer_store(ptr.handle, val.handle, boundary_check, cache, eviction) + return self.tensor(self._annotate_debug_statement(op, _debug_statement_source, _debug_statement_id), tl.void) - def _store_legacy(self, ptr, val, mask, boundary_check, cache, eviction): + def _store_legacy(self, ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source=None, + _debug_statement_id=None): # Store by a tensor of pointers or a pointer of scalar: `block_type>` or `pointer_type<>` if not ptr.type.scalar.is_ptr(): raise ValueError(f"Unsupported ptr type {ptr.type.__repr__()} in `tl.store`") @@ -1279,14 +1290,15 @@ def _store_legacy(self, ptr, val, mask, boundary_check, cache, eviction): # Build IR if mask is None: - return self.tensor(self.builder.create_store(ptr.handle, val.handle, cache, eviction), tl.void) - if not mask.type.scalar.is_bool(): - raise ValueError("Mask must have boolean scalar type") - return self.tensor(self.builder.create_masked_store(ptr.handle, val.handle, mask.handle, cache, eviction), - tl.void) + op = self.builder.create_store(ptr.handle, val.handle, cache, eviction) + else: + if not mask.type.scalar.is_bool(): + raise ValueError("Mask must have boolean scalar type") + op = self.builder.create_masked_store(ptr.handle, val.handle, mask.handle, cache, eviction) + return self.tensor(self._annotate_debug_statement(op, _debug_statement_source, _debug_statement_id), tl.void) def store(self, ptr: TensorTy, val: TensorTy, mask: Optional[TensorTy], boundary_check, cache_modifier: str, - eviction_policy: str) -> TensorTy: + eviction_policy: str, _debug_statement_source=None, _debug_statement_id=None) -> TensorTy: # Cache and eviction options cache = self._str_to_store_cache_modifier(cache_modifier) eviction = self._str_to_eviction_policy(eviction_policy) @@ -1296,10 +1308,12 @@ def store(self, ptr: TensorTy, val: TensorTy, mask: Optional[TensorTy], boundary if ptr.type.is_ptr() and ptr.type.element_ty.is_block(): # Store by a block pointer: `pointer_type>` - return self._store_block_pointer(ptr, val, mask, boundary_check, cache, eviction) + return self._store_block_pointer(ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source, + _debug_statement_id) else: # Store by a tensor of pointers or a pointer of scalar: `block_type>` or `pointer_type<>` - return self._store_legacy(ptr, val, mask, boundary_check, cache, eviction) + return self._store_legacy(ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source, + _debug_statement_id) ######### # atomic @@ -1801,6 +1815,19 @@ def max_constancy(self, x: TensorTy, values: List[int]) -> TensorTy: def debug_barrier(self) -> TensorTy: return self.tensor(self.builder.create_barrier(), tl.void) + def debug_collect_start(self, level: int, addr_level: Optional[int]) -> TensorTy: + if not hasattr(self.builder, "create_debug_collect_begin"): + raise RuntimeError("FlagTree debugger support is not available in this build; " + "rebuild with -DFLAGTREE_ENABLE_DEBUGGER=ON") + handle = self.builder.create_debug_collect_begin(level, -1 if addr_level is None else int(addr_level)) + return self.tensor(handle, tl.void) + + def debug_collect_end(self) -> TensorTy: + if not hasattr(self.builder, "create_debug_collect_end"): + raise RuntimeError("FlagTree debugger support is not available in this build; " + "rebuild with -DFLAGTREE_ENABLE_DEBUGGER=ON") + return self.tensor(self.builder.create_debug_collect_end(), tl.void) + def device_print(self, prefix: str, args: List[TensorTy], hex: bool) -> TensorTy: # It makes sense visually for prefix to end in ": "; make it so. Also, # non-empty prefixes should start with " ". diff --git a/python/triton/runtime/debug_collect_runtime.py b/python/triton/runtime/debug_collect_runtime.py new file mode 100644 index 000000000..050743f7f --- /dev/null +++ b/python/triton/runtime/debug_collect_runtime.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +from dataclasses import dataclass +import importlib +import os +from typing import Any + + +@dataclass +class DebugLaunchContext: + debug_kernel_id: int + hidden_arg: int + handle: Any | None = None + metadata: Any | None = None + + +@dataclass +class DebugExportedRun: + debug_kernel_id: int + raw_buffer: bytes + meta: dict[str, Any] | None = None + runtime_metadata: dict[str, Any] | None = None + decoded: dict[str, Any] | None = None + report: str = "" + + +class DebugCollectRuntime: + """Debugger runtime bridge for post-kernel export and host decode.""" + + def __init__(self) -> None: + self._exported_runs: list[DebugExportedRun] = [] + + @staticmethod + def _binding(): + return importlib.import_module("triton._C.libtriton").debugger + + @staticmethod + def _metadata_to_dict(metadata: Any) -> dict[str, Any]: + if metadata is None: + return {} + if isinstance(metadata, dict): + result = dict(metadata) + else: + result = {} + asdict = getattr(metadata, "_asdict", None) + if callable(asdict): + result.update(dict(asdict())) + elif hasattr(metadata, "__dict__"): + result.update(dict(vars(metadata))) + + for key in ( + "debug_enabled", + "debug_kernel_id", + "debug_protocol_version", + "debug_record_level", + "debug_addr_level", + "debug_export_mode", + "debug_record_capacity", + "debug_record_size", + "debug_records_per_instance", + "debug_record_layout", + "debug_record_plan", + "debug_full_dump_payload_bytes_per_instance", + "debug_full_dump_plan", + "debug_backend_name", + "debug_target_name", + "debug_metadata_json", + "name", + "backend_name", + "target_name", + "target", + ): + try: + value = getattr(metadata, key) + except AttributeError: + continue + if not key.startswith("debug_") and key in result: + continue + result[key] = value + return result + + @staticmethod + def _target_to_name(target: Any) -> str: + if isinstance(target, dict): + arch = target.get("arch") + if arch is not None: + return str(arch) + backend = target.get("backend") + if backend is not None: + return str(backend) + return "" + arch = getattr(target, "arch", None) + if arch is not None: + return str(arch) + backend = getattr(target, "backend", None) + if backend is not None: + return str(backend) + return "" + + @staticmethod + def _target_backend(target: Any) -> str: + if isinstance(target, dict): + backend = target.get("backend") + return "" if backend is None else str(backend) + backend = getattr(target, "backend", None) + return "" if backend is None else str(backend) + + @staticmethod + def _normalize_backend_name(backend_name: Any) -> str: + name = "" if backend_name is None else str(backend_name) + return name + + @staticmethod + def _normalize_launch_metadata(metadata_dict: dict[str, Any]) -> dict[str, Any]: + target = metadata_dict.get("target") + target_name = (metadata_dict.get("debug_target_name") or metadata_dict.get("target_name") + or DebugCollectRuntime._target_to_name(target)) + backend_name = (metadata_dict.get("debug_backend_name") or os.environ.get("FLAGTREE_BACKEND") + or metadata_dict.get("backend_name") or DebugCollectRuntime._target_backend(target)) + metadata_dict["debug_backend_name"] = DebugCollectRuntime._normalize_backend_name(backend_name) + metadata_dict["debug_target_name"] = str(target_name) + return metadata_dict + + def prepare(self, metadata: Any, stream, runtime_metadata: Any | None = None) -> DebugLaunchContext: + metadata_dict = self._normalize_launch_metadata(self._metadata_to_dict(metadata)) + runtime_metadata_dict = dict(runtime_metadata or {}) + if metadata_dict.get("debug_record_layout"): + runtime_metadata_dict.setdefault("record_layout", metadata_dict["debug_record_layout"]) + if metadata_dict.get("debug_record_plan") is not None: + runtime_metadata_dict.setdefault("record_plan", metadata_dict["debug_record_plan"]) + if metadata_dict.get("debug_records_per_instance") is not None: + runtime_metadata_dict.setdefault("records_per_instance", int(metadata_dict["debug_records_per_instance"])) + if metadata_dict.get("debug_full_dump_plan") is not None: + runtime_metadata_dict.setdefault("full_dump_plan", metadata_dict["debug_full_dump_plan"]) + if (int(metadata_dict.get("debug_record_level", 1)) == 2 + and int(metadata_dict.get("debug_full_dump_payload_bytes_per_instance", 0)) > 0 + and metadata_dict.get("debug_full_dump_plan")): + from triton.runtime import debugger as process_debugger + + if process_debugger.get_output_dir() is None: + raise RuntimeError("level-2 debugger full dump requires debugger output_dir") + handle = self._binding().prepare_launch(metadata_dict, int(stream or 0), runtime_metadata_dict) + return DebugLaunchContext( + debug_kernel_id=int(metadata_dict.get("debug_kernel_id", 0)), + hidden_arg=int(handle.hidden_arg_value), + handle=handle, + metadata=metadata, + ) + + def hidden_arg(self, ctx: DebugLaunchContext) -> int: + return int(ctx.hidden_arg) + + def export(self, ctx: DebugLaunchContext, stream) -> DebugExportedRun: + del stream + if ctx.handle is None: + return DebugExportedRun(debug_kernel_id=ctx.debug_kernel_id, raw_buffer=b"") + + exported = ctx.handle.finish() + ctx.handle = None + metadata_dict = self._normalize_launch_metadata(self._metadata_to_dict(ctx.metadata)) + if (int(metadata_dict.get("debug_record_level", 1)) == 2 + and int(metadata_dict.get("debug_full_dump_payload_bytes_per_instance", 0)) > 0 + and metadata_dict.get("debug_full_dump_plan")): + from triton.runtime import debugger as process_debugger + + finalized = process_debugger._finalize_exported_run( # noqa: SLF001 + dict(exported), metadata_dict) + run = DebugExportedRun( + debug_kernel_id=ctx.debug_kernel_id, + raw_buffer=bytes(finalized.get("raw_buffer", b"")), + meta=dict(finalized.get("meta", {})), + runtime_metadata=dict(finalized.get("runtime_metadata", {})), + decoded=dict(finalized.get("decoded", {})), + report=str(finalized.get("report", "")), + ) + self._exported_runs.append(run) + return run + + raw_buffer = bytes(exported.get("raw_buffer", b"")) + decoded = self._binding().decode_exported_run(exported) + + report = "" + metadata_json = getattr(ctx.metadata, "debug_metadata_json", None) + if metadata_json: + report = self._binding().render_text_report(exported, str(metadata_json)) + + run = DebugExportedRun( + debug_kernel_id=ctx.debug_kernel_id, + raw_buffer=raw_buffer, + meta=dict(exported.get("meta", {})), + runtime_metadata=dict(exported.get("runtime_metadata", {})), + decoded=dict(decoded), + report=report, + ) + self._exported_runs.append(run) + return run + + def release(self, ctx: DebugLaunchContext) -> None: + if ctx.handle is not None: + ctx.handle.release() + ctx.handle = None + + def clear_exported_runs(self) -> None: + self._exported_runs.clear() + + def peek_exported_runs(self) -> list[DebugExportedRun]: + return list(self._exported_runs) + + def take_exported_runs(self) -> list[DebugExportedRun]: + runs = list(self._exported_runs) + self._exported_runs.clear() + return runs + + +default_debug_collect_runtime = DebugCollectRuntime() diff --git a/python/triton/runtime/debugger.py b/python/triton/runtime/debugger.py new file mode 100644 index 000000000..26bbdbc02 --- /dev/null +++ b/python/triton/runtime/debugger.py @@ -0,0 +1,892 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +import hashlib +import importlib +import inspect +import json +import os +from pathlib import Path +import pprint +import re +import sys +from typing import Any, Callable, Mapping, Optional, Sequence + + +@dataclass(frozen=True) +class PreparedKernelLaunch: + kernel_args: tuple[int, ...] = () + finalize: Optional[Callable[[Optional[BaseException]], None]] = None + + +@dataclass(frozen=True) +class DebuggerConfig: + enabled: bool = False + record_level: int = 1 + addr_level: int = 0 + export_mode: str = "POST_KERNEL_EXPORT" + record_capacity: int = 1024 + export_on_error: bool = False + output_dir: str | None = "/tmp/flagtree_debugger_manual" + export_raw_records: bool = False + runtime_metadata_builder: Optional[Callable[[Any, Any, Sequence[Any]], Any]] = None + export_handler: Optional[Callable[[dict[str, Any]], None]] = None + + +_DEFAULT_OUTPUT_DIR = Path("/tmp/flagtree_debugger_manual") +_DEFAULT_RECORD_CAPACITY = 1024 +_DEFAULT_ADDR_LEVEL = 0 +_DEFAULT_EXPORT_MODE = "POST_KERNEL_EXPORT" +_DEFAULT_EXPORT_ON_ERROR = False +_USE_CURRENT_CONFIG = object() +_USE_CURRENT_OUTPUT_DIR = _USE_CURRENT_CONFIG +_launch_prepare_hook = None +_output_dir: Path | None = _DEFAULT_OUTPUT_DIR +_record_capacity = _DEFAULT_RECORD_CAPACITY +_export_mode = _DEFAULT_EXPORT_MODE +_export_on_error = _DEFAULT_EXPORT_ON_ERROR +_raw_record_export_enabled = False +_active_config = DebuggerConfig() +_exported_runs: list[dict[str, Any]] = [] +_CONFIG_KEYS = frozenset({ + "output_dir", + "record_capacity", + "export_mode", + "export_on_error", + "export_raw_records", +}) +_DISABLED_BUILD_MESSAGE = ("FlagTree debugger support is not available in this build; rebuild with " + "-DFLAGTREE_ENABLE_DEBUGGER=ON") + + +def _normalize_kernel_args(kernel_args: Any) -> tuple[int, ...]: + if kernel_args is None: + return () + if isinstance(kernel_args, int): + return (int(kernel_args), ) + if isinstance(kernel_args, Sequence): + return tuple(int(arg) for arg in kernel_args) + raise TypeError("debugger launch hook must return an int, a sequence of ints, or " + "PreparedKernelLaunch") + + +def _wrap_launch_prepare_hook(hook: Callable[..., Any]) -> Callable[..., Any]: + signature = inspect.signature(hook) + positional = [ + parameter for parameter in signature.parameters.values() if parameter.kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ) + ] + accepts_varargs = any(parameter.kind == inspect.Parameter.VAR_POSITIONAL + for parameter in signature.parameters.values()) + if accepts_varargs or len(positional) >= 4: + return hook + if len(positional) == 3: + return lambda metadata, stream, launch_metadata, kernel_args: hook(metadata, stream, launch_metadata) + raise TypeError("debugger launch hook must accept " + "(metadata, stream, launch_metadata[, kernel_args])") + + +def _load_binding(): + return importlib.import_module("triton._C.libtriton").debugger + + +def is_available() -> bool: + try: + libtriton = importlib.import_module("triton._C.libtriton") + except (ImportError, OSError): + return False + native_debugger = getattr(libtriton, "debugger", None) + native_passes = getattr(getattr(libtriton, "passes", None), "flagtree_debug", None) + return native_debugger is not None and native_passes is not None + + +def _require_available() -> None: + if not is_available(): + raise RuntimeError(_DISABLED_BUILD_MESSAGE) + + +def _normalize_output_dir(path: Any) -> Path | None: + if path is None: + return None + return Path(os.fspath(path)).expanduser() + + +def configure(config: Mapping[str, Any] | None = None, **kwargs: Any) -> None: + """Update debugger defaults used by the next ``enable_debug(level=...)``. + + Supported keys are ``output_dir``, ``record_capacity``, ``export_mode``, + ``export_on_error``, and ``export_raw_records``. Keys not provided keep + their current values. + """ + global _output_dir, _record_capacity, _export_mode + global _export_on_error, _raw_record_export_enabled + + updates = {} + if config is not None: + updates.update(dict(config)) + updates.update(kwargs) + + unknown = sorted(set(updates) - _CONFIG_KEYS) + if unknown: + raise TypeError(f"unknown debugger config key(s): {', '.join(unknown)}") + + if "output_dir" in updates: + _output_dir = _normalize_output_dir(updates["output_dir"]) + if "record_capacity" in updates: + capacity = int(updates["record_capacity"]) + if capacity <= 0: + raise ValueError("debugger record capacity must be positive") + _record_capacity = capacity + if "export_mode" in updates: + _export_mode = _normalize_export_mode(updates["export_mode"]) + if "export_on_error" in updates: + _export_on_error = bool(updates["export_on_error"]) + if "export_raw_records" in updates: + _raw_record_export_enabled = bool(updates["export_raw_records"]) + + +def reset_config() -> None: + """Restore debugger defaults used by ``enable_debug(level=...)``.""" + configure( + output_dir=_DEFAULT_OUTPUT_DIR, + record_capacity=_DEFAULT_RECORD_CAPACITY, + export_mode=_DEFAULT_EXPORT_MODE, + export_on_error=_DEFAULT_EXPORT_ON_ERROR, + export_raw_records=False, + ) + + +def get_config() -> dict[str, Any]: + """Return the current debugger defaults.""" + return { + "output_dir": get_output_dir(), + "record_capacity": int(_record_capacity), + "export_mode": str(_export_mode), + "export_on_error": bool(_export_on_error), + "export_raw_records": bool(_raw_record_export_enabled), + } + + +def set_output_dir(path: str | os.PathLike | None) -> None: + """Compatibility wrapper for ``configure(output_dir=...)``.""" + configure(output_dir=path) + + +def get_output_dir() -> str | None: + """Return the current automatic report export directory, or ``None``.""" + if _output_dir is None: + return None + return str(_output_dir) + + +def _normalize_export_mode(export_mode: str | int) -> str: + if isinstance(export_mode, int): + return "STREAMING_EXPORT" if export_mode == 2 else "POST_KERNEL_EXPORT" + normalized = str(export_mode).strip().replace("-", "_").upper() + if normalized in {"STREAMING", "STREAMING_EXPORT"}: + return "STREAMING_EXPORT" + return "POST_KERNEL_EXPORT" + + +def _normalize_addr_level(addr_level: int) -> int: + value = int(addr_level) + if value < 0 or value > 2: + raise ValueError("debugger addr_level must be 0, 1, or 2") + return value + + +def _derive_kernel_id(metadata_dict: dict[str, Any]) -> int: + kernel_hash = metadata_dict.get("hash") + if isinstance(kernel_hash, str) and kernel_hash: + kernel_id = int(kernel_hash[:8], 16) + return kernel_id or 1 + + digest = hashlib.sha256(repr(sorted(metadata_dict.items())).encode("utf-8")).hexdigest() + kernel_id = int(digest[:8], 16) + return kernel_id or 1 + + +def _target_to_name(target: Any) -> str: + if isinstance(target, dict): + arch = target.get("arch") + if arch is not None: + return str(arch) + backend = target.get("backend") + if backend is not None: + return str(backend) + return "" + arch = getattr(target, "arch", None) + if arch is not None: + return str(arch) + backend = getattr(target, "backend", None) + if backend is not None: + return str(backend) + return "" + + +def _target_backend(target: Any) -> str: + if isinstance(target, dict): + backend = target.get("backend") + return "" if backend is None else str(backend) + backend = getattr(target, "backend", None) + return "" if backend is None else str(backend) + + +def _normalize_backend_name(backend_name: Any) -> str: + name = "" if backend_name is None else str(backend_name) + return name + + +def _safe_filename_component(value: Any, fallback: str) -> str: + text = "" if value is None else str(value) + text = text.strip() or fallback + safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", text) + safe = safe.strip("._-") + return safe or fallback + + +def _current_script_stem() -> str: + script = sys.argv[0] if sys.argv else "" + if not script: + return "interactive" + stem = Path(script).stem + return _safe_filename_component(stem, "interactive") + + +def _exported_run_meta(exported_run: dict[str, Any]) -> dict[str, Any]: + meta = exported_run.get("meta") + return meta if isinstance(meta, dict) else {} + + +def _render_export_summary(exported_run: dict[str, Any], decoded: dict[str, Any] | None, + metadata_dict: dict[str, Any]) -> str: + meta = _exported_run_meta(exported_run) + raw_buffer = exported_run.get("raw_buffer", b"") + raw_size = len(raw_buffer) if hasattr(raw_buffer, "__len__") else 0 + lines = [ + "FlagTree Debug Export", + f"kernel_name: {metadata_dict.get('debug_kernel_name') or metadata_dict.get('name') or ''}", + f"kernel_id: {meta.get('kernel_id', metadata_dict.get('debug_kernel_id', 0))}", + f"run_id: {meta.get('run_id', '')}", + f"backend: {metadata_dict.get('debug_backend_name', '')}", + f"target: {metadata_dict.get('debug_target_name', '')}", + f"raw_buffer_bytes: {raw_size}", + ] + if decoded is not None: + lines.extend([ + "", + "Decoded Header", + pprint.pformat(decoded.get("header", {}), sort_dicts=True), + ]) + return "\n".join(lines) + + +def _build_report_path(output_dir: Path, exported_run: dict[str, Any], metadata_dict: dict[str, Any]) -> Path: + meta = _exported_run_meta(exported_run) + script_stem = _current_script_stem() + kernel_name = _safe_filename_component( + metadata_dict.get("debug_kernel_name") or metadata_dict.get("name"), + "kernel", + ) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3] + run_id = _safe_filename_component(meta.get("run_id", "0"), "0") + return output_dir / f"{script_stem}_{kernel_name}_{timestamp}_run{run_id}.txt" + + +def _render_raw_records(exported_run: dict[str, Any], decoded: dict[str, Any], metadata_dict: dict[str, Any]) -> str: + meta = _exported_run_meta(exported_run) + lines = [ + "FlagTree Debug Raw Records", + f"kernel_name: {metadata_dict.get('debug_kernel_name') or metadata_dict.get('name') or ''}", + f"kernel_id: {meta.get('kernel_id', metadata_dict.get('debug_kernel_id', 0))}", + f"run_id: {meta.get('run_id', '')}", + "", + "Decoded Header", + pprint.pformat(decoded.get("header", {}), sort_dicts=True), + "", + "Decoded Records", + pprint.pformat(decoded.get("records", []), sort_dicts=True), + ] + return "\n".join(lines) + + +def _is_full_dump_run(metadata_dict: dict[str, Any]) -> bool: + return (int(metadata_dict.get("debug_record_level", 1)) == 2 + and int(metadata_dict.get("debug_full_dump_payload_bytes_per_instance", 0)) > 0 + and bool(metadata_dict.get("debug_full_dump_plan"))) + + +def _record_level_id(value: Any) -> int: + if isinstance(value, int): + return 2 if value == 2 else 1 + text = str(value) + return 2 if text in {"2", "LEVEL_TENSOR_FULL"} else 1 + + +def _export_mode_id(value: Any) -> int: + if isinstance(value, int): + return 2 if value == 2 else 1 + text = str(value) + return 2 if text == "STREAMING_EXPORT" else 1 + + +def _empty_debug_raw_buffer(metadata_dict: Mapping[str, Any]) -> bytes: + record_size = int(metadata_dict.get("debug_record_size", 64) or 64) + fields = ( + 0, # writeIdx + 0, # capacity + 0, # overflowCount + 0, # flags + record_size, + 32, # payloadOffset + 0, + 0, + ) + return b"".join(int(field).to_bytes(4, "little", signed=False) for field in fields) + + +def _metadata_only_exported_run(metadata_dict: dict[str, Any], runtime_metadata: Mapping[str, Any]) -> dict[str, Any]: + return { + "meta": { + "run_id": len(_exported_runs) + 1, + "device_id": int(metadata_dict.get("debug_device_id", 0) or 0), + "kernel_id": int(metadata_dict.get("debug_kernel_id", _derive_kernel_id(metadata_dict))), + "protocol_version": int(metadata_dict.get("debug_protocol_version", 2) or 2), + "record_level": _record_level_id(metadata_dict.get("debug_record_level", 1)), + "export_mode": _export_mode_id(metadata_dict.get("debug_export_mode", _active_config.export_mode)), + "backend_kind": 0, + }, + "runtime_metadata": dict(runtime_metadata), + "raw_buffer": _empty_debug_raw_buffer(metadata_dict), + } + + +def _npy_dtype_descriptor(dtype: str) -> str: + if dtype == "float32": + return " int: + if dtype in {"float32"}: + return 4 + if dtype in {"float64", "int64", "uint64"}: + return 8 + raise ValueError(f"unsupported debugger artifact dtype: {dtype}") + + +def _write_npy(path: Path, payload: bytes, dtype: str, shape: Sequence[int]) -> None: + descr = _npy_dtype_descriptor(dtype) + dims = tuple(int(dim) for dim in shape) + if any(dim < 0 for dim in dims): + raise ValueError("debugger artifact shape cannot contain negative dimensions") + element_count = 1 + for dim in dims: + element_count *= dim + expected_bytes = element_count * _npy_dtype_element_bytes(dtype) + if len(payload) != expected_bytes: + raise RuntimeError("debugger artifact payload size does not match planned dtype/shape") + if len(dims) == 1: + shape_repr = f"({dims[0]},)" + else: + shape_repr = "(" + ", ".join(str(dim) for dim in dims) + ")" + header = f"{{'descr': '{descr}', 'fortran_order': False, 'shape': {shape_repr}, }}" + header_bytes = header.encode("latin1") + padding = (16 - ((10 + len(header_bytes) + 1) % 16)) % 16 + header_bytes = header_bytes + b" " * padding + b"\n" + if len(header_bytes) > 0xFFFF: + raise ValueError("debugger artifact .npy header is too large") + path.write_bytes(b"\x93NUMPY\x01\x00" + len(header_bytes).to_bytes(2, "little") + header_bytes + payload) + + +def _full_dump_plan_by_record(metadata_dict: dict[str, Any]) -> dict[int, dict[str, Any]]: + plan = metadata_dict.get("debug_full_dump_plan") or [] + result = {} + for entry in plan: + if not isinstance(entry, Mapping): + continue + result[int(entry.get("record_index", 0))] = dict(entry) + return result + + +def _record_index_for_slot(slot_index: int, runtime_metadata: Mapping[str, Any]) -> int: + records_per_instance = int(runtime_metadata.get("records_per_instance") or 0) + if records_per_instance <= 0: + return slot_index + return slot_index % records_per_instance + + +def _write_full_dump_artifacts(report_path: Path, exported_run: dict[str, Any], decoded: dict[str, Any], + metadata_dict: dict[str, Any]) -> list[dict[str, Any]]: + runtime_metadata = dict(exported_run.get("runtime_metadata") or {}) + plan_by_record = _full_dump_plan_by_record(metadata_dict) + raw_buffer = bytes(exported_run.get("raw_buffer", b"")) + header = decoded.get("header", {}) + if int(header.get("overflow_count", 0)) != 0 or int(header.get("flags", 0)) & 1: + raise RuntimeError("level-2 debugger full dump cannot export from an overflowed debug buffer") + + artifact_dir = report_path.with_suffix("") + artifact_dir = artifact_dir.with_name(f"{artifact_dir.name}_artifacts") + artifact_dir.mkdir(parents=True, exist_ok=True) + + artifacts: list[dict[str, Any]] = [] + records = decoded.get("records", []) + for slot_index, record in enumerate(records): + if not isinstance(record, Mapping) or record.get("record_kind") != "FULL_VALUE": + continue + record_index = _record_index_for_slot(slot_index, runtime_metadata) + plan = plan_by_record.get(record_index) + if plan is None: + raise RuntimeError(f"missing full-dump plan for record_index={record_index}") + payload_offset = int(record.get("payload_offset", 0)) + payload_length = int(record.get("payload_length", 0)) + if payload_length <= 0: + raise RuntimeError(f"empty full-dump payload for record_index={record_index}") + if payload_offset < 0 or payload_offset + payload_length > len(raw_buffer): + raise RuntimeError(f"full-dump payload range is outside raw buffer for record_index={record_index}") + payload = raw_buffer[payload_offset:payload_offset + payload_length] + dtype = str(plan.get("artifact_dtype", "")) + shape = plan.get("shape") or [int(plan.get("element_count", 0))] + kind = str(plan.get("kind", "value")) + op_id = int(record.get("op_id", plan.get("op_id", 0))) + instance_id = int(record.get("logical_instance_id", 0)) + stem = (f"op{op_id}_inst{instance_id}_rec{record_index}_" + f"{_safe_filename_component(kind, 'dump')}.npy") + artifact_path = artifact_dir / stem + _write_npy(artifact_path, payload, dtype, shape) + artifacts.append({ + "op_id": op_id, + "logical_instance_id": instance_id, + "record_index": record_index, + "kind": kind, + "source": plan.get("source", ""), + "artifact_dtype": dtype, + "shape": list(shape), + "payload_offset": payload_offset, + "payload_length": payload_length, + "path": str(artifact_path), + }) + + expected_records = len(plan_by_record) + if expected_records and not artifacts: + raise RuntimeError("level-2 debugger did not produce any full-dump artifacts") + + index_path = artifact_dir / "tensor_index.json" + index = { + "kernel_name": metadata_dict.get("debug_kernel_name") or metadata_dict.get("name") or "", + "kernel_id": metadata_dict.get("debug_kernel_id", 0), + "run_id": _exported_run_meta(exported_run).get("run_id", 0), + "artifacts": artifacts, + } + index_path.write_text(json.dumps(index, indent=2, sort_keys=True)) + runtime_metadata["full_dump_artifacts"] = artifacts + exported_run["runtime_metadata"] = runtime_metadata + exported_run["full_dump_artifact_dir"] = str(artifact_dir) + exported_run["full_dump_index_path"] = str(index_path) + return artifacts + + +def _finalize_exported_run(exported_run: dict[str, Any], metadata_dict: dict[str, Any]) -> dict[str, Any]: + exported_run["debug_kernel_name"] = str(metadata_dict.get("debug_kernel_name") or metadata_dict.get("name") or "") + tracked_table = metadata_dict.get("debug_tracked_table") + if isinstance(tracked_table, Sequence) and not isinstance(tracked_table, (str, bytes, bytearray)): + exported_run["debug_tracked_table"] = list(tracked_table) + + binding = _load_binding() + decoded = binding.decode_exported_run(exported_run) + exported_run["decoded"] = decoded + + output_dir = _output_dir + report_path = None + if output_dir is not None: + output_dir.mkdir(parents=True, exist_ok=True) + report_path = _build_report_path(output_dir, exported_run, metadata_dict) + + if _is_full_dump_run(metadata_dict): + if report_path is None: + raise RuntimeError("level-2 debugger full dump requires debugger output_dir") + _write_full_dump_artifacts(report_path, exported_run, decoded, metadata_dict) + + summary = _render_export_summary(exported_run, decoded, metadata_dict) + report = "" + json_report = "" + op_log_report = "" + op_log_json_report = "" + metadata_json = metadata_dict.get("debug_metadata_json") + if metadata_json: + render_text_statement_report = getattr(binding, "render_text_statement_report", None) + if callable(render_text_statement_report): + report = render_text_statement_report(exported_run, str(metadata_json)) + else: + report = binding.render_text_report(exported_run, str(metadata_json)) + exported_run["report"] = report + + render_text_op_log_report = getattr(binding, "render_text_op_log_report", None) + if callable(render_text_op_log_report): + op_log_report = render_text_op_log_report(exported_run, str(metadata_json)) + exported_run["op_log_report"] = op_log_report + + render_json_statement_report = getattr(binding, "render_json_statement_report", None) + render_json_report = (render_json_statement_report if callable(render_json_statement_report) else getattr( + binding, "render_json_report", None)) + if callable(render_json_report): + json_report = render_json_report(exported_run, str(metadata_json)) + exported_run["json_report"] = json_report + + render_json_op_log_report = getattr(binding, "render_json_op_log_report", None) + if callable(render_json_op_log_report): + op_log_json_report = render_json_op_log_report(exported_run, str(metadata_json)) + exported_run["op_log_json_report"] = op_log_json_report + + report_text = summary + if report: + report_text += "\n\n" + report_text += report.lstrip("\n") + + op_log_report_text = summary + if op_log_report: + op_log_report_text += "\n\n" + op_log_report_text += op_log_report.lstrip("\n") + + if report_path is not None: + report_path.write_text(report_text) + exported_run["report_path"] = str(report_path) + if op_log_report: + op_log_report_path = report_path.with_name(f"{report_path.stem}_op_log.txt") + op_log_report_path.write_text(op_log_report_text) + exported_run["op_log_report_path"] = str(op_log_report_path) + if json_report: + json_report_path = report_path.with_suffix(".json") + json_report_path.write_text(json_report) + exported_run["json_report_path"] = str(json_report_path) + if op_log_json_report: + op_log_json_report_path = report_path.with_name(f"{report_path.stem}_op_log.json") + op_log_json_report_path.write_text(op_log_json_report) + exported_run["op_log_json_report_path"] = str(op_log_json_report_path) + if _active_config.export_raw_records: + raw_records_path = report_path.with_name(f"{report_path.stem}_raw_records.txt") + raw_records_path.write_text(_render_raw_records(exported_run, decoded, metadata_dict)) + exported_run["raw_records_path"] = str(raw_records_path) + + return exported_run + + +def _normalize_device_id(device: Any) -> int: + if isinstance(device, int): + return device + index = getattr(device, "index", None) + if isinstance(index, int): + return index + return 0 + + +def _metadata_to_dict(metadata: Any) -> dict[str, Any]: + if metadata is None: + return {} + if isinstance(metadata, dict): + return dict(metadata) + asdict = getattr(metadata, "_asdict", None) + if callable(asdict): + return dict(asdict()) + if hasattr(metadata, "__dict__"): + return dict(vars(metadata)) + return {} + + +def _materialize_launch_metadata(launch_metadata: Any) -> Any: + if launch_metadata is None or isinstance(launch_metadata, dict): + return launch_metadata + getter = getattr(launch_metadata, "get", None) + if callable(getter): + try: + return getter() + except TypeError: + return launch_metadata + return launch_metadata + + +def _normalize_launch_grid(launch_metadata: Any) -> tuple[int, int, int] | None: + if not isinstance(launch_metadata, Mapping): + return None + grid = launch_metadata.get("grid") + if grid is None: + return None + values = tuple(int(dim) for dim in grid) + if not values: + return None + return ( + values[0], + values[1] if len(values) > 1 else 1, + values[2] if len(values) > 2 else 1, + ) + + +def _build_launch_metadata_dict(metadata: Any) -> dict[str, Any]: + metadata_dict = _metadata_to_dict(metadata) + target = metadata_dict.get("target") + target_name = _target_to_name(target) + backend_name = (metadata_dict.get("debug_backend_name") or os.environ.get("FLAGTREE_BACKEND") + or metadata_dict.get("backend_name") or _target_backend(target) or "") + backend_name = _normalize_backend_name(backend_name) + + launch_dict = dict(metadata_dict) + launch_dict["debug_enabled"] = True + launch_dict["debug_protocol_version"] = int(metadata_dict.get("debug_protocol_version", 2)) + launch_dict["debug_record_level"] = int(metadata_dict.get("debug_record_level", _active_config.record_level)) + launch_dict["debug_addr_level"] = _normalize_addr_level( + metadata_dict.get("debug_addr_level", _active_config.addr_level)) + launch_dict["debug_export_mode"] = _normalize_export_mode( + metadata_dict.get("debug_export_mode", _active_config.export_mode)) + launch_dict["debug_record_capacity"] = int( + metadata_dict.get("debug_record_capacity", _active_config.record_capacity)) + launch_dict["debug_record_size"] = int(metadata_dict.get("debug_record_size", 32)) + launch_dict["debug_kernel_id"] = int(metadata_dict.get("debug_kernel_id", _derive_kernel_id(metadata_dict))) + launch_dict["debug_kernel_name"] = str(metadata_dict.get("debug_kernel_name", metadata_dict.get("name", ""))) + launch_dict["debug_backend_name"] = str(backend_name) + launch_dict["debug_target_name"] = str(metadata_dict.get("debug_target_name", target_name)) + try: + from triton.runtime.driver import driver + + launch_dict["debug_device_id"] = _normalize_device_id(driver.active.get_current_device()) + except Exception: + launch_dict["debug_device_id"] = 0 + return launch_dict + + +def _default_launch_prepare_hook(metadata: Any, stream: int, launch_metadata: Any, + kernel_args: Sequence[Any]) -> PreparedKernelLaunch: + launch_metadata = _materialize_launch_metadata(launch_metadata) + metadata_dict = _build_launch_metadata_dict(metadata) + runtime_metadata = {} + if _active_config.runtime_metadata_builder is not None: + built_runtime_metadata = _active_config.runtime_metadata_builder(metadata, launch_metadata, kernel_args) + if built_runtime_metadata is not None: + runtime_metadata = dict(built_runtime_metadata) + grid = _normalize_launch_grid(launch_metadata) + if grid is not None: + runtime_metadata.setdefault("grid", grid) + records_per_instance = metadata_dict.get("debug_records_per_instance") + if records_per_instance is not None: + runtime_metadata.setdefault("records_per_instance", int(records_per_instance)) + if metadata_dict.get("debug_record_layout"): + runtime_metadata.setdefault("record_layout", metadata_dict["debug_record_layout"]) + if metadata_dict.get("debug_record_plan") is not None: + runtime_metadata.setdefault("record_plan", metadata_dict["debug_record_plan"]) + if _is_full_dump_run(metadata_dict): + if _output_dir is None: + raise RuntimeError("level-2 debugger full dump requires debugger output_dir") + runtime_metadata.setdefault("full_dump_plan", metadata_dict.get("debug_full_dump_plan", [])) + + handle = _load_binding().prepare_launch(metadata_dict, int(stream), runtime_metadata) + + def finalize(error: Optional[BaseException]) -> None: + if error is not None and not _active_config.export_on_error: + handle.release() + return + + exported_run = handle.finish() + exported_run = _finalize_exported_run(exported_run, metadata_dict) + _exported_runs.append(exported_run) + if _active_config.export_handler is not None: + _active_config.export_handler(exported_run) + + return PreparedKernelLaunch( + kernel_args=(int(handle.hidden_arg_value), ), + finalize=finalize, + ) + + +def prepare_metadata_only_kernel_launch( + metadata: Any, + stream: int, + launch_metadata: Any = None, + kernel_args: Optional[Sequence[Any]] = None, +) -> Optional[PreparedKernelLaunch]: + del stream + enabled = bool(getattr(metadata, "debug_enabled", False)) or _active_config.enabled + if not enabled: + return None + + launch_metadata = _materialize_launch_metadata(launch_metadata) + metadata_dict = _build_launch_metadata_dict(metadata) + runtime_metadata: dict[str, Any] = {} + if _active_config.runtime_metadata_builder is not None: + built_runtime_metadata = _active_config.runtime_metadata_builder(metadata, launch_metadata, + tuple(kernel_args or ())) + if built_runtime_metadata is not None: + runtime_metadata = dict(built_runtime_metadata) + grid = _normalize_launch_grid(launch_metadata) + if grid is not None: + runtime_metadata.setdefault("grid", grid) + runtime_metadata.setdefault( + "records_per_instance", + int(metadata_dict.get("debug_records_per_instance", 0) or 0), + ) + if metadata_dict.get("debug_record_layout"): + runtime_metadata.setdefault("record_layout", metadata_dict["debug_record_layout"]) + if metadata_dict.get("debug_record_plan") is not None: + runtime_metadata.setdefault("record_plan", metadata_dict["debug_record_plan"]) + + def finalize(error: Optional[BaseException]) -> None: + if error is not None and not _active_config.export_on_error: + return + + exported_run = _metadata_only_exported_run(metadata_dict, runtime_metadata) + exported_run = _finalize_exported_run(exported_run, metadata_dict) + _exported_runs.append(exported_run) + if _active_config.export_handler is not None: + _active_config.export_handler(exported_run) + + return PreparedKernelLaunch(kernel_args=(), finalize=finalize) + + +def register_launch_prepare_hook(hook: Callable[..., Any]) -> None: + global _launch_prepare_hook + _launch_prepare_hook = _wrap_launch_prepare_hook(hook) + + +def clear_launch_prepare_hook() -> None: + global _launch_prepare_hook + _launch_prepare_hook = None + + +def is_active() -> bool: + return _active_config.enabled + + +def current_compile_config() -> dict[str, Any]: + if not _active_config.enabled: + return {} + return { + "debug_enabled": True, + "debug_protocol_version": 2, + "debug_record_level": int(_active_config.record_level), + "debug_addr_level": int(_active_config.addr_level), + "debug_export_mode": _normalize_export_mode(_active_config.export_mode), + "debug_record_capacity": int(_active_config.record_capacity), + } + + +def activate( + *, + level: int | None = None, + addr_level: int = _DEFAULT_ADDR_LEVEL, + record_level: int | None = None, + export_mode: Any = _USE_CURRENT_CONFIG, + record_capacity: Any = _USE_CURRENT_CONFIG, + export_on_error: Any = _USE_CURRENT_CONFIG, + output_dir: Any = _USE_CURRENT_OUTPUT_DIR, + export_raw_records: Any = _USE_CURRENT_CONFIG, + runtime_metadata_builder: Optional[Callable[[Any, Any, Sequence[Any]], Any]] = None, + export_handler: Optional[Callable[[dict[str, Any]], None]] = None, +) -> None: + global _active_config + _require_available() + if output_dir is not _USE_CURRENT_OUTPUT_DIR: + configure(output_dir=output_dir) + + if level is not None and record_level is not None: + raise TypeError("use either level or record_level, not both") + effective_level = record_level if record_level is not None else level + if effective_level is None: + effective_level = 1 + effective_addr_level = _normalize_addr_level(addr_level) + effective_export_mode = (_export_mode + if export_mode is _USE_CURRENT_CONFIG else _normalize_export_mode(export_mode)) + effective_record_capacity = (_record_capacity if record_capacity is _USE_CURRENT_CONFIG else int(record_capacity)) + if effective_record_capacity <= 0: + raise ValueError("debugger record capacity must be positive") + effective_export_on_error = (_export_on_error if export_on_error is _USE_CURRENT_CONFIG else bool(export_on_error)) + effective_export_raw_records = (_raw_record_export_enabled + if export_raw_records is _USE_CURRENT_CONFIG else bool(export_raw_records)) + + _active_config = DebuggerConfig( + enabled=True, + record_level=int(effective_level), + addr_level=effective_addr_level, + export_mode=effective_export_mode, + record_capacity=effective_record_capacity, + export_on_error=effective_export_on_error, + output_dir=get_output_dir(), + export_raw_records=effective_export_raw_records, + runtime_metadata_builder=runtime_metadata_builder, + export_handler=export_handler, + ) + register_launch_prepare_hook(_default_launch_prepare_hook) + + try: + import triton + + triton.knobs.compilation.instrumentation_mode = "debugger" + except Exception: + pass + + +def deactivate() -> None: + global _active_config + clear_launch_prepare_hook() + _active_config = DebuggerConfig() + + try: + import triton + + if str(triton.knobs.compilation.instrumentation_mode).startswith("debugger"): + triton.knobs.compilation.instrumentation_mode = "" + except Exception: + pass + + +def clear_exported_runs() -> None: + _exported_runs.clear() + + +def peek_exported_runs() -> list[dict[str, Any]]: + return list(_exported_runs) + + +def take_exported_runs() -> list[dict[str, Any]]: + exported_runs = list(_exported_runs) + _exported_runs.clear() + return exported_runs + + +def prepare_kernel_launch(metadata: Any, stream: int, launch_metadata: Any = None, + kernel_args: Optional[Sequence[Any]] = None) -> Optional[PreparedKernelLaunch]: + enabled = bool(getattr(metadata, "debug_enabled", False)) or _active_config.enabled + if not enabled: + return None + + if _launch_prepare_hook is None: + raise RuntimeError("debug-enabled kernel launch requires " + "triton.runtime.debugger.register_launch_prepare_hook(...)") + + prepared = _launch_prepare_hook( + metadata, + int(stream), + launch_metadata, + tuple(kernel_args or ()), + ) + if prepared is None: + raise RuntimeError("debugger launch prepare hook returned None for a debug-enabled " + "kernel") + if isinstance(prepared, PreparedKernelLaunch): + return PreparedKernelLaunch( + kernel_args=_normalize_kernel_args(prepared.kernel_args), + finalize=prepared.finalize, + ) + return PreparedKernelLaunch(kernel_args=_normalize_kernel_args(prepared)) + + +def finalize_prepared_launch(prepared: Optional[PreparedKernelLaunch], error: Optional[BaseException] = None) -> None: + if prepared is None or prepared.finalize is None: + return + prepared.finalize(error) diff --git a/python/triton/runtime/interpreter.py b/python/triton/runtime/interpreter.py index aaacc9db2..4d74f9294 100644 --- a/python/triton/runtime/interpreter.py +++ b/python/triton/runtime/interpreter.py @@ -715,6 +715,12 @@ def create_barrier(self): # Triton's barrier applies to each program in a grid, so it's a no-op in the interpreter pass + def create_debug_collect_begin(self, level, addr_level=-1): + del level, addr_level + + def create_debug_collect_end(self): + pass + def create_make_block_ptr(self, base, shape, strides, offsets, block_shape, order): # Create new offsets to avoid modifying the original new_offsets = [offset.clone() for offset in offsets] diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index e66b7b749..a4e44fef8 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -26,6 +26,15 @@ T = TypeVar("T") + +def _apply_compilation_instrumentation_mode(kwargs): + if "instrumentation_mode" in kwargs: + return + mode = getattr(knobs.compilation, "instrumentation_mode", "") + if mode: + kwargs["instrumentation_mode"] = str(mode) + + # ----------------------------------------------------------------------------- # Dependencies Finder # ----------------------------------------------------------------------------- @@ -717,6 +726,7 @@ def _pack_args(self, backend, kwargs, bound_args, specialization, options): def run(self, *args, grid, warmup, **kwargs): kwargs["debug"] = kwargs.get("debug", self.debug) or knobs.runtime.debug + _apply_compilation_instrumentation_mode(kwargs) # parse options device = driver.active.get_current_device() diff --git a/python/triton/spec/ascend/compiler/code_generator.py b/python/triton/spec/ascend/compiler/code_generator.py index f6a023476..d04b38d4a 100644 --- a/python/triton/spec/ascend/compiler/code_generator.py +++ b/python/triton/spec/ascend/compiler/code_generator.py @@ -485,6 +485,53 @@ def set_value(self, name: str, value: Union[base_value, constexpr]) -> None: self.lscope[name] = value self.local_defs[name] = value + def _statement_source(self, node): + source = None + if hasattr(node, "lineno"): + try: + source = ast.get_source_segment(self.jit_fn.src, node) + except Exception: + source = None + if source is None and hasattr(ast, "unparse"): + try: + source = ast.unparse(node) + except Exception: + source = None + return "" if source is None else " ".join(source.strip().split()) + + def _statement_id(self, node): + line = int(self.begin_line + getattr(node, "lineno", 0)) + col = int(getattr(node, "col_offset", 0)) + return max(0, min(line * 1000 + min(col, 999), (1 << 31) - 1)) + + def _annotate_assignment_value(self, node, name, value, statement_source): + if not _is_triton_tensor(value): + return + handle = getattr(value, "handle", None) + if handle is None or not hasattr(handle, "set_attr"): + return + if statement_source: + handle.set_attr("flagtree.debug.triton_statement", self.builder.get_string_attr(statement_source)) + handle.set_attr("flagtree.debug.statement_result_name", self.builder.get_string_attr(str(name))) + handle.set_attr("flagtree.debug.statement_id", self.builder.get_int32_attr(self._statement_id(node))) + + def _annotate_assignment_target(self, node, target, value, statement_source): + if isinstance(target, ast.Name): + self._annotate_assignment_value(node, target.id, value, statement_source) + elif isinstance(target, ast.Tuple): + for child, child_value in zip(target.elts, value.values): + self._annotate_assignment_target(node, child, child_value, statement_source) + + def _annotate_expression_value(self, node, value, statement_source): + if not _is_triton_tensor(value): + return + handle = getattr(value, "handle", None) + if handle is None or not hasattr(handle, "set_attr"): + return + if statement_source: + handle.set_attr("flagtree.debug.triton_statement", self.builder.get_string_attr(statement_source)) + handle.set_attr("flagtree.debug.statement_id", self.builder.get_int32_attr(self._statement_id(node))) + def _get_insertion_point_and_loc(self, builder=None): # XXX: this is a hack to get the location of the insertion point. # The insertion point's location could be invalid sometimes, @@ -727,6 +774,7 @@ def _sanitize_value(value): values = _sanitize_value(self.visit(node.value)) else: values = _sanitize_value(self.visit(node.value)) + self._annotate_assignment_target(node, target, values, self._statement_source(node)) self.assignTarget(target, values) def visit_AugAssign(self, node): @@ -1403,6 +1451,10 @@ def call_Function(self, node, fn, args, kws): extra_kwargs["_semantic"] = self.semantic if '_generator' in sig.parameters: extra_kwargs['_generator'] = self + if '_debug_statement_source' in sig.parameters: + extra_kwargs['_debug_statement_source'] = self._statement_source(node) + if '_debug_statement_id' in sig.parameters: + extra_kwargs['_debug_statement_id'] = self._statement_id(node) try: if fn.__name__ == 'copy': # extract tle hints from the generator to identify if node in the tle hints scope @@ -1540,7 +1592,9 @@ def visit_Attribute(self, node): def visit_Expr(self, node): node.value._is_unused = True - ast.NodeVisitor.generic_visit(self, node) + value = self.visit(node.value) + self._annotate_expression_value(node, value, self._statement_source(node)) + return value def visit_NoneType(self, node): return None diff --git a/python/triton/spec/ascend/compiler/compiler.py b/python/triton/spec/ascend/compiler/compiler.py index 505141ac9..47f25a24f 100644 --- a/python/triton/spec/ascend/compiler/compiler.py +++ b/python/triton/spec/ascend/compiler/compiler.py @@ -13,6 +13,7 @@ from ..tools.disasm import get_sass from .errors import MLIRCompilationError from pathlib import Path +import importlib import re import functools import os @@ -280,6 +281,12 @@ def compile(src, target=None, options=None, _env_vars=None): **options.__dict__, **env_vars, } + # Debug instrumentation consumes target identity in the TTIR stage. Ascend + # options do not define backend_name, so populate it from GPUTarget before + # instrumentation selects the runtime transfer backend. + metadata["backend_name"] = str(target.backend) + metadata["debug_backend_name"] = str(target.backend) + metadata["debug_target_name"] = str(target.arch) metadata["triton_version"] = __version__ # run compilation pipeline and populate metadata stages = dict() @@ -513,10 +520,27 @@ def run(self): return self._run def launch_metadata(self, grid, stream, *args): - if knobs.runtime.launch_enter_hook is None: + debugger_active = False + try: + debugger_active = importlib.import_module("triton.runtime.debugger").is_active() + except Exception: + pass + if (knobs.runtime.launch_enter_hook is None and not getattr(self.metadata, "debug_enabled", False) + and not debugger_active): return None self._init_handles() - ret = LazyDict({"name": self.name, "function": self.function, "stream": stream}) + grid_size = len(grid) + normalized_grid = ( + int(grid[0]), + int(grid[1]) if grid_size > 1 else 1, + int(grid[2]) if grid_size > 2 else 1, + ) + ret = LazyDict({ + "name": self.name, + "function": self.function, + "stream": stream, + "grid": normalized_grid, + }) if not isinstance(self.src, ASTSource) or self.src.fn.launch_metadata is None: return ret arg_dict = {name: arg for name, arg in zip(self.src.fn.arg_names, args)} diff --git a/python/triton/spec/ascend/language/core.py b/python/triton/spec/ascend/language/core.py index 0a8daeee1..632a5b10e 100644 --- a/python/triton/spec/ascend/language/core.py +++ b/python/triton/spec/ascend/language/core.py @@ -2158,7 +2158,8 @@ def store_tensor_descriptor(desc: tensor_descriptor_base, offsets: Sequence[cons @_tensor_member_fn @builtin -def store(pointer, value, mask=None, boundary_check=(), cache_modifier="", eviction_policy="", _semantic=None): +def store(pointer, value, mask=None, boundary_check=(), cache_modifier="", eviction_policy="", _semantic=None, + _debug_statement_source=None, _debug_statement_id=None): """ Store a tensor of data into memory locations defined by `pointer`. @@ -2204,7 +2205,8 @@ def store(pointer, value, mask=None, boundary_check=(), cache_modifier="", evict mask = _semantic.to_tensor(mask) cache_modifier = _unwrap_if_constexpr(cache_modifier) eviction_policy = _unwrap_if_constexpr(eviction_policy) - return _semantic.store(pointer, value, mask, boundary_check, cache_modifier, eviction_policy) + return _semantic.store(pointer, value, mask, boundary_check, cache_modifier, eviction_policy, + _debug_statement_source=_debug_statement_source, _debug_statement_id=_debug_statement_id) @builtin @@ -2864,6 +2866,27 @@ def debug_barrier(_semantic=None): return _semantic.debug_barrier() +@builtin +def debug_collect_start(level=1, addr_level=None, _semantic=None): + '''Begin a FlagTree debug collect region.''' + level_v = level.value if isinstance(level, constexpr) else level + if not isinstance(level_v, int): + raise TypeError("tl.debug_collect_start: level must be an integer") + + addr_level_v = addr_level.value if isinstance(addr_level, constexpr) else addr_level + if addr_level_v is not None and not isinstance(addr_level_v, int): + raise TypeError("tl.debug_collect_start: addr_level must be an integer") + if addr_level_v is not None and not 0 <= addr_level_v <= 2: + raise ValueError("tl.debug_collect_start: addr_level must be 0, 1, or 2") + return _semantic.debug_collect_start(int(level_v), addr_level_v) + + +@builtin +def debug_collect_end(_semantic=None): + '''End a FlagTree debug collect region.''' + return _semantic.debug_collect_end() + + @builtin def multiple_of(input, values, _semantic=None): """ diff --git a/python/triton/spec/ascend/language/semantic.py b/python/triton/spec/ascend/language/semantic.py index 1c43eef24..baf421fd6 100644 --- a/python/triton/spec/ascend/language/semantic.py +++ b/python/triton/spec/ascend/language/semantic.py @@ -1286,7 +1286,17 @@ def descriptor_scatter(self, desc, value: TensorTy, x_offsets, y_offset) -> Tens self.builder.create_descriptor_scatter(desc.handle, value.handle, x_offsets.handle, y_offset) return self.tensor(None, tl.void) - def _store_block_pointer(self, ptr, val, mask, boundary_check, cache, eviction): + def _annotate_debug_statement(self, op, source, statement_id): + if op is None or not hasattr(op, "set_attr"): + return op + if source: + op.set_attr("flagtree.debug.triton_statement", self.builder.get_string_attr(source)) + if statement_id is not None: + op.set_attr("flagtree.debug.statement_id", self.builder.get_int32_attr(statement_id)) + return op + + def _store_block_pointer(self, ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source=None, + _debug_statement_id=None): # Store by a block pointer: `pointer_type>` # Block pointers can not have the `mask` argument if mask is not None: @@ -1311,10 +1321,11 @@ def _store_block_pointer(self, ptr, val, mask, boundary_check, cache, eviction): val = self.cast(val, elt_ty) # Build IR - return self.tensor( - self.builder.create_tensor_pointer_store(ptr.handle, val.handle, boundary_check, cache, eviction), tl.void) + op = self.builder.create_tensor_pointer_store(ptr.handle, val.handle, boundary_check, cache, eviction) + return self.tensor(self._annotate_debug_statement(op, _debug_statement_source, _debug_statement_id), tl.void) - def _store_legacy(self, ptr, val, mask, boundary_check, cache, eviction): + def _store_legacy(self, ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source=None, + _debug_statement_id=None): # Store by a tensor of pointers or a pointer of scalar: `block_type>` or `pointer_type<>` if not ptr.type.scalar.is_ptr(): raise ValueError(f"Unsupported ptr type {ptr.type.__repr__()} in `tl.store`") @@ -1352,14 +1363,15 @@ def _store_legacy(self, ptr, val, mask, boundary_check, cache, eviction): # Build IR if mask is None: - return self.tensor(self.builder.create_store(ptr.handle, val.handle, cache, eviction), tl.void) - if not mask.type.scalar.is_bool(): - raise ValueError("Mask must have boolean scalar type") - return self.tensor(self.builder.create_masked_store(ptr.handle, val.handle, mask.handle, cache, eviction), - tl.void) + op = self.builder.create_store(ptr.handle, val.handle, cache, eviction) + else: + if not mask.type.scalar.is_bool(): + raise ValueError("Mask must have boolean scalar type") + op = self.builder.create_masked_store(ptr.handle, val.handle, mask.handle, cache, eviction) + return self.tensor(self._annotate_debug_statement(op, _debug_statement_source, _debug_statement_id), tl.void) def store(self, ptr: TensorTy, val: TensorTy, mask: Optional[TensorTy], boundary_check, cache_modifier: str, - eviction_policy: str) -> TensorTy: + eviction_policy: str, _debug_statement_source=None, _debug_statement_id=None) -> TensorTy: # Cache and eviction options cache = self._str_to_store_cache_modifier(cache_modifier) eviction = self._str_to_eviction_policy(eviction_policy) @@ -1369,10 +1381,12 @@ def store(self, ptr: TensorTy, val: TensorTy, mask: Optional[TensorTy], boundary if ptr.type.is_ptr() and ptr.type.element_ty.is_block(): # Store by a block pointer: `pointer_type>` - return self._store_block_pointer(ptr, val, mask, boundary_check, cache, eviction) + return self._store_block_pointer(ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source, + _debug_statement_id) else: # Store by a tensor of pointers or a pointer of scalar: `block_type>` or `pointer_type<>` - return self._store_legacy(ptr, val, mask, boundary_check, cache, eviction) + return self._store_legacy(ptr, val, mask, boundary_check, cache, eviction, _debug_statement_source, + _debug_statement_id) ######### # atomic @@ -1875,6 +1889,19 @@ def max_constancy(self, x: TensorTy, values: List[int]) -> TensorTy: def debug_barrier(self) -> TensorTy: return self.tensor(self.builder.create_barrier(), tl.void) + def debug_collect_start(self, level: int, addr_level: Optional[int]) -> TensorTy: + if not hasattr(self.builder, "create_debug_collect_begin"): + raise RuntimeError("FlagTree debugger support is not available in this build; " + "rebuild with -DFLAGTREE_ENABLE_DEBUGGER=ON") + handle = self.builder.create_debug_collect_begin(level, -1 if addr_level is None else int(addr_level)) + return self.tensor(handle, tl.void) + + def debug_collect_end(self) -> TensorTy: + if not hasattr(self.builder, "create_debug_collect_end"): + raise RuntimeError("FlagTree debugger support is not available in this build; " + "rebuild with -DFLAGTREE_ENABLE_DEBUGGER=ON") + return self.tensor(self.builder.create_debug_collect_end(), tl.void) + def device_print(self, prefix: str, args: List[TensorTy], hex: bool) -> TensorTy: # It makes sense visually for prefix to end in ": "; make it so. Also, # non-empty prefixes should start with " ". diff --git a/python/triton/spec/ascend/runtime/jit.py b/python/triton/spec/ascend/runtime/jit.py index 1cdaf1dfc..6a3e643bd 100644 --- a/python/triton/spec/ascend/runtime/jit.py +++ b/python/triton/spec/ascend/runtime/jit.py @@ -26,6 +26,61 @@ T = TypeVar("T") + +def _apply_compilation_instrumentation_mode(kwargs): + if "instrumentation_mode" in kwargs: + return + mode = getattr(knobs.compilation, "instrumentation_mode", "") + if mode: + kwargs["instrumentation_mode"] = str(mode) + + +def _disable_flagtree_debug_hidden_arg(kernel): + run = getattr(kernel, "run", None) + if run is not None: + run.debug_launch_hidden_arg = False + run.debug_ctrl_ptr = 0 + + +def _prepare_flagtree_debug_launch(kernel, stream, launch_metadata, kernel_args): + records_per_instance = int(getattr(kernel.metadata, "debug_records_per_instance", 0)) + if (not getattr(kernel.metadata, "debug_launch_hidden_arg", False) or records_per_instance <= 0): + _disable_flagtree_debug_hidden_arg(kernel) + if not getattr(kernel.metadata, "debug_enabled", False): + return None + from triton.runtime import debugger as flagtree_debugger + + return flagtree_debugger.prepare_metadata_only_kernel_launch(kernel.metadata, stream, launch_metadata, + kernel_args) + + from triton.runtime import debugger as flagtree_debugger + + prepared = flagtree_debugger.prepare_kernel_launch(kernel.metadata, stream, launch_metadata, kernel_args) + if prepared is not None and prepared.kernel_args: + kernel.run.debug_ctrl_ptr = int(prepared.kernel_args[0]) + return prepared + + +def _finalize_flagtree_debug_launch(prepared, error): + if prepared is None: + return + from triton.runtime import debugger as flagtree_debugger + + final_error = error + if final_error is None: + try: + import torch_npu + + torch_npu.npu.synchronize() + except BaseException as exc: + final_error = exc + try: + flagtree_debugger.finalize_prepared_launch(prepared, final_error) + finally: + if error is None and final_error is not None: + raise final_error + + # ----------------------------------------------------------------------------- # Dependencies Finder # ----------------------------------------------------------------------------- @@ -708,6 +763,7 @@ def _pack_args(self, backend, kwargs, bound_args, specialization, options): def run(self, *args, grid, warmup, **kwargs): kwargs["debug"] = kwargs.get("debug", self.debug) or knobs.runtime.debug + _apply_compilation_instrumentation_mode(kwargs) # parse options device = driver.active.get_current_device() @@ -753,9 +809,18 @@ def run(self, *args, grid, warmup, **kwargs): if hasattr(kernel, "result"): kernel = kernel.result() # launch kernel - launch_metadata = kernel.launch_metadata(grid, stream, *bound_args.values()) - kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata, - knobs.runtime.launch_enter_hook, knobs.runtime.launch_exit_hook, *bound_args.values()) + kernel_args = list(bound_args.values()) + launch_metadata = kernel.launch_metadata(grid, stream, *kernel_args) + prepared_debug_launch = _prepare_flagtree_debug_launch(kernel, stream, launch_metadata, kernel_args) + launch_error = None + try: + kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata, + knobs.runtime.launch_enter_hook, knobs.runtime.launch_exit_hook, *kernel_args) + except BaseException as exc: + launch_error = exc + raise + finally: + _finalize_flagtree_debug_launch(prepared_debug_launch, launch_error) return kernel def repr(self, _): diff --git a/test/lit.cfg.py b/test/lit.cfg.py index 0890e609b..2ff6ad423 100644 --- a/test/lit.cfg.py +++ b/test/lit.cfg.py @@ -42,7 +42,7 @@ # test_exec_root: The root path where tests should be run. config.test_exec_root = os.path.join(config.triton_obj_root, 'test') -config.triton_tools_dir = os.path.join(config.triton_obj_root, 'bin') +config.triton_tools_dir = getattr(config, 'triton_tools_dir', os.path.join(config.triton_obj_root, 'bin')) config.filecheck_dir = os.path.join(config.triton_obj_root, 'bin', 'FileCheck') # FileCheck -enable-var-scope is enabled by default in MLIR test diff --git a/third_party/Debugger/CMakeLists.txt b/third_party/Debugger/CMakeLists.txt new file mode 100644 index 000000000..37c92a270 --- /dev/null +++ b/third_party/Debugger/CMakeLists.txt @@ -0,0 +1,4 @@ +# FlagTree Debugger is optional so regular FlagTree builds can omit all of its +# native sources, generated IR, runtime support, and bindings. +add_subdirectory(include/Debugger) +add_subdirectory(lib) diff --git a/third_party/Debugger/README.md b/third_party/Debugger/README.md new file mode 100644 index 000000000..e70b88337 --- /dev/null +++ b/third_party/Debugger/README.md @@ -0,0 +1,505 @@ +# FlagTree Debugger + +`third_party/Debugger` 是 FlagTree debugger 的独立可选模块，包含编译期插桩、 +运行期导出、record 解码和报告生成等实现。跨模块接口和协议定义以 +`third_party/Debugger/include/Debugger` 为准。 + +本文面向 debugger 使用者和维护者，说明功能目的、用户接口、输出格式、可采集 +指标、运行示例以及简要设计架构。 + +## 编译开关 + +Debugger 由 CMake option `FLAGTREE_ENABLE_DEBUGGER` 控制，默认开启。容器内可用 +以下命令分别构建启用或禁用版本： + +```bash +FLAGTREE_ENABLE_DEBUGGER=ON bash build.sh --rebuild +FLAGTREE_ENABLE_DEBUGGER=OFF bash build.sh --rebuild +``` + +禁用时不会生成或编译 Debugger dialect、passes、runtime、native binding 和测试 +target；普通 `import triton` 与非 debugger kernel 保持可用。禁用版本调用 +`triton.enable_debug()` 或 `tl.debug_collect_start/end` 会明确提示使用 +`-DFLAGTREE_ENABLE_DEBUGGER=ON` 重新构建。 + +可运行示例位于 `third_party/Debugger/examples/`，精简的 FlagGems 回归样例位于 +`third_party/Debugger/samples/`。生成的 kernel、报告、cache 和复制的第三方源码 +不进入仓库。 + +## 目的 + +FlagTree debugger 用于在 Triton kernel 内采集指定代码区域的运行期调试信息。 +它将编译期静态信息与运行期动态记录关联起来，用于定位 kernel 内部 operation +的数值状态、内存访问状态和 full dump 数据。 + +## 用户接口 + +### Python 配置接口 + +用户在 Python 侧通过 `triton.runtime.debugger` 配置输出目录、record 容量和导出 +选项，并通过 `triton.enable_debug(...)` 开启 debugger 编译和运行流程。 + +```python +import triton +from triton.runtime import debugger + +debugger.configure( + output_dir="/tmp/flagtree_debugger_example", + record_capacity=4096, + export_raw_records=False, +) +triton.enable_debug(level=1, addr_level=1) +``` + +常用接口： + +- `debugger.configure(...)`：配置输出目录、record 容量和导出选项。 +- `debugger.get_config()`：读取当前 debugger 配置。 +- `debugger.reset_config()`：恢复默认配置。 +- `triton.enable_debug(level=..., addr_level=...)`：开启后续 kernel 的 debugger + pipeline。 +- `triton.disable_debug()`：关闭 debugger pipeline。 + +### Triton JIT 采集接口 + +Python 侧 `triton.enable_debug(...)` 仅开启 debugger pipeline。实际采集范围由 +Triton JIT 函数内的 `tl.debug_collect_start(...)` 和 `tl.debug_collect_end()` 界定。 + +```python +import triton +import triton.language as tl + + +@triton.jit +def kernel(x_ptr, y_ptr, n: tl.constexpr, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n + + tl.debug_collect_start(level=1, addr_level=1) + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.abs(x) + tl.store(y_ptr + offsets, y, mask=mask) + tl.debug_collect_end() +``` + +采集等级： + +- `level=1`：采集 summary 指标。 +- `level=2`：采集 full tensor value。 +- `addr_level=1`：采集 memory address summary。 +- `addr_level=2`：采集 full memory address。 + +## 输出 + +### level 1 输出示例 + +level 1 通常生成主文本报告、主 JSON 报告和 IR op 级日志报告。例如： + +```text +test_debug_abs_kernel_aiv_20260629_193456_497_run1.txt +test_debug_abs_kernel_aiv_20260629_193456_497_run1.json +test_debug_abs_kernel_aiv_20260629_193456_497_run1_op_log.txt +test_debug_abs_kernel_aiv_20260629_193456_497_run1_op_log.json +``` + +通用格式： + +```text +