[AArch64][GlobalISel] Clamp vectors of p0 when legalizing G_LOAD/G_STORE

[paquette]

Summary

We had a rule for <n x s64> but not one for <n x p0>. As a result, we'd fall back on like <5 x p0> or whatever.

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

getActionDefinitionsBuilder(G_LOAD)
       // Lower any any-extending loads left into G_ANYEXT and G_LOAD
       .lowerIf([=](const LegalityQuery &Query) {
         return Query.Types[0] != Query.MMODescrs[0].MemoryTy;
       })
       .clampMaxNumElements(0, s8, 16)
       .clampMaxNumElements(0, s16, 8)
       .clampMaxNumElements(0, s32, 4)
       .clampMaxNumElements(0, s64, 2)
+      .clampMaxNumElements(0, p0, 2)
       .customIf(IsPtrVecPred)
       .scalarizeIf(typeIs(0, v2s16), 0);
 
   getActionDefinitionsBuilder(G_STORE)
       .legalForTypesWithMemDesc({{s8, p0, s8, 8},
                                  {s16, p0, s8, 8}, // truncstorei8 from s16
                                  {s32, p0, s8, 8}, // truncstorei8 from s32
                                  {s64, p0, s8, 8}, // truncstorei8 from s64
getActionDefinitionsBuilder(G_STORE)
         return Query.Types[0].isScalar() &&
                Query.Types[0] != Query.MMODescrs[0].MemoryTy;
       })
       // Maximum: sN * k = 128
       .clampMaxNumElements(0, s8, 16)
       .clampMaxNumElements(0, s16, 8)
       .clampMaxNumElements(0, s32, 4)
       .clampMaxNumElements(0, s64, 2)
+      .clampMaxNumElements(0, p0, 2)
       .lowerIfMemSizeNotPow2()
       .customIf(IsPtrVecPred)
       .scalarizeIf(typeIs(0, v2s16), 0);
 
   // Constants
   getActionDefinitionsBuilder(G_CONSTANT)
       .legalFor({p0, s8, s16, s32, s64})
       .widenScalarToNextPow2(0)

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

bb.1:
     ; CHECK: %ext:_(s64) = G_AND [[ANYEXT]], [[C]]
     ; CHECK: $x0 = COPY %ext(s64)
     ; CHECK: RET_ReallyLR implicit $x0
     %ptr:_(p0) = COPY $x0
     %load:_(s1) = G_LOAD %ptr(p0) :: (load (s1))
     %ext:_(s64) = G_ZEXT %load
     $x0 = COPY %ext
     RET_ReallyLR implicit $x0
