[om-linker] Set up boilerplate; parse and concat modules (#5693)

This commit is based on firld work by Will. This commit adds a boilerplate and
minimal implementation for `om-linker` which aims to link separated objectmodel
IRs into a single IR.

With this PR `om-linker` is able to read input MLIR files, and concat into
a single MLIR file without flattening ModuleOp. Currently it doesn't perform
any symbol resolution. For instance, we link following IR

```a.mlir
    om.class @A(%arg: i1) {}
```

```b.mlir
    om.class.extern @A(%arg: i1) {}
    om.class @b(%arg: i2) {}
```

`$ om-linker a.mlir b.mlir`

```mlir
module {
  module {
    om.class @A(%arg: i1) {}
  }
  module {
    om.class.extern @A(%arg: i1) {}
    om.class @b(%arg: i2) {}
  }
}
```

The actual linking pass will be added in a follow-up PR.

Co-authored-by: Will Dietz <will.dietz@sifive.com>

test/CMakeLists.txt

@@ -30,6 +30,7 @@ set(CIRCT_TEST_DEPENDS
   handshake-runner
   firtool
   hlstool
+  om-linker
   )
 
 if (CIRCT_GTEST_AVAILABLE)

test/lit.cfg.py

@@ -58,7 +58,7 @@ tool_dirs = [
 tools = [
     'firtool', 'circt-as', 'circt-dis', 'circt-opt', 'circt-reduce',
     'circt-translate', 'circt-capi-ir-test', 'circt-capi-om-test', 'esi-tester',
-    'hlstool', 'arcilator'
+    'hlstool', 'arcilator', 'om-linker'
 ]
 
 # Enable Verilator if it has been detected.

test/om-linker/Inputs/a.mlir

@@ -0,0 +1,4 @@
+module {
+  om.class @A(%arg: i1) {
+  }
+}

test/om-linker/Inputs/b.mlir

@@ -0,0 +1,6 @@
+module {
+  om.class.extern @A(%arg: i1) {
+  }
+  om.class @B(%arg: i2) {
+  }
+}

test/om-linker/commandline.mlir

@@ -0,0 +1,6 @@
+// RUN: om-linker --help | FileCheck %s --implicit-check-not='{{[Oo]}}ptions:'
+
+// CHECK: OVERVIEW: OM linker 
+// CHECK: General {{[Oo]}}ptions
+// CHECK: Generic Options
+// CHECK: om-linker Options

test/om-linker/link.mlir

@@ -0,0 +1,13 @@
+// RUN: om-linker %S/Inputs/a.mlir %S/Inputs/b.mlir | FileCheck %s
+// CHECK:      module {
+// CHECK-NEXT:   module {
+// CHECK-NEXT:     om.class @A(%arg: i1) {
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT:   module {
+// CHECK-NEXT:     om.class.extern @A(%arg: i1) {
+// CHECK-NEXT:     }
+// CHECK-NEXT:     om.class @B(%arg: i2) {
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }

tools/CMakeLists.txt

@@ -12,5 +12,6 @@ add_subdirectory(esi)
 add_subdirectory(handshake-runner)
 add_subdirectory(firtool)
 add_subdirectory(llhd-sim)
+add_subdirectory(om-linker)
 add_subdirectory(py-split-input-file)
 add_subdirectory(hlstool)

tools/om-linker/CMakeLists.txt

@@ -0,0 +1,19 @@
+set(LLVM_LINK_COMPONENTS
+  Support
+)
+
+add_llvm_tool(om-linker
+  om-linker.cpp
+)
+llvm_update_compile_flags(om-linker)
+target_link_libraries(om-linker PRIVATE
+  CIRCTOM
+  CIRCTSupport
+
+  MLIRBytecodeReader
+  MLIRBytecodeWriter
+  MLIRParser
+  MLIRSupport
+  MLIRIR
+  MLIROptLib
+)

tools/om-linker/om-linker.cpp

@@ -0,0 +1,202 @@
+//===- om-linker.cpp - An utility for linking objectmodel  ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements 'om-linker', which links separated OM dialect IRs into a
+// single IR.
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Dialect/OM/OMDialect.h"
+#include "circt/Support/Version.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Threading.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/FileUtilities.h"
+#include "mlir/Support/Timing.h"
+#include "mlir/Support/ToolUtilities.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
+
+using namespace llvm;
+using namespace mlir;
+using namespace circt;
+
+static cl::OptionCategory mainCategory("om-linker Options");
+
+static cl::list<std::string> inputFilenames(cl::Positional, cl::OneOrMore,
+                                            cl::desc("<input files>"),
+                                            cl::cat(mainCategory));
+
+static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
+                                           cl::value_desc("filename"),
+                                           cl::init("-"),
+                                           cl::cat(mainCategory));
+
+static cl::opt<bool>
+    emitBytecode("emit-bytecode",
+                 cl::desc("Emit bytecode when generating MLIR output"),
+                 cl::init(false), cl::cat(mainCategory));
+
+/// Check output stream before writing bytecode to it.
+/// Warn and return true if output is known to be displayed.
+static bool checkBytecodeOutputToConsole(raw_ostream &os) {
+  if (os.is_displayed()) {
+    errs() << "WARNING: You're attempting to print out a bytecode file.\n"
+              "This is inadvisable as it may cause display problems. If\n"
+              "you REALLY want to taste MLIR bytecode first-hand, you\n"
+              "can force output with the `-f' option.\n\n";
+    return true;
+  }
+  return false;
+}
+
+static cl::opt<bool> force("f", cl::desc("Enable binary output on terminals"),
+                           cl::init(false), cl::cat(mainCategory));
+
+/// Print the operation to the specified stream, emitting bytecode when
+/// requested and politely avoiding dumping to terminal unless forced.
+static LogicalResult printOp(Operation *op, raw_ostream &os) {
+  if (emitBytecode && (force || !checkBytecodeOutputToConsole(os)))
+    return writeBytecodeToFile(op, os,
+                               mlir::BytecodeWriterConfig(getCirctVersion()));
+  op->print(os);
+  return success();
+}
+
+/// This implements the top-level logic for the om-linker command, invoked once
+/// command line options are parsed and LLVM/MLIR are all set up and ready to
+/// go.
+static LogicalResult executeOMLinker(MLIRContext &context) {
+  // Create the timing manager we use to sample execution times.
+  DefaultTimingManager tm;
+  applyDefaultTimingManagerCLOptions(tm);
+  auto ts = tm.getRootScope();
+
+  // Set up the input files.
+  struct ParsedInput {
+    StringRef name;
+    OwningOpRef<ModuleOp> mod;
+    SourceMgr mgr;
+  };
+
+  struct LeakyInputs {
+    SmallVector<ParsedInput> inputs;
+    LeakyInputs(size_t n) { inputs.resize(n); }
+    ~LeakyInputs() {
+      // Leak
+      for (auto &input : inputs)
+        input.mod.release();
+    };
+  };
+
+  auto numFiles = inputFilenames.size();
+  LeakyInputs leakyInputs(numFiles);
+  auto &srcs = leakyInputs.inputs;
+  auto parserTimer = ts.nest("Parsing inputs");
+  auto loadFile = [&](size_t i) -> LogicalResult {
+    auto &s = srcs[i];
+    s.name = inputFilenames[i];
+    auto fileParseTimer = parserTimer.nest(s.name);
+    s.mod = parseSourceFile<ModuleOp>(s.name, s.mgr, &context);
+    if (!s.mod)
+      return failure();
+    return success();
+  };
+
+  if (failed(failableParallelForEachN(&context, 0, numFiles, loadFile))) {
+    errs() << "error reading inputs\n";
+    return failure();
+  }
+
+  // This is the result module we are linking into.
+  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
+
+  // Concat input modules without flatting ModuleOp. For example we construct
+  // an IR like this:
+  // ```
+  // module {
+  //   module {
+  //     om.class @A(%arg: i1) {}
+  //   }
+  //   module {
+  //     om.class.extern @A(%arg: i1) {}
+  //     om.class @B(%arg: i2) {}
+  //   }
+  // }
+  // ```
+  // Actual linking will be performed by the subsequent pass pipeline.
+  auto builder = OpBuilder::atBlockEnd(module.get().getBody());
+  for (auto &in : srcs)
+    builder.insert(in.mod.get());
+
+  // TODO: Apply a linking pass here.
+
+  // Create the output file.
+  std::string errorMessage;
+  auto outputFile = openOutputFile(outputFilename, &errorMessage);
+  if (!outputFile) {
+    errs() << errorMessage << "\n";
+    return failure();
+  }
+
+  // Dump output.
+  if (failed(printOp(module.get(), outputFile->os())))
+    return failure();
+
+  // If the result succeeded and we're emitting a file, close it.
+  outputFile->keep();
+
+  // We intentionally "leak" the Module into the MLIRContext instead of
+  // deallocating it.  There is no need to deallocate it right before process
+  // exit.
+  (void)module.release();
+  return success();
+}
+
+/// Main driver for om-linker command.  This sets up LLVM and MLIR, and parses
+/// command line options before passing off to 'executeOMLinker'.  This is set
+/// up so we can `exit(0)` at the end of the program to avoid teardown of the
+/// MLIRContext and modules inside of it (reducing compile time).
+int main(int argc, char **argv) {
+  InitLLVM y(argc, argv);
+
+  // Set the bug report message to indicate users should file issues on
+  // llvm/circt and not llvm/llvm-project.
+  setBugReportMsg(circtBugReportMsg);
+
+  // Hide default LLVM options, other than for this tool.
+  // MLIR options are added below.
+  cl::HideUnrelatedOptions(mainCategory);
+
+  // Register any pass manager command line options.
+  registerMLIRContextCLOptions();
+  registerPassManagerCLOptions();
+  registerDefaultTimingManagerCLOptions();
+  registerAsmPrinterCLOptions();
+  cl::AddExtraVersionPrinter(
+      [](raw_ostream &os) { os << getCirctVersion() << '\n'; });
+  // Parse pass names in main to ensure static initialization completed.
+  cl::ParseCommandLineOptions(argc, argv, "OM linker\n");
+
+  MLIRContext context;
+  // Register OM dialect.
+  context.loadDialect<om::OMDialect>();
+
+  // Do the guts of the om-linker process.
+  auto result = executeOMLinker(context);
+
+  // Use "exit" instead of return'ing to signal completion.  This avoids
+  // invoking the MLIRContext destructor, which spends a bunch of time
+  // deallocating memory etc which process exit will do for us.
+  exit(failed(result));
+}