[llvm] ef84452 - [DAGCombiner] Be more careful about looking through extends and truncates in mergeTruncStores. (#91375)

[via llvm-commits]

Author: Craig Topper
Date: 2024-05-07T21:17:50-07:00
New Revision: ef84452571b8e8f4a38a173e6adf6a5ecbbde97e

URL: https://github.com/llvm/llvm-project/commit/ef84452571b8e8f4a38a173e6adf6a5ecbbde97e
DIFF: https://github.com/llvm/llvm-project/commit/ef84452571b8e8f4a38a173e6adf6a5ecbbde97e.diff

LOG: [DAGCombiner] Be more careful about looking through extends and truncates in mergeTruncStores. (#91375)

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.

Added: 
    

Modified: 
    llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
    llvm/test/CodeGen/AArch64/pr90936.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 05ab6e2e4820..e835bd950a7b 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,18 @@ SDValue DAGCombiner::mergeTruncStores(StoreSDNode *N) {
     }
 
     // Stores must share the same source value with 
diff erent 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 cd816cdbf735..3ed8468b37f4 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