Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 29, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5b14559
[DAG] Improve known bits of Zext/Sext loads with range metadata
davemgreen Feb 20, 2024
3a34658
Use getFixedSizeInBits
davemgreen Feb 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 29 additions & 19 deletions llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3645,32 +3645,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().getFixedSizeInBits()
: 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);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you need to guard against FP type ext loads?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

else
Known = Known0;
assert(Known.getBitWidth() == BitWidth);
return Known;
}
break;
}
Expand Down
4 changes: 1 addition & 3 deletions llvm/test/CodeGen/AArch64/setcc_knownbits.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

; CHECK-NEXT: ret
entry:
%0 = load i32, ptr %pll, align 4
Expand Down
68 changes: 68 additions & 0 deletions llvm/unittests/CodeGen/AArch64SelectionDAGTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,13 @@
//
//===----------------------------------------------------------------------===//

#include "llvm/Analysis/MemoryLocation.h"
#include "llvm/Analysis/OptimizationRemarkEmitter.h"
#include "llvm/AsmParser/Parser.h"
#include "llvm/CodeGen/MachineModuleInfo.h"
#include "llvm/CodeGen/SelectionDAG.h"
#include "llvm/CodeGen/TargetLowering.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/KnownBits.h"
#include "llvm/Support/SourceMgr.h"
Expand Down Expand Up @@ -728,4 +730,70 @@ TEST_F(AArch64SelectionDAGTest, ReplaceAllUsesWith) {
EXPECT_EQ(DAG->getPCSections(New.getNode()), MD);
}

TEST_F(AArch64SelectionDAGTest, computeKnownBits_extload_known01) {
SDLoc Loc;
auto Int8VT = EVT::getIntegerVT(Context, 8);
auto Int32VT = EVT::getIntegerVT(Context, 32);
auto Int64VT = EVT::getIntegerVT(Context, 64);
auto Ptr = DAG->getConstant(0, Loc, Int64VT);
auto PtrInfo =
MachinePointerInfo::getFixedStack(DAG->getMachineFunction(), 0);
AAMDNodes AA;
MDBuilder MDHelper(*DAG->getContext());
MDNode *Range = MDHelper.createRange(APInt(8, 0), APInt(8, 2));
MachineMemOperand *MMO = DAG->getMachineFunction().getMachineMemOperand(
PtrInfo, MachineMemOperand::MOLoad, 8, Align(8), AA, Range);

auto ALoad = DAG->getExtLoad(ISD::EXTLOAD, Loc, Int32VT, DAG->getEntryNode(),
Ptr, Int8VT, MMO);
KnownBits Known = DAG->computeKnownBits(ALoad);
EXPECT_EQ(Known.Zero, APInt(32, 0xfe));
EXPECT_EQ(Known.One, APInt(32, 0));

auto ZLoad = DAG->getExtLoad(ISD::ZEXTLOAD, Loc, Int32VT, DAG->getEntryNode(),
Ptr, Int8VT, MMO);
Known = DAG->computeKnownBits(ZLoad);
EXPECT_EQ(Known.Zero, APInt(32, 0xfffffffe));
EXPECT_EQ(Known.One, APInt(32, 0));

auto SLoad = DAG->getExtLoad(ISD::SEXTLOAD, Loc, Int32VT, DAG->getEntryNode(),
Ptr, Int8VT, MMO);
Known = DAG->computeKnownBits(SLoad);
EXPECT_EQ(Known.Zero, APInt(32, 0xfffffffe));
EXPECT_EQ(Known.One, APInt(32, 0));
}

TEST_F(AArch64SelectionDAGTest, computeKnownBits_extload_knownnegative) {
SDLoc Loc;
auto Int8VT = EVT::getIntegerVT(Context, 8);
auto Int32VT = EVT::getIntegerVT(Context, 32);
auto Int64VT = EVT::getIntegerVT(Context, 64);
auto Ptr = DAG->getConstant(0, Loc, Int64VT);
auto PtrInfo =
MachinePointerInfo::getFixedStack(DAG->getMachineFunction(), 0);
AAMDNodes AA;
MDBuilder MDHelper(*DAG->getContext());
MDNode *Range = MDHelper.createRange(APInt(8, 0xf0), APInt(8, 0xff));
MachineMemOperand *MMO = DAG->getMachineFunction().getMachineMemOperand(
PtrInfo, MachineMemOperand::MOLoad, 8, Align(8), AA, Range);

auto ALoad = DAG->getExtLoad(ISD::EXTLOAD, Loc, Int32VT, DAG->getEntryNode(),
Ptr, Int8VT, MMO);
KnownBits Known = DAG->computeKnownBits(ALoad);
EXPECT_EQ(Known.Zero, APInt(32, 0));
EXPECT_EQ(Known.One, APInt(32, 0xf0));

auto ZLoad = DAG->getExtLoad(ISD::ZEXTLOAD, Loc, Int32VT, DAG->getEntryNode(),
Ptr, Int8VT, MMO);
Known = DAG->computeKnownBits(ZLoad);
EXPECT_EQ(Known.Zero, APInt(32, 0xffffff00));
EXPECT_EQ(Known.One, APInt(32, 0x000000f0));

auto SLoad = DAG->getExtLoad(ISD::SEXTLOAD, Loc, Int32VT, DAG->getEntryNode(),
Ptr, Int8VT, MMO);
Known = DAG->computeKnownBits(SLoad);
EXPECT_EQ(Known.Zero, APInt(32, 0));
EXPECT_EQ(Known.One, APInt(32, 0xfffffff0));
}

} // end namespace llvm