[AMDGPU] Don't unify divergent exit nodes with musttail calls

[shiltian]

Fixes SWDEV-512254.

[shiltian]

• [AMDGPU] Don't unify divergent exit nodes with musttail calls #126395 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: Shilei Tian (shiltian)

Changes

Fixes SWDEV-512254.

---

Full diff: https://github.com/llvm/llvm-project/pull/126395.diff

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp (+4-1)
• (added) llvm/test/CodeGen/AMDGPU/do-not-unify-divergent-exit-nodes-with-musttail.ll (+25)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp b/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp
index fda2a38c2464e00..d087fbc86545c99 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUUnifyDivergentExitNodes.cpp
@@ -215,7 +215,10 @@ bool AMDGPUUnifyDivergentExitNodesImpl::run(Function &F, DominatorTree *DT,
       PDT.roots(), [&](auto BB) { return !isUniformlyReached(UA, *BB); });
 
   for (BasicBlock *BB : PDT.roots()) {
-    if (isa<ReturnInst>(BB->getTerminator())) {
+    if (auto *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {
+      auto *CI = dyn_cast_or_null<CallInst>(RI->getPrevNode());
+      if (CI && CI->isMustTailCall())
+        continue;
       if (HasDivergentExitBlock)
         ReturningBlocks.push_back(BB);
     } else if (isa<UnreachableInst>(BB->getTerminator())) {
diff --git a/llvm/test/CodeGen/AMDGPU/do-not-unify-divergent-exit-nodes-with-musttail.ll b/llvm/test/CodeGen/AMDGPU/do-not-unify-divergent-exit-nodes-with-musttail.ll
new file mode 100644
index 000000000000000..2f367d64351089c
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/do-not-unify-divergent-exit-nodes-with-musttail.ll
@@ -0,0 +1,25 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=amdgpu-unify-divergent-exit-nodes -mtriple=amdgcn-amd-amdhsa -mcpu=gfx90a %s -o - | FileCheck %s
+
+define void @spill_sgpr_with_tail_call() {
+; CHECK-LABEL: define void @spill_sgpr_with_tail_call(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[L1:%.*]] = load i1, ptr null, align 1
+; CHECK-NEXT:    br i1 [[L1]], label %[[SW_C:.*]], label %[[SW_D:.*]]
+; CHECK:       [[SW_D]]:
+; CHECK-NEXT:    musttail call void null()
+; CHECK-NEXT:    ret void
+; CHECK:       [[SW_C]]:
+; CHECK-NEXT:    ret void
+;
+  %L1 = load i1, ptr null, align 1
+  br i1 %L1, label %SW_C, label %SW_D
+
+SW_D:
+  musttail call void null()
+  ret void
+
+SW_C:
+  ret void
+}
+

[arsenm]

We have an issue if we have a musttail indirect call with a divergent call target