Merge pull request #24102 from gottesmm/pr-cdde379c40f94a6fd49361a2b347388665705682

[constant-prop] When replacing uses of destructure_tuple, only RAUW i…

Author: gottesmm

lib/SILOptimizer/Utils/ConstantFolding.cpp

@@ -113,7 +113,7 @@ APInt swift::constantFoldCast(APInt val, const BuiltinInfo &BI) {
   SrcTy->castTo<BuiltinIntegerType>()->getGreatestWidth();
   uint32_t DestBitWidth =
   DestTy->castTo<BuiltinIntegerType>()->getGreatestWidth();
-  
+
   APInt CastResV;
   if (SrcBitWidth == DestBitWidth) {
     return val;
@@ -1761,7 +1761,46 @@ ConstantFolder::processWorkList() {
               FoldedUsers.insert(TI);
           }
 
-          // We were able to fold, so all users should use the new folded value.
+          // We were able to fold, so all users should use the new folded
+          // value. If we don't have any such users, continue.
+          //
+          // NOTE: The reason why we check if our result has uses is that if
+          // User is a MultipleValueInstruction an infinite loop can result if
+          // User has a result different than the one at Index that we can not
+          // constant fold and if C's defining instruction is an aggregate that
+          // defines an operand of I.
+          //
+          // As an elucidating example, consider the following SIL:
+          //
+          //   %w = integer_literal $Builtin.Word, 1
+          //   %0 = struct $Int (%w : $Builtin.Word)                     (*)
+          //   %1 = apply %f() : $@convention(thin) () -> @owned Klass
+          //   %2 = tuple (%0 : $Int, %1 : $Klass)
+          //   (%3, %4) = destructure_tuple %2 : $(Int, Klass)
+          //   store %4 to [init] %mem2: %*Klass
+          //
+          // Without this check, we would infinite loop by processing our
+          // worklist as follows:
+          //
+          // 1. We visit %w and add %0 to the worklist unconditionally since it
+          //    is a StructInst.
+          //
+          // 2. We visit %0 and then since %2 is a tuple, we add %2 to the
+          //    worklist unconditionally.
+          //
+          // 3. We visit %2 and see that it has a destructure_tuple user. We see
+          //    that we can simplify %3 -> %0, but cannot simplify %4. This
+          //    means that if we just assume success if we can RAUW %3 without
+          //    checking if we will actually replace anything, we will add %0's
+          //    defining instruction (*) to the worklist. Q.E.D.
+          //
+          // In contrast, if we realize that RAUWing %3 does nothing and skip
+          // it, we exit the worklist as expected.
+          SILValue r = User->getResult(Index);
+          if (r->use_empty())
+            continue;
+
+          // Otherwise, do the RAUW.
           User->getResult(Index)->replaceAllUsesWith(C);
 
           // The new constant could be further folded now, add it to the

test/SILOptimizer/constant_propagation_ownership.sil

@@ -4,6 +4,10 @@
 import Swift
 import Builtin
 
+//////////////////
+// Declarations //
+//////////////////
+
 struct UInt {
   var value: Builtin.Word
 }
@@ -20,6 +24,15 @@ struct UInt64 {
   var value: Builtin.Int64
 }
 
+class Klass {}
+
+sil @klass_allocator : $@convention(method) (@thick Klass.Type) -> @owned Klass
+sil @generic_user : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
+
+///////////
+// Tests //
+///////////
+
 sil [ossa] @count_leading_zeros_corner_case : $@convention(thin) () -> Builtin.Int64 {
 bb0:
  %zero64 = integer_literal $Builtin.Int64, 0
@@ -1132,3 +1145,39 @@ bb2(%8 : @owned $AnSubObject):
 bb3(%11 : $Builtin.Int1):
   return %11 : $Builtin.Int1
 }
+
+sil [ossa] @do_not_infinite_loop : $@convention(thin) () -> () {
+bb0:
+  %0 = integer_literal $Builtin.Word, 1
+  %1 = struct $Int (%0 : $Builtin.Word)
+  %2 = metatype $@thick Klass.Type
+  %3 = function_ref @klass_allocator : $@convention(method) (@thick Klass.Type) -> @owned Klass
+  %4 = apply %3(%2) : $@convention(method) (@thick Klass.Type) -> @owned Klass
+  // We can never in constant propagation forward %5 today since we would need
+  // to know that we are eliminting all of its uses. But the pass doesn't
+  // understand that today and so destroy_value %5 in bb2 stops our ability to
+  // optimize.
+  %5 = tuple (%1 : $Int, %4 : $Klass)
+  cond_br undef, bb1, bb2
+
+bb1:
+  %6 = alloc_stack $(Int, Klass)
+  %7 = tuple_element_addr %6 : $*(Int, Klass), 0
+  %8 = tuple_element_addr %6 : $*(Int, Klass), 1
+  (%9, %10) = destructure_tuple %5 : $(Int, Klass)
+  store %1 to [trivial] %7 : $*Int
+  store %10 to [init] %8 : $*Klass
+  %13 = function_ref @generic_user : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
+  %14 = apply %13<(Int, Klass)>(%6) : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
+  destroy_addr %6 : $*(Int, Klass)
+  dealloc_stack %6 : $*(Int, Klass)
+  br bb3
+
+bb2:
+  destroy_value %5 : $(Int, Klass)
+  br bb3
+
+bb3:
+  %9999 = tuple()
+  return %9999 : $()
+}