AMDGPU/GlobalISel: Legalize G_BUILD_VECTOR v2s16

[arsenm]

Summary

Handle it the same way as G_BUILD_VECTOR_TRUNC. Arguably only
G_BUILD_VECTOR_TRUNC should be legal for this, but G_BUILD_VECTOR will
probably be more convenient in most cases.

[kerbowa]

LGTM

[arsenm]

r371440

lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

getActionDefinitionsBuilder(Op)
           return (LitTy.getScalarSizeInBits() < 16);
         },
         LegalizeMutations::widenScalarOrEltToNextPow2(LitTyIdx, 16))
       .moreElementsIf(isSmallOddVector(BigTyIdx), oneMoreElement(BigTyIdx))
       .widenScalarToNextPow2(BigTyIdx, 32);
 
   }
 
-  getActionDefinitionsBuilder(G_BUILD_VECTOR)
-      .legalForCartesianProduct(AllS32Vectors, {S32})
-      .legalForCartesianProduct(AllS64Vectors, {S64})
-      .clampNumElements(0, V16S32, V16S32)
-      .clampNumElements(0, V2S64, V8S64)
-      .minScalarSameAs(1, 0)
-      .legalIf(isRegisterType(0))
-      .minScalarOrElt(0, S32);
+  auto &BuildVector = getActionDefinitionsBuilder(G_BUILD_VECTOR)
+    .clampNumElements(0, V2S64, V8S64);
+
+  if (ST.hasScalarPackInsts())
+    BuildVector.legalFor({V2S16, S32});
+
+  BuildVector
 
   if (ST.hasScalarPackInsts()) {
     getActionDefinitionsBuilder(G_BUILD_VECTOR_TRUNC)
       .legalFor({V2S16, S32})
       .lower();
   } else {
     getActionDefinitionsBuilder(G_BUILD_VECTOR_TRUNC)
       .lower();

lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

case AMDGPU::G_ZEXT: {
 
     MRI.setRegBank(DstReg, *SrcBank);
     MRI.setRegBank(Ext.getReg(0), *SrcBank);
     MRI.setRegBank(ShiftAmt.getReg(0), *SrcBank);
     MRI.setRegBank(Shl.getReg(0), *SrcBank);
     MI.eraseFromParent();
     return;
   }
+  case AMDGPU::G_BUILD_VECTOR:
   case AMDGPU::G_BUILD_VECTOR_TRUNC: {
+    Register DstReg = MI.getOperand(0).getReg();
+    LLT DstTy = MRI.getType(DstReg);
+    if (DstTy != LLT::vector(2, 16))
+      break;
+
     assert(MI.getNumOperands() == 3 && empty(OpdMapper.getVRegs(0)));
     substituteSimpleCopyRegs(OpdMapper, 1);
     substituteSimpleCopyRegs(OpdMapper, 2);
 
-    Register DstReg = MI.getOperand(0).getReg();
     const RegisterBank *DstBank = getRegBank(DstReg, MRI, *TRI);
     if (DstBank == &AMDGPU::SGPRRegBank)
       break; // Can use S_PACK_* instructions.
 
     MachineIRBuilder B(MI);
 
     Register Lo = MI.getOperand(1).getReg();
     Register Hi = MI.getOperand(2).getReg();
+    const LLT S32 = LLT::scalar(32);
 
     const RegisterBank *BankLo = getRegBank(Lo, MRI, *TRI);
     const RegisterBank *BankHi = getRegBank(Hi, MRI, *TRI);
 
-    const LLT S32 = LLT::scalar(32);
-    auto MaskLo = B.buildConstant(S32, 0xffff);
-    MRI.setRegBank(MaskLo.getReg(0), *BankLo);
+    Register ZextLo;
+    Register ShiftHi;
+
+    if (Opc == AMDGPU::G_BUILD_VECTOR) {
+      ZextLo = B.buildZExt(S32, Lo).getReg(0);
+      MRI.setRegBank(ZextLo, *BankLo);
+
+      Register ZextHi = B.buildZExt(S32, Hi).getReg(0);
+      MRI.setRegBank(ZextHi, *BankHi);
 
       auto ShiftAmt = B.buildConstant(S32, 16);
       MRI.setRegBank(ShiftAmt.getReg(0), *BankHi);
 
-    auto ShiftHi = B.buildShl(S32, Hi, ShiftAmt);
-    MRI.setRegBank(ShiftHi.getReg(0), *BankHi);
+      ShiftHi = B.buildShl(S32, ZextHi, ShiftAmt).getReg(0);
+      MRI.setRegBank(ShiftHi, *BankHi);
+    } else {
+      Register MaskLo = B.buildConstant(S32, 0xffff).getReg(0);
+      MRI.setRegBank(MaskLo, *BankLo);
 
-    auto Masked = B.buildAnd(S32, Lo, MaskLo);
-    MRI.setRegBank(Masked.getReg(0), *BankLo);
+      auto ShiftAmt = B.buildConstant(S32, 16);
+      MRI.setRegBank(ShiftAmt.getReg(0), *BankHi);
+
+      ShiftHi = B.buildShl(S32, Hi, ShiftAmt).getReg(0);
+      MRI.setRegBank(ShiftHi, *BankHi);
+
+      ZextLo = B.buildAnd(S32, Lo, MaskLo).getReg(0);
+      MRI.setRegBank(ZextLo, *BankLo);
+    }
 
-    auto Or = B.buildOr(S32, Masked, ShiftHi);
+    auto Or = B.buildOr(S32, ZextLo, ShiftHi);
     MRI.setRegBank(Or.getReg(0), *DstBank);
 
     B.buildBitcast(DstReg, Or);
     MI.eraseFromParent();
     return;
   }
   case AMDGPU::G_EXTRACT_VECTOR_ELT:
     applyDefaultMapping(OpdMapper);
case AMDGPU::G_EXTRACT: {
     unsigned BankID = getRegBankID(MI.getOperand(1).getReg(), MRI, *TRI);
     unsigned DstSize = getSizeInBits(MI.getOperand(0).getReg(), MRI, *TRI);
     unsigned SrcSize = getSizeInBits(MI.getOperand(1).getReg(), MRI, *TRI);
     OpdsMapping[0] = AMDGPU::getValueMapping(BankID, DstSize);
     OpdsMapping[1] = AMDGPU::getValueMapping(BankID, SrcSize);
     OpdsMapping[2] = nullptr;
     break;
   }
-  case AMDGPU::G_MERGE_VALUES:
   case AMDGPU::G_BUILD_VECTOR:
+  case AMDGPU::G_BUILD_VECTOR_TRUNC: {
+    LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
+    if (DstTy == LLT::vector(2, 16)) {
+      unsigned DstSize = DstTy.getSizeInBits();
+      unsigned SrcSize = MRI.getType(MI.getOperand(1).getReg()).getSizeInBits();
+      unsigned Src0BankID = getRegBankID(MI.getOperand(1).getReg(), MRI, *TRI);
+      unsigned Src1BankID = getRegBankID(MI.getOperand(2).getReg(), MRI, *TRI);
+      unsigned DstBankID = regBankUnion(Src0BankID, Src1BankID);
+
+      OpdsMapping[0] = AMDGPU::getValueMapping(DstBankID, DstSize);
+      OpdsMapping[1] = AMDGPU::getValueMapping(Src0BankID, SrcSize);
+      OpdsMapping[2] = AMDGPU::getValueMapping(Src1BankID, SrcSize);
+      break;
+    }
+
+    LLVM_FALLTHROUGH;
+  }
+  case AMDGPU::G_MERGE_VALUES:
   case AMDGPU::G_CONCAT_VECTORS: {
     unsigned Bank = isSALUMapping(MI) ?
       AMDGPU::SGPRRegBankID : AMDGPU::VGPRRegBankID;
     unsigned DstSize = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
     unsigned SrcSize = MRI.getType(MI.getOperand(1).getReg()).getSizeInBits();
 
     OpdsMapping[0] = AMDGPU::getValueMapping(Bank, DstSize);
     // Op1 and Dst should use the same register bank.
     for (unsigned i = 1, e = MI.getNumOperands(); i != e; ++i)
       OpdsMapping[i] = AMDGPU::getValueMapping(Bank, SrcSize);
     break;
   }
-  case AMDGPU::G_BUILD_VECTOR_TRUNC: {
-    assert(MI.getNumOperands() == 3);
-
-    unsigned DstSize = MRI.getType(MI.getOperand(0).getReg()).getSizeInBits();
-    unsigned SrcSize = MRI.getType(MI.getOperand(1).getReg()).getSizeInBits();
-    unsigned Src0BankID = getRegBankID(MI.getOperand(1).getReg(), MRI, *TRI);
-    unsigned Src1BankID = getRegBankID(MI.getOperand(2).getReg(), MRI, *TRI);
-    unsigned DstBankID = regBankUnion(Src0BankID, Src1BankID);
-
-    OpdsMapping[0] = AMDGPU::getValueMapping(DstBankID, DstSize);
-    OpdsMapping[1] = AMDGPU::getValueMapping(Src0BankID, SrcSize);
-    OpdsMapping[2] = AMDGPU::getValueMapping(Src1BankID, SrcSize);
-    break;
-  }
   case AMDGPU::G_BITCAST:
   case AMDGPU::G_INTTOPTR:
   case AMDGPU::G_PTRTOINT:
   case AMDGPU::G_CTLZ:
   case AMDGPU::G_CTLZ_ZERO_UNDEF:
   case AMDGPU::G_CTTZ:
   case AMDGPU::G_CTTZ_ZERO_UNDEF:
   case AMDGPU::G_CTPOP:

test/CodeGen/AMDGPU/GlobalISel/regbankselect-build-vector.v2s16.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -mcpu=gfx900 -run-pass=regbankselect -regbankselect-fast -verify-machineinstrs -o - %s | FileCheck %s
+# RUN: llc -march=amdgcn -mcpu=gfx900 -run-pass=regbankselect -regbankselect-greedy -verify-machineinstrs -o - %s | FileCheck %s
+
+---
+name: build_vector_v2s16_s32_ss
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0, $sgpr1
+    ; CHECK-LABEL: name: build_vector_v2s16_s32_ss
+    ; CHECK: [[COPY:%[0-9]+]]:sgpr(s32) = COPY $sgpr0
+    ; CHECK: [[COPY1:%[0-9]+]]:sgpr(s32) = COPY $sgpr1
+    ; CHECK: [[TRUNC:%[0-9]+]]:sgpr(s16) = G_TRUNC [[COPY]](s32)
+    ; CHECK: [[TRUNC1:%[0-9]+]]:sgpr(s16) = G_TRUNC [[COPY1]](s32)
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:sgpr(<2 x s16>) = G_BUILD_VECTOR [[TRUNC]](s16), [[TRUNC1]](s16)
+    %0:_(s32) = COPY $sgpr0
+    %1:_(s32) = COPY $sgpr1
+    %2:_(s16) = G_TRUNC %0
+    %3:_(s16) = G_TRUNC %1
+    %4:_(<2 x s16>) = G_BUILD_VECTOR %2, %3
+...
+
+---
+name: build_vector_v2s16_s32_sv
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $sgpr0, $vgpr0
+
+    ; CHECK-LABEL: name: build_vector_v2s16_s32_sv
+    ; CHECK: [[COPY:%[0-9]+]]:sgpr(s32) = COPY $sgpr0
+    ; CHECK: [[COPY1:%[0-9]+]]:vgpr(s32) = COPY $vgpr0
+    ; CHECK: [[TRUNC:%[0-9]+]]:sgpr(s16) = G_TRUNC [[COPY]](s32)
+    ; CHECK: [[TRUNC1:%[0-9]+]]:vgpr(s16) = G_TRUNC [[COPY1]](s32)
+    ; CHECK: [[ZEXT:%[0-9]+]]:sgpr(s32) = G_ZEXT [[TRUNC]](s16)
+    ; CHECK: [[ZEXT1:%[0-9]+]]:vgpr(s32) = G_ZEXT [[TRUNC1]](s16)
+    ; CHECK: [[C:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 16
+    ; CHECK: [[SHL:%[0-9]+]]:vgpr(s32) = G_SHL [[ZEXT1]], [[C]](s32)
+    ; CHECK: [[OR:%[0-9]+]]:vgpr(s32) = G_OR [[ZEXT]], [[SHL]]
+    ; CHECK: [[BITCAST:%[0-9]+]]:vgpr(<2 x s16>) = G_BITCAST [[OR]](s32)
+    %0:_(s32) = COPY $sgpr0
+    %1:_(s32) = COPY $vgpr0
+    %2:_(s16) = G_TRUNC %0
+    %3:_(s16) = G_TRUNC %1
+    %4:_(<2 x s16>) = G_BUILD_VECTOR %2, %3
+...
+
+---
+name: build_vector_v2s16_s32_vs
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0, $sgpr0
+    ; CHECK-LABEL: name: build_vector_v2s16_s32_vs
+    ; CHECK: [[COPY:%[0-9]+]]:vgpr(s32) = COPY $vgpr0
+    ; CHECK: [[COPY1:%[0-9]+]]:sgpr(s32) = COPY $sgpr0
+    ; CHECK: [[TRUNC:%[0-9]+]]:vgpr(s16) = G_TRUNC [[COPY]](s32)
+    ; CHECK: [[TRUNC1:%[0-9]+]]:sgpr(s16) = G_TRUNC [[COPY1]](s32)
+    ; CHECK: [[ZEXT:%[0-9]+]]:vgpr(s32) = G_ZEXT [[TRUNC]](s16)
+    ; CHECK: [[ZEXT1:%[0-9]+]]:sgpr(s32) = G_ZEXT [[TRUNC1]](s16)
+    ; CHECK: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 16
+    ; CHECK: [[SHL:%[0-9]+]]:sgpr(s32) = G_SHL [[ZEXT1]], [[C]](s32)
+    ; CHECK: [[OR:%[0-9]+]]:vgpr(s32) = G_OR [[ZEXT]], [[SHL]]
+    ; CHECK: [[BITCAST:%[0-9]+]]:vgpr(<2 x s16>) = G_BITCAST [[OR]](s32)
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s32) = COPY $sgpr0
+    %2:_(s16) = G_TRUNC %0
+    %3:_(s16) = G_TRUNC %1
+    %4:_(<2 x s16>) = G_BUILD_VECTOR %2, %3
+...
+
+---
+name: build_vector_v2s16_s32_vv
+legalized: true
+
+body: |
+  bb.0:
+    liveins: $vgpr0, $vgpr1
+    ; CHECK-LABEL: name: build_vector_v2s16_s32_vv
+    ; CHECK: [[COPY:%[0-9]+]]:vgpr(s32) = COPY $vgpr0
+    ; CHECK: [[COPY1:%[0-9]+]]:vgpr(s32) = COPY $vgpr1
+    ; CHECK: [[TRUNC:%[0-9]+]]:vgpr(s16) = G_TRUNC [[COPY]](s32)
+    ; CHECK: [[TRUNC1:%[0-9]+]]:vgpr(s16) = G_TRUNC [[COPY1]](s32)
+    ; CHECK: [[ZEXT:%[0-9]+]]:vgpr(s32) = G_ZEXT [[TRUNC]](s16)
+    ; CHECK: [[ZEXT1:%[0-9]+]]:vgpr(s32) = G_ZEXT [[TRUNC1]](s16)
+    ; CHECK: [[C:%[0-9]+]]:vgpr(s32) = G_CONSTANT i32 16
+    ; CHECK: [[SHL:%[0-9]+]]:vgpr(s32) = G_SHL [[ZEXT1]], [[C]](s32)
+    ; CHECK: [[OR:%[0-9]+]]:vgpr(s32) = G_OR [[ZEXT]], [[SHL]]
+    ; CHECK: [[BITCAST:%[0-9]+]]:vgpr(<2 x s16>) = G_BITCAST [[OR]](s32)
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s32) = COPY $vgpr1
+    %2:_(s16) = G_TRUNC %0
+    %3:_(s16) = G_TRUNC %1
+    %4:_(<2 x s16>) = G_BUILD_VECTOR %2, %3
+...