Revert D85686892 (#530)

Summary:
Pull Request resolved: #530

This diff reverts D85686892
/data/users/daohang/fbtriton/third_party/tlx/tutorials/blackwell-gemm-ws.py:116:44: error: Result has an invalid layout: Could not determine the number of warps per CTA. Operation is not in a context with `ttg.num-warps`.
                    result = tlx.local_load(acc_tmem_subslice1)

Depends on D85686892

Reviewed By: dshi7

Differential Revision: D85779224

fbshipit-source-id: 464f5efe297067606e21d3b54a41c5aade777ae1

Author: Dark Knight

python/src/ir.cc

@@ -88,23 +88,6 @@ llvm::raw_ostream &mlir_dumps_or_dbgs() {
   }
 }
 
-// Function to parse a comma-separated string into a vector of C-style strings
-llvm::SmallVector<const char *, 3>
-parseCommaSeparatedValues(const std::string &input,
-                          llvm::SmallVector<std::string, 3> &storage) {
-  llvm::SmallVector<StringRef, 3> split;
-  llvm::SmallVector<const char *, 3> result;
-  StringRef(input.c_str()).split(split, ',');
-  llvm::transform(split, std::back_inserter(result), [&storage](StringRef str) {
-    // StringRefs are not always null-terminated.
-    // The purpose for this storage pattern is to
-    // produce a collection of C-strings that are.
-    storage.push_back(str.str());
-    return storage.back().c_str();
-  });
-  return result;
-}
-
 // Run the pass manager under a source manager diagnostic handler, which
 // enables emitted MLIR diagnostics to directly reference Python source
 // code. This diagnostic handler supports filtering diagnostic info by
@@ -141,43 +124,6 @@ struct TritonSourceMgrDiagnosticHandler : public SourceMgrDiagnosticHandler {
   llvm::SourceMgr sourceMgr;
 };
 
-TritonSourceMgrDiagnosticHandler
-setupTritonDiagnosticHandler(MLIRContext *context) {
-  bool showOperations = false, showStacktraces = false, showRemarks = false,
-       showWarnings = false;
-
-  if (auto enableDiagnostics =
-          triton::tools::getStrEnv("MLIR_ENABLE_DIAGNOSTICS");
-      !enableDiagnostics.empty()) {
-    llvm::SmallVector<std::string, 3> storage;
-    parseCommaSeparatedValues(enableDiagnostics, storage);
-    for (auto &str : storage) {
-      if (str == "warnings") {
-        showWarnings = true;
-      } else if (str == "remarks") {
-        showRemarks = true;
-      } else if (str == "stacktraces") {
-        showStacktraces = true;
-      } else if (str == "operations") {
-        showOperations = true;
-      }
-      // we show errors by default, so no need to set it
-    }
-  }
-
-  DiagnosticSeverity minSeverity =
-      showWarnings ? DiagnosticSeverity::Warning : DiagnosticSeverity::Error;
-  minSeverity = showRemarks ? DiagnosticSeverity::Remark : minSeverity;
-
-  context->printOpOnDiagnostic(showOperations);
-  context->printStackTraceOnDiagnostic(showStacktraces);
-  if (showStacktraces) {
-    context->disableMultithreading();
-  }
-
-  return TritonSourceMgrDiagnosticHandler(context, minSeverity);
-}
-
 std::string locationToString(Location loc) {
   std::string str;
   llvm::raw_string_ostream os(str);
@@ -186,6 +132,23 @@ std::string locationToString(Location loc) {
   return str;
 }
 
+// Function to parse a comma-separated string into a vector of C-style strings
+llvm::SmallVector<const char *, 3>
+parseCommaSeparatedValues(const std::string &input,
+                          llvm::SmallVector<std::string, 3> &storage) {
+  llvm::SmallVector<StringRef, 3> split;
+  llvm::SmallVector<const char *, 3> result;
+  StringRef(input.c_str()).split(split, ',');
+  llvm::transform(split, std::back_inserter(result), [&storage](StringRef str) {
+    // StringRefs are not always null-terminated.
+    // The purpose for this storage pattern is to
+    // produce a collection of C-strings that are.
+    storage.push_back(str.str());
+    return storage.back().c_str();
+  });
+  return result;
+}
+
 void outputWarning(Location loc, const std::string &msg) {
   std::string locStr = locationToString(loc);
 
@@ -731,12 +694,7 @@ void init_triton_ir(py::module &&m) {
       .def("walk",
            [](ModuleOp &self, const std::function<void(Operation *)> &fn) {
              self.walk(fn);
-           })
-      .def("verify_with_diagnostics", [](ModuleOp &self) {
-        TritonSourceMgrDiagnosticHandler handler =
-            setupTritonDiagnosticHandler(self.getContext());
-        return succeeded(verify(self.getOperation()));
-      });
+           });
 
   m.def("make_attr", [](const std::vector<int> &values, MLIRContext &context) {
     return mlir::cast<Attribute>(DenseIntElementsAttr::get(
@@ -1965,8 +1923,42 @@ void init_triton_ir(py::module &&m) {
               self.enableTiming();
             }
 
-            TritonSourceMgrDiagnosticHandler diagHandler =
-                setupTritonDiagnosticHandler(context);
+            // setting up diagnostics
+            bool showOperations = false, showStacktraces = false,
+                 showRemarks = false, showWarnings = false;
+
+            if (auto enableDiagnostics =
+                    triton::tools::getStrEnv("MLIR_ENABLE_DIAGNOSTICS");
+                !enableDiagnostics.empty()) {
+              llvm::SmallVector<std::string, 3> storage;
+              parseCommaSeparatedValues(enableDiagnostics, storage);
+              for (auto &str : storage) {
+                if (str == "warnings") {
+                  showWarnings = true;
+                } else if (str == "remarks") {
+                  showRemarks = true;
+                } else if (str == "stacktraces") {
+                  showStacktraces = true;
+                } else if (str == "operations") {
+                  showOperations = true;
+                }
+                // we show errors by default, so no need to set it
+              }
+            }
+
+            DiagnosticSeverity minSeverity = showWarnings
+                                                 ? DiagnosticSeverity::Warning
+                                                 : DiagnosticSeverity::Error;
+            minSeverity =
+                showRemarks ? DiagnosticSeverity::Remark : minSeverity;
+
+            TritonSourceMgrDiagnosticHandler diagHandler(context, minSeverity);
+
+            context->printOpOnDiagnostic(showOperations);
+            context->printStackTraceOnDiagnostic(showStacktraces);
+            if (showStacktraces) {
+              context->disableMultithreading();
+            }
             if (failed(self.run(mod.getOperation())))
               throw std::runtime_error("PassManager::run failed");
           },

python/triton/compiler/code_generator.py

@@ -1602,11 +1602,7 @@ def ast_to_ttir(fn, src, context, options, codegen_fns, module_map, module=None)
                               jit_fn=fn, is_kernel=True, file_name=file_name, begin_line=begin_line, options=options,
                               codegen_fns=codegen_fns, module_map=module_map, module=module, is_gluon=fn.is_gluon())
     generator.visit(fn.parse())
-    module = generator.module
+    ret = generator.module
     # module takes ownership of the context
-    module.context = context
-    if not module.verify_with_diagnostics():
-        if not fn.is_gluon():
-            print(module)
-        raise RuntimeError("error encountered during parsing")
-    return module
+    ret.context = context
+    return ret