github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[BACKEND] Fix enable_debug implementation. (#1876)

Browse Source 
Print before every pass and after failures if MLIR_ENABLE_DUMP is set.

Co-authored-by: Keren Zhou <kerenzhou@openai.com>
This commit is contained in:
Christian Sigg
2023-07-17 03:50:30 +02:00
committed by GitHub
parent bcfd990a88
commit 80c6e39716
1 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/src/triton.cc
View File

@@ -1524,19 +1524,21 @@ void init_triton_ir(py::module &&m) {
.def(py::init<mlir::MLIRContext *>())
.def("enable_debug",
[](mlir::PassManager &self) {
if (!::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP"))
return;
self.getContext()->disableMultithreading();
auto printingFlags = mlir::OpPrintingFlags();
printingFlags.elideLargeElementsAttrs(16);
printingFlags.enableDebugInfo();
auto print_always = [](mlir::Pass *, mlir::Operation *) {
return true;
};
self.enableIRPrinting(
/*shouldPrintBeforePass=*/nullptr,
/*shouldPrintAfterPass=*/
[](mlir::Pass *pass, mlir::Operation *) {
return ::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP");
},
/*printModuleScope=*/false,
/*printAfterOnlyOnChange=*/true,
/*printAfterOnlyOnFailure*/ false, llvm::dbgs(),
printingFlags);
/*shouldPrintBeforePass=*/print_always,
/*shouldPrintAfterPass=*/print_always,
/*printModuleScope=*/true,
/*printAfterOnlyOnChange=*/false,
/*printAfterOnlyOnFailure*/ true, llvm::dbgs(), printingFlags);
})
.def("run",
[](mlir::PassManager &self, mlir::ModuleOp &mod) {