[MLIR][mlir-link] Make createCompositeModule part of LinkerInterface

Author: xlauko

mlir/include/mlir/Linker/IRMover.h

@@ -19,15 +19,15 @@ namespace mlir::link {
 
 class IRMover {
 public:
-  IRMover(Operation *composite);
+  IRMover(mlir::ModuleOp composite) : composite(composite) {}
 
   ModuleOp getComposite() { return composite; }
   MLIRContext *getContext() { return composite->getContext(); }
 
   Error move(OwningOpRef<Operation *> src, ArrayRef<GlobalValue> valuesToLink);
 
 private:
-  ModuleOp composite;
+  mlir::ModuleOp composite;
 };
 
 } // namespace mlir::link

mlir/include/mlir/Linker/Linker.h

@@ -75,13 +75,12 @@ class Linker {
     bool internalize = false;
   };
 
-  Linker(ModuleOp composite, const LinkerConfig &config);
+  Linker(const LinkerConfig &config, MLIRContext *context)
+      : config(config), context(context) {}
 
-  MLIRContext *getContext() { return mover.getContext(); }
+  MLIRContext *getContext() { return context; }
 
-  LogicalResult linkInModule(OwningOpRef<Operation *> src,
-                             unsigned flags = None,
-                             InternalizeCallbackFn internalizeCallback = {});
+  LogicalResult linkInModule(OwningOpRef<ModuleOp> src, unsigned flags = None);
 
   unsigned getFlags() const;
 
@@ -92,9 +91,20 @@ class Linker {
   /// OverrideFromSrc flag set
   LinkFileConfig firstFileConfig(unsigned fileFlags = None) const;
 
+  OwningOpRef<ModuleOp> takeModule() { return std::move(composite); }
+
+  LogicalResult emitFileError(const Twine &fileName, const Twine &message) {
+    return emitError("Error processing file '" + fileName + "': " + message);
+  }
+
+  LogicalResult emitError(const Twine &message) {
+    return mlir::emitError(UnknownLoc::get(context), message);
+  }
+
 private:
   const LinkerConfig &config;
-  IRMover mover;
+  MLIRContext *context;
+  OwningOpRef<ModuleOp> composite;
 };
 
 } // namespace mlir::link

mlir/include/mlir/Linker/LinkerInterface.h

@@ -74,6 +74,11 @@ class LinkerInterface : public DialectInterface::Base<LinkerInterface> {
                              const LinkerSummaryState *state) const {
     return failure();
   }
+
+  virtual OwningOpRef<ModuleOp> createCompositeModule(ModuleOp src) const {
+    return ModuleOp::create(
+        FileLineColLoc::get(src.getContext(), "composite", 0, 0));
+  }
 };
 
 struct LinkableOp {

mlir/lib/Linker/IRMover.cpp

@@ -617,11 +617,8 @@ Error MLIRLinker::run() {
 }
 } // namespace
 
-IRMover::IRMover(Operation *composite) : composite(composite) {}
-
 Error IRMover::move(OwningOpRef<Operation *> src,
                     ArrayRef<GlobalValue> valuesToLink) {
-
   MLIRLinker linker(composite, std::move(src), valuesToLink);
   return linker.run();
 }

mlir/lib/Linker/Linker.cpp

