[AArch64] Fix stores of zero values

[niravd]

Summary

Now that DAGCombiner may merge multiple non-consecutive stores in a single pass, the simple one-use analysis for when to split a store
into stores of XZR for pairing purposes is no longer sufficient. Improve check to capture these cases.

[niravd]

Ping.

[niravd]

Ping.

[gberry]

Hi Nirav,

I did a test run with this patch to see what the extent of the changes are and it seems to get into an infinite loop during ISel in quite a few cases. For example, when compiling a test from the llvm test-suite like so:
clang -DSMALL_PROBLEM_SIZE=1 -DNDEBUG -O3 -c MultiSource/Benchmarks/TSVC/Packing-dbl/tsc.c

Also, could you fix the summary of this patch? The first sentence seems incomplete.

[niravd]

Fix infinite loop

[niravd]

In D33518#778381, @niravd wrote:

Fix infinite loop

[niravd]

There are some simplifications to BaseOffsetDifference16 which can be done once cleanups, but I'd like to get this in before that.

[niravd]

Rewrite to use BaseIndexOffset

[niravd]

Ping.

[gberry]

This doesn't even compile for me:

FAILED: /local/mnt/gberry/infra/tools/clang-tools/clang-x86-494-x++ -DGTEST_HAS_RTTI=0 -DLLVM_BUILD_GLOBAL_ISEL -D_DEBUG -D_GNU_SOURCE -DSTDC_CONSTANT_MACROS -DSTDC_FORMAT_MACROS -D__\
STDC_LIMIT_MACROS -Ilib/Target/AArch64 -I/local/mnt/gberry/src/llvm/lib/Target/AArch64 -Iinclude -I/local/mnt/gberry/src/llvm/include -O2 -fPIC -fPIC -fvisibility-inlines-hidden -Werror -W\
error=date-time -std=c++11 -Wall -W -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wmissing-field-initializers -pedantic -Wno-long-long -Wcovered-switch-default -Wnon-virtual-dtor -Wdel\
ete-non-virtual-dtor -Wstring-conversion -fcolor-diagnostics -ffunction-sections -fdata-sections -O3 -UNDEBUG -fno-exceptions -fno-rtti -MMD -MT lib/Target/AArch64/CMakeFiles/LLVMAArch6\
4CodeGen.dir/AArch64ISelLowering.cpp.o -MF lib/Target/AArch64/CMakeFiles/LLVMAArch64CodeGen.dir/AArch64ISelLowering.cpp.o.d -o lib/Target/AArch64/CMakeFiles/LLVMAArch64CodeGen.dir/AArch64IS\
elLowering.cpp.o -c /local/mnt/gberry/src/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
/local/mnt/gberry/src/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp:9298:52: error: too few arguments to function call, expected 2, have 1

auto AddrDecomp = BaseIndexOffset::match(Addr);
                  ~~~~~~~~~~~~~~~~~~~~~~     ^

/local/mnt/gberry/src/llvm/include/llvm/CodeGen/SelectionDAGAddressAnalysis.h:60:3: note: 'match' declared here

static BaseIndexOffset match(SDValue Ptr, const SelectionDAG &DAG);
^

/local/mnt/gberry/src/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp:9300:64: error: too few arguments to function call, expected 2, have 1

auto OtherAddrDecomp = BaseIndexOffset::match(OtherAddr);
                       ~~~~~~~~~~~~~~~~~~~~~~          ^

/local/mnt/gberry/src/llvm/include/llvm/CodeGen/SelectionDAGAddressAnalysis.h:60:3: note: 'match' declared here

static BaseIndexOffset match(SDValue Ptr, const SelectionDAG &DAG);
^

[niravd]

Rebase to current.

[gberry]

I noticed a few code size regressions looking through the diffs of generated code for this change which were caused by:

• the offsets of the stores not having the right alignment to allow for stp x opcodes to be formed
• the offsets of the stores being too large for stp x, but not too large for stp q
• the offsets of the stores being so large that adds are needed to compute the address for str x, but not so large that adds are needed to compute the address for str q

lib/Target/AArch64/AArch64ISelLowering.cpp

9335

Would it not be better to store the BaseIndexOffset in STAddrs instead of re-computing it for every iteration of the StVal->uses() loop?

