[DAG] Improve known bits of Zext/Sext loads with range metadata #80829

davemgreen · 2024-02-06T12:56:10Z

This extends the known bits for extending loads which have range metadata, handling the range metadata on the original memory type, extending that to the correct BitWidth.

llvmbot · 2024-02-06T12:56:40Z

@llvm/pr-subscribers-backend-aarch64

Author: David Green (davemgreen)

Changes

This extends the known bits for extending loads which have range metadata, handling the range metadata on the original memory type, extending that to the correct BitWidth.

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

2 Files Affected:

(modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+29-19)
(modified) llvm/test/CodeGen/AArch64/setcc_knownbits.ll (+1-3)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 3c1343836187a..eec09219ad01a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3612,32 +3612,42 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
           }
         }
       }
-    } else if (ISD::isZEXTLoad(Op.getNode()) && Op.getResNo() == 0) {
-      // If this is a ZEXTLoad and we are looking at the loaded value.
-      EVT VT = LD->getMemoryVT();
-      unsigned MemBits = VT.getScalarSizeInBits();
-      Known.Zero.setBitsFrom(MemBits);
-    } else if (const MDNode *Ranges = LD->getRanges()) {
-      EVT VT = LD->getValueType(0);
-
-      // TODO: Handle for extending loads
-      if (LD->getExtensionType() == ISD::NON_EXTLOAD) {
+    } else if (Op.getResNo() == 0) {
+      KnownBits Known0(!LD->getMemoryVT().isScalableVT()
+                           ? LD->getMemoryVT().getSizeInBits()
+                           : BitWidth);
+      EVT VT = Op.getValueType();
+      // Fill in any known bits from range information. There are 3 types being
+      // used. The results VT (same vector elt size as BitWidth), the loaded
+      // MemoryVT (which may or may not be vector) and the range VTs original
+      // type. The range matadata needs the full range (i.e
+      // MemoryVT().getSizeInBits()), which is truncated to the correct elt size
+      // if it is know. These are then extended to the original VT sizes below.
+      if (const MDNode *MD = LD->getRanges()) {
+        computeKnownBitsFromRangeMetadata(*MD, Known0);
         if (VT.isVector()) {
           // Handle truncation to the first demanded element.
           // TODO: Figure out which demanded elements are covered
           if (DemandedElts != 1 || !getDataLayout().isLittleEndian())
             break;
+          Known0 = Known0.trunc(BitWidth);
+        }
+      }
 
-          // Handle the case where a load has a vector type, but scalar memory
-          // with an attached range.
-          EVT MemVT = LD->getMemoryVT();
-          KnownBits KnownFull(MemVT.getSizeInBits());
+      if (LD->getMemoryVT().isVector())
+        Known0 = Known0.trunc(LD->getMemoryVT().getScalarSizeInBits());
 
-          computeKnownBitsFromRangeMetadata(*Ranges, KnownFull);
-          Known = KnownFull.trunc(BitWidth);
-        } else
-          computeKnownBitsFromRangeMetadata(*Ranges, Known);
-      }
+      // Extend the Known bits from memory to the size of the result.
+      if (ISD::isZEXTLoad(Op.getNode()))
+        Known = Known0.zext(BitWidth);
+      else if (ISD::isSEXTLoad(Op.getNode()))
+        Known = Known0.sext(BitWidth);
+      else if (ISD::isEXTLoad(Op.getNode()))
+        Known = Known0.anyext(BitWidth);
+      else
+        Known = Known0;
+      assert(Known.getBitWidth() == BitWidth);
+      return Known;
     }
     break;
   }
