Handle mark_dependence in ForEachLoopUnroll

meg-gupta · meg-gupta · commit 61ce3a3af2f9 · 2023-10-27T14:19:21.000-07:00
diff --git a/lib/SILOptimizer/LoopTransforms/ForEachLoopUnroll.cpp b/lib/SILOptimizer/LoopTransforms/ForEachLoopUnroll.cpp
@@ -302,6 +302,11 @@ void ArrayInfo::classifyUsesOfArray(SILValue arrayValue) {
     // above as the array would be passed indirectly.
     if (isFixLifetimeUseOfArray(user, arrayValue))
       continue;
+    if (auto *MDI = dyn_cast<MarkDependenceInst>(user)) {
+      if (MDI->getBase() == arrayValue) {
+        continue;
+      }
+    }
     // Check if this is a forEach call on the array.
     if (TryApplyInst *forEachCall = isForEachUseOfArray(user, arrayValue)) {
       forEachCalls.insert(forEachCall);
diff --git a/test/SILOptimizer/constant_evaluator_test.sil b/test/SILOptimizer/constant_evaluator_test.sil
@@ -1241,6 +1241,7 @@ bb0:
   %1 = function_ref @$ss27_allocateUninitializedArrayySayxG_BptBwlF : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
   %2 = apply %1<String>(%0) : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
   (%3, %4) = destructure_tuple %2 : $(Array<String>, Builtin.RawPointer)
+  %5 = mark_dependence %4 : $Builtin.RawPointer on %3 : $Array<String>
   return %3 : $Array<String>
 } // CHECK: Returns Array<String>
   // CHECK: size: 0 contents []
diff --git a/test/SILOptimizer/for_each_loop_unroll_test.sil b/test/SILOptimizer/for_each_loop_unroll_test.sil
@@ -12,12 +12,42 @@ sil [_semantics "sequence.forEach"] @forEach : $@convention(method) <τ_0_0 wher
 
 sil @forEachBody : $@convention(thin) (@in_guaranteed Builtin.Int64) -> @error any Error
 
+// CHECK-LABEL: forEachLoopUnrollTest
+// CHECK: [[LIT1:%[0-9]+]] = integer_literal $Builtin.Int64, 15
+// CHECK: [[LIT2:%[0-9]+]] = integer_literal $Builtin.Int64, 27
+// CHECK: [[BODYCLOSURE:%[0-9]+]] = thin_to_thick_function
+// CHECK-NOT: forEach
+// CHECK: [[STACK:%[0-9]+]] = alloc_stack $Builtin.Int64
+// CHECK: store [[LIT1]] to [trivial] [[STACK]]
+// CHECK: try_apply [[BODYCLOSURE]]([[STACK]]) : $@noescape @callee_guaranteed (@in_guaranteed Builtin.Int64) -> @error any Error, normal [[NORMAL:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
+
+// CHECK: [[ERROR]]([[ERRPARAM:%[0-9]+]] : @owned $any Error):
+// CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM]] : $any Error)
+
+// CHECK: [[NORMAL]](%{{.*}} : $()):
+// CHECK: store [[LIT2]] to [trivial] [[STACK]] : $*Builtin.Int64
+// CHECK: try_apply [[BODYCLOSURE]]([[STACK]]) : $@noescape @callee_guaranteed (@in_guaranteed Builtin.Int64) -> @error any Error, normal [[NORMAL2:bb[0-9]+]], error [[ERROR2:bb[0-9]+]]
+
+// CHECK: [[ERROR2]]([[ERRPARAM2:%[0-9]+]] : @owned $any Error):
+// CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM2]] : $any Error)
+
+// CHECK: [[NORMAL2]](%{{.*}} : $()):
+// CHECK: dealloc_stack [[STACK]]
+// Note that the temporary alloc_stack of the array created for the forEach call
+// will be cleaned up when the forEach call is removed.
+// CHECK: destroy_value
+
+// CHECK: [[ERROR3]]([[ERRPARAM3:%[0-9]+]] : @owned $any Error):
+// CHECK: dealloc_stack [[STACK]]
+// CHECK: unreachable
+
 sil hidden [ossa] @forEachLoopUnrollTest : $@convention(thin) () -> () {
 bb0:
   %0 = integer_literal $Builtin.Word, 2
   %1 = function_ref @_allocateUninitializedArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
   %2 = apply %1<Builtin.Int64>(%0) : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
-  (%3, %4) = destructure_tuple %2 : $(Array<Builtin.Int64>, Builtin.RawPointer)
+  (%3, %4a) = destructure_tuple %2 : $(Array<Builtin.Int64>, Builtin.RawPointer)
+  %4 = mark_dependence %4a : $Builtin.RawPointer on %3 : $Array<Builtin.Int64>
   %5 = pointer_to_address %4 : $Builtin.RawPointer to [strict] $*Builtin.Int64
   %6 = integer_literal $Builtin.Int64, 15
   store %6 to [trivial] %5 : $*Builtin.Int64
@@ -46,26 +76,35 @@ bb1(%32 : $()):
 bb2(%39 : @owned $Error):
   unreachable
 }
-// CHECK-LABEL: forEachLoopUnrollTest
-// CHECK: [[LIT1:%[0-9]+]] = integer_literal $Builtin.Int64, 15
-// CHECK: [[LIT2:%[0-9]+]] = integer_literal $Builtin.Int64, 27
+
+sil @forEachBody2 : $@convention(thin) (@in_guaranteed @callee_guaranteed @substituted <A> () -> @out A for <Int>) -> @error any Error
+
+// CHECK-LABEL: nonTrivialForEachLoopUnrollTest
+// CHECK: [[ELEM1:%[0-9]+]] = copy_value %0
+// CHECK-NEXT: store %0 to [init] %{{.*}} : $*@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
+// CHECK: [[ELEM2:%[0-9]+]] = copy_value %1
+// CHECK-NEXT: store %1 to [init] %{{.*}} : $*@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
 // CHECK: [[BODYCLOSURE:%[0-9]+]] = thin_to_thick_function
 // CHECK-NOT: forEach
-// CHECK: [[STACK:%[0-9]+]] = alloc_stack $Builtin.Int64
-// CHECK: store [[LIT1]] to [trivial] [[STACK]]
-// CHECK: try_apply [[BODYCLOSURE]]([[STACK]]) : $@noescape @callee_guaranteed (@in_guaranteed Builtin.Int64) -> @error any Error, normal [[NORMAL:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
+// CHECK: [[STACK:%[0-9]+]] = alloc_stack $@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
+// CHECK: [[ELEM1BORROW:%[0-9]+]] = begin_borrow [[ELEM1]]
+// CHECK: [[SB1:%.*]] = store_borrow [[ELEM1BORROW]] to [[STACK]]
+// CHECK: try_apply [[BODYCLOSURE]]([[SB1]]) : $@noescape @callee_guaranteed (@in_guaranteed @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>) -> @error any Error, normal [[NORMAL:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
 
 // CHECK: [[ERROR]]([[ERRPARAM:%[0-9]+]] : @owned $any Error):
 // CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM]] : $any Error)
 
 // CHECK: [[NORMAL]](%{{.*}} : $()):
-// CHECK: store [[LIT2]] to [trivial] [[STACK]] : $*Builtin.Int64
-// CHECK: try_apply [[BODYCLOSURE]]([[STACK]]) : $@noescape @callee_guaranteed (@in_guaranteed Builtin.Int64) -> @error any Error, normal [[NORMAL2:bb[0-9]+]], error [[ERROR2:bb[0-9]+]]
+// CHECK: end_borrow [[ELEM1BORROW]]
+// CHECK: [[ELEM2BORROW:%[0-9]+]] = begin_borrow [[ELEM2]]
+// CHECK: [[SB2:%.*]] = store_borrow [[ELEM2BORROW]] to [[STACK]]
+// CHECK: try_apply [[BODYCLOSURE]]([[SB2]]) : $@noescape @callee_guaranteed (@in_guaranteed @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>) -> @error any Error, normal [[NORMAL2:bb[0-9]+]], error [[ERROR2:bb[0-9]+]]
 
 // CHECK: [[ERROR2]]([[ERRPARAM2:%[0-9]+]] : @owned $any Error):
 // CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM2]] : $any Error)
 
 // CHECK: [[NORMAL2]](%{{.*}} : $()):
+// CHECK: end_borrow [[ELEM2BORROW]]
 // CHECK: dealloc_stack [[STACK]]
 // Note that the temporary alloc_stack of the array created for the forEach call
 // will be cleaned up when the forEach call is removed.
@@ -74,15 +113,13 @@ bb2(%39 : @owned $Error):
 // CHECK: [[ERROR3]]([[ERRPARAM3:%[0-9]+]] : @owned $any Error):
 // CHECK: dealloc_stack [[STACK]]
 // CHECK: unreachable
-
-sil @forEachBody2 : $@convention(thin) (@in_guaranteed @callee_guaranteed @substituted <A> () -> @out A for <Int>) -> @error any Error
-
 sil hidden [ossa] @nonTrivialForEachLoopUnrollTest : $@convention(thin) (@owned @callee_guaranteed @substituted <A> () -> @out A for <Int>, @owned @callee_guaranteed @substituted <A> () -> @out A for <Int>) -> () {
 bb0(%0: @owned $@callee_guaranteed @substituted <A> () -> @out A for <Int>, %1: @owned $@callee_guaranteed @substituted <A> () -> @out A for <Int>):
   %2 = integer_literal $Builtin.Word, 2
   %3 = function_ref @_allocateUninitializedArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
   %4 = apply %3<() -> Int>(%2) : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
-  (%5, %6) = destructure_tuple %4 : $(Array<()->Int>, Builtin.RawPointer)
+  (%5, %6a) = destructure_tuple %4 : $(Array<()->Int>, Builtin.RawPointer)
+  %6 = mark_dependence %6a : $Builtin.RawPointer on %5 : $Array<() -> Int>
   %7 = pointer_to_address %6 : $Builtin.RawPointer to [strict] $*@callee_guaranteed @substituted <A> () -> @out A for <Int>
   store %0 to [init] %7 : $*@callee_guaranteed @substituted <A> () -> @out A for <Int>
   %12 = integer_literal $Builtin.Word, 1
@@ -109,47 +146,17 @@ bb1(%32 : $()):
 bb2(%39 : @owned $Error):
   unreachable
 }
-// CHECK-LABEL: nonTrivialForEachLoopUnrollTest
-// CHECK: [[ELEM1:%[0-9]+]] = copy_value %0
-// CHECK-NEXT: store %0 to [init] %{{.*}} : $*@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
-// CHECK: [[ELEM2:%[0-9]+]] = copy_value %1
-// CHECK-NEXT: store %1 to [init] %{{.*}} : $*@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
-// CHECK: [[BODYCLOSURE:%[0-9]+]] = thin_to_thick_function
-// CHECK-NOT: forEach
-// CHECK: [[STACK:%[0-9]+]] = alloc_stack $@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>
-// CHECK: [[ELEM1BORROW:%[0-9]+]] = begin_borrow [[ELEM1]]
-// CHECK: [[SB1:%.*]] = store_borrow [[ELEM1BORROW]] to [[STACK]]
-// CHECK: try_apply [[BODYCLOSURE]]([[SB1]]) : $@noescape @callee_guaranteed (@in_guaranteed @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>) -> @error any Error, normal [[NORMAL:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
-
-// CHECK: [[ERROR]]([[ERRPARAM:%[0-9]+]] : @owned $any Error):
-// CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM]] : $any Error)
-
-// CHECK: [[NORMAL]](%{{.*}} : $()):
-// CHECK: end_borrow [[ELEM1BORROW]]
-// CHECK: [[ELEM2BORROW:%[0-9]+]] = begin_borrow [[ELEM2]]
-// CHECK: [[SB2:%.*]] = store_borrow [[ELEM2BORROW]] to [[STACK]]
-// CHECK: try_apply [[BODYCLOSURE]]([[SB2]]) : $@noescape @callee_guaranteed (@in_guaranteed @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <Int>) -> @error any Error, normal [[NORMAL2:bb[0-9]+]], error [[ERROR2:bb[0-9]+]]
-
-// CHECK: [[ERROR2]]([[ERRPARAM2:%[0-9]+]] : @owned $any Error):
-// CHECK: br [[ERROR3:bb[0-9]+]]([[ERRPARAM2]] : $any Error)
-
-// CHECK: [[NORMAL2]](%{{.*}} : $()):
-// CHECK: end_borrow [[ELEM2BORROW]]
-// CHECK: dealloc_stack [[STACK]]
-// Note that the temporary alloc_stack of the array created for the forEach call
-// will be cleaned up when the forEach call is removed.
-// CHECK: destroy_value
-
-// CHECK: [[ERROR3]]([[ERRPARAM3:%[0-9]+]] : @owned $any Error):
-// CHECK: dealloc_stack [[STACK]]
-// CHECK: unreachable
 
+// CHECK-LABEL: @checkIndirectFixLifetimeUsesAreIgnored
+// CHECK-NOT: function_ref @forEach : $@convention(method) <τ_0_0 where τ_0_0 : Sequence> (@noescape @callee_guaranteed (@in_guaranteed τ_0_0.Element) -> @error any Error, @in_guaranteed τ_0_0) -> @error any Error
+// CHECK: end sil function 'checkIndirectFixLifetimeUsesAreIgnored'
 sil hidden [ossa] @checkIndirectFixLifetimeUsesAreIgnored : $@convention(thin) () -> () {
 bb0:
   %0 = integer_literal $Builtin.Word, 2
   %1 = function_ref @_allocateUninitializedArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
   %2 = apply %1<Builtin.Int64>(%0) : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
-  (%3, %4) = destructure_tuple %2 : $(Array<Builtin.Int64>, Builtin.RawPointer)
+  (%3, %4a) = destructure_tuple %2 : $(Array<Builtin.Int64>, Builtin.RawPointer)
+  %4 = mark_dependence %4a : $Builtin.RawPointer on %3 : $Array<Builtin.Int64>
   %5 = pointer_to_address %4 : $Builtin.RawPointer to [strict] $*Builtin.Int64
   %6 = integer_literal $Builtin.Int64, 15
   store %6 to [trivial] %5 : $*Builtin.Int64
@@ -184,10 +191,10 @@ bb1(%32 : $()):
 bb2(%39 : @owned $Error):
   unreachable
 }
-// CHECK-LABEL: @checkIndirectFixLifetimeUsesAreIgnored
-// CHECK-NOT: function_ref @forEach : $@convention(method) <τ_0_0 where τ_0_0 : Sequence> (@noescape @callee_guaranteed (@in_guaranteed τ_0_0.Element) -> @error any Error, @in_guaranteed τ_0_0) -> @error any Error
-// CHECK: end sil function 'checkIndirectFixLifetimeUsesAreIgnored'
 
+// CHECK-LABEL: @testUnrollOfArrayWithPhiArguments
+// CHECK-NOT: function_ref @forEach : $@convention(method) <τ_0_0 where τ_0_0 : Sequence> (@noescape @callee_guaranteed (@in_guaranteed τ_0_0.Element) -> @error any Error, @in_guaranteed τ_0_0) -> @error any Error
+// CHECK: end sil function 'testUnrollOfArrayWithPhiArguments'
 sil hidden [ossa] @testUnrollOfArrayWithPhiArguments : $@convention(thin) () -> () {
 bb0:
   %0 = integer_literal $Builtin.Int64, 57
@@ -197,7 +204,8 @@ bb1(%arg : $Builtin.Int64):
   %10 = integer_literal $Builtin.Word, 1
   %11 = function_ref @_allocateUninitializedArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
   %12 = apply %11<Builtin.Int64>(%10) : $@convention(thin) <τ_0_0> (Builtin.Word) -> (@owned Array<τ_0_0>, Builtin.RawPointer)
-  (%13, %14) = destructure_tuple %12 : $(Array<Builtin.Int64>, Builtin.RawPointer)
+  (%13, %14a) = destructure_tuple %12 : $(Array<Builtin.Int64>, Builtin.RawPointer)
+  %14 = mark_dependence %14a : $Builtin.RawPointer on %13 : $Array<Builtin.Int64>
   %15 = pointer_to_address %14 : $Builtin.RawPointer to [strict] $*Builtin.Int64
   store %arg to [trivial] %15 : $*Builtin.Int64
   br bb2(%arg : $Builtin.Int64)
@@ -224,7 +232,4 @@ bb3(%32 : $()):
 bb4(%39 : @owned $Error):
   unreachable
 }
-// CHECK-LABEL: @testUnrollOfArrayWithPhiArguments
-// CHECK-NOT: function_ref @forEach : $@convention(method) <τ_0_0 where τ_0_0 : Sequence> (@noescape @callee_guaranteed (@in_guaranteed τ_0_0.Element) -> @error any Error, @in_guaranteed τ_0_0) -> @error any Error
-// CHECK: end sil function 'testUnrollOfArrayWithPhiArguments'
 
diff --git a/test/SILOptimizer/for_each_loop_unroll_test.swift b/test/SILOptimizer/for_each_loop_unroll_test.swift
@@ -64,8 +64,10 @@ func unrollLetArrayLiteralWithClosures(i: Int32, j: Int32) {
   a.forEach { print($0()) }
   // CHECK: [[ALLOCATE:%[0-9]+]] = function_ref @$ss27_allocateUninitializedArrayySayxG_BptBwlF
   // CHECK: [[ARRAYTUP:%[0-9]+]] = apply [[ALLOCATE]]<() -> Int32>
+  // CHECK: [[ARRAYVAL:%[0-9]+]] =  tuple_extract [[ARRAYTUP]] : $(Array<() -> Int32>, Builtin.RawPointer), 0
   // CHECK: [[STORAGEPTR:%[0-9]+]] =  tuple_extract [[ARRAYTUP]] : $(Array<() -> Int32>, Builtin.RawPointer), 1
-  // CHECK: [[STORAGEADDR:%[0-9]+]] = pointer_to_address [[STORAGEPTR]]
+  // CHECK: [[MDI:%[0-9]+]] = mark_dependence [[STORAGEPTR]] : $Builtin.RawPointer on [[ARRAYVAL]] : $Array<() -> Int32>
+  // CHECK: [[STORAGEADDR:%[0-9]+]] = pointer_to_address [[MDI]]
   // CHECK: store [[CLOSURE1:%[0-9]+]] to [[STORAGEADDR]]
   // CHECK: [[INDEX1:%[0-9]+]] = index_addr [[STORAGEADDR]]
   // CHECK: store [[CLOSURE2:%[0-9]+]] to [[INDEX1]]
