Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[Reassociate] Use UniformityInfo to group uniform operands together (#198507)

Browse source 
On targets with branch divergence (TTI.hasBranchDivergence()), boost the
rank of divergent operands in ReassociateExpression() so they sort
towards the root of the expression tree. This clusters uniform operands
at the leaves, letting the backend lower the inner sub-expression to a
scalar ALU instruction (s_mul/s_or/s_and) instead of a vector.

Ref:
https://github.com/llvm/llvm-project/pull/175167#issuecomment-4476483692
This commit is contained in:
Pankaj Dwivedi 2026-07-23 11:21:20 +05:30 committed by GitHub
parent 22cb41c9ad
commit 87304516a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 173 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
llvm/include/llvm/Transforms/Scalar/Reassociate.h
View file

@ -25,6 +25,7 @@
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Analysis/UniformityAnalysis.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/ValueHandle.h"
@ -95,9 +96,12 @@ protected:
DenseMap<std::pair<Value *, Value *>, PairMapValue> PairMap[NumBinaryOps];
bool MadeChange;
UniformityInfo *UA = nullptr;
public:
LLVM_ABI PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
LLVM_ABI PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
LLVM_ABI PreservedAnalyses runImpl(Function &F, UniformityInfo &UI);
private:
void BuildRankMap(Function &F, ReversePostOrderTraversal<Function *> &RPOT);

54
llvm/lib/Transforms/Scalar/Reassociate.cpp
View file

@ -2396,6 +2396,31 @@ void ReassociatePass::ReassociateExpression(BinaryOperator *I) {
LLVM_DEBUG(dbgs() << "RAIn:\t"; PrintOps(I, Ops); dbgs() << '\n');
// Boost the rank of divergent operands so they sort towards the root of the
// expression tree, clustering uniform operands together at the leaves. On
// targets without divergence UniformityInfo is empty and this is a no-op.
//
// Example: (uniform1 + divergent) + uniform2
// -> (uniform1 + uniform2) + divergent
if (UA && Ops.size() > 2) {
constexpr unsigned DivergentRankOffset = 1U << 28;
BasicBlock *ParentBB = I->getParent();
for (ValueEntry &Entry : Ops) {
if (isa<Constant>(Entry.Op))
continue;
bool Divergent = false;
for (const Use &U : Entry.Op->uses()) {
Instruction *Usr = dyn_cast<Instruction>(U.getUser());
if (Usr && Usr->getParent() == ParentBB) {
Divergent = UA->isDivergentAtUse(U);
break;
}
}
if (Divergent)
Entry.Rank += DivergentRankOffset;
}
}
// Now that we have linearized the tree to a list and have gathered all of
// the operands and their ranks, sort the operands by their rank. Use a
// stable_sort so that values with equal ranks will have their relative
@ -2645,7 +2670,17 @@ ReassociatePass::BuildPairMap(ReversePostOrderTraversal<Function *> &RPOT) {
}
}
PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
PreservedAnalyses ReassociatePass::run(Function &F,
FunctionAnalysisManager &AM) {
// UniformityInfo is empty (and cheap) on targets without branch divergence,
// so request it unconditionally.
UniformityInfo &UI = AM.getResult<UniformityInfoAnalysis>(F);
return runImpl(F, UI);
}
PreservedAnalyses ReassociatePass::runImpl(Function &F, UniformityInfo &UI) {
UA = &UI;
// Get the functions basic blocks in Reverse Post Order. This order is used by
// BuildRankMap to pre calculate ranks correctly. It also excludes dead basic
// blocks (it has been seen that the analysis in this pass could hang when
@ -2708,11 +2743,12 @@ PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
}
}
// We are done with the rank map and pair map.
// We are done with the rank map, pair map, and uniformity info.
RankMap.clear();
ValueRankMap.clear();
for (auto &Entry : PairMap)
Entry.clear();
UA = nullptr;
if (MadeChange) {
PreservedAnalyses PA;
@ -2739,13 +2775,16 @@ public:
if (skipFunction(F))
return false;
FunctionAnalysisManager DummyFAM;
auto PA = Impl.run(F, DummyFAM);
UniformityInfo &UI =
getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
PreservedAnalyses PA = Impl.runImpl(F, UI);
return !PA.areAllPreserved();
}
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.setPreservesCFG();
AU.addRequired<UniformityInfoWrapperPass>();
AU.addPreserved<AAResultsWrapperPass>();
AU.addPreserved<BasicAAWrapperPass>();
AU.addPreserved<GlobalsAAWrapperPass>();
@ -2756,8 +2795,11 @@ public:
char ReassociateLegacyPass::ID = 0;
INITIALIZE_PASS(ReassociateLegacyPass, "reassociate",
"Reassociate expressions", false, false)
INITIALIZE_PASS_BEGIN(ReassociateLegacyPass, "reassociate",
"Reassociate expressions", false, false)
INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
INITIALIZE_PASS_END(ReassociateLegacyPass, "reassociate",
"Reassociate expressions", false, false)
// Public interface to the Reassociate pass
FunctionPass *llvm::createReassociatePass() {

113
llvm/test/CodeGen/AMDGPU/reassoc-scalar.ll
View file

@ -125,5 +125,118 @@ bb:
ret void
}
; Multiply: (x * tid) * y -- backend does NOT reassociate mul to scalar.
; Both muls use VALU even though x and y are uniform.
; The reassociate pass groups the uniform operands (x * y) so this
; lowers to s_mul + v_mul; see reassoc-uniform-e2e.ll.
; GCN-LABEL: reassoc_mul_i32:
; GCN-NOT: s_mul
; GCN: v_mul_lo_u32
; GCN: v_mul_lo_u32
define void @reassoc_mul_i32(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%mul1 = mul i32 %x, %tid
%mul2 = mul i32 %mul1, %y
store i32 %mul2, ptr addrspace(1) %arg, align 4
ret void
}
; Multiply with uniforms pre-grouped: (x * y) * tid -- produces s_mul + v_mul.
; GCN-LABEL: reassoc_mul_i32_uniform_grouped:
; GCN: s_mul_i32 [[MUL1:s[0-9]+]], s{{[0-9]+}}, s{{[0-9]+}}
; GCN: v_mul_lo_u32 v{{[0-9]+}}, [[MUL1]], v{{[0-9]+}}
define void @reassoc_mul_i32_uniform_grouped(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%mul1 = mul i32 %x, %y
%mul2 = mul i32 %mul1, %tid
store i32 %mul2, ptr addrspace(1) %arg, align 4
ret void
}
; OR: (x | tid) | y -- backend does NOT reassociate or to scalar.
; Both ors use VALU.
; The reassociate pass groups the uniform operands (x | y) so this
; lowers to s_or + v_or; see reassoc-uniform-e2e.ll.
; GCN-LABEL: reassoc_or_i32:
; GCN-NOT: s_or
; GCN: v_or_b32_e32
; GCN: v_or_b32_e32
define void @reassoc_or_i32(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%or1 = or i32 %x, %tid
%or2 = or i32 %or1, %y
store i32 %or2, ptr addrspace(1) %arg, align 4
ret void
}
; OR with uniforms pre-grouped: (x | y) | tid -- produces s_or + v_or.
; GCN-LABEL: reassoc_or_i32_uniform_grouped:
; GCN: s_or_b32 [[OR1:s[0-9]+]], s{{[0-9]+}}, s{{[0-9]+}}
; GCN: v_or_b32_e32 v{{[0-9]+}}, [[OR1]], v{{[0-9]+}}
define void @reassoc_or_i32_uniform_grouped(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%or1 = or i32 %x, %y
%or2 = or i32 %or1, %tid
store i32 %or2, ptr addrspace(1) %arg, align 4
ret void
}
; AND: (x & tid) & y -- backend does NOT reassociate and to scalar.
; Both ands use VALU.
; The reassociate pass groups the uniform operands (x & y) so this
; lowers to s_and + v_and; see reassoc-uniform-e2e.ll.
; GCN-LABEL: reassoc_and_i32:
; GCN-NOT: s_and
; GCN: v_and_b32_e32
; GCN: v_and_b32_e32
define void @reassoc_and_i32(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%and1 = and i32 %x, %tid
%and2 = and i32 %and1, %y
store i32 %and2, ptr addrspace(1) %arg, align 4
ret void
}
; AND with uniforms pre-grouped: (x & y) & tid -- produces s_and + v_and.
; GCN-LABEL: reassoc_and_i32_uniform_grouped:
; GCN: s_and_b32 [[AND1:s[0-9]+]], s{{[0-9]+}}, s{{[0-9]+}}
; GCN: v_and_b32_e32 v{{[0-9]+}}, [[AND1]], v{{[0-9]+}}
define void @reassoc_and_i32_uniform_grouped(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%and1 = and i32 %x, %y
%and2 = and i32 %and1, %tid
store i32 %and2, ptr addrspace(1) %arg, align 4
ret void
}
; XOR: (x ^ tid) ^ y -- backend ALREADY reassociates xor to scalar.
; The DAG combiner handles this, producing s_xor + v_xor for both variants.
; GCN-LABEL: reassoc_xor_i32:
; GCN: s_xor_b32 [[XOR1:s[0-9]+]], s{{[0-9]+}}, s{{[0-9]+}}
; GCN: v_xor_b32_e32 v{{[0-9]+}}, [[XOR1]], v{{[0-9]+}}
define void @reassoc_xor_i32(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%xor1 = xor i32 %x, %tid
%xor2 = xor i32 %xor1, %y
store i32 %xor2, ptr addrspace(1) %arg, align 4
ret void
}
; Multi-use: intermediate (x * tid) is stored, blocks reassociation.
; Both muls use VALU. The reassociate pass cannot group the uniforms here
; because the intermediate result is used more than once.
; GCN-LABEL: reassoc_mul_i32_multiuse:
; GCN-NOT: s_mul
; GCN: v_mul_lo_u32
; GCN: v_mul_lo_u32
define void @reassoc_mul_i32_multiuse(ptr addrspace(1) inreg %arg, i32 inreg %x, i32 inreg %y) {
%tid = tail call i32 @llvm.amdgcn.workitem.id.x()
%mul1 = mul i32 %x, %tid
%mul2 = mul i32 %mul1, %y
store volatile i32 %mul1, ptr addrspace(1) %arg, align 4
store volatile i32 %mul2, ptr addrspace(1) %arg, align 4
ret void
}
declare i32 @llvm.amdgcn.workitem.id.x()
declare i32 @llvm.amdgcn.workitem.id.y()

