[DAGCombiner] Fix load-store forwarding of indexed loads.

[niravd]

Summary

Handle extra output from index loads in cases where we wish to
forward a load value directly from a preceeding store.

Fixes PR39571.

[peter.smith]

Thanks very much for the quick update. I can confirm that the original test file compiles successfully. I'll take a look at the change tomorrow (need to leave office).

[peter.smith]

This is looking good. I've got one style nit that I'll leave to your discretion. I think it is probably worth tidying up the test case a little more though. I've made a proposal in the comment.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12865 ↗

(On Diff #173171)

Taking up one of the comments from the original review for D49200. The use of auto could be considered a bit excessive here. https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

For example use bool for IsSub, ISD::NodeType for Opc and SDValue for Idx. I don't have a strong opinion about this personally but it might be more consistent with the other code in the file.

llvm/test/CodeGen/ARM/pr39571.ll

1 ↗

(On Diff #173171)

I've made an attempt to strip down the test case a bit. I've also upped the target Arm architecture to a more recent one. It seems like the original test passed with armv7a due to the availability of unaligned access, but it will fail if -mno-unaligned-access is used.

I've changed the variable names to remove the association with the original program.

; RUN: llc < %s -mtriple armv7a-unknown-linux-gnueabi -mattr=+strict-align

; Avoid crash from forwarding indexed-loads back to store.
%struct.anon = type { %struct.ma*, %struct.mb }
%struct.ma = type { i8 }
%struct.mb = type { i8, i8 }
%struct.anon.0 = type { %struct.anon.1 }
%struct.anon.1 = type { %struct.ds }
%struct.ds = type <{ i8, %union.ie }>
%union.ie = type { %struct.ib }
%struct.ib = type { i8, i8, i16 }

@a = common dso_local local_unnamed_addr global %struct.anon* null, align 4
@b = common dso_local local_unnamed_addr global %struct.anon.0 zeroinitializer, align 1

; Function Attrs: norecurse nounwind
define dso_local void @func() local_unnamed_addr {
entry:
  %0 = load %struct.anon*, %struct.anon** @a, align 4
  %ad = getelementptr inbounds %struct.anon, %struct.anon* %0, i32 0, i32 0
  %1 = load %struct.ma*, %struct.ma** %ad, align 4
  %c.sroa.0.0..sroa_idx = getelementptr inbounds %struct.ma, %struct.ma* %1, i32 0, i32 0
  %c.sroa.0.0.copyload = load i8, i8* %c.sroa.0.0..sroa_idx, align 1
  %cb = getelementptr inbounds %struct.anon, %struct.anon* %0, i32 0, i32 1
  %band = getelementptr inbounds %struct.anon, %struct.anon* %0, i32 0, i32 1, i32 1
  store i8 %c.sroa.0.0.copyload, i8* %band, align 4
  store i8 6, i8* getelementptr inbounds (%struct.anon.0, %struct.anon.0* @b, i32 0, i32 0, i32 0, i32 1, i32 0, i32 0), align 1
  store i8 2, i8* getelementptr inbounds (%struct.anon.0, %struct.anon.0* @b, i32 0, i32 0, i32 0, i32 1, i32 0, i32 1), align 1
  %2 = bitcast %struct.mb* %cb to i32*
  %3 = load i32, i32* bitcast (i8* getelementptr inbounds (%struct.anon.0, %struct.anon.0* @b, i32 0, i32 0, i32 0, i32 1, i32 0, i32 0) to i32*), align 1
  store i32 %3, i32* %2, align 1
  ret void
}

[niravd]

Fold in Peter's comments.

[peter.smith]

Thanks for the update. Looks good to me.

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

SDValue DAGCombiner::ForwardStoreValueToDirectLoad(LoadSDNode *LD) {
   if (DAG.getDataLayout().isBigEndian())
     Offset = (STMemType.getStoreSizeInBits() -
               LDMemType.getStoreSizeInBits()) / 8 - Offset;
 
   // Check that the stored value cover all bits that are loaded.
   bool STCoversLD =
       (Offset >= 0) &&
       (Offset * 8 + LDMemType.getSizeInBits() <= STMemType.getSizeInBits());
+
+  auto ReplaceLd = [&](LoadSDNode *LD, SDValue Val, SDValue Chain) -> SDValue {
+    if (LD->isIndexed()) {
+      bool IsSub = (LD->getAddressingMode() == ISD::PRE_DEC ||
+                    LD->getAddressingMode() == ISD::POST_DEC);
+      unsigned Opc = IsSub ? ISD::SUB : ISD::ADD;
+      SDValue Idx = DAG.getNode(Opc, SDLoc(LD), LD->getOperand(1).getValueType(),
+                             LD->getOperand(1), LD->getOperand(2));
+      SDValue Ops[] = {Val, Idx, Chain};
+      return CombineTo(LD, Ops, 3);
+    }
+    return CombineTo(LD, Val, Chain);
+  };
+
   if (!STCoversLD)
     return SDValue();
 
   // Memory as copy space (potentially masked).
   if (Offset == 0 && LDType == STType && STMemType == LDMemType) {
     // Simple case: Direct non-truncating forwarding
     if (LDType.getSizeInBits() == LDMemType.getSizeInBits())
-      return CombineTo(LD, ST->getValue(), Chain);
+      return ReplaceLd(LD, ST->getValue(), Chain);
     // Can we model the truncate and extension with an and mask?
     if (STType.isInteger() && LDMemType.isInteger() && !STType.isVector() &&
         !LDMemType.isVector() && LD->getExtensionType() != ISD::SEXTLOAD) {
       // Mask to size of LDMemType
       auto Mask =
           DAG.getConstant(APInt::getLowBitsSet(STType.getSizeInBits(),
                                                STMemType.getSizeInBits()),
                           SDLoc(ST), STType);
       auto Val = DAG.getNode(ISD::AND, SDLoc(LD), LDType, ST->getValue(), Mask);
-      return CombineTo(LD, Val, Chain);
+      return ReplaceLd(LD, Val, Chain);
     }
   }
 
   // TODO: Deal with nonzero offset.
   if (LD->getBasePtr().isUndef() || Offset != 0)
     return SDValue();
   // Model necessary truncations / extenstions.
   SDValue Val;
   // Truncate Value To Stored Memory Size.
   do {
     if (!getTruncatedStoreValue(ST, Val))
       continue;
     if (!isTypeLegal(LDMemType))
       continue;
     if (STMemType != LDMemType) {
       // TODO: Support vectors? This requires extract_subvector/bitcast.
       if (!STMemType.isVector() && !LDMemType.isVector() &&
           STMemType.isInteger() && LDMemType.isInteger())
         Val = DAG.getNode(ISD::TRUNCATE, SDLoc(LD), LDMemType, Val);
       else
         continue;
     }
     if (!extendLoadedValueToExtension(LD, Val))
       continue;
-    return CombineTo(LD, Val, Chain);
+    return ReplaceLd(LD, Val, Chain);
   } while (false);
 
   // On failure, cleanup dead nodes we may have created.
   if (Val->use_empty())
     deleteAndRecombine(Val.getNode());
   return SDValue();
 }
 

llvm/trunk/test/CodeGen/ARM/pr39571.ll

+; RUN: llc < %s -mtriple armv4t-unknown-linux-gnueabi -mattr=+strict-align
+
+; Avoid crash from forwarding indexed-loads back to store.
+%struct.anon = type { %struct.ma*, %struct.mb }
+%struct.ma = type { i8 }
+%struct.mb = type { i8, i8 }
+%struct.anon.0 = type { %struct.anon.1 }
+%struct.anon.1 = type { %struct.ds }
+%struct.ds = type <{ i8, %union.ie }>
+%union.ie = type { %struct.ib }
+%struct.ib = type { i8, i8, i16 }
+
+@a = common dso_local local_unnamed_addr global %struct.anon* null, align 4
+@b = common dso_local local_unnamed_addr global %struct.anon.0 zeroinitializer, align 1
+
+; Function Attrs: norecurse nounwind
+define dso_local void @func() local_unnamed_addr {
+entry:
+  %0 = load %struct.anon*, %struct.anon** @a, align 4
+  %ad = getelementptr inbounds %struct.anon, %struct.anon* %0, i32 0, i32 0
+  %1 = load %struct.ma*, %struct.ma** %ad, align 4
+  %c.sroa.0.0..sroa_idx = getelementptr inbounds %struct.ma, %struct.ma* %1, i32 0, i32 0
+  %c.sroa.0.0.copyload = load i8, i8* %c.sroa.0.0..sroa_idx, align 1
+  %cb = getelementptr inbounds %struct.anon, %struct.anon* %0, i32 0, i32 1
+  %band = getelementptr inbounds %struct.anon, %struct.anon* %0, i32 0, i32 1, i32 1
+  store i8 %c.sroa.0.0.copyload, i8* %band, align 4
+  store i8 6, i8* getelementptr inbounds (%struct.anon.0, %struct.anon.0* @b, i32 0, i32 0, i32 0, i32 1, i32 0, i32 0), align 1
+  store i8 2, i8* getelementptr inbounds (%struct.anon.0, %struct.anon.0* @b, i32 0, i32 0, i32 0, i32 1, i32 0, i32 1), align 1
+  %2 = bitcast %struct.mb* %cb to i32*
+  %3 = load i32, i32* bitcast (i8* getelementptr inbounds (%struct.anon.0, %struct.anon.0* @b, i32 0, i32 0, i32 0, i32 1, i32 0, i32 0) to i32*), align 1
+  store i32 %3, i32* %2, align 1
+  ret void
+}