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

303

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

mlir/test/lit.cfg.py

131

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

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/IR/OperationSupport.h

public:
 
   /// Enable or disable printing of debug information (based on `enable`). If
   /// 'prettyForm' is set to true, debug information is printed in a more
   /// readable 'pretty' form. Note: The IR generated with 'prettyForm' is not
   /// parsable.
   OpPrintingFlags &enableDebugInfo(bool enable = true, bool prettyForm = false);
 
   /// Always print operations in the generic form.
-  OpPrintingFlags &printGenericOpForm();
+  OpPrintingFlags &printGenericOpForm(bool enable = true);
 
   /// Skip printing regions.
   OpPrintingFlags &skipRegions(bool skip = true);
 
   /// Do not verify the operation when using custom operation printers.
   OpPrintingFlags &assumeVerified();
 
   /// Use local scope when printing the operation. This allows for using the

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

public:
 
   /// Set whether to run the verifier after each transformation pass.
   MlirOptMainConfig &verifyPasses(bool verify) {
     verifyPassesFlag = verify;
     return *this;
   }
   bool shouldVerifyPasses() const { return verifyPassesFlag; }
 
+  /// Set whether to run the verifier after each transformation pass.
+  MlirOptMainConfig &verifyRoundtrip(bool verify) {
+    verifyRoundtripFlag = verify;
+    return *this;
+  }
+  bool shouldVerifyRoundtrip() const { return verifyRoundtripFlag; }
+
 protected:
   /// Allow operation with no registered dialects.
   /// This option is for convenience during testing only and discouraged in
   /// general.
   bool allowUnregisteredDialectsFlag = false;
 
   /// Configuration for the debugging hooks.
   tracing::DebugConfig debugConfig;
protected:
   bool useExplicitModuleFlag = false;
 
   /// Set whether to check that emitted diagnostics match `expected-*` lines on
   /// the corresponding line. This is meant for implementing diagnostic tests.
   bool verifyDiagnosticsFlag = false;
 
   /// Run the verifier after each transformation pass.
   bool verifyPassesFlag = true;
+
+  /// Verify that the input IR round-trips perfectly.
+  bool verifyRoundtripFlag = false;
 };
 
 /// This defines the function type used to setup the pass manager. This can be
 /// used to pass in a callback to setup a default pass pipeline to be applied on
 /// the loaded IR.
 using PassPipelineFn = llvm::function_ref<LogicalResult(PassManager &pm)>;
 
 /// Perform the core processing behind `mlir-opt`.

mlir/lib/IR/AsmPrinter.cpp

 OpPrintingFlags &OpPrintingFlags::enableDebugInfo(bool enable,
                                                   bool prettyForm) {
   printDebugInfoFlag = enable;
   printDebugInfoPrettyFormFlag = prettyForm;
   return *this;
 }
 
 /// Always print operations in the generic form.
-OpPrintingFlags &OpPrintingFlags::printGenericOpForm() {
-  printGenericOpFormFlag = true;
+OpPrintingFlags &OpPrintingFlags::printGenericOpForm(bool enable) {
+  printGenericOpFormFlag = enable;
   return *this;
 }
 
 /// Always skip Regions.
 OpPrintingFlags &OpPrintingFlags::skipRegions(bool skip) {
   skipRegionsFlag = skip;
   return *this;
 }

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

static cl::opt<bool, /*ExternalStorage=*/true> verifyDiagnostics(
                  "expected-* lines on the corresponding line"),
         cl::location(verifyDiagnosticsFlag), cl::init(false));
 
     static cl::opt<bool, /*ExternalStorage=*/true> verifyPasses(
         "verify-each",
         cl::desc("Run the verifier after each transformation pass"),
         cl::location(verifyPassesFlag), cl::init(true));
 