1
llvm/test/Other/new-pm-defaults.ll
View file

@ -157,6 +157,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: LoopAnalysis
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis

1
llvm/test/Other/new-pm-thinlto-postlink-defaults.ll
View file

@ -93,6 +93,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: LoopAnalysis
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis

1
llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll
View file

@ -82,6 +82,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis
; CHECK-O-NEXT: Running pass: LoopSimplifyPass

1
llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll
View file

@ -89,6 +89,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis
; CHECK-O-NEXT: Running pass: LoopSimplifyPass

1
llvm/test/Other/new-pm-thinlto-prelink-defaults.ll
View file

@ -119,6 +119,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: LoopAnalysis
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis

1
llvm/test/Other/new-pm-thinlto-prelink-pgo-defaults.ll
View file

@ -129,6 +129,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis
; CHECK-O-NEXT: Running pass: LoopSimplifyPass

1
llvm/test/Other/new-pm-thinlto-prelink-samplepgo-defaults.ll
View file

@ -94,6 +94,7 @@
; CHECK-O23-NEXT: Running pass: TailCallElimPass
; CHECK-O-NEXT: Running pass: SimplifyCFGPass
; CHECK-O-NEXT: Running pass: ReassociatePass
; CHECK-O-NEXT: Running analysis: UniformityInfoAnalysis
; CHECK-O23-NEXT: Running pass: ConstraintEliminationPass
; CHECK-O23-NEXT: Running analysis: ScalarEvolutionAnalysis
; CHECK-O-NEXT: Running pass: LoopSimplifyPass