[MLIR][mlir-link] Simplify inference of linker flags for the first module

Author: xlauko

mlir/include/mlir/Linker/Linker.h

@@ -60,17 +60,10 @@ class Linker {
 
   MLIRContext *getContext() { return context; }
 
-  LogicalResult linkInModule(OwningOpRef<ModuleOp> src, unsigned flags = None);
+  LogicalResult linkInModule(OwningOpRef<ModuleOp> src);
 
   unsigned getFlags() const;
 
-  // Infer how to link file from linker config
-  LinkFileConfig linkFileConfig(unsigned fileFlags = None) const;
-
-  /// The first file is linked without internalization and with the
-  /// OverrideFromSrc flag set
-  LinkFileConfig firstFileConfig(unsigned fileFlags = None) const;
-
   OwningOpRef<ModuleOp> takeModule() { return std::move(composite); }
 
   LogicalResult emitFileError(const Twine &fileName, const Twine &message) {

mlir/lib/Linker/Linker.cpp

@@ -268,13 +268,17 @@ LogicalResult ModuleLinker::run() {
   return success();
 }
 
-LogicalResult Linker::linkInModule(OwningOpRef<ModuleOp> src, unsigned flags) {
+LogicalResult Linker::linkInModule(OwningOpRef<ModuleOp> src) {
+  unsigned flags = getFlags();
+
   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());
+    // We always override from source for the first module.
+    flags &= Linker::OverrideFromSrc;
   }
 
   IRMover mover(composite.get());
@@ -290,12 +294,3 @@ unsigned Linker::getFlags() const {
 
   return flags;
 }
-
-Linker::LinkFileConfig Linker::linkFileConfig(unsigned fileFlags) const {
-  return {.flags = fileFlags};
-}
-
-Linker::LinkFileConfig Linker::firstFileConfig(unsigned fileFlags) const {
-  // Filter out flags that don't apply to the first file we load.
-  return {.flags = fileFlags & Linker::OverrideFromSrc};
-}

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

@@ -153,7 +153,6 @@ const LinkerCLOptions &createLinkerConfigFromCLOptions() {
 /// This class encapsulates all the file handling logic for linker.
 class FileProcessor {
 public:
-  using FileConfig = Linker::LinkFileConfig;
   using OwningMemoryBuffer = std::unique_ptr<MemoryBuffer>;
 
   explicit FileProcessor(Linker &linker, raw_ostream &os, StringRef inMarker,
@@ -162,49 +161,39 @@ class FileProcessor {
 
   /// Process and link multiple input files
   LogicalResult linkFiles(const std::vector<std::string> fileNames) {
-
-    unsigned flags = linker.getFlags();
-    FileConfig config = linker.firstFileConfig(flags);
-
     for (StringRef fileName : fileNames) {
-      if (failed(processFile(fileName, config)))
+      if (failed(processFile(fileName)))
         return failure();
-
-      // Update config for subsequent files
-      config = linker.linkFileConfig(flags);
     }
 
     return success();
   }
 
 private:
-  /// Process a single file
-  LogicalResult processFile(StringRef fileName, FileConfig config) {
-    // Open input file
+  LogicalResult processFile(StringRef fileName) {
     std::string errorMessage;
     auto input = openInputFile(fileName, &errorMessage);
     if (!input)
       return linker.emitFileError(fileName, errorMessage);
 
     // Process each file chunk
-    if (failed(processFile(std::move(input), config)))
+    if (failed(processFile(std::move(input))))
       return linker.emitFileError(fileName, "Failed to process input file");
 
     return success();
   }
 
-  /// Process a single file buffer
-  LogicalResult processFile(OwningMemoryBuffer file, FileConfig config) {
+  /// Process a single input file, potentially containing multiple modules
+  LogicalResult processFile(OwningMemoryBuffer file) {
     return splitAndProcessBuffer(
         std::move(file),
-        [config, this](OwningMemoryBuffer chunk, raw_ostream & /* os */) {
-          return processFileChunk(std::move(chunk), config);
+        [this](OwningMemoryBuffer chunk, raw_ostream & /* os */) {
+          return processFileChunk(std::move(chunk));
         },
         os, inMarker, outMarker);
   }
 
-  /// Process a single input file, potentially containing multiple modules
-  LogicalResult processFileChunk(OwningMemoryBuffer buffer, FileConfig config) {
+  LogicalResult processFileChunk(OwningMemoryBuffer buffer) {
     auto sourceMgr = std::make_shared<SourceMgr>();
     sourceMgr->AddNewSourceBuffer(std::move(buffer), SMLoc());
 
@@ -214,7 +203,7 @@ class FileProcessor {
 
     // Parse the source file
     OwningOpRef<Operation *> op =
-        parseSourceFileForTool(sourceMgr, ctx, true /*insertImplicitModule*/);
+        parseSourceFileForTool(sourceMgr, ctx, /*insertImplicitModule=*/true);
     ctx->enableMultithreading(wasThreadingEnabled);
 
     if (!op)
@@ -228,7 +217,7 @@ class FileProcessor {
     // TBD: internalization
     OwningOpRef<ModuleOp> mod = cast<ModuleOp>(op.release());
     // Link the parsed operation
-    return linker.linkInModule(std::move(mod), config.flags);
+    return linker.linkInModule(std::move(mod));
   }
 
   Linker &linker;