[Mlir-commits] [mlir] [mlir][transform] Allow passing various library files to interpreter. (PR #67120)

[Ingo Müller]

================
@@ -611,53 +613,131 @@ LogicalResult transform::detail::interpreterBaseRunOnOperationImpl(
 
 LogicalResult transform::detail::interpreterBaseInitializeImpl(
     MLIRContext *context, StringRef transformFileName,
-    StringRef transformLibraryFileName,
+    ArrayRef<std::string> transformLibraryPaths,
     std::shared_ptr<OwningOpRef<ModuleOp>> &sharedTransformModule,
     std::shared_ptr<OwningOpRef<ModuleOp>> &transformLibraryModule,
     function_ref<std::optional<LogicalResult>(OpBuilder &, Location)>
         moduleBuilder) {
-  OwningOpRef<ModuleOp> parsedTransformModule;
-  if (failed(parseTransformModuleFromFile(context, transformFileName,
-                                          parsedTransformModule)))
-    return failure();
-  if (parsedTransformModule && failed(mlir::verify(*parsedTransformModule)))
-    return failure();
+  auto unknownLoc = UnknownLoc::get(context);
 
-  OwningOpRef<ModuleOp> parsedLibraryModule;
-  if (failed(parseTransformModuleFromFile(context, transformLibraryFileName,
-                                          parsedLibraryModule)))
-    return failure();
-  if (parsedLibraryModule && failed(mlir::verify(*parsedLibraryModule)))
-    return failure();
+  // Parse module from file.
+  OwningOpRef<ModuleOp> moduleFromFile;
+  {
+    auto loc = FileLineColLoc::get(context, transformFileName, 0, 0);
+    if (failed(parseTransformModuleFromFile(context, transformFileName,
+                                            moduleFromFile)))
+      return emitError(loc) << "failed to parse transform module";
+    if (moduleFromFile && failed(mlir::verify(*moduleFromFile)))
+      return emitError(loc) << "failed to verify transform module";
+  }
 
-  if (parsedTransformModule) {
-    sharedTransformModule = std::make_shared<OwningOpRef<ModuleOp>>(
-        std::move(parsedTransformModule));
+  // Assemble list of library files.
+  SmallVector<std::string> libraryFileNames;
+  for (const std::string &path : transformLibraryPaths) {
+    auto loc = FileLineColLoc::get(context, transformFileName, 0, 0);
+
+    if (llvm::sys::fs::is_regular_file(path)) {
+      LLVM_DEBUG(DBGS() << "Adding '" << path << "' to list of files\n");
+      libraryFileNames.push_back(path);
+      continue;
+    }
+
+    if (!llvm::sys::fs::is_directory(path)) {
+      return emitError(loc)
+             << "'" << path << "' is neither a file nor a directory";
+    }
+
+    LLVM_DEBUG(DBGS() << "Looking for files in '" << path << "':\n");
+
+    std::error_code ec;
+    for (llvm::sys::fs::directory_iterator it(path, ec), itEnd;
+         it != itEnd && !ec; it.increment(ec)) {
+      const std::string &fileName = it->path();
+
+      if (it->type() != llvm::sys::fs::file_type::regular_file) {
+        LLVM_DEBUG(DBGS() << "  Skipping non-regular file '" << fileName
+                          << "'\n");
+        continue;
+      }
+
+      if (!StringRef(fileName).endswith(".mlir")) {
+        LLVM_DEBUG(DBGS() << "  Skipping '" << fileName
+                          << "' because it does not end with '.mlir'\n");
+        continue;
+      }
+
+      LLVM_DEBUG(DBGS() << "  Adding '" << fileName << "' to list of files\n");
+      libraryFileNames.push_back(fileName);
+    }
+
+    if (ec)
+      return emitError(loc) << "error while opening files in '" << path
+                            << "': " << ec.message();
+  }
----------------
ingomueller-net wrote:

Good idea! Done in latest commit.

https://github.com/llvm/llvm-project/pull/67120