Swift Optimizer: add infrastructure for module passes.

To add a module pass in `Passes.def` use the new `SWIFT_MODULE_PASS` macro.
On the swift side, create a `ModulePass`.
It’s run function receives a `ModulePassContext`, which provides access to all functions of a module.
But it doesn't provide any APIs to modify functions.
In order to modify a function, a module pass must use `ModulePassContext.transform(function:)`.

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/PassManager/CMakeLists.txt

@@ -7,6 +7,7 @@
 # See http://swift.org/CONTRIBUTORS.txt for Swift project authors
 
 swift_compiler_sources(Optimizer
+  ModulePassContext.swift
   Passes.swift
   PassContext.swift
   PassRegistration.swift)

SwiftCompilerSources/Sources/Optimizer/PassManager/ModulePassContext.swift

@@ -0,0 +1,51 @@
+//===--- ModulePassContext.swift ------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+import OptimizerBridging
+
+/// The context which is passed to a `ModulePass`'s run-function.
+///
+/// It provides access to all functions of a module, but it doesn't provide any
+/// APIs to modify functions.
+/// In order to modify a function, a module pass must use `transform(function:)`.
+struct ModulePassContext {
+  let _bridged: BridgedPassContext
+
+  struct FunctionList : CollectionLikeSequence, IteratorProtocol {
+    private var currentFunction: Function?
+    
+    fileprivate init(first: Function?) { currentFunction = first }
+
+    mutating func next() -> Function? {
+      if let f = currentFunction {
+        currentFunction = PassContext_nextFunctionInModule(f.bridged).function
+        return f
+      }
+      return nil
+    }
+  }
+  
+  var functions: FunctionList {
+    FunctionList(first: PassContext_firstFunctionInModule(_bridged).function)
+  }
+  
+  /// Run a closure with a `PassContext` for a function, which allows to modify that function.
+  ///
+  /// Only a single `transform` can be alive at the same time, i.e. it's not allowed to nest
+  /// calls to `transform`.
+  func transform(function: Function, _ runOnFunction: (PassContext) -> ()) {
+    PassContext_beginTransformFunction(function.bridged, _bridged)
+    runOnFunction(PassContext(_bridged: _bridged))
+    PassContext_endTransformFunction(_bridged);
+  }
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -21,6 +21,14 @@ public func initializeSwiftModules() {
   initializeSwiftParseModules()
 }
 
+private func registerPass(
+      _ pass: ModulePass,
+      _ runFn: @escaping (@convention(c) (BridgedPassContext) -> ())) {
+  pass.name._withStringRef { nameStr in
+    SILPassManager_registerModulePass(nameStr, runFn)
+  }
+}
+
 private func registerPass(
       _ pass: FunctionPass,
       _ runFn: @escaping (@convention(c) (BridgedFunctionPassCtxt) -> ())) {

SwiftCompilerSources/Sources/Optimizer/PassManager/Passes.swift

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 import SIL
+import OptimizerBridging
 
 struct FunctionPass {
 
@@ -47,3 +48,19 @@ struct InstructionPass<InstType: Instruction> {
   }
 }
 
+struct ModulePass {
+
+  let name: String
+  let runFunction: (ModulePassContext) -> ()
+
+  public init(name: String,
+              _ runFunction: @escaping (ModulePassContext) -> ()) {
+    self.name = name
+    self.runFunction = runFunction
+  }
+
+  func run(_ bridgedCtxt: BridgedPassContext) {
+    let context = ModulePassContext(_bridged: bridgedCtxt)
+    runFunction(context)
+  }
+}

include/swift/SILOptimizer/OptimizerBridging.h

@@ -77,12 +77,14 @@ typedef struct {
   void * _Nonnull functionInfo;
 } BridgedRCIdentityFunctionInfo;
 
+typedef void (* _Nonnull BridgedModulePassRunFn)(BridgedPassContext);
 typedef void (* _Nonnull BridgedFunctionPassRunFn)(BridgedFunctionPassCtxt);
 typedef void (* _Nonnull BridgedInstructionPassRunFn)(BridgedInstructionPassCtxt);
 
+void SILPassManager_registerModulePass(llvm::StringRef name,
+                                       BridgedModulePassRunFn runFn);
 void SILPassManager_registerFunctionPass(llvm::StringRef name,
                                          BridgedFunctionPassRunFn runFn);
-
 void SILCombine_registerInstructionPass(llvm::StringRef name,
                                         BridgedInstructionPassRunFn runFn);
 
@@ -156,6 +158,15 @@ BridgedSubstitutionMap
 PassContext_getContextSubstitutionMap(BridgedPassContext context,
                                       BridgedType bridgedType);
 
+void PassContext_beginTransformFunction(BridgedFunction function,
+                                        BridgedPassContext ctxt);
+void PassContext_endTransformFunction(BridgedPassContext ctxt);
+
+OptionalBridgedFunction
+PassContext_firstFunctionInModule(BridgedPassContext context);
+OptionalBridgedFunction
+PassContext_nextFunctionInModule(BridgedFunction function);
+
 OptionalBridgedFunction
 PassContext_loadFunction(BridgedPassContext context, llvm::StringRef name);
 

include/swift/SILOptimizer/PassManager/PassManager.h

@@ -108,17 +108,35 @@ class SwiftPassInvocation {
   /// Called by the pass when changes are made to the SIL.
   void notifyChanges(SILAnalysis::InvalidationKind invalidationKind);
 
+  /// Called by the pass manager before the pass starts running.
+  void startModulePassRun(SILModuleTransform *transform);
+
   /// Called by the pass manager before the pass starts running.
   void startFunctionPassRun(SILFunctionTransform *transform);
 
   /// Called by the SILCombiner before the instruction pass starts running.
   void startInstructionPassRun(SILInstruction *inst);
 
+  /// Called by the pass manager when the pass has finished.
+  void finishedModulePassRun();
+
   /// Called by the pass manager when the pass has finished.
   void finishedFunctionPassRun();
 
   /// Called by the SILCombiner when the instruction pass has finished.
   void finishedInstructionPassRun();
+  
+  void beginTransformFunction(SILFunction *function) {
+    assert(!this->function && transform && "not running a module pass");
+    this->function = function;
+  }
+
+  void endTransformFunction() {
+    assert(function && "beginTransformFunction not called");
+    function = nullptr;
+    assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
+    assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
+  }
 };
 
 /// The SIL pass manager.

include/swift/SILOptimizer/PassManager/Passes.def

@@ -43,6 +43,17 @@
 #define IRGEN_PASS(Id, Tag, Description) PASS(Id, Tag, Description)
 #endif
 
+/// SWIFT_MODULE_PASS(Id, Tag, Description)
+///   This macro follows the same conventions as PASS(Id, Tag, Description),
+///   but is used for module passes which are implemented in Swift.
+///
+///   No further code is need on the C++ side. On the swift swift a module
+///   pass with the same name must be registered with 'registerPass()'.
+///
+#ifndef SWIFT_MODULE_PASS
+#define SWIFT_MODULE_PASS(Id, Tag, Description) PASS(Id, Tag, Description)
+#endif
+
 /// SWIFT_FUNCTION_PASS(Id, Tag, Description)
 ///   This macro follows the same conventions as PASS(Id, Tag, Description),
 ///   but is used for function passes which are implemented in Swift.
@@ -468,6 +479,7 @@ SWIFT_INSTRUCTION_PASS_WITH_LEGACY(StrongRetainInst, "simplify-strong_retain")
 SWIFT_INSTRUCTION_PASS_WITH_LEGACY(StrongReleaseInst, "simplify-strong_release")
 
 #undef IRGEN_PASS
+#undef SWIFT_MODULE_PASS
 #undef SWIFT_FUNCTION_PASS
 #undef SWIFT_FUNCTION_PASS_WITH_LEGACY
 #undef SWIFT_INSTRUCTION_PASS

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -727,11 +727,14 @@ void SILPassManager::runModulePass(unsigned TransIdx) {
     verifyAnalyses();
   }
 
+  swiftPassInvocation.startModulePassRun(SMT);
+
   llvm::sys::TimePoint<> StartTime = std::chrono::system_clock::now();
   assert(analysesUnlocked() && "Expected all analyses to be unlocked!");
   SMT->run();
   assert(analysesUnlocked() && "Expected all analyses to be unlocked!");
   Mod->flushDeletedInsts();
+  swiftPassInvocation.finishedModulePassRun();
 
   std::chrono::nanoseconds duration = std::chrono::system_clock::now() - StartTime;
   totalPassRuntime += duration;
@@ -1279,6 +1282,12 @@ void SwiftPassInvocation::freeNodeSet(NodeSet *set) {
   }
 }
 
+void SwiftPassInvocation::startModulePassRun(SILModuleTransform *transform) {
+  assert(!this->function && !this->transform && "a pass is already running");
+  this->function = nullptr;
+  this->transform = transform;
+}
+
 void SwiftPassInvocation::startFunctionPassRun(SILFunctionTransform *transform) {
   assert(!this->function && !this->transform && "a pass is already running");
   this->function = transform->getFunction();
@@ -1290,6 +1299,12 @@ void SwiftPassInvocation::startInstructionPassRun(SILInstruction *inst) {
          "running instruction pass on wrong function");
 }
 
+void SwiftPassInvocation::finishedModulePassRun() {
+  endPassRunChecks();
+  assert(!function && transform && "not running a pass");
+  transform = nullptr;
+}
+
 void SwiftPassInvocation::finishedFunctionPassRun() {
   endPassRunChecks();
   assert(function && transform && "not running a pass");
@@ -1533,6 +1548,31 @@ PassContext_getContextSubstitutionMap(BridgedPassContext context,
   return {type.getASTType()->getContextSubstitutionMap(m, ntd).getOpaqueValue()};
 }
 
+void PassContext_beginTransformFunction(BridgedFunction function, BridgedPassContext ctxt) {
+  castToPassInvocation(ctxt)->beginTransformFunction(castToFunction(function));
+}
+
+void PassContext_endTransformFunction(BridgedPassContext ctxt) {
+  castToPassInvocation(ctxt)->endTransformFunction();
+}
+
+OptionalBridgedFunction
+PassContext_firstFunctionInModule(BridgedPassContext context) {
+  SILModule *mod = castToPassInvocation(context)->getPassManager()->getModule();
+  if (mod->getFunctions().empty())
+    return {nullptr};
+  return {&*mod->getFunctions().begin()};
+}
+
+OptionalBridgedFunction
+PassContext_nextFunctionInModule(BridgedFunction function) {
+  auto *f = castToFunction(function);
+  auto nextIter = std::next(f->getIterator());
+  if (nextIter == f->getModule().getFunctions().end())
+    return {nullptr};
+  return {&*nextIter};
+}
+
 OptionalBridgedFunction
 PassContext_loadFunction(BridgedPassContext context, StringRef name) {
   SILModule *mod = castToPassInvocation(context)->getPassManager()->getModule();

lib/SILOptimizer/PassManager/Passes.cpp

@@ -217,9 +217,30 @@ void swift::runSILLoweringPasses(SILModule &Module) {
 }
 
 /// Registered briged pass run functions.
-static llvm::StringMap<BridgedFunctionPassRunFn> bridgedPassRunFunctions;
+static llvm::StringMap<BridgedModulePassRunFn> bridgedModulePassRunFunctions;
+static llvm::StringMap<BridgedFunctionPassRunFn> bridgedFunctionPassRunFunctions;
 static bool passesRegistered = false;
 
+/// Runs a bridged module pass.
+///
+/// \p runFunction is a cache for the run function, so that it has to be looked
+/// up only once in bridgedPassRunFunctions.
+static void runBridgedModulePass(BridgedModulePassRunFn &runFunction,
+                                 SILPassManager *passManager,
+                                 StringRef passName) {
+  if (!runFunction) {
+    runFunction = bridgedModulePassRunFunctions[passName];
+    if (!runFunction) {
+      if (passesRegistered) {
+        llvm::errs() << "Swift pass " << passName << " is not registered\n";
+        abort();
+      }
+      return;
+    }
+  }
+  runFunction({passManager->getSwiftPassInvocation()});
+}
+
 /// Runs a bridged function pass.
 ///
 /// \p runFunction is a cache for the run function, so that it has to be looked
@@ -228,7 +249,7 @@ static void runBridgedFunctionPass(BridgedFunctionPassRunFn &runFunction,
                                    SILPassManager *passManager,
                                    SILFunction *f, StringRef passName) {
   if (!runFunction) {
-    runFunction = bridgedPassRunFunctions[passName];
+    runFunction = bridgedFunctionPassRunFunctions[passName];
     if (!runFunction) {
       if (passesRegistered) {
         llvm::errs() << "Swift pass " << passName << " is not registered\n";
@@ -245,9 +266,15 @@ static void runBridgedFunctionPass(BridgedFunctionPassRunFn &runFunction,
 }
 
 // Called from initializeSwiftModules().
+void SILPassManager_registerModulePass(llvm::StringRef name,
+                                       BridgedModulePassRunFn runFn) {
+  bridgedModulePassRunFunctions[name] = runFn;
+  passesRegistered = true;
+}
+
 void SILPassManager_registerFunctionPass(llvm::StringRef name,
                                          BridgedFunctionPassRunFn runFn) {
-  bridgedPassRunFunctions[name] = runFn;
+  bridgedFunctionPassRunFunctions[name] = runFn;
   passesRegistered = true;
 }
 
@@ -263,6 +290,16 @@ BridgedFunctionPassRunFn ID##Pass::runFunction = nullptr;                  \
 #define PASS(ID, TAG, DESCRIPTION)
 #define SWIFT_INSTRUCTION_PASS(INST, TAG)
 
+#define SWIFT_MODULE_PASS(ID, TAG, DESCRIPTION) \
+class ID##Pass : public SILModuleTransform {                               \
+  static BridgedModulePassRunFn runFunction;                               \
+  void run() override {                                                    \
+    runBridgedModulePass(runFunction, PM, TAG);                            \
+  }                                                                        \
+};                                                                         \
+BridgedModulePassRunFn ID##Pass::runFunction = nullptr;                    \
+SILTransform *swift::create##ID() { return new ID##Pass(); }               \
+
 #define SWIFT_FUNCTION_PASS(ID, TAG, DESCRIPTION) \
 SWIFT_FUNCTION_PASS_COMMON(ID, TAG)                                        \
 SILTransform *swift::create##ID() { return new ID##Pass(); }               \