+    static cl::opt<bool, /*ExternalStorage=*/true> verifyRoundtrip(
+        "verify-roundtrip",
+        cl::desc("Round-trip the IR after parsing and ensure it succeeds"),
+        cl::location(verifyRoundtripFlag), cl::init(false));
+
     static cl::list<std::string> passPlugins(
         "load-pass-plugin", cl::desc("Load passes from plugin library"));
     /// Set the callback to load a pass plugin.
     passPlugins.setCallback([&](const std::string &pluginPath) {
       auto plugin = PassPlugin::load(pluginPath);
       if (!plugin) {
         errs() << "Failed to load passes from '" << pluginPath
                << "'. Request ignored.\n";
if (!plugin) {
       errs() << "Failed to load dialect plugin from '" << pluginPath
              << "'. Request ignored.\n";
       return;
     };
     plugin.get().registerDialectRegistryCallbacks(registry);
   });
 }
 
+LogicalResult loadIRDLDialects(StringRef irdlFile, MLIRContext &ctx) {
+  DialectRegistry registry;
+  registry.insert<irdl::IRDLDialect>();
+  ctx.appendDialectRegistry(registry);
+
+  // Set up the input file.
+  std::string errorMessage;
+  std::unique_ptr<MemoryBuffer> file = openInputFile(irdlFile, &errorMessage);
+  if (!file) {
+    emitError(UnknownLoc::get(&ctx)) << errorMessage;
+    return failure();
+  }
+
+  // Give the buffer to the source manager.
+  // This will be picked up by the parser.
+  SourceMgr sourceMgr;
+  sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
+
+  SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &ctx);
+
+  // Parse the input file.
+  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, &ctx));
+
+  // Load IRDL dialects.
+  return irdl::loadDialects(module.get());
+}
+
+// Return success if the module can correctly round-trip. This intended to test
+// that the custom printers/parsers are complete.
+static LogicalResult doVerifyRoundTrip(Operation *op,
+                                       const MlirOptMainConfig &config,
+                                       bool useBytecode) {
+  // We use a new context to avoid resource handle renaming issue in the diff.
+  MLIRContext roundtripContext;
+  OwningOpRef<Operation *> roundtripModule;
+  roundtripContext.appendDialectRegistry(
+      op->getContext()->getDialectRegistry());
+  if (op->getContext()->allowsUnregisteredDialects())
+    roundtripContext.allowUnregisteredDialects();
+  StringRef irdlFile = config.getIrdlFile();
+  if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, roundtripContext)))
+    return failure();
+
+  // Print a first time with custom format (or bytecode) and parse it back to
+  // the roundtripModule.
+  {
+    std::string buffer;
+    llvm::raw_string_ostream ostream(buffer);
+    if (useBytecode) {
+      if (failed(writeBytecodeToFile(op, ostream))) {
+        op->emitOpError() << "failed to write bytecode, cannot verify round-trip.\n";
+        return failure();
+      }
+    } else {
+      op->print(ostream,
+                OpPrintingFlags().printGenericOpForm(false).enableDebugInfo());
+    }
+    FallbackAsmResourceMap fallbackResourceMap;
+    ParserConfig parseConfig(&roundtripContext, /*verifyAfterParse=*/true,
+                             &fallbackResourceMap);
+    roundtripModule =
+        parseSourceString<Operation *>(ostream.str(), parseConfig);
+    if (!roundtripModule) {
+      op->emitOpError() << "failed to parse bytecode back, cannot verify round-trip.\n";
+      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);
+    op->print(ostreamref,
+              OpPrintingFlags().printGenericOpForm().enableDebugInfo());
+    llvm::raw_string_ostream ostreamrndtrip(roundtrip);
+    roundtripModule.get()->print(
+        ostreamrndtrip,
+        OpPrintingFlags().printGenericOpForm().enableDebugInfo());
+  }
+  if (reference != roundtrip) {
+    // TODO implement a diff.
+    return op->emitOpError() << "roundTrip testing failed!\n<<<<<<\n"
+                             << reference << "\n>>>>>\n"
+                             << roundtrip;
+  }
+
+  return success();
+}
+
+static LogicalResult doVerifyRoundTrip(Operation *op,
+                                       const MlirOptMainConfig &config) {
+  // Textual round-trip isn't fully robust at the moment (for example implicit
+  // terminator are losing location informations).
+
+  return doVerifyRoundTrip(op, config, /*useBytecode=*/true);
+}
+
 /// 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,
