Skip to content

Commit 5762912

Browse files
aschwaighoferaschwaighofer
authored
Merge pull request #80639 from aschwaighofer/store_once_async_context_addr
IRGen: Don't stored the current async frame to llvm.swift.async.context.addr twice
2 parents a1430ab + 0bb30d6 commit 5762912

File tree

6 files changed

+57
-13
lines changed

6 files changed

+57
-13
lines changed

‎lib/IRGen/GenCall.cpp

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -262,10 +262,14 @@ llvm::CallInst *IRGenFunction::emitSuspendAsyncCall(
262262
llvm::Function *projectFn = cast<llvm::Function>(
263263
(cast<llvm::Constant>(args[2])->stripPointerCasts()));
264264
auto *fnTy = projectFn->getFunctionType();
265-
266-
llvm::Value *context =
265+
llvm::Value *callerContext = nullptr;
266+
if (projectFn == getOrCreateResumePrjFn()) {
267+
callerContext = popAsyncContext(calleeContext);
268+
} else {
269+
callerContext =
267270
Builder.CreateCallWithoutDbgLoc(fnTy, projectFn, {calleeContext});
268-
storeCurrentAsyncContext(context);
271+
}
272+
storeCurrentAsyncContext(callerContext);
269273
}
270274

271275
return id;

‎lib/IRGen/GenFunc.cpp

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2810,8 +2810,7 @@ void irgen::emitBlockHeader(IRGenFunction &IGF,
28102810
IGF.Builder.CreateStructGEP(headerAddr, 4, layout));
28112811
}
28122812

