[mlir-translate] Support parsing operations other than 'builtin.module' as top-level

[rkayaith]

Summary

This adds a '--no-implicit-module' option, which disables the insertion
of a top-level 'builtin.module' during parsing.

The translation APIs are also updated to take/return 'Operation*'
instead of 'ModuleOp', to allow other operation types to be used. To
simplify translations which are restricted to specific operation types,
'TranslateFromMLIRRegistration' has an overload which performs the
necessary cast and error checking.

[rkayaith]

mlir/test/Target/SPIRV/module.mlir

1

I only updated one test, but actually all the spirv translate tests work with --no-implicit-module. The translation code could be simplified a bit if this was the only mode that needed to be supported, since the IR walks to find a spv.module could be removed

[rkayaith]

rebase

[rkayaith]

fix build bot?

[rkayaith]

mlir/lib/Tools/mlir-translate/Translation.cpp

96–99

Putting the cl option here doesn't seem great, but it's hard to plumb it to the parse call otherwise. Open to other ideas.

[rkayaith]

friendly ping :)

[rkayaith]

rebase

[mehdi_amini]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

I think if would be more friendly if we could register these with a type safe trampoline. Something that automatically dyn_cast the operation* to the Op expected by the call back and print an error message and fails before calling the translation call back in case of mismatch.

[rkayaith]

fix standalone-translate

[rriddle]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1185–1188

[rkayaith]

address comment

[rkayaith]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

I can give this a try but I don't think any of the upstream translations would actually use this at the moment, unless I changed the spirv translations to only work with --no-implicit-module.

[mehdi_amini]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

I'm not sure you got my point here, let me try to reformulate: we should make it so that the current translation registration don't need to be updated with their current callback.

They are passing a lambda that is [](ModuleOp module, raw_ostream &output) { and they should be able to continue to do that.

The behavior of the system is that it should take a function with the signature LogicalResult(T module, raw_ostream &output) and then check after parsing that the parsed operation is actually T or error out.

[mehdi_amini]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

It reminds of the logic we use with walk([] (T op) {...}) where we check that we match T before calling the callback: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Visitors.h#L227-L229

[rkayaith]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

Sure, what I meant though is that I don't see why any of the upstream translations would want to continue taking a ModuleOp:

• mlir-to-llvm translation supports generic Operation *: https://github.com/llvm/llvm-project/blob/0849f0a5a38a16f52d7c698b7b024bd58942fab7/mlir/include/mlir/Target/LLVMIR/Export.h#L28
• mlir-to-cpp does too: https://github.com/llvm/llvm-project/blob/0849f0a5a38a16f52d7c698b7b024bd58942fab7/mlir/include/mlir/Target/Cpp/CppEmitter.h#L29
• spirv translation uses spirv::ModuleOp: https://github.com/llvm/llvm-project/blob/14b9505be9d10513b4ca0a22a73a158bf28d0959/mlir/include/mlir/Target/SPIRV/Serialization.h#L35. because of the implicit module it has to traverse the builtin.module to actually find a spirv.module, but this traversal can be done on Operation * as well.

The api you proposed makes sense to me, but is there an issue if it ends up unused upstream?

[mehdi_amini]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

Can we default to --no-implicit-module for translations?

[rriddle]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

Could we just remove it altogether? Translations are way less common than passes.

[rkayaith]

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

if I flip the default I only get 24 failures (out of 81). but I think it would be confusing if the default doesn't match mlir-opt.

[rkayaith]

add a TranslateFromMLIRRegistration constructor that handles casting to concrete op type

mlir/include/mlir/Tools/mlir-translate/Translation.h

49

I added a TranslateFromMLIRRegistration constructor that handles the casting, and for now I switched the spirv translation back to only working for builtin.module so that something actually uses it. and a patch here to switch it to work on spirv.module instead if we want that: https://reviews.llvm.org/D135819

[rriddle]

mlir/include/mlir/Tools/mlir-translate/Translation.h

94–102

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

1321–1324

We should likely special case conversions from OwningOpRef<FooOp> to OwningOpRef<Operation *>.

[rkayaith]

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

1321–1324

so something like this on OwningOpRef?

operator OwningOpRef<Operation *>() && { return release().getOperation(); }

[rkayaith]

update description, use OwningOpRef conversion, address comment

[Peiming]

mlir/lib/Tools/mlir-translate/Translation.cpp

96–99

This seems break some customized mlir-opt tools (which calls MlirOptMain) and need to register the translation here as they both register the cl option, leading to Option 'no-implicit-module' registered more than once!

[mehdi_amini]

mlir/lib/Tools/mlir-translate/Translation.cpp

96–99

Yes a static isn't great here, that said it's not clear to me why you'd have global translate registrations in a mlir-opt tool?

[Peiming]

mlir/lib/Tools/mlir-translate/Translation.cpp

96–99

I will offer you more details offline :-)

