[llvm] AMDGPU/GlobalISel: RegBankLegalize rules for G_INSERT_VECTOR_ELT (PR #191296)

[Petar Avramovic via llvm-commits]

================
@@ -979,6 +979,109 @@ bool RegBankLegalizeHelper::lowerExtrVecEltTo32(MachineInstr &MI) {
   return true;
 }
 
+bool RegBankLegalizeHelper::lowerInsVecEltToSel(MachineInstr &MI) {
+  // Lower insert vector element to a compare-select chain:
+  //   for i in 0..N-1:
+  //     result[i] = (idx == i) ? elt : srcVec[i]
+  //   dst = merge(result[0..N-1])
+  //
+  // VGPR B64 requires splitting to lo/hi s32 pairs since there is no
+  // v_cndmask_b64. SGPR B64/B32 and VGPR B32 can be handled natively.
+  Register Dst = MI.getOperand(0).getReg();
+  Register Src = MI.getOperand(1).getReg();
+  Register Elt = MI.getOperand(2).getReg();
+  Register Idx = MI.getOperand(3).getReg();
+
+  LLT VecTy = MRI.getType(Src);
+  LLT ScalarTy = VecTy.getScalarType();
+  unsigned NumElts = VecTy.getNumElements();
+  bool IsSGPR = (MRI.getRegBank(Src) == SgprRB);
+  SmallVector<Register, 16> Selects;
+
+  MachineRegisterInfo::VRegAttrs SrcRB_EltTy = {MRI.getRegBank(Src), ScalarTy};
+  MachineRegisterInfo::VRegAttrs CmpTy = IsSGPR ? SgprRB_S32 : VccRB_S1;
+  MachineRegisterInfo::VRegAttrs SgprRB_IdxTy = {SgprRB, MRI.getType(Idx)};
+
+  if (!IsSGPR && ScalarTy.getSizeInBits() == 64) {
+    // VGPR B64: split to 32-bit lo/hi since there is no v_cndmask_b64.
+    auto Unmerge = B.buildUnmerge(VgprRB_S32, Src);
+    auto EltUnmerge = B.buildUnmerge(VgprRB_S32, Elt);
+    for (unsigned I = 0; I < NumElts; ++I) {
+      auto IdxConst = B.buildConstant(SgprRB_IdxTy, I);
+      auto Cmp = B.buildICmp(CmpInst::ICMP_EQ, VccRB_S1, Idx, IdxConst);
+      Selects.push_back(B.buildSelect(VgprRB_S32, Cmp, EltUnmerge.getReg(0),
+                                      Unmerge.getReg(2 * I))
+                            .getReg(0));
+      Selects.push_back(B.buildSelect(VgprRB_S32, Cmp, EltUnmerge.getReg(1),
+                                      Unmerge.getReg(2 * I + 1))
+                            .getReg(0));
----------------
petar-avramovic wrote:

Can put EltUnmerge.getReg(0) in local variable to avoid getting it each iteration

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