Skip to content

Commit fba46ee

Browse files
meg-guptameg-gupta
committed
Fix edgecase in AllocBoxToStack handling of local apply
rdar://66542551
1 parent 9bbe6e7 commit fba46ee

File tree

3 files changed

+228
-15
lines changed

3 files changed

+228
-15
lines changed

‎lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
#define DEBUG_TYPE "allocbox-to-stack"
1414
#include "swift/AST/DiagnosticsSIL.h"
15+
#include "swift/Basic/BlotMapVector.h"
1516
#include "swift/SIL/ApplySite.h"
1617
#include "swift/SIL/Dominance.h"
1718
#include "swift/SIL/SILArgument.h"
@@ -978,8 +979,7 @@ specializeApplySite(SILOptFunctionBuilder &FuncBuilder, ApplySite Apply,
978979
}
979980

980981
static void rewriteApplySites(AllocBoxToStackState &pass) {
981-
llvm::DenseMap<ApplySite, ArgIndexList> IndexMap;
982-
llvm::SmallVector<ApplySite, 8> AppliesToSpecialize;
982+
swift::SmallBlotMapVector<ApplySite, ArgIndexList, 8> AppliesToSpecialize;
983983
ArgIndexList Indices;
984984

985985
// Build a map from the ApplySite to the indices of the operands
@@ -993,29 +993,39 @@ static void rewriteApplySites(AllocBoxToStackState &pass) {
993993
Indices.clear();
994994
Indices.push_back(CalleeArgIndexNumber);
995995

996-
auto iterAndSuccess = IndexMap.try_emplace(Apply, Indices);
997996
// AllocBoxStack opt promotes boxes passed to a chain of applies when it is
998997
// safe to do so. All such applies have to be specialized to take pointer
999998
// arguments instead of box arguments. This has to be done in dfs order.
1000-
// Since PromotedOperands is already populated in dfs order by
1001-
// `recursivelyFindBoxOperandsPromotableToAddress`. Build
1002-
// `AppliesToSpecialize` in the same order.
1003-
if (iterAndSuccess.second) {
1004-
AppliesToSpecialize.push_back(Apply);
1005-
} else {
1006-
iterAndSuccess.first->second.push_back(CalleeArgIndexNumber);
999+
1000+
// PromotedOperands is already populated in dfs order by
1001+
// `recursivelyFindBoxOperandsPromotableToAddress` w.r.t a single alloc_box.
1002+
// AppliesToSpecialize is then populated in the order of PromotedOperands.
1003+
// If multiple alloc_boxes are passed to the same apply instruction, then
1004+
// the apply instruction can appear multiple times in AppliesToSpecialize.
1005+
// Only its last appearance is maintained and previous appearances are
1006+
// blotted.
1007+
auto iterAndSuccess =
1008+
AppliesToSpecialize.insert(std::make_pair(Apply, Indices));
1009+
if (!iterAndSuccess.second) {
1010+
// Blot the previously inserted apply and insert at the end with updated
1011+
// indices
1012+
auto OldIndices = iterAndSuccess.first->getValue().second;
1013+
OldIndices.push_back(CalleeArgIndexNumber);
1014+
AppliesToSpecialize.erase(iterAndSuccess.first);
1015+
AppliesToSpecialize.insert(std::make_pair(Apply, OldIndices));
10071016
}
10081017
}
10091018

10101019
// Clone the referenced function of each ApplySite, removing the
10111020
// operands that we will not need, and remove the existing
10121021
// ApplySite.
10131022
SILOptFunctionBuilder FuncBuilder(*pass.T);
1014-
for (auto &Apply : AppliesToSpecialize) {
1015-
auto It = IndexMap.find(Apply);
1016-
assert(It != IndexMap.end());
1017-
auto &Indices = It->second;
1018-
1023+
for (auto &It : AppliesToSpecialize) {
1024+
if (!It.hasValue()) {
1025+
continue;
1026+
}
1027+
auto Apply = It.getValue().first;
1028+
auto Indices = It.getValue().second;
10191029
// Sort the indices and unique them.
10201030
sortUnique(Indices);
10211031

‎test/SILOptimizer/allocboxtostack_localapply.swift

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,3 +118,67 @@ public func testrecur() -> Int {
118118
}
119119
return bas() + bar()
120120
}
121+
122+
// Test to make sure AppliesToSpecialize in AllocBoxToStack is populated correctly when there are common function calls for mutiple allox_boxes.
123+
// Order of function calls constructed in PromotedOperands: bar common bas common.
124+
// AppliesToSpecialize should have the order: bar bas common.
125+
// Only then, the functions get specialized correctly, and we won't see an assert in checkNoPromotedBoxInApply.
126+
// CHECK-LABEL: sil [noinline] @$s26allocboxtostack_localapply8testdfs1SiyF :
127+
// CHECK-NOT : alloc_box ${ var Int }, var, name "x"
128+
// CHECK-NOT : alloc_box ${ var Int }, var, name "y"
129+
// CHECK-LABEL:} // end sil function '$s26allocboxtostack_localapply8testdfs1SiyF'
130+
@inline(never)
131+
public func testdfs1() -> Int {
132+
var x = 0
133+
var y = 0
134+
@inline(never)
135+
func bar() -> Int {
136+
return x
137+
}
138+
@inline(never)
139+
func bas() -> Int {
140+
return y
141+
}
142+
@inline(never)
143+
func common() -> Int {
144+
return bar() + bas()
145+
}
146+
return common()
147+
}
148+
149+
// Test to make sure we don't optimize the case when we have an inner common function call for multiple boxes.
150+
// We don't optimize this case now, because we don't have additional logic to correctly construct AppliesToSpecialize
151+
// Order of function calls constructed in PromotedOperands: bar innercommon local1 bas innercommon local2
152+
// AppliesToSpecialize should have the order: bar bas innercommon local1 local2
153+
// Since we don't maintain any tree like data structure with more info on the call tree, this is not possible to contruct today
154+
// CHECK-LABEL: sil [noinline] @$s26allocboxtostack_localapply8testdfs2SiyF :
155+
// CHECK: alloc_box ${ var Int }, var, name "x"
156+
// CHECK: alloc_box ${ var Int }, var, name "y"
157+
// CHECK-LABEL:} // end sil function '$s26allocboxtostack_localapply8testdfs2SiyF'
158+
@inline(never)
159+
public func testdfs2() -> Int {
160+
var x = 0
161+
var y = 0
162+
@inline(never)
163+
func bar() -> Int {
164+
return x
165+
}
166+
@inline(never)
167+
func bas() -> Int {
168+
return y
169+
}
170+
@inline(never)
171+
func innercommon() -> Int {
172+
return bar() + bas()
173+
}
174+
@inline(never)
175+
func local1() -> Int {
176+
return innercommon()
177+
}
178+
@inline(never)
179+
func local2() -> Int {
180+
return innercommon()
181+
}
182+
return local1() + local2()
183+
}
184+

‎test/SILOptimizer/allocboxtostack_localapply_ossa.sil

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -349,3 +349,142 @@ bb2:
349349
unwind
350350
}
351351

352+
struct Int {
353+
var _value: Builtin.Int64
354+
}
355+
356+
// Test to make sure AppliesToSpecialize in AllocBoxToStack is populated correctly when there are common function calls for mutiple allox_boxes.
357+
// Order of function calls constructed in PromotedOperands: bar common bas common.
358+
// AppliesToSpecialize should have the order: bar bas common.
359+
// Only then, the functions get specialized correctly, and we won't see an assert in checkNoPromotedBoxInApply.
360+
// CHECK-LABEL: sil [noinline] [ossa] @$testdfs1 :
361+
// CHECK-NOT : alloc_box ${ var Int64 }, var, name "x"
362+
// CHECK-NOT : alloc_box ${ var Int64 }, var, name "y"
363+
// CHECK-LABEL:} // end sil function '$testdfs1'
364+
sil [noinline] [ossa] @$testdfs1 : $@convention(thin) () -> Int64 {
365+
bb0:
366+
%0 = alloc_box ${ var Int64 }, var, name "x"
367+
%1 = project_box %0 : ${ var Int64 }, 0
368+
%2 = integer_literal $Builtin.Int64, 0
369+
%3 = struct $Int64 (%2 : $Builtin.Int64)
370+
store %3 to [trivial] %1 : $*Int64
371+
%5 = alloc_box ${ var Int64 }, var, name "y"
372+
%6 = project_box %5 : ${ var Int64 }, 0
373+
%7 = integer_literal $Builtin.Int64, 0
374+
%8 = struct $Int64 (%7 : $Builtin.Int64)
375+
store %8 to [trivial] %6 : $*Int64
376+
%10 = function_ref @$testdfs1common : $@convention(thin) (@guaranteed { var Int64 }, @guaranteed { var Int64 }) -> Int64
377+
%11 = apply %10(%0, %5) : $@convention(thin) (@guaranteed { var Int64 }, @guaranteed { var Int64 }) -> Int64
378+
destroy_value %5 : ${ var Int64 }
379+
destroy_value %0 : ${ var Int64 }
380+
return %11 : $Int64
381+
}
382+
383+
sil private [noinline] [ossa] @$testdfs1common : $@convention(thin) (@guaranteed { var Int64 }, @guaranteed { var Int64 }) -> Int64 {
384+
bb0(%0 : @guaranteed ${ var Int64 }, %1 : @guaranteed ${ var Int64 }):
385+
%proj1 = project_box %0 : ${ var Int64 }, 0
386+
%proj2 = project_box %1 : ${ var Int64 }, 0
387+
%barfunc = function_ref @$testdfs1bar : $@convention(thin) (@guaranteed { var Int64 }) -> Int64
388+
%res1 = apply %barfunc(%0) : $@convention(thin) (@guaranteed { var Int64 }) -> Int64
389+
%basfunc = function_ref @$testdfs1bas : $@convention(thin) (@guaranteed { var Int64 }) -> Int64
390+
%res2 = apply %basfunc(%1) : $@convention(thin) (@guaranteed { var Int64 }) -> Int64
391+
%func = function_ref @$blackhole : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
392+
%tmp1 = apply %func<Int64>(%proj1) : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
393+
%tmp2 = apply %func<Int64>(%proj2) : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
394+
%res = load [trivial] %proj1 : $*Int64
395+
return %res : $Int64
396+
}
397+
398+
sil private [noinline] [ossa] @$testdfs1bar : $@convention(thin) (@guaranteed { var Int64 }) -> Int64 {
399+
bb0(%0 : @guaranteed ${ var Int64 }):
400+
%1 = project_box %0 : ${ var Int64 }, 0
401+
%4 = load [trivial] %1 : $*Int64
402+
return %4 : $Int64
403+
}
404+
405+
sil private [noinline] [ossa] @$testdfs1bas : $@convention(thin) (@guaranteed { var Int64 }) -> Int64 {
406+
bb0(%0 : @guaranteed ${ var Int64 }):
407+
%1 = project_box %0 : ${ var Int64 }, 0
408+
%4 = load [trivial] %1 : $*Int64
409+
return %4 : $Int64
410+
}
411+
412+
// Test to make sure we don't optimize the case when we have an inner common function call for multiple boxes.
413+
// We don't optimize this case now, because we don't have additional logic to correctly construct AppliesToSpecialize
414+
// Order of function calls constructed in PromotedOperands: bar innercommon local1 bas innercommon local2
415+
// AppliesToSpecialize should have the order: bar bas innercommon local1 local2
416+
// Since we don't maintain any tree like data structure with more info on the call tree, this is not possible to contruct today
417+
// CHECK-LABEL: sil [noinline] @$s26allocboxtostack_localapply8testdfs2SiyF :
418+
// CHECK: alloc_box ${ var Int }, var, name "x"
419+
// CHECK: alloc_box ${ var Int }, var, name "y"
420+
// CHECK-LABEL:} // end sil function '$s26allocboxtostack_localapply8testdfs2SiyF'
421+
sil [noinline] [ossa] @$testdfs2 : $@convention(thin) () -> Int {
422+
bb0:
423+
%0 = alloc_box ${ var Int }, var, name "x"
424+
%1 = project_box %0 : ${ var Int }, 0
425+
%2 = integer_literal $Builtin.Int64, 0
426+
%3 = struct $Int (%2 : $Builtin.Int64)
427+
store %3 to [trivial] %1 : $*Int
428+
%5 = alloc_box ${ var Int }, var, name "y"
429+
%6 = project_box %5 : ${ var Int }, 0
430+
%7 = integer_literal $Builtin.Int64, 0
431+
%8 = struct $Int (%7 : $Builtin.Int64)
432+
store %8 to [trivial] %6 : $*Int
433+
%10 = function_ref @$testdfs2local1 : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
434+
%11 = apply %10(%0, %5) : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
435+
%12 = function_ref @$testdfs2local2 : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
436+
%13 = apply %12(%0, %5) : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
437+
%14 = struct_extract %11 : $Int, #Int._value
438+
%15 = struct_extract %13 : $Int, #Int._value
439+
%16 = integer_literal $Builtin.Int1, -1
440+
%17 = builtin "sadd_with_overflow_Int64"(%14 : $Builtin.Int64, %15 : $Builtin.Int64, %16 : $Builtin.Int1) : $(Builtin.Int64, Builtin.Int1)
441+
(%18, %19) = destructure_tuple %17 : $(Builtin.Int64, Builtin.Int1)
442+
cond_fail %19 : $Builtin.Int1, "arithmetic overflow"
443+
%21 = struct $Int (%18 : $Builtin.Int64)
444+
destroy_value %5 : ${ var Int }
445+
destroy_value %0 : ${ var Int }
446+
return %21 : $Int
447+
}
448+
449+
sil private [noinline] [ossa] @$testdfs2bar : $@convention(thin) (@guaranteed { var Int }) -> Int {
450+
bb0(%0 : @guaranteed ${ var Int }):
451+
%1 = project_box %0 : ${ var Int }, 0
452+
%4 = load [trivial] %1 : $*Int
453+
return %4 : $Int
454+
}
455+
456+
sil private [noinline] [ossa] @$testdfs2bas : $@convention(thin) (@guaranteed { var Int }) -> Int {
457+
bb0(%0 : @guaranteed ${ var Int }):
458+
%1 = project_box %0 : ${ var Int }, 0
459+
%4 = load [trivial] %1 : $*Int
460+
return %4 : $Int
461+
}
462+
463+
sil private [noinline] [ossa] @$testdfs2innercommon : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int {
464+
bb0(%0 : @guaranteed ${ var Int }, %1 : @guaranteed ${ var Int }):
465+
%2 = project_box %0 : ${ var Int }, 0
466+
%4 = project_box %1 : ${ var Int }, 0
467+
%8 = function_ref @$testdfs2bar : $@convention(thin) (@guaranteed { var Int }) -> Int
468+
%9 = apply %8(%0) : $@convention(thin) (@guaranteed { var Int }) -> Int
469+
%11 = function_ref @$testdfs2bas : $@convention(thin) (@guaranteed { var Int }) -> Int
470+
%12 = apply %11(%1) : $@convention(thin) (@guaranteed { var Int }) -> Int
471+
return %12 : $Int
472+
}
473+
474+
sil private [noinline] [ossa] @$testdfs2local1 : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int {
475+
bb0(%0 : @guaranteed ${ var Int }, %1 : @guaranteed ${ var Int }):
476+
%2 = project_box %0 : ${ var Int }, 0
477+
%4 = project_box %1 : ${ var Int }, 0
478+
%7 = function_ref @$testdfs2innercommon : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
479+
%8 = apply %7(%0, %1) : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
480+
return %8 : $Int
481+
}
482+
483+
sil private [noinline] [ossa] @$testdfs2local2 : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int {
484+
bb0(%0 : @guaranteed ${ var Int }, %1 : @guaranteed ${ var Int }):
485+
%2 = project_box %0 : ${ var Int }, 0
486+
%4 = project_box %1 : ${ var Int }, 0
487+
%7 = function_ref @$testdfs2innercommon : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
488+
%8 = apply %7(%0, %1) : $@convention(thin) (@guaranteed { var Int }, @guaranteed { var Int }) -> Int
489+
return %8 : $Int
490+
}

0 commit comments

Comments
 (0)