[rkayaith]

mlir/lib/Tools/mlir-translate/Translation.cpp

96–99

I think D136561 should let you work around this in your tool by skipping the option registration

[mehdi_amini]

The bot is broken, are you on it already? https://lab.llvm.org/buildbot/#/builders/61/builds/34355

[mehdi_amini]

mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp:67:30: error: use 'template' keyword to treat 'dyn_cast_or_null' as a dependent template name
                            .dyn_cast_or_null<arith::FastMathFlagsAttr>();
                             ^
                             template

[rkayaith]

In D134237#3886861, @mehdi_amini wrote:

The bot is broken, are you on it already? https://lab.llvm.org/buildbot/#/builders/61/builds/34355

i think you meant to comment on D126305

mlir/include/mlir/Tools/mlir-translate/Translation.h

 class SourceMgr;
 class StringRef;
 } // namespace llvm
 
 namespace mlir {
 class DialectRegistry;
 struct LogicalResult;
 class MLIRContext;
-class ModuleOp;
+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<OwningOpRef<ModuleOp>(
+/// should create a new MLIR 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
+/// implementation should create a new MLIR Operation in the given context. If
 /// source-related error reporting is required from within the function, use
 /// TranslateSourceMgrToMLIRFunction instead.
 using TranslateStringRefToMLIRFunction =
-    std::function<OwningOpRef<ModuleOp>(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.
+/// outputs the result to a stream. It is allowed to modify the operation.
 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 *)>;

 struct TranslateRegistration {
   TranslateRegistration(llvm::StringRef name,
                         const TranslateFunction &function);
 };
 /// \}
 
 /// A command line parser for translation functions.
 struct TranslationParser : public llvm::cl::parser<const TranslateFunction *> {
   TranslationParser(llvm::cl::Option &opt);
 
   void printOptionInfo(const llvm::cl::Option &o,
                        size_t globalWidth) const override;
 };
 
 } // namespace mlir
 
 #endif // MLIR_TOOLS_MLIRTRANSLATE_TRANSLATION_H

mlir/lib/Target/Cpp/TranslateRegistration.cpp

 void registerToCppTranslation() {
   static llvm::cl::opt<bool> declareVariablesAtTop(
       "declare-variables-at-top",
       llvm::cl::desc("Declare variables at top when emitting C/C++"),
       llvm::cl::init(false));
 
   TranslateFromMLIRRegistration reg(
       "mlir-to-cpp",
-      [](ModuleOp module, raw_ostream &output) {
+      [](Operation *op, raw_ostream &output) {
         return emitc::translateToCpp(
-            module, output,
+            op, output,
             /*declareVariablesAtTop=*/declareVariablesAtTop);
       },
       [](DialectRegistry &registry) {
         // clang-format off
         registry.insert<arith::ArithmeticDialect,
                         cf::ControlFlowDialect,
                         emitc::EmitCDialect,
                         func::FuncDialect,
                         math::MathDialect,
                         scf::SCFDialect>();
         // clang-format on
       });
 }
 
 } // namespace mlir

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.
-OwningOpRef<ModuleOp> translateLLVMIRToModule(llvm::SourceMgr &sourceMgr,
-                                              MLIRContext *context) {
+static 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);
+  OwningOpRef<ModuleOp> module =
+      translateLLVMIRToModule(std::move(llvmModule), context);
+  return module.release().getOperation();
 }
 
 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 *op, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
