[GlobalISel] Add widenScalar support for G_CONCAT_VECTORS and use for legalization v2s32 G_ICMPs

[aemerson]

Summary

Also adds FewerElements support for G_TRUNC. The switch case was present but the type index was always being rejected by the helper.

[arsenm]

llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

165–166

I'm not sure why all of these lines changed?

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

3941–3947

I think this deserves to be split into a separate patch

3945

This doesn't do anything, I assume this should have a NewNarrowTy =

llvm/test/CodeGen/AArch64/GlobalISel/legalize-vector-icmp.mir

1959

Can you add a case with an <8 x s32> icmp result type

[aemerson]

I've decided to use custom legalization for this particular case in D88191, abandoning this change but I might revive it in another form later.

llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

public:
   /// def by inserting a G_BITCAST from \p CastTy
   void bitcastDst(MachineInstr &MI, LLT CastTy, unsigned OpIdx);
 
   /// Widen \p OrigReg to \p WideTy by merging to a wider type, padding with
   /// G_IMPLICIT_DEF, and producing dead results.
   Register widenWithUnmerge(LLT WideTy, Register OrigReg);
 
 private:
-  LegalizeResult
-  widenScalarMergeValues(MachineInstr &MI, unsigned TypeIdx, LLT WideTy);
-  LegalizeResult
-  widenScalarUnmergeValues(MachineInstr &MI, unsigned TypeIdx, LLT WideTy);
-  LegalizeResult
-  widenScalarExtract(MachineInstr &MI, unsigned TypeIdx, LLT WideTy);
-  LegalizeResult
-  widenScalarInsert(MachineInstr &MI, unsigned TypeIdx, LLT WideTy);
-  LegalizeResult
-  widenScalarAddSubShlSat(MachineInstr &MI, unsigned TypeIdx, LLT WideTy);
+  LegalizeResult widenScalarConcatVectors(MachineInstr &MI, unsigned TypeIdx,
+                                          LLT WideTy);
+  LegalizeResult widenScalarMergeValues(MachineInstr &MI, unsigned TypeIdx,
+                                        LLT WideTy);
+  LegalizeResult widenScalarUnmergeValues(MachineInstr &MI, unsigned TypeIdx,
+                                          LLT WideTy);
+  LegalizeResult widenScalarExtract(MachineInstr &MI, unsigned TypeIdx,
+                                    LLT WideTy);
+  LegalizeResult widenScalarInsert(MachineInstr &MI, unsigned TypeIdx,
+                                   LLT WideTy);
+  LegalizeResult widenScalarAddSubShlSat(MachineInstr &MI, unsigned TypeIdx,
+                                         LLT WideTy);
 
   /// Helper function to split a wide generic register into bitwise blocks with
   /// the given Type (which implies the number of blocks needed). The generic
   /// registers created are appended to Ops, starting at bit 0 of Reg.
   void extractParts(Register Reg, LLT Ty, int NumParts,
                     SmallVectorImpl<Register> &VRegs);
 
   /// Version which handles irregular splits.

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

void LegalizerHelper::bitcastDst(MachineInstr &MI, LLT CastTy, unsigned OpIdx) {
   MachineOperand &MO = MI.getOperand(OpIdx);
   Register CastDst = MRI.createGenericVirtualRegister(CastTy);
   MIRBuilder.setInsertPt(MIRBuilder.getMBB(), ++MIRBuilder.getInsertPt());
   MIRBuilder.buildBitcast(MO, CastDst);
   MO.setReg(CastDst);
 }
 
 LegalizerHelper::LegalizeResult
