[DivRemPairs] Avoid RAUW pitfalls (PR42823)

[lebedev.ri]

Summary

DivRemPairs internally creates two maps:

• {sign, divident, divisor} -> div instruction
• {sign, divident, divisor} -> rem instruction

Then it iterates over rem map, and looks if there is an entry
in div map with the same key. Then depending on some internal logic
it may RAUW rem instruction with something else.

But if that rem instruction is an input to other div/rem,
then it was used as a key in these maps, so the old value (used in key)
is now dandling, because RAUW didn't update those maps.
And we can't even RAUW map keys in general, there's ValueMap,
but we don't have a single Value as key...

The bug was discovered via D65298, and the test there exists.
Now, i'm not sure how to expose this issue in trunk.
The bug is clearly there if i change the map keys to be AssertingVH/PoisoningVH,
but i guess this didn't miscompiled anything thus far?
I really don't think this is benin without that patch.

The fix is actually rather straight-forward - instead of trying to somehow
shoe-horn ValueMap here (doesn't fit, key isn't just Value), or writing a new
ValueMap with key being a struct of Values, we can just have an intermediate
data structure - a vector, each entry containing matching Div, Rem pair,
and pre-filling it before doing any modifications.
This way we won't need to query map after doing RAUW, so no bug is possible.

[spatel]

Sorry about the bug - this is why I suggested in the previous review that someone with a better grasp of C++ / data structures have a look. :)
In the original code, if we delay the RAUW by inserting instructions in a ToBeKilled worklist in the loop (and only doing the erasing as a post-processing step) is there still a bug?

llvm/lib/Transforms/Scalar/DivRemPairs.cpp

62

spelling: getDivident --> getDividend

[lebedev.ri]

Updated - fix typo, no need for TrackingVH, simply using worklist was sufficient to fix it.

[spatel]

LGTM - should the test from https://bugs.llvm.org/show_bug.cgi?id=42823 be included with this patch even if does not fail yet without D65298?

llvm/lib/Transforms/Scalar/DivRemPairs.cpp

209–213

Add to this code comment to explain why we need to preserve the Sub in RemInst?

[lebedev.ri]

In D65451#1608076, @spatel wrote:

LGTM

Thank you for the review.

should the test from https://bugs.llvm.org/show_bug.cgi?id=42823 be included with this patch even if does not fail yet without D65298?

I precommitted them already.

