Add a `-verify-roundtrip` option to `mlir-opt` intended to validate custom printer/parser completeness

[mehdi_amini]

Summary

Running:

MLIR_OPT_CHECK_IR_ROUNDTRIP=1 ninja check-mlir

will now exercises all of our test with a round-trip to bytecode and a comparison for equality.

[mehdi_amini]

I wonder if we should just make it the default?

[ftynse]

mlir/lib/Support/MlirOptMain.cpp

92 ↗

(On Diff #300448)

Do we want to still run passes in this case? I'd just return doVerifyRoundTrip if the flag is set.

[ftynse]

D90179 should fix the bug this found

[rriddle]

mlir/lib/Support/MlirOptMain.cpp

50 ↗

(On Diff #300448)

Can you emit an error here to give context(i.e. that these errors are from round trip) to the errors that may come out of parseSourceString?

64 ↗

(On Diff #300448)

nit: Add braces here, or move the comment.

65 ↗

(On Diff #300448)

nit: failed to round-trip resultant IR?

87 ↗

(On Diff #300448)

Why the newline here?

[rriddle]

In D90088#2351672, @mehdi_amini wrote:

I wonder if we should just make it the default?

IMO, it depends on the user running mlir-opt. If I'm iterating locally, I generally wouldn't want these additional cost/complexity of running this. For tests that run automatically, that default is probably fine. I would strongly prefer that we keep mlir-opt defaults useful for users iterating locally and generally using mlir-opt as a development tool. For testing, could we just have another frontend(alias of mlir-opt?) that has better defaults for testing use cases? We've had conversations of sane defaults in the past, and it seems like the truly contentious ones always have this divide between local development and testing.

[jpienaar]

In D90088#2354323, @rriddle wrote:

In D90088#2351672, @mehdi_amini wrote:

I wonder if we should just make it the default?

IMO, it depends on the user running mlir-opt. If I'm iterating locally, I generally wouldn't want these additional cost/complexity of running this. For tests that run automatically, that default is probably fine. I would strongly prefer that we keep mlir-opt defaults useful for users iterating locally and generally using mlir-opt as a development tool. For testing, could we just have another frontend(alias of mlir-opt?) that has better defaults for testing use cases? We've had conversations of sane defaults in the past, and it seems like the truly contentious ones always have this divide between local development and testing.

I'd also be pro not making it default from orthogonality point of view. Having all tests verifying printing/parsing vs having dedicated tests seems a wrong tradeoff. Yes having it default would catch cases not covered, but the fact that they aren't covered by unit tests is the problematic part that has to be addressed IMHO.

mlir/include/mlir/Support/MlirOptMain.h

47 ↗

(On Diff #300448)

I don't know at which point we refactor this into a options struct, but 6 bools in a row is quite easy to mess up.

mlir/lib/Support/MlirOptMain.cpp

40 ↗

(On Diff #300448)

Nit: this need not verify custom, depending on other flags passed (unless you set generic form to false on first print I believe the command line option for whether set or not takes preference) this just verifies printing and parsing.

77 ↗

(On Diff #300448)

Mmm, we now have 3 bools in a row here, thats a bit error prone

92 ↗

(On Diff #300448)

Or potentially rename to alsoVerifyRoundTrip to avoid confusion?

206 ↗

(On Diff #300448)

Roundtrip

[mehdi_amini]

Rebase (WIP)

[mehdi_amini]

Add support for bytcode, fix location management which also uncover issues with singleblockimplicitterminator dropping locations :(

[mehdi_amini]

Rebase and use bytecode

[mehdi_amini]

Fix a bug that prevented round-trip from actually runnning with bytecode!

[mehdi_amini]

Perform round-trip in a new context to avoid issues with resource handles

[jpienaar]

Looks good to help flushing these issues out.

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

124

roundtripped module differs from reference? (The above are also errors in roundTrip testing)

mlir/test/lit.cfg.py

131 ↗

(On Diff #525480)

Is else needed?

[mehdi_amini]

mlir/lib/Support/MlirOptMain.cpp

87 ↗

(On Diff #300448)

Got there in a rebase/conflict fixing.

92 ↗

(On Diff #300448)

It seems less surprising to me to run passes the user would have specified than silently ignore them?
(We could also error out if passes are provided, but we don't for -verifyDiagnostic I think)

mlir/test/lit.cfg.py

131 ↗

(On Diff #525480)

I don't know what happens if we put mlir-opt twice here. That said I think listing our tools in the tools has no effect right now when we don't substitute

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

 /// against emitted diagnostics.
 /// - verifyPasses enables the IR verifier in-between each pass in the pipeline.
 /// - allowUnregisteredDialects allows to parse and create operation without
 /// registering the Dialect in the MLIRContext.
 /// - preloadDialectsInContext will trigger the upfront loading of all
 ///   dialects from the global registry in the MLIRContext. This option is
 ///   deprecated and will be removed soon.
 /// - emitBytecode will generate bytecode output instead of text.
-LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
-                          std::unique_ptr<llvm::MemoryBuffer> buffer,
-                          const PassPipelineCLParser &passPipeline,
-                          DialectRegistry &registry, bool splitInputFile,
-                          bool verifyDiagnostics, bool verifyPasses,
-                          bool allowUnregisteredDialects,
-                          bool preloadDialectsInContext = false,
-                          bool emitBytecode = false);
+LogicalResult MlirOptMain(
+    llvm::raw_ostream &outputStream, std::unique_ptr<llvm::MemoryBuffer> buffer,
+    const PassPipelineCLParser &passPipeline, DialectRegistry &registry,
+    bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
+    bool allowUnregisteredDialects, bool preloadDialectsInContext = false,
+    bool emitBytecode = false, bool verifyRoundTrip = false);
 
 /// Support a callback to setup the pass manager.
 /// - passManagerSetupFn is the callback invoked to setup the pass manager to
 ///   apply on the loaded IR.
-LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
-                          std::unique_ptr<llvm::MemoryBuffer> buffer,
-                          PassPipelineFn passManagerSetupFn,
-                          DialectRegistry &registry, bool splitInputFile,
-                          bool verifyDiagnostics, bool verifyPasses,
-                          bool allowUnregisteredDialects,
-                          bool preloadDialectsInContext = false,
-                          bool emitBytecode = false);
+LogicalResult MlirOptMain(
+    llvm::raw_ostream &outputStream, std::unique_ptr<llvm::MemoryBuffer> buffer,
+    PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
+    bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
+    bool allowUnregisteredDialects, bool preloadDialectsInContext = false,
+    bool emitBytecode = false, bool verifyRoundTrip = false);
 
 /// Implementation for tools like `mlir-opt`.
 /// - toolName is used for the header displayed by `--help`.
 /// - registry should contain all the dialects that can be parsed in the source.
 /// - preloadDialectsInContext will trigger the upfront loading of all
 ///   dialects from the global registry in the MLIRContext. This option is
 ///   deprecated and will be removed soon.
 LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName,

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 using namespace llvm;
 
+// Return success if the module can correctly round-trip. This intended to test
+// that the custom printers/parsers are complete.
+static LogicalResult doVerifyRoundTrip(ModuleOp moduleOp, ParserConfig &config) {
+  OwningOpRef<ModuleOp> roundtripModule;
+  // Print a first time and parse it back to the roundtripModule.
+  {
+    std::string buffer;
+    llvm::raw_string_ostream ostream(buffer);
+    moduleOp.print(ostream);
+    roundtripModule = parseSourceString<ModuleOp>(ostream.str(), config);
+    if (!roundtripModule)
+      return failure();
+  }
+
+  // Print in the generic form for the reference module and the round-tripped
+  // one and compare the outputs.
+  std::string reference, roundtrip;
+  {
+    llvm::raw_string_ostream ostreamref(reference);
+    moduleOp.print(ostreamref, OpPrintingFlags().printGenericOpForm());
+    llvm::raw_string_ostream ostreamrndtrip(roundtrip);
+    roundtripModule->print(ostreamrndtrip,
+                           OpPrintingFlags().printGenericOpForm());
+  }
+  if (reference != roundtrip)
+    // TODO implement a diff.
+    return moduleOp.emitOpError() << "roundTrip testing failed!";
+
+  return success();
+}
+
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
 /// This typically parses the main source file, runs zero or more optimization
 /// passes, then prints the output.
 ///
 static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
-                                    bool verifyPasses, SourceMgr &sourceMgr,
-                                    MLIRContext *context,
+                                    bool verifyPasses, bool verifyRoundTrip,
+                                    SourceMgr &sourceMgr, MLIRContext *context,
                                     PassPipelineFn passManagerSetupFn,
                                     bool emitBytecode) {
   DefaultTimingManager tm;
   applyDefaultTimingManagerCLOptions(tm);
   TimingScope timing = tm.getRootScope();
 
   // Disable multi-threading when parsing the input file. This removes the
   // unnecessary/costly context synchronization when parsing.
static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
   // Prepare the parser config, and attach any useful/necessary resource
   // handlers.
   ParserConfig config(context);
   attachPassReproducerAsmResource(config, pm, wasThreadingEnabled);
 
   // Parse the input file and reset the context threading state.
   TimingScope parserTiming = timing.nest("Parser");
   OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, config));
+  
+  // Perform round-trip verification if requested
+  if (verifyRoundTrip && failed(doVerifyRoundTrip(*module, config)))
+    return failure();
+
   context->enableMultithreading(wasThreadingEnabled);
   if (!module)
     return failure();
   parserTiming.stop();
 
   // Callback to build the pipeline.
   if (failed(passManagerSetupFn(pm)))
     return failure();
 
   // Run the pipeline.
   if (failed(pm.run(*module)))
     return failure();
 
   // Print the output.
   TimingScope outputTiming = timing.nest("Output");
   if (emitBytecode) {
     writeBytecodeToFile(module->getOperation(), os);
   } else {
     module->print(os);
     os << '\n';
   }
   return success();
 }
 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static LogicalResult
 processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
               bool verifyDiagnostics, bool verifyPasses,
               bool allowUnregisteredDialects, bool preloadDialectsInContext,
-              bool emitBytecode, PassPipelineFn passManagerSetupFn,
-              DialectRegistry &registry, llvm::ThreadPool *threadPool) {
+              bool emitBytecode, bool verifyRoundTrip,
+              PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
+              llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   if (threadPool)
     context.setThreadPool(*threadPool);
 
   // Parse the input file.
   if (preloadDialectsInContext)
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
   if (verifyDiagnostics)
     context.printOpOnDiagnostic(false);
   context.getDebugActionManager().registerActionHandler<DebugCounter>();
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
   if (!verifyDiagnostics) {
     SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
-    return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
-                          &context, passManagerSetupFn, emitBytecode);
+    return performActions(os, verifyDiagnostics, verifyPasses, verifyRoundTrip,
+                          sourceMgr, &context, passManagerSetupFn,
+                          emitBytecode);
   }
 
   SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
 
   // Do any processing requested by command line flags.  We don't care whether
   // these actions succeed or fail, we only care what diagnostics they produce
   // and whether they match our expectations.
-  (void)performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
-                       passManagerSetupFn, emitBytecode);
+  (void)performActions(os, verifyDiagnostics, verifyPasses, verifyRoundTrip,
+                       sourceMgr, &context, passManagerSetupFn, emitBytecode);
 
   // Verify the diagnostic handler to make sure that each of the diagnostics
   // matched.
   return sourceMgrHandler.verify();
 }
 
 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 std::unique_ptr<MemoryBuffer> buffer,
                                 PassPipelineFn passManagerSetupFn,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
                                 bool preloadDialectsInContext,
-                                bool emitBytecode) {
+                                bool emitBytecode, bool verifyRoundtrip) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   // We use an explicit threadpool to avoid creating and joining/destroying
   // threads for each of the split.
   ThreadPool *threadPool = nullptr;
 
   // Create a temporary context for the sake of checking if
   // --mlir-disable-threading was passed on the command line.
   // We use the thread-pool this context is creating, and avoid
   // creating any thread when disabled.
   MLIRContext threadPoolCtx;
   if (threadPoolCtx.isMultithreadingEnabled())
     threadPool = &threadPoolCtx.getThreadPool();
 
   auto chunkFn = [&](std::unique_ptr<MemoryBuffer> chunkBuffer,
                      raw_ostream &os) {
     return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                          verifyPasses, allowUnregisteredDialects,
-                         preloadDialectsInContext, emitBytecode,
+                         preloadDialectsInContext, emitBytecode,verifyRoundtrip, 
                          passManagerSetupFn, registry, threadPool);
   };
   return splitAndProcessBuffer(std::move(buffer), chunkFn, outputStream,
                                splitInputFile, /*insertMarkerInOutput=*/true);
 }
 
 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 std::unique_ptr<MemoryBuffer> buffer,
                                 const PassPipelineCLParser &passPipeline,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
                                 bool preloadDialectsInContext,
-                                bool emitBytecode) {
+                                bool emitBytecode, bool verifyRoundTrip) {
   auto passManagerSetupFn = [&](PassManager &pm) {
     auto errorHandler = [&](const Twine &msg) {
       emitError(UnknownLoc::get(pm.getContext())) << msg;
       return failure();
     };
     return passPipeline.addToPipeline(pm, errorHandler);
   };
   return MlirOptMain(outputStream, std::move(buffer), passManagerSetupFn,
                      registry, splitInputFile, verifyDiagnostics, verifyPasses,
                      allowUnregisteredDialects, preloadDialectsInContext,
-                     emitBytecode);
+                     emitBytecode, verifyRoundTrip);
 }
 
 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                                 DialectRegistry &registry,
                                 bool preloadDialectsInContext) {
   static cl::opt<std::string> inputFilename(
       cl::Positional, cl::desc("<input file>"), cl::init("-"));
 
   static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
                                              cl::value_desc("filename"),
                                              cl::init("-"));
 
   static cl::opt<bool> splitInputFile(
       "split-input-file",
       cl::desc("Split the input file into pieces and process each "
                "chunk independently"),
       cl::init(false));
 
   static cl::opt<bool> verifyDiagnostics(
       "verify-diagnostics",
       cl::desc("Check that emitted diagnostics match "
                "expected-* lines on the corresponding line"),
       cl::init(false));
 
+  static cl::opt<bool> verifyRoundTrip(
+      "verify-roundtrip",
+      cl::desc(
+          "Check that the input IR can rountrip before executing the pipeline"),
+      cl::init(false));
+
   static cl::opt<bool> verifyPasses(
       "verify-each",
       cl::desc("Run the verifier after each transformation pass"),
       cl::init(true));
 
   static cl::opt<bool> allowUnregisteredDialects(
       "allow-unregistered-dialect",
       cl::desc("Allow operation with no registered dialects"), cl::init(false));
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
   if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
                          splitInputFile, verifyDiagnostics, verifyPasses,
                          allowUnregisteredDialects, preloadDialectsInContext,
-                         emitBytecode)))
+                         emitBytecode, verifyRoundTrip)))
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
   return success();
 }

