[DAGCombiner] Be more careful about looking through extends and trunc…

[topperc]

…ates in mergeTruncStores.

Previously we recursively looked through extends and truncates on both SourceValue and WideVal.

SourceValue is the largest source found for each of the stores we are combining. WideVal is the source for the current store.

Previously we could incorrectly look through a (zext (trunc X)) pair and incorrectly believe X to be a good source.

I think we could also look through a zext on one store and a sext on another store and arbitrarily pick one of the extends as the final source.

With this patch we only look through one level of extend or truncate. And we don't look through extends/truncs on both SourceValue and WideVal at the same time.

This may lose some optimization cases, but keeps everything we had tests for.

Fixes #90936.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-selectiondag

Author: Craig Topper (topperc)

Changes

…ates in mergeTruncStores.

Previously we recursively looked through extends and truncates on both SourceValue and WideVal.

SourceValue is the largest source found for each of the stores we are combining. WideVal is the source for the current store.

Previously we could incorrectly look through a (zext (trunc X)) pair and incorrectly believe X to be a good source.

I think we could also look through a zext on one store and a sext on another store and arbitrarily pick one of the extends as the final source.

With this patch we only look through one level of extend or truncate. And we don't look through extends/truncs on both SourceValue and WideVal at the same time.

This may lose some optimization cases, but keeps everything we had tests for.

Fixes #90936.

---

Full diff: https://github.com/llvm/llvm-project/pull/91375.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+11-9)
• (modified) llvm/test/CodeGen/AArch64/pr90936.ll (+6-2)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 05ab6e2e48206..547a571ea24ac 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -8728,15 +8728,16 @@ static std::optional<bool> isBigEndian(const ArrayRef<int64_t> ByteOffsets,
   return BigEndian;
 }
 
+// Look through one layer of truncate or extend.
 static SDValue stripTruncAndExt(SDValue Value) {
   switch (Value.getOpcode()) {
   case ISD::TRUNCATE:
   case ISD::ZERO_EXTEND:
   case ISD::SIGN_EXTEND:
   case ISD::ANY_EXTEND:
-    return stripTruncAndExt(Value.getOperand(0));
+    return Value.getOperand(0);
   }
-  return Value;
+  return SDValue();
 }
 
 /// Match a pattern where a wide type scalar value is stored by several narrow
@@ -8849,16 +8850,17 @@ SDValue DAGCombiner::mergeTruncStores(StoreSDNode *N) {
     }
 
     // Stores must share the same source value with different offsets.
-    // Truncate and extends should be stripped to get the single source value.
     if (!SourceValue)
       SourceValue = WideVal;
-    else if (stripTruncAndExt(SourceValue) != stripTruncAndExt(WideVal))
-      return SDValue();
-    else if (SourceValue.getValueType() != WideVT) {
-      if (WideVal.getValueType() == WideVT ||
-          WideVal.getScalarValueSizeInBits() >
-              SourceValue.getScalarValueSizeInBits())
+    else if (SourceValue != WideVal) {
+      // Truncate and extends can be stripped to see if the values are related.
+      if (stripTruncAndExt(SourceValue) != WideVal &&
+          stripTruncAndExt(WideVal) != SourceValue)
+        return SDValue();
+
+      if (WideVal.getScalarValueSizeInBits() > SourceValue.getScalarValueSizeInBits())
         SourceValue = WideVal;
+
       // Give up if the source value type is smaller than the store size.
       if (SourceValue.getScalarValueSizeInBits() < WideVT.getScalarSizeInBits())
         return SDValue();
diff --git a/llvm/test/CodeGen/AArch64/pr90936.ll b/llvm/test/CodeGen/AArch64/pr90936.ll
index cd816cdbf7351..3ed8468b37f4e 100644
--- a/llvm/test/CodeGen/AArch64/pr90936.ll
+++ b/llvm/test/CodeGen/AArch64/pr90936.ll
@@ -22,8 +22,12 @@ bb:
 define void @g(i32 %arg, ptr %arg1) {
 ; CHECK-LABEL: g:
 ; CHECK:       // %bb.0: // %bb
-; CHECK-NEXT:    and w8, w0, #0xff
-; CHECK-NEXT:    str w8, [x1]
+; CHECK-NEXT:    lsr w8, w0, #8
+; CHECK-NEXT:    lsr w9, w0, #16
+; CHECK-NEXT:    strb w0, [x1]
+; CHECK-NEXT:    strb wzr, [x1, #3]
+; CHECK-NEXT:    strb w8, [x1, #1]
+; CHECK-NEXT:    strb w9, [x1, #2]
 ; CHECK-NEXT:    ret
 bb:
   %i = trunc i32 %arg to i8

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[arsenm]

It probably isn't helpful to look through multiple levels of extends. They would most likely have combined away anyway. Maybe this should be using ByteProvider?