[llvm] [SLSR] Prove disjoint instruction reuse using an annotation-free IR (PR #218308)
Nick Riasanovsky via llvm-commits
llvm-commits at lists.llvm.org
Sun Aug 23 20:39:50 PDT 2026
https://github.com/njriasan created https://github.com/llvm/llvm-project/pull/218308
Allow SLSR to reuse instructions derived from or disjoint when disjointness can be proven independently of poison-generating annotations. The pass clones the function, removes poison-generating annotations, and checks disjointness against this scrubbed IR. This permits structurally safe reuse while rejecting proofs that depend on annotations such as nneg. Includes positive and negative regression tests covering both cases.
>From 3aaba058ba87ca58acc5c75cb945aec3739837b9 Mon Sep 17 00:00:00 2001
From: Nick Riasanovsky <njriasan at meta.com>
Date: Sun, 23 Aug 2026 15:32:38 -0700
Subject: [PATCH] [SLSR] Prove disjoint reuse on scrubbed IR
Allow SLSR to reuse instructions derived from or disjoint when disjointness can be independently proved after removing poison-generating annotations.
Differential Revision: https://phabricator.intern.facebook.com/D117020071
---
.../Scalar/StraightLineStrengthReduce.cpp | 113 +++++++++++++++++-
.../StraightLineStrengthReduce/slsr-mul.ll | 48 ++++++++
2 files changed, 160 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Transforms/Scalar/StraightLineStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/StraightLineStrengthReduce.cpp
index 8826127dd7ae3..08f941c712499 100644
--- a/llvm/lib/Transforms/Scalar/StraightLineStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/StraightLineStrengthReduce.cpp
@@ -70,9 +70,12 @@
#include "llvm/Transforms/Scalar/StraightLineStrengthReduce.h"
#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TargetTransformInfo.h"
@@ -83,6 +86,7 @@
#include "llvm/IR/Dominators.h"
#include "llvm/IR/GetElementPtrTypeIterator.h"
#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
@@ -96,11 +100,14 @@
#include "llvm/Support/DebugCounter.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
#include <cassert>
#include <cstdint>
#include <limits>
#include <list>
+#include <memory>
#include <queue>
#include <vector>
@@ -402,6 +409,11 @@ class StraightLineStrengthReduce {
// Rewrites candidate C with respect to Basis.
void rewriteCandidate(const Candidate &C);
+ bool canReuseInstruction(const SCEV *S, Instruction *I,
+ SmallVectorImpl<Instruction *> &DropInsts);
+ bool isDisjointAfterScrubbing(const PossiblyDisjointInst *PDI);
+ void initializeScrubbedFunction(const Function &F);
+
// Emit code that computes the "bump" from Basis to C.
static Value *emitBump(const Candidate &Basis, const Candidate &C,
IRBuilder<> &Builder, const DataLayout *DL);
@@ -412,6 +424,12 @@ class StraightLineStrengthReduce {
TargetTransformInfo *TTI = nullptr;
std::list<Candidate> Candidates;
+ std::unique_ptr<Function> ScrubbedFunction;
+ ValueToValueMapTy ScrubbedValues;
+ std::unique_ptr<DominatorTree> ScrubbedDT;
+ std::unique_ptr<AssumptionCache> ScrubbedAC;
+ DenseMap<const Instruction *, bool> ScrubbedDisjointCache;
+
// Map from SCEV to instructions that represent the value,
// instructions are sorted in depth-first order.
DenseMap<const SCEV *, SmallSetVector<Instruction *, 2>> SCEVToInsts;
@@ -716,6 +734,90 @@ bool StraightLineStrengthReduce::isSimilar(Candidate &C, Candidate &Basis,
Basis.CandidateKind == C.CandidateKind;
}
+void StraightLineStrengthReduce::initializeScrubbedFunction(const Function &F) {
+ if (ScrubbedFunction)
+ return;
+
+ ScrubbedFunction.reset(
+ Function::Create(F.getFunctionType(), GlobalValue::ExternalLinkage));
+ auto NewArg = ScrubbedFunction->arg_begin();
+ for (const Argument &Arg : F.args()) {
+ NewArg->setName(Arg.getName());
+ ScrubbedValues[&Arg] = &*NewArg++;
+ }
+
+ SmallVector<ReturnInst *, 4> Returns;
+ CloneFunctionInto(ScrubbedFunction.get(), &F, ScrubbedValues,
+ CloneFunctionChangeType::LocalChangesOnly, Returns);
+ for (Instruction &I : instructions(*ScrubbedFunction))
+ I.dropPoisonGeneratingAnnotations();
+
+ ScrubbedDT = std::make_unique<DominatorTree>(*ScrubbedFunction);
+ ScrubbedAC = std::make_unique<AssumptionCache>(*ScrubbedFunction);
+}
+
+bool StraightLineStrengthReduce::isDisjointAfterScrubbing(
+ const PossiblyDisjointInst *PDI) {
+ auto *OrigI = cast<Instruction>(PDI);
+ auto It = ScrubbedDisjointCache.find(OrigI);
+ if (It != ScrubbedDisjointCache.end())
+ return It->second;
+
+ initializeScrubbedFunction(*OrigI->getFunction());
+ auto *ScrubbedI = cast<Instruction>(
+ ScrubbedValues.lookup(const_cast<Instruction *>(OrigI)));
+ bool Proven = haveNoCommonBitsSet(
+ ScrubbedI->getOperand(0), ScrubbedI->getOperand(1),
+ SimplifyQuery(*DL, ScrubbedDT.get(), ScrubbedAC.get(), ScrubbedI));
+ ScrubbedDisjointCache[OrigI] = Proven;
+ return Proven;
+}
+
+bool StraightLineStrengthReduce::canReuseInstruction(
+ const SCEV *S, Instruction *I, SmallVectorImpl<Instruction *> &DropInsts) {
+ if (programUndefinedIfPoison(I))
+ return true;
+
+ SmallPtrSet<const Value *, 8> PoisonVals;
+ SE->getPoisonGeneratingValues(PoisonVals, S);
+
+ SmallVector<Value *> Worklist;
+ SmallPtrSet<Value *, 8> Visited;
+ Worklist.push_back(I);
+ while (!Worklist.empty()) {
+ Value *V = Worklist.pop_back_val();
+ if (!Visited.insert(V).second)
+ continue;
+ if (Visited.size() > 16)
+ return false;
+ if (PoisonVals.contains(V) || isGuaranteedNotToBePoison(V))
+ continue;
+
+ auto *CurI = dyn_cast<Instruction>(V);
+ if (!CurI)
+ return false;
+
+ if (auto *PDI = dyn_cast<PossiblyDisjointInst>(CurI))
+ if (PDI->isDisjoint() && !isDisjointAfterScrubbing(PDI))
+ return false;
+
+ if (auto *II = dyn_cast<IntrinsicInst>(CurI);
+ II && II->getIntrinsicID() == Intrinsic::vscale)
+ continue;
+
+ if (canCreatePoison(cast<Operator>(CurI),
+ /*ConsiderFlagsAndMetadata=*/false))
+ return false;
+
+ if (CurI->hasPoisonGeneratingAnnotations())
+ DropInsts.push_back(CurI);
+
+ for (Value *Op : CurI->operands())
+ Worklist.push_back(Op);
+ }
+ return true;
+}
+
// Try to find a Delta that C can reuse Basis to rewrite.
// Set C.Delta, C.Basis, and C.DeltaKind if found.
// Return true if found a constant delta.
@@ -1053,7 +1155,7 @@ void StraightLineStrengthReduce::allocateCandidatesAndFindBasis(
// same instruction. The DropList is stored on the Candidate so
// candidatePredicate can drop the flags when a rewrite is being done.
if (!EnablePoisonReuseGuard ||
- SE->canReuseInstruction(SE->getSCEV(I), I, Candidates.back().DropList)) {
+ canReuseInstruction(SE->getSCEV(I), I, Candidates.back().DropList)) {
CandidateDict.add(Candidates.back());
}
}
@@ -1364,6 +1466,15 @@ bool StraightLineStrengthReduceLegacyPass::runOnFunction(Function &F) {
bool StraightLineStrengthReduce::runOnFunction(Function &F) {
LLVM_DEBUG(dbgs() << "SLSR on Function: " << F.getName() << "\n");
+ // Candidate discovery may drop poison annotations, so snapshot the entry IR.
+ for (Instruction &I : instructions(F)) {
+ auto *PDI = dyn_cast<PossiblyDisjointInst>(&I);
+ if (PDI && PDI->isDisjoint()) {
+ initializeScrubbedFunction(F);
+ break;
+ }
+ }
+
// Traverse the dominator tree in the depth-first order. This order makes sure
// all bases of a candidate are in Candidates when we process it.
for (const auto Node : depth_first(DT))
diff --git a/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll b/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll
index 485c666a8779b..1873410155548 100644
--- a/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll
+++ b/llvm/test/Transforms/StraightLineStrengthReduce/slsr-mul.ll
@@ -68,6 +68,54 @@ define void @or(i32 %a, i32 %s) {
ret void
}
+; Reuse a basis containing `or disjoint` when disjointness follows from the
+; operands themselves. The reused instruction has the annotation dropped so
+; it is no more poisonous than the SCEV expression it replaces.
+define void @or_disjoint_structural(i32 %a, i32 %s) {
+; CHECK-LABEL: @or_disjoint_structural(
+; CHECK-NEXT: [[SHL:%.*]] = shl i32 [[A:%.*]], 4
+; CHECK-NEXT: [[OR:%.*]] = or i32 [[SHL]], 3
+; CHECK-NEXT: [[MUL0:%.*]] = mul i32 [[OR]], [[S:%.*]]
+; CHECK-NEXT: call void @foo(i32 [[MUL0]])
+; CHECK-NEXT: [[MUL1:%.*]] = add i32 [[MUL0]], [[S]]
+; CHECK-NEXT: call void @foo(i32 [[MUL1]])
+; CHECK-NEXT: ret void
+ %shl = shl i32 %a, 4
+ %or = or disjoint i32 %shl, 3
+ %mul0 = mul i32 %or, %s
+ call void @foo(i32 %mul0)
+
+ %add = add i32 %shl, 4
+ %mul1 = mul i32 %add, %s
+ call void @foo(i32 %mul1)
+ ret void
+}
+
+; Do not use poison-generating annotations to prove that an `or disjoint`
+; basis is safe. Without `nneg`, bit 7 of %zext may overlap %bit.
+define void @or_disjoint_annotation(i32 %a, i8 %b, i32 %s) {
+; CHECK-LABEL: @or_disjoint_annotation(
+; CHECK-NEXT: [[BIT:%.*]] = and i32 [[A:%.*]], 128
+; CHECK-NEXT: [[ZEXT:%.*]] = zext nneg i8 [[B:%.*]] to i32
+; CHECK-NEXT: [[OR:%.*]] = or disjoint i32 [[BIT]], [[ZEXT]]
+; CHECK-NEXT: [[MUL0:%.*]] = mul i32 [[OR]], [[S:%.*]]
+; CHECK-NEXT: call void @foo(i32 [[MUL0]])
+; CHECK-NEXT: [[ADD:%.*]] = add i32 [[BIT]], [[ZEXT]]
+; CHECK-NEXT: [[MUL1:%.*]] = mul i32 [[ADD]], [[S]]
+; CHECK-NEXT: call void @foo(i32 [[MUL1]])
+; CHECK-NEXT: ret void
+ %bit = and i32 %a, 128
+ %zext = zext nneg i8 %b to i32
+ %or = or disjoint i32 %bit, %zext
+ %mul0 = mul i32 %or, %s
+ call void @foo(i32 %mul0)
+
+ %add = add i32 %bit, %zext
+ %mul1 = mul i32 %add, %s
+ call void @foo(i32 %mul1)
+ ret void
+}
+
; foo(a * b)
; foo((a + 1) * b)
; foo(a * (b + 1))
More information about the llvm-commits
mailing list