[MLIR][mlir-link] Add option to sort output symbols by name

Author: xlauko

mlir/include/mlir/Linker/Linker.h

@@ -64,7 +64,7 @@ class Linker {
 
   /// Perform linking and materialization in the destination module.
   /// Returns the linked module.
-  OwningOpRef<ModuleOp> link();
+  OwningOpRef<ModuleOp> link(bool sortSymbols = false);
 
   MLIRContext *getContext() { return context; }
 

mlir/include/mlir/Linker/LinkerInterface.h

@@ -139,6 +139,9 @@ struct SymbolLinkerInterface : LinkerInterface {
   /// Determines if the given operation is eligible for linking.
   virtual bool canBeLinked(Operation *op) const = 0;
 
+  /// Returns the symbol for the given operation.
+  virtual StringRef getSymbol(Operation *op) const = 0;
+
   /// Checks if an operation conflicts with existing linked operations.
   virtual ConflictPair findConflict(Operation *src) const = 0;
 

mlir/lib/Dialect/LLVMIR/IR/LLVMLinkerInterface.cpp

@@ -217,6 +217,8 @@ struct LLVMSymbolLinkerInterface
     return isa<LLVM::GlobalOp>(op) || isa<LLVM::LLVMFuncOp>(op);
   }
 
+  StringRef getSymbol(Operation *op) const override { return symbol(op); }
+
   ConflictPair findConflict(Operation *src) const override {
     assert(canBeLinked(src) && "expected linkable operation");
 

mlir/lib/Linker/Linker.cpp

@@ -69,10 +69,33 @@ LogicalResult Linker::process(ModuleOp src, unsigned flags) {
   return iface->process(src, flags);
 }
 
-OwningOpRef<ModuleOp> Linker::link() {
+OwningOpRef<ModuleOp> Linker::link(bool sortSymbols) {
   ModuleOp mod = composite.get();
   if (failed(getModuleLinkerInterface(mod)->link(mod)))
     return nullptr;
+
+  if (sortSymbols) {
+    std::vector<Operation *> symbols;
+
+    mod->walk([&](Operation *op) {
+      if (auto iface = dyn_cast<SymbolLinkerInterface>(op->getDialect())) {
+        if (iface->canBeLinked(op)) {
+          symbols.push_back(op);
+        }
+      }
+    });
+
+    llvm::stable_sort(symbols, [](Operation *lhs, Operation *rhs) {
+      auto lhsSym = cast<SymbolLinkerInterface>(lhs->getDialect());
+      auto rhsSym = cast<SymbolLinkerInterface>(rhs->getDialect());
+      return lhsSym->getSymbol(lhs) < rhsSym->getSymbol(rhs);
+    });
+
+    for (Operation *symbol : llvm::reverse(symbols)) {
+      symbol->moveBefore(&mod.front());
+    }
+  }
+
   return std::move(composite);
 }
 

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

@@ -69,6 +69,13 @@ struct LinkerCLOptions : public LinkerOptions {
   }
   StringRef outputSplitMarker() const { return outputSplitMarkerFlag; }
 
+  /// Sort symbols in the output module.
+  LinkerCLOptions &sortSymbols(bool sort) {
+    sortSymbolsFlag = sort;
+    return *this;
+  }
+  bool shouldSortSymbols() const { return sortSymbolsFlag; }
+
   /// Creates and initializes a LinkerOptions from command line options.
   /// These options are static but use ExternalStorage to initialize the
   /// members of the LinkerOptions class.
@@ -104,6 +111,10 @@ struct LinkerCLOptions : public LinkerOptions {
         }),
         cl::location(splitInputFileFlag), cl::init(""), cl::cat(getCategory()));
 
+    static cl::opt<bool, /*ExternalStorage=*/true> clSortSymbols(
+        "sort-symbols", cl::desc("Sort symbols in the output module"),
+        cl::location(sortSymbolsFlag), cl::init(false), cl::cat(getCategory()));
+
     static cl::opt<std::string, /*ExternalStorage=*/true> clOutputSplitMarker(
         "output-split-marker",
         cl::desc("Split marker to use for merging the ouput"),
@@ -136,6 +147,9 @@ struct LinkerCLOptions : public LinkerOptions {
   /// Show the registered dialects before trying to load the input file.
   bool showDialectsFlag = false;
 
+  /// Sort symbols in the output module.
+  bool sortSymbolsFlag = false;
+
   /// Split the input file based on the given marker into chunks and process
   /// each chunk independently. Input is not split if empty.
   std::string splitInputFileFlag = "";
@@ -265,7 +279,7 @@ LogicalResult mlir::MlirLinkMain(int argc, char **argv,
   if (failed(proc.linkFiles(options.inputFiles)))
     return failure();
 
-  OwningOpRef<ModuleOp> composite = linker.link();
+  OwningOpRef<ModuleOp> composite = linker.link(options.shouldSortSymbols());
   if (failed(verify(composite.get(), true))) {
     return composite->emitError("verification after linking failed");
   }