Skip to content

Commit 20ae468

Browse files
committed
Address review comments
1 parent 27478f6 commit 20ae468

File tree

6 files changed

+161
-166
lines changed

6 files changed

+161
-166
lines changed

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

Lines changed: 116 additions & 121 deletions
Large diffs are not rendered by default.

SwiftCompilerSources/Sources/Optimizer/Utilities/Test.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ public func registerOptimizerTests() {
160160
enclosingValuesTest,
161161
forwardingDefUseTest,
162162
forwardingUseDefTest,
163-
gatherCallSiteTest,
163+
getPullbackClosureInfoTest,
164164
interiorLivenessTest,
165165
lifetimeDependenceRootTest,
166166
lifetimeDependenceScopeTest,

test/AutoDiff/SILOptimizer/closure_specialization/multi_bb_bte.sil

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -46,15 +46,15 @@ bb0(%0 : $Float, %1 : $_AD__$s4test5mul42yS2fSgF_bb2__Pred__src_0_wrt_0, %2 : $@
4646
sil hidden @$s4test5mul42yS2fSgFTJrSpSr : $@convention(thin) (Optional<Float>) -> (Float, @owned @callee_guaranteed (Float) -> Optional<Float>.TangentVector) {
4747
bb0(%0 : $Optional<Float>):
4848
//=========== Test callsite and closure gathering logic ===========//
49-
specify_test "closure_specialize_gather_call_site"
49+
specify_test "autodiff_closure_specialize_get_pullback_closure_info"
5050
// TRUNNER-LABEL: Specializing closures in function: $s4test5mul42yS2fSgFTJrSpSr
51-
// TRUNNER: PartialApply call site: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]]) : $@convention(thin) (Float, @owned _AD__$s4test5mul42yS2fSgF_bb2__Pred__src_0_wrt_0, @owned @callee_guaranteed (Float) -> Float) -> Optional<Float>.TangentVector
51+
// TRUNNER: PartialApply of pullback: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]]) : $@convention(thin) (Float, @owned _AD__$s4test5mul42yS2fSgF_bb2__Pred__src_0_wrt_0, @owned @callee_guaranteed (Float) -> Float) -> Optional<Float>.TangentVector
5252
// TRUNNER-NEXT: Passed in closures:
5353
// TRUNNER-NEXT: 1. %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]]) : $@convention(thin) (Float, Float, Float) -> (Float, Float)
5454
// TRUNNER-EMPTY:
5555

5656
//=========== Test specialized function signature and body ===========//
57-
specify_test "closure_specialize_specialized_function_signature_and_body"
57+
specify_test "autodiff_closure_specialize_specialized_function_signature_and_body"
5858
// TRUNNER-LABEL: Generated specialized function: $s4test5mul42yS2fSgFTJpSpSr073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktj1_k5FZSf_K6SfcfU_S2fTf1nnc_n
5959
// CHECK: sil private [signature_optimized_thunk] [always_inline] @$s4test5mul42yS2fSgFTJpSpSr073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktj1_k5FZSf_K6SfcfU_S2fTf1nnc_n : $@convention(thin) (Float, @owned _AD__$s4test5mul42yS2fSgF_bb2__Pred__src_0_wrt_0, Float, Float) -> Optional<Float>.TangentVector {
6060
// CHECK: bb0(%0 : $Float, %1 : $_AD__$s4test5mul42yS2fSgF_bb2__Pred__src_0_wrt_0, %2 : $Float, %3 : $Float):
@@ -71,7 +71,7 @@ bb0(%0 : $Optional<Float>):
7171
// CHECK: return %[[#A11]] : $Optional<Float>.TangentVector
7272

7373
//=========== Test rewritten body ===========//
74-
specify_test "closure_specialize_rewritten_caller_body"
74+
specify_test "autodiff_closure_specialize_rewritten_caller_body"
7575
// TRUNNER-LABEL: Rewritten caller body for: $s4test5mul42yS2fSgFTJrSpSr:
7676
// CHECK: sil hidden @$s4test5mul42yS2fSgFTJrSpSr : $@convention(thin) (Optional<Float>) -> (Float, @owned @callee_guaranteed (Float) -> Optional<Float>.TangentVector) {
7777
// CHECK: bb1(%2 : $Float):
@@ -300,16 +300,16 @@ bb6(%122 : $Builtin.FPIEEE32):
300300
sil hidden @$s4test5ClassV6methodSfyFTJrSpSr : $@convention(method) (Class) -> (Float, @owned @callee_guaranteed (Float) -> Class.TangentVector) {
301301
bb0(%0 : $Class):
302302
//=========== Test callsite and closure gathering logic ===========//
303-
specify_test "closure_specialize_gather_call_site"
303+
specify_test "autodiff_closure_specialize_get_pullback_closure_info"
304304
// TRUNNER-LABEL: Specializing closures in function: $s4test5ClassV6methodSfyFTJrSpSr
305-
// TRUNNER: PartialApply call site: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]], %[[#C42:]]) : $@convention(thin) (Float, @owned _AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Class.TangentVector
305+
// TRUNNER: PartialApply of pullback: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]], %[[#C42:]]) : $@convention(thin) (Float, @owned _AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0, @owned @callee_guaranteed (Float) -> Float, @owned @callee_guaranteed (Float) -> (Float, Float)) -> Class.TangentVector
306306
// TRUNNER-NEXT: Passed in closures:
307307
// TRUNNER-NEXT: 1. %[[#]] = partial_apply [callee_guaranteed] %[[#C7:]](%[[#C34:]], %[[#]]) : $@convention(thin) (Float, Float, Float) -> (Float, Float)
308308
// TRUNNER-NEXT: 2. %[[#C42]] = partial_apply [callee_guaranteed] %[[#C7]](%[[#C34]], %[[#]]) : $@convention(thin) (Float, Float, Float) -> (Float, Float)
309309
// TRUNNER-EMPTY:
310310

311311
//=========== Test specialized function signature and body ===========//
312-
specify_test "closure_specialize_specialized_function_signature_and_body"
312+
specify_test "autodiff_closure_specialize_specialized_function_signature_and_body"
313313
// TRUNNER-LABEL: Generated specialized function: $s4test5ClassV6methodSfyFTJpSpSr073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktk1_l5FZSf_L6SfcfU_S2fAES2fTf1nncc_n
314314
// CHECK: sil private @$s4test5ClassV6methodSfyFTJpSpSr073$sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktk1_l5FZSf_L6SfcfU_S2fAES2fTf1nncc_n : $@convention(thin) (Float, @owned _AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0, Float, Float, Float, Float) -> Class.TangentVector {
315315
// CHECK: bb0(%0 : $Float, %1 : $_AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0, %2 : $Float, %3 : $Float, %4 : $Float, %5 : $Float):
@@ -330,7 +330,7 @@ bb0(%0 : $Class):
330330
// TRUNNER: strong_release %[[#D9]] : $@callee_guaranteed (Float) -> Float
331331

332332
//=========== Test rewritten body ===========//
333-
specify_test "closure_specialize_rewritten_caller_body"
333+
specify_test "autodiff_closure_specialize_rewritten_caller_body"
334334
// TRUNNER-LABEL: Rewritten caller body for: $s4test5ClassV6methodSfyFTJrSpSr:
335335
// CHECK: sil hidden @$s4test5ClassV6methodSfyFTJrSpSr : $@convention(method) (Class) -> (Float, @owned @callee_guaranteed (Float) -> Class.TangentVector) {
336336
// CHECK: bb0(%0 : $Class):
@@ -569,17 +569,17 @@ sil hidden @$s4test14cond_tuple_varyS2fFTJrSpSr : $@convention(thin) (Float) ->
569569
[global: ]
570570
bb0(%0 : $Float):
571571
//=========== Test callsite and closure gathering logic ===========//
572-
specify_test "closure_specialize_gather_call_site"
572+
specify_test "autodiff_closure_specialize_get_pullback_closure_info"
573573
// TRUNNER-LABEL: Specializing closures in function: $s4test14cond_tuple_varyS2fFTJrSpSr
574-
// TRUNNER: PartialApply call site: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#F1:]], %[[#F2:]], %[[#F3:]]) : $@convention(thin) (Float, @owned _AD__$s4test14cond_tuple_varyS2fF_bb3__Pred__src_0_wrt_0, @owned @callee_guaranteed (Float) -> (Float, Float), @owned @callee_guaranteed (Float) -> (Float, Float), @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float
574+
// TRUNNER: PartialApply of pullback: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#F1:]], %[[#F2:]], %[[#F3:]]) : $@convention(thin) (Float, @owned _AD__$s4test14cond_tuple_varyS2fF_bb3__Pred__src_0_wrt_0, @owned @callee_guaranteed (Float) -> (Float, Float), @owned @callee_guaranteed (Float) -> (Float, Float), @owned @callee_guaranteed (Float) -> (Float, Float)) -> Float
575575
// TRUNNER-NEXT: Passed in closures:
576576
// TRUNNER-NEXT: 1. %[[#F1]] = thin_to_thick_function %[[#]] : $@convention(thin) (Float) -> (Float, Float) to $@callee_guaranteed (Float) -> (Float, Float)
577577
// TRUNNER-NEXT: 2. %[[#F2]] = thin_to_thick_function %[[#]] : $@convention(thin) (Float) -> (Float, Float) to $@callee_guaranteed (Float) -> (Float, Float)
578578
// TRUNNER-NEXT: 3. %[[#F3]] = thin_to_thick_function %[[#]] : $@convention(thin) (Float) -> (Float, Float) to $@callee_guaranteed (Float) -> (Float, Float)
579579
// TRUNNER-EMPTY:
580580

581581
//=========== Test specialized function signature and body ===========//
582-
specify_test "closure_specialize_specialized_function_signature_and_body"
582+
specify_test "autodiff_closure_specialize_specialized_function_signature_and_body"
583583
// TRUNNER-LABEL: Generated specialized function: $s4test14cond_tuple_varyS2fFTJpSpSr067$sSf16_DifferentiationE7_vjpAdd3lhs3rhsSf5value_Sf_SftSfc8pullbacktl1_m5FZSf_M6SfcfU_0ef1_g4E12_i16Subtract3lhs3rhsk1_l1_mnl1_mo1_mP2U_ACTf1nnccc_n
584584
// CHECK: sil private @$s4test14cond_tuple_varyS2fFTJpSpSr067$sSf16_DifferentiationE7_vjpAdd3lhs3rhsSf5value_Sf_SftSfc8pullbacktl1_m5FZSf_M6SfcfU_0ef1_g4E12_i16Subtract3lhs3rhsk1_l1_mnl1_mo1_mP2U_ACTf1nnccc_n : $@convention(thin) (Float, @owned _AD__$s4test14cond_tuple_varyS2fF_bb3__Pred__src_0_wrt_0) -> Float {
585585
// CHECK: bb0(%0 : $Float, %1 : $_AD__$s4test14cond_tuple_varyS2fF_bb3__Pred__src_0_wrt_0):
@@ -609,7 +609,7 @@ bb0(%0 : $Float):
609609
// CHECK: {{^}}bb3{{.*}}:
610610

611611
//=========== Test rewritten body ===========//
612-
specify_test "closure_specialize_rewritten_caller_body"
612+
specify_test "autodiff_closure_specialize_rewritten_caller_body"
613613
// TRUNNER-LABEL: Rewritten caller body for: $s4test14cond_tuple_varyS2fFTJrSpSr:
614614
// CHECK: sil hidden @$s4test14cond_tuple_varyS2fFTJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) {
615615
// CHECK: bb0(%0 : $Float):

test/AutoDiff/SILOptimizer/closure_specialization/multi_bb_no_bte1.sil

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -93,15 +93,15 @@ bb0(%0 : $Float, %1 : $@callee_guaranteed (Float) -> Class.TangentVector):
9393
sil hidden @$s4test13methodWrapperySfAA5ClassVFTJrSpSr : $@convention(thin) (Class) -> (Float, @owned @callee_guaranteed (Float) -> Class.TangentVector) {
9494
bb0(%0 : $Class):
9595
//=========== Test callsite and closure gathering logic ===========//
96-
specify_test "closure_specialize_gather_call_site"
96+
specify_test "autodiff_closure_specialize_get_pullback_closure_info"
9797
// TRUNNER-LABEL: Specializing closures in function: $s4test13methodWrapperySfAA5ClassVFTJrSpSr
98-
// TRUNNER: PartialApply call site: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#A36:]]) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Class.TangentVector) -> Class.TangentVector
98+
// TRUNNER: PartialApply of pullback: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#A36:]]) : $@convention(thin) (Float, @owned @callee_guaranteed (Float) -> Class.TangentVector) -> Class.TangentVector
9999
// TRUNNER-NEXT: Passed in closures:
100100
// TRUNNER-NEXT: 1. %[[#A36]] = partial_apply [callee_guaranteed] %[[#]](%[[#]]) : $@convention(thin) (Float, @owned _AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0) -> Class.TangentVector
101101
// TRUNNER-EMPTY:
102102

103103
//=========== Test specialized function signature and body ===========//
104-
specify_test "closure_specialize_specialized_function_signature_and_body"
104+
specify_test "autodiff_closure_specialize_specialized_function_signature_and_body"
105105
// TRUNNER-LABEL: Generated specialized function: $s4test13methodWrapperySfAA5ClassVFTJpSpSr08$s4test5D19V6methodSfyFTJpSpSr4main05_AD__edfG24F_bb3__Pred__src_0_wrt_0OTf1nc_n
106106
// CHECK: sil private [signature_optimized_thunk] [always_inline] @$s4test13methodWrapperySfAA5ClassVFTJpSpSr08$s4test5D19V6methodSfyFTJpSpSr4main05_AD__edfG24F_bb3__Pred__src_0_wrt_0OTf1nc_n : $@convention(thin) (Float, @owned _AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0) -> Class.TangentVector {
107107
// CHECK: bb0(%0 : $Float, %1 : $_AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0):
@@ -114,7 +114,7 @@ bb0(%0 : $Class):
114114
// CHECK: return %[[#B4]]
115115

116116
//=========== Test rewritten body ===========//
117-
specify_test "closure_specialize_rewritten_caller_body"
117+
specify_test "autodiff_closure_specialize_rewritten_caller_body"
118118
// TRUNNER-LABEL: Rewritten caller body for: $s4test13methodWrapperySfAA5ClassVFTJrSpSr:
119119
// CHECK: sil hidden @$s4test13methodWrapperySfAA5ClassVFTJrSpSr : $@convention(thin) (Class) -> (Float, @owned @callee_guaranteed (Float) -> Class.TangentVector) {
120120
// CHECK: bb3(%[[#C33:]] : $Float, %[[#C34:]] : $_AD__$s4test5ClassV6methodSfyF_bb3__Pred__src_0_wrt_0):

test/AutoDiff/SILOptimizer/closure_specialization/multi_bb_no_bte2.sil

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,15 +50,15 @@ bb0(%0 : $Float, %1 : $@callee_guaranteed (@guaranteed Array<Float>.Differentiab
5050
sil hidden @$s4test4sum1ySfSaySfG_ACtFTJrSSpSr : $@convention(thin) (@guaranteed Array<Float>, @guaranteed Array<Float>) -> (Float, @owned @callee_guaranteed (Float) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView)) {
5151
bb0(%0 : $Array<Float>, %1 : $Array<Float>):
5252
//=========== Test callsite and closure gathering logic ===========//
53-
specify_test "closure_specialize_gather_call_site"
53+
specify_test "autodiff_closure_specialize_get_pullback_closure_info"
5454
// TRUNNER-LABEL: Specializing closures in function: $s4test4sum1ySfSaySfG_ACtFTJrSSpSr
55-
// TRUNNER: PartialApply call site: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]]) : $@convention(thin) (Float, @owned @callee_guaranteed (@guaranteed Array<Float>.DifferentiableView) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView), @owned @callee_guaranteed (@in_guaranteed Float) -> @owned Array<Float>.DifferentiableView) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView)
55+
// TRUNNER: PartialApply of pullback: %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]]) : $@convention(thin) (Float, @owned @callee_guaranteed (@guaranteed Array<Float>.DifferentiableView) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView), @owned @callee_guaranteed (@in_guaranteed Float) -> @owned Array<Float>.DifferentiableView) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView)
5656
// TRUNNER-NEXT: Passed in closures:
5757
// TRUNNER-NEXT: 1. %[[#]] = partial_apply [callee_guaranteed] %[[#]](%[[#]], %[[#]]) : $@convention(thin) (@in_guaranteed Float, @guaranteed Array<Float>, Int) -> @owned Array<Float>.DifferentiableView
5858
// TRUNNER-EMPTY:
5959

6060
//=========== Test specialized function signature and body ===========//
61-
specify_test "closure_specialize_specialized_function_signature_and_body"
61+
specify_test "autodiff_closure_specialize_specialized_function_signature_and_body"
6262
// TRUNNER-LABEL: Generated specialized function: $s4test4sum1ySfSaySfG_ACtFTJpSSpSr055$sSfSa16_DifferentiationAA14DifferentiableRzlE0B4ViewVyd8_GIegno_D10AEIegyo_TRSfSa01_F0AE0H0RzlE0hL0VySf_GIegno_Tf1nnc_n0ce1_fghi39E13_vjpSubscript5indexx5value_SaA2aBRzljkl48Vy13TangentVectorQz_GAIc8pullbacktSi_tFAKL_yAjiaV7FSf_TG5ACSiTf1nnc_n
6363
// CHECK: sil private [signature_optimized_thunk] [always_inline] @$s4test4sum1ySfSaySfG_ACtFTJpSSpSr055$sSfSa16_DifferentiationAA14DifferentiableRzlE0B4ViewVyd8_GIegno_D10AEIegyo_TRSfSa01_F0AE0H0RzlE0hL0VySf_GIegno_Tf1nnc_n0ce1_fghi39E13_vjpSubscript5indexx5value_SaA2aBRzljkl48Vy13TangentVectorQz_GAIc8pullbacktSi_tFAKL_yAjiaV7FSf_TG5ACSiTf1nnc_n : $@convention(thin) (Float, @owned @callee_guaranteed (@guaranteed Array<Float>.DifferentiableView) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView), @owned Array<Float>, Int) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView) {
6464
// CHECK: bb0(%0 : $Float, %1 : $@callee_guaranteed (@guaranteed Array<Float>.DifferentiableView) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView), %2 : $Array<Float>, %3 : $Int):
@@ -75,7 +75,7 @@ bb0(%0 : $Array<Float>, %1 : $Array<Float>):
7575
// CHECK: return %[[#D17]] : $(Array<Float>.DifferentiableView, Array<Float>.DifferentiableView)
7676

7777
//=========== Test rewritten body ===========//
78-
specify_test "closure_specialize_rewritten_caller_body"
78+
specify_test "autodiff_closure_specialize_rewritten_caller_body"
7979
// TRUNNER-LABEL: Rewritten caller body for: $s4test4sum1ySfSaySfG_ACtFTJrSSpSr:
8080
// CHECK: sil hidden @$s4test4sum1ySfSaySfG_ACtFTJrSSpSr : $@convention(thin) (@guaranteed Array<Float>, @guaranteed Array<Float>) -> (Float, @owned @callee_guaranteed (Float) -> (@owned Array<Float>.DifferentiableView, @owned Array<Float>.DifferentiableView)) {
8181
// CHECK: bb8:

0 commit comments

Comments
 (0)