2813-
llvm::Value *
2814-
IRGenFunction::emitAsyncResumeProjectContext(llvm::Value *calleeContext) {
2813+
llvm::Value *IRGenFunction::popAsyncContext(llvm::Value *calleeContext) {
28152814
auto addr = Builder.CreateBitOrPointerCast(calleeContext, IGM.Int8PtrPtrTy);
28162815
Address callerContextAddr(addr, IGM.Int8PtrTy, IGM.getPointerAlignment());
28172816
llvm::Value *callerContext = Builder.CreateLoad(callerContextAddr);
@@ -2820,6 +2819,13 @@ IRGenFunction::emitAsyncResumeProjectContext(llvm::Value *calleeContext) {
28202819
PointerAuthInfo::emit(*this, schema, addr, PointerAuthEntity());
28212820
callerContext = emitPointerAuthAuth(*this, callerContext, authInfo);
28222821
}
2822+
return callerContext;
2823+
}
2824+
2825+
llvm::Value *
2826+
IRGenFunction::emitAsyncResumeProjectContext(llvm::Value *calleeContext) {
2827+
auto callerContext = popAsyncContext(calleeContext);
2828+
28232829
// TODO: remove this once all platforms support lowering the intrinsic.
28242830
// At the time of this writing only arm64 supports it.
28252831
if (IGM.TargetInfo.canUseSwiftAsyncContextAddrIntrinsic()) {

‎lib/IRGen/IRGenFunction.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,7 @@ class IRGenFunction {
196196

197197
llvm::Value *emitAsyncResumeProjectContext(llvm::Value *callerContextAddr);
198198
llvm::Function *getOrCreateResumePrjFn();
199+
llvm::Value *popAsyncContext(llvm::Value *calleeContext);
199200
llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,
200201
ArrayRef<llvm::Value *> args);
201202
llvm::Function *createAsyncDispatchFn(const FunctionPointer &fnPtr,

‎test/Distributed/distributed_actor_accessor_thunks_32bit.swift

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,8 @@ public distributed actor MyOtherActor {
145145
// CHECK-SAME: %T27distributed_actor_accessors7MyActorC* [[ACTOR]])
146146

147147
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, %swift.error* } [[THUNK_RESULT]], 0
148-
// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
148+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
149+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
149150
// CHECK: {{.*}} = call i1 (i8*, i1, ...) @llvm.coro.end.async({{.*}}, %swift.context* {{.*}}, %swift.error* {{.*}})
150151

151152
/// ---> Thunk and distributed method accessor for `simple2`
@@ -178,7 +179,8 @@ public distributed actor MyOtherActor {
178179
// CHECK-SAME: %T27distributed_actor_accessors7MyActorC* [[ACTOR]])
179180

180181
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i32, i32, i32, %swift.error* } [[THUNK_RESULT]], 0
181-
// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
182+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
183+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
182184

183185
/// Initialize the result buffer with values produced by the thunk
184186

@@ -247,7 +249,8 @@ public distributed actor MyOtherActor {
247249
// CHECK-SAME: %T27distributed_actor_accessors7MyActorC* [[ACTOR]])
248250

249251
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 0
250-
// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
252+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
253+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
251254
// CHECK: [[INT_RES:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 1
252255
// CHECK: %._value = getelementptr inbounds %TSi, %TSi* [[TYPED_RESULT_BUFF]], i32 0, i32 0
253256
// CHECK: store i32 [[INT_RES]], i32* %._value
@@ -292,7 +295,8 @@ public distributed actor MyOtherActor {
292295

293296
// CHECK: [[THUNK_RESULT:%.*]] = call { i8*, i32, %swift.error* } (i32, i8*, i8*, ...) @llvm.coro.suspend.async.sl_p0i8i32p0s_swift.errorss({{.*}}, %swift.context* {{.*}}, i32 [[NATIVE_ENUM_VAL]], i32 [[NATIVE_INT_VAL]], %T27distributed_actor_accessors7MyActorC* {{.*}})
294297
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 0
295-
// CHECK-NEXT: {{.*}} = call i8* @__swift_async_resume_project_context(i8* [[TASK_REF]])
298+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
299+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
296300
// CHECK: [[ENUM_RESULT:%.*]] = extractvalue { i8*, i32, %swift.error* } [[THUNK_RESULT]], 1
297301
// CHECK: [[NATIVE_RESULT_PTR:%.*]] = bitcast %T27distributed_actor_accessors9IndirectEO* [[TYPED_RESULT_BUFF]] to i32*
298302
// CHECK-NEXT: store i32 [[ENUM_RESULT]], i32* [[NATIVE_RESULT_PTR]]

‎test/Distributed/distributed_actor_accessor_thunks_64bit.swift

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,8 @@ public distributed actor MyOtherActor {
141141
// CHECK-SAME: ptr [[ACTOR]])
142142

143143
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, ptr } [[THUNK_RESULT]], 0
144-
// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
144+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
145+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
145146
// CHECK: {{.*}} = call i1 (ptr, i1, ...) @llvm.coro.end.async({{.*}}, ptr {{.*}}, ptr {{.*}})
146147

147148
/// ---> Thunk and distributed method accessor for `simple2`
@@ -172,7 +173,8 @@ public distributed actor MyOtherActor {
172173
// CHECK-SAME: ptr [[ACTOR]])
173174

174175
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, i64, ptr, ptr } [[THUNK_RESULT]], 0
175-
// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
176+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
177+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
176178

177179
/// Extract information about `String` from result and call `end`
178180

@@ -222,7 +224,8 @@ public distributed actor MyOtherActor {
222224
// CHECK-SAME: ptr [[ACTOR]])
223225

224226
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 0
225-
// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
227+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
228+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
226229
// CHECK: [[INT_RES:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 1
227230
// CHECK: %._value = getelementptr inbounds %TSi, ptr [[RESULT_BUFF]], i32 0, i32 0
228231
// CHECK: store i64 [[INT_RES]], ptr %._value
@@ -262,7 +265,8 @@ public distributed actor MyOtherActor {
262265

263266
// CHECK: [[THUNK_RESULT:%.*]] = call { ptr, i64, ptr } (i32, ptr, ptr, ...) @llvm.coro.suspend.async.sl_p0i64p0s({{.*}}, ptr {{.*}}, i64 [[NATIVE_ENUM_VAL]], i64 [[NATIVE_INT_VAL]], ptr {{.*}})
264267
// CHECK-NEXT: [[TASK_REF:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 0
265-
// CHECK-NEXT: {{.*}} = call ptr @__swift_async_resume_project_context(ptr [[TASK_REF]])
268+
// CHECK-NEXT: [[CALLER_ASYNC_CTXT:%.*]] = load ptr, ptr [[TASK_REF]]
269+
// CHECK-NEXT: store ptr [[CALLER_ASYNC_CTXT]], ptr
266270
// CHECK: [[ENUM_RESULT:%.*]] = extractvalue { ptr, i64, ptr } [[THUNK_RESULT]], 1
267271
// CHECK: store i64 [[ENUM_RESULT]], ptr [[RESULT_BUFF]]
268272

‎test/IRGen/async/async_context_projection_arm64e.swift

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// RUN: %target-swift-frontend -emit-irgen %s | %FileCheck %s
2+
3+
// REQUIRES: CPU=arm64e
4+
5+
public func callee() async -> Int {
6+
return 5
7+
}
8+
9+
public func caller() async {
10+
let r = await callee()
11+
print(r)
12+
}
13+
14+
15+
// CHECK: define swifttailcc void @"$s31async_context_projection_arm64e6calleryyYaF"(ptr swiftasync %0)
16+
// CHECK: [[CURR_ASYNC_CTXT:%.*]] = alloca ptr
17+
// CHECK: [[SUSPEND_RES:%.*]] = call { ptr, i64 } (i32, ptr, ptr, ...) @llvm.coro.suspend.async.sl_p0i64s(i32 0, ptr {{.*}}, ptr @__swift_async_resume_project_context, ptr @"$s31async_context_projection_arm64e6calleryyYaF{{.*}}", ptr @"$s31async_context_projection_arm64e6calleeSiyYaF", ptr {{.*}})
18+
// CHECK: [[CALLEE_CTXT:%.*]] = extractvalue { ptr, i64 } [[SUSPEND_RES]], 0
19+
// CHECK: [[CALLER_CTXT:%.*]] = load ptr, ptr [[CALLEE_CTXT]]
20+
// CHECK: [[CALLER_CTXT_ADDR:%.*]] = ptrtoint ptr [[CALLEE_CTXT]] to i64
21+
// CHECK: [[DIS:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[CALLER_CTXT_ADDR]], i64 48546)
22+
// CHECK: [[CALLER_CTXT_PTRAUTH:%.*]] = ptrtoint ptr [[CALLER_CTXT]] to i64
23+
// CHECK: [[CALLER_CTXT_AUTH:%.*]] = call i64 @llvm.ptrauth.auth(i64 [[CALLER_CTXT_PTRAUTH]], i32 2, i64 [[DIS]])
24+
// CHECK: [[CALLER_CTXT_PTR:%.*]] = inttoptr i64 [[CALLER_CTXT_AUTH]] to ptr
25+
// CHECK: store ptr [[CALLER_CTXT_PTR]], ptr [[CURR_ASYNC_CTXT]]

0 commit comments

Comments
 (0)