Skip to content

[CodeGen] Make LoadExtActions address-space aware #162407

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[CodeGen] Make LoadExtActions address-space aware #162407

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
99e87a3
Initial conversion of LoadExtActions to map
QuantumSegfault Oct 8, 2025
ebcf9ba
Finish making address-space sensitive, and fix AMDGPU backend
QuantumSegfault Oct 8, 2025
5091372
Make per address-space actions a sparse map overriding the original A…
QuantumSegfault Oct 8, 2025
8f51e32
Touch up formatting
QuantumSegfault Oct 8, 2025
63a5187
Fix problem causing AArch masked_ldst.ll fail
QuantumSegfault Oct 21, 2025
113ba86
Reformat
QuantumSegfault Oct 21, 2025
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
16 changes: 12 additions & 4 deletions llvm/include/llvm/CodeGen/BasicTTIImpl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
#include "llvm/Analysis/TargetTransformInfoImpl.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/ISDOpcodes.h"
#include "llvm/CodeGen/SelectionDAGNodes.h"
#include "llvm/CodeGen/TargetLowering.h"
#include "llvm/CodeGen/TargetSubtargetInfo.h"
#include "llvm/CodeGen/ValueTypes.h"
Expand Down Expand Up @@ -1244,9 +1245,15 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
EVT LoadVT = EVT::getEVT(Src);
unsigned LType =
((Opcode == Instruction::ZExt) ? ISD::ZEXTLOAD : ISD::SEXTLOAD);
if (DstLT.first == SrcLT.first &&
TLI->isLoadExtLegal(LType, ExtVT, LoadVT))
return 0;

if (I && isa<LoadInst>(I->getOperand(0))) {
auto *LI = cast<LoadInst>(I->getOperand(0));

if (DstLT.first == SrcLT.first &&
TLI->isLoadExtLegal(LType, ExtVT, LoadVT,
LI->getPointerAddressSpace()))
return 0;
}
}
break;
case Instruction::AddrSpaceCast:
Expand Down Expand Up @@ -1531,7 +1538,8 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
if (Opcode == Instruction::Store)
LA = getTLI()->getTruncStoreAction(LT.second, MemVT);
else
LA = getTLI()->getLoadExtAction(ISD::EXTLOAD, LT.second, MemVT);
LA = getTLI()->getLoadExtAction(ISD::EXTLOAD, LT.second, MemVT,
AddressSpace);

