Skip to content

[SPIR-V] Move ASSIGN_TYPE generation to PostLegalizer #169696

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
s-perron merged 4 commits into from
Dec 4, 2025
Merged

[SPIR-V] Move ASSIGN_TYPE generation to PostLegalizer #169696

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7de5b73
[SPIR-V] Move ASSIGN_TYPE generation to PostLegalizer
s-perron Oct 30, 2025
8745bb2
Remove redudant return.
s-perron Dec 1, 2025
990957d
Move debugging code.
s-perron Dec 1, 2025
6bc36c0
Merge branch 'main' into spirv-move-assign-type
s-perron Dec 1, 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
36 changes: 16 additions & 20 deletions llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -533,14 +533,11 @@ static Register buildBuiltinVariableLoad(
return LoadedRegister;
}

/// Helper external function for inserting ASSIGN_TYPE instuction between \p Reg
/// and its definition, set the new register as a destination of the definition,
/// assign SPIRVType to both registers. If SpirvTy is provided, use it as
/// SPIRVType in ASSIGN_TYPE, otherwise create it from \p Ty. Defined in
/// SPIRVPreLegalizer.cpp.
extern void insertAssignInstr(Register Reg, Type *Ty, SPIRVType *SpirvTy,
SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI);
/// Helper external function for assigning SPIRVType to a register, ensuring the
/// register class and type are set in MRI. Defined in SPIRVPreLegalizer.cpp.
extern void updateRegType(Register Reg, Type *Ty, SPIRVType *SpirvTy,
SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI);

