pull the ballot argument to all the match users

Author: PankajDwivedi-25

llvm/lib/Target/AMDGPU/AMDGPUUniformIntrinsicCombine.cpp

@@ -96,12 +96,12 @@ AMDGPUUniformIntrinsicCombinePass::run(Function &F,
 bool AMDGPUUniformIntrinsicCombineImpl::run(Function &F) {
   bool IsChanged{false};
   Module *M = F.getParent();
-
+  
   // If none of the relevant intrinsics are declared, return early.
-  // if (!M->getFunction(Intrinsic::getName(Intrinsic::amdgcn_permlane64)) &&
-  //     !M->getFunction(Intrinsic::getName(Intrinsic::amdgcn_readfirstlane)) &&
-  //     !M->getFunction(Intrinsic::getName(Intrinsic::amdgcn_readlane)) &&
-  //     !M->getFunction(Intrinsic::getName(Intrinsic::amdgcn_ballot))) {
+  // if (!Intrinsic::getDeclarationIfExists(M, Intrinsic::amdgcn_permlane64, {}) &&
+  //     !Intrinsic::getDeclarationIfExists(M, Intrinsic::amdgcn_readfirstlane, {}) &&
+  //     !Intrinsic::getDeclarationIfExists(M, Intrinsic::amdgcn_readlane, {}) &&
+  //     !Intrinsic::getDeclarationIfExists(M, Intrinsic::amdgcn_ballot, {})) {
   //   return false;
   // }
 
@@ -139,28 +139,21 @@ bool AMDGPUUniformIntrinsicCombineImpl::optimizeUniformIntrinsicInst(
     LLVM_DEBUG(dbgs() << "Found uniform ballot intrinsic: " << II << "\n");
 
     // Look for a direct `icmp eq` use of the ballot result.
-    // FIXME: replace all the uses?
-    auto It = llvm::find_if(II.users(), [&](User *U) {
-      return match(U, m_ICmp(m_Specific(&II), m_Zero()));
-    });
-
-    // Check if a match was found
-    if (It == II.user_end())
-      return false;
-
-    // Extract the matching `icmp` instruction
-    ICmpInst *ICmp = dyn_cast<ICmpInst>(*It);
-    IRBuilder<> Builder(ICmp);
-
-    // Convert ballot argument to match `icmp` operand type (i64)
-    Value *ConvertedSrc = Builder.CreateZExtOrTrunc(Src, II.getType());
-
-    LLVM_DEBUG(dbgs() << "Replacing ballot result in icmp: " << *ICmp
-                      << " with " << *ConvertedSrc << "\n");
-
-    // Replace `%ballot` in `icmp` with `ConvertedSrc`
-    ICmp->setOperand(0, ConvertedSrc);
-    return true;
+    bool Changed = false;
+    for (User *U : make_early_inc_range(II.users())) {
+      if (match(U, m_ICmp(m_Specific(&II), m_Zero()))) {
+        ICmpInst *ICmp = dyn_cast<ICmpInst>(U);
+        IRBuilder<> Builder(ICmp);
+        Value *ConvertedSrc = Builder.CreateZExtOrTrunc(Src, II.getType());
+
+        LLVM_DEBUG(dbgs() << "Replacing ballot result in icmp: " << *ICmp
+                          << " with " << *ConvertedSrc << "\n");
+
+        ICmp->setOperand(0, ConvertedSrc);
+        Changed = true;
+      }
+    }
+    return Changed;
   }
   }
   return false;

llvm/test/CodeGen/AMDGPU/amdgpu-simplify-uniform-waterfall.ll

@@ -122,6 +122,53 @@ exit:
   ret void
 }
 
+define protected amdgpu_kernel void @trivial_waterfall_multiple_icmp(ptr addrspace(1) %out) {
+; PASS-CHECK-LABEL: define protected amdgpu_kernel void @trivial_waterfall_multiple_icmp(
+; PASS-CHECK-SAME: ptr addrspace(1) [[OUT:%.*]]) #[[ATTR0]] {
+; PASS-CHECK-NEXT:  [[ENTRY:.*]]:
+; PASS-CHECK-NEXT:    br label %[[WHILE:.*]]
+; PASS-CHECK:       [[WHILE]]:
+; PASS-CHECK-NEXT:    [[DONE:%.*]] = phi i1 [ false, %[[ENTRY]] ], [ true, %[[IF:.*]] ]
+; PASS-CHECK-NEXT:    [[NOT_DONE:%.*]] = xor i1 [[DONE]], true
+; PASS-CHECK-NEXT:    [[BALLOT:%.*]] = tail call i64 @llvm.amdgcn.ballot.i64(i1 [[NOT_DONE]])
+; PASS-CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[NOT_DONE]] to i64
+; PASS-CHECK-NEXT:    [[IS_DONE_1:%.*]] = icmp eq i64 [[TMP1]], 0
+; PASS-CHECK-NEXT:    [[TMP0:%.*]] = zext i1 [[NOT_DONE]] to i64
+; PASS-CHECK-NEXT:    [[IS_DONE_3:%.*]] = icmp eq i64 [[TMP0]], 0
+; PASS-CHECK-NEXT:    br i1 [[IS_DONE_1]], label %[[EXIT:.*]], label %[[IF]]
+; PASS-CHECK:       [[IF]]:
+; PASS-CHECK-NEXT:    store i32 5, ptr addrspace(1) [[OUT]], align 4
+; PASS-CHECK-NEXT:    [[TMP2:%.*]] = zext i1 [[NOT_DONE]] to i64
+; PASS-CHECK-NEXT:    [[IS_DONE_4:%.*]] = icmp eq i64 [[TMP2]], 0
+; PASS-CHECK-NEXT:    br label %[[WHILE]]
+; PASS-CHECK:       [[EXIT]]:
+; PASS-CHECK-NEXT:    ret void
+;
+; DCE-CHECK-LABEL: define protected amdgpu_kernel void @trivial_waterfall_multiple_icmp(
+; DCE-CHECK-SAME: ptr addrspace(1) [[OUT:%.*]]) #[[ATTR0]] {
+; DCE-CHECK-NEXT:  [[ENTRY:.*:]]
+; DCE-CHECK-NEXT:    store i32 5, ptr addrspace(1) [[OUT]], align 4
+; DCE-CHECK-NEXT:    ret void
+;
+entry:
+  br label %while
+
+while:
+  %done = phi i1 [ 0, %entry ], [ 1, %if ]
+  %not_done = xor i1 %done, true
+  %ballot = tail call i64 @llvm.amdgcn.ballot.i64(i1 %not_done)
+  %is_done_1 = icmp eq i64 %ballot, 0
+  %is_done_2 = icmp eq i64 %ballot, 0
+  br i1 %is_done_1, label %exit, label %if
+
+if:
+  store i32 5, ptr addrspace(1) %out
+  %is_done_3 = icmp eq i64 %ballot, 0
+  br label %while
+
+exit:
+  ret void
+}
 
 declare i64 @llvm.amdgcn.ballot.i64(i1) #1
 !6 = !{i64 690}