Reland "Redesign Straight-Line Strength Reduction (SLSR) (#162930)"

llvm/lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

@@ -495,6 +495,11 @@ class StraightLineStrengthReduce {
     return S;
   }
 
+  bool candidatePredicate(Candidate *Basis, Candidate &C, Candidate::DKind K);
+
+  bool searchFrom(const CandidateDictTy::BBToCandsTy &BBToCands, Candidate &C,
+                  Candidate::DKind K);
+
   // Get the nearest instruction before CI that represents the value of S,
   // return nullptr if no instruction is associated with S or S is not a
   // reusable expression.
@@ -538,7 +543,8 @@ class StraightLineStrengthReduce {
 
   Candidate *pickRewriteCandidate(Instruction *I) const;
   void sortCandidateInstructions();
-  static Constant *getIndexDelta(Candidate &C, Candidate &Basis);
+  Value *getDelta(const Candidate &C, const Candidate &Basis,
+                  Candidate::DKind K) const;
   static bool isSimilar(Candidate &C, Candidate &Basis, Candidate::DKind K);
 
   // Add Basis -> C in DependencyGraph and propagate
@@ -579,8 +585,7 @@ inline raw_ostream &operator<<(raw_ostream &OS,
   return OS;
 }
 
-LLVM_DUMP_METHOD
-inline raw_ostream &
+[[maybe_unused]] LLVM_DUMP_METHOD inline raw_ostream &
 operator<<(raw_ostream &OS, const StraightLineStrengthReduce::DeltaInfo &DI) {
   OS << "Cand: " << *DI.Cand << "\n";
   OS << "Delta Kind: ";
@@ -625,14 +630,25 @@ static void unifyBitWidth(APInt &A, APInt &B) {
     B = B.sext(A.getBitWidth());
 }
 
-Constant *StraightLineStrengthReduce::getIndexDelta(Candidate &C,
-                                                    Candidate &Basis) {
-  APInt Idx = C.Index->getValue(), BasisIdx = Basis.Index->getValue();
-  unifyBitWidth(Idx, BasisIdx);
-  APInt IndexDelta = Idx - BasisIdx;
-  IntegerType *DeltaType =
-      IntegerType::get(C.Ins->getContext(), IndexDelta.getBitWidth());
-  return ConstantInt::get(DeltaType, IndexDelta);
+Value *StraightLineStrengthReduce::getDelta(const Candidate &C,
+                                            const Candidate &Basis,
+                                            Candidate::DKind K) const {
+  if (K == Candidate::IndexDelta) {
+    APInt Idx = C.Index->getValue();
+    APInt BasisIdx = Basis.Index->getValue();
+    unifyBitWidth(Idx, BasisIdx);
+    APInt IndexDelta = Idx - BasisIdx;
+    IntegerType *DeltaType =
+        IntegerType::get(C.Ins->getContext(), IndexDelta.getBitWidth());
+    return ConstantInt::get(DeltaType, IndexDelta);
+  } else if (K == Candidate::BaseDelta || K == Candidate::StrideDelta) {
+    const SCEV *BasisPart =
+        (K == Candidate::BaseDelta) ? Basis.Base : Basis.StrideSCEV;
+    const SCEV *CandPart = (K == Candidate::BaseDelta) ? C.Base : C.StrideSCEV;
+    const SCEV *Diff = SE->getMinusSCEV(CandPart, BasisPart);
+    return getNearestValueOfSCEV(Diff, C.Ins);
+  }
+  return nullptr;
 }
 
 bool StraightLineStrengthReduce::isSimilar(Candidate &C, Candidate &Basis,
@@ -654,106 +670,104 @@ bool StraightLineStrengthReduce::isSimilar(Candidate &C, Candidate &Basis,
          Basis.CandidateKind == C.CandidateKind;
 }
 
-void StraightLineStrengthReduce::setBasisAndDeltaFor(Candidate &C) {
-  auto SearchFrom = [this, &C](const CandidateDictTy::BBToCandsTy &BBToCands,
-                               auto IsTarget) -> bool {
-    // Search dominating candidates by walking the immediate-dominator chain
-    // from the candidate's defining block upward. Visiting blocks in this
-    // order ensures we prefer the closest dominating basis.
-    const BasicBlock *BB = C.Ins->getParent();
-    while (BB) {
-      auto It = BBToCands.find(BB);
-      if (It != BBToCands.end())
-        for (Candidate *Basis : reverse(It->second))
-          if (IsTarget(Basis))
-            return true;
-
-      const DomTreeNode *Node = DT->getNode(BB);
-      if (!Node)
-        break;
-      Node = Node->getIDom();
-      BB = Node ? Node->getBlock() : nullptr;
-    }
+// 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.
+// Return false if not found or the delta is not a constant.
+bool StraightLineStrengthReduce::candidatePredicate(Candidate *Basis,
+                                                    Candidate &C,
+                                                    Candidate::DKind K) {
+  SmallVector<Instruction *> DropPoisonGeneratingInsts;
+  // Ensure the IR of Basis->Ins is not more poisonous than its SCEV.
+  if (!isSimilar(C, *Basis, K) ||
+      !SE->canReuseInstruction(SE->getSCEV(Basis->Ins), Basis->Ins,
+                               DropPoisonGeneratingInsts))
+    return false;
+
+  assert(DT->dominates(Basis->Ins, C.Ins));
+  Value *Delta = getDelta(C, *Basis, K);
+  if (!Delta)
     return false;
-  };
 
-  // Priority:
-  // Constant Delta from Index > Constant Delta from Base >
-  // Constant Delta from Stride > Variable Delta from Base or Stride
-  // TODO: Change the priority to align with the cost model.
+  // IndexDelta rewrite is not always profitable, e.g.,
+  // X = B + 8 * S
+  // Y = B + S,
+  // rewriting Y to X - 7 * S is probably a bad idea.
+  // So, we need to check if the rewrite form's computation efficiency
+  // is better than the original form.
+  if (K == Candidate::IndexDelta &&
+      !C.isProfitableRewrite(Delta, Candidate::IndexDelta))
+    return false;
 
-  // First, look for a constant index-diff basis
-  if (const auto *IndexDeltaCandidates =
-          CandidateDict.getCandidatesWithDeltaKind(C, Candidate::IndexDelta)) {
-    bool FoundConstDelta =
-        SearchFrom(*IndexDeltaCandidates, [&DT = DT, &C](Candidate *Basis) {
-          if (isSimilar(C, *Basis, Candidate::IndexDelta)) {
-            assert(DT->dominates(Basis->Ins, C.Ins));
-            auto *Delta = getIndexDelta(C, *Basis);
-            if (!C.isProfitableRewrite(Delta, Candidate::IndexDelta))
-              return false;
-            C.Basis = Basis;
-            C.DeltaKind = Candidate::IndexDelta;
-            C.Delta = Delta;
-            LLVM_DEBUG(dbgs() << "Found delta from Index " << *C.Delta << "\n");
-            return true;
-          }
-          return false;
-        });
-    if (FoundConstDelta)
-      return;
+  // If there is a Delta that we can reuse Basis to rewrite C,
+  // clean up DropPoisonGeneratingInsts returned by successful
+  // SE->canReuseInstruction()
+  for (Instruction *I : DropPoisonGeneratingInsts)
+    I->dropPoisonGeneratingAnnotations();
+
+  // Record delta if none has been found yet, or the new delta is
+  // a constant that is better than the existing delta.
+  if (!C.Delta || isa<ConstantInt>(Delta)) {
+    C.Delta = Delta;
+    C.Basis = Basis;
+    C.DeltaKind = K;
   }
+  return isa<ConstantInt>(C.Delta);
+}
 
-  // No constant-index-diff basis found. look for the best possible base-diff
-  // or stride-diff basis
-  // Base/Stride diffs not supported for form (B + i) * S
-  if (C.CandidateKind == Candidate::Mul)
-    return;
-
-  auto For = [this, &C](Candidate::DKind K) {
-    // return true if find a Basis with constant delta and stop searching,
-    // return false if did not find a Basis or the delta is not a constant
-    // and continue searching for a Basis with constant delta
-    return [K, this, &C](Candidate *Basis) -> bool {
-      if (!isSimilar(C, *Basis, K))
-        return false;
+// return true if find a Basis with constant delta and stop searching,
+// return false if did not find a Basis or the delta is not a constant
+// and continue searching for a Basis with constant delta
+bool StraightLineStrengthReduce::searchFrom(
+    const CandidateDictTy::BBToCandsTy &BBToCands, Candidate &C,
+    Candidate::DKind K) {
 
-      assert(DT->dominates(Basis->Ins, C.Ins));
-      const SCEV *BasisPart =
-          (K == Candidate::BaseDelta) ? Basis->Base : Basis->StrideSCEV;
-      const SCEV *CandPart =
-          (K == Candidate::BaseDelta) ? C.Base : C.StrideSCEV;
-      const SCEV *Diff = SE->getMinusSCEV(CandPart, BasisPart);
-      Value *AvailableVal = getNearestValueOfSCEV(Diff, C.Ins);
-      if (!AvailableVal)
-        return false;
+  // Stride delta rewrite on Mul form is usually non-profitable, and Base
+  // delta rewrite sometimes is profitable, so we do not support them on Mul.
+  if (C.CandidateKind == Candidate::Mul && K != Candidate::IndexDelta)
+    return false;
 
-      // Record delta if none has been found yet, or the new delta is
-      // a constant that is better than the existing delta.
-      if (!C.Delta || isa<ConstantInt>(AvailableVal)) {
-        C.Delta = AvailableVal;
-        C.Basis = Basis;
-        C.DeltaKind = K;
-      }
-      return isa<ConstantInt>(C.Delta);
-    };
-  };
+  // Search dominating candidates by walking the immediate-dominator chain
+  // from the candidate's defining block upward. Visiting blocks in this
+  // order ensures we prefer the closest dominating basis.
+  const BasicBlock *BB = C.Ins->getParent();
+  while (BB) {
+    auto It = BBToCands.find(BB);
+    if (It != BBToCands.end())
+      for (Candidate *Basis : reverse(It->second))
+        if (candidatePredicate(Basis, C, K))
+          return true;
+
+    const DomTreeNode *Node = DT->getNode(BB);
+    if (!Node)
+      break;
+    Node = Node->getIDom();
+    BB = Node ? Node->getBlock() : nullptr;
+  }
+  return false;
+}
 
+void StraightLineStrengthReduce::setBasisAndDeltaFor(Candidate &C) {
   if (const auto *BaseDeltaCandidates =
-          CandidateDict.getCandidatesWithDeltaKind(C, Candidate::BaseDelta)) {
-    if (SearchFrom(*BaseDeltaCandidates, For(Candidate::BaseDelta))) {
+          CandidateDict.getCandidatesWithDeltaKind(C, Candidate::BaseDelta))
+    if (searchFrom(*BaseDeltaCandidates, C, Candidate::BaseDelta)) {
       LLVM_DEBUG(dbgs() << "Found delta from Base: " << *C.Delta << "\n");
       return;
     }
-  }
 
   if (const auto *StrideDeltaCandidates =
-          CandidateDict.getCandidatesWithDeltaKind(C, Candidate::StrideDelta)) {
-    if (SearchFrom(*StrideDeltaCandidates, For(Candidate::StrideDelta))) {
+          CandidateDict.getCandidatesWithDeltaKind(C, Candidate::StrideDelta))
+    if (searchFrom(*StrideDeltaCandidates, C, Candidate::StrideDelta)) {
       LLVM_DEBUG(dbgs() << "Found delta from Stride: " << *C.Delta << "\n");
       return;
     }
-  }
+
+  if (const auto *IndexDeltaCandidates =
+          CandidateDict.getCandidatesWithDeltaKind(C, Candidate::IndexDelta))
+    if (searchFrom(*IndexDeltaCandidates, C, Candidate::IndexDelta)) {
+      LLVM_DEBUG(dbgs() << "Found delta from Index: " << *C.Delta << "\n");
+      return;
+    }
 
   // If we did not find a constant delta, we might have found a variable delta
   if (C.Delta) {
@@ -792,7 +806,8 @@ auto StraightLineStrengthReduce::compressPath(Candidate &C,
     Candidate *NextRoot = Root->Basis;
     if (C.Base == NextRoot->Base && C.StrideSCEV == NextRoot->StrideSCEV &&
         isSimilar(C, *NextRoot, Candidate::IndexDelta)) {
-      ConstantInt *CI = cast<ConstantInt>(getIndexDelta(C, *NextRoot));
+      ConstantInt *CI =
+          cast<ConstantInt>(getDelta(C, *NextRoot, Candidate::IndexDelta));
       if (CI->isZero() || CI->isOne() || isa<SCEVConstant>(C.StrideSCEV)) {
         Root = NextRoot;
         NewKind = Candidate::IndexDelta;
@@ -1276,8 +1291,7 @@ bool StraightLineStrengthReduce::runOnFunction(Function &F) {
   // Build the dependency graph and sort candidate instructions from dependency
   // roots to leaves
   for (auto &C : Candidates) {
-    if (DependencyGraph.find(C.Ins) == DependencyGraph.end())
-      DependencyGraph[C.Ins] = {};
+    DependencyGraph.try_emplace(C.Ins);
     addDependency(C, C.Basis);
   }
   sortCandidateInstructions();

llvm/test/CodeGen/AMDGPU/promote-constOffset-to-imm.ll

@@ -723,6 +723,7 @@ define hidden amdgpu_kernel void @clmem_read(ptr addrspace(1)  %buffer) {
 ; GFX90A-NEXT:    v_add_co_u32_e32 v2, vcc, 0x2800, v1
 ; GFX90A-NEXT:    v_addc_co_u32_e32 v3, vcc, 0, v3, vcc
 ; GFX90A-NEXT:    v_mov_b32_e32 v1, 0x7f
+; GFX90A-NEXT:    v_pk_mov_b32 v[4:5], 0, 0
 ; GFX90A-NEXT:    s_movk_i32 s2, 0xf000
 ; GFX90A-NEXT:    s_movk_i32 s3, 0x1000
 ; GFX90A-NEXT:    s_movk_i32 s4, 0x2000