Also, do you have a use case/benchmark where this helps performance/code size significantly to justify the n^2 checking here?

[niravd]

In D33518#813376, @gberry wrote:

I noticed a few code size regressions looking through the diffs of generated code for this change which were caused by:

• the offsets of the stores not having the right alignment to allow for stp x opcodes to be formed
• the offsets of the stores being too large for stp x, but not too large for stp q
• the offsets of the stores being so large that adds are needed to compute the address for str x, but not so large that adds are needed to compute the address for str q

Hmm. Dealing with offsets makes this harder. It seems like we could achieve the same rough goal in a more straight forward-manner post-DAG after load-store pairing has happened.

lib/Target/AArch64/AArch64ISelLowering.cpp

while (--NumVecElts) {
     NewST1 = DAG.getStore(NewST1.getValue(0), DL, SplatVal, OffsetPtr,
                           PtrInfo.getWithOffset(Offset), Alignment,
                           St.getMemOperand()->getFlags());
     Offset += EltOffset;
   }
   return NewST1;
 }
 
+// Check if addresses X and Y differ by 16 in one direction or another.
+static bool BaseOffsetDifference16(SDValue X, SDValue Y) {
+  int64_t PartialOffset = 0;
+  if (X.getOpcode() == ISD::SIGN_EXTEND)
+    X = X->getOperand(0);
+  if (Y.getOpcode() == ISD::SIGN_EXTEND)
+    Y = Y->getOperand(0);
+  if (X.getOpcode() == ISD::ADD && isa<ConstantSDNode>(X->getOperand(1))) {
+    PartialOffset -= cast<ConstantSDNode>(X->getOperand(1))->getSExtValue();
+    X = X->getOperand(0);
+  }
+  if (Y.getOpcode() == ISD::ADD && isa<ConstantSDNode>(Y->getOperand(1))) {
+    PartialOffset += cast<ConstantSDNode>(Y->getOperand(1))->getSExtValue();
+    Y = Y->getOperand(0);
+  }
+  while (X.getOpcode() == ISD::ADD && Y.getOpcode() == ISD::ADD) {
+    if (X->getOperand(0) == Y->getOperand(0)) {
+      X = X->getOperand(1);
+      Y = Y->getOperand(1);
+    } else if (X->getOperand(1) == Y->getOperand(1)) {
+      X = X->getOperand(0);
+      Y = Y->getOperand(0);
+    }
+  }
+  if (isa<ConstantSDNode>(X) && isa<ConstantSDNode>(Y)) {
+    PartialOffset -= cast<ConstantSDNode>(X)->getSExtValue();
+    PartialOffset += cast<ConstantSDNode>(Y)->getSExtValue();
+    return (PartialOffset == 16) || (PartialOffset == -16);
+  }
+  if (X == Y)
+    return (PartialOffset == 16) || (PartialOffset == -16);
+  return false;
+}
+
 /// Replace a splat of zeros to a vector store by scalar stores of WZR/XZR.  The
 /// load store optimizer pass will merge them to store pair stores.  This should
 /// be better than a movi to create the vector zero followed by a vector store
 /// if the zero constant is not re-used, since one instructions and one register
 /// live range will be removed.
 ///
 /// For example, the final generated code should be:
 ///
if (!(((NumVecElts == 2 || NumVecElts == 3) &&
          VT.getVectorElementType().getSizeInBits() == 64) ||
         ((NumVecElts == 2 || NumVecElts == 3 || NumVecElts == 4) &&
          VT.getVectorElementType().getSizeInBits() == 32)))
     return SDValue();
 
   if (StVal.getOpcode() != ISD::BUILD_VECTOR)
     return SDValue();
 
