[DAG] Fix constant store folding to handle non-byte sizes.

[niravd]

Summary

Avoid crashes from zero-byte values due to sub-byte store sizes.

[courbet]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15436 ↗

(On Diff #188177)

This is overly optimistic: for example, this will return true for STBitSize = 1 and ChainBitSize= 2 if both are at offset 0

What about passing size in bits to BaseIndexOffset::contains() and letting it do the computation in bits:

bool BaseIndexOffset::contains(int64_t BitSize, const BaseIndexOffset &Other,
                               int64_t OtherBitSize, const SelectionDAG &DAG,
                               int64_t &Offset) const {
  if (!equalBaseIndex(Other, DAG, Offset))
    return false;
  if (Offset >= 0) {
    // Other is after *this:
    // [-------*this---------]
    //            [---Other--]
    // ==Offset==>
    return 8 * Offset + OtherBitSize <= BitSize;
  }
  // Other starts strictly before *this, it cannot be fully contained.
  //    [-------*this---------]
  // [--Other--]
  return false;
}

[niravd]

Convert to bit-level sizing for contains. I'm leaving equalBaseIndex as byte-level for now, but a followup NFC patch is probably warranted.

[courbet]

llvm/include/llvm/CodeGen/SelectionDAGAddressAnalysis.h

63 ↗

(On Diff #188218)

please update the comment.

70 ↗

(On Diff #188218)

BitOffset ?

[uabelho]

Nice! I ran a new bunch of llvm-stress tests with this patch and now I don't see the crashes anymore.
Please push.

[uabelho]

Added a few nits about too long lines.

llvm/include/llvm/CodeGen/SelectionDAGAddressAnalysis.h

65 ↗

(On Diff #188218)

long line

69 ↗

(On Diff #188218)

long line

llvm/lib/CodeGen/SelectionDAG/SelectionDAGAddressAnalysis.cpp

139 ↗

(On Diff #188218)

long line

llvm/trunk/include/llvm/CodeGen/SelectionDAGAddressAnalysis.h

public:
   bool equalBaseIndex(const BaseIndexOffset &Other,
                       const SelectionDAG &DAG) const {
     int64_t Off;
     return equalBaseIndex(Other, DAG, Off);
   }
 
   // Returns true if `Other` (with size `OtherSize`) can be proven to be fully
   // contained in `*this` (with size `Size`).
-  bool contains(int64_t Size, const BaseIndexOffset &Other, int64_t OtherSize,
-                const SelectionDAG &DAG) const {
-    int64_t Offset;
-    return contains(Size, Other, OtherSize, DAG, Offset);
+  bool contains(const SelectionDAG &DAG, int64_t BitSize,
+                const BaseIndexOffset &Other, int64_t OtherBitSize,
+                int64_t &BitOffset) const;
+
+  bool contains(const SelectionDAG &DAG, int64_t BitSize,
+                const BaseIndexOffset &Other, int64_t OtherBitSize) const {
+    int64_t BitOffset;
+    return contains(DAG, BitSize, Other, OtherBitSize, BitOffset);
   }
 
-  bool contains(int64_t Size, const BaseIndexOffset &Other, int64_t OtherSize,
-                const SelectionDAG &DAG, int64_t &Offset) const;
-
   // Returns true `BasePtr0` and `BasePtr1` can be proven to alias/not alias, in
   // which case `IsAlias` is set to true/false.
   static bool computeAliasing(const BaseIndexOffset &BasePtr0,
                               const int64_t NumBytes0,
                               const BaseIndexOffset &BasePtr1,
                               const int64_t NumBytes1, const SelectionDAG &DAG,
                               bool &IsAlias);
 

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (ST->isUnindexed() && !ST->isVolatile() && ST1->isUnindexed() &&
         // same location, then the store is dead/noop.
         return Chain;
       }
 
       if (OptLevel != CodeGenOpt::None && ST1->hasOneUse() &&
           !ST1->getBasePtr().isUndef()) {
         const BaseIndexOffset STBase = BaseIndexOffset::match(ST, DAG);
         const BaseIndexOffset ChainBase = BaseIndexOffset::match(ST1, DAG);
-        unsigned STByteSize = ST->getMemoryVT().getSizeInBits() / 8;
-        unsigned ChainByteSize = ST1->getMemoryVT().getSizeInBits() / 8;
+        unsigned STBitSize = ST->getMemoryVT().getSizeInBits();
+        unsigned ChainBitSize = ST1->getMemoryVT().getSizeInBits();
         // If this is a store who's preceding store to a subset of the current
         // location and no one other node is chained to that store we can
         // effectively drop the store. Do not remove stores to undef as they may
         // be used as data sinks.
-        if (STBase.contains(STByteSize, ChainBase, ChainByteSize, DAG)) {
+        if (STBase.contains(DAG, STBitSize, ChainBase, ChainBitSize)) {
           CombineTo(ST1, ST1->getChain());
           return SDValue();
         }
 
         // If ST stores to a subset of preceding store's write set, we may be
         // able to fold ST's value into the preceding stored value. As we know
         // the other uses of ST1's chain are unconcerned with ST, this folding
         // will not affect those nodes.
-        int64_t Offset;
-        if (ChainBase.contains(ChainByteSize, STBase, STByteSize, DAG,
-                               Offset)) {
+        int64_t BitOffset;
+        if (ChainBase.contains(DAG, ChainBitSize, STBase, STBitSize,
+                               BitOffset)) {
           SDValue ChainValue = ST1->getValue();
           if (auto *C1 = dyn_cast<ConstantSDNode>(ChainValue)) {
             if (auto *C = dyn_cast<ConstantSDNode>(Value)) {
               APInt Val = C1->getAPIntValue();
-              APInt InsertVal = C->getAPIntValue().zextOrTrunc(STByteSize * 8);
+              APInt InsertVal = C->getAPIntValue().zextOrTrunc(STBitSize);
               // FIXME: Handle Big-endian mode.
               if (!DAG.getDataLayout().isBigEndian()) {
-                Val.insertBits(InsertVal, Offset * 8);
+                Val.insertBits(InsertVal, BitOffset);
                 SDValue NewSDVal =
                     DAG.getConstant(Val, SDLoc(C), ChainValue.getValueType(),
                                     C1->isTargetOpcode(), C1->isOpaque());
                 SDNode *NewST1 = DAG.UpdateNodeOperands(
                     ST1, ST1->getChain(), NewSDVal, ST1->getOperand(2),
                     ST1->getOperand(3));
                 return CombineTo(ST, SDValue(NewST1, 0));
               }

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAGAddressAnalysis.cpp

if ((BasePtr0.getIndex() == BasePtr1.getIndex() || (IsFI0 != IsFI1) ||
        (IsGV0 != IsGV1) || (IsCV0 != IsCV1)) &&
       (IsFI0 || IsGV0 || IsCV0) && (IsFI1 || IsGV1 || IsCV1)) {
     IsAlias = false;
     return true;
   }
   return false; // Cannot determine whether the pointers alias.
 }
 
-bool BaseIndexOffset::contains(int64_t Size, const BaseIndexOffset &Other,
-                               int64_t OtherSize, const SelectionDAG &DAG,
-                               int64_t &Offset) const {
+bool BaseIndexOffset::contains(const SelectionDAG &DAG, int64_t BitSize,
+                               const BaseIndexOffset &Other,
+                               int64_t OtherBitSize, int64_t &BitOffset) const {
+  int64_t Offset;
   if (!equalBaseIndex(Other, DAG, Offset))
     return false;
   if (Offset >= 0) {
     // Other is after *this:
     // [-------*this---------]
     //            [---Other--]
     // ==Offset==>
-    return Offset + OtherSize <= Size;
+    BitOffset = 8 * Offset;
+    return BitOffset + OtherBitSize <= BitSize;
   }
   // Other starts strictly before *this, it cannot be fully contained.
   //    [-------*this---------]
   // [--Other--]
   return false;
 }
 
 /// Parses tree in Ptr for base, index, offset addresses.

llvm/trunk/test/CodeGen/X86/constant-combines.ll

entry:
   %6 = trunc i64 %5 to i32
   %7 = bitcast i32 %6 to float
   %8 = fmul float %7, 0.000000e+00
   %9 = bitcast float* %1 to i32*
   store i32 %6, i32* %9, align 4
   store float %8, float* %0, align 4
   ret void
 }
+
+
+define void @bitstore_fold() {
+; CHECK-LABEL: bitstore_fold:
+; CHECK:       # %bb.0: # %BB
+; CHECK-NEXT:    movl $-2, 0
+; CHECK-NEXT:    retq
+BB:
+   store i32 -1, i32* null
+   store i1 false, i1* null
+   ret void
+}