RedundantLoadElimination: support redundant loads of array elements with non-constant index

SwiftCompilerSources/Sources/Optimizer/Analysis/AliasAnalysis.swift

@@ -334,7 +334,7 @@ struct AliasAnalysis {
     case .stack, .global, .argument, .storeBorrow:
       // Those access bases cannot be interior pointers of a borrowed value
       return .noEffects
-    case .pointer, .unidentified, .yield:
+    case .pointer, .index, .unidentified, .yield:
       // We don't know anything about this address -> get the conservative effects
       return defaultEffects(of: endBorrow, on: memLoc)
     case .box, .class, .tail:

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceDiagnostics.swift

@@ -337,7 +337,7 @@ private struct LifetimeVariable {
     case .pointer(let ptrToAddr):
       self.varDecl = nil
       self.sourceLoc = ptrToAddr.location.sourceLoc
-    case .unidentified:
+    case .index, .unidentified:
       self.varDecl = nil
       self.sourceLoc = nil
     }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/RedundantLoadElimination.swift

@@ -142,7 +142,7 @@ private extension LoadInst {
   }
 
   func isRedundant(complexityBudget: inout Int, _ context: FunctionPassContext) -> DataflowResult {
-    return isRedundant(at: address.accessPath, complexityBudget: &complexityBudget, context)
+    return isRedundant(at: address.constantAccessPath, complexityBudget: &complexityBudget, context)
   }
 
   func isRedundant(at accessPath: AccessPath, complexityBudget: inout Int, _ context: FunctionPassContext) -> DataflowResult {
@@ -282,7 +282,7 @@ private func provideValue(
   from availableValue: AvailableValue,
   _ context: FunctionPassContext
 ) -> Value {
-  let projectionPath = availableValue.address.accessPath.getMaterializableProjection(to: load.address.accessPath)!
+  let projectionPath = availableValue.address.constantAccessPath.getMaterializableProjection(to: load.address.constantAccessPath)!
 
   switch load.loadOwnership {
   case .unqualified:
@@ -483,7 +483,7 @@ private struct InstructionScanner {
         // This happens if the load is in a loop.
         return .available
       }
-      let precedingLoadPath = precedingLoad.address.accessPath
+      let precedingLoadPath = precedingLoad.address.constantAccessPath
       if precedingLoadPath.getMaterializableProjection(to: accessPath) != nil {
         availableValues.append(.viaLoad(precedingLoad))
         return .available
@@ -500,7 +500,7 @@ private struct InstructionScanner {
       if precedingStore.source is Undef {
         return .overwritten
       }
-      let precedingStorePath = precedingStore.destination.accessPath
+      let precedingStorePath = precedingStore.destination.constantAccessPath
       if precedingStorePath.getMaterializableProjection(to: accessPath) != nil {
         availableValues.append(.viaStore(precedingStore))
         return .available

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -467,7 +467,7 @@ private extension AccessBase {
         return .decidedInCaller(arg)
       case .yield, .pointer:
         return .unknown
-      case .unidentified:
+      case .index, .unidentified:
         // In the rare case of an unidentified access, just ignore it.
         // This should not happen in regular SIL, anyway.
         return .no

SwiftCompilerSources/Sources/Optimizer/TestPasses/AccessDumper.swift

@@ -67,8 +67,16 @@ private func printAccessInfo(address: Value) {
     print("  Scope: base")
   }
 
-  print("  Base: \(ap.base)")
-  print("  Path: \"\(ap.projectionPath)\"")
+  let constAp = address.constantAccessPath
+  if constAp == ap {
+    print("  Base: \(ap.base)")
+    print("  Path: \"\(ap.projectionPath)\"")
+  } else {
+    print("  nonconst-base: \(ap.base)")
+    print("  nonconst-path: \"\(ap.projectionPath)\"")
+    print("  const-base: \(constAp.base)")
+    print("  const-path: \"\(constAp.projectionPath)\"")
+  }
 
   var arw = AccessStoragePathVisitor()
   if !arw.visitAccessStorageRoots(of: ap) {
@@ -79,18 +87,40 @@ private func printAccessInfo(address: Value) {
 private func checkAliasInfo(forArgumentsOf apply: ApplyInst, expectDistinct: Bool) {
   let address1 = apply.arguments[0]
   let address2 = apply.arguments[1]
-  let path1 = address1.accessPath
-  let path2 = address2.accessPath
 
+  checkIsDistinct(path1: address1.accessPath,
+                  path2: address2.accessPath,
+                  expectDistinct: expectDistinct,
+                  instruction: apply)
+
+  if !expectDistinct {
+    // Also check all combinations with the constant variant of access paths.
+    // Note: we can't do that for "isDistinct" because "isDistinct" might be more conservative in one of the variants.
+    checkIsDistinct(path1: address1.constantAccessPath,
+                    path2: address2.constantAccessPath,
+                    expectDistinct: false,
+                    instruction: apply)
+    checkIsDistinct(path1: address1.accessPath,
+                    path2: address2.constantAccessPath,
+                    expectDistinct: false,
+                    instruction: apply)
+    checkIsDistinct(path1: address1.constantAccessPath,
+                    path2: address2.accessPath,
+                    expectDistinct: false,
+                    instruction: apply)
+  }
+}
+
+private func checkIsDistinct(path1: AccessPath, path2: AccessPath, expectDistinct: Bool, instruction: Instruction) {
   if path1.isDistinct(from: path2) != expectDistinct {
-    print("wrong isDistinct result of \(apply)")
+    print("wrong isDistinct result of \(instruction)")
   } else if path2.isDistinct(from: path1) != expectDistinct {
-    print("wrong reverse isDistinct result of \(apply)")
+    print("wrong reverse isDistinct result of \(instruction)")
   } else {
     return
   }
-  
+
   print("in function")
-  print(apply.parentFunction)
+  print(instruction.parentFunction)
   fatalError()
 }

SwiftCompilerSources/Sources/Optimizer/Utilities/AddressUtils.swift

@@ -464,7 +464,7 @@ enum AddressOwnershipLiveRange : CustomStringConvertible {
       }
     case .storeBorrow(let sb):
       return computeValueLiveRange(of: sb.source, context)
-    case .pointer, .unidentified:
+    case .pointer, .index, .unidentified:
       return nil
     }
   }

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -366,7 +366,7 @@ extension LifetimeDependence.Scope {
         return nil
       }
       self = scope
-    case .pointer, .unidentified:
+    case .pointer, .index, .unidentified:
       self = .unknown(address)
     }
   }
@@ -1094,7 +1094,7 @@ extension LifetimeDependenceDefUseWalker {
       break
     case .yield:
       return storeToYieldDependence(address: address, of: operand)
-    case .global, .class, .tail, .storeBorrow, .pointer, .unidentified:
+    case .global, .class, .tail, .storeBorrow, .pointer, .index, .unidentified:
       // An address produced by .storeBorrow should never be stored into.
       //
       // TODO: allow storing an immortal value into a global.

SwiftCompilerSources/Sources/SIL/Utilities/AccessUtils.swift

@@ -78,6 +78,11 @@ public enum AccessBase : CustomStringConvertible, Hashable {
   /// An address which is derived from a `Builtin.RawPointer`.
   case pointer(PointerToAddressInst)
 
+  // The result of an `index_addr` with a non-constant index.
+  // This can only occur in access paths returned by `Value.constantAccessPath`.
+  // In "regular" access paths such `index_addr` projections are contained in the `projectionPath` (`i*`).
+  case index(IndexAddrInst)
+
   /// The access base is some SIL pattern which does not fit into any other case.
   /// This should be a very rare situation.
   case unidentified
@@ -115,6 +120,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
       case .yield(let result): return "yield - \(result)"
       case .storeBorrow(let sb): return "storeBorrow - \(sb)"
       case .pointer(let p):    return "pointer - \(p)"
+      case .index(let ia):     return "index - \(ia)"
     }
   }
 
@@ -123,7 +129,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
     switch self {
       case .class, .tail:
         return true
-      case .box, .stack, .global, .argument, .yield, .storeBorrow, .pointer, .unidentified:
+      case .box, .stack, .global, .argument, .yield, .storeBorrow, .pointer, .index, .unidentified:
         return false
     }
   }
@@ -134,7 +140,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
       case .box(let pbi):      return pbi.box
       case .class(let rea):    return rea.instance
       case .tail(let rta):     return rta.instance
-      case .stack, .global, .argument, .yield, .storeBorrow, .pointer, .unidentified:
+      case .stack, .global, .argument, .yield, .storeBorrow, .pointer, .index, .unidentified:
         return nil
     }
   }
@@ -162,7 +168,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
     switch self {
       case .class(let rea):    return rea.fieldIsLet
       case .global(let g):     return g.isLet
-      case .box, .stack, .tail, .argument, .yield, .storeBorrow, .pointer, .unidentified:
+      case .box, .stack, .tail, .argument, .yield, .storeBorrow, .pointer, .index, .unidentified:
         return false
     }
   }
@@ -174,7 +180,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
     case .class(let rea):    return rea.instance.referenceRoot is AllocRefInstBase
     case .tail(let rta):     return rta.instance.referenceRoot is AllocRefInstBase
     case .stack, .storeBorrow: return true
-    case .global, .argument, .yield, .pointer, .unidentified:
+    case .global, .argument, .yield, .pointer, .index, .unidentified:
       return false
     }
   }
