[Coroutines] Don't optimize readnone function before we split coroutine (4/5)

[ChuanqiXu]

Summary

The revision tries to fix https://github.com/llvm/llvm-project/issues/47177. The key reason for this bug is similar to https://github.com/llvm/llvm-project/issues/47179. It is about the introduction of coroutine breaks the assumption that a function could only be executed in one thread. Also the fix is similar too. I add a wrapper for readnone function to block the optimizations. The main difference is that we missed optimization chances for readnone function if we enabled coroutine. And this would be addressed in following patch.

[ChuanqiXu]

Emit coro_readnone attribute in the frontend.

[rjmccall]

clang/lib/CodeGen/CGCall.cpp

2132 ↗

(On Diff #425479)

This is a global setting and will affect every function in every file that has coroutines enabled, which is presumably every file compiled with -std=c++20 or later. So you're radically changing optimization for a ton of C++ code that doesn't use coroutines.

It should be fine to just add coro_readnone instead of readnone to call sites from coroutine bodies, right? Oh, but I guess we could have a readnone call in an inlined function body. Can we just suppress this kind of code motion in coroutine bodies, or is that too invasive to the optimizer?

Actually, maybe your change to this patch was a misunderstanding. What I was trying to say in Discourse was that I wasn't sure that your early pass approach would work for thread-local variables, because addresses of thread-locals are LLVM constants and can get moved around implicitly; you need the frontend to be involved there so that there's something hooked into the right place in the emitted function. But I think it should work for this readnone annotation.

[ChuanqiXu]

Emit coro_readnone in CoroEarly pass instead of frontend.

[ChuanqiXu]

clang/lib/CodeGen/CGCall.cpp

2132 ↗

(On Diff #425479)

Actually, maybe your change to this patch was a misunderstanding. What I was trying to say in Discourse was that I wasn't sure that your early pass approach would work for thread-local variables, because addresses of thread-locals are LLVM constants and can get moved around implicitly; you need the frontend to be involved there so that there's something hooked into the right place in the emitted function. But I think it should work for this readnone annotation.

Yeah... I misunderstood your words. You were talking about TLS variable only but I thought you mentioned both.

This is a global setting and will affect every function in every file that has coroutines enabled, which is presumably every file compiled with -std=c++20 or later. So you're radically changing optimization for a ton of C++ code that doesn't use coroutines.

Yeah... it is a little better in current version since it affects codes with coroutines only. So it wouldn't affect codes compiled with -std=c++20 or later if they don't use coroutines. For optimizations, I was imaging we could re-enable the optimization by running EarlyCSE pass after CoroCleanup pass in https://reviews.llvm.org/D124364. But as you said, the optimization would be changed in the proposal.

Can we just suppress this kind of code motion in coroutine bodies, or is that too invasive to the optimizer?

Yes. According to the previous discussion (https://discourse.llvm.org/t/rfc-coroutine-and-pthread-self/56985), people don't like the direction to insert checks for coroutines around the passes. It is a burden for other developers to understand and remember they may handle a coroutine.

It should be fine to just add coro_readnone instead of readnone to call sites from coroutine bodies, right? Oh, but I guess we could have a readnone call in an inlined function body.

I am trying to do that. I think we could replace coro_readnone with readnone when we are inlining a call into an unlowered coroutine. However, I meet a technical problem that it doesn't work if we don't remove readnone for the function declaration...
I would mark this patch as plan changes before I could get a conclusion if it is possible.

---

I think we could go to https://reviews.llvm.org/D124361 first, it doesn't depend on this one and it wouldn't meet the optimization changing problem.

[ChuanqiXu]

clang/lib/CodeGen/CGCall.cpp

2132 ↗

(On Diff #425479)

replace coro_readnone with readnone

replace readnone with coro_readnone

[ChuanqiXu]

https://reviews.llvm.org/D125292

llvm/include/llvm/AsmParser/LLToken.h

enum Kind {
   kw_sanitize_address,
   kw_sanitize_hwaddress,
   kw_sanitize_memtag,
   kw_builtin,
   kw_byval,
   kw_inalloca,
   kw_cold,
   kw_convergent,
+  kw_coro_readnone,
   kw_dereferenceable,
   kw_dereferenceable_or_null,
   kw_disable_sanitizer_instrumentation,
   kw_elementtype,
   kw_inaccessiblememonly,
   kw_inaccessiblemem_or_argmemonly,
   kw_inlinehint,
   kw_inreg,

llvm/include/llvm/Bitcode/LLVMBitCodes.h

enum AttributeKindCodes {
   ATTR_KIND_VSCALE_RANGE = 74,
   ATTR_KIND_SWIFT_ASYNC = 75,
   ATTR_KIND_NO_SANITIZE_COVERAGE = 76,
   ATTR_KIND_ELEMENTTYPE = 77,
   ATTR_KIND_DISABLE_SANITIZER_INSTRUMENTATION = 78,
   ATTR_KIND_NO_SANITIZE_BOUNDS = 79,
   ATTR_KIND_ALLOC_ALIGN = 80,
   ATTR_KIND_ALLOCATED_POINTER = 81,
+  ATTR_KIND_CORO_READ_NONE = 82,
 };
 
 enum ComdatSelectionKindCodes {
   COMDAT_SELECTION_KIND_ANY = 1,
   COMDAT_SELECTION_KIND_EXACT_MATCH = 2,
   COMDAT_SELECTION_KIND_LARGEST = 3,
   COMDAT_SELECTION_KIND_NO_DUPLICATES = 4,
   COMDAT_SELECTION_KIND_SAME_SIZE = 5,

llvm/include/llvm/IR/Attributes.td

 def OptimizeNone : EnumAttr<"optnone", [FnAttr]>;
 
 /// Similar to byval but without a copy.
 def Preallocated : TypeAttr<"preallocated", [FnAttr, ParamAttr]>;
 
 /// Function does not access memory.
 def ReadNone : EnumAttr<"readnone", [FnAttr, ParamAttr]>;
 
+/// Helper attribute to block optimizations for ReadNone attribute
+/// during the lowering of coroutines to avoid misoptimizations.
+/// This should be used internally only.
+def CoroReadNone : EnumAttr<"coro_readnone", [FnAttr]>;
+
 /// Function only reads from memory.
 def ReadOnly : EnumAttr<"readonly", [FnAttr, ParamAttr]>;
 
 /// Return value is always equal to this argument.
 def Returned : EnumAttr<"returned", [ParamAttr]>;
 
 /// Parameter is required to be a trivial constant.
 def ImmArg : EnumAttr<"immarg", [ParamAttr]>;

llvm/lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(allocalign);
   KEYWORD(allocsize);
   KEYWORD(argmemonly);
   KEYWORD(builtin);
   KEYWORD(byval);
   KEYWORD(inalloca);
   KEYWORD(cold);
   KEYWORD(convergent);
+  KEYWORD(coro_readnone);
   KEYWORD(dereferenceable);
   KEYWORD(dereferenceable_or_null);
   KEYWORD(disable_sanitizer_instrumentation);
   KEYWORD(elementtype);
   KEYWORD(inaccessiblememonly);
   KEYWORD(inaccessiblemem_or_argmemonly);
   KEYWORD(inlinehint);
   KEYWORD(inreg);

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
   case bitc::ATTR_KIND_BY_VAL:
     return Attribute::ByVal;
   case bitc::ATTR_KIND_IN_ALLOCA:
     return Attribute::InAlloca;
   case bitc::ATTR_KIND_COLD:
     return Attribute::Cold;
   case bitc::ATTR_KIND_CONVERGENT:
     return Attribute::Convergent;
+  case bitc::ATTR_KIND_CORO_READ_NONE:
+    return Attribute::CoroReadNone;
   case bitc::ATTR_KIND_DISABLE_SANITIZER_INSTRUMENTATION:
     return Attribute::DisableSanitizerInstrumentation;
   case bitc::ATTR_KIND_ELEMENTTYPE:
     return Attribute::ElementType;
   case bitc::ATTR_KIND_INACCESSIBLEMEM_ONLY:
     return Attribute::InaccessibleMemOnly;
   case bitc::ATTR_KIND_INACCESSIBLEMEM_OR_ARGMEMONLY:
     return Attribute::InaccessibleMemOrArgMemOnly;

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
   case Attribute::ByVal:
     return bitc::ATTR_KIND_BY_VAL;
   case Attribute::Convergent:
     return bitc::ATTR_KIND_CONVERGENT;
   case Attribute::InAlloca:
     return bitc::ATTR_KIND_IN_ALLOCA;
   case Attribute::Cold:
     return bitc::ATTR_KIND_COLD;
+  case Attribute::CoroReadNone:
+    return bitc::ATTR_KIND_CORO_READ_NONE;
   case Attribute::DisableSanitizerInstrumentation:
     return bitc::ATTR_KIND_DISABLE_SANITIZER_INSTRUMENTATION;
   case Attribute::Hot:
     return bitc::ATTR_KIND_HOT;
   case Attribute::ElementType:
     return bitc::ATTR_KIND_ELEMENTTYPE;
   case Attribute::InaccessibleMemOnly:
     return bitc::ATTR_KIND_INACCESSIBLEMEM_ONLY;

llvm/lib/Target/DirectX/DXILWriter/DXILBitcodeWriter.cpp

uint64_t DXILBitcodeWriter::getAttrKindEncoding(Attribute::AttrKind Kind) {
   case Attribute::ByVal:
     return bitc::ATTR_KIND_BY_VAL;
   case Attribute::Convergent:
     return bitc::ATTR_KIND_CONVERGENT;
   case Attribute::InAlloca:
     return bitc::ATTR_KIND_IN_ALLOCA;
   case Attribute::Cold:
     return bitc::ATTR_KIND_COLD;
+  case Attribute::CoroReadNone:
+    return bitc::ATTR_KIND_CORO_READ_NONE;
   case Attribute::DisableSanitizerInstrumentation:
     return bitc::ATTR_KIND_DISABLE_SANITIZER_INSTRUMENTATION;
   case Attribute::Hot:
     return bitc::ATTR_KIND_HOT;
   case Attribute::ElementType:
     return bitc::ATTR_KIND_ELEMENTTYPE;
   case Attribute::InaccessibleMemOnly:
     return bitc::ATTR_KIND_INACCESSIBLEMEM_ONLY;

llvm/lib/Transforms/Coroutines/CoroCleanup.cpp

if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
         Target->replaceAllUsesWith(NewFuncPtrStruct);
         break;
       }
 
       II->eraseFromParent();
     }
   }
 
+  // Convert CoroReadNone markers to normal ReadNone attribute to enable
+  // further optimizations.
+  if (F.hasFnAttribute(Attribute::CoroReadNone)) {
+    F.removeFnAttr(Attribute::CoroReadNone);
+    F.addFnAttr(Attribute::ReadNone);
+    for (auto *U : F.users())
+      if (auto *CB = dyn_cast<CallBase>(U))
+        CB->addFnAttr(Attribute::ReadNone);
+  }
+
   // After replacement were made we can cleanup the function body a little.
   simplifyCFG(F);
 }
 
 static bool declaresCoroCleanupIntrinsics(const Module &M) {
   return coro::declaresIntrinsics(
       M, {"llvm.coro.alloc", "llvm.coro.begin", "llvm.coro.subfn.addr",
           "llvm.coro.free", "llvm.coro.id", "llvm.coro.id.retcon",

llvm/lib/Transforms/Coroutines/CoroEarly.cpp

return coro::declaresIntrinsics(
           "llvm.coro.suspend"});
 }
 
 PreservedAnalyses CoroEarlyPass::run(Module &M, ModuleAnalysisManager &) {
   if (!declaresCoroEarlyIntrinsics(M))
     return PreservedAnalyses::all();
 
   Lowerer L(M);
-  for (auto &F : M)
+  for (auto &F : M) {
     L.lowerEarlyIntrinsics(F);
 
+    // This is a workaround for the bug about pthread_self() in coroutine.
+    // See https://github.com/llvm/llvm-project/issues/47177 for the background.
+    // The reason behind the bug is that pthread_self() is marked as
+    // __attribute__((__const__)) which would be converted to `readnone` when
+    // get lowered to LLVM IR. The readnone function without parameters implies
+    // that all the calls to the function would return the same result. So the
+    // compiler would optimize the following code:
+    //
+    //   auto a = pthread_self();
+    //   co_await something();
+    //   auto b = pthread_self();
+    //
+    // to
+    //
+    //   auto a = pthread_self();
+    //   co_await something();
+    //   // replace uses of b with a
+    //
+    // The transformation is incorrect in case the coroutine might resume in
+    // another thread. The key reason for the bug is the abuse of
+    // __attribute__((__const__)) for thread identification. However, it is not
+    // easy to fix the problem in the library side (we don't know if there are
+    // other similar problems in other places) and ask the end user to update
+    // the library. So we choose to block the optimization before we split
+    // coroutine by replacing readnone attribute as a placeholder so that we
+    // could rewrite readnone later. Here we filted LLVM intrinsics since we
+    // could fix LLVM intrinsics easily in case there are similar problems.
+    if (F.hasFnAttribute(Attribute::ReadNone) && !F.isIntrinsic()) {
+      F.removeFnAttr(Attribute::ReadNone);
+      // We don't need to add coro_readnone attribute for calls since
+      // coro_readnone wouldn't do anything special.
+      for (auto *U : F.users())
+        if (auto *CB = dyn_cast<CallBase>(U))
+          CB->removeFnAttr(Attribute::ReadNone);
+      F.addFnAttr(Attribute::CoroReadNone);
+    }
+  }
+
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }

llvm/test/Transforms/Coroutines/coro-readnone-01.ll

+; Tests that the readnone function attribute could be lowered correctly by
+; CoroEarly pass.
+; RUN: opt < %s -S -passes=coro-early -opaque-pointers | FileCheck %s
+
+define ptr @f() "coroutine.presplit" {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  %j = call i32 @readnone_func() readnone
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [i8 0, label %resume
+                                         i8 1, label %cleanup]
+resume:
+  %i = call i32 @readnone_func() readnone
+  %cmp = icmp eq i32 %i, %j
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  br label %cleanup
+
+diff:
+  call void @print_diff()
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; CHECK: %j = call i32 @readnone_func()
+; CHECK: %i = call i32 @readnone_func()
+; CEHCK: declare i32 @readnone_func() #[[ATTR_NUM:[0-9]+]]
+; attributes #[[ATTR_NUM]] = { coro_readnone }
+
+declare i32 @readnone_func() readnone
+
+declare void @print_same()
+declare void @print_diff()
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare void @free(i8*)

llvm/test/Transforms/Coroutines/coro-readnone-02.ll

+; Tests that the coro_readnone function attribute could be lowered correctly by
+; CoroCleanup pass.
+; RUN: opt < %s -S -passes=coro-cleanup -opaque-pointers | FileCheck %s
+define ptr @f() #0 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, ptr null, ptr @f, ptr null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call ptr @malloc(i32 %size)
+  %hdl = call ptr @llvm.coro.begin(token %id, ptr %alloc) #5
+  %j = call i32 @readnone_func()
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [
+    i8 0, label %resume
+    i8 1, label %cleanup
+  ]
+
+resume:                                           ; preds = %entry
+  %i = call i32 @readnone_func()
+  %cmp = icmp eq i32 %i, %j
+  br i1 %cmp, label %same, label %diff
+
+same:                                             ; preds = %resume
+  call void @print_same()
+  br label %cleanup
+
+diff:                                             ; preds = %resume
+  call void @print_diff()
+  br label %cleanup
+
+cleanup:                                          ; preds = %diff, %same, %entry
+  %mem = call ptr @llvm.coro.free(token %id, ptr %hdl)
+  call void @free(ptr %mem)
+  br label %suspend
+
+suspend:                                          ; preds = %cleanup, %entry
+  %0 = call i1 @llvm.coro.end(ptr %hdl, i1 false) #5
+  ret ptr %hdl
+}
+
+; CHECK: %j = call i32 @readnone_func() #[[ATTR_NUM:[0-9]+]]
+; CHECK: %i = call i32 @readnone_func() #[[ATTR_NUM]]
+; CEHCK: declare i32 @readnone_func() #[[ATTR_NUM]]
+; attributes #[[ATTR_NUM]] = { coro_readnone }
+
+declare i32 @readnone_func() coro_readnone
+declare void @print_same()
+declare void @print_diff()
+
+declare ptr @llvm.coro.free(token, ptr nocapture readonly)
+declare i32 @llvm.coro.size.i32()
+declare i8 @llvm.coro.suspend(token, i1)
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
+declare i1 @llvm.coro.alloc(token)
+declare ptr @llvm.coro.begin(token, ptr writeonly)
+declare i1 @llvm.coro.end(ptr, i1)
+declare noalias ptr @malloc(i32)
+declare void @free(ptr)

llvm/test/Transforms/Coroutines/coro-readnone-03.ll

+; Tests that the readnone function which cross suspend points wouldn't misoptimized
+; RUN: opt < %s -S -passes='default<O3>' -opaque-pointers | FileCheck %s
+
+define ptr @f() "coroutine.presplit" {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %size = call i32 @llvm.coro.size.i32()
+  %alloc = call i8* @malloc(i32 %size)
+  %hdl = call i8* @llvm.coro.begin(token %id, i8* %alloc)
+  %j = call i32 @readnone_func() readnone
+  %sus_result = call i8 @llvm.coro.suspend(token none, i1 false)
+  switch i8 %sus_result, label %suspend [i8 0, label %resume
+                                         i8 1, label %cleanup]
+resume:
+  %i = call i32 @readnone_func() readnone
+  %cmp = icmp eq i32 %i, %j
+  br i1 %cmp, label %same, label %diff
+
+same:
+  call void @print_same()
+  br label %cleanup
+
+diff:
+  call void @print_diff()
+  br label %cleanup
+
+cleanup:
+  %mem = call i8* @llvm.coro.free(token %id, i8* %hdl)
+  call void @free(i8* %mem)
+  br label %suspend
+
+suspend:
+  call i1 @llvm.coro.end(i8* %hdl, i1 0)
+  ret i8* %hdl
+}
+
+; CHECK-LABEL: f.resume(
+; CHECK: br i1 %cmp, label %same, label %diff
+; CHECK-EMPTY:
+; CHECK-NEXT: same:
+; CHECK-NEXT:   call void @print_same()
+; CHECK-NEXT:   br label
+; CHECK-EMPTY:
+; CHECK-NEXT: diff:
+; CHECK-NEXT:   call void @print_diff()
+; CHECK-NEXT:   br label
+
+declare i32 @readnone_func() readnone
+
+declare void @print_same()
+declare void @print_diff()
+declare i8* @llvm.coro.free(token, i8*)
+declare i32 @llvm.coro.size.i32()
+declare i8  @llvm.coro.suspend(token, i1)
+
+declare token @llvm.coro.id(i32, i8*, i8*, i8*)
+declare i1 @llvm.coro.alloc(token)
+declare i8* @llvm.coro.begin(token, i8*)
+declare i1 @llvm.coro.end(i8*, i1)
+
+declare noalias i8* @malloc(i32)
+declare void @free(i8*)