[mlir] Add fallback support for parsing/printing unknown external resources

[rriddle]

Summary

This is necessary/useful for building generic tooling that can roundtrip external
resources without needing to explicitly handle them. For example, this allows
for viewing the resources encoded within a bytecode file without having to
explicitly know how to process them (e.g. making it easier to interact with a
reproducer encoded in bytecode).

[jpienaar]

Nice, thanks

mlir/lib/IR/AsmPrinter.cpp

1311

We could now use c++17 if initializer form:

if (FailureOr<AsmResourceBlob> blob = entry.parseAsBlob(); succeeded(blob)) 
  resources.emplace_back(entry.getKey(), std::move(*blob));
else return failure();
break;

(unsure if more readable yet, so just an aside)

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

1303

And I think per discussion this may be WAI.

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

69

I'd expand on passthroughs here to make behavior clearer.

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1311

I'm not sure it adds much here (it's the same amount of lines), but thanks for the reminder. I think this will come in handy a lot moving forward.

[bondhugula]

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

1327–1329

Just update the ctor to optionally take this resource map and avoid the attach call here and in MlirOptMain.cpp?

mlir/include/mlir/Bytecode/BytecodeWriter.h

public:
   template <typename CallableT>
   std::enable_if_t<std::is_convertible<
       CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value>
   attachResourcePrinter(StringRef name, CallableT &&printFn) {
     attachResourcePrinter(AsmResourcePrinter::fromCallable(
         name, std::forward<CallableT>(printFn)));
   }
 
+  /// Attach resource printers to the AsmState for the fallback resources
+  /// in the given map.
+  void attachFallbackResourcePrinter(FallbackAsmResourceMap &map) {
+    for (auto &printer : map.getPrinters())
+      attachResourcePrinter(std::move(printer));
+  }
+
 private:
   /// A pointer to allocated storage for the impl state.
   std::unique_ptr<Impl> impl;
 };
 
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/AsmState.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ASMSTATE_H_
 #define MLIR_IR_ASMSTATE_H_
 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/StringMap.h"
 
 #include <memory>
 
 namespace mlir {
 class AsmResourcePrinter;
 class AsmDialectResourceHandle;
 class Operation;
 
static std::unique_ptr<AsmResourcePrinter> fromCallable(StringRef name,
     };
     return std::make_unique<Printer>(name, std::forward<CallableT>(printFn));
   }
 
 private:
   std::string name;
 };
 