-        auto llvmModule = translateModuleToLLVMIR(module, llvmContext);
+        auto llvmModule = translateModuleToLLVMIR(op, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       },
       [](DialectRegistry &registry) {
         registry.insert<DLTIDialect, func::FuncDialect>();
         registerAllToLLVMIRTranslations(registry);
       });
 }
 } // namespace mlir

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

llvmModule->getOrInsertFunction("free", builder.getVoidTy(),
                                   builder.getInt8PtrTy());
 
   return llvmModule;
 }
 
 std::unique_ptr<llvm::Module>
 mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
                               StringRef name) {
   if (!satisfiesLLVMModule(module))
-    return nullptr;
+    return module->emitOpError("can not be translated to an LLVMIR module"),
+           nullptr;
 
   std::unique_ptr<llvm::Module> llvmModule =
       prepareLLVMModule(module, llvmContext, name);
   if (!llvmModule)
     return nullptr;
 
   LLVM::ensureDistinctSuccessors(module);
 
   ModuleTranslation translator(module, std::move(llvmModule));

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 OwningOpRef<ModuleOp> 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 {};
 
-  OwningOpRef<ModuleOp> module(ModuleOp::create(FileLineColLoc::get(
-      context, input->getBufferIdentifier(), /*line=*/0, /*column=*/0)));
-  module->getBody()->push_front(spirvModule.release());
+  auto module = ModuleOp::create(FileLineColLoc::get(
+      context, input->getBufferIdentifier(), /*line=*/0, /*column=*/0));
+  module.getBody()->push_front(spirvModule.release());
 
-  return module;
+  return module.getOperation();
 }
 
 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 *op, raw_ostream &output) {
+  if (!op)
     return failure();
 
   SmallVector<uint32_t, 0> binary;
 
   SmallVector<spirv::ModuleOp, 1> spirvModules;
-  module.walk([&](spirv::ModuleOp op) { spirvModules.push_back(op); });
+  op->walk([&](spirv::ModuleOp op) { spirvModules.push_back(op); });
 
   if (spirvModules.empty())
-    return module.emitError("found no 'spirv.module' op");
+    return op->emitError("found no 'spv.module' op");
 
   if (spirvModules.size() != 1)
-    return module.emitError("found more than one 'spirv.module' op");
+    return op->emitError("found more than one 'spv.module' op");
 
   if (failed(spirv::serialize(spirvModules[0], binary)))
     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);
