[AMDGPU][GlobalISel] Allow forming s16 U/SBFX pre-regbankselect

Make s16 G_U/SBFX legal and widen them in RegBankSelect.
This allows the set of BFX formation combines to work on s16 types.

Author: Pierre-vh

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

@@ -2069,10 +2069,13 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
       .minScalar(0, S32)
       .lower();
 
+  // Only {S32, S32} or {S32, S64} should ever reach codegen.
+  // We allow S/UBFX for S16 so the combiner can form them before
+  // RegBankSelect, and RegBankSelect will then legalize them correctly.
   getActionDefinitionsBuilder({G_SBFX, G_UBFX})
-      .legalFor({{S32, S32}, {S64, S32}})
-      .clampScalar(1, S32, S32)
-      .clampScalar(0, S32, S64)
+      .legalFor({{S16, S16}, {S32, S32}, {S64, S32}})
+      .clampScalar(1, S16, S32)
+      .clampScalar(0, S16, S64)
       .widenScalarToNextPow2(0)
       .scalarize(0);
 

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

@@ -1485,7 +1485,9 @@ bool AMDGPURegisterBankInfo::applyMappingBFE(MachineIRBuilder &B,
   Register DstReg = MI.getOperand(0).getReg();
   LLT Ty = MRI.getType(DstReg);
 
+  const LLT S64 = LLT::scalar(64);
   const LLT S32 = LLT::scalar(32);
+  const LLT S16 = LLT::scalar(16);
 
   unsigned FirstOpnd = isa<GIntrinsic>(MI) ? 2 : 1;
   Register SrcReg = MI.getOperand(FirstOpnd).getReg();
@@ -1495,6 +1497,18 @@ bool AMDGPURegisterBankInfo::applyMappingBFE(MachineIRBuilder &B,
   const RegisterBank *DstBank =
     OpdMapper.getInstrMapping().getOperandMapping(0).BreakDown[0].RegBank;
   if (DstBank == &AMDGPU::VGPRRegBank) {
+    if (Ty == S16) {
+      ApplyRegBankMapping ApplyBank(B, *this, MRI, &AMDGPU::VGPRRegBank);
+      B.setInsertPt(B.getMBB(), MI);
+      LegalizerHelper Helper(B.getMF(), ApplyBank, B);
+
+      Helper.widenScalarDst(MI, S32);
+      Helper.widenScalarSrc(MI, S32, 1, AMDGPU::G_ANYEXT);
+      Helper.widenScalarSrc(MI, S32, 2, AMDGPU::G_ZEXT);
+      Helper.widenScalarSrc(MI, S32, 3, AMDGPU::G_ZEXT);
+      return true;
+    }
+
     if (Ty == S32)
       return true;
 
@@ -1554,6 +1568,11 @@ bool AMDGPURegisterBankInfo::applyMappingBFE(MachineIRBuilder &B,
 
   ApplyRegBankMapping ApplyBank(B, *this, MRI, &AMDGPU::SGPRRegBank);
 
+  if (Ty == S16) {
+    OffsetReg = B.buildAnyExtOrTrunc(S32, OffsetReg).getReg(0);
+    WidthReg = B.buildAnyExtOrTrunc(S32, WidthReg).getReg(0);
+  }
+
   // Ensure the high bits are clear to insert the offset.
   auto OffsetMask = B.buildConstant(S32, maskTrailingOnes<unsigned>(6));
   auto ClampOffset = B.buildAnd(S32, OffsetReg, OffsetMask);
@@ -1568,13 +1587,21 @@ bool AMDGPURegisterBankInfo::applyMappingBFE(MachineIRBuilder &B,
 
   // TODO: It might be worth using a pseudo here to avoid scc clobber and
   // register class constraints.
-  unsigned Opc = Ty == S32 ? (Signed ? AMDGPU::S_BFE_I32 : AMDGPU::S_BFE_U32) :
-                             (Signed ? AMDGPU::S_BFE_I64 : AMDGPU::S_BFE_U64);
+  unsigned Opc = (Ty != S64) ? (Signed ? AMDGPU::S_BFE_I32 : AMDGPU::S_BFE_U32)
+                             : (Signed ? AMDGPU::S_BFE_I64 : AMDGPU::S_BFE_U64);
 
-  auto MIB = B.buildInstr(Opc, {DstReg}, {SrcReg, MergedInputs});
+  Register BFEDst = DstReg;
+  if (Ty == S16) {
+    BFEDst = MRI.createGenericVirtualRegister(S32);
+    MRI.setRegBank(BFEDst, AMDGPU::SGPRRegBank);
+  }
+  auto MIB = B.buildInstr(Opc, {BFEDst}, {SrcReg, MergedInputs});
   if (!constrainSelectedInstRegOperands(*MIB, *TII, *TRI, *this))
     llvm_unreachable("failed to constrain BFE");
 
+  if (BFEDst != DstReg)
+    B.buildZExtOrTrunc(DstReg, BFEDst);
+
   MI.eraseFromParent();
   return true;
 }