+/// A fallback map containing external resources not explicitly handled by
+/// another parser/printer.
+class FallbackAsmResourceMap {
+public:
+  /// This class represents an opaque resource.
+  struct OpaqueAsmResource {
+    OpaqueAsmResource(StringRef key,
+                      std::variant<AsmResourceBlob, bool, std::string> value)
+        : key(key.str()), value(std::move(value)) {}
+
+    /// The key identifying the resource.
+    std::string key;
+    /// An opaque value for the resource, whose variant values align 1-1 with
+    /// the kinds defined in AsmResourceEntryKind.
+    std::variant<AsmResourceBlob, bool, std::string> value;
+  };
+
+  /// Return a parser than can be used for parsing entries for the given
+  /// identifier key.
+  AsmResourceParser &getParserFor(StringRef key);
+
+  /// Build a set of resource printers to print the resources within this map.
+  std::vector<std::unique_ptr<AsmResourcePrinter>> getPrinters();
+
+private:
+  struct ResourceCollection : public AsmResourceParser {
+    ResourceCollection(StringRef name) : AsmResourceParser(name) {}
+
+    /// Parse a resource into this collection.
+    LogicalResult parseResource(AsmParsedResourceEntry &entry) final;
+
+    /// Build the resources held by this collection.
+    void buildResources(Operation *op, AsmResourceBuilder &builder) const;
+
+    /// The set of resources parsed into this collection.
+    SmallVector<OpaqueAsmResource> resources;
+  };
+
+  /// The set of opaque resources.
+  llvm::MapVector<std::string, std::unique_ptr<ResourceCollection>,
+                  llvm::StringMap<unsigned>>
+      keyToResources;
+};
+
 //===----------------------------------------------------------------------===//
 // ParserConfig
 //===----------------------------------------------------------------------===//
 
 /// This class represents a configuration for the MLIR assembly parser. It
 /// contains all of the necessary state to parse a MLIR source file.
 class ParserConfig {
 public:
-  ParserConfig(MLIRContext *context) : context(context) {
+  /// Construct a parser configuration with the given context.
+  /// `fallbackResourceMap` is an optional fallback handler that can be used to
+  /// parse external resources not explicitly handled by another parser.
+  ParserConfig(MLIRContext *context,
+               FallbackAsmResourceMap *fallbackResourceMap = nullptr)
+      : context(context), fallbackResourceMap(fallbackResourceMap) {
     assert(context && "expected valid MLIR context");
   }
 
   /// Return the MLIRContext to be used when parsing.
   MLIRContext *getContext() const { return context; }
 
   /// Return the resource parser registered to the given name, or nullptr if no
   /// parser with `name` is registered.
   AsmResourceParser *getResourceParser(StringRef name) const {
     auto it = resourceParsers.find(name);
-    return it == resourceParsers.end() ? nullptr : it->second.get();
+    if (it != resourceParsers.end())
+      return it->second.get();
+    if (fallbackResourceMap)
+      return &fallbackResourceMap->getParserFor(name);
+    return nullptr;
   }
 
   /// Attach the given resource parser.
   void attachResourceParser(std::unique_ptr<AsmResourceParser> parser) {
     StringRef name = parser->getName();
     auto it = resourceParsers.try_emplace(name, std::move(parser));
     (void)it;
     assert(it.second &&
            "resource parser already registered with the given name");
   }
 
   /// Attach the given callable resource parser with the given name.
   template <typename CallableT>
   std::enable_if_t<std::is_convertible<
       CallableT, function_ref<LogicalResult(AsmParsedResourceEntry &)>>::value>
   attachResourceParser(StringRef name, CallableT &&parserFn) {
     attachResourceParser(AsmResourceParser::fromCallable(
         name, std::forward<CallableT>(parserFn)));
   }
 
 private:
   MLIRContext *context;
   DenseMap<StringRef, std::unique_ptr<AsmResourceParser>> resourceParsers;
+  FallbackAsmResourceMap *fallbackResourceMap;
 };
 
 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 
 /// This class provides management for the lifetime of the state used when
 /// printing the IR. It allows for alleviating the cost of recomputing the
public:
   template <typename CallableT>
   std::enable_if_t<std::is_convertible<
       CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value>
   attachResourcePrinter(StringRef name, CallableT &&printFn) {
     attachResourcePrinter(AsmResourcePrinter::fromCallable(
         name, std::forward<CallableT>(printFn)));
   }
 
+  /// Attach resource printers to the AsmState for the fallback resources
+  /// in the given map.
+  void attachFallbackResourcePrinter(FallbackAsmResourceMap &map) {
+    for (auto &printer : map.getPrinters())
+      attachResourcePrinter(std::move(printer));
+  }
+
   /// Returns a map of dialect resources that were referenced when using this
   /// state to print IR.
   DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> &
   getDialectResources() const;
 
 private:
   AsmState() = delete;
 

mlir/lib/IR/AsmPrinter.cpp

StringRef mlir::toString(AsmResourceEntryKind kind) {
   case AsmResourceEntryKind::Bool:
     return "bool";
   case AsmResourceEntryKind::String:
     return "string";
   }
   llvm_unreachable("unknown AsmResourceEntryKind");
 }
 
