[GlobalISel] LegalizerHelper: Implement fewerElementsVector for G_LOAD/G_STORE

[volkan]

Rebased.

[dsanders]

LGTM. The only comment I have is an observation that odd-length vectors generally don't fit well into fewerElements at the moment but that's not for this patch

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

1115

We don't need to deal with it for this patch but: It occurs to me that extractParts can never change an odd-length vector into an even-length one. The effect of this is that most odd-length vectors can only unmerge into scalars (there's a few exceptions like 9 into 3 but only a few within the normal range of vector lengths) yet it would be more efficient to fewerElements a 5-element vector into a 4-element and a 1-element vector so that we can still use vectors for the majority of it.

[volkan]

• Replaced buildMerge with buildBuildVector/buildConcatVectors.
• Updated the test to avoid generating G_CONCAT_VECTORS as it's not legal for v4s32/v2s64.

[dsanders]

LGTM

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

 }
 
 LegalizerHelper::LegalizeResult
 LegalizerHelper::fewerElementsVector(MachineInstr &MI, unsigned TypeIdx,
                                      LLT NarrowTy) {
   // FIXME: Don't know how to handle secondary types yet.
   if (TypeIdx != 0)
     return UnableToLegalize;
+
+  MIRBuilder.setInstr(MI);
   switch (MI.getOpcode()) {
   default:
     return UnableToLegalize;
   case TargetOpcode::G_ADD: {
     unsigned NarrowSize = NarrowTy.getSizeInBits();
     unsigned DstReg = MI.getOperand(0).getReg();
     unsigned Size = MRI.getType(DstReg).getSizeInBits();
     int NumParts = Size / NarrowSize;
     // FIXME: Don't know how to handle the situation where the small vectors
     // aren't all the same size yet.
     if (Size % NarrowSize != 0)
       return UnableToLegalize;
 
-    MIRBuilder.setInstr(MI);
-
     SmallVector<unsigned, 2> Src1Regs, Src2Regs, DstRegs;
     extractParts(MI.getOperand(1).getReg(), NarrowTy, NumParts, Src1Regs);
     extractParts(MI.getOperand(2).getReg(), NarrowTy, NumParts, Src2Regs);
 
     for (int i = 0; i < NumParts; ++i) {
       unsigned DstReg = MRI.createGenericVirtualRegister(NarrowTy);
       MIRBuilder.buildAdd(DstReg, Src1Regs[i], Src2Regs[i]);
       DstRegs.push_back(DstReg);
     }
 
     MIRBuilder.buildMerge(DstReg, DstRegs);
     MI.eraseFromParent();
     return Legalized;
   }
+  case TargetOpcode::G_LOAD:
+  case TargetOpcode::G_STORE: {
+    bool IsLoad = MI.getOpcode() == TargetOpcode::G_LOAD;
+    unsigned ValReg = MI.getOperand(0).getReg();
+    unsigned AddrReg = MI.getOperand(1).getReg();
+    unsigned NarrowSize = NarrowTy.getSizeInBits();
+    unsigned Size = MRI.getType(ValReg).getSizeInBits();
+    unsigned NumParts = Size / NarrowSize;
+
+    SmallVector<unsigned, 8> NarrowRegs;
+    if (!IsLoad)
+      extractParts(ValReg, NarrowTy, NumParts, NarrowRegs);
+
+    const LLT OffsetTy =
+        LLT::scalar(MRI.getType(AddrReg).getScalarSizeInBits());
+    MachineFunction &MF = *MI.getMF();
+    MachineMemOperand *MMO = *MI.memoperands_begin();
+    for (unsigned Idx = 0; Idx < NumParts; ++Idx) {
+      unsigned Adjustment = Idx * NarrowTy.getSizeInBits() / 8;
+      unsigned Alignment = MinAlign(MMO->getAlignment(), Adjustment);
+      unsigned NewAddrReg = 0;
+      MIRBuilder.materializeGEP(NewAddrReg, AddrReg, OffsetTy, Adjustment);
+      MachineMemOperand &NewMMO = *MF.getMachineMemOperand(
+          MMO->getPointerInfo().getWithOffset(Adjustment), MMO->getFlags(),
+          NarrowTy.getSizeInBits() / 8, Alignment);
+      if (IsLoad) {
+        unsigned Dst = MRI.createGenericVirtualRegister(NarrowTy);
+        NarrowRegs.push_back(Dst);
+        MIRBuilder.buildLoad(Dst, NewAddrReg, NewMMO);
+      } else {
+        MIRBuilder.buildStore(NarrowRegs[Idx], NewAddrReg, NewMMO);
+      }
+    }
+    if (IsLoad)
+      MIRBuilder.buildMerge(ValReg, NarrowRegs);
+    MI.eraseFromParent();
+    return Legalized;
+  }
   }
 }
 
 LegalizerHelper::LegalizeResult
 LegalizerHelper::lowerBitCount(MachineInstr &MI, unsigned TypeIdx, LLT Ty) {
   unsigned Opc = MI.getOpcode();
   auto &TII = *MI.getMF()->getSubtarget().getInstrInfo();
   auto isLegalOrCustom = [this](const LegalityQuery &Q) {

test/CodeGen/AArch64/GlobalISel/legalize-load-fewerElts.mir

-# RUN: llc -march=aarch64 -o - -run-pass=legalizer -global-isel-abort=0 -debug-only=legalizer 2>&1 %s | FileCheck %s
-# REQUIRES: asserts
-
-# CHECK: Legalize Machine IR for: load_v4s32
-# CHECK-NEXT: %{{[0-9]+}}:_(<4 x s32>) = G_LOAD %{{[0-9]+}}:_(p0)
-# CHECK-NEXT: Reduce number of elements
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=aarch64 -o - -run-pass=legalizer %s | FileCheck %s
 ---
 name:            load_v4s32
 legalized:       false
 tracksRegLiveness: true
 body:             |
   bb.1:
     liveins: $x0
 
+    ; CHECK-LABEL: name: load_v4s32
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[LOAD:%[0-9]+]]:_(<2 x s32>) = G_LOAD [[COPY]](p0) :: (load 8, align 16)
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
+    ; CHECK: [[GEP:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[C]](s64)
+    ; CHECK: [[LOAD1:%[0-9]+]]:_(<2 x s32>) = G_LOAD [[GEP]](p0) :: (load 8)
+    ; CHECK: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[LOAD1]](<2 x s32>)
+    ; CHECK: $w0 = COPY [[UV1]](s32)
     %0:_(p0) = COPY $x0
-    %1:_(<4 x s32>) = G_LOAD %0(p0) :: (load 16, align 4)
+    %1:_(<4 x s32>) = G_LOAD %0(p0) :: (load 16)
     %2:_(s32), %3:_(s32), %4:_(s32), %5:_(s32) = G_UNMERGE_VALUES %1(<4 x s32>)
     $w0 = COPY %5(s32)
 
 ...
-
-# Make sure we are able to scalarize v2s64.
-# CHECK: Legalize Machine IR for: load_v2s64
-# CHECK-NEXT: %{{[0-9]+}}:_(<2 x s64>) = G_LOAD %{{[0-9]+}}:_(p0)
-# CHECK-NEXT: Reduce number of elements
 ---
 name:            load_v2s64
 legalized:       false
 tracksRegLiveness: true
 body:             |
   bb.1:
     liveins: $x0
 
+    ; CHECK-LABEL: name: load_v2s64
+    ; CHECK: liveins: $x0
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[LOAD:%[0-9]+]]:_(s64) = G_LOAD [[COPY]](p0) :: (load 8, align 16)
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
+    ; CHECK: [[GEP:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[C]](s64)
+    ; CHECK: [[LOAD1:%[0-9]+]]:_(s64) = G_LOAD [[GEP]](p0) :: (load 8)
+    ; CHECK: $x0 = COPY [[LOAD1]](s64)
     %0:_(p0) = COPY $x0
     %1:_(<2 x s64>) = G_LOAD %0(p0) :: (load 16)
     %2:_(s64), %3:_(s64) = G_UNMERGE_VALUES %1(<2 x s64>)
     $x0 = COPY %3(s64)
 
 ...

test/CodeGen/AArch64/GlobalISel/legalize-load-store.mir

bb.0.entry:
     ; CHECK: $x0 = COPY [[BITCAST]](s64)
     ; CHECK: [[LOAD6:%[0-9]+]]:_(s64) = G_LOAD [[COPY]](p0) :: (load 8, align 16)
     ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
     ; CHECK: [[GEP:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[C]](s64)
     ; CHECK: [[LOAD7:%[0-9]+]]:_(s64) = G_LOAD [[GEP]](p0) :: (load 8)
     ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[LOAD6]](s64), [[LOAD7]](s64)
     ; CHECK: [[TRUNC:%[0-9]+]]:_(s64) = G_TRUNC [[MV]](s128)
     ; CHECK: $x0 = COPY [[TRUNC]](s64)
+    ; CHECK: [[LOAD8:%[0-9]+]]:_(<2 x s32>) = G_LOAD [[COPY]](p0) :: (load 8, align 16)
+    ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
+    ; CHECK: [[GEP1:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[C1]](s64)
+    ; CHECK: [[LOAD9:%[0-9]+]]:_(<2 x s32>) = G_LOAD [[GEP1]](p0) :: (load 8)
+    ; CHECK: [[MV1:%[0-9]+]]:_(<4 x s32>) = G_MERGE_VALUES [[LOAD8]](<2 x s32>), [[LOAD9]](<2 x s32>)
+    ; CHECK: [[BITCAST1:%[0-9]+]]:_(s128) = G_BITCAST [[MV1]](<4 x s32>)
+    ; CHECK: [[TRUNC1:%[0-9]+]]:_(s64) = G_TRUNC [[BITCAST1]](s128)
+    ; CHECK: $x0 = COPY [[TRUNC1]](s64)
     %0:_(p0) = COPY $x0
     %1:_(s1) = G_LOAD %0(p0) :: (load 1)
     %2:_(s32) = G_ANYEXT %1(s1)
     $w0 = COPY %2(s32)
     %3:_(s8) = G_LOAD %0(p0) :: (load 1)
     %4:_(s32) = G_ANYEXT %3(s8)
     $w0 = COPY %4(s32)
     %5:_(s16) = G_LOAD %0(p0) :: (load 2)
     %6:_(s32) = G_ANYEXT %5(s16)
     $w0 = COPY %6(s32)
     %7:_(s32) = G_LOAD %0(p0) :: (load 4)
     $w0 = COPY %4(s32)
     %8:_(s64) = G_LOAD %0(p0) :: (load 8)
     $x0 = COPY %8(s64)
     %9:_(p0) = G_LOAD %0(p0) :: (load 8)
     %10:_(s64) = G_PTRTOINT %9(p0)
     $x0 = COPY %10(s64)
     %11:_(<2 x s32>) = G_LOAD %0(p0) :: (load 8)
     %12:_(s64) = G_BITCAST %11(<2 x s32>)
     $x0 = COPY %12(s64)
     %13:_(s128) = G_LOAD %0(p0) :: (load 16)
     %14:_(s64) = G_TRUNC %13(s128)
     $x0 = COPY %14(s64)
+    %15:_(<4 x s32>) = G_LOAD %0(p0) :: (load 16)
+    %16:_(s128) = G_BITCAST %15(<4 x s32>)
+    %17:_(s64) = G_TRUNC %16(s128)
+    $x0 = COPY %17(s64)
 ...
 
 ---
 name:            test_store
 body: |
   bb.0.entry:
     liveins: $x0, $w1
 
bb.0.entry:
     ; CHECK: G_STORE [[TRUNC1]](s8), [[COPY]](p0) :: (store 1)
     ; CHECK: [[TRUNC2:%[0-9]+]]:_(s16) = G_TRUNC [[COPY1]](s32)
     ; CHECK: G_STORE [[TRUNC2]](s16), [[COPY]](p0) :: (store 2)
     ; CHECK: G_STORE [[COPY1]](s32), [[COPY]](p0) :: (store 4)
     ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[COPY]](p0)
     ; CHECK: G_STORE [[PTRTOINT]](s64), [[COPY]](p0) :: (store 8)
     ; CHECK: G_STORE [[COPY]](p0), [[COPY]](p0) :: (store 8)
     ; CHECK: [[PTRTOINT1:%[0-9]+]]:_(s64) = G_PTRTOINT [[COPY]](p0)
+    ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[PTRTOINT1]](s64), [[PTRTOINT1]](s64)
     ; CHECK: G_STORE [[PTRTOINT1]](s64), [[COPY]](p0) :: (store 8, align 16)
     ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
     ; CHECK: [[GEP:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[C1]](s64)
     ; CHECK: G_STORE [[PTRTOINT1]](s64), [[GEP]](p0) :: (store 8)
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<4 x s32>) = G_BITCAST [[MV]](s128)
+    ; CHECK: [[UV:%[0-9]+]]:_(<2 x s32>), [[UV1:%[0-9]+]]:_(<2 x s32>) = G_UNMERGE_VALUES [[BITCAST]](<4 x s32>)
+    ; CHECK: G_STORE [[UV]](<2 x s32>), [[COPY]](p0) :: (store 8, align 16)
+    ; CHECK: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
+    ; CHECK: [[GEP1:%[0-9]+]]:_(p0) = G_GEP [[COPY]], [[C2]](s64)
+    ; CHECK: G_STORE [[UV1]](<2 x s32>), [[GEP1]](p0) :: (store 8)
     %0:_(p0) = COPY $x0
     %1:_(s32) = COPY $w1
     %2:_(s1) = G_TRUNC %1(s32)
     G_STORE %2(s1), %0(p0) :: (store 1)
     %3:_(s8) = G_TRUNC %1(s32)
     G_STORE %3(s8), %0(p0) :: (store 1)
     %4:_(s16) = G_TRUNC %1(s32)
     G_STORE %4(s16), %0(p0) :: (store 2)
     G_STORE %1(s32), %0(p0) :: (store 4)
     %5:_(s64) = G_PTRTOINT %0(p0)
     G_STORE %5(s64), %0(p0) :: (store 8)
     G_STORE %0(p0), %0(p0) :: (store 8)
     %6:_(s64) = G_PTRTOINT %0(p0)
     %7:_(s128) = G_MERGE_VALUES %6(s64), %6
     G_STORE %7(s128), %0(p0) :: (store 16)
+    %8:_(<4 x s32>) = G_BITCAST %7(s128)
+    G_STORE %8(<4 x s32>), %0(p0) :: (store 16)
 ...