llvm/include/llvm/Transforms/Utils/BypassSlowDivision.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_BYPASSSLOWDIVISION_H
 #define LLVM_TRANSFORMS_UTILS_BYPASSSLOWDIVISION_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/IR/ValueHandle.h"
 #include <cstdint>
 
 namespace llvm {
 
 class BasicBlock;
 class Value;
 
 struct DivRemMapKey {
   bool SignedOp;
-  Value *Dividend;
-  Value *Divisor;
+  AssertingVH<Value> Dividend;
+  AssertingVH<Value> Divisor;
 
   DivRemMapKey(bool InSignedOp, Value *InDividend, Value *InDivisor)
       : SignedOp(InSignedOp), Dividend(InDividend), Divisor(InDivisor) {}
 };
 
 template <> struct DenseMapInfo<DivRemMapKey> {
   static bool isEqual(const DivRemMapKey &Val1, const DivRemMapKey &Val2) {
     return Val1.SignedOp == Val2.SignedOp && Val1.Dividend == Val2.Dividend &&
            Val1.Divisor == Val2.Divisor;
   }
 
   static DivRemMapKey getEmptyKey() {
     return DivRemMapKey(false, nullptr, nullptr);
   }
 
   static DivRemMapKey getTombstoneKey() {
     return DivRemMapKey(true, nullptr, nullptr);
   }
 
   static unsigned getHashValue(const DivRemMapKey &Val) {
-    return (unsigned)(reinterpret_cast<uintptr_t>(Val.Dividend) ^
-                      reinterpret_cast<uintptr_t>(Val.Divisor)) ^
+    return (unsigned)(reinterpret_cast<uintptr_t>(
+                          static_cast<Value *>(Val.Dividend)) ^
+                      reinterpret_cast<uintptr_t>(
+                          static_cast<Value *>(Val.Divisor))) ^
            (unsigned)Val.SignedOp;
   }
 };
 
 /// This optimization identifies DIV instructions in a BB that can be
 /// profitably bypassed and carried out with a shorter, faster divide.
 ///
 /// This optimization may add basic blocks immediately after BB; for obvious
 /// reasons, you shouldn't pass those blocks to bypassSlowDivision.
 bool bypassSlowDivision(
     BasicBlock *BB, const DenseMap<unsigned int, unsigned int> &BypassWidth);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_BYPASSSLOWDIVISION_H

llvm/lib/Transforms/Scalar/DivRemPairs.cpp

 
 #define DEBUG_TYPE "div-rem-pairs"
 STATISTIC(NumPairs, "Number of div/rem pairs");
 STATISTIC(NumHoisted, "Number of instructions hoisted");
 STATISTIC(NumDecomposed, "Number of instructions decomposed");
 DEBUG_COUNTER(DRPCounter, "div-rem-pairs-transform",
               "Controls transformations in div-rem-pairs pass");
 
-/// Find matching pairs of integer div/rem ops (they have the same numerator,
-/// denominator, and signedness). If they exist in different basic blocks, bring
-/// them together by hoisting or replace the common division operation that is
-/// implicit in the remainder:
-/// X % Y <--> X - ((X / Y) * Y).
-///
-/// We can largely ignore the normal safety and cost constraints on speculation
-/// of these ops when we find a matching pair. This is because we are already
-/// guaranteed that any exceptions and most cost are already incurred by the
-/// first member of the pair.
-///
-/// Note: This transform could be an oddball enhancement to EarlyCSE, GVN, or
-/// SimplifyCFG, but it's split off on its own because it's different enough
-/// that it doesn't quite match the stated objectives of those passes.
-static bool optimizeDivRem(Function &F, const TargetTransformInfo &TTI,
-                           const DominatorTree &DT) {
-  bool Changed = false;
+/// A thin wrapper to store two values that we matched as div-rem pair.
+/// We want this extra indirection to avoid dealing with RAUW'ing the map keys.
+struct DivRemPairWorklistEntry {
+  /// The actual udiv/sdiv instruction. Source of truth.
+  TrackingVH<Instruction> DivInst;
+
+  /// The instruction that we have matched as a remainder instruction.
+  /// Should only be used as Value, don't introspect it.
+  TrackingVH<Instruction> RemInst;
+
+  DivRemPairWorklistEntry(Instruction *DivInst_, Instruction *RemInst_)
+      : DivInst(DivInst_), RemInst(RemInst_) {
+    assert((DivInst->getOpcode() == Instruction::UDiv ||
+            DivInst->getOpcode() == Instruction::SDiv) &&
+           "Not a division.");
+    assert(DivInst->getType() == RemInst->getType() && "Types should match.");
+    // We can't check anything else about remainder instruction,
+    // it's not strictly required to be a urem/srem.
+  }
+
+  /// The type for this pair, identical for both the div and rem.
+  Type *getType() const { return DivInst->getType(); }
+
+  /// Is this pair signed or unsigned?
+  bool isSigned() const { return DivInst->getOpcode() == Instruction::SDiv; }
+
+  /// In this pair, what are the divident and divisor?
+  Value *getDivident() const { return DivInst->getOperand(0); }
+  Value *getDivisor() const { return DivInst->getOperand(1); }
+};
+using DivRemWorklistTy = SmallVector<DivRemPairWorklistEntry, 4>;
 
+/// Find matching pairs of integer div/rem ops (they have the same numerator,
+/// denominator, and signedness). Place those pairs into a worklist for further
+/// processing. This indirection is needed because we have to use TrackingVH<>
+/// because we will be doing RAUW, and if one of the rem instructions we change
+/// happens to be an input to another div/rem in the maps, we'd have problems.
+static DivRemWorklistTy getWorklist(Function &F) {
   // Insert all divide and remainder instructions into maps keyed by their
   // operands and opcode (signed or unsigned).
   DenseMap<DivRemMapKey, Instruction *> DivMap;
   // Use a MapVector for RemMap so that instructions are moved/inserted in a
   // deterministic order.
   MapVector<DivRemMapKey, Instruction *> RemMap;
   for (auto &BB : F) {
     for (auto &I : BB) {
       if (I.getOpcode() == Instruction::SDiv)
         DivMap[DivRemMapKey(true, I.getOperand(0), I.getOperand(1))] = &I;
       else if (I.getOpcode() == Instruction::UDiv)
         DivMap[DivRemMapKey(false, I.getOperand(0), I.getOperand(1))] = &I;
       else if (I.getOpcode() == Instruction::SRem)
         RemMap[DivRemMapKey(true, I.getOperand(0), I.getOperand(1))] = &I;
       else if (I.getOpcode() == Instruction::URem)
         RemMap[DivRemMapKey(false, I.getOperand(0), I.getOperand(1))] = &I;
     }
   }
 
+  // We'll accumulate the matching pairs of div-rem instructions here.
+  DivRemWorklistTy Worklist;
+
   // We can iterate over either map because we are only looking for matched
   // pairs. Choose remainders for efficiency because they are usually even more
   // rare than division.
   for (auto &RemPair : RemMap) {
     // Find the matching division instruction from the division map.
     Instruction *DivInst = DivMap[RemPair.first];
     if (!DivInst)
       continue;
 
-    // We have a matching pair of div/rem instructions. If one dominates the
-    // other, hoist and/or replace one.
+    // We have a matching pair of div/rem instructions.
     NumPairs++;
     Instruction *RemInst = RemPair.second;
-    bool IsSigned = DivInst->getOpcode() == Instruction::SDiv;
-    bool HasDivRemOp = TTI.hasDivRemOp(DivInst->getType(), IsSigned);
+
+    // Place it in the worklist.
+    Worklist.emplace_back(DivInst, RemInst);
+  }
+
+  return Worklist;
+}
+
+/// Find matching pairs of integer div/rem ops (they have the same numerator,
+/// denominator, and signedness). If they exist in different basic blocks, bring
+/// them together by hoisting or replace the common division operation that is
+/// implicit in the remainder:
+/// X % Y <--> X - ((X / Y) * Y).
+///
+/// We can largely ignore the normal safety and cost constraints on speculation
+/// of these ops when we find a matching pair. This is because we are already
+/// guaranteed that any exceptions and most cost are already incurred by the
+/// first member of the pair.
+///
+/// Note: This transform could be an oddball enhancement to EarlyCSE, GVN, or
+/// SimplifyCFG, but it's split off on its own because it's different enough
+/// that it doesn't quite match the stated objectives of those passes.
+static bool optimizeDivRem(Function &F, const TargetTransformInfo &TTI,
+                           const DominatorTree &DT) {
+  bool Changed = false;
+
+  // Get the matching pairs of div-rem instructions. We want this extra
+  // indirection to avoid dealing with having to RAUW the keys of the maps.
+  const DivRemWorklistTy Worklist = getWorklist(F);
+
+  // Process each entry in the worklist.
+  for (const DivRemPairWorklistEntry &E : Worklist) {
+    bool HasDivRemOp = TTI.hasDivRemOp(E.getType(), E.isSigned());
+
+    Instruction *DivInst = E.DivInst;
+    Instruction *RemInst = E.RemInst;
 
     // If the target supports div+rem and the instructions are in the same block
     // already, there's nothing to do. The backend should handle this. If the
     // target does not support div+rem, then we will decompose the rem.
     if (HasDivRemOp && RemInst->getParent() == DivInst->getParent())
       continue;
 
     bool DivDominates = DT.dominates(DivInst, RemInst);
if (HasDivRemOp) {
         RemInst->moveAfter(DivInst);
       else
         DivInst->moveAfter(RemInst);
       NumHoisted++;
     } else {
       // The target does not have a single div/rem operation. Decompose the
       // remainder calculation as:
       // X % Y --> X - ((X / Y) * Y).
-      Value *X = RemInst->getOperand(0);
-      Value *Y = RemInst->getOperand(1);
+      Value *X = E.getDivident();
+      Value *Y = E.getDivisor();
       Instruction *Mul = BinaryOperator::CreateMul(DivInst, Y);
       Instruction *Sub = BinaryOperator::CreateSub(X, Mul);
 
       // If the remainder dominates, then hoist the division up to that block:
       //
       // bb1:
       //   %rem = srem %x, %y
       // bb2:
if (HasDivRemOp) {
       // If the div and rem are in the same block, we do the same transform,
       // but any code movement would be within the same block.
 
       if (!DivDominates)
         DivInst->moveBefore(RemInst);
       Mul->insertAfter(RemInst);
       Sub->insertAfter(Mul);
 
       // Now kill the explicit remainder. We have replaced it with:
       // (sub X, (mul (div X, Y), Y)
       RemInst->replaceAllUsesWith(Sub);
       RemInst->eraseFromParent();
       NumDecomposed++;
     }
     Changed = true;
   }
 
   return Changed;
 }
 
 // Pass manager boilerplate below here.
struct DivRemPairsLegacyPass : public FunctionPass {
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
     auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
     auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     return optimizeDivRem(F, TTI, DT);
   }
 };
-}
+} // namespace
 
 char DivRemPairsLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(DivRemPairsLegacyPass, "div-rem-pairs",
                       "Hoist/decompose integer division and remainder", false,
                       false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_END(DivRemPairsLegacyPass, "div-rem-pairs",
                     "Hoist/decompose integer division and remainder", false,