[InstCombine] Implement generic assume-based uniformity optimization

TejaX-Alaghari · TejaX-Alaghari · commit b692ae947711 · 2025-10-05T17:21:20.000+05:30
Implement a comprehensive generic optimization for assume intrinsics that extracts
uniformity information and optimizes dominated uses. The optimization recognizes
multiple patterns that establish value uniformity and replaces dominated uses with
uniform constants.

Addresses uniformity analysis optimization opportunities identified in
AMDGPU ballot/readfirstlane + assume patterns for improved code generation
through constant propagation.
diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md
@@ -1,74 +1,4 @@
-# LLVM Project AI Coding Agent Instructions
-
-## Architecture Overview
-
-LLVM is a compiler infrastructure with modular components:
-- **Core LLVM** (`llvm/`): IR processing, optimizations, code generation
-- **Clang** (`clang/`): C/C++/Objective-C frontend 
-- **LLD** (`lld/`): Linker
-- **libc++** (`libcxx/`): C++ standard library
-- **Target backends** (`llvm/lib/Target/{AMDGPU,X86,ARM,...}/`): Architecture-specific code generation
-
-## Essential Development Workflows
-
-### Build System (CMake + Ninja)
-```bash
-# Configure with common options for development
-cmake -G Ninja -S llvm-project/llvm -B build \
-  -DCMAKE_BUILD_TYPE=RelWithDebInfo \
-  -DLLVM_ENABLE_PROJECTS="clang;lld" \
-  -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" \
-  -DLLVM_ENABLE_ASSERTIONS=ON
-
-# Build and install
-cmake --build build
-cmake --install build --prefix install/
-```
-
-### Testing with LIT
-- Use `opt < file.ll -passes=instcombine -S | FileCheck %s` pattern for IR transforms
-- Test files go in `llvm/test/Transforms/{PassName}/` with `.ll` extension
-- Always include both positive and negative test cases
-- Use `CHECK-LABEL:` for function boundaries, `CHECK-NEXT:` for strict sequence
-
-### Key Patterns for Transforms
-
-**InstCombine Pattern** (`llvm/lib/Transforms/InstCombine/`):
-- Implement in `InstCombine*.cpp` using visitor pattern (`visitCallInst`, `visitBinaryOperator`)
-- Use `PatternMatch.h` matchers: `match(V, m_Add(m_Value(X), m_ConstantInt()))`
-- Return `nullptr` for no change, modified instruction, or replacement
-- Add to worklist with `Worklist.pushValue()` for dependent values
-
-**Target-Specific Intrinsics**:
-- AMDGPU: `@llvm.amdgcn.*` intrinsics in `llvm/include/llvm/IR/IntrinsicsAMDGPU.td`
-- Pattern: `if (II->getIntrinsicID() == Intrinsic::amdgcn_ballot)`
-
-## Code Quality Standards
-
-### Control Flow & Debug Info
-When modifying control flow, ensure changes don't corrupt:
-- Performance profiling data (branch weights, call counts)
-- Debug information for branches and calls
-- Exception handling unwind information
-
-### Target-Specific Considerations
-- **AMDGPU**: Wavefront uniformity analysis affects ballot intrinsics
-- **X86**: Vector width and ISA feature dependencies
-- Use `TargetTransformInfo` for cost models and capability queries
-
-### Testing Requirements
-- Every optimization needs regression tests showing before/after IR
-- Include edge cases: constants, undef, poison values
-- Test target-specific intrinsics with appropriate triple
-- Use `; RUN: opt < %s -passes=... -S | FileCheck %s` format
-
-## Common Development Pitfalls
-- Don't assume instruction operand order without checking `isCommutative()`
-- Verify type compatibility before creating new instructions
-- Consider poison/undef propagation in optimizations
-- Check for side effects before eliminating instructions
-
-## Pass Pipeline Context
-- InstCombine runs early and multiple times in the pipeline
-- Subsequent passes like SimplifyCFG will clean up control flow
-- Use `replaceAllUsesWith()` carefully to maintain SSA form
+When performing a code review, pay close attention to code modifying a function's
+control flow. Could the change result in the corruption of performance profile
+data? Could the change result in invalid debug information, in particular for
+branches and calls?
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
@@ -1322,12 +1322,7 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
     if (isa<PoisonValue>(Arg))
       return IC.replaceInstUsesWith(II, PoisonValue::get(II.getType()));
 