if (LA != TargetLowering::Legal && LA != TargetLowering::Custom) {
// This is a vector load/store for some illegal type that is scalarized.
Expand Down
67 changes: 48 additions & 19 deletions llvm/include/llvm/CodeGen/TargetLowering.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1472,27 +1472,38 @@ class LLVM_ABI TargetLoweringBase {
/// Return how this load with extension should be treated: either it is legal,
/// needs to be promoted to a larger size, needs to be expanded to some other
/// code sequence, or the target has a custom expander for it.
LegalizeAction getLoadExtAction(unsigned ExtType, EVT ValVT,
EVT MemVT) const {
if (ValVT.isExtended() || MemVT.isExtended()) return Expand;
unsigned ValI = (unsigned) ValVT.getSimpleVT().SimpleTy;
unsigned MemI = (unsigned) MemVT.getSimpleVT().SimpleTy;
LegalizeAction getLoadExtAction(unsigned ExtType, EVT ValVT, EVT MemVT,
unsigned AddrSpace) const {
Comment on lines +1475 to +1476
Copy link
Contributor

@arsenm arsenm Nov 25, 2025

Choose a reason for hiding this comment

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

This is too narrowly targeted at ext loads - really we need to know the address space for the legality of all loads (and stores). This has been a big DAG pain point since always (e.g., #90714). Every memory instruction legality query should be a function of the types, address space, and alignment.

std::map probably isn't the right way to encode this information either. In GlobalISel you have a chain of arbitrary legalization predicates that can check the address space.
It is a full 24-bit value, it just happens the common cases targets care about are small numbers. You can construct a system that wants to handle an arbitrary value range. The system today mostly treats non-0, but not explicitly handled address spaces as 0, which is kind of a nice property.

Copy link
Contributor Author

@QuantumSegfault QuantumSegfault Nov 25, 2025

Choose a reason for hiding this comment

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

Yeah. The idea was to get it worked out for loads, then move on to (at least) truncating stores. But yeah, they all share the same fundamental flaw

Copy link
Contributor Author

@QuantumSegfault QuantumSegfault Nov 25, 2025

Choose a reason for hiding this comment

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

Yeah, I noticed that in GISel. That's actually exactly why I started #157161.

If not a map, then it seems that maybe getLoadExtAction and the likes should just be a virtual hook? Defaults to Legal, and we just move any exceptions into code in each target? Though that seems like a step backwards...

Did you have something particular in mind?

Copy link
Contributor

@arsenm arsenm Nov 25, 2025

Choose a reason for hiding this comment

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

I was thinking something like:


virtual bool isLoadLegalImpl(ValVT, MemVT, AddrSpace, Align) {  return true; 
}
bool isLoadLegal(ValVT, MemVT, AddrSpace, Align) {
  Action = getOperationAction(ISD::LOAD, ...
  if (Action == Legal)
    return allowsMemoryAccess();

  if (Action == Custom)
    return isLoadLegalImpl(ValVT, MemVT, AddrSpace, Align);
  return false;
}

Copy link
Contributor Author

@QuantumSegfault QuantumSegfault Nov 25, 2025

Choose a reason for hiding this comment

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

So "Custom" is now overloaded to either mean conditionally legal, or actually custom lowered? Could be worse I suppose...

Do you have any thoughts on combining getOperationAction(ISD::LOAD... and getLoadExtAction into a unified getLoadAction that takes ext type, memory and val types, addrspace, and align. Have a target hook that can optionally return an override Action. Otherwise fallback to the existing data-driven definitions. Such that the existing setOperationAction and setLoadExtAction work, but special cases can customized. Probably unify uses of setOperationAction(ISD::LOAD and setLoadExtAction, and rename the later to be setLoadAction that handles both (since NON_EXTLOAD is a valid ext load type).

That way Custom and Legal remain distinct, there's still a data-driven API for simple existing cases, and we open the doors to more complex predicates. Not quite as elegant as GISel, but an improvement over the status-quo?

if (ValVT.isExtended() || MemVT.isExtended())
return Expand;
unsigned ValI = (unsigned)ValVT.getSimpleVT().SimpleTy;
unsigned MemI = (unsigned)MemVT.getSimpleVT().SimpleTy;
assert(ExtType < ISD::LAST_LOADEXT_TYPE && ValI < MVT::VALUETYPE_SIZE &&
MemI < MVT::VALUETYPE_SIZE && "Table isn't big enough!");
unsigned Shift = 4 * ExtType;

uint64_t OverrideKey = ((uint64_t)(ValI & 0xFF) << 40) |
((uint64_t)(MemI & 0xFF) << 32) |
(uint64_t)AddrSpace;

if (LoadExtActionOverrides.count(OverrideKey)) {
return (LegalizeAction)((LoadExtActionOverrides.at(OverrideKey) >> Shift) & 0xf);
}
return (LegalizeAction)((LoadExtActions[ValI][MemI] >> Shift) & 0xf);
}

/// Return true if the specified load with extension is legal on this target.
bool isLoadExtLegal(unsigned ExtType, EVT ValVT, EVT MemVT) const {
return getLoadExtAction(ExtType, ValVT, MemVT) == Legal;
bool isLoadExtLegal(unsigned ExtType, EVT ValVT, EVT MemVT,
unsigned AddrSpace) const {
return getLoadExtAction(ExtType, ValVT, MemVT, AddrSpace) == Legal;
}

/// Return true if the specified load with extension is legal or custom
/// on this target.
bool isLoadExtLegalOrCustom(unsigned ExtType, EVT ValVT, EVT MemVT) const {
return getLoadExtAction(ExtType, ValVT, MemVT) == Legal ||
getLoadExtAction(ExtType, ValVT, MemVT) == Custom;
bool isLoadExtLegalOrCustom(unsigned ExtType, EVT ValVT, EVT MemVT,
unsigned AddrSpace) const {
return getLoadExtAction(ExtType, ValVT, MemVT, AddrSpace) == Legal ||
getLoadExtAction(ExtType, ValVT, MemVT, AddrSpace) == Custom;
}

/// Same as getLoadExtAction, but for atomic loads.
Expand Down Expand Up @@ -2634,23 +2645,38 @@ class LLVM_ABI TargetLoweringBase {
/// Indicate that the specified load with extension does not work with the
/// specified type and indicate what to do about it.
void setLoadExtAction(unsigned ExtType, MVT ValVT, MVT MemVT,
LegalizeAction Action) {
LegalizeAction Action, unsigned AddrSpace = ~0) {
assert(ExtType < ISD::LAST_LOADEXT_TYPE && ValVT.isValid() &&
MemVT.isValid() && "Table isn't big enough!");
assert((unsigned)Action < 0x10 && "too many bits for bitfield array");

unsigned Shift = 4 * ExtType;
LoadExtActions[ValVT.SimpleTy][MemVT.SimpleTy] &= ~((uint16_t)0xF << Shift);
LoadExtActions[ValVT.SimpleTy][MemVT.SimpleTy] |= (uint16_t)Action << Shift;

if (AddrSpace == ~((unsigned)0)) {
LoadExtActions[ValVT.SimpleTy][MemVT.SimpleTy] &=
~((uint16_t)0xF << Shift);
LoadExtActions[ValVT.SimpleTy][MemVT.SimpleTy] |= (uint16_t)Action
<< Shift;
} else {
uint64_t OverrideKey = ((uint64_t)(ValVT.SimpleTy & 0xFF) << 40) |
((uint64_t)(MemVT.SimpleTy & 0xFF) << 32) |
(uint64_t)AddrSpace;
uint16_t &OverrideVal = LoadExtActionOverrides[OverrideKey];

OverrideVal &= ~((uint16_t)0xF << Shift);
OverrideVal |= (uint16_t)Action << Shift;
}
}
void setLoadExtAction(ArrayRef<unsigned> ExtTypes, MVT ValVT, MVT MemVT,
LegalizeAction Action) {
LegalizeAction Action, unsigned AddrSpace = ~0) {
for (auto ExtType : ExtTypes)
setLoadExtAction(ExtType, ValVT, MemVT, Action);
setLoadExtAction(ExtType, ValVT, MemVT, Action, AddrSpace);
}
void setLoadExtAction(ArrayRef<unsigned> ExtTypes, MVT ValVT,
ArrayRef<MVT> MemVTs, LegalizeAction Action) {
ArrayRef<MVT> MemVTs, LegalizeAction Action,
unsigned AddrSpace = ~0) {
for (auto MemVT : MemVTs)
setLoadExtAction(ExtTypes, ValVT, MemVT, Action);
setLoadExtAction(ExtTypes, ValVT, MemVT, Action, AddrSpace);
}

/// Let target indicate that an extending atomic load of the specified type
Expand Down Expand Up @@ -3126,7 +3152,7 @@ class LLVM_ABI TargetLoweringBase {
LType = ISD::SEXTLOAD;
}

return isLoadExtLegal(LType, VT, LoadVT);
return isLoadExtLegal(LType, VT, LoadVT, Load->getPointerAddressSpace());
}

/// Return true if any actual instruction that defines a value of type FromTy
Expand Down Expand Up @@ -3753,8 +3779,11 @@ class LLVM_ABI TargetLoweringBase {
/// For each load extension type and each value type, keep a LegalizeAction
/// that indicates how instruction selection should deal with a load of a
/// specific value type and extension type. Uses 4-bits to store the action
/// for each of the 4 load ext types.
/// for each of the 4 load ext types. These actions can be specified for each
/// address space.
uint16_t LoadExtActions[MVT::VALUETYPE_SIZE][MVT::VALUETYPE_SIZE];
using LoadExtActionOverrideMap = std::map<uint64_t, uint16_t>;
LoadExtActionOverrideMap LoadExtActionOverrides;

/// Similar to LoadExtActions, but for atomic loads. Only Legal or Expand
/// (default) values are supported.
Expand Down
3 changes: 2 additions & 1 deletion llvm/lib/CodeGen/CodeGenPrepare.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7347,7 +7347,8 @@ bool CodeGenPrepare::optimizeLoadExt(LoadInst *Load) {

// Reject cases that won't be matched as extloads.
if (!LoadResultVT.bitsGT(TruncVT) || !TruncVT.isRound() ||
!TLI->isLoadExtLegal(ISD::ZEXTLOAD, LoadResultVT, TruncVT))
!TLI->isLoadExtLegal(ISD::ZEXTLOAD, LoadResultVT, TruncVT,
Load->getPointerAddressSpace()))
return false;

IRBuilder<> Builder(Load->getNextNode());
Expand Down
Loading
Loading