Address feedback on the location of the opt

TejaX-Alaghari · TejaX-Alaghari · commit 08355f9a05de · 2025-10-08T16:19:13.000+05:30
- Remove redundant const propagration (assume equality opt) from InstCombine.
- Moved assume(ballot(cmp) == -1) optimization from InstCombine to GVN.
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
@@ -1322,7 +1322,12 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
     if (isa<PoisonValue>(Arg))
       return IC.replaceInstUsesWith(II, PoisonValue::get(II.getType()));
 
-    // For Wave32 targets, convert i64 ballot to i32 ballot + zext
+    if (auto *Src = dyn_cast<ConstantInt>(Arg)) {
+      if (Src->isZero()) {
+        // amdgcn.ballot(i1 0) is zero.
+        return IC.replaceInstUsesWith(II, Constant::getNullValue(II.getType()));
+      }
+    }
     if (ST->isWave32() && II.getType()->getIntegerBitWidth() == 64) {
       // %b64 = call i64 ballot.i64(...)
       // =>
@@ -1336,15 +1341,6 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
       Call->takeName(&II);
       return IC.replaceInstUsesWith(II, Call);
     }
-
-    if (auto *Src = dyn_cast<ConstantInt>(Arg)) {
-      if (Src->isZero()) {
-        // amdgcn.ballot(i1 0) is zero.
-        return IC.replaceInstUsesWith(II, Constant::getNullValue(II.getType()));
-      }
-      // Note: ballot(true) is NOT constant folded because the result depends
-      // on the active lanes in the wavefront, not just the condition value.
-    }
     break;
   }
   case Intrinsic::amdgcn_wavefrontsize: {
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3540,81 +3540,6 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       }
     }
 
-    // Basic assume equality optimization: assume(x == c) -> replace uses of x
-    // with c
-    if (auto *ICmp = dyn_cast<ICmpInst>(IIOperand)) {
-      if (ICmp->getPredicate() == ICmpInst::ICMP_EQ) {
-        Value *LHS = ICmp->getOperand(0);
-        Value *RHS = ICmp->getOperand(1);
-        Value *Variable = nullptr;
-        Constant *ConstantVal = nullptr;
-
-        if (auto *C = dyn_cast<Constant>(RHS)) {
-          Variable = LHS;
-          ConstantVal = C;
-        } else if (auto *C = dyn_cast<Constant>(LHS)) {
-          Variable = RHS;
-          ConstantVal = C;
-        }
-
-        if (Variable && ConstantVal && Variable->hasUseList()) {
-          SmallVector<Use *, 8> Uses;
-          for (Use &U : Variable->uses()) {
-            if (auto *UseInst = dyn_cast<Instruction>(U.getUser())) {
-              if (UseInst != ICmp &&
-                  isValidAssumeForContext(II, UseInst, &DT)) {
-                Uses.push_back(&U);
-              }
-            }
-          }
-
-          for (Use *U : Uses) {
-            U->set(ConstantVal);
-            Worklist.pushValue(U->getUser());
-          }
-
-          if (!Uses.empty()) {
-            Worklist.pushValue(Variable);
-          }
-        }
-      }
-    }
-
-    // Optimize AMDGPU ballot patterns in assumes:
-    // assume(ballot(cmp) == -1) means cmp is true on all active lanes
-    // We can replace uses of cmp with true
-    Value *BallotInst;
-    if (match(IIOperand, m_SpecificICmp(ICmpInst::ICMP_EQ, m_Value(BallotInst),
-                                        m_AllOnes()))) {
-      if (auto *IntrCall = dyn_cast<IntrinsicInst>(BallotInst)) {
-        if (IntrCall->getIntrinsicID() == Intrinsic::amdgcn_ballot) {
-          Value *BallotArg = IntrCall->getArgOperand(0);
-          if (BallotArg->getType()->isIntegerTy(1) && BallotArg->hasUseList()) {
-            // Find uses and replace with true
-            SmallVector<Use *, 8> Uses;
-            for (Use &U : BallotArg->uses()) {
-              if (auto *UseInst = dyn_cast<Instruction>(U.getUser())) {
-                if (UseInst != IntrCall &&
-                    isValidAssumeForContext(II, UseInst, &DT)) {
-                  Uses.push_back(&U);
-                }
-              }
-            }
-
-            // Replace uses with true
-            for (Use *U : Uses) {
-              U->set(ConstantInt::getTrue(BallotArg->getType()));
-              Worklist.pushValue(U->getUser());
-            }
-
-            if (!Uses.empty()) {
-              Worklist.pushValue(BallotArg);
-            }
-          }
-        }
-      }
-    }
-
     // If there is a dominating assume with the same condition as this one,
     // then this one is redundant, and should be removed.
     KnownBits Known(1);
diff --git a/llvm/lib/Transforms/Scalar/GVN.cpp b/llvm/lib/Transforms/Scalar/GVN.cpp
@@ -54,6 +54,7 @@
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
@@ -2206,6 +2207,23 @@ bool GVNPass::processAssumeIntrinsic(AssumeInst *IntrinsicI) {
           std::swap(CmpLHS, CmpRHS);
       }
 