-    if (auto *Src = dyn_cast<ConstantInt>(Arg)) {
-      if (Src->isZero()) {
-        // amdgcn.ballot(i1 0) is zero.
-        return IC.replaceInstUsesWith(II, Constant::getNullValue(II.getType()));
-      }
-    }
+    // For Wave32 targets, convert i64 ballot to i32 ballot + zext
     if (ST->isWave32() && II.getType()->getIntegerBitWidth() == 64) {
       // %b64 = call i64 ballot.i64(...)
       // =>
@@ -1341,6 +1336,15 @@ 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
@@ -85,8 +85,6 @@ using namespace PatternMatch;
 
 STATISTIC(NumSimplified, "Number of library calls simplified");
 
-
-
 static cl::opt<unsigned> GuardWideningWindow(
     "instcombine-guard-widening-window",
     cl::init(3),
@@ -2989,20 +2987,6 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     }
     break;
   }
-  case Intrinsic::amdgcn_ballot: {
-    // Optimize ballot intrinsics when the condition is known to be uniform
-    Value *Condition = II->getArgOperand(0);
-    
-    // If the condition is a constant, we can evaluate the ballot directly
-    if (auto *ConstCond = dyn_cast<ConstantInt>(Condition)) {
-      // ballot(true) -> -1 (all lanes active)
-      // ballot(false) -> 0 (no lanes active)
-      uint64_t Result = ConstCond->isOne() ? ~0ULL : 0ULL;
-      return replaceInstUsesWith(*II, ConstantInt::get(II->getType(), Result));
-    }
-    
-    break;
-  }
   case Intrinsic::ldexp: {
     // ldexp(ldexp(x, a), b) -> ldexp(x, a + b)
     //
@@ -3556,8 +3540,6 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       }
     }
 
-
-
     // 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);
@@ -3571,7 +3553,9 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       return eraseInstFromFunction(*II);
     }
 
-
+    // Try to extract uniformity information from the assume and optimize
+    // dominated uses of any variables that are established as uniform.
+    optimizeAssumedUniformValues(cast<AssumeInst>(II));
 
     // Update the cache of affected values for this assumption (we might be
     // here because we just simplified the condition).
@@ -5026,3 +5010,117 @@ InstCombinerImpl::transformCallThroughTrampoline(CallBase &Call,
   Call.setCalledFunction(FTy, NestF);
   return &Call;
 }