diff --git a/llvm/test/CodeGen/AArch64/setcc_knownbits.ll b/llvm/test/CodeGen/AArch64/setcc_knownbits.ll
index 9e9c814be0266..60d6187e3423d 100644
--- a/llvm/test/CodeGen/AArch64/setcc_knownbits.ll
+++ b/llvm/test/CodeGen/AArch64/setcc_knownbits.ll
@@ -21,9 +21,7 @@ define noundef i1 @logger(i32 noundef %logLevel, ptr %ea, ptr %pll) {
 ; CHECK-NEXT:    ret
 ; CHECK-NEXT:  .LBB1_2: // %land.rhs
 ; CHECK-NEXT:    ldr x8, [x1]
-; CHECK-NEXT:    ldrb w8, [x8]
-; CHECK-NEXT:    cmp w8, #0
-; CHECK-NEXT:    cset w0, ne
+; CHECK-NEXT:    ldrb w0, [x8]
 ; CHECK-NEXT:    ret
 entry:
   %0 = load i32, ptr %pll, align 4

llvmbot · 2024-02-06T12:56:40Z

@llvm/pr-subscribers-llvm-selectiondag

Author: David Green (davemgreen)

Changes

This extends the known bits for extending loads which have range metadata, handling the range metadata on the original memory type, extending that to the correct BitWidth.

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

2 Files Affected:

(modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+29-19)
(modified) llvm/test/CodeGen/AArch64/setcc_knownbits.ll (+1-3)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 3c1343836187a..eec09219ad01a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -3612,32 +3612,42 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
           }
         }
       }
-    } else if (ISD::isZEXTLoad(Op.getNode()) && Op.getResNo() == 0) {
-      // If this is a ZEXTLoad and we are looking at the loaded value.
-      EVT VT = LD->getMemoryVT();
-      unsigned MemBits = VT.getScalarSizeInBits();
-      Known.Zero.setBitsFrom(MemBits);
-    } else if (const MDNode *Ranges = LD->getRanges()) {
-      EVT VT = LD->getValueType(0);
-
-      // TODO: Handle for extending loads
-      if (LD->getExtensionType() == ISD::NON_EXTLOAD) {
+    } else if (Op.getResNo() == 0) {
+      KnownBits Known0(!LD->getMemoryVT().isScalableVT()
+                           ? LD->getMemoryVT().getSizeInBits()
+                           : BitWidth);
+      EVT VT = Op.getValueType();
+      // Fill in any known bits from range information. There are 3 types being
+      // used. The results VT (same vector elt size as BitWidth), the loaded
+      // MemoryVT (which may or may not be vector) and the range VTs original
+      // type. The range matadata needs the full range (i.e
+      // MemoryVT().getSizeInBits()), which is truncated to the correct elt size
+      // if it is know. These are then extended to the original VT sizes below.
+      if (const MDNode *MD = LD->getRanges()) {
+        computeKnownBitsFromRangeMetadata(*MD, Known0);
         if (VT.isVector()) {
           // Handle truncation to the first demanded element.
           // TODO: Figure out which demanded elements are covered
           if (DemandedElts != 1 || !getDataLayout().isLittleEndian())
             break;
+          Known0 = Known0.trunc(BitWidth);
+        }
+      }
 
-          // Handle the case where a load has a vector type, but scalar memory
-          // with an attached range.
-          EVT MemVT = LD->getMemoryVT();
-          KnownBits KnownFull(MemVT.getSizeInBits());
+      if (LD->getMemoryVT().isVector())
+        Known0 = Known0.trunc(LD->getMemoryVT().getScalarSizeInBits());
 
-          computeKnownBitsFromRangeMetadata(*Ranges, KnownFull);
-          Known = KnownFull.trunc(BitWidth);
-        } else
-          computeKnownBitsFromRangeMetadata(*Ranges, Known);
-      }
+      // Extend the Known bits from memory to the size of the result.
+      if (ISD::isZEXTLoad(Op.getNode()))
+        Known = Known0.zext(BitWidth);
+      else if (ISD::isSEXTLoad(Op.getNode()))
+        Known = Known0.sext(BitWidth);
+      else if (ISD::isEXTLoad(Op.getNode()))
+        Known = Known0.anyext(BitWidth);
+      else
+        Known = Known0;
+      assert(Known.getBitWidth() == BitWidth);
+      return Known;
     }
     break;
   }