+AsmResourceParser &FallbackAsmResourceMap::getParserFor(StringRef key) {
+  std::unique_ptr<ResourceCollection> &collection = keyToResources[key.str()];
+  if (!collection)
+    collection = std::make_unique<ResourceCollection>(key);
+  return *collection;
+}
+
+std::vector<std::unique_ptr<AsmResourcePrinter>>
+FallbackAsmResourceMap::getPrinters() {
+  std::vector<std::unique_ptr<AsmResourcePrinter>> printers;
+  for (auto &it : keyToResources) {
+    ResourceCollection *collection = it.second.get();
+    auto buildValues = [=](Operation *op, AsmResourceBuilder &builder) {
+      return collection->buildResources(op, builder);
+    };
+    printers.emplace_back(
+        AsmResourcePrinter::fromCallable(collection->getName(), buildValues));
+  }
+  return printers;
+}
+
+LogicalResult FallbackAsmResourceMap::ResourceCollection::parseResource(
+    AsmParsedResourceEntry &entry) {
+  switch (entry.getKind()) {
+  case AsmResourceEntryKind::Blob: {
+    FailureOr<AsmResourceBlob> blob = entry.parseAsBlob();
+    if (failed(blob))
+      return failure();
+    resources.emplace_back(entry.getKey(), std::move(*blob));
+    break;
+  }
+  case AsmResourceEntryKind::Bool: {
+    FailureOr<bool> value = entry.parseAsBool();
+    if (failed(value))
+      return failure();
+    resources.emplace_back(entry.getKey(), *value);
+    break;
+  }
+  case AsmResourceEntryKind::String: {
+    FailureOr<std::string> str = entry.parseAsString();
+    if (failed(str))
+      return failure();
+    resources.emplace_back(entry.getKey(), std::move(*str));
+    break;
+  }
+  }
+  return success();
+}
+
+void FallbackAsmResourceMap::ResourceCollection::buildResources(
+    Operation *op, AsmResourceBuilder &builder) const {
+  for (const auto &entry : resources) {
+    if (const auto *value = std::get_if<AsmResourceBlob>(&entry.value))
+      builder.buildBlob(entry.key, *value);
+    else if (const auto *value = std::get_if<bool>(&entry.value))
+      builder.buildBool(entry.key, *value);
+    else if (const auto *value = std::get_if<std::string>(&entry.value))
+      builder.buildString(entry.key, *value);
+    else
+      llvm_unreachable("unknown AsmResourceEntryKind");
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 class AsmStateImpl {
 public:
void Value::printAsOperand(raw_ostream &os, AsmState &state) {
   // Currently, region arguments can be shadowed when printing the main
   // operation. If the IR hasn't been printed, this will produce the old SSA
   // name and not the shadowed name.
   state.getImpl().getSSANameState().printValueID(*this, /*printResultNo=*/true,
                                                  os);
 }
 
 void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
-  // If this is a top level operation, we also print aliases.
-  if (!getParent() && !printerFlags.shouldUseLocalScope()) {
-    AsmState state(this, printerFlags);
-    state.getImpl().initializeAliases(this);
-    print(os, state);
-    return;
-  }
-
   // Find the operation to number from based upon the provided flags.
   Operation *op = this;
   bool shouldUseLocalScope = printerFlags.shouldUseLocalScope();
   do {
     // If we are printing local scope, stop at the first operation that is
     // isolated from above.
     if (shouldUseLocalScope && op->hasTrait<OpTrait::IsIsolatedFromAbove>())
       break;
 
     // Otherwise, traverse up to the next parent.
     Operation *parentOp = op->getParentOp();
     if (!parentOp)
       break;
     op = parentOp;
   } while (true);
 
   AsmState state(op, printerFlags);
   print(os, state);
 }
 void Operation::print(raw_ostream &os, AsmState &state) {
   OperationPrinter printer(os, state.getImpl());
-  if (!getParent() && !state.getPrinterFlags().shouldUseLocalScope())
+  if (!getParent() && !state.getPrinterFlags().shouldUseLocalScope()) {
+    state.getImpl().initializeAliases(this);
     printer.printTopLevelOperation(this);
-  else
+  } else {
     printer.print(this);
   }
+}
 
 void Operation::dump() {
   print(llvm::errs(), OpPrintingFlags().useLocalScope());
   llvm::errs() << "\n";
 }
 
 void Block::print(raw_ostream &os) {
   Operation *parentOp = getParentOp();

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

struct MLIRDocument {
 
   /// The high level parser state used to find definitions and references within
   /// the source file.
   AsmParserState asmState;
 
   /// The container for the IR parsed from the input file.
   Block parsedIR;
 
+  /// A collection of external resources, which we want to propagate up to the
+  /// user.
+  FallbackAsmResourceMap fallbackResourceMap;
+
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
 };
 } // namespace
 
 MLIRDocument::MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
                            StringRef contents,
                            std::vector<lsp::Diagnostic> &diagnostics) {
   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
     diagnostics.push_back(getLspDiagnoticFromDiag(sourceMgr, diag, uri));
   });
 
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
     return;
   }
 
