AMDGPU/GlobalISel: Fix splitting 64-bit extensions

[arsenm]

Summary

This was replicating the low bits into the high bits, rather than
applying the appropriate extension type.

[foad]

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

2315–2317

Surely the src is always i1 here, so the high part of dst will always be 0 (or undef for anyext)?

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

2315–2317

Not if it’s G_SEXT

[foad]

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

2315–2317

OK, but if it's G_ZEXT the high part is zero and if it's G_SEXT the high part is the same as the low part (and if it's G_ANYEXT the high part is undef). So in no case do you need either a shift or another select.

[arsenm]

Special case extend from s1 sext case. The zext case emits a constant 0, but I think this depends on whether it's using the CSEMIRbuilder or not. There's not much difference between copy of 0 or second copy of 0

[arsenm]

e8f6b0e583cbf97a08c26b2219faaa5e433a67ab

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

static void reinsertVectorIndexAdd(MachineIRBuilder &B,
   auto MaterializedOffset = B.buildConstant(S32, ConstOffset);
 
   auto Add = B.buildAdd(S32, WaterfallIdx, MaterializedOffset);
   MRI.setRegBank(MaterializedOffset.getReg(0), AMDGPU::SGPRRegBank);
   MRI.setRegBank(Add.getReg(0), AMDGPU::SGPRRegBank);
   IdxUseInstr.getOperand(OpIdx).setReg(Add.getReg(0));
 }
 
+/// Implement extending a 32-bit value to a 64-bit value. \p Lo32Reg is the
+/// original 32-bit source value (to be inserted in the low part of the combined
+/// 64-bit result), and \p Hi32Reg is the high half of the combined 64-bit
+/// value.
+static void extendLow32IntoHigh32(MachineIRBuilder &B,
+                                  Register Hi32Reg, Register Lo32Reg,
+                                  unsigned ExtOpc,
+                                  const RegisterBank &RegBank,
+                                  bool IsBooleanSrc = false) {
+  if (ExtOpc == AMDGPU::G_ZEXT) {
+    B.buildConstant(Hi32Reg, 0);
+  } else if (ExtOpc == AMDGPU::G_SEXT) {
+    if (IsBooleanSrc) {
+      // If we know the original source was an s1, the high half is the same as
+      // the low.
+      B.buildCopy(Hi32Reg, Lo32Reg);
+    } else {
+      // Replicate sign bit from 32-bit extended part.
+      auto ShiftAmt = B.buildConstant(LLT::scalar(32), 31);
+      B.getMRI()->setRegBank(ShiftAmt.getReg(0), RegBank);
+      B.buildAShr(Hi32Reg, Lo32Reg, ShiftAmt);
+    }
+  } else {
+    assert(ExtOpc == AMDGPU::G_ANYEXT && "not an integer extension");
+    B.buildUndef(Hi32Reg);
+  }
+}
+
 void AMDGPURegisterBankInfo::applyMappingImpl(
     const OperandsMapper &OpdMapper) const {
   MachineInstr &MI = OpdMapper.getMI();
   unsigned Opc = MI.getOpcode();
   MachineRegisterInfo &MRI = OpdMapper.getMRI();
   switch (Opc) {
   case AMDGPU::G_PHI: {
     Register DstReg = MI.getOperand(0).getReg();
case AMDGPU::G_ANYEXT: {
     LLT DstTy = MRI.getType(DstReg);
     if (DstTy.isScalar() &&
         SrcBank != &AMDGPU::SGPRRegBank &&
         SrcBank != &AMDGPU::VCCRegBank &&
         // FIXME: Should handle any type that round to s64 when irregular
         // breakdowns supported.
         DstTy.getSizeInBits() == 64 &&
         SrcTy.getSizeInBits() <= 32) {
-      const LLT S32 = LLT::scalar(32);
       SmallVector<Register, 2> DefRegs(OpdMapper.getVRegs(0));
 
       // Extend to 32-bit, and then extend the low half.
       if (Signed) {
         // TODO: Should really be buildSExtOrCopy
         B.buildSExtOrTrunc(DefRegs[0], SrcReg);
-
-        // Replicate sign bit from 32-bit extended part.
-        auto ShiftAmt = B.buildConstant(S32, 31);
-        MRI.setRegBank(ShiftAmt.getReg(0), *SrcBank);
-        B.buildAShr(DefRegs[1], DefRegs[0], ShiftAmt);
       } else if (Opc == AMDGPU::G_ZEXT) {
         B.buildZExtOrTrunc(DefRegs[0], SrcReg);
-        B.buildConstant(DefRegs[1], 0);
       } else {
         B.buildAnyExtOrTrunc(DefRegs[0], SrcReg);
-        B.buildUndef(DefRegs[1]);
       }
 
+      extendLow32IntoHigh32(B, DefRegs[1], DefRegs[0], Opc, *SrcBank);
       MRI.setRegBank(DstReg, *SrcBank);
       MI.eraseFromParent();
       return;
     }
 
     if (SrcTy != LLT::scalar(1))
       return;
 
     // It is not legal to have a legalization artifact with a VCC source. Rather
-    // than introducing a copy, insert the selcet we would have to select the
+    // than introducing a copy, insert the select we would have to select the
     // copy to.
     if (SrcBank == &AMDGPU::VCCRegBank) {
       SmallVector<Register, 2> DefRegs(OpdMapper.getVRegs(0));
 
       const RegisterBank *DstBank = &AMDGPU::VGPRRegBank;
 
       unsigned DstSize = DstTy.getSizeInBits();
       // 64-bit select is SGPR only
       const bool UseSel64 = DstSize > 32 &&
         SrcBank->getID() == AMDGPU::SGPRRegBankID;
 
       // TODO: Should s16 select be legal?
       LLT SelType = UseSel64 ? LLT::scalar(64) : LLT::scalar(32);
       auto True = B.buildConstant(SelType, Signed ? -1 : 1);
       auto False = B.buildConstant(SelType, 0);
 
       MRI.setRegBank(True.getReg(0), *DstBank);
       MRI.setRegBank(False.getReg(0), *DstBank);
       MRI.setRegBank(DstReg, *DstBank);
 
       if (DstSize > 32) {
         B.buildSelect(DefRegs[0], SrcReg, True, False);
-        B.buildCopy(DefRegs[1], DefRegs[0]);
+        extendLow32IntoHigh32(B, DefRegs[1], DefRegs[0], Opc, *SrcBank, true);
       } else if (DstSize < 32) {
         auto Sel = B.buildSelect(SelType, SrcReg, True, False);
         MRI.setRegBank(Sel.getReg(0), *DstBank);
         B.buildTrunc(DstReg, Sel);
       } else {
         B.buildSelect(DstReg, SrcReg, True, False);
       }
 
       MI.eraseFromParent();
       return;
     }
 

llvm/test/CodeGen/AMDGPU/GlobalISel/regbankselect-anyext.mir

bb.0:
     liveins: $vgpr0, $vgpr1
     ; CHECK-LABEL: name: anyext_s1_to_s64_vcc
     ; CHECK: [[COPY:%[0-9]+]]:vgpr(s32) = COPY $vgpr0
     ; CHECK: [[COPY1:%[0-9]+]]:vgpr(s32) = COPY $vgpr1
     ; CHECK: [[ICMP:%[0-9]+]]:vcc(s1) = G_ICMP intpred(eq), [[COPY]](s32), [[COPY1]]
     ; CHECK: [[C:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 1
     ; CHECK: [[C1:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 0
     ; CHECK: [[SELECT:%[0-9]+]]:vgpr(s32) = G_SELECT [[ICMP]](s1), [[C]], [[C1]]
-    ; CHECK: [[COPY2:%[0-9]+]]:vgpr(s32) = COPY [[SELECT]](s32)
-    ; CHECK: [[MV:%[0-9]+]]:vgpr(s64) = G_MERGE_VALUES [[SELECT]](s32), [[COPY2]](s32)
+    ; CHECK: [[DEF:%[0-9]+]]:vgpr(s32) = G_IMPLICIT_DEF
+    ; CHECK: [[MV:%[0-9]+]]:vgpr(s64) = G_MERGE_VALUES [[SELECT]](s32), [[DEF]](s32)
     %0:_(s32) = COPY $vgpr0
     %1:_(s32) = COPY $vgpr1
     %2:_(s1) = G_ICMP intpred(eq), %0, %1
     %3:_(s64) = G_ANYEXT %2
 ...
 
 ---
 name: anyext_s1_to_s16_sgpr

llvm/test/CodeGen/AMDGPU/GlobalISel/regbankselect-zext.mir

bb.0:
     liveins: $vgpr0, $vgpr1
     ; CHECK-LABEL: name: zext_s1_to_s64_vcc
     ; CHECK: [[COPY:%[0-9]+]]:vgpr(s32) = COPY $vgpr0
     ; CHECK: [[COPY1:%[0-9]+]]:vgpr(s32) = COPY $vgpr1
     ; CHECK: [[ICMP:%[0-9]+]]:vcc(s1) = G_ICMP intpred(eq), [[COPY]](s32), [[COPY1]]
     ; CHECK: [[C:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 1
     ; CHECK: [[C1:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 0
     ; CHECK: [[SELECT:%[0-9]+]]:vgpr(s32) = G_SELECT [[ICMP]](s1), [[C]], [[C1]]
-    ; CHECK: [[COPY2:%[0-9]+]]:vgpr(s32) = COPY [[SELECT]](s32)
-    ; CHECK: [[MV:%[0-9]+]]:vgpr(s64) = G_MERGE_VALUES [[SELECT]](s32), [[COPY2]](s32)
+    ; CHECK: [[C2:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 0
+    ; CHECK: [[MV:%[0-9]+]]:vgpr(s64) = G_MERGE_VALUES [[SELECT]](s32), [[C2]](s32)
     %0:_(s32) = COPY $vgpr0
     %1:_(s32) = COPY $vgpr1
     %2:_(s1) = G_ICMP intpred(eq), %0, %1
     %3:_(s64) = G_ZEXT %2
 ...
 
 ---
 name: zext_s1_to_s16_sgpr