-  // If the zero constant has more than one use then the vector store could be
-  // better since the constant mov will be amortized and stp q instructions
-  // should be able to be formed.
-  if (!StVal.hasOneUse())
+  // If the zero constant has more than one use then the vector store
+  // could be better since the constant mov will be amortized and stp
+  // q instructions should be able to be formed. The exception to this
+  // is when all uses of StVal are non-consecutive ST16 instances. In
+  // this case, it is better to replace each instance with stp xzr, xzr.
+
+  if (!StVal.hasOneUse()) {
+    if (VT.getVectorElementType().getSizeInBits() != 64)
+      return SDValue();
+    SmallVector<SDValue, 8> STAddrs;
+    for (auto *U : StVal->uses()) {
+      if ((U->getOpcode() != ISD::STORE) ||
+          (U->getOperand(1).getValueType().getSizeInBits() != 16 * 8))
+        return SDValue();
+      SDValue Addr = U->getOperand(2);
+      bool HasConsecutive = false;
+      for (SDValue &OtherAddr : STAddrs)
+        if (BaseOffsetDifference16(Addr, OtherAddr)) {
+          HasConsecutive = true;
+          break;
+        }
+      if (HasConsecutive)
         return SDValue();
+      STAddrs.push_back(Addr);
+    }
+  }
 
   // If the immediate offset of the address operand is too large for the stp
   // instruction, then bail out.
   if (DAG.isBaseWithConstantOffset(St.getBasePtr())) {
     int64_t Offset = St.getBasePtr()->getConstantOperandVal(1);
     if (Offset < -512 || Offset > 504)
       return SDValue();
   }

test/CodeGen/AArch64/ldst-opt.ll

 ; CHECK: // %entry
 ; CHECK-NEXT: movi v[[REG:[0-9]]].2d, #0000000000000000
 ; CHECK-NEXT: stp q[[REG]], q[[REG]], [x{{[0-9]+}}]
 ; CHECK-NEXT: ret
 entry:
   store <4 x double> zeroinitializer, <4 x double>* %p
   ret void
 }
+
+; Verify that non-consecutive merges do not generate q0
+define void @merge_multiple_128bit_stores(i64* %p) {
+; CHECK-LABEL: merge_multiple_128bit_stores
+; CHECK: // %entry
+; CHECK-NEXT: stp xzr, xzr, [x{{[0-9]+}}]
+; CHECK-NEXT: stp xzr, xzr, [x{{[0-9]+}}, #24]
+; CHECK-NEXT: stp xzr, xzr, [x{{[0-9]+}}, #48]
+; CHECK-NEXT: ret
+entry:
+  store i64 0, i64* %p
+  %p1 = getelementptr i64, i64* %p, i64 1
+  store i64 0, i64* %p1
+  %p3 = getelementptr i64, i64* %p, i64 3
+  store i64 0, i64* %p3
+  %p4 = getelementptr i64, i64* %p, i64 4
+  store i64 0, i64* %p4
+  %p6 = getelementptr i64, i64* %p, i64 6
+  store i64 0, i64* %p6
+  %p7 = getelementptr i64, i64* %p, i64 7
+  store i64 0, i64* %p7
+  ret void
+}
+
+; Verify that large stores generate stp q
+define void @merge_multiple_128bit_stores_consec(i64* %p) {
+; CHECK-LABEL: merge_multiple_128bit_stores_consec
+; CHECK: // %entry
+; NOSTRICTALIGN-NEXT: movi v[[REG:[0-9]]].2d, #0000000000000000
+; NOSTRICTALIGN-NEXT: stp q[[REG]], q[[REG]], [x{{[0-9]+}}]
+; NOSTRICTALIGN-NEXT: stp q[[REG]], q[[REG]], [x{{[0-9]+}}, #32]
+; STRICTALIGN-NEXT: stp  xzr, xzr, [x0]
+; STRICTALIGN-NEXT: stp  xzr, xzr, [x0, #16]
+; STRICTALIGN-NEXT: stp  xzr, xzr, [x0, #32]
+; STRICTALIGN-NEXT: stp  xzr, xzr, [x0, #48]
+; CHECK-NEXT: ret
+entry:
+  store i64 0, i64* %p
+  %p1 = getelementptr i64, i64* %p, i64 1
+  store i64 0, i64* %p1
+  %p2 = getelementptr i64, i64* %p, i64 2
+  store i64 0, i64* %p2
+  %p3 = getelementptr i64, i64* %p, i64 3
+  store i64 0, i64* %p3
+  %p4 = getelementptr i64, i64* %p, i64 4
+  store i64 0, i64* %p4
+  %p5 = getelementptr i64, i64* %p, i64 5
+  store i64 0, i64* %p5
+  %p6 = getelementptr i64, i64* %p, i64 6
+  store i64 0, i64* %p6
+  %p7 = getelementptr i64, i64* %p, i64 7
+  store i64 0, i64* %p7
+  ret void
+}
+
+