diff --git a/llvm/test/CodeGen/AArch64/setcc_knownbits.ll b/llvm/test/CodeGen/AArch64/setcc_knownbits.ll
index 9e9c814be0266..60d6187e3423d 100644
--- a/llvm/test/CodeGen/AArch64/setcc_knownbits.ll
+++ b/llvm/test/CodeGen/AArch64/setcc_knownbits.ll
@@ -21,9 +21,7 @@ define noundef i1 @logger(i32 noundef %logLevel, ptr %ea, ptr %pll) {
 ; CHECK-NEXT:    ret
 ; CHECK-NEXT:  .LBB1_2: // %land.rhs
 ; CHECK-NEXT:    ldr x8, [x1]
-; CHECK-NEXT:    ldrb w8, [x8]
-; CHECK-NEXT:    cmp w8, #0
-; CHECK-NEXT:    cset w0, ne
+; CHECK-NEXT:    ldrb w0, [x8]
 ; CHECK-NEXT:    ret
 entry:
   %0 = load i32, ptr %pll, align 4

arsenm · 2024-02-06T13:04:02Z

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

+      else if (ISD::isSEXTLoad(Op.getNode()))
+        Known = Known0.sext(BitWidth);
+      else if (ISD::isEXTLoad(Op.getNode()))
+        Known = Known0.anyext(BitWidth);


Do you need to guard against FP type ext loads?

It looks like we don't create extending loads with ISD opcode FP_EXTEND - we'll just create a normal load followed by a separate FP_EXTEND node. Perhaps there is a way to assert this, i.e. assert for FP types that memory VT == results VT?

Hmm. As far as I can tell fp loads will not have range data, and so if we have a anyext fp load then we should get a Known(LD->getMemoryVT().getFixedSizeInBits()).anyext(Bitwidth), which should be fine.
@arsenm any reason you think this might need to be guarded? Thanks.

I'm not sure MMO metadata is properly handled when bit casts are folded into load/store

arsenm · 2024-02-06T13:05:06Z

llvm/test/CodeGen/AArch64/setcc_knownbits.ll

-; CHECK-NEXT:    ldrb w8, [x8]
-; CHECK-NEXT:    cmp w8, #0
-; CHECK-NEXT:    cset w0, ne
+; CHECK-NEXT:    ldrb w0, [x8]


Not enough test updates? I would expect a diff from the zext/sext/anyext cases

Perhaps it's possible to use the range metadata to write some explicit tests for sext/anyext?

Yeah I need to add some better tests. Whilst this seems to come up a fair amount during C++ code (bools often have range metadata applied), it can be difficult to test is simple cases as the sext(load !range) is often already optimized before it creates the load sext !range. It might require C++ testing.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

david-arm · 2024-02-09T14:56:18Z

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

-      if (LD->getExtensionType() == ISD::NON_EXTLOAD) {
+    } else if (Op.getResNo() == 0) {
+      KnownBits Known0(!LD->getMemoryVT().isScalableVT()
+                           ? LD->getMemoryVT().getSizeInBits()


nit: Can you call getFixedSizeInBits here instead given we know it's fixed-width? I think it's a bit clearer and it does have a compile-time advantage too because it avoids going through the more expensive TypeSize->uint64_t cast operator, which is a function call.

david-arm · 2024-02-09T15:00:11Z

llvm/test/CodeGen/AArch64/setcc_knownbits.ll

-; CHECK-NEXT:    ldrb w8, [x8]
-; CHECK-NEXT:    cmp w8, #0
-; CHECK-NEXT:    cset w0, ne
+; CHECK-NEXT:    ldrb w0, [x8]


Perhaps it's possible to use the range metadata to write some explicit tests for sext/anyext?

This extends the known bits for extending loads which have range metadata, handling the range metadata on the original memory type, extending that to the correct BitWidths.

davemgreen · 2024-02-20T17:38:05Z

I've added a couple of C++ tests to make sure this is tested. Thanks.

david-arm

Thanks for the changes @davemgreen! The patch LGTM. There was just one comment about FP extending loads. I left a comment about possibly adding an assert - do you think this it's worth doing?

david-arm · 2024-02-21T09:36:08Z

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

+      else if (ISD::isSEXTLoad(Op.getNode()))
+        Known = Known0.sext(BitWidth);
+      else if (ISD::isEXTLoad(Op.getNode()))
+        Known = Known0.anyext(BitWidth);


It looks like we don't create extending loads with ISD opcode FP_EXTEND - we'll just create a normal load followed by a separate FP_EXTEND node. Perhaps there is a way to assert this, i.e. assert for FP types that memory VT == results VT?

arsenm

would be good to get the equivalent in globalisel

davemgreen · 2024-02-29T12:52:58Z

would be good to get the equivalent in globalisel

Yeah I'll add it to the list. Thanks.

…tadata Similar to llvm#80829 for GlobalISel.

…tadata (#86431) Similar to #80829 for GlobalISel.

davemgreen requested review from RKSimon and topperc February 6, 2024 12:56

llvmbot added backend:AArch64 llvm:SelectionDAG SelectionDAGISel as well labels Feb 6, 2024

arsenm reviewed Feb 6, 2024

View reviewed changes

nikic reviewed Feb 6, 2024

View reviewed changes

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp Show resolved Hide resolved

david-arm reviewed Feb 9, 2024

View reviewed changes

[DAG] Improve known bits of Zext/Sext loads with range metadata

5b14559

This extends the known bits for extending loads which have range metadata, handling the range metadata on the original memory type, extending that to the correct BitWidths.

davemgreen force-pushed the gh-a64-extloadrange branch from 8972e9f to 5b14559 Compare February 20, 2024 17:20

Use getFixedSizeInBits

3a34658

david-arm reviewed Feb 21, 2024

View reviewed changes

arsenm approved these changes Feb 28, 2024

View reviewed changes

davemgreen merged commit dbca8a4 into llvm:main Feb 29, 2024
4 checks passed

davemgreen deleted the gh-a64-extloadrange branch February 29, 2024 12:53

davemgreen added a commit to davemgreen/llvm-project that referenced this pull request Mar 24, 2024

[GlobalISel] Add Knownbits for G_LOAD/ZEXTLOAD/SEXTLOAD with range me…

0d066fb

…tadata Similar to llvm#80829 for GlobalISel.

davemgreen mentioned this pull request Mar 24, 2024

[GlobalISel] Add Knownbits for G_LOAD/ZEXTLOAD/SEXTLOAD with range metadata #86431

Merged

davemgreen added a commit that referenced this pull request Mar 26, 2024

[GlobalISel] Add Knownbits for G_LOAD/ZEXTLOAD/SEXTLOAD with range me…

47f4a07

…tadata (#86431) Similar to #80829 for GlobalISel.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[DAG] Improve known bits of Zext/Sext loads with range metadata #80829

[DAG] Improve known bits of Zext/Sext loads with range metadata #80829

davemgreen commented Feb 6, 2024

llvmbot commented Feb 6, 2024

llvmbot commented Feb 6, 2024

arsenm Feb 6, 2024

david-arm Feb 21, 2024

davemgreen Feb 21, 2024

arsenm Feb 28, 2024

arsenm Feb 6, 2024

david-arm Feb 9, 2024

davemgreen Feb 11, 2024

david-arm Feb 9, 2024

david-arm Feb 9, 2024

davemgreen commented Feb 20, 2024

david-arm left a comment

david-arm Feb 21, 2024

arsenm left a comment

davemgreen commented Feb 29, 2024

[DAG] Improve known bits of Zext/Sext loads with range metadata #80829

[DAG] Improve known bits of Zext/Sext loads with range metadata #80829

Conversation

davemgreen commented Feb 6, 2024

llvmbot commented Feb 6, 2024

llvmbot commented Feb 6, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

davemgreen commented Feb 20, 2024

david-arm left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

arsenm left a comment

Choose a reason for hiding this comment

davemgreen commented Feb 29, 2024