Only print out diagnostic messages if an environment variable is set

This prevents crashes in test_core.py due to too many diagnostics emitted in llvm/llvm-project#78228
It should also speed up compile times, as we can use multithreading, and avoid handling diagnostic messages.

PiperOrigin-RevId: 607293980

Author: gflegar

third_party/triton/cl607293980.patch

@@ -0,0 +1,76 @@
+Remove once b/325453581 is fixed
+diff --git a/include/triton/Tools/Sys/GetEnv.hpp b/include/triton/Tools/Sys/GetEnv.hpp
+--- a/include/triton/Tools/Sys/GetEnv.hpp
++++ b/include/triton/Tools/Sys/GetEnv.hpp
+@@ -32,7 +32,7 @@ namespace triton {
+ const std::set<std::string> ENV_VARS = {
+     "DISABLE_MMA_V3",     "TRITON_DISABLE_LINE_INFO", "DISABLE_FAST_REDUCTION",
+     "ENABLE_TMA",         "MLIR_ENABLE_DUMP",         "LLVM_IR_ENABLE_DUMP",
+-    "AMDGCN_ENABLE_DUMP", "DISABLE_LLVM_OPT"};
++    "AMDGCN_ENABLE_DUMP", "DISABLE_LLVM_OPT", "MLIR_ENABLE_DIAGNOSTICS"};
+ 
+ namespace tools {
+ 
+diff --git a/python/src/ir.cc b/python/src/ir.cc
+--- a/python/src/ir.cc
++++ b/python/src/ir.cc
+@@ -1551,29 +1551,36 @@ void init_triton_ir(py::module &&m) {
+       .def("enable_debug",
+            [](mlir::PassManager &self) {
+              auto *context = self.getContext();
+-             context->printOpOnDiagnostic(true);
+-             context->printStackTraceOnDiagnostic(true);
+-             context->disableMultithreading();
+-             context->getDiagEngine().registerHandler(
+-                 [](mlir::Diagnostic &diag) {
+-                   llvm::outs() << diag << "\n";
+-                   return mlir::success();
+-                 });
+-
+-             if (!::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP"))
+-               return;
+-             auto printingFlags = mlir::OpPrintingFlags();
+-             printingFlags.elideLargeElementsAttrs(16);
+-             printingFlags.enableDebugInfo();
+-             auto print_always = [](mlir::Pass *, mlir::Operation *) {
+-               return true;
+-             };
+-             self.enableIRPrinting(
+-                 /*shouldPrintBeforePass=*/print_always,
+-                 /*shouldPrintAfterPass=*/print_always,
+-                 /*printModuleScope=*/true,
+-                 /*printAfterOnlyOnChange=*/false,
+-                 /*printAfterOnlyOnFailure*/ true, llvm::dbgs(), printingFlags);
++             bool have_diagnostics =
++                ::triton::tools::getBoolEnv("MLIR_ENABLE_DIAGNOSTICS");
++             bool have_dump = ::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP");
++             if (have_diagnostics || have_dump) {
++                context->disableMultithreading();
++             }
++             if (have_diagnostics) {
++              context->printOpOnDiagnostic(true);
++              context->printStackTraceOnDiagnostic(true);
++              context->getDiagEngine().registerHandler(
++                  [](mlir::Diagnostic &diag) {
++                    llvm::outs() << diag << "\n";
++                    return mlir::success();
++                  });
++             }
++             if (have_dump) {
++              auto printingFlags = mlir::OpPrintingFlags();
++              printingFlags.elideLargeElementsAttrs(16);
++              printingFlags.enableDebugInfo();
++              auto print_always = [](mlir::Pass *, mlir::Operation *) {
++                return true;
++              };
++              self.enableIRPrinting(
++                  /*shouldPrintBeforePass=*/print_always,
++                  /*shouldPrintAfterPass=*/print_always,
++                  /*printModuleScope=*/true,
++                  /*printAfterOnlyOnChange=*/false,
++                  /*printAfterOnlyOnFailure*/ true, llvm::dbgs(),
++                  printingFlags);
++             }
+            })
+       .def("run", [](mlir::PassManager &self, mlir::ModuleOp &mod) {
+         // TODO: maybe dump module to file and print error for better

third_party/triton/workspace.bzl

@@ -16,5 +16,6 @@ def repo():
         patch_file = [
             # Upstream this in the next integrate
             #"//third_party/triton:cl602997103.patch"
+            "//third_party/triton:cl607293980.patch",
         ],
     )