[BACKEND] Standardize IR print options (#6440)

Previously, when we performed IR relocation, we didn’t output locations
of the ModuleOp.

What we printed to the output stream didn’t include locations, but what
we stored in the file still did, resulting in mismatches.

Now that we use a single function to control the format of the output
IR, we always output locations.

Author: Jokeren

python/src/ir.cc

@@ -218,12 +218,18 @@ struct ConsoleReproducerStream : public mlir::ReproducerStream {
   raw_ostream &os() override { return llvm::errs(); }
 };
 
-static ReproducerStreamFactory makeConsoleReproducer() {
+ReproducerStreamFactory makeConsoleReproducer() {
   return [](std::string &error) -> std::unique_ptr<ReproducerStream> {
     return std::make_unique<ConsoleReproducerStream>();
   };
 }
 
+OpPrintingFlags getOpPrintingFlags() {
+  auto printingFlags = OpPrintingFlags();
+  printingFlags.enableDebugInfo();
+  return printingFlags;
+}
+
 } // anonymous namespace
 
 /*****************************************************************************/
@@ -511,8 +517,7 @@ void init_triton_ir(py::module &&m) {
            [](OpState &self) -> std::string {
              std::string str;
              llvm::raw_string_ostream os(str);
-             auto printingFlags = OpPrintingFlags();
-             printingFlags.enableDebugInfo();
+             auto printingFlags = getOpPrintingFlags();
              self->print(os, printingFlags);
              return str;
            })
@@ -583,8 +588,7 @@ void init_triton_ir(py::module &&m) {
            [](ModuleOp &self) -> std::string {
              std::string str;
              llvm::raw_string_ostream os(str);
-             auto printingFlags = OpPrintingFlags();
-             printingFlags.enableDebugInfo();
+             auto printingFlags = getOpPrintingFlags();
              self.print(os, printingFlags);
              return str;
            })
@@ -658,9 +662,9 @@ void init_triton_ir(py::module &&m) {
            })
       .def("create_location_snapshot",
            [](ModuleOp &self, const std::string &fileName) -> void {
-             generateLocationsFromIR(/*raw_ostream=*/llvm::nulls(),
-                                     /*fileName=*/fileName,
-                                     /*op=*/self, /*flags=*/{});
+             auto printingFlags = getOpPrintingFlags();
+             if (failed(generateLocationsFromIR(fileName, self, printingFlags)))
+               throw std::runtime_error("Failed to create location snapshot");
            })
       .def("walk",
            [](ModuleOp &self, const std::function<void(Operation *)> &fn) {
@@ -1780,9 +1784,8 @@ void init_triton_ir(py::module &&m) {
              }
              if (haveDump) {
                context->disableMultithreading();
-               auto printingFlags = OpPrintingFlags();
+               auto printingFlags = getOpPrintingFlags();
                printingFlags.elideLargeElementsAttrs(16);
-               printingFlags.enableDebugInfo();
                auto printAlways = [funcToDump](Pass *, Operation *op) -> bool {
                  if (funcToDump.empty())
                    return true;