+      // Optimize AMDGPU ballot pattern: assume(ballot(cmp) == -1) or
+      // assume(ballot(cmp) == exec_mask). This implies cmp is true on all
+      // active lanes and hence can be replaced with true.
+      if (isa<IntrinsicInst>(CmpLHS) && isa<Constant>(CmpRHS)) {
+        auto *IntrCall = cast<IntrinsicInst>(CmpLHS);
+        // Check if CmpLHS is a ballot intrinsic
+        if (IntrCall->getIntrinsicID() ==
+            Intrinsic::AMDGCNIntrinsics::amdgcn_ballot) {
+          Value *BallotArg = IntrCall->getArgOperand(0);
+          if (BallotArg->getType()->isIntegerTy(1) &&
+              (match(CmpRHS, m_AllOnes()) || !isa<Constant>(CmpRHS))) {
+            CmpLHS = BallotArg;
+            CmpRHS = ConstantInt::getTrue(BallotArg->getType());
+          }
+        }
+      }
+
       // Handle degenerate case where we either haven't pruned a dead path or a
       // removed a trivial assume yet.
       if (isa<Constant>(CmpLHS) && isa<Constant>(CmpRHS))
diff --git a/llvm/test/Transforms/GVN/assume-equal.ll b/llvm/test/Transforms/GVN/assume-equal.ll
@@ -387,6 +387,35 @@ define i8 @assume_ptr_eq_same_prov(ptr %p, i64 %x) {
   ret i8 %v
 }
 
+; Test AMDGPU ballot pattern optimization
+; assume(ballot(cmp) == -1) means cmp is true on all active lanes
+; so uses of cmp can be replaced with true
+define void @assume_ballot_uniform(i32 %x) {
+; CHECK-LABEL: @assume_ballot_uniform(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[BALLOT:%.*]] = call i64 @llvm.amdgcn.ballot.i64(i1 [[CMP]])
+; CHECK-NEXT:    [[ALL:%.*]] = icmp eq i64 [[BALLOT]], -1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[ALL]])
+; CHECK-NEXT:    br i1 true, label [[FOO:%.*]], label [[BAR:%.*]]
+; CHECK:       foo:
+; CHECK-NEXT:    ret void
+; CHECK:       bar:
+; CHECK-NEXT:    ret void
+;
+  %cmp = icmp eq i32 %x, 0
+  %ballot = call i64 @llvm.amdgcn.ballot.i64(i1 %cmp)
+  %all = icmp eq i64 %ballot, -1
+  call void @llvm.assume(i1 %all)
+  br i1 %cmp, label %foo, label %bar
+
+foo:
+  ret void
+
+bar:
+  ret void
+}
+
+declare i64 @llvm.amdgcn.ballot.i64(i1)
 declare noalias ptr @_Znwm(i64)
 declare void @_ZN1AC1Ev(ptr)
 declare void @llvm.assume(i1)
diff --git a/llvm/test/Transforms/InstCombine/amdgpu-ballot-constant-fold.ll b/llvm/test/Transforms/InstCombine/amdgpu-ballot-constant-fold.ll
diff --git a/llvm/test/Transforms/InstCombine/assume.ll b/llvm/test/Transforms/InstCombine/assume.ll
@@ -82,7 +82,7 @@ define i32 @simple(i32 %a) #1 {
 ; CHECK-LABEL: @simple(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], 4
 ; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
-; CHECK-NEXT:    ret i32 4
+; CHECK-NEXT:    ret i32 [[A]]
 ;
   %cmp = icmp eq i32 %a, 4
   tail call void @llvm.assume(i1 %cmp)
@@ -1034,35 +1034,6 @@ define i1 @neg_assume_trunc_eq_one(i8 %x) {
   ret i1 %q
 }
 
-; Test AMDGPU ballot pattern optimization
-; assume(ballot(cmp) == -1) means cmp is true on all active lanes
-; so uses of cmp can be replaced with true
-define void @assume_ballot_uniform(i32 %x) {
-; CHECK-LABEL: @assume_ballot_uniform(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[BALLOT:%.*]] = call i64 @llvm.amdgcn.ballot.i64(i1 [[CMP]])
-; CHECK-NEXT:    [[ALL:%.*]] = icmp eq i64 [[BALLOT]], -1
-; CHECK-NEXT:    call void @llvm.assume(i1 [[ALL]])
-; CHECK-NEXT:    br i1 true, label [[FOO:%.*]], label [[BAR:%.*]]
-; CHECK:       foo:
-; CHECK-NEXT:    ret void
-; CHECK:       bar:
-; CHECK-NEXT:    ret void
-;
-  %cmp = icmp eq i32 %x, 0
-  %ballot = call i64 @llvm.amdgcn.ballot.i64(i1 %cmp)
-  %all = icmp eq i64 %ballot, -1
-  call void @llvm.assume(i1 %all)
-  br i1 %cmp, label %foo, label %bar
-
-foo:
-  ret void
-
-bar:
-  ret void
-}
-
-declare i64 @llvm.amdgcn.ballot.i64(i1)
 declare void @use(i1)
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