+      [](Operation *op, raw_ostream &output) {
+        return serializeModule(op, output);
       },
       [](DialectRegistry &registry) {
         registry.insert<spirv::SPIRVDialect>();
       });
 }
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Round-trip registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult roundTripModule(ModuleOp srcModule, bool emitDebugInfo,
+static LogicalResult roundTripModule(Operation *src, bool emitDebugInfo,
                                      raw_ostream &output) {
   SmallVector<uint32_t, 0> binary;
-  MLIRContext *context = srcModule.getContext();
-  auto spirvModules = srcModule.getOps<spirv::ModuleOp>();
+  MLIRContext *context = src->getContext();
+  SmallVector<spirv::ModuleOp, 1> spirvModules;
+  src->walk([&](spirv::ModuleOp op) { spirvModules.push_back(op); });
 
-  if (spirvModules.begin() == spirvModules.end())
-    return srcModule.emitError("found no 'spirv.module' op");
+  if (spirvModules.empty())
+    return src->emitError("found no 'spv.module' op");
 
-  if (std::next(spirvModules.begin()) != spirvModules.end())
-    return srcModule.emitError("found more than one 'spirv.module' op");
+  if (spirvModules.size() > 1)
+    return src->emitError("found more than one 'spv.module' op");
 
   spirv::SerializationOptions options;
   options.emitDebugInfo = emitDebugInfo;
   if (failed(spirv::serialize(*spirvModules.begin(), binary, options)))
     return failure();
 
   MLIRContext deserializationContext(context->getDialectRegistry());
   // TODO: we should only load the required dialects instead of all dialects.
static LogicalResult roundTripModule(Operation *src, bool emitDebugInfo,
 
   return mlir::success();
 }
 
 namespace mlir {
 void registerTestRoundtripSPIRV() {
   TranslateFromMLIRRegistration roundtrip(
       "test-spirv-roundtrip",
-      [](ModuleOp module, raw_ostream &output) {
-        return roundTripModule(module, /*emitDebugInfo=*/false, output);
+      [](Operation *op, raw_ostream &output) {
+        return roundTripModule(op, /*emitDebugInfo=*/false, output);
       },
       [](DialectRegistry &registry) {
         registry.insert<spirv::SPIRVDialect>();
       });
 }
 
 void registerTestRoundtripDebugSPIRV() {
   TranslateFromMLIRRegistration roundtrip(
       "test-spirv-roundtrip-debug",
-      [](ModuleOp module, raw_ostream &output) {
-        return roundTripModule(module, /*emitDebugInfo=*/true, output);
+      [](Operation *op, raw_ostream &output) {
+        return roundTripModule(op, /*emitDebugInfo=*/true, output);
       },
       [](DialectRegistry &registry) {
         registry.insert<spirv::SPIRVDialect>();
       });
 }
 } // namespace mlir

mlir/lib/Tools/mlir-translate/Translation.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-translate/Translation.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Parser/Parser.h"
+#include "mlir/Tools/ParseUtilties.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Translation Registry
 //===----------------------------------------------------------------------===//
 

 //===----------------------------------------------------------------------===//
 
 // 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) {
-    OwningOpRef<ModuleOp> module = function(sourceMgr, context);
-    if (!module || failed(verify(*module)))
+    OwningOpRef<Operation *> op = function(sourceMgr, context);
+    if (!op || failed(verify(*op)))
       return failure();
-    module->print(output);
+    op.get()->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,
     const std::function<void(DialectRegistry &)> &dialectRegistration) {
+
+  static llvm::cl::opt<bool> noImplicitModule{
+      "no-implicit-module",
+      llvm::cl::desc("Disable the parsing of an implicit top-level module op"),
+      llvm::cl::init(false)};
+
   registerTranslation(name, [function, dialectRegistration](
                                 llvm::SourceMgr &sourceMgr, raw_ostream &output,
                                 MLIRContext *context) {
     DialectRegistry registry;
     dialectRegistration(registry);
     context->appendDialectRegistry(registry);
-    auto module = parseSourceFile<ModuleOp>(sourceMgr, context);
-    if (!module || failed(verify(*module)))
+    OwningOpRef<Operation *> op =
+        parseSourceFileForTool(sourceMgr, context, !noImplicitModule);
+    if (!op || failed(verify(*op)))
       return failure();
-    return function(module.get(), output);
+    return function(op.get(), output);
   });
 }
 
 //===----------------------------------------------------------------------===//
 // Translation Parser
 //===----------------------------------------------------------------------===//
 
 TranslationParser::TranslationParser(llvm::cl::Option &opt)

mlir/test/Target/LLVMIR/invalid-module.mlir

+// RUN: mlir-translate -verify-diagnostics -mlir-to-llvmir --no-implicit-module %s
+
+// expected-error@below {{'llvm.func' op can not be translated to an LLVMIR module}}
+llvm.func @foo() {
+  llvm.return
+}

mlir/test/Target/SPIRV/module.mlir

-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -test-spirv-roundtrip --no-implicit-module -split-input-file %s | FileCheck %s
 
 // CHECK:      spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
 // CHECK-NEXT:   spirv.func @foo() "Inline" {
 // CHECK-NEXT:     spirv.Return
 // CHECK-NEXT:   }
 // CHECK-NEXT: }
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {