Merge pull request #81131 from eeckstein/copyaddr-in-rle-6.2

[6.2] RedundantLoadElimination: support replacing a redundant `copy_addr` with a `store`

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/RedundantLoadElimination.swift

@@ -12,7 +12,7 @@
 
 import SIL
 
-/// Replaces redundant load instructions with already available values.
+/// Replaces redundant `load` or `copy_addr` instructions with already available values.
 ///
 /// A load is redundant if the loaded value is already available at that point.
 /// This can be via a preceding store to the same address:
@@ -52,6 +52,9 @@ import SIL
 ///     %f2 = load %fa2
 ///     %2 = struct (%f1, %f2)
 ///
+/// This works in a similar fashion for `copy_addr`. If the source value of the `copy_addr` is
+/// already available, the `copy_addr` is replaced by a `store` of the available value.
+///
 /// The algorithm is a data flow analysis which starts at the original load and searches
 /// for preceding stores or loads by following the control flow in backward direction.
 /// The preceding stores and loads provide the "available values" with which the original
@@ -99,7 +102,7 @@ func eliminateRedundantLoads(in function: Function,
     while let i = inst {
       defer { inst = i.previous }
 
-      if let load = inst as? LoadInst {
+      if let load = inst as? LoadingInstruction {
         if !context.continueWithNextSubpassRun(for: load) {
           return changed
         }
@@ -116,7 +119,59 @@ func eliminateRedundantLoads(in function: Function,
   return changed
 }
 
-private func tryEliminate(load: LoadInst, complexityBudget: inout Int, _ context: FunctionPassContext) -> Bool {
+/// Either a `load` or a `copy_addr` (which is equivalent to a load+store).
+private protocol LoadingInstruction: Instruction {
+  var address: Value { get }
+  var type: Type { get }
+  var ownership: Ownership { get }
+  var loadOwnership: LoadInst.LoadOwnership { get }
+  var canLoadValue: Bool { get }
+  func trySplit(_ context: FunctionPassContext) -> Bool
+  func materializeLoadForReplacement(_ context: FunctionPassContext) -> LoadInst
+}
+
+extension LoadInst : LoadingInstruction {
+  // We know that the type is loadable because - well - this is a load.
+  var canLoadValue: Bool { true }
+
+  // Nothing to materialize, because this is already a `load`.
+  func materializeLoadForReplacement(_ context: FunctionPassContext) -> LoadInst { return self }
+}
+
+extension CopyAddrInst : LoadingInstruction {
+  var address: Value { source }
+  var type: Type { address.type.objectType }
+  var typeIsLoadable: Bool { type.isLoadable(in: parentFunction) }
+
+  var ownership: Ownership {
+    if !parentFunction.hasOwnership || type.isTrivial(in: parentFunction) {
+      return .none
+    }
+    // Regardless of if the copy is taking or copying, the loaded value is an owned value.
+    return .owned
+  }
+
+  var canLoadValue: Bool {
+    if !source.type.isLoadable(in: parentFunction) {
+      // Although the original load's type is loadable (obviously), it can be projected-out
+      // from the copy_addr's type which might be not loadable.
+      return false
+    }
+    if !parentFunction.hasOwnership {
+      if !isTakeOfSrc || !isInitializationOfDest {
+        // For simplicity, bail if we would have to insert compensating retains and releases.
+        return false
+      }
+    }
+    return true
+  }
+
+  func materializeLoadForReplacement(_ context: FunctionPassContext) -> LoadInst {
+    return replaceWithLoadAndStore(context).load
+  }
+}
+
+private func tryEliminate(load: LoadingInstruction, complexityBudget: inout Int, _ context: FunctionPassContext) -> Bool {
   switch load.isRedundant(complexityBudget: &complexityBudget, context) {
   case .notRedundant:
     return false
@@ -136,9 +191,12 @@ private func tryEliminate(load: LoadInst, complexityBudget: inout Int, _ context
   }
 }
 
-private extension LoadInst {
+private extension LoadingInstruction {
 
   func isEligibleForElimination(in variant: RedundantLoadEliminationVariant, _ context: FunctionPassContext) -> Bool {
+    if !canLoadValue {
+      return false
+    }
     switch variant {
     case .mandatory, .mandatoryInGlobalInit:
       if loadOwnership == .take {
@@ -171,20 +229,6 @@ private extension LoadInst {
     return true
   }
 
-  enum DataflowResult {
-    case notRedundant
-    case redundant([AvailableValue])
-    case maybePartiallyRedundant(AccessPath)
-
-    init(notRedundantWith subPath: AccessPath?) {
-      if let subPath = subPath {
-        self = .maybePartiallyRedundant(subPath)
-      } else {
-        self = .notRedundant
-      }
-    }
-  }
-
   func isRedundant(complexityBudget: inout Int, _ context: FunctionPassContext) -> DataflowResult {
     return isRedundant(at: address.constantAccessPath, complexityBudget: &complexityBudget, context)
   }
@@ -285,7 +329,7 @@ private extension LoadInst {
   }
 }
 
-private func replace(load: LoadInst, with availableValues: [AvailableValue], _ context: FunctionPassContext) {
+private func replace(load: LoadingInstruction, with availableValues: [AvailableValue], _ context: FunctionPassContext) {
   var ssaUpdater = SSAUpdater(function: load.parentFunction,
                               type: load.type, ownership: load.ownership, context)
 
@@ -318,14 +362,16 @@ private func replace(load: LoadInst, with availableValues: [AvailableValue], _ c
     newValue = ssaUpdater.getValue(inMiddleOf: load.parentBlock)
   }
 
+  let originalLoad = load.materializeLoadForReplacement(context)
+
   // Make sure to keep dependencies valid after replacing the load
-  insertMarkDependencies(for: load, context)
+  insertMarkDependencies(for: originalLoad, context)
 
-  load.replace(with: newValue, context)
+  originalLoad.replace(with: newValue, context)
 }
 
 private func provideValue(
-  for load: LoadInst,
+  for load: LoadingInstruction,
   from availableValue: AvailableValue,
   _ context: FunctionPassContext
 ) -> Value {
@@ -341,9 +387,9 @@ private func provideValue(
                                                         builder: availableValue.getBuilderForProjections(context))
   case .take:
     if projectionPath.isEmpty {
-      return shrinkMemoryLifetime(from: load, to: availableValue, context)
+      return shrinkMemoryLifetime(to: availableValue, context)
     } else {
-      return shrinkMemoryLifetimeAndSplit(from: load, to: availableValue, projectionPath: projectionPath, context)
+      return shrinkMemoryLifetimeAndSplit(to: availableValue, projectionPath: projectionPath, context)
     }
   }
 }
@@ -366,7 +412,7 @@ private func insertMarkDependencies(for load: LoadInst, _ context: FunctionPassC
 private struct MarkDependenceInserter : AddressUseDefWalker {
   let load: LoadInst
   let context: FunctionPassContext
-  
+
   mutating func walkUp(address: Value, path: UnusedWalkingPath) -> WalkResult {
     if let mdi = address as? MarkDependenceInst {
       let builder = Builder(after: load, context)
@@ -375,7 +421,7 @@ private struct MarkDependenceInserter : AddressUseDefWalker {
     }
     return walkUpDefault(address: address, path: path)
   }
-  
+
   mutating func rootDef(address: Value, path: UnusedWalkingPath) -> WalkResult {
     return .continueWalk
   }
@@ -392,7 +438,7 @@ private struct MarkDependenceInserter : AddressUseDefWalker {
 ///     ...
 ///     // replace %2 with %1
 ///
-private func shrinkMemoryLifetime(from load: LoadInst, to availableValue: AvailableValue, _ context: FunctionPassContext) -> Value {
+private func shrinkMemoryLifetime(to availableValue: AvailableValue, _ context: FunctionPassContext) -> Value {
   switch availableValue {
   case .viaLoad(let availableLoad):
     assert(availableLoad.loadOwnership == .copy)
@@ -442,7 +488,7 @@ private func shrinkMemoryLifetime(from load: LoadInst, to availableValue: Availa
 ///     ...
 ///     // replace %3 with %1
 ///
-private func shrinkMemoryLifetimeAndSplit(from load: LoadInst, to availableValue: AvailableValue, projectionPath: SmallProjectionPath, _ context: FunctionPassContext) -> Value {
+private func shrinkMemoryLifetimeAndSplit(to availableValue: AvailableValue, projectionPath: SmallProjectionPath, _ context: FunctionPassContext) -> Value {
   switch availableValue {
   case .viaLoad(let availableLoad):
     assert(availableLoad.loadOwnership == .copy)
@@ -462,6 +508,20 @@ private func shrinkMemoryLifetimeAndSplit(from load: LoadInst, to availableValue
   }
 }
 
+private enum DataflowResult {
+  case notRedundant
+  case redundant([AvailableValue])
+  case maybePartiallyRedundant(AccessPath)
+
+  init(notRedundantWith subPath: AccessPath?) {
+    if let subPath = subPath {
+      self = .maybePartiallyRedundant(subPath)
+    } else {
+      self = .notRedundant
+    }
+  }
+}
+
 /// Either a `load` or `store` which is preceding the original load and provides the loaded value.
 private enum AvailableValue {
   case viaLoad(LoadInst)
@@ -505,7 +565,7 @@ private extension Array where Element == AvailableValue {
   func replaceCopyAddrsWithLoadsAndStores(_ context: FunctionPassContext) -> [AvailableValue] {
     return map {
       if case .viaCopyAddr(let copyAddr) = $0 {
-        return .viaStore(copyAddr.replaceWithLoadAndStore(context))
+        return .viaStore(copyAddr.replaceWithLoadAndStore(context).store)
       } else {
         return $0
       }
@@ -514,15 +574,15 @@ private extension Array where Element == AvailableValue {
 }
 
 private struct InstructionScanner {
-  private let load: LoadInst
+  private let load: LoadingInstruction
   private let accessPath: AccessPath
   private let storageDefBlock: BasicBlock?
   private let aliasAnalysis: AliasAnalysis
 
   private(set) var potentiallyRedundantSubpath: AccessPath? = nil
   private(set) var availableValues = Array<AvailableValue>()
 
-  init(load: LoadInst, accessPath: AccessPath, _ aliasAnalysis: AliasAnalysis) {
+  init(load: LoadingInstruction, accessPath: AccessPath, _ aliasAnalysis: AliasAnalysis) {
     self.load = load
     self.accessPath = accessPath
     self.storageDefBlock = accessPath.base.reference?.referenceRoot.parentBlock
@@ -616,7 +676,7 @@ private struct InstructionScanner {
         potentiallyRedundantSubpath = precedingStorePath
       }
 
-    case let preceedingCopy as CopyAddrInst where preceedingCopy.canProvideValue:
+    case let preceedingCopy as CopyAddrInst where preceedingCopy.canLoadValue:
       let copyPath = preceedingCopy.destination.constantAccessPath
       if copyPath.getMaterializableProjection(to: accessPath) != nil {
         availableValues.append(.viaCopyAddr(preceedingCopy))
@@ -712,20 +772,3 @@ private struct Liverange {
     return false
   }
 }
-
-private extension CopyAddrInst {
-  var canProvideValue: Bool {
-    if !source.type.isLoadable(in: parentFunction) {
-      // Although the original load's type is loadable (obviously), it can be projected-out 
-      // from the copy_addr's type which might be not loadable.
-      return false
-    }
-    if !parentFunction.hasOwnership {
-      if !isTakeOfSrc || !isInitializationOfDest {
-        // For simplicity, bail if we would have to insert compensating retains and releases.
-        return false
-      }
-    }
-    return true
-  }
-}

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -941,27 +941,74 @@ extension CheckedCastAddrBranchInst {
 
 extension CopyAddrInst {
   @discardableResult
-  func replaceWithLoadAndStore(_ context: some MutatingContext) -> StoreInst {
-    let loadOwnership: LoadInst.LoadOwnership
-    let storeOwnership: StoreInst.StoreOwnership
-    if parentFunction.hasOwnership {
-      if source.type.isTrivial(in: parentFunction) {
-        loadOwnership = .trivial
-        storeOwnership = .trivial
-      } else {
-        loadOwnership = isTakeOfSrc ? .take : .copy
-        storeOwnership = isInitializationOfDest ? .initialize : .assign
+  func trySplit(_ context: FunctionPassContext) -> Bool {
+    let builder = Builder(before: self, context)
+    if source.type.isStruct {
+      if (source.type.nominal as! StructDecl).hasUnreferenceableStorage {
+        return false
       }
-    } else {
-      loadOwnership = .unqualified
-      storeOwnership = .unqualified
+      guard let fields = source.type.getNominalFields(in: parentFunction) else {
+        return false
+      }
+      for idx in 0..<fields.count {
+        let srcFieldAddr = builder.createStructElementAddr(structAddress: source, fieldIndex: idx)
+        let destFieldAddr = builder.createStructElementAddr(structAddress: destination, fieldIndex: idx)
+        builder.createCopyAddr(from: srcFieldAddr, to: destFieldAddr,
+                               takeSource: isTake(for: srcFieldAddr), initializeDest: isInitializationOfDest)
+      }
+      context.erase(instruction: self)
+      return true
+    } else if source.type.isTuple {
+      let builder = Builder(before: self, context)
+      for idx in 0..<source.type.tupleElements.count {
+        let srcFieldAddr = builder.createTupleElementAddr(tupleAddress: source, elementIndex: idx)
+        let destFieldAddr = builder.createTupleElementAddr(tupleAddress: destination, elementIndex: idx)
+        builder.createCopyAddr(from: srcFieldAddr, to: destFieldAddr,
+                               takeSource: isTake(for: srcFieldAddr), initializeDest: isInitializationOfDest)
+      }
+      context.erase(instruction: self)
+      return true
     }
-    
+    return false
+  }
+
+  private func isTake(for fieldValue: Value) -> Bool {
+    return isTakeOfSrc && !fieldValue.type.objectType.isTrivial(in: parentFunction)
+  }
+
+  @discardableResult
+  func replaceWithLoadAndStore(_ context: some MutatingContext) -> (load: LoadInst, store: StoreInst) {
     let builder = Builder(before: self, context)
-    let value = builder.createLoad(fromAddress: source, ownership: loadOwnership)
-    let store = builder.createStore(source: value, destination: destination, ownership: storeOwnership)
+    let load = builder.createLoad(fromAddress: source, ownership: loadOwnership)
+    let store = builder.createStore(source: load, destination: destination, ownership: storeOwnership)
     context.erase(instruction: self)
-    return store
+    return (load, store)
+  }
+
+  var loadOwnership: LoadInst.LoadOwnership {
+    if !parentFunction.hasOwnership {
+      return .unqualified
+    }
+    if type.isTrivial(in: parentFunction) {
+      return .trivial
+    }
+    if isTakeOfSrc {
+      return .take
+    }
+    return .copy
+  }
+
+  var storeOwnership: StoreInst.StoreOwnership {
+    if !parentFunction.hasOwnership {
+      return .unqualified
+    }
+    if type.isTrivial(in: parentFunction) {
+      return .trivial
+    }
+    if isInitializationOfDest {
+      return .initialize
+    }
+    return .assign
   }
 }
 

test/AutoDiff/compiler_crashers_fixed/issue-62608-conflicting-debug-info.swift

@@ -38,7 +38,7 @@ public func pullback<T>(
 
 // CHECK-LABEL: sil private @$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU_TJpSUpSr :
 // CHECK: alloc_stack $Double, var, name "derivative of 'element' in scope at {{.*}} (scope #3)"
-// CHECK: debug_value %{{.*}} : $Builtin.FPIEEE64, var, (name "derivative of 'element' in scope at {{.*}} (scope #1)"
+// CHECK: debug_value %{{.*}} : $Builtin.FPIEEE64, var, (name "derivative of 'element' in scope at {{.*}} (scope #{{.*}})"
 
 public extension Array where Element: Differentiable {
     @inlinable

test/SILGen/reference_bindings.swift

@@ -35,10 +35,10 @@ func doSomething() {}
 // SIL:   [[BOX:%.*]] = alloc_stack [var_decl] $Int, var, name "x"
 // SIL:   [[INOUT_BOX:%.*]] = alloc_stack [var_decl] $Int, var, name "x2"
 // SIL:   [[ACCESS:%.*]] = begin_access [modify] [static] [[BOX]]
-// SIL:   copy_addr [take] [[ACCESS]] to [init] [[INOUT_BOX]]
+// SIL:   store {{%.*}} to [[INOUT_BOX]]
 // SIL:   [[FUNC:%.*]] = function_ref @$s18reference_bindings11doSomethingyyF : $@convention(thin) () -> ()
 // SIL:   apply [[FUNC]]()
-// SIL:   copy_addr [take] [[INOUT_BOX]] to [init] [[ACCESS]]
+// SIL:   store {{%.*}} to [[ACCESS]]
 // SIL:   end_access [[ACCESS]]
 // SIL: } // end sil function '$s18reference_bindings13testBindToVaryyF'
 func testBindToVar() {