@@ -33,14 +33,6 @@ class ModuleLinker {
   /// For symbol clashes, prefer those from src.
   unsigned flags;
 
-  /// List of global value names that should be internalized.
-  StringSet<> internalize;
-
-  /// Function that will perform the actual internalization. The reason for a
-  /// callback is that the linker cannot call internalizeModule without
-  /// creating a circular dependency between IPO and the linker.
-  InternalizeCallbackFn internalizeCallback;
-
   ModuleOp getSourceModule() { return cast<ModuleOp>(src.get()); }
 
   bool shouldOverrideFromSrc() const { return flags & Linker::OverrideFromSrc; }
@@ -85,8 +77,7 @@ class ModuleLinker {
 public:
   ModuleLinker(IRMover &mover, OwningOpRef<Operation *> src, unsigned flags,
                InternalizeCallbackFn internalizeCallback = {})
-      : mover(mover), src(std::move(src)), flags(flags),
-        internalizeCallback(std::move(internalizeCallback)) {}
+      : mover(mover), src(std::move(src)), flags(flags) {}
   LogicalResult run();
 };
 
@@ -235,7 +226,6 @@ bool ModuleLinker::linkIfNeeded(GlobalValue gv,
 
 LogicalResult ModuleLinker::run() {
   LLVM_DEBUG(llvm::dbgs() << "ModuleLinker::run" << "\n");
-  auto dst = mover.getComposite();
   auto src = getSourceModule();
 
   std::vector<Operation *> gvToClone;
@@ -263,14 +253,6 @@ LogicalResult ModuleLinker::run() {
   //   llvm_unreachable("unimplemented");
   // }
 
-  if (internalizeCallback) {
-    for (GlobalValue gvl : valuesToLink) {
-      StringRef name = gvl.getLinkedName();
-      LLVM_DEBUG(llvm::dbgs() << "Internalizing: " << name << "\n");
-      internalize.insert(name);
-    }
-  }
-
   bool hasErrors = false;
   // TODO: We are moving whatever the local src points to here (this->src), so
   // it can't be touched past this point.
@@ -284,20 +266,20 @@ LogicalResult ModuleLinker::run() {
   if (hasErrors)
     return failure();
 
-  if (internalizeCallback) {
-    internalizeCallback(dst, internalize);
-  }
-
   return success();
 }
 
-Linker::Linker(ModuleOp composite, const LinkerConfig &cfg)
-    : config(cfg), mover(composite) {}
+LogicalResult Linker::linkInModule(OwningOpRef<ModuleOp> src, unsigned flags) {
+  if (!composite) {
+    auto interface =
+        dyn_cast_or_null<LinkerInterface>(src->getOperation()->getDialect());
+    if (!interface)
+      return emitError("Module does not have a linker interface");
+    composite = interface->createCompositeModule(src.get());
+  }
 
-LogicalResult Linker::linkInModule(OwningOpRef<Operation *> src, unsigned flags,
-                                   InternalizeCallbackFn internalizeCallback) {
-  ModuleLinker modLinker(mover, std::move(src), flags,
-                         std::move(internalizeCallback));
+  IRMover mover(composite.get());
+  ModuleLinker modLinker(mover, std::move(src), flags);
   return modLinker.run();
 }
 

mlir/lib/Tools/mlir-link/MlirLinkMain.cpp

@@ -189,11 +189,11 @@ class FileProcessor {
     std::string errorMessage;
     auto input = openInputFile(fileName, &errorMessage);
     if (!input)
-      return emitFileError(fileName, errorMessage);
+      return linker.emitFileError(fileName, errorMessage);
 
     // Process each file chunk
     if (failed(processFile(std::move(input), config)))
-      return emitFileError(fileName, "Failed to process input file");
+      return linker.emitFileError(fileName, "Failed to process input file");
 
     return success();
   }
@@ -220,26 +220,20 @@ class FileProcessor {
     // Parse the source file
     OwningOpRef<Operation *> op =
         parseSourceFileForTool(sourceMgr, ctx, true /*insertImplicitModule*/);
-
     ctx->enableMultithreading(wasThreadingEnabled);
 
     if (!op)
-      return emitError("Failed to parse source file");
+      return linker.emitError("Failed to parse source file");
+
+    if (!isa<ModuleOp>(op.get()))
+      return op->emitError("Expected a ModuleOp");
 
     // TBD: symbol promotion
 
     // TBD: internalization
-
+    OwningOpRef<ModuleOp> mod = cast<ModuleOp>(op.release());
     // Link the parsed operation
-    return linker.linkInModule(std::move(op), config.flags);
-  }
-
-  LogicalResult emitFileError(const Twine &fileName, const Twine &message) {
-    return emitError("Error processing file '" + fileName + "': " + message);
-  }
-
-  LogicalResult emitError(const Twine &message) {
-    return mlir::emitError(mlir::UnknownLoc::get(linker.getContext()), message);
+    return linker.linkInModule(std::move(mod), config.flags);
   }
 
   Linker &linker;
@@ -269,22 +263,14 @@ LogicalResult mlir::MlirLinkMain(int argc, char **argv,
   MLIRContext context(registry);
   context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());
 
-  // Create composite module
-  OwningOpRef<ModuleOp> composite = [&context]() {
-    OpBuilder builder(&context);
-    return OwningOpRef<ModuleOp>(builder.create<ModuleOp>(
-        FileLineColLoc::get(&context, "mlir-link", 0, 0)));
-  }();
-
-  Linker linker(composite.get(), config);
+  Linker linker(config, &context);
 
   // Prepare output file
   std::string errorMessage;
   auto out = openOutputFile(config.outputFile, &errorMessage);
 
   if (!out) {
-    errs() << errorMessage;
-    return failure();
+    return linker.emitError("Failed to open output file: " + errorMessage);
   }
 
   StringRef inMarker = config.inputSplitMarker();
@@ -296,10 +282,11 @@ LogicalResult mlir::MlirLinkMain(int argc, char **argv,
   if (failed(proc.linkFiles(config.inputFiles)))
     return failure();
 
-  // TODO: Remove
+  OwningOpRef<ModuleOp> composite = linker.takeModule();
   if (failed(verify(composite.get(), true))) {
-    llvm::outs() << "Verify failed\n";
+    return composite->emitError("verification after linking failed");
   }
+
   composite->print(out->os());
   out->keep();