[mlir] Change translation API to allow custom top-level operations.

[stellaraccident]

Summary

• Changes hooks to operate on Operation* instead of ModuleOp.
• Adds a TranslateFromMLIRRegistration optional parserFunction argument allowing the parser to be configured (i.e. to build a custom top-level operation).
• This allowed the removal of some top-level op munging in the SPIRV translations (verified by tests that required no modification).
• As discussed in https://llvm.discourse.group/t/de-privileging-moduleop-in-translation-apis/3733

[rriddle]

mlir/include/mlir/Translation.h

80

It would be nice if users could pass in a translate function that took Op Foo, and have it automatically do the type checking.

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

29

nit: Drop trivial braces.

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

131–132

Can you add a constructor for OwningOpRef<Operation *> that allows for taking in OwningOpRef<Foo> &&? (You should be able to write it just using std::is_convertible<OpTy, Foo>). That would remove all of these weird conversion dances.

mlir/include/mlir/Translation.h

 class StringRef;
 } // namespace llvm
 
 namespace mlir {
 class DialectRegistry;
 struct LogicalResult;
 class MLIRContext;
 class ModuleOp;
-class OwningModuleRef;
+class Operation;
+template <typename OpTy>
+class OwningOpRef;
 
 /// Interface of the function that translates the sources managed by `sourceMgr`
 /// to MLIR. The source manager has at least one buffer. The implementation
-/// should create a new MLIR ModuleOp in the given context and return a pointer
-/// to it, or a nullptr in case of any error.
-using TranslateSourceMgrToMLIRFunction =
-    std::function<OwningModuleRef(llvm::SourceMgr &sourceMgr, MLIRContext *)>;
+/// should create a new MLIR root operation in the given context and return a
+/// pointer to it, or a nullptr in case of any error.
+using TranslateSourceMgrToMLIRFunction = std::function<OwningOpRef<Operation *>(
+    llvm::SourceMgr &sourceMgr, MLIRContext *)>;
 
 /// Interface of the function that translates the given string to MLIR. The
-/// implementation should create a new MLIR ModuleOp in the given context. If
-/// source-related error reporting is required from within the function, use
+/// implementation should create a new MLIR root operation in the given context.
+/// If source-related error reporting is required from within the function, use
 /// TranslateSourceMgrToMLIRFunction instead.
 using TranslateStringRefToMLIRFunction =
-    std::function<OwningModuleRef(llvm::StringRef, MLIRContext *)>;
+    std::function<OwningOpRef<Operation *>(llvm::StringRef, MLIRContext *)>;
 
 /// Interface of the function that translates MLIR to a different format and
 /// outputs the result to a stream. It is allowed to modify the module.
 using TranslateFromMLIRFunction =
-    std::function<LogicalResult(ModuleOp, llvm::raw_ostream &output)>;
+    std::function<LogicalResult(Operation *, llvm::raw_ostream &output)>;
 
 /// Interface of the function that performs file-to-file translation involving
 /// MLIR. The input file is held in the given MemoryBuffer; the output file
 /// should be written to the given raw_ostream. The implementation should create
 /// all MLIR constructs needed during the process inside the given context. This
 /// can be used for round-tripping external formats through the MLIR system.
 using TranslateFunction = std::function<LogicalResult(
     llvm::SourceMgr &sourceMgr, llvm::raw_ostream &output, MLIRContext *)>;
 
 /// Use Translate[ToMLIR|FromMLIR]Registration as an initializer that
 /// registers a function and associates it with name. This requires that a
 /// translation has not been registered to a given name.
 ///
 /// Usage:
 ///
 ///   // At file scope.
 ///   namespace mlir {
 ///   void registerTRexToMLIRRegistration() {
 ///     TranslateToMLIRRegistration Unused(&MySubCommand, [] { ... });
 ///   }
 ///   } // namespace mlir
 ///
+/// Note that the TranslateFROMMLIRRegistration usage has additional optional
+/// arguments that are expected to be useful in special situations:
+///   - dialectRegistration: A callback function to register dialects.
+///   - parseFunction: A callback to perform parsing (to be used if the top
+///     level module is not the builtin ModuleOp or special parser setup is
+///     required).
 /// \{
 struct TranslateToMLIRRegistration {
   TranslateToMLIRRegistration(llvm::StringRef name,
                               const TranslateSourceMgrToMLIRFunction &function);
   TranslateToMLIRRegistration(llvm::StringRef name,
                               const TranslateStringRefToMLIRFunction &function);
 };
 
 struct TranslateFromMLIRRegistration {
+  using ParseFunction = OwningOpRef<Operation *>(llvm::SourceMgr &sourceMgr,
+                                                 MLIRContext *context);
+  static OwningOpRef<Operation *> parseBuiltinModule(llvm::SourceMgr &sourceMgr,
+                                                     MLIRContext *context);
+
   TranslateFromMLIRRegistration(
       llvm::StringRef name, const TranslateFromMLIRFunction &function,
       std::function<void(DialectRegistry &)> dialectRegistration =
-          [](DialectRegistry &) {});
+          [](DialectRegistry &) {},
+      ParseFunction parseFunction =
+          TranslateFromMLIRRegistration::parseBuiltinModule);
 };
 struct TranslateRegistration {
   TranslateRegistration(llvm::StringRef name,
                         const TranslateFunction &function);
 };
 /// \}
 
 /// A command line parser for translation functions.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

