Merge pull request #60072 from eeckstein/walkers-in-release-devirtualizer

Some improvements in WalkUtils and use the ValueUseDefWalker in the ReleaseDevirtualizer pass

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ComputeEffects.swift

@@ -13,7 +13,6 @@
 import SIL
 
 fileprivate typealias Selection = ArgumentEffect.Selection
-fileprivate typealias Path = ArgumentEffect.Path
 
 /// Computes effects for function arguments.
 ///
@@ -34,7 +33,7 @@ let computeEffects = FunctionPass(name: "compute-effects", {
   var argsWithDefinedEffects = getArgIndicesWithDefinedEffects(of: function)
 
   struct IgnoreRecursiveCallVisitor : EscapeInfoVisitor {
-    func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+    func visitUse(operand: Operand, path: EscapePath) -> UseResult {
       return isOperandOfRecursiveCall(operand) ? .ignore : .continueWalk
     }
   }
@@ -50,17 +49,17 @@ let computeEffects = FunctionPass(name: "compute-effects", {
     if argsWithDefinedEffects.contains(arg.index) { continue }
 
     // First check: is the argument (or a projected value of it) escaping at all?
-    if !escapeInfo.isEscapingWhenWalkingDown(object: arg, path: Path(.anything)) {
-      let selectedArg = Selection(arg, pathPattern: Path(.anything))
+    if !escapeInfo.isEscapingWhenWalkingDown(object: arg, path: SmallProjectionPath(.anything)) {
+      let selectedArg = Selection(arg, pathPattern: SmallProjectionPath(.anything))
       newEffects.push(ArgumentEffect(.notEscaping, selectedArg: selectedArg))
       continue
     }
 
     // Now compute effects for two important cases:
     //   * the argument itself + any value projections, and...
-    if addArgEffects(context: context, arg, argPath: Path(), to: &newEffects, returnInst) {
+    if addArgEffects(context: context, arg, argPath: SmallProjectionPath(), to: &newEffects, returnInst) {
       //   * single class indirections
-      _ = addArgEffects(context: context, arg, argPath: Path(.anyValueFields).push(.anyClassField),
+      _ = addArgEffects(context: context, arg, argPath: SmallProjectionPath(.anyValueFields).push(.anyClassField),
                         to: &newEffects, returnInst)
     }
   }
@@ -75,7 +74,7 @@ let computeEffects = FunctionPass(name: "compute-effects", {
 
 /// Returns true if an argument effect was added.
 private
-func addArgEffects(context: PassContext, _ arg: FunctionArgument, argPath ap: Path,
+func addArgEffects(context: PassContext, _ arg: FunctionArgument, argPath ap: SmallProjectionPath,
                    to newEffects: inout Stack<ArgumentEffect>,
                    _ returnInst: ReturnInst?) -> Bool {
   // Correct the path if the argument is not a class reference itself, but a value type
@@ -91,18 +90,18 @@ func addArgEffects(context: PassContext, _ arg: FunctionArgument, argPath ap: Pa
     var toSelection: Selection?
     var returnInst: ReturnInst?
 
-    mutating func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+    mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
       if operand.instruction == returnInst {
         // The argument escapes to the function return
-        if state.followStores {
+        if path.followStores {
           // The escaping path must not introduce a followStores.
           return .abort
         }
         if let ta = toSelection {
           if ta.value != .returnValue { return .abort }
-          toSelection = Selection(.returnValue, pathPattern: path.merge(with: ta.pathPattern))
+          toSelection = Selection(.returnValue, pathPattern: path.projectionPath.merge(with: ta.pathPattern))
         } else {
-          toSelection = Selection(.returnValue, pathPattern: path)
+          toSelection = Selection(.returnValue, pathPattern: path.projectionPath)
         }
         return .ignore
       }
@@ -112,21 +111,21 @@ func addArgEffects(context: PassContext, _ arg: FunctionArgument, argPath ap: Pa
       return .continueWalk
     }
 
-    mutating func visitDef(def: Value, path: Path, state: State) -> DefResult {
+    mutating func visitDef(def: Value, path: EscapePath) -> DefResult {
       guard let destArg = def as? FunctionArgument else {
         return .continueWalkUp
       }
       // The argument escapes to another argument (e.g. an out or inout argument)
-      if state.followStores {
+      if path.followStores {
         // The escaping path must not introduce a followStores.
         return .abort
       }
       let argIdx = destArg.index
       if let ta = toSelection {
         if ta.value != .argument(argIdx) { return .abort }
-        toSelection = Selection(.argument(argIdx), pathPattern: path.merge(with: ta.pathPattern))
+        toSelection = Selection(.argument(argIdx), pathPattern: path.projectionPath.merge(with: ta.pathPattern))
       } else {
-        toSelection = Selection(.argument(argIdx), pathPattern: path)
+        toSelection = Selection(.argument(argIdx), pathPattern: path.projectionPath)
       }
       return .walkDown
     }
@@ -197,7 +196,7 @@ private func isOperandOfRecursiveCall(_ op: Operand) -> Bool {
 /// there are no other arguments or escape points than `fromArgument`. Also, the
 /// path at the `fromArgument` must match with `fromPath`.
 private
-func isExclusiveEscape(context: PassContext, fromArgument: Argument, fromPath: Path, to toSelection: Selection,
+func isExclusiveEscape(context: PassContext, fromArgument: Argument, fromPath: SmallProjectionPath, to toSelection: Selection,
                        _ returnInst: ReturnInst) -> Bool {
   switch toSelection.value {
 
@@ -207,25 +206,25 @@ func isExclusiveEscape(context: PassContext, fromArgument: Argument, fromPath: P
       let fromArgument: Argument
       let toSelection: Selection
       let returnInst: ReturnInst
-      let fromPath: Path
+      let fromPath: SmallProjectionPath
 
-      mutating func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+      mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
         if operand.instruction == returnInst {
-          if state.followStores { return .abort }
-          if path.matches(pattern: toSelection.pathPattern) {
+          if path.followStores { return .abort }
+          if path.projectionPath.matches(pattern: toSelection.pathPattern) {
             return .ignore
           }
           return .abort
         }
         return .continueWalk
       }
 
-      mutating func visitDef(def: Value, path: Path, state: State) -> DefResult {
+      mutating func visitDef(def: Value, path: EscapePath) -> DefResult {
         guard let arg = def as? FunctionArgument else {
           return .continueWalkUp
         }
-        if state.followStores { return .abort }
-        if arg == fromArgument && path.matches(pattern: fromPath) {
+        if path.followStores { return .abort }
+        if arg == fromArgument && path.projectionPath.matches(pattern: fromPath) {
           return .walkDown
         }
         return .abort
@@ -240,17 +239,17 @@ func isExclusiveEscape(context: PassContext, fromArgument: Argument, fromPath: P
   case .argument(let toArgIdx):
     struct IsExclusiveArgumentEscapeVisitor : EscapeInfoVisitor {
       let fromArgument: Argument
-      let fromPath: Path
+      let fromPath: SmallProjectionPath
       let toSelection: Selection
       let toArg: FunctionArgument
 
-      mutating func visitDef(def: Value, path: Path, state: State) -> DefResult {
+      mutating func visitDef(def: Value, path: EscapePath) -> DefResult {
         guard let arg = def as? FunctionArgument else {
           return .continueWalkUp
         }
-        if state.followStores { return .abort }
-        if arg == fromArgument && path.matches(pattern: fromPath) { return .walkDown }
-        if arg == toArg && path.matches(pattern: toSelection.pathPattern) { return .walkDown }
+        if path.followStores { return .abort }
+        if arg == fromArgument && path.projectionPath.matches(pattern: fromPath) { return .walkDown }
+        if arg == toArg && path.projectionPath.matches(pattern: toSelection.pathPattern) { return .walkDown }
         return .abort
       }
     }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/EscapeInfoDumper.swift

@@ -25,19 +25,19 @@ let escapeInfoDumper = FunctionPass(name: "dump-escape-info", {
   struct Visitor : EscapeInfoVisitor {
     var results: Set<String> =  Set()
 
-    mutating func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+    mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
       if operand.instruction is ReturnInst {
-        results.insert("return[\(path)]")
+        results.insert("return[\(path.projectionPath)]")
         return .ignore
       }
       return .continueWalk
     }
 
-    mutating func visitDef(def: Value, path: Path, state: State) -> DefResult {
+    mutating func visitDef(def: Value, path: EscapePath) -> DefResult {
       guard let arg = def as? FunctionArgument else {
         return .continueWalkUp
       }
-      results.insert("arg\(arg.index)[\(path)]")
+      results.insert("arg\(arg.index)[\(path.projectionPath)]")
       return .walkDown
     }
   }
@@ -96,7 +96,7 @@ let addressEscapeInfoDumper = FunctionPass(name: "dump-addr-escape-info", {
 
   struct Visitor : EscapeInfoVisitor {
     let apply: Instruction
-    mutating func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+    mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
       let user = operand.instruction
       if user == apply {
         return .abort

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ReleaseDevirtualizer.swift

@@ -74,12 +74,11 @@ private func tryDevirtualizeReleaseOfObject(
   _ deallocStackRef: DeallocStackRefInst
 ) {
   let allocRefInstruction = deallocStackRef.allocRef
-  var root = release.operands[0].value
-  while let newRoot = stripRCIdentityPreservingInsts(root) {
-    root = newRoot
-  }
 
-  if root != allocRefInstruction {
+  // Check if the release instruction right before the `dealloc_stack_ref` really releases
+  // the allocated object (and not something else).
+  var finder = FindAllocationOfRelease(allocation: allocRefInstruction)
+  if !finder.allocationIsRoot(of: release.operand) {
     return
   }
 
@@ -109,119 +108,58 @@ private func tryDevirtualizeReleaseOfObject(
   context.erase(instruction: release)
 }
 
-private func stripRCIdentityPreservingInsts(_ value: Value) -> Value? {
-  guard let inst = value as? Instruction else { return nil }
-
-  switch inst {
-  // First strip off RC identity preserving casts.
-  case is UpcastInst,
-       is UncheckedRefCastInst,
-       is InitExistentialRefInst,
-       is OpenExistentialRefInst,
-       is RefToBridgeObjectInst,
-       is BridgeObjectToRefInst,
-       is ConvertFunctionInst,
-       is UncheckedEnumDataInst:
-    return inst.operands[0].value
-
-  // Then if we have a struct_extract that is extracting a non-trivial member
-  // from a struct with no other non-trivial members, a ref count operation on
-  // the struct is equivalent to a ref count operation on the extracted
-  // member. Strip off the extract.
-  case let sei as StructExtractInst where sei.isFieldOnlyNonTrivialField:
-    return sei.operand
-
-  // If we have a struct or tuple instruction with only one non-trivial operand, the
-  // only reference count that can be modified is the non-trivial operand. Return
-  // the non-trivial operand.
-  case is StructInst, is TupleInst:
-    return inst.uniqueNonTrivialOperand
-
-  // If we have an enum instruction with a payload, strip off the enum to
-  // expose the enum's payload.
-  case let ei as EnumInst where !ei.operands.isEmpty:
-    return ei.operand
-
-  // If we have a tuple_extract that is extracting the only non trivial member
-  // of a tuple, a retain_value on the tuple is equivalent to a retain_value on
-  // the extracted value.
-  case let tei as TupleExtractInst where tei.isEltOnlyNonTrivialElt:
-    return tei.operand
-
-  default:
-    return nil
-  }
-}
+// Up-walker to find the root of a release instruction.
+private struct FindAllocationOfRelease : ValueUseDefWalker {
+  private let allocInst: AllocRefInstBase
+  private var allocFound = false
 
-private extension Instruction {
-  /// Search the operands of this tuple for a unique non-trivial elt. If we find
-  /// it, return it. Otherwise return `nil`.
-  var uniqueNonTrivialOperand: Value? {
-    var candidateElt: Value?
-    let function = self.function
-
-    for op in operands {
-      if !op.value.type.isTrivial(in: function) {
-        if candidateElt == nil {
-          candidateElt = op.value
-          continue
-        }
+  var walkUpCache = WalkerCache<SmallProjectionPath>()
 
-        // Otherwise, we have two values that are non-trivial. Bail.
-        return nil
-      }
-    }
+  init(allocation: AllocRefInstBase) { allocInst = allocation }
 
-    return candidateElt
+  /// The top-level entry point: returns true if the root of `value` is the `allocInst`.
+  mutating func allocationIsRoot(of value: Value) -> Bool {
+    let path = SmallProjectionPath().push(.anyValueFields)
+    return walkUp(value: value, path: path) != .abortWalk &&
+           allocFound
   }
-}
 
-private extension TupleExtractInst {
-  var isEltOnlyNonTrivialElt: Bool {
-    let function = self.function
-
-    if type.isTrivial(in: function) {
-      return false
+  mutating func rootDef(value: Value, path: SmallProjectionPath) -> WalkResult {
+    if value == allocInst {
+      allocFound = true
+      return .continueWalk
     }
-
-    let opType = operand.type
-
-    var nonTrivialEltsCount = 0
-    for elt in opType.tupleElements {
-      if elt.isTrivial(in: function) {
-        nonTrivialEltsCount += 1
-      }
-
-      if nonTrivialEltsCount > 1 {
-        return false
-      }
-    }
-
-    return true
+    return .abortWalk
   }
-}
-
-private extension StructExtractInst {
-  var isFieldOnlyNonTrivialField: Bool {
-    let function = self.function
-
-    if type.isTrivial(in: function) {
-      return false
-    }
 
-    let structType = operand.type
-
-    var nonTrivialFieldsCount = 0
-    for field in structType.getNominalFields(in: function) {
-      if field.isTrivial(in: function) {
-        nonTrivialFieldsCount += 1
-      }
-
-      if nonTrivialFieldsCount > 1 {
-        return false
+  // This function is called for `struct` and `tuple` instructions in case the `path` doesn't select
+  // a specific operand but all operands.
+  mutating func walkUpAllOperands(of def: Instruction, path: SmallProjectionPath) -> WalkResult {
+  
+    // Instead of walking up _all_ operands (which would be the default behavior), require that only a
+    // _single_ operand is not trivial and walk up that operand.
+    // We need to check this because the released value must not contain multiple copies of the
+    // allocated object. We can only replace a _single_ release with the destructor call and not
+    // multiple releases of the same object. E.g.
+    //
+    //    %x = alloc_ref [stack] $X
+    //    strong_retain %x
+    //    %t = tuple (%x, %x)
+    //    release_value %t       // -> releases %x two times!
+    //    dealloc_stack_ref %x
+    //
+    var nonTrivialOperandFound = false
+    for operand in def.operands {
+      if !operand.value.type.isTrivial(in: def.function) {
+        if nonTrivialOperandFound {
+          return .abortWalk
+        }
+        nonTrivialOperandFound = true
+        if walkUp(value: operand.value, path: path) == .abortWalk {
+          return .abortWalk
+        }
       }
     }
-
-    return true
+    return .continueWalk
   }
 }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/StackPromotion.swift

@@ -201,7 +201,7 @@ private func getDominatingBlockOfAllUsePoints(context: PassContext,
   struct Visitor : EscapeInfoVisitor {
     var dominatingBlock: BasicBlock
     let domTree: DominatorTree
-    mutating func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+    mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
       let defBlock = operand.value.definingBlock
       if defBlock.dominates(dominatingBlock, domTree) {
         dominatingBlock = defBlock
@@ -232,7 +232,7 @@ func computeInnerAndOuterLiferanges(instruction: SingleValueInstruction, in domB
       self.domTree = domTree
     }
 
-    mutating func visitUse(operand: Operand, path: Path, state: State) -> UseResult {
+    mutating func visitUse(operand: Operand, path: EscapePath) -> UseResult {
       let user = operand.instruction
       if innerRange.blockRange.begin.dominates(user.block, domTree) {
         innerRange.insert(user)