New SIL Loop Invariant Code Motion pass

SwiftCompilerSources/Sources/Optimizer/Analysis/CMakeLists.txt

@@ -9,6 +9,7 @@
 swift_compiler_sources(Optimizer
   AliasAnalysis.swift
   CalleeAnalysis.swift
+  LoopTree.swift
   DeadEndBlocksAnalysis.swift
   DominatorTree.swift
   PostDominatorTree.swift)

SwiftCompilerSources/Sources/Optimizer/Analysis/DominatorTree.swift

@@ -15,6 +15,31 @@ import OptimizerBridging
 
 struct DominatorTree {
   let bridged: BridgedDomTree
+
+  func getChildren(of block: BasicBlock) -> DomChildren {
+    return DomChildren(bridgedDomTree: bridged, bb: block)
+  }
+}
+
+struct DomChildren: BridgedRandomAccessCollection {
+  private let bridgedDomTree: BridgedDomTree
+  private let block: BasicBlock
+
+  let count: Int
+
+  var startIndex: Int { return 0 }
+  var endIndex: Int { return count }
+
+  init(bridgedDomTree: BridgedDomTree, bb: BasicBlock) {
+    self.bridgedDomTree = bridgedDomTree
+    self.block = bb
+    self.count = bridgedDomTree.getNumberOfChildren(bb.bridged)
+  }
+
+  subscript(_ index: Int) -> BasicBlock {
+    assert(index >= startIndex && index < endIndex)
+    return bridgedDomTree.getChildAt(block.bridged, index).block
+  }
 }
 
 extension BasicBlock {

SwiftCompilerSources/Sources/Optimizer/Analysis/LoopTree.swift

@@ -0,0 +1,213 @@
+//===--- LoopTree.swift ---------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2025 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
+
+/// Describes top level loops.
+struct LoopTree {
+  fileprivate let bridged: BridgedLoopTree
+
+  let loops: TopLevelLoopArray
+
+  init(bridged: BridgedLoopTree, context: FunctionPassContext) {
+    self.bridged = bridged
+    self.loops = TopLevelLoopArray(bridged)
+  }
+}
+
+/// Describes a loop with its children.
+struct Loop {
+  private let bridged: BridgedLoop
+
+  let innerLoops: LoopArray
+  let loopBlocks: LoopBlocks
+
+  var exitingAndLatchBlocks: some Sequence<BasicBlock> {
+    return header.predecessors.lazy
+      .filter { predecessor in
+        contains(block: predecessor) && !isLoopExiting(loopBlock: predecessor)
+      } + exitingBlocks
+  }
+
+  /// Exit blocks of the loop.
+  ///
+  /// - Note: Some exit blocks will be duplicated if the loop has critical edges.
+  var exitBlocks: some Sequence<BasicBlock> {
+    return loopBlocks.lazy
+      .flatMap(\.successors)
+      .filter { !contains(block: $0) }
+  }
+
+  var exitingBlocks: some Sequence<BasicBlock> {
+    return loopBlocks.lazy
+      .filter { isLoopExiting(loopBlock: $0) }
+  }
+
+  init(bridged: BridgedLoop) {
+    self.bridged = bridged
+    self.innerLoops = LoopArray(bridged)
+    self.loopBlocks = LoopBlocks(bridged)
+  }
+
+  var preheader: BasicBlock? {
+    bridged.getPreheader().block
+  }
+
+  var header: BasicBlock {
+    bridged.getHeader().block
+  }
+
+  /// Returns `true` if the loop has exactly one exit block.
+  var hasSingleExitBlock: Bool {
+    return exitBlocks.singleElement != nil
+  }
+
+  var hasNoExitBlocks: Bool {
+    return exitBlocks.isEmpty
+  }
+
+  private func isLoopExiting(loopBlock: BasicBlock) -> Bool {
+    return loopBlock.successors.contains { !contains(block: $0) }
+  }
+
+  func getBlocksThatDominateAllExitingAndLatchBlocks(_ context: FunctionPassContext) -> [BasicBlock] {
+    var result: [BasicBlock] = []
+    var cachedExitingAndLatchBlocks = Stack<BasicBlock>(context)
+    var workList = BasicBlockWorklist(context)
+    defer {
+      cachedExitingAndLatchBlocks.deinitialize()
+      workList.deinitialize()
+    }
+
+    cachedExitingAndLatchBlocks.append(contentsOf: exitingAndLatchBlocks)
+    workList.pushIfNotVisited(header)
+
+    while let block = workList.pop() {
+      guard cachedExitingAndLatchBlocks.allSatisfy({ exitBlock in
+        return block.dominates(exitBlock, context.dominatorTree)
+      }) else {
+        continue
+      }
+
+      result.append(block)
+
+      workList.pushIfNotVisited(contentsOf: context.dominatorTree.getChildren(of: block))
+    }
+
+    return result
+  }
+
+  func contains(block: BasicBlock) -> Bool {
+    return bridged.contains(block.bridged)
+  }
+
+  func splitCriticalExitingAndBackEdges(_ context: FunctionPassContext) {
+    for exitingOrLatchBlock in exitingAndLatchBlocks {
+      for (index, succesor) in exitingOrLatchBlock.successors.enumerated() where !contains(block: succesor) {
+        splitCriticalEdge(
+          from: exitingOrLatchBlock.terminator.parentBlock,
+          toEdgeIndex: index,
+          dominatorTree: context.dominatorTree,
+          loopTree: context.loopTree,
+          context
+        )
+      }
+    }
+  }
+}
+
+struct TopLevelLoopArray: BridgedRandomAccessCollection {
+  private let bridgedLoopTree: BridgedLoopTree
+
+  let count: Int
+
+  var startIndex: Int { return 0 }
+  var endIndex: Int { return count }
+
+  init(_ bridgedLoopTree: BridgedLoopTree) {
+    self.bridgedLoopTree = bridgedLoopTree
+    self.count = bridgedLoopTree.getTopLevelLoopCount()
+  }
+
+  subscript(_ index: Int) -> Loop {
+    assert(index >= startIndex && index < endIndex)
+    return Loop(bridged: bridgedLoopTree.getLoop(index))
+  }
+}
+
+struct LoopArray: BridgedRandomAccessCollection {
+  private let bridgedLoop: BridgedLoop
+
+  let count: Int
+
+  var startIndex: Int { return 0 }
+  var endIndex: Int { return count }
+
+  init(_ bridgedLoop: BridgedLoop) {
+    self.bridgedLoop = bridgedLoop
+    self.count = bridgedLoop.getInnerLoopCount()
+  }
+
+  subscript(_ index: Int) -> Loop {
+    assert(index >= startIndex && index < endIndex)
+    return Loop(bridged: bridgedLoop.getInnerLoop(index))
+  }
+}
+
+struct LoopBlocks: BridgedRandomAccessCollection {
+  private let bridgedLoop: BridgedLoop
+
+  let count: Int
+
+  var startIndex: Int { return 0 }
+  var endIndex: Int { return count }
+
+  init(_ bridgedLoop: BridgedLoop) {
+    self.bridgedLoop = bridgedLoop
+    self.count = bridgedLoop.getBasicBlockCount()
+  }
+
+  subscript(_ index: Int) -> BasicBlock {
+    assert(index >= startIndex && index < endIndex)
+    return bridgedLoop.getBasicBlock(index).block
+  }
+}
+
+func splitEdge(
+  from block: BasicBlock,
+  toEdgeIndex: Int,
+  dominatorTree: DominatorTree,
+  loopTree: LoopTree,
+  _ context: some MutatingContext
+) -> BasicBlock {
+  let result = loopTree.bridged.splitEdge(block.bridged, toEdgeIndex, dominatorTree.bridged).block
+
+  context.notifyBranchesChanged()
+  return result
+}
+
+/// If the specified edge is critical, the function returns inserted block. Otherwise returns `nil`.
+@discardableResult
+func splitCriticalEdge(
+  from block: BasicBlock,
+  toEdgeIndex: Int,
+  dominatorTree: DominatorTree,
+  loopTree: LoopTree,
+  _ context: some MutatingContext
+) -> BasicBlock? {
+  guard block.isCriticalEdge(edgeIndex: toEdgeIndex) else {
+    return nil
+  }
+
+  return splitEdge(from: block, toEdgeIndex: toEdgeIndex, dominatorTree: dominatorTree, loopTree: loopTree, context)
+}

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -311,10 +311,9 @@ private struct FunctionSpecializations {
       // This can happen if a previous run of the pass already created this specialization.
       return
     }
-    let cloner = SpecializationCloner(emptySpecializedFunction: specializedFunc, context)
-    cloner.cloneFunctionBody(from: original)
-
     context.buildSpecializedFunction(specializedFunction: specializedFunc) { (specializedFunc, specContext) in
+      cloneFunction(from: original, toEmpty: specializedFunc, specContext)
+
       replaceBoxWithStackArguments(in: specializedFunc, original: original, specContext)
     }
     context.notifyNewFunction(function: specializedFunc, derivedFrom: original)

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

@@ -25,6 +25,7 @@ swift_compiler_sources(Optimizer
   LifetimeDependenceDiagnostics.swift
   LifetimeDependenceInsertion.swift
   LifetimeDependenceScopeFixup.swift
+  LoopInvariantCodeMotion.swift
   ObjectOutliner.swift
   ObjCBridgingOptimization.swift
   MergeCondFails.swift

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -247,9 +247,10 @@ private func getOrCreateSpecializedFunction(
   context.buildSpecializedFunction(
     specializedFunction: specializedFunction,
     buildFn: { (emptySpecializedFunction, functionPassContext) in
-      let closureSpecCloner = SpecializationCloner(
-        emptySpecializedFunction: emptySpecializedFunction, functionPassContext)
+      var closureSpecCloner = Cloner(
+        cloneToEmptyFunction: emptySpecializedFunction, functionPassContext)
       closureSpecCloner.cloneAndSpecializeFunctionBody(using: pullbackClosureInfo)
+      closureSpecCloner.deinitialize()
     })
 
   return (specializedFunction, false)
@@ -731,15 +732,14 @@ private func markConvertedAndReabstractedClosuresAsUsed(
   }
 }
 
-extension SpecializationCloner {
+extension Cloner where Context == FunctionPassContext {
   fileprivate func cloneAndSpecializeFunctionBody(using pullbackClosureInfo: PullbackClosureInfo) {
     self.cloneEntryBlockArgsWithoutOrigClosures(usingOrigCalleeAt: pullbackClosureInfo)
 
     let (allSpecializedEntryBlockArgs, closureArgIndexToAllClonedReleasableClosures) =
       cloneAllClosures(at: pullbackClosureInfo)
 
-    self.cloneFunctionBody(
-      from: pullbackClosureInfo.pullbackFn, entryBlockArguments: allSpecializedEntryBlockArgs)
+    self.cloneFunctionBody(from: pullbackClosureInfo.pullbackFn, entryBlockArguments: allSpecializedEntryBlockArgs)
 
     self.insertCleanupCodeForClonedReleasableClosures(
       from: pullbackClosureInfo,
@@ -750,8 +750,8 @@ extension SpecializationCloner {
     usingOrigCalleeAt pullbackClosureInfo: PullbackClosureInfo
   ) {
     let originalEntryBlock = pullbackClosureInfo.pullbackFn.entryBlock
-    let clonedFunction = self.cloned
-    let clonedEntryBlock = self.entryBlock
+    let clonedFunction = self.targetFunction
+    let clonedEntryBlock = self.getOrCreateEntryBlock()
 
     originalEntryBlock.arguments
       .enumerated()
@@ -782,7 +782,8 @@ extension SpecializationCloner {
     )
   {
     func entryBlockArgsWithOrigClosuresSkipped() -> [Value?] {
-      var clonedNonClosureEntryBlockArgs = self.entryBlock.arguments.makeIterator()
+      let clonedEntryBlock = self.getOrCreateEntryBlock()
+      var clonedNonClosureEntryBlockArgs = clonedEntryBlock.arguments.makeIterator()
 
       return pullbackClosureInfo.pullbackFn
         .entryBlock
@@ -823,8 +824,8 @@ extension SpecializationCloner {
   {
     let (origToClonedValueMap, capturedArgRange) = self.addEntryBlockArgs(
       forValuesCapturedBy: closureArgDesc)
-    let clonedFunction = self.cloned
-    let clonedEntryBlock = self.entryBlock
+    let clonedFunction = self.targetFunction
+    let clonedEntryBlock = self.getOrCreateEntryBlock()
     let clonedClosureArgs = Array(clonedEntryBlock.arguments[capturedArgRange])
 
     let builder =
@@ -853,8 +854,8 @@ extension SpecializationCloner {
     -> (origToClonedValueMap: [HashableValue: Value], capturedArgRange: Range<Int>)
   {
     var origToClonedValueMap: [HashableValue: Value] = [:]
-    let clonedFunction = self.cloned
-    let clonedEntryBlock = self.entryBlock
+    let clonedFunction = self.targetFunction
+    let clonedEntryBlock = self.getOrCreateEntryBlock()
 
     let capturedArgRangeStart = clonedEntryBlock.arguments.count
 
@@ -908,8 +909,8 @@ extension SpecializationCloner {
       }
     }
 
-    if self.context.needFixStackNesting {
-      self.context.fixStackNesting(in: self.cloned)
+    if (self.context.needFixStackNesting) {
+      self.context.fixStackNesting(in: targetFunction)
     }
   }
 }
@@ -1425,12 +1426,10 @@ private struct PullbackClosureInfo {
   }
 
   func specializedCalleeName(_ context: FunctionPassContext) -> String {
-    let closureArgs = Array(self.closureArgDescriptors.map { $0.closure })
-    let closureIndices = Array(self.closureArgDescriptors.map { $0.closureArgIndex })
-
-    return context.mangle(
-      withClosureArguments: closureArgs, closureArgIndices: closureIndices,
-      from: pullbackFn)
+    let closureArgs = Array(self.closureArgDescriptors.map {
+      (argumentIndex: $0.closureArgIndex, argumentValue: $0.closure)
+    })
+    return context.mangle(withClosureArguments: closureArgs, from: pullbackFn)
   }
 }