[MLIR][mlir-link] Generate composite module

trailofbits · Jan 22, 2025 · f21b247 · f21b247
1 parent bcc180d
commit f21b247
Showing 1 changed file with 36 additions and 2 deletions.
diff --git a/mlir/lib/Tools/mlir-link/MlirLinkMain.cpp b/mlir/lib/Tools/mlir-link/MlirLinkMain.cpp
@@ -7,15 +7,29 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-link/MlirLinkMain.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/Support/FileUtilities.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/WithColor.h"
 
 using namespace mlir;
 using namespace llvm;
 
+OwningOpRef<ModuleOp> makeCompositeModule(MLIRContext *context) {
+  OpBuilder builder(context);
+  ModuleOp op =
+      builder.create<ModuleOp>(FileLineColLoc::get(context, "mlir-link", 0, 0));
+  OwningOpRef<ModuleOp> composite(op);
+  return composite;
+}
+
 LogicalResult mlir::MlirLinkMain(int argc, char **argv,
                                  DialectRegistry &registry) {
   static cl::OptionCategory linkCategory("Link options");
@@ -28,13 +42,33 @@ LogicalResult mlir::MlirLinkMain(int argc, char **argv,
       "o", cl::desc("Override output filename"), cl::init("-"),
       cl::value_desc("filename"), cl::cat(linkCategory));
 
-  static ExitOnError ExitOnErr;
+  static cl::opt<bool> verbose(
+      "v", cl::desc("Print information about actions taken"),
+      cl::cat(linkCategory));
+
+  static ExitOnError exitOnErr;
 
   InitLLVM y(argc, argv);
-  ExitOnErr.setBanner(std::string(argv[0]) + ": ");
+  exitOnErr.setBanner(std::string(argv[0]) + ": ");
 
   cl::HideUnrelatedOptions({&linkCategory, &getColorCategory()});
   cl::ParseCommandLineOptions(argc, argv, "mlir linker\n");
 
+  MLIRContext context(registry);
+  auto composite = makeCompositeModule(&context);
+
+  std::string errorMessage;
+
+  auto output = openOutputFile(outputFilename, &errorMessage);
+  if (!output) {
+    errs() << errorMessage;
+    return failure();
+  }
+
+  if (verbose)
+    errs() << "Writing linked module to '" << outputFilename << "'\n";
+
+  composite.get()->print(output->os());
+  output->keep();
   return success();
 }