[LICM] Don't hoist threadlocals within presplit coroutines

[ChuanqiXu]

Summary

Close https://github.com/llvm/llvm-project/issues/63022

This is the following of https://reviews.llvm.org/D135550, which is discussed in https://discourse.llvm.org/t/rfc-unify-memory-effect-attributes/65579. In my imagination, we could fix the issue fundamentally after we introduces new memory kind thread id. But I am not very sure if we can fix the issue fundamentally in time.

Besides that, I think the correctness is the most important. So it should not be bad to land this given it is innocent.

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

1228 ↗

(On Diff #526884)

Why is this specific to threadlocal_address? Don't you have the same problem with all readnone/readonly functions, which might be based on thread-local state internally?

[ChuanqiXu]

Address comments.

[ChuanqiXu]

llvm/lib/Transforms/Scalar/LICM.cpp

1228 ↗

(On Diff #526884)

Oh, sorry. This was an oversight.

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

1228 ↗

(On Diff #527702)

Can you please move this to the existing Behavior check below?

[ChuanqiXu]

Address comments.

[ChuanqiXu]

llvm/lib/Transforms/Scalar/LICM.cpp

1228 ↗

(On Diff #527702)

Done. Sorry for not noticing it.

[fplk]

llvm/lib/Transforms/Scalar/LICM.cpp

1240 ↗

(On Diff #528683)

I'm wondering if the new check could be unified with this check to simplify code; they both seem to check for onlyReadsMemory. I didn't double-check yet, but from the naming it seems that onlyReadsMemory and doesNotAccessMemory are mutually exclusive so we can safely move the check to below.

[nikic]

LGTM

llvm/lib/Transforms/Scalar/LICM.cpp

1231 ↗

(On Diff #528683)

We -> we

1240 ↗

(On Diff #528683)

doesNotAccessMemory() also implies onlyReadsMemory(). So I think it's okay to do the separate check that covers both branches.

llvm/test/Transforms/LICM/sink-with-coroutine.ll

55

Please precommit the new test together with a run of update_test_checks.py on the file.

[ChuanqiXu]

Thanks for reviewing!

llvm/test/Transforms/LICM/sink-with-coroutine.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; Verifies that LICM is disabled for loops that contains coro.suspend.
 ; RUN: opt -S < %s -passes=licm | FileCheck %s
 
 define i64 @licm(i64 %n) #0 {
-; CHECK-LABEL: @licm(
+; CHECK-LABEL: define i64 @licm
+; CHECK-SAME: (i64 [[N:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[P:%.*]] = alloca i64, align 8
 ; CHECK-NEXT:    br label [[BB0:%.*]]
 ; CHECK:       bb0:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[I:%.*]] = phi i64 [ 0, [[BB0]] ], [ [[T5:%.*]], [[AWAIT_READY:%.*]] ]
 ; CHECK-NEXT:    [[T5]] = add i64 [[I]], 1
 ; CHECK-NEXT:    [[SUSPEND:%.*]] = call i8 @llvm.coro.suspend(token none, i1 false)
 ; CHECK-NEXT:    switch i8 [[SUSPEND]], label [[BB2:%.*]] [
 ; CHECK-NEXT:    i8 0, label [[AWAIT_READY]]
 ; CHECK-NEXT:    ]
 ; CHECK:       await.ready:
 ; CHECK-NEXT:    store i64 1, ptr [[P]], align 4
-; CHECK-NEXT:    [[T6:%.*]] = icmp ult i64 [[T5]], [[N:%.*]]
+; CHECK-NEXT:    [[T6:%.*]] = icmp ult i64 [[T5]], [[N]]
 ; CHECK-NEXT:    br i1 [[T6]], label [[LOOP]], label [[BB2]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    [[RES:%.*]] = call i1 @llvm.coro.end(ptr null, i1 false)
 ; CHECK-NEXT:    ret i64 0
 ;
 entry:
   %p = alloca i64
   br label %bb0
await.ready:
   %t6 = icmp ult i64 %t5, %n
   br i1 %t6, label %loop, label %bb2
 
 bb2:
   %res = call i1 @llvm.coro.end(ptr null, i1 false)
   ret i64 0
 }
 
+@tls = thread_local global i32 0
+
+define i64 @hoist_threadlocal() presplitcoroutine {
+; CHECK-LABEL: define i64 @hoist_threadlocal
+; CHECK-SAME: () #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[P:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    br label [[LOOP_PREHEADER:%.*]]
+; CHECK:       loop.preheader:
+; CHECK-NEXT:    [[THREAD_LOCAL_0:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @tls)
+; CHECK-NEXT:    [[THREAD_LOCAL_1:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @tls)
+; CHECK-NEXT:    [[CMP_0:%.*]] = icmp eq ptr [[THREAD_LOCAL_0]], [[THREAD_LOCAL_1]]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I:%.*]] = phi i64 [ 0, [[LOOP_PREHEADER]] ], [ [[I_NEXT:%.*]], [[LOOP_END:%.*]] ]
+; CHECK-NEXT:    [[I_NEXT]] = add i64 [[I]], 1
+; CHECK-NEXT:    [[READONLY_0:%.*]] = call ptr @readonly_funcs()
+; CHECK-NEXT:    [[SUSPEND:%.*]] = call i8 @llvm.coro.suspend(token none, i1 false)
+; CHECK-NEXT:    switch i8 [[SUSPEND]], label [[EXIT:%.*]] [
+; CHECK-NEXT:    i8 0, label [[AWAIT_READY:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       await.ready:
+; CHECK-NEXT:    [[READONLY_1:%.*]] = call ptr @readonly_funcs()
+; CHECK-NEXT:    [[CMP_1:%.*]] = icmp eq ptr [[READONLY_0]], [[READONLY_1]]
+; CHECK-NEXT:    [[CMP:%.*]] = and i1 [[CMP_0]], [[CMP_1]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[NOT_REACHABLE:%.*]], label [[LOOP_END]]
+; CHECK:       not.reachable:
+; CHECK-NEXT:    call void @not.reachable()
+; CHECK-NEXT:    br label [[LOOP_END]]
+; CHECK:       loop.end:
+; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT]], label [[FOR_BODY]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[RES:%.*]] = call i1 @llvm.coro.end(ptr null, i1 false)
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  %p = alloca i64
+  br label %loop.preheader
+
+loop.preheader:
+  br label %for.body
+
+for.body:
+  %i = phi i64 [ 0, %loop.preheader ], [ %i.next, %loop.end ]
+  %i.next = add i64 %i, 1
+  %thread_local.0 = call ptr @llvm.threadlocal.address(ptr @tls)
+  %readonly.0 = call ptr @readonly_funcs()
+  %suspend = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %suspend, label %exit [
+  i8 0, label %await.ready
+  ]
+
+await.ready:
+  %thread_local.1 = call ptr @llvm.threadlocal.address(ptr @tls)
+  %readonly.1 = call ptr @readonly_funcs()
+  %cmp.0 = icmp eq ptr %thread_local.0, %thread_local.1
+  %cmp.1 = icmp eq ptr %readonly.0, %readonly.1
+  %cmp = and i1 %cmp.0, %cmp.1
+  br i1 %cmp, label %not.reachable, label %loop.end
+
+not.reachable:
+  call void @not.reachable()
+  br label %loop.end
+
+loop.end:
+  %loop.end.cond = icmp ugt i64 %i.next, 5
+  br i1 %cmp, label %exit, label %for.body
+
+exit:
+  %res = call i1 @llvm.coro.end(ptr null, i1 false)
+  ret i64 0
+}
+
 declare i8  @llvm.coro.suspend(token, i1)
 declare i1  @llvm.coro.end(ptr, i1)
+declare nonnull ptr @readonly_funcs() readonly
+declare nonnull ptr @llvm.threadlocal.address(ptr nonnull) nounwind readnone willreturn
+declare void @not.reachable()