if (failed(deserializer.processFunction(&f)))
       return {};
   }
 
   return module;
 }
 
 // Deserializes the LLVM bitcode stored in `input` into an MLIR module in the
 // LLVM dialect.
-OwningModuleRef translateLLVMIRToModule(llvm::SourceMgr &sourceMgr,
+OwningOpRef<Operation *> translateLLVMIRToModule(llvm::SourceMgr &sourceMgr,
                                                  MLIRContext *context) {
   llvm::SMDiagnostic err;
   llvm::LLVMContext llvmContext;
   std::unique_ptr<llvm::Module> llvmModule = llvm::parseIR(
       *sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID()), err, llvmContext);
   if (!llvmModule) {
     std::string errStr;
     llvm::raw_string_ostream errStream(errStr);
     err.print(/*ProgName=*/"", errStream);
     emitError(UnknownLoc::get(context)) << errStream.str();
     return {};
   }
-  return translateLLVMIRToModule(std::move(llvmModule), context);
+  return OwningOpRef<Operation *>(
+      translateLLVMIRToModule(std::move(llvmModule), context).release());
 }
 
 namespace mlir {
 void registerFromLLVMIRTranslation() {
   TranslateToMLIRRegistration fromLLVM(
       "import-llvm", [](llvm::SourceMgr &sourceMgr, MLIRContext *context) {
         return ::translateLLVMIRToModule(sourceMgr, context);
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

 #include "llvm/IR/Module.h"
 
 using namespace mlir;
 
 namespace mlir {
 void registerToLLVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
       "mlir-to-llvmir",
-      [](ModuleOp module, raw_ostream &output) {
+      [](Operation *root, raw_ostream &output) -> LogicalResult {
+        auto module = dyn_cast<ModuleOp>(root);
+        if (!module) {
+          return emitError(root->getLoc()) << "expected module op";
+        }
         llvm::LLVMContext llvmContext;
         auto llvmModule = translateModuleToLLVMIR(module, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       },
       [](DialectRegistry &registry) {
         registerAllToLLVMIRTranslations(registry);
       });
 }
 } // namespace mlir

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Deserialization registration
 //===----------------------------------------------------------------------===//
 
 // Deserializes the SPIR-V binary module stored in the file named as
 // `inputFilename` and returns a module containing the SPIR-V module.
-static OwningModuleRef deserializeModule(const llvm::MemoryBuffer *input,
-                                         MLIRContext *context) {
+static OwningOpRef<Operation *>
+deserializeModule(const llvm::MemoryBuffer *input, MLIRContext *context) {
   context->loadDialect<spirv::SPIRVDialect>();
 
   // Make sure the input stream can be treated as a stream of SPIR-V words
   auto start = input->getBufferStart();
   auto size = input->getBufferSize();
   if (size % sizeof(uint32_t) != 0) {
     emitError(UnknownLoc::get(context))
         << "SPIR-V binary module must contain integral number of 32-bit words";
     return {};
   }
 
   auto binary = llvm::makeArrayRef(reinterpret_cast<const uint32_t *>(start),
                                    size / sizeof(uint32_t));
 
   OwningOpRef<spirv::ModuleOp> spirvModule =
       spirv::deserialize(binary, context);
-  if (!spirvModule)
-    return {};
-
-  OwningModuleRef module(ModuleOp::create(FileLineColLoc::get(
-      context, input->getBufferIdentifier(), /*line=*/0, /*column=*/0)));
-  module->getBody()->push_front(spirvModule.release());
-
-  return module;
+  return OwningOpRef<Operation *>(spirvModule.release());
 }
 
 namespace mlir {
 void registerFromSPIRVTranslation() {
   TranslateToMLIRRegistration fromBinary(
       "deserialize-spirv",
       [](llvm::SourceMgr &sourceMgr, MLIRContext *context) {
         assert(sourceMgr.getNumBuffers() == 1 && "expected one buffer");
         return deserializeModule(
             sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID()), context);
       });
 }
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Serialization registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult serializeModule(ModuleOp module, raw_ostream &output) {
-  if (!module)
+static LogicalResult serializeModule(Operation *root, raw_ostream &output) {
+  if (!root)
     return failure();
 
+  auto spirvModule = dyn_cast<spirv::ModuleOp>(root);
+  if (!spirvModule)
+    return root->emitError("expected a spirv.module");
   SmallVector<uint32_t, 0> binary;
-
-  SmallVector<spirv::ModuleOp, 1> spirvModules;
-  module.walk([&](spirv::ModuleOp op) { spirvModules.push_back(op); });
-
-  if (spirvModules.empty())
-    return module.emitError("found no 'spv.module' op");
-
-  if (spirvModules.size() != 1)
-    return module.emitError("found more than one 'spv.module' op");
-
-  if (failed(
-          spirv::serialize(spirvModules[0], binary, /*emitDebuginfo=*/false)))
+  if (failed(spirv::serialize(spirvModule, binary, /*emitDebuginfo=*/false)))
     return failure();
 
   output.write(reinterpret_cast<char *>(binary.data()),
                binary.size() * sizeof(uint32_t));
 
   return mlir::success();
 }
 
 namespace mlir {
 void registerToSPIRVTranslation() {
   TranslateFromMLIRRegistration toBinary(
-      "serialize-spirv",
-      [](ModuleOp module, raw_ostream &output) {
-        return serializeModule(module, output);
-      },
-      [](DialectRegistry &registry) {
+      "serialize-spirv", serializeModule, [](DialectRegistry &registry) {
         registry.insert<spirv::SPIRVDialect>();
       });
 }
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Round-trip registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult roundTripModule(ModuleOp srcModule, bool emitDebugInfo,
+static LogicalResult roundTripModule(Operation *srcRoot, bool emitDebugInfo,
                                      raw_ostream &output) {
   SmallVector<uint32_t, 0> binary;
-  MLIRContext *context = srcModule.getContext();
-  auto spirvModules = srcModule.getOps<spirv::ModuleOp>();
+  MLIRContext *context = srcRoot->getContext();
+  auto srcSpirvModule = dyn_cast<spirv::ModuleOp>(srcRoot);
+  if (!srcSpirvModule)
+    return srcRoot->emitError("expected spirv.module op");
 
-  if (spirvModules.begin() == spirvModules.end())
-    return srcModule.emitError("found no 'spv.module' op");
-
-  if (std::next(spirvModules.begin()) != spirvModules.end())
-    return srcModule.emitError("found more than one 'spv.module' op");
-
-  if (failed(spirv::serialize(*spirvModules.begin(), binary, emitDebugInfo)))
+  if (failed(spirv::serialize(srcSpirvModule, binary, emitDebugInfo)))
     return failure();
 
   MLIRContext deserializationContext(context->getDialectRegistry());
   // TODO: we should only load the required dialects instead of all dialects.
   deserializationContext.loadAllAvailableDialects();
   // Then deserialize to get back a SPIR-V module.
-  OwningOpRef<spirv::ModuleOp> spirvModule =
+  OwningOpRef<spirv::ModuleOp> destSpirvModule =
       spirv::deserialize(binary, &deserializationContext);
-  if (!spirvModule)
+  if (!destSpirvModule)
     return failure();
 
-  // Wrap around in a new MLIR module.
-  OwningModuleRef dstModule(ModuleOp::create(
-      FileLineColLoc::get(&deserializationContext,
-                          /*filename=*/"", /*line=*/0, /*column=*/0)));
-  dstModule->getBody()->push_front(spirvModule.release());
-  dstModule->print(output);
+  destSpirvModule->print(output);
+  return success();
+}
 
-  return mlir::success();
+static OwningOpRef<Operation *> parseSPIRVModule(llvm::SourceMgr &sourceMgr,
+                                                 MLIRContext *context) {
+  return OwningOpRef<Operation *>(
+      parseSourceFile<spirv::ModuleOp>(sourceMgr, context).release());
 }
 
 namespace mlir {
 void registerTestRoundtripSPIRV() {
   TranslateFromMLIRRegistration roundtrip(
       "test-spirv-roundtrip",
-      [](ModuleOp module, raw_ostream &output) {
-        return roundTripModule(module, /*emitDebugInfo=*/false, output);
+      [](Operation *root, raw_ostream &output) {
+        return roundTripModule(root, /*emitDebugInfo=*/false, output);
       },
-      [](DialectRegistry &registry) {
-        registry.insert<spirv::SPIRVDialect>();
-      });
+      [](DialectRegistry &registry) { registry.insert<spirv::SPIRVDialect>(); },
+      parseSPIRVModule);
 }
 
 void registerTestRoundtripDebugSPIRV() {
   TranslateFromMLIRRegistration roundtrip(
       "test-spirv-roundtrip-debug",
-      [](ModuleOp module, raw_ostream &output) {
-        return roundTripModule(module, /*emitDebugInfo=*/true, output);
+      [](Operation *root, raw_ostream &output) {
+        return roundTripModule(root, /*emitDebugInfo=*/true, output);
       },
-      [](DialectRegistry &registry) {
-        registry.insert<spirv::SPIRVDialect>();
-      });
+      [](DialectRegistry &registry) { registry.insert<spirv::SPIRVDialect>(); },
+      parseSPIRVModule);
 }
 } // namespace mlir

mlir/lib/Translation/Translation.cpp

 //===----------------------------------------------------------------------===//
 
 // Puts `function` into the to-MLIR translation registry unless there is already
 // a function registered for the same name.
 static void registerTranslateToMLIRFunction(
     StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
   auto wrappedFn = [function](llvm::SourceMgr &sourceMgr, raw_ostream &output,
                               MLIRContext *context) {
-    OwningModuleRef module = function(sourceMgr, context);
-    if (!module || failed(verify(*module)))
+    OwningOpRef<Operation *> root = function(sourceMgr, context);
+    if (!root || failed(verify(*root)))
       return failure();
-    module->print(output);
+    (*root)->print(output);
     return success();
   };
   registerTranslation(name, wrappedFn);
 }
 
 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
     StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
   registerTranslateToMLIRFunction(name, function);

 }
 
 //===----------------------------------------------------------------------===//
 // Translation from MLIR
 //===----------------------------------------------------------------------===//
 
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
     StringRef name, const TranslateFromMLIRFunction &function,
-    std::function<void(DialectRegistry &)> dialectRegistration) {
-  registerTranslation(name, [function, dialectRegistration](
+    std::function<void(DialectRegistry &)> dialectRegistration,
+    ParseFunction parseFunction) {
+  registerTranslation(name, [function, dialectRegistration, parseFunction](
                                 llvm::SourceMgr &sourceMgr, raw_ostream &output,
                                 MLIRContext *context) {
     DialectRegistry registry;
     dialectRegistration(registry);
     context->appendDialectRegistry(registry);
-    auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
-    if (!module)
+    OwningOpRef<Operation *> root = parseFunction(sourceMgr, context);
+    if (!root)
       return failure();
-    return function(module.get(), output);
+    return function(root.get(), output);
   });
 }
 
+OwningOpRef<Operation *>
+TranslateFromMLIRRegistration::parseBuiltinModule(llvm::SourceMgr &sourceMgr,
+                                                  MLIRContext *context) {
+  return OwningOpRef<Operation *>(
+      parseSourceFile(sourceMgr, context).release());
+}
+
 //===----------------------------------------------------------------------===//
 // Translation Parser
 //===----------------------------------------------------------------------===//
 
 TranslationParser::TranslationParser(llvm::cl::Option &opt)
     : llvm::cl::parser<const TranslateFunction *>(opt) {
   for (const auto &kv : getTranslationRegistry())
     addLiteralOption(kv.first(), &kv.second, kv.first());