Address feedback:1. Add the condition to make sure ballot width matches wave sixe

2. Modify logic to make the InferredCondValue non-optional
3. Simplified cmp arg match logic
4. Added missing braces and comments indents
5. Removed O2 checks from the test cases

Author: TejaX-Alaghari

llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp

@@ -1353,9 +1353,10 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
     if (!isa<Instruction>(Arg))
       break;
 
-    // Skip if ballot width < wave size (e.g., ballot.i32 on wave64).
-    if (ST->isWave64() && II.getType()->getIntegerBitWidth() == 32)
-      break; // ballot.i32 on wave64 captures only lanes [0:31]
+    // Skip if ballot width doesn't match wave size.
+    unsigned WavefrontSize = ST->getWavefrontSize();
+    if (WavefrontSize != II.getType()->getIntegerBitWidth())
+      break;
 
     for (auto &AssumeVH : IC.getAssumptionCache().assumptionsFor(&II)) {
       if (!AssumeVH)
@@ -1364,37 +1365,37 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
       auto *Assume = cast<AssumeInst>(AssumeVH);
       Value *Cond = Assume->getArgOperand(0);
 
-      // Check if assume condition is an equality comparison.
-      auto *ICI = dyn_cast<ICmpInst>(Cond);
+      // Pattern match: assume(icmp eq ballot, CompareValue)
+      ICmpInst *ICI = dyn_cast<ICmpInst>(Cond);
       if (!ICI || ICI->getPredicate() != ICmpInst::ICMP_EQ)
         continue;
 
-      // Extract the ballot and the value being compared against it.
-      Value *LHS = ICI->getOperand(0), *RHS = ICI->getOperand(1);
-      Value *CompareValue = (LHS == &II) ? RHS : (RHS == &II) ? LHS : nullptr;
-      if (!CompareValue)
+      Value *CompareValue;
+      if (!match(ICI, m_c_ICmp(m_Specific(&II), m_Value(CompareValue))))
         continue;
 
-      // Determine the constant value of the ballot's condition argument.
-      std::optional<bool> InferredCondValue;
+      // Determine the inferred value of the ballot's condition argument.
+      bool InferredCondValue;
       if (auto *CI = dyn_cast<ConstantInt>(CompareValue)) {
-        // ballot(x) == -1 means all lanes have x = true.
-        if (CI->isMinusOne())
+        if (CI->isMinusOne()) {
+          // ballot(x) == -1 means all lanes have x = true.
           InferredCondValue = true;
-        // ballot(x) == 0 means all lanes have x = false.
-        else if (CI->isZero())
+        } else if (CI->isZero()) {
+          // ballot(x) == 0 means all lanes have x = false.
           InferredCondValue = false;
+        } else {
+          continue;
+        }
       } else if (match(CompareValue,
                        m_Intrinsic<Intrinsic::amdgcn_ballot>(m_One()))) {
         // ballot(x) == ballot(true) means x = true (EXEC mask comparison).
         InferredCondValue = true;
-      }
-
-      if (!InferredCondValue)
+      } else {
         continue;
+      }
 
       Constant *ReplacementValue =
-          ConstantInt::getBool(Arg->getContext(), *InferredCondValue);
+          ConstantInt::getBool(Arg->getContext(), InferredCondValue);
 
       // Replace uses of the condition argument dominated by the assume.
       bool Changed = false;

llvm/test/Transforms/InstCombine/AMDGPU/llvm.amdgcn.ballot-assume-wave32.ll

@@ -1,6 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -mtriple=amdgcn-amd-amdhsa -mattr=+wavefrontsize32 -passes=instcombine -S | FileCheck %s
-; RUN: opt < %s -mtriple=amdgcn-amd-amdhsa -mattr=+wavefrontsize32 -O2 -S | FileCheck %s --check-prefix=O2
 ;
 ; Wave32-specific tests for ballot-assume optimizations.
 ; - ballot.i32 should optimize (captures all 32 lanes)
@@ -27,9 +26,6 @@ define amdgpu_kernel void @wave32_ballot_i32_all_lanes(i32 %x, ptr addrspace(1)
 ; CHECK:       bar:
 ; CHECK-NEXT:    ret void
 ;
-; O2-LABEL: @wave32_ballot_i32_all_lanes(
-; O2-NEXT:  common.ret:
-; O2-NEXT:    ret void
 ;
   %cmp = icmp eq i32 %x, 0
   %ballot = call i32 @llvm.amdgcn.ballot.i32(i1 %cmp)

llvm/test/Transforms/InstCombine/AMDGPU/llvm.amdgcn.ballot-assume-wave64.ll

@@ -1,6 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -mtriple=amdgcn-amd-amdhsa -mattr=+wavefrontsize64 -passes=instcombine -S | FileCheck %s
-; RUN: opt < %s -mtriple=amdgcn-amd-amdhsa -mattr=+wavefrontsize64 -O2 -S | FileCheck %s --check-prefix=O2
 ;
 ; Wave64-specific tests for ballot-assume optimizations.
 ; - ballot.i64 should optimize (captures all 64 lanes)
@@ -27,9 +26,6 @@ define amdgpu_kernel void @wave64_ballot_i64_all_lanes(i32 %x, ptr addrspace(1)
 ; CHECK:       bar:
 ; CHECK-NEXT:    ret void
 ;
-; O2-LABEL: @wave64_ballot_i64_all_lanes(
-; O2-NEXT:  common.ret:
-; O2-NEXT:    ret void
 ;
   %cmp = icmp eq i32 %x, 0
   %ballot = call i64 @llvm.amdgcn.ballot.i64(i1 %cmp)