From 01024575a523a862b9795d43d57e6e87d4e35fe8 Mon Sep 17 00:00:00 2001 From: Mogball Date: Thu, 14 Nov 2024 14:33:10 -0800 Subject: [PATCH 1/5] [Triton] Generate local MLIR reproducers when possible By setting a reproducer path, the pass manager will dump a standard MLIR reproducer before each pass manager invocation. This PR also enables additional local crash reproducer generation (to the same path set through the env var), which tries to narrow down the specific pass that failed, if the pass pipeline fails at any point. --- python/src/ir.cc | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/src/ir.cc b/python/src/ir.cc index 0b8f1c550ad1..1306fc6e75b9 100644 --- a/python/src/ir.cc +++ b/python/src/ir.cc @@ -1707,7 +1707,15 @@ void init_triton_ir(py::module &&m) { auto anchorName = self.getOpAnchorName(); auto passes = self.getPasses(); Operation *op = mod.getOperation(); + // Save a reproducer for the current pass manager invocation + // immediately. makeReproducer(anchorName, passes, op, reproducerPath); + // But if the pass manager crashes, attempt to generate a local + // reproducer. + mod.getContext()->disableMultithreading(); + self.enableCrashReproducerGeneration( + ("crash." + reproducerPath).str(), + /*genLocalReproducer=*/true); } if (triton::tools::getBoolEnv("TRITON_ENABLE_LLVM_DEBUG")) { From 44e410623e63a466ace98ad584a93aac61e2a857 Mon Sep 17 00:00:00 2001 From: Mogball Date: Thu, 14 Nov 2024 14:35:57 -0800 Subject: [PATCH 2/5] document the env var in the README --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 18e46403e726..3bdf9742da5f 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,9 @@ For detailed instructions on how to debug Triton's frontend, please refer to thi kernels. Use `MLIR_ENABLE_DUMP=kernelName` to dump for a specific kernel only. - Triton cache can interfere with the dump. In cases where `MLIR_ENABLE_DUMP=1` does not work, try cleaning your triton cache: `rm -r ~/.triton/cache/*` - `LLVM_IR_ENABLE_DUMP=1` dumps the IR before every pass run over the LLVM IR. +- `TRITON_REPRODUCER_PATH=` will generate an MLIR file at `` + before each MLIR compiler stage. If any of the stages crash, an additional `crash.` + file is generated containing a local crash reproducer. - `TRITON_INTERPRET=1` uses the Triton interpreter instead of running on the GPU. You can insert Python breakpoints in your kernel code! - `TRITON_ENABLE_LLVM_DEBUG=1` passes `-debug` to LLVM, printing a lot of From 316fd8c6ed2359c43303260e8243dcc3562898e2 Mon Sep 17 00:00:00 2001 From: Mogball Date: Thu, 14 Nov 2024 14:36:57 -0800 Subject: [PATCH 3/5] it's not a stringref --- python/src/ir.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/src/ir.cc b/python/src/ir.cc index 1306fc6e75b9..75ac7149f05f 100644 --- a/python/src/ir.cc +++ b/python/src/ir.cc @@ -1713,9 +1713,8 @@ void init_triton_ir(py::module &&m) { // But if the pass manager crashes, attempt to generate a local // reproducer. mod.getContext()->disableMultithreading(); - self.enableCrashReproducerGeneration( - ("crash." + reproducerPath).str(), - /*genLocalReproducer=*/true); + self.enableCrashReproducerGeneration("crash." + reproducerPath, + /*genLocalReproducer=*/true); } if (triton::tools::getBoolEnv("TRITON_ENABLE_LLVM_DEBUG")) { From 9ff1d3b19e2ed0ee09d278f91d38fcd90ef3893c Mon Sep 17 00:00:00 2001 From: Mogball Date: Thu, 14 Nov 2024 14:42:56 -0800 Subject: [PATCH 4/5] override regular reproducer with the local crash repro --- README.md | 6 +++--- python/src/ir.cc | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 3bdf9742da5f..4b724d2a1bbd 100644 --- a/README.md +++ b/README.md @@ -176,9 +176,9 @@ For detailed instructions on how to debug Triton's frontend, please refer to thi kernels. Use `MLIR_ENABLE_DUMP=kernelName` to dump for a specific kernel only. - Triton cache can interfere with the dump. In cases where `MLIR_ENABLE_DUMP=1` does not work, try cleaning your triton cache: `rm -r ~/.triton/cache/*` - `LLVM_IR_ENABLE_DUMP=1` dumps the IR before every pass run over the LLVM IR. -- `TRITON_REPRODUCER_PATH=` will generate an MLIR file at `` - before each MLIR compiler stage. If any of the stages crash, an additional `crash.` - file is generated containing a local crash reproducer. +- `TRITON_REPRODUCER_PATH=` will generate an MLIR reproducer file + at `` before each MLIR compiler stage. If any of the stages fail, + `` will be a local MLIR reproducer captured right before the failing pass. - `TRITON_INTERPRET=1` uses the Triton interpreter instead of running on the GPU. You can insert Python breakpoints in your kernel code! - `TRITON_ENABLE_LLVM_DEBUG=1` passes `-debug` to LLVM, printing a lot of diff --git a/python/src/ir.cc b/python/src/ir.cc index 75ac7149f05f..8cee50a25d6f 100644 --- a/python/src/ir.cc +++ b/python/src/ir.cc @@ -1711,9 +1711,9 @@ void init_triton_ir(py::module &&m) { // immediately. makeReproducer(anchorName, passes, op, reproducerPath); // But if the pass manager crashes, attempt to generate a local - // reproducer. + // reproducer instead. mod.getContext()->disableMultithreading(); - self.enableCrashReproducerGeneration("crash." + reproducerPath, + self.enableCrashReproducerGeneration(reproducerPath, /*genLocalReproducer=*/true); } From 9008d289a0866c020c80d5eb06d755efb9e33519 Mon Sep 17 00:00:00 2001 From: Mogball Date: Thu, 14 Nov 2024 14:49:56 -0800 Subject: [PATCH 5/5] add a source diagnostic handler --- python/src/ir.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/src/ir.cc b/python/src/ir.cc index 8cee50a25d6f..e7322c4fd232 100644 --- a/python/src/ir.cc +++ b/python/src/ir.cc @@ -1747,6 +1747,12 @@ void init_triton_ir(py::module &&m) { self.enableTiming(); } + // Run the pass manager under a source manager diagnostic handler, which + // enables emitted MLIR diagnostics to directly reference Python source + // code. + llvm::SourceMgr sourceMgr; + SourceMgrDiagnosticHandler diagHandler(sourceMgr, mod.getContext(), + llvm::errs()); if (failed(self.run(mod.getOperation()))) throw std::runtime_error("PassManager::run failed"); });