[NFC][SPIRV] Introduce function to handle 64 bits overflow (#193088)

Some intrinsics require special logic to deal with 64bits, usually
because vulkan doesn't allow 64 overloads. This patch introduces a
function that can be reused in some of those intrinsics to deal with
this specific case.

Requested here
https://github.com/llvm/llvm-project/pull/193068#discussion_r3113352858
as part of https://github.com/llvm/llvm-project/issues/192756
This commit is contained in:
joaosaffran
2026-04-22 15:21:03 -07:00
committed by GitHub
parent 107701b3e6
commit e68d91afdf

View File

@@ -34,6 +34,7 @@
#include "llvm/IR/IntrinsicsSPIRV.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
#include <functional>
#include <optional>
#define DEBUG_TYPE "spirv-isel"
@@ -144,11 +145,6 @@ private:
MachineInstr &I, Register SrcReg,
unsigned BitSetOpcode, bool SwapPrimarySide) const;
bool selectFirstBitSet64Overflow(Register ResVReg, SPIRVTypeInst ResType,
MachineInstr &I, Register SrcReg,
unsigned BitSetOpcode,
bool SwapPrimarySide) const;
bool selectGlobalValue(Register ResVReg, MachineInstr &I,
const MachineInstr *Init = nullptr) const;
@@ -257,10 +253,6 @@ private:
MachineInstr &I, Register SrcReg,
unsigned Opcode) const;
bool selectPopCount64Overflow(Register ResVReg, SPIRVTypeInst ResType,
MachineInstr &I, Register SrcReg,
unsigned Opcode) const;
template <bool Signed>
bool selectDot4AddPacked(Register ResVReg, SPIRVTypeInst ResType,
MachineInstr &I) const;
@@ -510,6 +502,13 @@ private:
unsigned ComponentCount,
MachineInstr &I,
SPIRVTypeInst I32Type) const;
bool
handle64BitOverflow(Register ResVReg, SPIRVTypeInst ResType, MachineInstr &I,
Register SrcReg, unsigned int Opcode,
std::function<bool(Register, SPIRVTypeInst,
MachineInstr &, Register, unsigned)>
CallbackFunction) const;
};
bool sampledTypeIsSignedInteger(const llvm::Type *HandleType) {
@@ -1685,74 +1684,6 @@ bool SPIRVInstructionSelector::selectPopCount32(Register ResVReg,
return selectOpWithSrcs(ResVReg, ResType, I, {SrcReg}, Opcode);
}
bool SPIRVInstructionSelector::selectPopCount64Overflow(
Register ResVReg, SPIRVTypeInst ResType, MachineInstr &I, Register SrcReg,
unsigned int Opcode) const {
unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
assert(ComponentCount < 5 && "Vec 5+ will generate invalid SPIR-V ops");
MachineIRBuilder MIRBuilder(I);
SPIRVTypeInst BaseType = GR.retrieveScalarOrVectorIntType(ResType);
SPIRVTypeInst I64Type = GR.getOrCreateSPIRVIntegerType(64, MIRBuilder);
SPIRVTypeInst I64x2Type =
GR.getOrCreateSPIRVVectorType(I64Type, 2, MIRBuilder, false);
SPIRVTypeInst Vec2ResType =
GR.getOrCreateSPIRVVectorType(BaseType, 2, MIRBuilder, false);
std::vector<Register> PartialRegs;
unsigned CurrentComponent = 0;
for (; CurrentComponent + 1 < ComponentCount; CurrentComponent += 2) {
Register PopCountResult =
MRI->createVirtualRegister(GR.getRegClass(I64x2Type));
auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
TII.get(SPIRV::OpVectorShuffle))
.addDef(PopCountResult)
.addUse(GR.getSPIRVTypeID(I64x2Type))
.addUse(SrcReg)
.addUse(SrcReg)
.addImm(CurrentComponent)
.addImm(CurrentComponent + 1);
MIB.constrainAllUses(TII, TRI, RBI);
Register SubVecReg =
MRI->createVirtualRegister(GR.getRegClass(Vec2ResType));
if (!selectPopCount64(SubVecReg, Vec2ResType, I, PopCountResult, Opcode))
return false;
PartialRegs.push_back(SubVecReg);
}
// On odd component counts we need to handle one more component
if (CurrentComponent != ComponentCount) {
bool ZeroAsNull = !STI.isShader();
Register FinalElemReg = MRI->createVirtualRegister(GR.getRegClass(I64Type));
Register ConstIntLastIdx = GR.getOrCreateConstInt(
ComponentCount - 1, I, BaseType, TII, ZeroAsNull);
if (!selectOpWithSrcs(FinalElemReg, I64Type, I, {SrcReg, ConstIntLastIdx},
SPIRV::OpVectorExtractDynamic))
return false;
Register FinalElemResReg =
MRI->createVirtualRegister(GR.getRegClass(BaseType));
if (!selectPopCount64(FinalElemResReg, BaseType, I, FinalElemReg, Opcode))
return false;
PartialRegs.push_back(FinalElemResReg);
}
// Join all the resulting registers back into the return type in order
// (ie i32x2, i32x2, i32x1 -> i32x5)
return selectOpWithSrcs(ResVReg, ResType, I, std::move(PartialRegs),
SPIRV::OpCompositeConstruct);
}
bool SPIRVInstructionSelector::selectPopCount64(Register ResVReg,
SPIRVTypeInst ResType,
MachineInstr &I,
@@ -1760,7 +1691,10 @@ bool SPIRVInstructionSelector::selectPopCount64(Register ResVReg,
unsigned Opcode) const {
unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
if (ComponentCount > 2)
return selectPopCount64Overflow(ResVReg, ResType, I, SrcReg, Opcode);
return handle64BitOverflow(
ResVReg, ResType, I, SrcReg, Opcode,
[this](Register R, SPIRVTypeInst T, MachineInstr &I, Register S,
unsigned O) { return this->selectPopCount64(R, T, I, S, O); });
MachineIRBuilder MIRBuilder(I);
@@ -3571,6 +3505,77 @@ bool SPIRVInstructionSelector::selectBitreverse16(Register ResVReg,
return selectOpWithSrcs(ResVReg, ResType, I, {ShiftReg}, ExtendOpcode);
}
bool SPIRVInstructionSelector::handle64BitOverflow(
Register ResVReg, SPIRVTypeInst ResType, MachineInstr &I, Register SrcReg,
unsigned int Opcode,
std::function<bool(Register, SPIRVTypeInst, MachineInstr &, Register,
unsigned)>
CallbackFunction) const {
SPIRVTypeInst BaseType = GR.retrieveScalarOrVectorIntType(ResType);
assert(BaseType->getOpcode() == SPIRV::OpTypeInt &&
"handle64BitOverflow should only be used for integer types");
unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
assert(ComponentCount < 5 && "Vec 5+ will generate invalid SPIR-V ops");
MachineIRBuilder MIRBuilder(I);
SPIRVTypeInst I64Type = GR.getOrCreateSPIRVIntegerType(64, MIRBuilder);
SPIRVTypeInst I64x2Type =
GR.getOrCreateSPIRVVectorType(I64Type, 2, MIRBuilder, false);
SPIRVTypeInst Vec2ResType =
GR.getOrCreateSPIRVVectorType(BaseType, 2, MIRBuilder, false);
std::vector<Register> PartialRegs;
unsigned CurrentComponent = 0;
for (; CurrentComponent + 1 < ComponentCount; CurrentComponent += 2) {
Register PopCountResult =
MRI->createVirtualRegister(GR.getRegClass(I64x2Type));
auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
TII.get(SPIRV::OpVectorShuffle))
.addDef(PopCountResult)
.addUse(GR.getSPIRVTypeID(I64x2Type))
.addUse(SrcReg)
.addUse(SrcReg)
.addImm(CurrentComponent)
.addImm(CurrentComponent + 1);
MIB.constrainAllUses(TII, TRI, RBI);
Register SubVecReg =
MRI->createVirtualRegister(GR.getRegClass(Vec2ResType));
if (!CallbackFunction(SubVecReg, Vec2ResType, I, PopCountResult, Opcode))
return false;
PartialRegs.push_back(SubVecReg);
}
// On odd component counts we need to handle one more component
if (CurrentComponent != ComponentCount) {
bool ZeroAsNull = !STI.isShader();
Register FinalElemReg = MRI->createVirtualRegister(GR.getRegClass(I64Type));
Register ConstIntLastIdx = GR.getOrCreateConstInt(
ComponentCount - 1, I, BaseType, TII, ZeroAsNull);
if (!selectOpWithSrcs(FinalElemReg, I64Type, I, {SrcReg, ConstIntLastIdx},
SPIRV::OpVectorExtractDynamic))
return false;
Register FinalElemResReg =
MRI->createVirtualRegister(GR.getRegClass(BaseType));
if (!CallbackFunction(FinalElemResReg, BaseType, I, FinalElemReg, Opcode))
return false;
PartialRegs.push_back(FinalElemResReg);
}
// Join all the resulting registers back into the return type in order
// (ie i32x2, i32x2, i32x1 -> i32x5)
return selectOpWithSrcs(ResVReg, ResType, I, std::move(PartialRegs),
SPIRV::OpCompositeConstruct);
}
bool SPIRVInstructionSelector::selectBitreverseNative(Register ResVReg,
SPIRVTypeInst ResType,
MachineInstr &I,
@@ -5916,83 +5921,6 @@ bool SPIRVInstructionSelector::selectFirstBitSet32(
return true;
}
bool SPIRVInstructionSelector::selectFirstBitSet64Overflow(
Register ResVReg, SPIRVTypeInst ResType, MachineInstr &I, Register SrcReg,
unsigned BitSetOpcode, bool SwapPrimarySide) const {
// SPIR-V allow vectors of size 2,3,4 only. Calling with a larger vectors
// requires creating a param register and return register with an invalid
// vector size. If that is resolved, then this function can be used for
// vectors of any component size.
unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
assert(ComponentCount < 5 && "Vec 5+ will generate invalid SPIR-V ops");
MachineIRBuilder MIRBuilder(I);
SPIRVTypeInst BaseType = GR.retrieveScalarOrVectorIntType(ResType);
SPIRVTypeInst I64Type = GR.getOrCreateSPIRVIntegerType(64, MIRBuilder);
SPIRVTypeInst I64x2Type =
GR.getOrCreateSPIRVVectorType(I64Type, 2, MIRBuilder, false);
SPIRVTypeInst Vec2ResType =
GR.getOrCreateSPIRVVectorType(BaseType, 2, MIRBuilder, false);
std::vector<Register> PartialRegs;
// Loops 0, 2, 4, ... but stops one loop early when ComponentCount is odd
unsigned CurrentComponent = 0;
for (; CurrentComponent + 1 < ComponentCount; CurrentComponent += 2) {
// This register holds the firstbitX result for each of the i64x2 vectors
// extracted from SrcReg
Register BitSetResult =
MRI->createVirtualRegister(GR.getRegClass(I64x2Type));
auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
TII.get(SPIRV::OpVectorShuffle))
.addDef(BitSetResult)
.addUse(GR.getSPIRVTypeID(I64x2Type))
.addUse(SrcReg)
.addUse(SrcReg)
.addImm(CurrentComponent)
.addImm(CurrentComponent + 1);
MIB.constrainAllUses(TII, TRI, RBI);
Register SubVecBitSetReg =
MRI->createVirtualRegister(GR.getRegClass(Vec2ResType));
if (!selectFirstBitSet64(SubVecBitSetReg, Vec2ResType, I, BitSetResult,
BitSetOpcode, SwapPrimarySide))
return false;
PartialRegs.push_back(SubVecBitSetReg);
}
// On odd component counts we need to handle one more component
if (CurrentComponent != ComponentCount) {
bool ZeroAsNull = !STI.isShader();
Register FinalElemReg = MRI->createVirtualRegister(GR.getRegClass(I64Type));
Register ConstIntLastIdx = GR.getOrCreateConstInt(
ComponentCount - 1, I, BaseType, TII, ZeroAsNull);
if (!selectOpWithSrcs(FinalElemReg, I64Type, I, {SrcReg, ConstIntLastIdx},
SPIRV::OpVectorExtractDynamic))
return false;
Register FinalElemBitSetReg =
MRI->createVirtualRegister(GR.getRegClass(BaseType));
if (!selectFirstBitSet64(FinalElemBitSetReg, BaseType, I, FinalElemReg,
BitSetOpcode, SwapPrimarySide))
return false;
PartialRegs.push_back(FinalElemBitSetReg);
}
// Join all the resulting registers back into the return type in order
// (ie i32x2, i32x2, i32x1 -> i32x5)
return selectOpWithSrcs(ResVReg, ResType, I, std::move(PartialRegs),
SPIRV::OpCompositeConstruct);
}
bool SPIRVInstructionSelector::selectFirstBitSet64(
Register ResVReg, SPIRVTypeInst ResType, MachineInstr &I, Register SrcReg,
unsigned BitSetOpcode, bool SwapPrimarySide) const {
@@ -6009,8 +5937,14 @@ bool SPIRVInstructionSelector::selectFirstBitSet64(
// operate on vectors with 2 or less components. When largers vectors are
// seen. Split them, recurse, then recombine them.
if (ComponentCount > 2) {
return selectFirstBitSet64Overflow(ResVReg, ResType, I, SrcReg,
BitSetOpcode, SwapPrimarySide);
auto Func = [this, SwapPrimarySide](Register ResVReg, SPIRVTypeInst ResType,
MachineInstr &I, Register SrcReg,
unsigned Opcode) -> bool {
return this->selectFirstBitSet64(ResVReg, ResType, I, SrcReg, Opcode,
SwapPrimarySide);
};
return handle64BitOverflow(ResVReg, ResType, I, SrcReg, BitSetOpcode, Func);
}
// 1. Split int64 into 2 pieces using a bitcast