[llvm] [DAGCombiner] Be more careful about looking through extends and trunc… (PR #91375)

[Craig Topper via llvm-commits]

https://github.com/topperc updated https://github.com/llvm/llvm-project/pull/91375

>From f22ab21ee49da71e7a42d8e575cf856e5e62001a Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Tue, 7 May 2024 11:14:52 -0700
Subject: [PATCH 1/2] [DAGCombiner] Be more careful about looking through
 extends and truncates 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.
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 20 ++++++++++---------
 llvm/test/CodeGen/AArch64/pr90936.ll          |  8 ++++++--
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 05ab6e2e48206f..547a571ea24ac8 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 cd816cdbf73512..3ed8468b37f4e5 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

>From 3f49f964449cdf37f5dc818e27888df905ed1c50 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Tue, 7 May 2024 11:26:20 -0700
Subject: [PATCH 2/2] fixup! clang-format

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 547a571ea24ac8..e835bd950a7be4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -8858,7 +8858,8 @@ SDValue DAGCombiner::mergeTruncStores(StoreSDNode *N) {
           stripTruncAndExt(WideVal) != SourceValue)
         return SDValue();
 
-      if (WideVal.getScalarValueSizeInBits() > SourceValue.getScalarValueSizeInBits())
+      if (WideVal.getScalarValueSizeInBits() >
+          SourceValue.getScalarValueSizeInBits())
         SourceValue = WideVal;
 
       // Give up if the source value type is smaller than the store size.