Add new loop invariant code motion.

Author: MAJKFL

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/LoopInvariantCodeMotion.swift

@@ -106,6 +106,7 @@ private func registerSwiftPasses() {
   registerPass(tempLValueElimination, { tempLValueElimination.run($0) })
   registerPass(generalClosureSpecialization, { generalClosureSpecialization.run($0) })
   registerPass(autodiffClosureSpecialization, { autodiffClosureSpecialization.run($0) })
+  registerPass(loopInvariantCodeMotionPass, { loopInvariantCodeMotionPass.run($0) })
 
   // Instruction passes
   registerForSILCombine(BeginBorrowInst.self,      { run(BeginBorrowInst.self, $0) })

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -484,6 +484,22 @@ extension Instruction {
     }
     return false
   }
+  
+  /// Returns true if `otherInst` is in the same block and is strictly dominated by this instruction or
+  /// the parent block of the instruction dominates parent block of `otherInst`.
+  func dominates(
+    _ otherInst: Instruction,
+    _ domTree: DominatorTree
+  ) -> Bool {
+    if parentBlock == otherInst.parentBlock {
+      return dominatesInSameBlock(otherInst)
+    } else {
+      return parentBlock.dominates(
+        otherInst.parentBlock,
+        domTree
+      )
+    }
+  }
 
   /// If this instruction uses a (single) existential archetype, i.e. it has a type-dependent operand,
   /// returns the concrete type if it is known.
