[llvm] [SelectionDAG] Prefer to use ATOMIC_LOAD extension type over getExtendForAtomicOps() in computeKnownBits/ComputeNumSignBits. (PR #136600)

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/136600

If an ATOMIC_LOAD has ZEXTLOAD/SEXTLOAD extension type we should trust that over getExtendForAtomicOps().

SystemZ is the only target that usesa setAtomicLoadExtAction and they return ANY_EXTEND from getExtendForAtomicOps(). So I'm not sure there's a way to get a contradiction currently.

Note, type legalization uses getExtendForAtomicOps() when promoting ATOMIC_LOAD so we may not need to check getExtendForAtomicOps() for ATOMIC_LOAD. I have not done much investigating of this.

>From 5ddb89f2fc49e2416d69ee512039eae0e746bf14 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Mon, 21 Apr 2025 12:19:44 -0700
Subject: [PATCH] [SelectionDAG] Prefer to use ATOMIC_LOAD extension type over
 getExtendForAtomicOps() in computeKnownBits/ComputeNumSignBits.

If an ATOMIC_LOAD has ZEXTLOAD/SEXTLOAD extension type we should trust that
over getExtendForAtomicOps().

SystemZ is the only target that usesa setAtomicLoadExtAction and they return
ANY_EXTEND from getExtendForAtomicOps(). So I'm not sure there's a way to get
a contradiction currently.

Note, type legalization uses getExtendForAtomicOps() when promoting ATOMIC_LOAD
so we may not even need to check getExtendForAtomicOps() for ATOMIC_LOAD. I
have not done much investigating of this.
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp | 40 ++++++++++++-------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 09d73633462b6..ba6c5d884d381 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4406,14 +4406,19 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
   case ISD::ATOMIC_LOAD_UMIN:
   case ISD::ATOMIC_LOAD_UMAX:
   case ISD::ATOMIC_LOAD: {
-    unsigned MemBits =
-        cast<AtomicSDNode>(Op)->getMemoryVT().getScalarSizeInBits();
     // If we are looking at the loaded value.
     if (Op.getResNo() == 0) {
-      if (TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
-        Known.Zero.setBitsFrom(MemBits);
-      else if (Op->getOpcode() == ISD::ATOMIC_LOAD &&
-               cast<AtomicSDNode>(Op)->getExtensionType() == ISD::ZEXTLOAD)
+      auto *AT = cast<AtomicSDNode>(Op);
+      unsigned MemBits = AT->getMemoryVT().getScalarSizeInBits();
+
+      // For atomic_load, prefer to use the extension type.
+      if (Op->getOpcode() == ISD::ATOMIC_LOAD) {
+        if (AT->getExtensionType() == ISD::ZEXTLOAD)
+          Known.Zero.setBitsFrom(MemBits);
+        else if (AT->getExtensionType() != ISD::SEXTLOAD &&
+                 TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
+          Known.Zero.setBitsFrom(MemBits);
+      } else if (TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
         Known.Zero.setBitsFrom(MemBits);
     }
     break;
@@ -5252,22 +5257,29 @@ unsigned SelectionDAG::ComputeNumSignBits(SDValue Op, const APInt &DemandedElts,
   case ISD::ATOMIC_LOAD_UMIN:
   case ISD::ATOMIC_LOAD_UMAX:
   case ISD::ATOMIC_LOAD: {
-    Tmp = cast<AtomicSDNode>(Op)->getMemoryVT().getScalarSizeInBits();
+    auto *AT = cast<AtomicSDNode>(Op);
     // If we are looking at the loaded value.
     if (Op.getResNo() == 0) {
+      Tmp = AT->getMemoryVT().getScalarSizeInBits();
       if (Tmp == VTBits)
         return 1; // early-out
-      if (TLI->getExtendForAtomicOps() == ISD::SIGN_EXTEND)
-        return VTBits - Tmp + 1;
-      if (TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
-        return VTBits - Tmp;
+
+      // For atomic_load, prefer to use the extension type.
       if (Op->getOpcode() == ISD::ATOMIC_LOAD) {
-        ISD::LoadExtType ETy = cast<AtomicSDNode>(Op)->getExtensionType();
-        if (ETy == ISD::SEXTLOAD)
+        switch (AT->getExtensionType()) {
+        default:
+          break;
+        case ISD::SEXTLOAD:
           return VTBits - Tmp + 1;
-        if (ETy == ISD::ZEXTLOAD)
+        case ISD::ZEXTLOAD:
           return VTBits - Tmp;
+        }
       }
+
+      if (TLI->getExtendForAtomicOps() == ISD::SIGN_EXTEND)
+        return VTBits - Tmp + 1;
+      if (TLI->getExtendForAtomicOps() == ISD::ZERO_EXTEND)
+        return VTBits - Tmp;
     }
     break;
   }