+  ParserConfig config(&context, &fallbackResourceMap);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
-  if (failed(parseAsmSourceFile(sourceMgr, &parsedIR, &context, &asmState))) {
+  if (failed(parseAsmSourceFile(sourceMgr, &parsedIR, config, &asmState))) {
     // If parsing failed, clear out any of the current state.
     parsedIR.clear();
     asmState = AsmParserState();
+    fallbackResourceMap = FallbackAsmResourceMap();
     return;
   }
 }
 
 //===----------------------------------------------------------------------===//
 // MLIRDocument: Definitions and References
 //===----------------------------------------------------------------------===//
 
if (parsedIR.empty()) {
           lsp::ErrorCode::RequestFailed);
     }
     return llvm::make_error<lsp::LSPError>(
         "expected a single top-level operation", lsp::ErrorCode::RequestFailed);
   }
 
   lsp::MLIRConvertBytecodeResult result;
   {
+    BytecodeWriterConfig writerConfig;
+    writerConfig.attachFallbackResourcePrinter(fallbackResourceMap);
+
     std::string rawBytecodeBuffer;
     llvm::raw_string_ostream os(rawBytecodeBuffer);
-    writeBytecodeToFile(&parsedIR.front(), os);
+    writeBytecodeToFile(&parsedIR.front(), os, writerConfig);
     result.output = llvm::encodeBase64(rawBytecodeBuffer);
   }
   return result;
 }
 
 //===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
lsp::MLIRServer::convertFromBytecode(const URIForFile &uri) {
   tempContext.allowUnregisteredDialects();
 
   // Collect any errors during parsing.
   std::string errorMsg;
   ScopedDiagnosticHandler diagHandler(
       &tempContext,
       [&](mlir::Diagnostic &diag) { errorMsg += diag.str() + "\n"; });
 
+  // Handling for external resources, which we want to propagate up to the user.
+  FallbackAsmResourceMap fallbackResourceMap;
+
+  // Setup the parser config.
+  ParserConfig parserConfig(&tempContext, &fallbackResourceMap);
+
   // Try to parse the given source file.
-  // TODO: This won't preserve external resources or the producer, we should try
-  // to fix this.
+  // TODO: This won't preserve the producer, we should try to fix this.
   Block parsedBlock;
-  if (failed(parseSourceFile(uri.file(), &parsedBlock, &tempContext))) {
+  if (failed(parseSourceFile(uri.file(), &parsedBlock, parserConfig))) {
     return llvm::make_error<lsp::LSPError>(
         "failed to parse bytecode source file: " + errorMsg,
         lsp::ErrorCode::RequestFailed);
   }
 
   // TODO: We currently expect a single top-level operation, but this could
   // conceptually be relaxed.
   if (!llvm::hasSingleElement(parsedBlock)) {
     return llvm::make_error<lsp::LSPError>(
         "expected bytecode to contain a single top-level operation",
         lsp::ErrorCode::RequestFailed);
   }
 
   // Print the module to a buffer.
   lsp::MLIRConvertBytecodeResult result;
   {
     // Extract the top-level op so that aliases get printed.
     // FIXME: We should be able to enable aliases without having to do this!
     OwningOpRef<Operation *> topOp = &parsedBlock.front();
     (*topOp)->remove();
 
+    AsmState state(*topOp,
+                   OpPrintingFlags().enableDebugInfo().assumeVerified());
+    state.attachFallbackResourcePrinter(fallbackResourceMap);
+
     llvm::raw_string_ostream os(result.output);
-    (*topOp)->print(os, OpPrintingFlags().enableDebugInfo().assumeVerified());
+    (*topOp)->print(os, state);
   }
   return std::move(result);
 }
 
 llvm::Expected<lsp::MLIRConvertBytecodeResult>
 lsp::MLIRServer::convertToBytecode(const URIForFile &uri) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt == impl->files.end()) {
     return llvm::make_error<lsp::LSPError>(
         "language server does not contain an entry for this source file",
         lsp::ErrorCode::RequestFailed);
   }
   return fileIt->second->convertToBytecode();
 }

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

