[llvm] [SPIRV]Implementing PopCount for 16 and 64 bits (PR #191283)
Farzon Lotfi via llvm-commits
llvm-commits at lists.llvm.org
Fri Apr 10 13:01:43 PDT 2026
=?utf-8?q?João?= Saffran <joaosaffranllvm at gmail.com>,Joao Saffran
<joaosaffranllvm at gmail.com>,Joao Saffran <joaosaffranllvm at gmail.com>,Joao
Saffran <joaosaffranllvm at gmail.com>,Joao Saffran <joaosaffranllvm at gmail.com>,Joao
Saffran <joaosaffranllvm at gmail.com>,Joao Saffran <joaosaffranllvm at gmail.com>,Joao
Saffran <joaosaffranllvm at gmail.com>,Joao Saffran <joaosaffranllvm at gmail.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/191283 at github.com>
================
@@ -1559,6 +1587,213 @@ bool SPIRVInstructionSelector::selectOpWithSrcs(Register ResVReg,
return true;
}
+bool SPIRVInstructionSelector::selectPopCount16(Register ResVReg,
+ SPIRVTypeInst ResType,
+ MachineInstr &I,
+ unsigned ExtOpcode,
+ unsigned Opcode) const {
+ Register OpReg = I.getOperand(1).getReg();
+ SPIRVTypeInst OpRegType = GR.getSPIRVTypeForVReg(OpReg);
+
+ unsigned NumElems = getVectorSizeOrOne(OpRegType);
+
+ MachineIRBuilder MIRBuilder(I);
+ SPIRVTypeInst I32Type = GR.getOrCreateSPIRVIntegerType(32, MIRBuilder);
+ SPIRVTypeInst I32VectorType =
+ GR.getOrCreateSPIRVVectorType(I32Type, NumElems, MIRBuilder, false);
+
+ bool IsVector = NumElems > 1;
+ SPIRVTypeInst ExtType = IsVector ? I32VectorType : I32Type;
+ Register ExtReg = MRI->createVirtualRegister(GR.getRegClass(ResType));
+ if (!selectOpWithSrcs(ExtReg, ExtType, I, {OpReg}, ExtOpcode))
+ return false;
+
+ Register PopCountReg = MRI->createVirtualRegister(GR.getRegClass(ExtType));
+ if (!selectPopCount32(PopCountReg, ExtType, I, ExtReg, Opcode))
+ return false;
+
+ return selectOpWithSrcs(ResVReg, ResType, I, {PopCountReg}, ExtOpcode);
+}
+
+bool SPIRVInstructionSelector::selectPopCount32(Register ResVReg,
+ SPIRVTypeInst ResType,
+ MachineInstr &I,
+ Register SrcReg,
+ unsigned Opcode) const {
+ 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,
+ Register SrcReg,
+ unsigned Opcode) const {
+ unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
+ if (ComponentCount > 2)
+ return selectPopCount64Overflow(ResVReg, ResType, I, SrcReg, Opcode);
+
+ bool ZeroAsNull = !STI.isShader();
+
+ MachineIRBuilder MIRBuilder(I);
+ SPIRVTypeInst I32Type = GR.getOrCreateSPIRVIntegerType(32, MIRBuilder);
+ SPIRVTypeInst VecI32Type = GR.getOrCreateSPIRVVectorType(
+ I32Type, 2 * ComponentCount, MIRBuilder, false);
+
+ Register BitcastReg = MRI->createVirtualRegister(GR.getRegClass(VecI32Type));
+ if (!selectOpWithSrcs(BitcastReg, VecI32Type, I, {SrcReg}, SPIRV::OpBitcast))
+ return false;
+
+ Register PopCountReg = MRI->createVirtualRegister(GR.getRegClass(VecI32Type));
+ if (!selectPopCount32(PopCountReg, VecI32Type, I, BitcastReg, Opcode))
+ return false;
+
+ bool IsScalarRes = ResType->getOpcode() != SPIRV::OpTypeVector;
+
+ Register HighReg, LowReg;
+ SPIRVTypeInst PartsType;
+
+ Register ConstIntZero =
+ GR.getOrCreateConstInt(0, I, I32Type, TII, ZeroAsNull);
+ Register ConstIntOne = GR.getOrCreateConstInt(1, I, I32Type, TII, ZeroAsNull);
+
+ if (IsScalarRes) {
+ PartsType = I32Type;
+ HighReg = MRI->createVirtualRegister(GR.getRegClass(PartsType));
+ LowReg = MRI->createVirtualRegister(GR.getRegClass(PartsType));
+
+ if (!selectOpWithSrcs(HighReg, I32Type, I, {PopCountReg, ConstIntOne},
+ SPIRV::OpVectorExtractDynamic))
+ return false;
+ if (!selectOpWithSrcs(LowReg, I32Type, I, {PopCountReg, ConstIntZero},
+ SPIRV::OpVectorExtractDynamic))
+ return false;
+ } else {
+ PartsType = GR.getOrCreateSPIRVVectorType(I32Type, ComponentCount,
+ MIRBuilder, false);
+
+ HighReg = MRI->createVirtualRegister(GR.getRegClass(PartsType));
+ LowReg = MRI->createVirtualRegister(GR.getRegClass(PartsType));
+
+ auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
+ TII.get(SPIRV::OpVectorShuffle))
+ .addDef(HighReg)
+ .addUse(GR.getSPIRVTypeID(PartsType))
+ .addUse(PopCountReg)
+ .addUse(PopCountReg);
+ for (unsigned J = 1; J < ComponentCount * 2; J += 2)
+ MIB.addImm(J);
+ MIB.constrainAllUses(TII, TRI, RBI);
+
+ MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
+ TII.get(SPIRV::OpVectorShuffle))
+ .addDef(LowReg)
+ .addUse(GR.getSPIRVTypeID(PartsType))
+ .addUse(PopCountReg)
+ .addUse(PopCountReg);
+ for (unsigned J = 0; J < ComponentCount * 2; J += 2)
+ MIB.addImm(J);
+ MIB.constrainAllUses(TII, TRI, RBI);
+ }
----------------
farzonl wrote:
I just threw my concerns into claude and it gave this suggestion, so take it with a grain of salt. Didn't independly test this code, but I think it reads more clearly as to the intent.
```cpp
bool SPIRVInstructionSelector::selectPopCount64(Register ResVReg,
SPIRVTypeInst ResType,
MachineInstr &I,
Register SrcReg,
unsigned Opcode) const {
unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
if (ComponentCount > 2)
return selectPopCount64Overflow(ResVReg, ResType, I, SrcReg, Opcode);
MachineIRBuilder MIRBuilder(I);
bool ZeroAsNull = !STI.isShader();
// ---- Types ----
SPIRVTypeInst I32Type =
GR.getOrCreateSPIRVIntegerType(32, MIRBuilder);
SPIRVTypeInst VecI32Type =
GR.getOrCreateSPIRVVectorType(I32Type, 2 * ComponentCount,
MIRBuilder, false);
// ---- Stage 1: 64-bit → 2x32-bit ----
Register Vec32 = createVReg(VecI32Type);
if (!selectOpWithSrcs(Vec32, VecI32Type, I, {SrcReg}, SPIRV::OpBitcast))
return false;
// ---- Stage 2: popcount per 32-bit lane ----
Register Pop32 = createVReg(VecI32Type);
if (!selectPopCount32(Pop32, VecI32Type, I, Vec32, Opcode))
return false;
// ---- Stage 3: split even/odd lanes ----
SplitParts Parts =
splitEvenOddLanes(Pop32, ResType, ComponentCount, I, I32Type);
// ---- Stage 4: sum high + low ----
unsigned OpAdd = Parts.IsScalar ? SPIRV::OpIAddS : SPIRV::OpIAddV;
Register Sum = createVReg(Parts.Type);
if (!selectOpWithSrcs(Sum, Parts.Type, I,
{Parts.High, Parts.Low}, OpAdd))
return false;
// ---- Stage 5: convert to final type ----
bool IsSigned = GR.isScalarOrVectorSigned(Parts.Type);
unsigned ConvOp = IsSigned ? SPIRV::OpSConvert : SPIRV::OpUConvert;
return selectOpWithSrcs(ResVReg, ResType, I, {Sum}, ConvOp);
}
```
https://github.com/llvm/llvm-project/pull/191283
More information about the llvm-commits
mailing list