+...
+---
+name:            load_store_6xp0
+alignment:       4
+tracksRegLiveness: true
+machineFunctionInfo: {}
+body:             |
+  bb.1:
+    liveins: $x0
+    ; CHECK-LABEL: name: load_store_6xp0
+    ; CHECK: liveins: $x0
+    ; CHECK: %ptr:_(p0) = COPY $x0
+    ; CHECK: [[LOAD:%[0-9]+]]:_(p0) = G_LOAD %ptr(p0) :: (load (p0), align 64)
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 8
+    ; CHECK: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C]](s64)
+    ; CHECK: [[LOAD1:%[0-9]+]]:_(p0) = G_LOAD [[PTR_ADD]](p0) :: (load (p0) from unknown-address + 8)
+    ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+    ; CHECK: [[PTR_ADD1:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C1]](s64)
+    ; CHECK: [[LOAD2:%[0-9]+]]:_(p0) = G_LOAD [[PTR_ADD1]](p0) :: (load (p0) from unknown-address + 16, align 16)
+    ; CHECK: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 24
+    ; CHECK: [[PTR_ADD2:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C2]](s64)
+    ; CHECK: [[LOAD3:%[0-9]+]]:_(p0) = G_LOAD [[PTR_ADD2]](p0) :: (load (p0) from unknown-address + 24)
+    ; CHECK: [[C3:%[0-9]+]]:_(s64) = G_CONSTANT i64 32
+    ; CHECK: [[PTR_ADD3:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C3]](s64)
+    ; CHECK: [[LOAD4:%[0-9]+]]:_(p0) = G_LOAD [[PTR_ADD3]](p0) :: (load (p0) from unknown-address + 32, align 32)
+    ; CHECK: [[C4:%[0-9]+]]:_(s64) = G_CONSTANT i64 40
+    ; CHECK: [[PTR_ADD4:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C4]](s64)
+    ; CHECK: [[LOAD5:%[0-9]+]]:_(p0) = G_LOAD [[PTR_ADD4]](p0) :: (load (p0) from unknown-address + 40)
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<2 x p0>) = G_BUILD_VECTOR [[LOAD]](p0), [[LOAD1]](p0)
+    ; CHECK: [[BUILD_VECTOR1:%[0-9]+]]:_(<2 x p0>) = G_BUILD_VECTOR [[LOAD2]](p0), [[LOAD3]](p0)
+    ; CHECK: [[BUILD_VECTOR2:%[0-9]+]]:_(<2 x p0>) = G_BUILD_VECTOR [[LOAD4]](p0), [[LOAD5]](p0)
+    ; CHECK: [[BITCAST:%[0-9]+]]:_(<2 x s64>) = G_BITCAST [[BUILD_VECTOR]](<2 x p0>)
+    ; CHECK: G_STORE [[BITCAST]](<2 x s64>), %ptr(p0) :: (store (<2 x s64>), align 64)
+    ; CHECK: [[PTR_ADD5:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C1]](s64)
+    ; CHECK: [[BITCAST1:%[0-9]+]]:_(<2 x s64>) = G_BITCAST [[BUILD_VECTOR1]](<2 x p0>)
+    ; CHECK: G_STORE [[BITCAST1]](<2 x s64>), [[PTR_ADD5]](p0) :: (store (<2 x s64>) into unknown-address + 16)
+    ; CHECK: [[PTR_ADD6:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr, [[C3]](s64)
+    ; CHECK: [[BITCAST2:%[0-9]+]]:_(<2 x s64>) = G_BITCAST [[BUILD_VECTOR2]](<2 x p0>)
+    ; CHECK: G_STORE [[BITCAST2]](<2 x s64>), [[PTR_ADD6]](p0) :: (store (<2 x s64>) into unknown-address + 32, align 32)
+    ; CHECK: RET_ReallyLR
+    %ptr:_(p0) = COPY $x0
+    %val:_(<6 x p0>) = G_LOAD %ptr(p0) :: (load (<6 x p0>))
+    G_STORE %val(<6 x p0>), %ptr(p0) :: (store (<6 x p0>))
+    RET_ReallyLR

llvm/test/CodeGen/AArch64/GlobalISel/legalize-phi.mir

 body:             |
   ; CHECK-LABEL: name: v4p0
   ; CHECK: bb.0:
   ; CHECK:   successors: %bb.1(0x50000000), %bb.2(0x30000000)
   ; CHECK:   liveins: $x0, $x1
   ; CHECK:   %ptr1:_(p0) = COPY $x1
   ; CHECK:   %ptr2:_(p0) = COPY $x0
   ; CHECK:   %cond:_(s1) = G_IMPLICIT_DEF
-  ; CHECK:   %val_1:_(<4 x p0>) = G_LOAD %ptr1(p0) :: (load (<4 x p0>))
-  ; CHECK:   [[UV:%[0-9]+]]:_(<2 x p0>), [[UV1:%[0-9]+]]:_(<2 x p0>) = G_UNMERGE_VALUES %val_1(<4 x p0>)
+  ; CHECK:   [[LOAD:%[0-9]+]]:_(<2 x s64>) = G_LOAD %ptr1(p0) :: (load (<2 x s64>), align 32)
+  ; CHECK:   [[BITCAST:%[0-9]+]]:_(<2 x p0>) = G_BITCAST [[LOAD]](<2 x s64>)
+  ; CHECK:   [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+  ; CHECK:   [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr1, [[C]](s64)
+  ; CHECK:   [[LOAD1:%[0-9]+]]:_(<2 x s64>) = G_LOAD [[PTR_ADD]](p0) :: (load (<2 x s64>) from unknown-address + 16)
+  ; CHECK:   [[BITCAST1:%[0-9]+]]:_(<2 x p0>) = G_BITCAST [[LOAD1]](<2 x s64>)
   ; CHECK:   G_BRCOND %cond(s1), %bb.2
   ; CHECK:   G_BR %bb.1
   ; CHECK: bb.1:
   ; CHECK:   successors: %bb.2(0x80000000)
-  ; CHECK:   [[LOAD:%[0-9]+]]:_(<4 x s64>) = G_LOAD %ptr2(p0) :: (load (<4 x s64>))
-  ; CHECK:   %val_2:_(<4 x p0>) = G_BITCAST [[LOAD]](<4 x s64>)
-  ; CHECK:   [[UV2:%[0-9]+]]:_(<2 x p0>), [[UV3:%[0-9]+]]:_(<2 x p0>) = G_UNMERGE_VALUES %val_2(<4 x p0>)
-  ; CHECK: bb.2:
-  ; CHECK:   [[PHI:%[0-9]+]]:_(<2 x p0>) = G_PHI [[UV2]](<2 x p0>), %bb.1, [[UV]](<2 x p0>), %bb.0
-  ; CHECK:   [[PHI1:%[0-9]+]]:_(<2 x p0>) = G_PHI [[UV3]](<2 x p0>), %bb.1, [[UV1]](<2 x p0>), %bb.0
-  ; CHECK:   %phi:_(<4 x p0>) = G_CONCAT_VECTORS [[PHI]](<2 x p0>), [[PHI1]](<2 x p0>)
-  ; CHECK:   %unmerge_1:_(<2 x p0>), %unmerge_2:_(<2 x p0>) = G_UNMERGE_VALUES %phi(<4 x p0>)
-  ; CHECK:   $q0 = COPY %unmerge_1(<2 x p0>)
-  ; CHECK:   $q1 = COPY %unmerge_2(<2 x p0>)
+  ; CHECK:   [[LOAD2:%[0-9]+]]:_(<2 x s64>) = G_LOAD %ptr2(p0) :: (load (<2 x s64>), align 32)
+  ; CHECK:   [[BITCAST2:%[0-9]+]]:_(<2 x p0>) = G_BITCAST [[LOAD2]](<2 x s64>)
+  ; CHECK:   [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+  ; CHECK:   [[PTR_ADD1:%[0-9]+]]:_(p0) = G_PTR_ADD %ptr2, [[C1]](s64)
+  ; CHECK:   [[LOAD3:%[0-9]+]]:_(<2 x s64>) = G_LOAD [[PTR_ADD1]](p0) :: (load (<2 x s64>) from unknown-address + 16)
+  ; CHECK:   [[BITCAST3:%[0-9]+]]:_(<2 x p0>) = G_BITCAST [[LOAD3]](<2 x s64>)
+  ; CHECK: bb.2:
+  ; CHECK:   [[PHI:%[0-9]+]]:_(<2 x p0>) = G_PHI [[BITCAST2]](<2 x p0>), %bb.1, [[BITCAST]](<2 x p0>), %bb.0
+  ; CHECK:   [[PHI1:%[0-9]+]]:_(<2 x p0>) = G_PHI [[BITCAST3]](<2 x p0>), %bb.1, [[BITCAST1]](<2 x p0>), %bb.0
+  ; CHECK:   $q0 = COPY [[PHI]](<2 x p0>)
+  ; CHECK:   $q1 = COPY [[PHI1]](<2 x p0>)
   ; CHECK:   RET_ReallyLR implicit $q0, implicit $q1
   bb.0:
     successors: %bb.1(0x50000000), %bb.2(0x30000000)
     liveins: $x0, $x1
 
     %ptr1:_(p0) = COPY $x1
     %ptr2:_(p0) = COPY $x0
     %cond:_(s1) = G_IMPLICIT_DEF