@@ -184,7 +190,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
     switch self {
       case .box, .class, .tail, .stack, .storeBorrow, .global:
         return true
-      case .argument, .yield, .pointer, .unidentified:
+      case .argument, .yield, .pointer, .index, .unidentified:
         return false
     }
   }
@@ -216,6 +222,8 @@ public enum AccessBase : CustomStringConvertible, Hashable {
       return sb1 == sb2
     case (.pointer(let p1), .pointer(let p2)):
       return p1 == p2
+    case (.index(let ia1), .index(let ia2)):
+      return ia1 == ia2
     default:
       return false
     }
@@ -258,7 +266,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
 
     switch (self, other) {
 
-    // First handle all pairs of the same kind (except `yield` and `pointer`).
+    // First handle all pairs of the same kind (except `yield`, `pointer` and `index`).
     case (.box(let pb), .box(let otherPb)):
       return pb.fieldIndex != otherPb.fieldIndex ||
         isDifferentAllocation(pb.box.referenceRoot, otherPb.box.referenceRoot) ||
@@ -303,7 +311,7 @@ public enum AccessBase : CustomStringConvertible, Hashable {
 
 /// An `AccessPath` is a pair of a `base: AccessBase` and a `projectionPath: Path`
 /// which denotes the offset of the access from the base in terms of projections.
-public struct AccessPath : CustomStringConvertible {
+public struct AccessPath : CustomStringConvertible, Hashable {
   public let base: AccessBase
 
   /// address projections only
@@ -475,6 +483,11 @@ public enum EnclosingScope {
 private struct AccessPathWalker : AddressUseDefWalker {
   var result = AccessPath.unidentified()
   var foundBeginAccess: BeginAccessInst?
+  let enforceConstantProjectionPath: Bool
+
+  init(enforceConstantProjectionPath: Bool = false) {
+    self.enforceConstantProjectionPath = enforceConstantProjectionPath
+  }
 
   mutating func walk(startAt address: Value, initialPath: SmallProjectionPath = SmallProjectionPath()) {
     if walkUp(address: address, path: Path(projectionPath: initialPath)) == .abortWalk {
@@ -533,9 +546,13 @@ private struct AccessPathWalker : AddressUseDefWalker {
   }
 
   mutating func walkUp(address: Value, path: Path) -> WalkResult {
-    if address is IndexAddrInst {
+    if let indexAddr = address as? IndexAddrInst {
+      if !(indexAddr.index is IntegerLiteralInst) && enforceConstantProjectionPath {
+        self.result = AccessPath(base: .index(indexAddr), projectionPath: path.projectionPath)
+        return .continueWalk
+      }
       // Track that we crossed an `index_addr` during the walk-up
-      return walkUpDefault(address: address, path: path.with(indexAddr: true))
+      return walkUpDefault(address: indexAddr, path: path.with(indexAddr: true))
     } else if path.indexAddr && !canBeOperandOfIndexAddr(address) {
       // An `index_addr` instruction cannot be derived from an address
       // projection. Bail out
@@ -565,6 +582,25 @@ extension Value {
     return walker.result
   }
 
+  /// Like `accessPath`, but ensures that the projectionPath only contains "constant" elements.
+  /// This means: if the access contains an `index_addr` projection with a non-constant index,
+  /// the `projectionPath` does _not_ contain the `index_addr`.
+  /// Instead, the `base` is an `AccessBase.index` which refers to the `index_addr`.
+  /// For example:
+  /// ```
+  ///    %1 = ref_tail_addr %some_reference
+  ///    %2 = index_addr %1, %some_non_const_value
+  ///    %3 = struct_element_addr %2, #field2
+  /// ```
+  /// `%3.accessPath`         = base: tail(`%1`),  projectionPath: `i*.s2`
+  /// `%3.constantAccessPath` = base: index(`%2`), projectionPath: `s2`
+  ///
+  public var constantAccessPath: AccessPath {
+    var walker = AccessPathWalker(enforceConstantProjectionPath: true)
+    walker.walk(startAt: self)
+    return walker.result
+  }
+
   public func getAccessPath(fromInitialPath: SmallProjectionPath) -> AccessPath {
     var walker = AccessPathWalker()
     walker.walk(startAt: self, initialPath: fromInitialPath)
@@ -637,7 +673,7 @@ extension ValueUseDefWalker where Path == SmallProjectionPath {
         return walkUp(value: rea.instance, path: path.push(.classField, index: rea.fieldIndex)) != .abortWalk
       case .tail(let rta):
         return walkUp(value: rta.instance, path: path.push(.tailElements, index: 0)) != .abortWalk
-      case .stack, .global, .argument, .yield, .storeBorrow, .pointer, .unidentified:
+      case .stack, .global, .argument, .yield, .storeBorrow, .pointer, .index, .unidentified:
         return false
     }
   }

test/SILOptimizer/accessutils.sil

@@ -626,3 +626,65 @@ bb0(%0 : @guaranteed $C):
   dealloc_stack %1 : $*C
   return %6 : $C
 }
+
+// CHECK-LABEL: Accesses for indexAddr
+// CHECK:       Value:   %6 = index_addr %3 : $*Int64, %4 : $Builtin.Word
+// CHECK:       Scope: base
+// CHECK:       Base: tail - %0 = argument of bb0 : $C
+// CHECK:       Path: "i4"
+// CHECK:         Storage: %0 = argument of bb0 : $C
+// CHECK:         Path: "ct.i4"
+// CHECK:          Value:   %7 = index_addr %3 : $*Int64, %5 : $Builtin.Word
+// CHECK:       Scope: base
+// CHECK:       Base: tail - %0 = argument of bb0 : $C
+// CHECK:       Path: "i5"
+// CHECK:         Storage: %0 = argument of bb0 : $C
+// CHECK:         Path: "ct.i5"
+// CHECK:          Value:   %8 = index_addr %3 : $*Int64, %1 : $Builtin.Word
+// CHECK:       Scope: base
+// CHECK:       nonconst-base: tail - %0 = argument of bb0 : $C
+// CHECK:       nonconst-path: "i*"
+// CHECK:       const-base: index -   %8 = index_addr %3 : $*Int64, %1 : $Builtin.Word
+// CHECK:       const-path: ""
+// CHECK:         Storage: %0 = argument of bb0 : $C
+// CHECK:         Path: "ct.i*"
+// CHECK:          Value:   %10 = struct_element_addr %9 : $*Int64, #Int64._value
+// CHECK:       Scope: base
+// CHECK:       nonconst-base: tail - %0 = argument of bb0 : $C
+// CHECK:       nonconst-path: "i*.s0"
+// CHECK:       const-base: index -   %9 = index_addr %3 : $*Int64, %2 : $Builtin.Word
+// CHECK:       const-path: "s0"
+// CHECK:         Storage: %0 = argument of bb0 : $C
+// CHECK:         Path: "ct.i*.s0"
+// CHECK-LABEL: End accesses for indexAddr
+sil @indexAddr : $@convention(thin) (@guaranteed C, Builtin.Word, Builtin.Word) -> () {
+bb0(%0 : $C, %1 : $Builtin.Word, %2 : $Builtin.Word):
+  %3 = ref_tail_addr %0 : $C, $Int64
+  %4 = integer_literal $Builtin.Word, 4
+  %5 = integer_literal $Builtin.Word, 5
+
+  %6 = index_addr %3 : $*Int64, %4 : $Builtin.Word
+  %7 = index_addr %3 : $*Int64, %5 : $Builtin.Word
+  %8 = index_addr %3 : $*Int64, %1 : $Builtin.Word
+  %9 = index_addr %3 : $*Int64, %2 : $Builtin.Word
+  %10 = struct_element_addr %9 : $*Int64, #Int64._value
+
+  %20 = load %6 : $*Int64
+  %21 = load %7 : $*Int64
+  %22 = load %8 : $*Int64
+  %23 = load %10 : $*Builtin.Int64
+
+  %isDistinct = function_ref @_isDistinct : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  %isNotDistinct = function_ref @_isNotDistinct : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+
+  apply %isDistinct<Int64, Int64>(%6, %7) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  apply %isNotDistinct<Int64, Int64>(%3, %6) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  apply %isNotDistinct<Int64, Int64>(%3, %8) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  apply %isNotDistinct<Int64, Builtin.Int64>(%3, %10) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  apply %isNotDistinct<Int64, Int64>(%6, %8) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  apply %isNotDistinct<Int64, Int64>(%8, %9) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+  apply %isNotDistinct<Int64, Builtin.Int64>(%8, %10) : $@convention(thin) <τ_0_0, τ_0_1> (@inout τ_0_0, @inout τ_0_1) -> ()
+
+  %200 = tuple ()
+  return %200 : $()
+}