[AArch64][GlobalISel] Legalize G_SHUFFLE_VECTOR with different lengths

[dzhidzhoev]

Summary

Legalize G_SHUFFLE_VECTOR having destination vector length greater than
source vector length by reshaping source vectors.

Partial implementation of SelectionDAGBuilder::visitShuffleVector.

[dzhidzhoev]

Rebase

[dzhidzhoev]

Rebase & ping

[paquette]

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

951

might want a TODO here so we can see it ought to be supported?

958

can we have a comment here which, at a high level, describes the general legalization algorithm/approach?

971

static_cast for greppability?

[aemerson]

I’m on vacation so haven’t done a review, but this seems to be something that needs to be done generically for all targets. I was considering doing this as some sort of canonicalization phase during translation but wasn’t decided.

@arsenm @foad any opinions on where we should do this?

[foad]

this seems to be something that needs to be done generically for all targets

It seems like something that could be available in LegalizerHelper for targets that want to use it. I don't see why it is necessarily required for all targets.

[aemerson]

In D132190#3758013, @foad wrote:

this seems to be something that needs to be done generically for all targets

It seems like something that could be available in LegalizerHelper for targets that want to use it. I don't see why it is necessarily required for all targets.

In that case I wonder if this could be done as a Lower action as an intermediate step before the legalizer tries to scalarize the shuffle.

[aemerson]

In D132190#3763576, @aemerson wrote:

In D132190#3758013, @foad wrote:

this seems to be something that needs to be done generically for all targets

It seems like something that could be available in LegalizerHelper for targets that want to use it. I don't see why it is necessarily required for all targets.

In that case I wonder if this could be done as a Lower action as an intermediate step before the legalizer tries to scalarize the shuffle.

But on second thought, doing so would probably make the target's legalizer rules more complex since we'd have to check for non-canonical forms with the legalizer predicates.

[arsenm]

This is a similarly complicated case to unaligned load/store access. My plan for that is to move all the logic into a generic lower function in LegalizerHelper, possibly with different implementation helper functions. However, the generic lower pass would be reliant on calling into some kind of TargetLowering hook to make decisions based on legal masks

[arsenm]

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

936

Nothing here really looks target specific, can this move to an implementation in LegalizerHelper? This looks like fewerElements to me

[dzhidzhoev]

Moved the implementation to LegalizerHelper.

[arsenm]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

4995 ↗

(On Diff #482278)

Maybe use small size 16?

5000 ↗

(On Diff #482278)

If we don't have a MIRBuilder helper to create an extract with a constant index already, there should be one

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

437

Some odd number of element cases would be nice

[dzhidzhoev]

Addressed @arsenm comments.

[kda]

This has broken UBSAN buildbot:
https://lab.llvm.org/buildbot/#/builders/85/builds/13003

I reverted it locally and it passed...
I will revert in 1 hour.

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

private:
   bool legalizeShlAshrLshr(MachineInstr &MI, MachineRegisterInfo &MRI,
                            MachineIRBuilder &MIRBuilder,
                            GISelChangeObserver &Observer) const;
 
   bool legalizeSmallCMGlobalValue(MachineInstr &MI, MachineRegisterInfo &MRI,
                                   MachineIRBuilder &MIRBuilder,
                                   GISelChangeObserver &Observer) const;
   bool legalizeVectorTrunc(MachineInstr &MI, LegalizerHelper &Helper) const;
+  bool legalizeShuffleVector(MachineInstr &MI, LegalizerHelper &Helper) const;
   bool legalizeBitfieldExtract(MachineInstr &MI, MachineRegisterInfo &MRI,
                                LegalizerHelper &Helper) const;
   bool legalizeRotate(MachineInstr &MI, MachineRegisterInfo &MRI,
                       LegalizerHelper &Helper) const;
   bool legalizeCTPOP(MachineInstr &MI, MachineRegisterInfo &MRI,
                      LegalizerHelper &Helper) const;
   bool legalizeAtomicCmpxchg128(MachineInstr &MI, MachineRegisterInfo &MRI,
                                 LegalizerHelper &Helper) const;
   bool legalizeCTTZ(MachineInstr &MI, LegalizerHelper &Helper) const;
   bool legalizeMemOps(MachineInstr &MI, LegalizerHelper &Helper) const;
   const AArch64Subtarget *ST;
 };
 } // End llvm namespace.
 #endif

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

AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
 
   getActionDefinitionsBuilder(G_CTTZ)
       .lowerIf(isVector(0))
       .clampScalar(0, s32, s64)
       .scalarSameSizeAs(1, 0)
       .customFor({s32, s64});
 
   getActionDefinitionsBuilder(G_SHUFFLE_VECTOR)
+      .customIf([=](const LegalityQuery &Query) {
+        return Query.Types[0].isVector() && Query.Types[1].isVector() &&
+               Query.Types[0].getNumElements() >
+                   Query.Types[1].getNumElements();
+      })
       .legalIf([=](const LegalityQuery &Query) {
         const LLT &DstTy = Query.Types[0];
         const LLT &SrcTy = Query.Types[1];
         // For now just support the TBL2 variant which needs the source vectors
         // to be the same size as the dest.
         if (DstTy != SrcTy)
           return false;
         return llvm::is_contained({v2s32, v4s32, v2s64, v2p0, v16s8, v8s16},
bool AArch64LegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
                                           MachineInstr &MI) const {
   MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
   MachineRegisterInfo &MRI = *MIRBuilder.getMRI();
   GISelChangeObserver &Observer = Helper.Observer;
   switch (MI.getOpcode()) {
   default:
     // No idea what to do.
     return false;
+  case TargetOpcode::G_SHUFFLE_VECTOR:
+    return legalizeShuffleVector(MI, Helper);
   case TargetOpcode::G_VAARG:
     return legalizeVaArg(MI, MRI, MIRBuilder);
   case TargetOpcode::G_LOAD:
   case TargetOpcode::G_STORE:
     return legalizeLoadStore(MI, MRI, MIRBuilder, Observer);
   case TargetOpcode::G_SHL:
   case TargetOpcode::G_ASHR:
   case TargetOpcode::G_LSHR:
auto Concat = MIRBuilder.buildConcatVectors(
       DstTy.changeElementSize(DstTy.getScalarSizeInBits() * 2), SplitSrcs);
 
   Helper.Observer.changingInstr(MI);
   MI.getOperand(1).setReg(Concat.getReg(0));
   Helper.Observer.changedInstr(MI);
   return true;
 }
 
+bool AArch64LegalizerInfo::legalizeShuffleVector(
+    MachineInstr &MI, LegalizerHelper &Helper) const {
+  assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
+
+  MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
+  MachineRegisterInfo &MRI = *MIRBuilder.getMRI();
+
+  ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();
+  unsigned MaskNumElts = Mask.size();
+  LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
+  LLT SrcTy = MRI.getType(MI.getOperand(1).getReg());
+  LLT DestEltTy = DstTy.getElementType();
+  Register DstReg = MI.getOperand(0).getReg();
+  unsigned SrcNumElts = SrcTy.getNumElements();
+
+  if (Mask.size() <= SrcNumElts)
+    return false;
+
+  unsigned PaddedMaskNumElts = alignTo(MaskNumElts, SrcNumElts);
+  unsigned NumConcat = PaddedMaskNumElts / SrcNumElts;
+  LLT PaddedTy = LLT::fixed_vector(PaddedMaskNumElts, DestEltTy);
+
+  auto Undef = MIRBuilder.buildUndef(SrcTy);
+  SmallVector<Register, 8> MOps1(NumConcat, Undef.getReg(0));
+  SmallVector<Register, 8> MOps2(NumConcat, Undef.getReg(0));
+  MOps1[0] = MI.getOperand(1).getReg();
+  MOps2[0] = MI.getOperand(2).getReg();
+
+  auto Src1 = MIRBuilder.buildConcatVectors(PaddedTy, MOps1);
+  auto Src2 = MIRBuilder.buildConcatVectors(PaddedTy, MOps2);
+
+  // Readjust mask for new input vector length.
+  SmallVector<int, 8> MappedOps(PaddedMaskNumElts, -1);
+  for (unsigned I = 0; I != MaskNumElts; ++I) {
+    int Idx = Mask[I];
+    if (Idx >= (int)SrcNumElts)
+      Idx -= SrcNumElts - PaddedMaskNumElts;
+    MappedOps[I] = Idx;
+  }
+
+  if (MaskNumElts != PaddedMaskNumElts) {
+    auto Shuffle =
+        MIRBuilder.buildShuffleVector(PaddedTy, Src1, Src2, MappedOps);
+
+    SmallVector<Register, 8> Elts(MaskNumElts);
+    for (unsigned I = 0; I < MaskNumElts; ++I) {
+      Elts[I] = MIRBuilder
+                    .buildExtractVectorElement(
+                        DestEltTy, Shuffle,
+                        MIRBuilder.buildConstant(LLT::scalar(64), I))
+                    .getReg(0);
+    }
+    MIRBuilder.buildBuildVector(DstReg, Elts);
+  } else {
+    MIRBuilder.buildShuffleVector(DstReg, Src1, Src2, MappedOps);
+  }
+
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AArch64LegalizerInfo::legalizeSmallCMGlobalValue(
     MachineInstr &MI, MachineRegisterInfo &MRI, MachineIRBuilder &MIRBuilder,
     GISelChangeObserver &Observer) const {
   assert(MI.getOpcode() == TargetOpcode::G_GLOBAL_VALUE);
   // We do this custom legalization to convert G_GLOBAL_VALUE into target ADRP +
   // G_ADD_LOW instructions.
   // By splitting this here, we can optimize accesses in the small code model by
   // folding in the G_ADD_LOW into the load/store offset.

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

bb.1:
     %14:_(s64) = G_LOAD %18(p0) :: (invariant load 8 from %fixed-stack.0)
     %1:_(<6 x s64>) = G_BUILD_VECTOR %9(s64), %10(s64), %11(s64), %12(s64), %13(s64), %14(s64)
     %2:_(p0) = COPY $x0
     %19:_(<6 x s64>) = G_SHUFFLE_VECTOR %0(<6 x s64>), %1, shufflemask(3, 4, 7, 0, 1, 11)
     G_STORE %19(<6 x s64>), %2(p0) :: (store (<6 x s64>), align 64)
     RET_ReallyLR
 
 ...
+---
+name:            shuffle_v4i32_v2i32
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $q0, $d1
+
+    ; CHECK-LABEL: name: shuffle_v4i32_v2i32
+    ; CHECK: liveins: $q0, $d1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<2 x s32>) = COPY $d1
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<4 x s32>) = G_CONCAT_VECTORS [[COPY]](<2 x s32>), [[DEF]](<2 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS1:%[0-9]+]]:_(<4 x s32>) = G_CONCAT_VECTORS [[COPY1]](<2 x s32>), [[DEF]](<2 x s32>)
+    ; CHECK-NEXT: [[SHUF:%[0-9]+]]:_(<4 x s32>) = G_SHUFFLE_VECTOR [[CONCAT_VECTORS]](<4 x s32>), [[CONCAT_VECTORS1]], shufflemask(0, 1, 4, 5)
+    ; CHECK-NEXT: $q0 = COPY [[SHUF]](<4 x s32>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %0:_(<2 x s32>) = COPY $d0
+    %1:_(<2 x s32>) = COPY $d1
+    %2:_(<4 x s32>) = G_SHUFFLE_VECTOR %0(<2 x s32>), %1, shufflemask(0, 1, 2, 3)
+    $q0 = COPY %2(<4 x s32>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            shuffle_v8i16_v4i16
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $d0, $d1
+
+    ; CHECK-LABEL: name: shuffle_v8i16_v4i16
+    ; CHECK: liveins: $d0, $d1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<4 x s16>) = COPY $d0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<4 x s16>) = COPY $d1
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(<4 x s16>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<8 x s16>) = G_CONCAT_VECTORS [[COPY]](<4 x s16>), [[DEF]](<4 x s16>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS1:%[0-9]+]]:_(<8 x s16>) = G_CONCAT_VECTORS [[COPY1]](<4 x s16>), [[DEF]](<4 x s16>)
+    ; CHECK-NEXT: [[SHUF:%[0-9]+]]:_(<8 x s16>) = G_SHUFFLE_VECTOR [[CONCAT_VECTORS]](<8 x s16>), [[CONCAT_VECTORS1]], shufflemask(11, 10, 9, 8, 3, 2, 1, 0)
+    ; CHECK-NEXT: $q0 = COPY [[SHUF]](<8 x s16>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %0:_(<4 x s16>) = COPY $d0
+    %1:_(<4 x s16>) = COPY $d1
+    %2:_(<8 x s16>) = G_SHUFFLE_VECTOR %0(<4 x s16>), %1, shufflemask(7, 6, 5, 4, 3, 2, 1, 0)
+    $q0 = COPY %2(<8 x s16>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            shuffle_v16i8_v8i8
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $d0, $d1
+
+    ; CHECK-LABEL: name: shuffle_v16i8_v8i8
+    ; CHECK: liveins: $d0, $d1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<8 x s8>) = COPY $d0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<8 x s8>) = COPY $d1
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(<8 x s8>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s8>) = G_CONCAT_VECTORS [[COPY]](<8 x s8>), [[DEF]](<8 x s8>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS1:%[0-9]+]]:_(<16 x s8>) = G_CONCAT_VECTORS [[COPY1]](<8 x s8>), [[DEF]](<8 x s8>)
+    ; CHECK-NEXT: [[SHUF:%[0-9]+]]:_(<16 x s8>) = G_SHUFFLE_VECTOR [[CONCAT_VECTORS]](<16 x s8>), [[CONCAT_VECTORS1]], shufflemask(7, 21, 6, 4, 5, 3, 0, 0, 0, 0, 0, 0, 0, 0, 23, 0)
+    ; CHECK-NEXT: $q0 = COPY [[SHUF]](<16 x s8>)
+    ; CHECK-NEXT: RET_ReallyLR implicit $q0
+    %0:_(<8 x s8>) = COPY $d0
+    %1:_(<8 x s8>) = COPY $d1
+    %2:_(<16 x s8>) = G_SHUFFLE_VECTOR %0(<8 x s8>), %1, shufflemask(7, 13, 6, 4, 5, 3, 0, 0, 0, 0, 0, 0, 0, 0, 15, 0)
+    $q0 = COPY %2(<16 x s8>)
+    RET_ReallyLR implicit $q0
+
+...
+---
+name:            size_shuffle_v6i32_v4i32
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $s0, $s1, $s2, $s3, $s4, $s5, $s6, $s7, $x0
+
+    ; CHECK-LABEL: name: size_shuffle_v6i32_v4i32
+    ; CHECK: liveins: $s0, $s1, $s2, $s3, $s4, $s5, $s6, $s7, $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $s0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(s32) = COPY $s1
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:_(s32) = COPY $s2
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:_(s32) = COPY $s3
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s32>) = G_BUILD_VECTOR [[COPY]](s32), [[COPY1]](s32), [[COPY2]](s32), [[COPY3]](s32)
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:_(s32) = COPY $s4
+    ; CHECK-NEXT: [[COPY5:%[0-9]+]]:_(s32) = COPY $s5
+    ; CHECK-NEXT: [[COPY6:%[0-9]+]]:_(s32) = COPY $s6
+    ; CHECK-NEXT: [[COPY7:%[0-9]+]]:_(s32) = COPY $s7
+    ; CHECK-NEXT: [[BUILD_VECTOR1:%[0-9]+]]:_(<4 x s32>) = G_BUILD_VECTOR [[COPY4]](s32), [[COPY5]](s32), [[COPY6]](s32), [[COPY7]](s32)
+    ; CHECK-NEXT: [[COPY8:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK-NEXT: [[SHUF:%[0-9]+]]:_(<4 x s32>) = G_SHUFFLE_VECTOR [[BUILD_VECTOR]](<4 x s32>), [[BUILD_VECTOR1]], shufflemask(3, 4, 7, 0)
+    ; CHECK-NEXT: [[SHUF1:%[0-9]+]]:_(<4 x s32>) = G_SHUFFLE_VECTOR [[BUILD_VECTOR]](<4 x s32>), [[BUILD_VECTOR1]], shufflemask(1, 5, undef, undef)
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+    ; CHECK-NEXT: [[EVEC:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[SHUF]](<4 x s32>), [[C]](s64)
+    ; CHECK-NEXT: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+    ; CHECK-NEXT: [[EVEC1:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[SHUF]](<4 x s32>), [[C1]](s64)
+    ; CHECK-NEXT: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 2
+    ; CHECK-NEXT: [[EVEC2:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[SHUF]](<4 x s32>), [[C2]](s64)
+    ; CHECK-NEXT: [[C3:%[0-9]+]]:_(s64) = G_CONSTANT i64 3
+    ; CHECK-NEXT: [[EVEC3:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[SHUF]](<4 x s32>), [[C3]](s64)
+    ; CHECK-NEXT: [[EVEC4:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[SHUF1]](<4 x s32>), [[C]](s64)
+    ; CHECK-NEXT: [[EVEC5:%[0-9]+]]:_(s32) = G_EXTRACT_VECTOR_ELT [[SHUF1]](<4 x s32>), [[C1]](s64)
+    ; CHECK-NEXT: [[BUILD_VECTOR2:%[0-9]+]]:_(<4 x s32>) = G_BUILD_VECTOR [[EVEC]](s32), [[EVEC1]](s32), [[EVEC2]](s32), [[EVEC3]](s32)
+    ; CHECK-NEXT: [[BUILD_VECTOR3:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[EVEC4]](s32), [[EVEC5]](s32)
+    ; CHECK-NEXT: G_STORE [[BUILD_VECTOR2]](<4 x s32>), [[COPY8]](p0) :: (store (<4 x s32>), align 32)
+    ; CHECK-NEXT: [[C4:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+    ; CHECK-NEXT: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY8]], [[C4]](s64)
+    ; CHECK-NEXT: G_STORE [[BUILD_VECTOR3]](<2 x s32>), [[PTR_ADD]](p0) :: (store (<2 x s32>) into unknown-address + 16, align 16)
+    ; CHECK-NEXT: RET_ReallyLR
+    %3:_(s32) = COPY $s0
+    %4:_(s32) = COPY $s1
+    %5:_(s32) = COPY $s2
+    %6:_(s32) = COPY $s3
+    %0:_(<4 x s32>) = G_BUILD_VECTOR %3(s32), %4(s32), %5(s32), %6(s32)
+    %7:_(s32) = COPY $s4
+    %8:_(s32) = COPY $s5
+    %9:_(s32) = COPY $s6
+    %10:_(s32) = COPY $s7
+    %1:_(<4 x s32>) = G_BUILD_VECTOR %7(s32), %8(s32), %9(s32), %10(s32)
+    %2:_(p0) = COPY $x0
+    %19:_(<6 x s32>) = G_SHUFFLE_VECTOR %0(<4 x s32>), %1, shufflemask(3, 4, 7, 0, 1, 5)
+    G_STORE %19(<6 x s32>), %2(p0) :: (store (<6 x s32>), align 32)
+    RET_ReallyLR
+
+...