+LegalizerHelper::widenScalarConcatVectors(MachineInstr &MI, unsigned TypeIdx,
+                                          LLT WideTy) {
+  if (TypeIdx != 1)
+    return UnableToLegalize;
+
+  Observer.changingInstr(MI);
+  for (unsigned SrcIdx = 1; SrcIdx < MI.getNumOperands(); ++SrcIdx)
+    widenScalarSrc(MI, WideTy, SrcIdx, TargetOpcode::G_ANYEXT);
+  LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
+  widenScalarDst(MI, DstTy.changeElementType(WideTy.getElementType()));
+  Observer.changedInstr(MI);
+  return Legalized;
+}
+
+LegalizerHelper::LegalizeResult
 LegalizerHelper::widenScalarMergeValues(MachineInstr &MI, unsigned TypeIdx,
                                         LLT WideTy) {
   if (TypeIdx != 1)
     return UnableToLegalize;
 
   Register DstReg = MI.getOperand(0).getReg();
   LLT DstTy = MRI.getType(DstReg);
   if (DstTy.isVector())
LegalizerHelper::widenScalarAddSubShlSat(MachineInstr &MI, unsigned TypeIdx,
   return Legalized;
 }
 
 LegalizerHelper::LegalizeResult
 LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
   switch (MI.getOpcode()) {
   default:
     return UnableToLegalize;
+  case TargetOpcode::G_CONCAT_VECTORS:
+    return widenScalarConcatVectors(MI, TypeIdx, WideTy);
   case TargetOpcode::G_EXTRACT:
     return widenScalarExtract(MI, TypeIdx, WideTy);
   case TargetOpcode::G_INSERT:
     return widenScalarInsert(MI, TypeIdx, WideTy);
   case TargetOpcode::G_MERGE_VALUES:
     return widenScalarMergeValues(MI, TypeIdx, WideTy);
   case TargetOpcode::G_UNMERGE_VALUES:
     return widenScalarUnmergeValues(MI, TypeIdx, WideTy);

 LegalizerHelper::LegalizeResult
 LegalizerHelper::fewerElementsVector(MachineInstr &MI, unsigned TypeIdx,
                                      LLT NarrowTy) {
   using namespace TargetOpcode;
 
   switch (MI.getOpcode()) {
   case G_IMPLICIT_DEF:
     return fewerElementsVectorImplicitDef(MI, TypeIdx, NarrowTy);
-  case G_TRUNC:
+  case G_TRUNC: {
+    LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
+    LLT NewNarrowTy = NarrowTy;
+    if (NarrowTy.isVector())
+      LLT::vector(NarrowTy.getNumElements(), DstTy.getElementType());
+    return reduceOperationWidth(MI, 0, NewNarrowTy);
+  }
   case G_AND:
   case G_OR:
   case G_XOR:
   case G_ADD:
   case G_SUB:
   case G_MUL:
   case G_PTR_ADD:
   case G_SMULH:

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder(G_ICMP)
           },
           0, 1)
       .minScalarOrEltIf(
           [=](const LegalityQuery &Query) { return Query.Types[1] == v2s16; },
           1, s32)
       .minScalarOrEltIf(
           [=](const LegalityQuery &Query) { return Query.Types[1] == v2p0; }, 0,
           s64)
-      .widenScalarOrEltToNextPow2(1);
+      .widenScalarOrEltToNextPow2(1)
+      .clampNumElements(0, v2s32, v4s32);
 
   getActionDefinitionsBuilder(G_FCMP)
       .legalFor({{s32, s32}, {s32, s64}})
       .clampScalar(0, s32, s32)
       .clampScalar(1, s32, s64)
       .widenScalarToNextPow2(1);
 
   // Extensions
if (SrcSize < 8 || !isPowerOf2_32(SrcSize))
       return false;
 
     return true;
   };
   getActionDefinitionsBuilder({G_ZEXT, G_SEXT, G_ANYEXT})
       .legalIf(ExtLegalFunc)
       .clampScalar(0, s64, s64); // Just for s128, others are handled above.
 
-  getActionDefinitionsBuilder(G_TRUNC).alwaysLegal();
+  getActionDefinitionsBuilder(G_TRUNC)
+      .minScalarOrEltIf(
+          [=](const LegalityQuery &Query) { return Query.Types[0].isVector(); },
+          0, s8)
+      .alwaysLegal();
 
   getActionDefinitionsBuilder(G_SEXT_INREG).legalFor({s32, s64}).lower();
 
   // FP conversions
   getActionDefinitionsBuilder(G_FPTRUNC)
       .legalFor(
           {{s16, s32}, {s16, s64}, {s32, s64}, {v4s16, v4s32}, {v2s32, v2s64}})
       .clampMaxNumElements(0, s32, 2);
getActionDefinitionsBuilder(G_SHUFFLE_VECTOR)
       // just want those lowered into G_BUILD_VECTOR
       .lowerIf([=](const LegalityQuery &Query) {
         return !Query.Types[1].isVector();
       })
       .clampNumElements(0, v4s32, v4s32)
       .clampNumElements(0, v2s64, v2s64);
 
   getActionDefinitionsBuilder(G_CONCAT_VECTORS)
-      .legalFor({{v4s32, v2s32}, {v8s16, v4s16}});
+      .legalFor({{v4s32, v2s32}, {v8s16, v4s16}, {v8s8, v4s8}})
+      .clampNumElements(0, v4s32, v4s32)
+      .minScalarOrElt(1, s8);
 
   getActionDefinitionsBuilder(G_JUMP_TABLE).legalFor({{p0}, {s64}});
 
   getActionDefinitionsBuilder(G_BRJT).legalIf([=](const LegalityQuery &Query) {
     return Query.Types[0] == p0 && Query.Types[1] == s64;
   });
 
   getActionDefinitionsBuilder(G_DYN_STACKALLOC).lower();

llvm/test/CodeGen/AArch64/GlobalISel/legalize-vector-icmp.mir

bb.1:
     %0:_(<2 x p0>) = COPY $q0
     %1:_(<2 x p0>) = COPY $q1
     %2:_(<2 x s1>) = G_ICMP intpred(eq), %0(<2 x p0>), %1
     %3:_(<2 x s32>) = G_ANYEXT %2(<2 x s1>)
     $d0 = COPY %3(<2 x s32>)
     RET_ReallyLR implicit $d0
 
 ...
+---
+name:            icmp_2xs32
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$q0' }
+  - { reg: '$q1' }
+  - { reg: '$q2' }
+  - { reg: '$q3' }
+body:             |
+  bb.1:
+    liveins: $q0, $q1, $q2, $q3
+
+    ; CHECK-LABEL: name: icmp_2xs32
+    ; CHECK: liveins: $q0, $q1, $q2, $q3
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(<4 x s32>) = COPY $q1
+    ; CHECK: [[COPY2:%[0-9]+]]:_(<4 x s32>) = COPY $q2
+    ; CHECK: [[COPY3:%[0-9]+]]:_(<4 x s32>) = COPY $q3
+    ; CHECK: [[ICMP:%[0-9]+]]:_(<4 x s32>) = G_ICMP intpred(eq), [[COPY]](<4 x s32>), [[COPY2]]
+    ; CHECK: [[ICMP1:%[0-9]+]]:_(<4 x s32>) = G_ICMP intpred(eq), [[COPY1]](<4 x s32>), [[COPY3]]
+    ; CHECK: [[TRUNC:%[0-9]+]]:_(<4 x s8>) = G_TRUNC [[ICMP]](<4 x s32>)
+    ; CHECK: [[TRUNC1:%[0-9]+]]:_(<4 x s8>) = G_TRUNC [[ICMP1]](<4 x s32>)
+    ; CHECK: [[CONCAT_VECTORS:%[0-9]+]]:_(<8 x s8>) = G_CONCAT_VECTORS [[TRUNC]](<4 x s8>), [[TRUNC1]](<4 x s8>)
+    ; CHECK: $d0 = COPY [[CONCAT_VECTORS]](<8 x s8>)
+    ; CHECK: RET_ReallyLR implicit $d0
+    %2:_(<4 x s32>) = COPY $q0
+    %3:_(<4 x s32>) = COPY $q1
+    %0:_(<8 x s32>) = G_CONCAT_VECTORS %2(<4 x s32>), %3(<4 x s32>)
+    %4:_(<4 x s32>) = COPY $q2
+    %5:_(<4 x s32>) = COPY $q3
+    %1:_(<8 x s32>) = G_CONCAT_VECTORS %4(<4 x s32>), %5(<4 x s32>)
+    %6:_(<8 x s1>) = G_ICMP intpred(eq), %0(<8 x s32>), %1
+    %7:_(<8 x s8>) = G_ANYEXT %6(<8 x s1>)
+    $d0 = COPY %7(<8 x s8>)
+    RET_ReallyLR implicit $d0
+...