@@ -583,36 +599,22 @@ extension StoreInst {
 extension LoadInst {
   @discardableResult
   func trySplit(_ context: FunctionPassContext) -> Bool {
-    var elements = [Value]()
-    let builder = Builder(before: self, context)
     if type.isStruct {
-      if (type.nominal as! StructDecl).hasUnreferenceableStorage {
-        return false
-      }
-      guard let fields = type.getNominalFields(in: parentFunction) else {
+      guard !(type.nominal as! StructDecl).hasUnreferenceableStorage,
+            let fields = type.getNominalFields(in: parentFunction) else {
         return false
       }
-      for idx in 0..<fields.count {
-        let fieldAddr = builder.createStructElementAddr(structAddress: address, fieldIndex: idx)
-        let splitLoad = builder.createLoad(fromAddress: fieldAddr, ownership: self.splitOwnership(for: fieldAddr))
-        elements.append(splitLoad)
-      }
-      let newStruct = builder.createStruct(type: self.type, elements: elements)
-      self.replace(with: newStruct, context)
+      
+      _ = splitStruct(fields: fields, context)
+      
       return true
     } else if type.isTuple {
-      var elements = [Value]()
-      let builder = Builder(before: self, context)
-      for idx in 0..<type.tupleElements.count {
-        let fieldAddr = builder.createTupleElementAddr(tupleAddress: address, elementIndex: idx)
-        let splitLoad = builder.createLoad(fromAddress: fieldAddr, ownership: self.splitOwnership(for: fieldAddr))
-        elements.append(splitLoad)
-      }
-      let newTuple = builder.createTuple(type: self.type, elements: elements)
-      self.replace(with: newTuple, context)
+      _ = splitTuple(context)
+      
       return true
+    } else {
+      return false
     }
-    return false
   }
 
   private func splitOwnership(for fieldValue: Value) -> LoadOwnership {
@@ -623,6 +625,70 @@ extension LoadInst {
       return fieldValue.type.isTrivial(in: parentFunction) ? .trivial : self.loadOwnership
     }
   }
+  
+  func trySplit(
+    alongPath projectionPath: SmallProjectionPath,
+    _ context: FunctionPassContext
+  ) -> [LoadInst]? {
+    if projectionPath.isEmpty {
+      return nil
+    }
+    
+    let (fieldKind, index, pathRemainder) = projectionPath.pop()
+    
+    var elements: [LoadInst]
+    
+    switch fieldKind {
+    case .structField where type.isStruct:
+      guard !(type.nominal as! StructDecl).hasUnreferenceableStorage,
+            let fields = type.getNominalFields(in: parentFunction) else {
+        return nil
+      }
+      
+      elements = splitStruct(fields: fields, context)
+    case .tupleField where type.isTuple:
+      elements = splitTuple(context)
+    default:
+      return nil
+    }
+    
+    if let recursiveSplitLoad = elements[index].trySplit(alongPath: pathRemainder, context) {
+      elements.remove(at: index)
+      elements += recursiveSplitLoad
+    }
+    
+    return elements
+  }
+  
+  private func splitStruct(fields: NominalFieldsArray, _ context: FunctionPassContext) -> [LoadInst] {
+    var elements = [LoadInst]()
+    let builder = Builder(before: self, context)
+    
+    for idx in 0..<fields.count {
+      let fieldAddr = builder.createStructElementAddr(structAddress: address, fieldIndex: idx)
+      let splitLoad = builder.createLoad(fromAddress: fieldAddr, ownership: self.splitOwnership(for: fieldAddr))
+      elements.append(splitLoad)
+    }
+    let newStruct = builder.createStruct(type: self.type, elements: elements)
+    self.replace(with: newStruct, context)
+    
+    return elements
+  }
+  
+  private func splitTuple(_ context: FunctionPassContext) -> [LoadInst] {
+    var elements = [LoadInst]()
+    let builder = Builder(before: self, context)
+    
+    for idx in 0..<type.tupleElements.count {
+      let fieldAddr = builder.createTupleElementAddr(tupleAddress: address, elementIndex: idx)
+      let splitLoad = builder.createLoad(fromAddress: fieldAddr, ownership: self.splitOwnership(for: fieldAddr))
+      elements.append(splitLoad)
+    }
+    let newTuple = builder.createTuple(type: self.type, elements: elements)
+    self.replace(with: newTuple, context)
+    
+    return elements
+  }
 }
 
 extension FunctionPassContext {

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -140,7 +140,7 @@ public struct Stack<Element> : CollectionLikeSequence {
   public mutating func deinitialize() { removeAll() }
 }
 
-extension Stack {
+public extension Stack {
   /// Mark a stack location for future iteration.
   ///
   /// TODO: Marker should be ~Escapable.
@@ -155,7 +155,7 @@ extension Stack {
     let low: Marker
     let high: Marker
 
-    init(in stack: Stack, low: Marker, high: Marker) {
+    public init(in stack: Stack, low: Marker, high: Marker) {
       if low.slab.data == nil {
         assert(low.index == 0, "invalid empty stack marker")
         // `low == nil` and `high == nil` is a valid empty segment,
@@ -173,7 +173,7 @@ extension Stack {
       self.high = high
     }
 
-    func makeIterator() -> Stack.Iterator {
+    public func makeIterator() -> Stack.Iterator {
       return Iterator(slab: low.slab, index: low.index,
                       lastSlab: high.slab, endIndex: high.index)
     }
@@ -219,3 +219,87 @@ extension Stack {
     }
   }
 }
+
+public struct StackWithCount<Element> : CollectionLikeSequence {
+  public private(set) var count = 0
+  private var underlyingStack: Stack<Element>
+  
+  public typealias Iterator = Stack<Element>.Iterator
+  
+  public init(_ context: some Context) {
+    self.underlyingStack = Stack<Element>(context)
+  }
+  
+  public func makeIterator() -> Stack<Element>.Iterator {
+    underlyingStack.makeIterator()
+  }
+  
+  public var first: Element? { underlyingStack.first }
+  public var last: Element? { underlyingStack.last }
+
+  public mutating func push(_ element: Element) {
+    count += 1
+    underlyingStack.push(element)
+  }
+
+  /// The same as `push` to provide an Array-like append API.
+  public mutating func append(_ element: Element) { push(element) }
+
+  public mutating func append<S: Sequence>(contentsOf other: S) where S.Element == Element {
+    for elem in other {
+      append(elem)
+    }
+  }
+
+  public var isEmpty: Bool { underlyingStack.isEmpty }
+
+  public mutating func pop() -> Element? {
+    if underlyingStack.isEmpty {
+      return nil
+    }
+    
+    count -= 1
+    return underlyingStack.pop()
+  }
+  
+  public mutating func removeAll() {
+    underlyingStack.removeAll()
+  }
+
+  /// TODO: once we have move-only types, make this a real deinit.
+  public mutating func deinitialize() { removeAll() }
+}
+
+public extension StackWithCount {
+  typealias Marker = Stack<Element>.Marker
+
+  struct Segment : CollectionLikeSequence {
+    var underlyingSegment: Stack<Element>.Segment
+    
+    public init(in stack: StackWithCount, low: Marker, high: Marker) {
+      underlyingSegment = Stack<Element>.Segment(in: stack.underlyingStack, low: low, high: high)
+    }
+    
+    public func makeIterator() -> StackWithCount.Iterator {
+      return underlyingSegment.makeIterator()
+    }
+  }
+  
+  var top: Marker { underlyingStack.top }
+
+  func assertValid(marker: Marker) { underlyingStack.assertValid(marker: marker) }
+
+  mutating func withMarker<R>(
+    _ body: (inout Stack<Element>, Marker) throws -> R) rethrows -> R {
+    return try underlyingStack.withMarker(body)
+  }
+
+  mutating func withMarker<R>(
+    pushElements body: (inout Stack<Element>) throws -> R,
+    withNewElements handleNewElements: ((Segment) -> ())
+  ) rethrows -> R {
+    return try underlyingStack.withMarker(pushElements: body) { [self] segment in
+      handleNewElements(Segment(in: self, low: segment.low, high: segment.high))
+    }
+  }
+}

SwiftCompilerSources/Sources/SIL/DataStructures/Stack.swift

@@ -505,6 +505,14 @@ public struct SideEffects : CustomStringConvertible, NoReflectionChildren {
     /// This is true when the function (or a callee, transitively) contains a
     /// deinit barrier instruction.
     public var isDeinitBarrier: Bool
+    
+    public static var noEffects: GlobalEffects {
+      return GlobalEffects(memory: .noEffects, ownership: .noEffects, allocates: false, isDeinitBarrier: false)
+    }
+    
+    public var isOnlyReading: Bool {
+      return !memory.write && ownership == .noEffects && !allocates && !isDeinitBarrier
+    }
 
     /// When called with default arguments, it creates an "effect-free" GlobalEffects.
     public init(memory: Memory = Memory(read: false, write: false),
@@ -643,6 +651,10 @@ public struct SideEffects : CustomStringConvertible, NoReflectionChildren {
       copy = copy || other.copy
       destroy = destroy || other.destroy
     }
+    
+    public static var noEffects: Ownership {
+      return Ownership(copy: false, destroy: false)
+    }
 
     public static var worstEffects: Ownership {
       Ownership(copy: true, destroy: true)

SwiftCompilerSources/Sources/SIL/Effects.swift

@@ -96,6 +96,11 @@ public class Instruction : CustomStringConvertible, Hashable {
     BridgedContext.moveInstructionBefore(bridged, otherInstruction.bridged)
     context.notifyInstructionsChanged()
   }
+  
+  public final func copy(before otherInstruction: Instruction, _ context: some MutatingContext) {
+    BridgedContext.copyInstructionBefore(bridged, otherInstruction.bridged)
+    context.notifyInstructionsChanged()
+  }
 
   public var mayTrap: Bool { false }
 
@@ -177,6 +182,10 @@ public class Instruction : CustomStringConvertible, Hashable {
   public static func ==(lhs: Instruction, rhs: Instruction) -> Bool {
     lhs === rhs
   }
+  
+  public func isIdenticalTo(_ otherInst: Instruction) -> Bool {
+    return bridged.isIdenticalTo(otherInst.bridged)
+  }
 
   public func hash(into hasher: inout Hasher) {
     hasher.combine(ObjectIdentifier(self))

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -392,6 +392,11 @@ public struct AccessPath : CustomStringConvertible, Hashable {
   public func isEqualOrContains(_ other: AccessPath) -> Bool {
     return getProjection(to: other) != nil
   }
+  
+  /// Returns true if this access contains `other` access and is not equal.
+  public func contains(_ other: AccessPath) -> Bool {
+    return !(getProjection(to: other)?.isEmpty ?? true)
+  }
 
   public var materializableProjectionPath: SmallProjectionPath? {
     if projectionPath.isMaterializable {

SwiftCompilerSources/Sources/SIL/Utilities/AccessUtils.swift

@@ -534,6 +534,18 @@ public struct SmallProjectionPath : Hashable, CustomStringConvertible, NoReflect
       return false
     }
   }
+  
+  public var isConstant: Bool {
+    let (kind, _, subPath) = pop()
+    switch kind {
+    case .root:
+      return true
+    case .structField, .tupleField, .enumCase, .classField, .existential, .indexedElement:
+      return subPath.isConstant
+    default:
+      return false
+    }
+  }
 }
 
 //===----------------------------------------------------------------------===//

SwiftCompilerSources/Sources/SIL/Utilities/SmallProjectionPath.swift

@@ -51,7 +51,7 @@
 ///   This macro follows the same conventions as PASS(Id, Tag, Description),
 ///   but is used for IRGen passes which are built outside of the
 ///   SILOptimizer library.
-///   
+///
 ///   An IRGen pass is created by IRGen and needs to be registered with the pass
 ///   manager dynamically.
 #ifndef IRGEN_PASS
@@ -148,6 +148,8 @@ PASS(TempRValueElimination, "temp-rvalue-elimination",
      "Remove short-lived immutable temporary copies")
 PASS(TempLValueElimination, "temp-lvalue-elimination",
      "Remove short-lived immutable temporary l-values")
+PASS(LoopInvariantCodeMotion, "loop-invariant-code-motion",
+     "New Loop Invariant Code Motion")
 
 // NOTE - ExperimentalSwiftBasedClosureSpecialization and AutodiffClosureSpecialization are a WIP
 PASS(ExperimentalSwiftBasedClosureSpecialization, "experimental-swift-based-closure-specialization",