ParserConfig parseConfig(context, /*verifyAfterParse=*/true,
                            &fallbackResourceMap);
   if (config.shouldRunReproducer())
     reproOptions.attachResourceParser(parseConfig);
 
   // Parse the input file and reset the context threading state.
   TimingScope parserTiming = timing.nest("Parser");
   OwningOpRef<Operation *> op = parseSourceFileForTool(
       sourceMgr, parseConfig, !config.shouldUseExplicitModule());
-  context->enableMultithreading(wasThreadingEnabled);
+  parserTiming.stop();
   if (!op)
     return failure();
-  parserTiming.stop();
+
+  // Perform round-trip verification if requested
+  if (config.shouldVerifyRoundtrip() &&
+      failed(doVerifyRoundTrip(op.get(), config)))
+    return failure();
+
+  context->enableMultithreading(wasThreadingEnabled);
 
   // Prepare the pass manager, applying command-line and reproducer options.
   PassManager pm(op.get()->getName(), PassManager::Nesting::Implicit);
   pm.enableVerifier(config.shouldVerifyPasses());
   if (failed(applyPassManagerCLOptions(pm)))
     return failure();
   pm.enableTiming(timing);
   if (config.shouldRunReproducer() && failed(reproOptions.apply(pm)))
return emitError(UnknownLoc::get(pm.getContext()))
            << "bytecode version while not emitting bytecode";
   AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,
                     &fallbackResourceMap);
   op.get()->print(os, asmState);
   os << '\n';
   return success();
 }
 
-LogicalResult loadIRDLDialects(StringRef irdlFile, MLIRContext &ctx) {
-  DialectRegistry registry;
-  registry.insert<irdl::IRDLDialect>();
-  ctx.appendDialectRegistry(registry);
-
-  // Set up the input file.
-  std::string errorMessage;
-  std::unique_ptr<MemoryBuffer> file = openInputFile(irdlFile, &errorMessage);
-  if (!file) {
-    emitError(UnknownLoc::get(&ctx)) << errorMessage;
-    return failure();
-  }
-
-  // Give the buffer to the source manager.
-  // This will be picked up by the parser.
-  SourceMgr sourceMgr;
-  sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
-
-  SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &ctx);
-
-  // Parse the input file.
-  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, &ctx));
-
-  // Load IRDL dialects.
-  return irdl::loadDialects(module.get());
-}
-
 /// 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,
                                    const MlirOptMainConfig &config,
                                    DialectRegistry &registry,
                                    llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.

mlir/test/lit.cfg.py

config.excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt',
                    'lit.cfg.py', 'lit.site.cfg.py']
 
 # Tweak the PATH to include the tools dir.
 llvm_config.with_environment('PATH', config.mlir_tools_dir, append_path=True)
 llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
 
 tool_dirs = [config.mlir_tools_dir, config.llvm_tools_dir]
 tools = [
-    'mlir-opt',
     'mlir-tblgen',
     'mlir-translate',
     'mlir-lsp-server',
     'mlir-capi-execution-engine-test',
     'mlir-capi-ir-test',
     'mlir-capi-llvm-test',
     'mlir-capi-pass-test',
     'mlir-capi-pdl-test',

 # On Windows the path to python could contains spaces in which case it needs to be provided in quotes.
 # This is the equivalent of how %python is setup in llvm/utils/lit/lit/llvm/config.py.
 elif "Windows" in config.host_os:
   python_executable = '"%s"' % (python_executable)
 tools.extend([
   ToolSubst('%PYTHON', python_executable, unresolved='ignore'),
 ])
 
+if "MLIR_OPT_CHECK_IR_ROUNDTRIP" in os.environ:
+  tools.extend([
+    ToolSubst('mlir-opt', 'mlir-opt --verify-roundtrip', unresolved='fatal'),
+  ])
+else:
+  tools.extend([
+    'mlir-opt',
+  ])
+
 llvm_config.add_tool_substitutions(tools, tool_dirs)
 
 
 # FileCheck -enable-var-scope is enabled by default in MLIR test
 # This option avoids to accidentally reuse variable across -LABEL match,
 # it can be explicitly opted-in by prefixing the variable name with $
 config.environment['FILECHECK_OPTS'] = "-enable-var-scope --allow-unused-prefixes=false"