// TODO: Move to TableGen.
static SPIRV::MemorySemantics::MemorySemantics
Expand Down Expand Up @@ -860,8 +857,8 @@ static bool buildAtomicRMWInst(const SPIRV::IncomingCall *Call, unsigned Opcode,
MIRBuilder.buildInstr(TargetOpcode::G_FNEG)
.addDef(NegValueReg)
.addUse(ValueReg);
insertAssignInstr(NegValueReg, nullptr, Call->ReturnType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
updateRegType(NegValueReg, nullptr, Call->ReturnType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
ValueReg = NegValueReg;
}
}
Expand Down Expand Up @@ -1285,8 +1282,8 @@ static bool generateGroupInst(const SPIRV::IncomingCall *Call,
MIRBuilder.buildICmp(
CmpInst::ICMP_NE, Arg0, BoolReg,
GR->buildConstantInt(0, MIRBuilder, BoolRegType, true));
insertAssignInstr(Arg0, nullptr, BoolType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
updateRegType(Arg0, nullptr, BoolType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
} else if (BoolRegType->getOpcode() != SPIRV::OpTypeBool) {
report_fatal_error("Expect a boolean argument");
}
Expand Down Expand Up @@ -1337,8 +1334,8 @@ static bool generateGroupInst(const SPIRV::IncomingCall *Call,
MIB.addUse(Call->Arguments[i]);
setRegClassIfNull(Call->Arguments[i], MRI, GR);
}
insertAssignInstr(VecReg, nullptr, VecType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
updateRegType(VecReg, nullptr, VecType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
}

// Build work/sub group instruction.
Expand Down Expand Up @@ -1603,8 +1600,7 @@ static bool genWorkgroupQuery(const SPIRV::IncomingCall *Call,

// If the index is dynamic, need check if it's < 3, and then use a select.
if (!IsConstantIndex) {
insertAssignInstr(Extracted, nullptr, PointerSizeType, GR, MIRBuilder,
*MRI);
updateRegType(Extracted, nullptr, PointerSizeType, GR, MIRBuilder, *MRI);

auto IndexType = GR->getSPIRVTypeForVReg(IndexRegister);
auto BoolType = GR->getOrCreateSPIRVBoolType(MIRBuilder, true);
Expand Down Expand Up @@ -1992,8 +1988,8 @@ static bool generateImageSizeQueryInst(const SPIRV::IncomingCall *Call,
.addUse(QueryResult)
.addImm(ExtractedComposite);
if (NewType != nullptr)
insertAssignInstr(Call->ReturnRegister, nullptr, NewType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
updateRegType(Call->ReturnRegister, nullptr, NewType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
} else {
// More than 1 component is expected, fill a new vector.
auto MIB = MIRBuilder.buildInstr(SPIRV::OpVectorShuffle)
Expand Down Expand Up @@ -2695,8 +2691,8 @@ static bool generateAsyncCopy(const SPIRV::IncomingCall *Call,
: buildConstantIntReg32(1, MIRBuilder, GR))
.addUse(EventReg);
if (NewType != nullptr)
insertAssignInstr(Call->ReturnRegister, nullptr, NewType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
updateRegType(Call->ReturnRegister, nullptr, NewType, GR, MIRBuilder,
MIRBuilder.getMF().getRegInfo());
return Res;
}
case SPIRV::OpGroupWaitEvents:
Expand Down
4 changes: 2 additions & 2 deletions llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,7 @@ class SPIRVGlobalRegistry : public SPIRVIRMapping {
}

// Get or create a SPIR-V type corresponding the given LLVM IR type,
// and map it to the given VReg by creating an ASSIGN_TYPE instruction.
// and map it to the given VReg.
SPIRVType *assignTypeToVReg(const Type *Type, Register VReg,
MachineIRBuilder &MIRBuilder,
SPIRV::AccessQualifier::AccessQualifier AQ,
Expand All @@ -290,7 +290,7 @@ class SPIRVGlobalRegistry : public SPIRVIRMapping {
const SPIRVInstrInfo &TII);

// In cases where the SPIR-V type is already known, this function can be
// used to map it to the given VReg via an ASSIGN_TYPE instruction.
// used to map it to the given VReg.
void assignSPIRVTypeToVReg(SPIRVType *Type, Register VReg,
const MachineFunction &MF);

Expand Down
94 changes: 62 additions & 32 deletions llvm/lib/Target/SPIRV/SPIRVPostLegalizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,15 @@
//===-- SPIRVPostLegalizer.cpp - ammend info after legalization -*- C++ -*-===//
//
// which may appear after the legalizer pass
//===-- SPIRVPostLegalizer.cpp - amend info after legalization -*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// The pass partially apply pre-legalization logic to new instructions inserted
// as a result of legalization:
// The pass partially applies pre-legalization logic to new instructions
// inserted as a result of legalization:
// - assigns SPIR-V types to registers for new instructions.
// - inserts ASSIGN_TYPE pseudo-instructions required for type folding.
//
//===----------------------------------------------------------------------===//

Expand All @@ -36,9 +35,9 @@ class SPIRVPostLegalizer : public MachineFunctionPass {

namespace llvm {
// Defined in SPIRVPreLegalizer.cpp.
extern void insertAssignInstr(Register Reg, Type *Ty, SPIRVType *SpirvTy,
SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI);
extern void updateRegType(Register Reg, Type *Ty, SPIRVType *SpirvTy,
SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI);
extern void processInstr(MachineInstr &MI, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI, SPIRVGlobalRegistry *GR,
SPIRVType *KnownResType);
Expand Down Expand Up @@ -314,6 +313,56 @@ static void registerSpirvTypeForNewInstructions(MachineFunction &MF,
}
}

static bool hasAssignType(Register Reg, MachineRegisterInfo &MRI) {
for (MachineInstr &UseInstr : MRI.use_nodbg_instructions(Reg)) {
if (UseInstr.getOpcode() == SPIRV::ASSIGN_TYPE) {
LLVM_DEBUG(dbgs() << " Instruction already has an ASSIGN_TYPE use: "
<< UseInstr);
return true;
}
}
return false;
}

static void generateAssignType(MachineInstr &MI, Register ResultRegister,
SPIRVType *ResultType, SPIRVGlobalRegistry *GR,
MachineRegisterInfo &MRI) {
LLVM_DEBUG(dbgs() << " Adding ASSIGN_TYPE for ResultRegister: "
<< printReg(ResultRegister, MRI.getTargetRegisterInfo())
<< " with type: " << *ResultType);
MachineIRBuilder MIB(MI);
updateRegType(ResultRegister, nullptr, ResultType, GR, MIB, MRI);

// Tablegen definition assumes SPIRV::ASSIGN_TYPE pseudo-instruction is
// present after each auto-folded instruction to take a type reference
// from.
Register NewReg =
MRI.createGenericVirtualRegister(MRI.getType(ResultRegister));
const auto *RegClass = GR->getRegClass(ResultType);
MRI.setRegClass(NewReg, RegClass);
MRI.setRegClass(ResultRegister, RegClass);

GR->assignSPIRVTypeToVReg(ResultType, ResultRegister, MIB.getMF());
// This is to make it convenient for Legalizer to get the SPIRVType
// when processing the actual MI (i.e. not pseudo one).
GR->assignSPIRVTypeToVReg(ResultType, NewReg, MIB.getMF());
// Copy MIFlags from Def to ASSIGN_TYPE instruction. It's required to
// keep the flags after instruction selection.
const uint32_t Flags = MI.getFlags();
MIB.buildInstr(SPIRV::ASSIGN_TYPE)
.addDef(ResultRegister)
.addUse(NewReg)
.addUse(GR->getSPIRVTypeID(ResultType))
.setMIFlags(Flags);
for (unsigned I = 0, E = MI.getNumDefs(); I != E; ++I) {
MachineOperand &MO = MI.getOperand(I);
if (MO.getReg() == ResultRegister) {
MO.setReg(NewReg);
break;
}
}
}

static void ensureAssignTypeForTypeFolding(MachineFunction &MF,
SPIRVGlobalRegistry *GR) {
LLVM_DEBUG(dbgs() << "Entering ensureAssignTypeForTypeFolding for function "
Expand All @@ -323,35 +372,16 @@ static void ensureAssignTypeForTypeFolding(MachineFunction &MF,
for (MachineInstr &MI : MBB) {
if (!isTypeFoldingSupported(MI.getOpcode()))
continue;
if (MI.getNumOperands() == 1 || !MI.getOperand(1).isReg())
continue;

LLVM_DEBUG(dbgs() << "Processing instruction: " << MI);

// Check uses of MI to see if it already has an use in SPIRV::ASSIGN_TYPE
bool HasAssignType = false;
Register ResultRegister = MI.defs().begin()->getReg();
// All uses of Result register
for (MachineInstr &UseInstr :
MRI.use_nodbg_instructions(ResultRegister)) {
if (UseInstr.getOpcode() == SPIRV::ASSIGN_TYPE) {
HasAssignType = true;
LLVM_DEBUG(dbgs() << " Instruction already has an ASSIGN_TYPE use: "
<< UseInstr);
break;
}
}
if (hasAssignType(ResultRegister, MRI))
continue;

if (!HasAssignType) {
Register ResultRegister = MI.defs().begin()->getReg();
SPIRVType *ResultType = GR->getSPIRVTypeForVReg(ResultRegister);
LLVM_DEBUG(
dbgs() << " Adding ASSIGN_TYPE for ResultRegister: "
<< printReg(ResultRegister, MRI.getTargetRegisterInfo())
<< " with type: " << *ResultType);
MachineIRBuilder MIB(MI);
insertAssignInstr(ResultRegister, nullptr, ResultType, GR, MIB, MRI);
}
SPIRVType *ResultType = GR->getSPIRVTypeForVReg(ResultRegister);
assert(ResultType);
generateAssignType(MI, ResultRegister, ResultType, GR, MRI);
}
}
}
Expand Down
83 changes: 16 additions & 67 deletions llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -425,52 +425,21 @@ static void setInsertPtAfterDef(MachineIRBuilder &MIB, MachineInstr *Def) {
}

namespace llvm {
void insertAssignInstr(Register Reg, Type *Ty, SPIRVType *SpvType,
SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI) {
void updateRegType(Register Reg, Type *Ty, SPIRVType *SpvType,
SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
MachineRegisterInfo &MRI) {
assert((Ty || SpvType) && "Either LLVM or SPIRV type is expected.");
MachineInstr *Def = MRI.getVRegDef(Reg);
setInsertPtAfterDef(MIB, Def);
if (!SpvType)
SpvType = GR->getOrCreateSPIRVType(Ty, MIB,
SPIRV::AccessQualifier::ReadWrite, true);

if (!isTypeFoldingSupported(Def->getOpcode())) {
// No need to generate SPIRV::ASSIGN_TYPE pseudo-instruction
if (!MRI.getRegClassOrNull(Reg))
MRI.setRegClass(Reg, GR->getRegClass(SpvType));
if (!MRI.getType(Reg).isValid())
MRI.setType(Reg, GR->getRegType(SpvType));
GR->assignSPIRVTypeToVReg(SpvType, Reg, MIB.getMF());
return;
}

// Tablegen definition assumes SPIRV::ASSIGN_TYPE pseudo-instruction is
// present after each auto-folded instruction to take a type reference from.
Register NewReg = MRI.createGenericVirtualRegister(MRI.getType(Reg));
const auto *RegClass = GR->getRegClass(SpvType);
MRI.setRegClass(NewReg, RegClass);
MRI.setRegClass(Reg, RegClass);

if (!MRI.getRegClassOrNull(Reg))
MRI.setRegClass(Reg, GR->getRegClass(SpvType));
if (!MRI.getType(Reg).isValid())
MRI.setType(Reg, GR->getRegType(SpvType));
GR->assignSPIRVTypeToVReg(SpvType, Reg, MIB.getMF());
// This is to make it convenient for Legalizer to get the SPIRVType
// when processing the actual MI (i.e. not pseudo one).
GR->assignSPIRVTypeToVReg(SpvType, NewReg, MIB.getMF());
// Copy MIFlags from Def to ASSIGN_TYPE instruction. It's required to keep
// the flags after instruction selection.
const uint32_t Flags = Def->getFlags();
MIB.buildInstr(SPIRV::ASSIGN_TYPE)
.addDef(Reg)
.addUse(NewReg)
.addUse(GR->getSPIRVTypeID(SpvType))
.setMIFlags(Flags);
for (unsigned I = 0, E = Def->getNumDefs(); I != E; ++I) {
MachineOperand &MO = Def->getOperand(I);
if (MO.getReg() == Reg) {
MO.setReg(NewReg);
break;
}
}
return;
}

void processInstr(MachineInstr &MI, MachineIRBuilder &MIB,
Expand Down Expand Up @@ -543,20 +512,18 @@ generateAssignInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
MachineInstr *Def = MRI.getVRegDef(Reg);
assert(Def && "Expecting an instruction that defines the register");
// G_GLOBAL_VALUE already has type info.
if (Def->getOpcode() != TargetOpcode::G_GLOBAL_VALUE &&
Def->getOpcode() != SPIRV::ASSIGN_TYPE)
insertAssignInstr(Reg, nullptr, AssignedPtrType, GR, MIB,
MF.getRegInfo());
if (Def->getOpcode() != TargetOpcode::G_GLOBAL_VALUE)
updateRegType(Reg, nullptr, AssignedPtrType, GR, MIB,
MF.getRegInfo());
ToErase.push_back(&MI);
} else if (isSpvIntrinsic(MI, Intrinsic::spv_assign_type)) {
Register Reg = MI.getOperand(1).getReg();
Type *Ty = getMDOperandAsType(MI.getOperand(2).getMetadata(), 0);
MachineInstr *Def = MRI.getVRegDef(Reg);
assert(Def && "Expecting an instruction that defines the register");
// G_GLOBAL_VALUE already has type info.
if (Def->getOpcode() != TargetOpcode::G_GLOBAL_VALUE &&
Def->getOpcode() != SPIRV::ASSIGN_TYPE)
insertAssignInstr(Reg, Ty, nullptr, GR, MIB, MF.getRegInfo());
if (Def->getOpcode() != TargetOpcode::G_GLOBAL_VALUE)
updateRegType(Reg, Ty, nullptr, GR, MIB, MF.getRegInfo());
ToErase.push_back(&MI);
} else if (MIOp == TargetOpcode::FAKE_USE && MI.getNumOperands() > 0) {
MachineInstr *MdMI = MI.getPrevNode();
Expand All @@ -581,20 +548,9 @@ generateAssignInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
MIOp == TargetOpcode::G_FCONSTANT ||
MIOp == TargetOpcode::G_BUILD_VECTOR) {
// %rc = G_CONSTANT ty Val
// ===>
// %cty = OpType* ty
// %rctmp = G_CONSTANT ty Val
// %rc = ASSIGN_TYPE %rctmp, %cty
// Ensure %rc has a valid SPIR-V type assigned in the Global Registry.
Register Reg = MI.getOperand(0).getReg();
bool NeedAssignType = true;
if (MRI.hasOneUse(Reg)) {
MachineInstr &UseMI = *MRI.use_instr_begin(Reg);
if (isSpvIntrinsic(UseMI, Intrinsic::spv_assign_type) ||
isSpvIntrinsic(UseMI, Intrinsic::spv_assign_name))
continue;
if (UseMI.getOpcode() == SPIRV::ASSIGN_TYPE)
NeedAssignType = false;
}
bool NeedAssignType = GR->getSPIRVTypeForVReg(Reg) == nullptr;
Type *Ty = nullptr;
if (MIOp == TargetOpcode::G_CONSTANT) {
auto TargetExtIt = TargetExtConstTypes.find(&MI);
Expand Down Expand Up @@ -639,13 +595,6 @@ generateAssignInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
if (const SPIRVType *ElemSpvType =
GR->getSPIRVTypeForVReg(MI.getOperand(1).getReg(), &MF))
ElemTy = const_cast<Type *>(GR->getTypeForSPIRVType(ElemSpvType));
if (!ElemTy) {
// There may be a case when we already know Reg's type.
MachineInstr *NextMI = MI.getNextNode();
if (!NextMI || NextMI->getOpcode() != SPIRV::ASSIGN_TYPE ||
NextMI->getOperand(1).getReg() != Reg)
llvm_unreachable("Unexpected opcode");
}
}
if (ElemTy)
Ty = VectorType::get(
Expand All @@ -655,7 +604,7 @@ generateAssignInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
NeedAssignType = false;
}
if (NeedAssignType)
insertAssignInstr(Reg, Ty, nullptr, GR, MIB, MRI);
updateRegType(Reg, Ty, nullptr, GR, MIB, MRI);
} else if (MIOp == TargetOpcode::G_GLOBAL_VALUE) {
propagateSPIRVType(&MI, GR, MRI, MIB);
}
Expand Down