[clang] [llvm] [HLSL] Implement elementwise firstbitlow builtin (PR #116858)

[Ashley Coleman via llvm-commits]

================
@@ -3178,98 +3178,74 @@ bool SPIRVInstructionSelector::selectFirstBitSet64Overflow(
     Register ResVReg, const SPIRVType *ResType, MachineInstr &I,
     Register SrcReg, unsigned BitSetOpcode, bool SwapPrimarySide) const {
 
+  // SPIR-V only allow vecs of size 2,3,4. Calling with a larger vec requires
+  // creating a param reg and return reg with an invalid vec size. If that is
+  // resolved then this function is valid for vectors of any component size.
   unsigned ComponentCount = GR.getScalarOrVectorComponentCount(ResType);
-  SPIRVType *BaseType = GR.retrieveScalarOrVectorIntType(ResType);
-  bool ZeroAsNull = STI.isOpenCLEnv();
-  Register ConstIntZero =
-      GR.getOrCreateConstInt(0, I, BaseType, TII, ZeroAsNull);
-  unsigned LeftComponentCount = ComponentCount / 2;
-  unsigned RightComponentCount = ComponentCount - LeftComponentCount;
-  bool LeftIsVector = LeftComponentCount > 1;
+  assert(ComponentCount < 5 && "Vec 5+ will generate invalid SPIR-V ops");
 
-  // Split the SrcReg in half into 2 smaller vec registers
-  // (ie i64x4 -> i64x2, i64x2)
+  bool ZeroAsNull = STI.isOpenCLEnv();
   MachineIRBuilder MIRBuilder(I);
-  SPIRVType *OpType = GR.getOrCreateSPIRVIntegerType(64, MIRBuilder);
-  SPIRVType *LeftVecOpType;
-  SPIRVType *LeftVecResType;
-  if (LeftIsVector) {
-    LeftVecOpType =
-        GR.getOrCreateSPIRVVectorType(OpType, LeftComponentCount, MIRBuilder);
-    LeftVecResType =
-        GR.getOrCreateSPIRVVectorType(BaseType, LeftComponentCount, MIRBuilder);
-  } else {
-    LeftVecOpType = OpType;
-    LeftVecResType = BaseType;
-  }
-
-  SPIRVType *RightVecOpType =
-      GR.getOrCreateSPIRVVectorType(OpType, RightComponentCount, MIRBuilder);
-  SPIRVType *RightVecResType =
-      GR.getOrCreateSPIRVVectorType(BaseType, RightComponentCount, MIRBuilder);
+  SPIRVType *BaseType = GR.retrieveScalarOrVectorIntType(ResType);
+  SPIRVType *I64Type = GR.getOrCreateSPIRVIntegerType(64, MIRBuilder);
+  SPIRVType *I64x2Type = GR.getOrCreateSPIRVVectorType(I64Type, 2, MIRBuilder);
+  SPIRVType *Vec2ResType =
+      GR.getOrCreateSPIRVVectorType(BaseType, 2, MIRBuilder);
 
-  Register LeftSideIn =
-      MRI->createVirtualRegister(GR.getRegClass(LeftVecOpType));
-  Register RightSideIn =
-      MRI->createVirtualRegister(GR.getRegClass(RightVecOpType));
+  std::vector<Register> PartialRegs;
 
-  bool Result;
+  // Loops 0, 2, 4, ... but stops one loop early when ComponentCount is odd
+  unsigned CurrentComponent = 0;
+  for (; CurrentComponent + 1 < ComponentCount; CurrentComponent += 2) {
+    Register SubVecReg = MRI->createVirtualRegister(GR.getRegClass(I64x2Type));
 
-  // Extract the left half from the SrcReg into LeftSideIn
-  // accounting for the special case when it only has one element
-  if (LeftIsVector) {
     auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
                        TII.get(SPIRV::OpVectorShuffle))
-                   .addDef(LeftSideIn)
-                   .addUse(GR.getSPIRVTypeID(LeftVecOpType))
+                   .addDef(SubVecReg)
+                   .addUse(GR.getSPIRVTypeID(I64x2Type))
                    .addUse(SrcReg)
                    // Per the spec, repeat the vector if only one vec is needed
                    .addUse(SrcReg);
 
-    for (unsigned J = 0; J < LeftComponentCount; J++) {
-      MIB.addImm(J);
-    }
+    MIB.addImm(CurrentComponent);
+    MIB.addImm(CurrentComponent + 1);
 
-    Result = MIB.constrainAllUses(TII, TRI, RBI);
-  } else {
-    Result =
-        selectOpWithSrcs(LeftSideIn, LeftVecOpType, I, {SrcReg, ConstIntZero},
-                         SPIRV::OpVectorExtractDynamic);
-  }
+    if (!MIB.constrainAllUses(TII, TRI, RBI))
+      return false;
 
-  // Extract the right half from the SrcReg into RightSideIn.
-  // Right will always be a vector since the only time one element is left is
-  // when Component == 3, and in that case Left is one element.
-  auto MIB = BuildMI(*I.getParent(), I, I.getDebugLoc(),
-                     TII.get(SPIRV::OpVectorShuffle))
-                 .addDef(RightSideIn)
-                 .addUse(GR.getSPIRVTypeID(RightVecOpType))
-                 .addUse(SrcReg)
-                 // Per the spec, repeat the vector if only one vec is needed
-                 .addUse(SrcReg);
+    Register SubVecBitSetReg =
----------------
V-FEXrt wrote:

uhh I think so? Its hard to give this thing a precise name.

This is the register that holds the firstbitlow/high result for each of the sub/split vectors extracted from the input vector

https://github.com/llvm/llvm-project/pull/116858