static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
 
   // Prepare the pass manager and apply any command line options.
   PassManager pm(context, OpPassManager::Nesting::Implicit);
   pm.enableVerifier(verifyPasses);
   applyPassManagerCLOptions(pm);
   pm.enableTiming(timing);
 
   // Prepare the parser config, and attach any useful/necessary resource
-  // handlers.
-  ParserConfig config(context);
+  // handlers. Unhandled external resources are treated as passthrough.
+  FallbackAsmResourceMap fallbackResourceMap;
+  ParserConfig config(context, &fallbackResourceMap);
   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));
   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);
+    BytecodeWriterConfig writerConfig;
+    writerConfig.attachFallbackResourcePrinter(fallbackResourceMap);
+
+    writeBytecodeToFile(module->getOperation(), os, writerConfig);
   } else {
-    module->print(os);
+    AsmState asmState(*module);
+    asmState.attachFallbackResourcePrinter(fallbackResourceMap);
+
+    module->print(os, asmState);
     os << '\n';
   }
   return success();
 }
 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static LogicalResult

mlir/test/IR/file-metadata-resources.mlir

 // RUN: mlir-opt %s -split-input-file | FileCheck %s
 
 // Check that we only preserve the blob that got referenced.
 // CHECK:      test: {
 // CHECK-NEXT:   blob1: "0x08000000010000000000000002000000000000000300000000000000"
 // CHECK-NEXT: }
 
+// Check that we properly preserve unknown external resources.
+// CHECK:      external: {
+// CHECK-NEXT:   blob: "0x08000000010000000000000002000000000000000300000000000000"
+// CHECK-NEXT:   bool: true
+// CHECK-NEXT:   string: "string"
+// CHECK-NEXT: }
+
 module attributes { test.blob_ref = #test.e1di64_elements<blob1> : tensor<*xi1>} {}
 
 {-#
   dialect_resources: {
     test: {
       blob1: "0x08000000010000000000000002000000000000000300000000000000",
       blob2: "0x08000000040000000000000005000000000000000600000000000000"
     }
+  },
+  external_resources: {
+    external: {
+      blob: "0x08000000010000000000000002000000000000000300000000000000",
+      bool: true,
+      string: "string"
+    }
   }
 #-}

mlir/test/IR/invalid-file-metadata.mlir

 // -----
 
 // expected-error@+3 {{expected ':'}}
 {-#
   external_resources: {
     entry "value"
   }
 #-}
-
-// -----
-
-// expected-warning@+3 {{ignoring unknown external resources for 'foobar'}}
-{-#
-  external_resources: {
-    foobar: {
-      entry: "foo"
-    }
-  }
-#-}