+
+/// Extract uniformity information from assume and optimize dominated uses.
+/// This works with any assume pattern that establishes value uniformity.
+void InstCombinerImpl::optimizeAssumedUniformValues(AssumeInst *Assume) {
+  Value *AssumedCondition = Assume->getArgOperand(0);
+  
+  // Map of uniform values to their uniform constants
+  SmallDenseMap<Value *, Constant *> UniformValues;
+  
+  // Pattern 1: assume(icmp eq (X, C)) -> X is uniform and equals C
+  if (auto *ICmp = dyn_cast<ICmpInst>(AssumedCondition)) {
+    if (ICmp->getPredicate() == ICmpInst::ICMP_EQ) {
+      Value *LHS = ICmp->getOperand(0);
+      Value *RHS = ICmp->getOperand(1);
+      
+      // X == constant -> X is uniform and equals constant
+      if (auto *C = dyn_cast<Constant>(RHS)) {
+        UniformValues[LHS] = C;
+      } else if (auto *C = dyn_cast<Constant>(LHS)) {
+        UniformValues[RHS] = C;
+      }
+      
+      // Handle intrinsic patterns in equality comparisons
+      // Pattern: assume(ballot(cmp) == -1) -> cmp is uniform and true
+      if (auto *IntrinsicCall = dyn_cast<IntrinsicInst>(LHS)) {
+        if (IntrinsicCall->getIntrinsicID() == Intrinsic::amdgcn_ballot) {
+          if (match(RHS, m_AllOnes())) {
+            Value *BallotArg = IntrinsicCall->getArgOperand(0);
+            if (BallotArg->getType()->isIntegerTy(1)) {
+              UniformValues[BallotArg] = ConstantInt::getTrue(BallotArg->getType());
+              
+              // Special case: if BallotArg is an equality comparison, 
+              // we know the operands are equal
+              if (auto *CmpInst = dyn_cast<ICmpInst>(BallotArg)) {
+                if (CmpInst->getPredicate() == ICmpInst::ICMP_EQ) {
+                  Value *CmpLHS = CmpInst->getOperand(0);
+                  Value *CmpRHS = CmpInst->getOperand(1);
+                  
+                  // If one operand is constant, the other is uniform and equals that constant
+                  if (auto *C = dyn_cast<Constant>(CmpRHS)) {
+                    UniformValues[CmpLHS] = C;
+                  } else if (auto *C = dyn_cast<Constant>(CmpLHS)) {
+                    UniformValues[CmpRHS] = C;
+                  }
+                  // TODO: Handle case where both operands are variables
+                }
+              }
+            }
+          }
+        } else if (IntrinsicCall->getIntrinsicID() == Intrinsic::amdgcn_readfirstlane) {
+          // assume(readfirstlane(x) == c) -> x is uniform and equals c
+          if (auto *C = dyn_cast<Constant>(RHS)) {
+            Value *ReadFirstLaneArg = IntrinsicCall->getArgOperand(0);
+            UniformValues[ReadFirstLaneArg] = C;
+          }
+        }
+      }
+      
+      // Handle the reverse case too
+      if (auto *IntrinsicCall = dyn_cast<IntrinsicInst>(RHS)) {
+        if (IntrinsicCall->getIntrinsicID() == Intrinsic::amdgcn_ballot) {
+          if (match(LHS, m_AllOnes())) {
+            Value *BallotArg = IntrinsicCall->getArgOperand(0);
+            if (BallotArg->getType()->isIntegerTy(1)) {
+              UniformValues[BallotArg] = ConstantInt::getTrue(BallotArg->getType());
+            }
+          }
+        } else if (IntrinsicCall->getIntrinsicID() == Intrinsic::amdgcn_readfirstlane) {
+          if (auto *C = dyn_cast<Constant>(LHS)) {
+            Value *ReadFirstLaneArg = IntrinsicCall->getArgOperand(0);
+            UniformValues[ReadFirstLaneArg] = C;
+          }
+        }
+      }
+    }
+  }
+  
+  // Pattern 2: assume(X) where X is i1 -> X is uniform and equals true  
+  if (AssumedCondition->getType()->isIntegerTy(1) && !isa<ICmpInst>(AssumedCondition)) {
+    UniformValues[AssumedCondition] = ConstantInt::getTrue(AssumedCondition->getType());
+  }
+  
+  // Now optimize dominated uses of all discovered uniform values
+  for (auto &[UniformValue, UniformConstant] : UniformValues) {
+    SmallVector<Use *, 8> DominatedUses;
+    
+    // Find all uses dominated by the assume
+    // Skip if the value doesn't have a use list (e.g., constants)
+    if (!UniformValue->hasUseList())
+      continue;
+      
+    for (Use &U : UniformValue->uses()) {
+      Instruction *UseInst = dyn_cast<Instruction>(U.getUser());
+      if (!UseInst || UseInst == Assume)
+        continue;
+        
+      // Critical: Check dominance using InstCombine's infrastructure  
+      if (isValidAssumeForContext(Assume, UseInst, &DT)) {
+        DominatedUses.push_back(&U);
+      }
+    }
+    
+    // Replace only dominated uses with the uniform constant
+    for (Use *U : DominatedUses) {
+      U->set(UniformConstant);
+      Worklist.pushValue(U->getUser());
+    }
+    
+    // Mark for further optimization if we made changes
+    if (!DominatedUses.empty()) {
+      Worklist.pushValue(UniformValue);
+    }
+  }
+}
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -124,8 +124,6 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
       BinaryOperator &I);
   Instruction *foldVariableSignZeroExtensionOfVariableHighBitExtract(
       BinaryOperator &OldAShr);
-  
-
   Instruction *visitAShr(BinaryOperator &I);
   Instruction *visitLShr(BinaryOperator &I);
   Instruction *commonShiftTransforms(BinaryOperator &I);
@@ -231,6 +229,9 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
 private:
   bool annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI);
   bool isDesirableIntType(unsigned BitWidth) const;
+  
+  /// Optimize uses of variables that are established as uniform by assume intrinsics.
+  void optimizeAssumedUniformValues(AssumeInst *Assume);
   bool shouldChangeType(unsigned FromBitWidth, unsigned ToBitWidth) const;
   bool shouldChangeType(Type *From, Type *To) const;
   Value *dyn_castNegVal(Value *V) const;
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