mlir/test/Dialect/Affine/ops.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -verify-roundtrip -allow-unregistered-dialect -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-op-generic | FileCheck -check-prefix=GENERIC %s
 
 // Check that the attributes for the affine operations are round-tripped.
 // Check that `affine.yield` is visible in the generic form.
 // CHECK-LABEL: @empty
 func.func @empty() {
   // CHECK: affine.for
   // CHECK-NEXT: } {some_attr = true}

mlir/test/Dialect/LLVMIR/roundtrip.mlir

-// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// RUN: mlir-opt %s -verify-roundtrip | FileCheck %s
 
 // CHECK-LABEL: func @ops
 // CHECK-SAME: (%[[I32:.*]]: i32, %[[FLOAT:.*]]: f32, %[[I8PTR1:.*]]: !llvm.ptr<i8>, %[[I8PTR2:.*]]: !llvm.ptr<i8>, %[[BOOL:.*]]: i1, %[[VI8PTR1:.*]]: !llvm.vec<2 x ptr<i8>>)
 func.func @ops(%arg0: i32, %arg1: f32,
           %arg2: !llvm.ptr<i8>, %arg3: !llvm.ptr<i8>,
           %arg4: i1, %arg5 : !llvm.vec<2x!llvm.ptr<i8>>) {
 // Integer arithmetic binary operations.
 //

mlir/test/Dialect/Linalg/roundtrip.mlir

-// RUN: mlir-opt %s | mlir-opt | FileCheck %s
-// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
+// RUN: mlir-opt -verify-roundtrip %s | FileCheck %s
 
 // TODO: Re-enable LLVM lowering test.
 //
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
 // DISABLED: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
 
 // CHECK-DAG: #[[$strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECK-DAG: #[[$strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>

mlir/test/Dialect/X86Vector/roundtrip.mlir

-// RUN: mlir-opt -verify-diagnostics %s | mlir-opt | FileCheck %s
+// RUN: mlir-opt -verify-diagnostics %s  -verify-roundtrip | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: func @avx512_mask_rndscale
 func.func @avx512_mask_rndscale(%a: vector<16xf32>, %b: vector<8xf64>, %i32: i32, %i16: i16, %i8: i8)
   -> (vector<16xf32>, vector<8xf64>)
 {
   // CHECK: x86vector.avx512.mask.rndscale {{.*}}: vector<16xf32>
   %0 = x86vector.avx512.mask.rndscale %a, %i32, %a, %i16, %i32 : vector<16xf32>
   // CHECK: x86vector.avx512.mask.rndscale {{.*}}: vector<8xf64>

mlir/test/IR/core-ops.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
-// Verify the printed output can be parsed.
-// RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -verify-roundtrip | FileCheck %s
 // Verify the generic form can be parsed.
 // RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
 
 // CHECK: #map0 = affine_map<(d0) -> (d0 + 1)>
 
 // CHECK: #map1 = affine_map<()[s0] -> (s0 + 1)>
 
 // CHECK-DAG: #[[$BASE_MAP0:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>

mlir/test/IR/recursive-type.mlir

-// RUN: mlir-opt %s -test-recursive-types | FileCheck %s
+// RUN: mlir-opt %s -test-recursive-types -verify-roundtrip | FileCheck %s
 
 // CHECK: !testrec = !test.test_rec<type_to_alias, test_rec<type_to_alias>>
 
 // CHECK-LABEL: @roundtrip
 func.func @roundtrip() {
   // CHECK: !test.test_rec<a, test_rec<b, test_type>>
   "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<a, test_rec<b, test_type>>
   // CHECK: !test.test_rec<c, test_rec<c>>

mlir/test/lib/Dialect/Test/TestDialect.cpp

 
 ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {
   // Parse the body region, and reuse the operand info as the argument info.
   Region *body = result.addRegion();
   return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
 }
 
 void AffineScopeOp::print(OpAsmPrinter &p) {
-  p << "test.affine_scope ";
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // Test parser.
 //===----------------------------------------------------------------------===//
 
 ParseResult ParseIntegerLiteralOp::parse(OpAsmParser &parser,