[ARM] Extend narrow values to allow using truncating scatters

[dmgreen]

Summary

As a minor adjustment to the existing lowering of offset scatters, this extends any smaller-than-legal vectors into full vectors using a zext, so that the truncating scatters can be used. Due to the way MVE legalizes the vectors this should be cheap in most situations, and will prevent the vector from being scalarized.

[samtebbs]

llvm/lib/Target/ARM/MVEGatherScatterLowering.cpp

643–644

What happens here if InputTy is an integer (as I assume is allowed by isIntOrIntVectorTy())?

[dmgreen]

llvm/lib/Target/ARM/MVEGatherScatterLowering.cpp

643–644

We will already have ruled out scalar integer types, we know InputTy will be some sort of vector. So isIntOrIntVectorTy is really just saying "is it an integer vector type"

[samtebbs]

Looks good to me. It's up to you if you would like to add a comment as mentioned in my reply.

llvm/lib/Target/ARM/MVEGatherScatterLowering.cpp

643–644

Cool. It might be worth adding this info as a brief comment.

llvm/lib/Target/ARM/MVEGatherScatterLowering.cpp

Value *MVEGatherScatterLowering::tryCreateMaskedScatterOffset(
   if (TruncInst *Trunc = dyn_cast<TruncInst>(Input)) {
     Value *PreTrunc = Trunc->getOperand(0);
     Type *PreTruncTy = PreTrunc->getType();
     if (PreTruncTy->getPrimitiveSizeInBits() == 128) {
       Input = PreTrunc;
       InputTy = PreTruncTy;
     }
   }
+  if (InputTy->getPrimitiveSizeInBits() < 128 &&
+      InputTy->isIntOrIntVectorTy()) {
+    // If we can't find a trunc to incorporate into the instruction, create an
+    // implicit one with a zext, so that we can still create a scatter. We know
+    // that the input type is 4x/8x/16x and of type i8/i16/i32, so any type
+    // smaller than 128 bits will divide evenly into a 128bit vector.
+    InputTy = InputTy->getWithNewBitWidth(
+        128 / cast<FixedVectorType>(InputTy)->getNumElements());
+    Input = Builder.CreateZExt(Input, InputTy);
+    LLVM_DEBUG(dbgs() << "masked scatters: Small input type, extending with:\n"
+                      << *Input << "\n");
+  }
   if (InputTy->getPrimitiveSizeInBits() != 128) {
     LLVM_DEBUG(dbgs() << "masked scatters: cannot create scatters for "
                          "non-standard input types. Expanding.\n");
     return nullptr;
   }
 
   Value *Offsets;
   int Scale;

llvm/test/CodeGen/Thumb2/mve-scatter-ind32-unscaled.ll

entry:
   %input.trunc = trunc <4 x i32> %input to <4 x i16>
   call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %input.trunc, <4 x i16*> %ptrs, i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
   ret void
 }
 
 define arm_aapcs_vfpcc void @trunc_signed_unscaled_i16_i8(i8* %base, <4 x i8>* %offptr, <4 x i16> %input) {
 ; CHECK-LABEL: trunc_signed_unscaled_i16_i8:
 ; CHECK:       @ %bb.0: @ %entry
-; CHECK-NEXT:    .save {r4, r5, r7, lr}
-; CHECK-NEXT:    push {r4, r5, r7, lr}
-; CHECK-NEXT:    vldrb.s32 q1, [r1]
-; CHECK-NEXT:    vmov r1, r3, d0
-; CHECK-NEXT:    vmov r4, r5, d1
-; CHECK-NEXT:    vadd.i32 q1, q1, r0
-; CHECK-NEXT:    vmov r0, r12, d2
-; CHECK-NEXT:    vmov r2, lr, d3
-; CHECK-NEXT:    strb r1, [r0]
-; CHECK-NEXT:    strb.w r3, [r12]
-; CHECK-NEXT:    strb r4, [r2]
-; CHECK-NEXT:    strb.w r5, [lr]
-; CHECK-NEXT:    pop {r4, r5, r7, pc}
+; CHECK-NEXT:    vmov.i32 q1, #0xff
+; CHECK-NEXT:    vldrb.s32 q2, [r1]
+; CHECK-NEXT:    vand q0, q0, q1
+; CHECK-NEXT:    vstrb.32 q0, [r0, q2]
+; CHECK-NEXT:    bx lr
 entry:
   %offs = load <4 x i8>, <4 x i8>* %offptr, align 1
   %offs.sext = sext <4 x i8> %offs to <4 x i32>
   %byte_ptrs = getelementptr inbounds i8, i8* %base, <4 x i32> %offs.sext
   %input.trunc = trunc <4 x i16> %input to <4 x i8>
   call void @llvm.masked.scatter.v4i8.v4p0i8(<4 x i8> %input.trunc, <4 x i8*> %byte_ptrs, i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
   ret void
 }
 
 define arm_aapcs_vfpcc void @trunc_unsigned_unscaled_i16_i8(i8* %base, <4 x i8>* %offptr, <4 x i16> %input) {
 ; CHECK-LABEL: trunc_unsigned_unscaled_i16_i8:
 ; CHECK:       @ %bb.0: @ %entry
-; CHECK-NEXT:    .save {r4, r5, r7, lr}
-; CHECK-NEXT:    push {r4, r5, r7, lr}
-; CHECK-NEXT:    vldrb.u32 q1, [r1]
-; CHECK-NEXT:    vmov r1, r3, d0
-; CHECK-NEXT:    vmov r4, r5, d1
-; CHECK-NEXT:    vadd.i32 q1, q1, r0
-; CHECK-NEXT:    vmov r0, r12, d2
-; CHECK-NEXT:    vmov r2, lr, d3
-; CHECK-NEXT:    strb r1, [r0]
-; CHECK-NEXT:    strb.w r3, [r12]
-; CHECK-NEXT:    strb r4, [r2]
-; CHECK-NEXT:    strb.w r5, [lr]
-; CHECK-NEXT:    pop {r4, r5, r7, pc}
+; CHECK-NEXT:    vmov.i32 q1, #0xff
+; CHECK-NEXT:    vldrb.u32 q2, [r1]
+; CHECK-NEXT:    vand q0, q0, q1
+; CHECK-NEXT:    vstrb.32 q0, [r0, q2]
+; CHECK-NEXT:    bx lr
 entry:
   %offs = load <4 x i8>, <4 x i8>* %offptr, align 1
   %offs.zext = zext <4 x i8> %offs to <4 x i32>
   %byte_ptrs = getelementptr inbounds i8, i8* %base, <4 x i32> %offs.zext
   %input.trunc = trunc <4 x i16> %input to <4 x i8>
   call void @llvm.masked.scatter.v4i8.v4p0i8(<4 x i8> %input.trunc, <4 x i8*> %byte_ptrs, i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
   ret void
 }
 
 declare void @llvm.masked.scatter.v4i8.v4p0i8(<4 x i8>, <4 x i8*>, i32, <4 x i1>)
 declare void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16>, <4 x i16*>, i32, <4 x i1>)
 declare void @llvm.masked.scatter.v4f16.v4p0f16(<4 x half>, <4 x half*>, i32, <4 x i1>)
 declare void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32>, <4 x i32*>, i32, <4 x i1>)
 declare void @llvm.masked.scatter.v4f32.v4p0f32(<4 x float>, <4 x float*>, i32, <4 x i1>)

llvm/test/CodeGen/Thumb2/mve-scatter-ind8-unscaled.ll

 entry:
   %offs = load <16 x i8>, <16 x i8>* %offptr, align 1
   %offs.zext = zext <16 x i8> %offs to <16 x i32>
   %ptrs = getelementptr inbounds i8, i8* %base, <16 x i32> %offs.zext
   call void @llvm.masked.scatter.v16i8.v16p0i8(<16 x i8> %input, <16 x i8*> %ptrs, i32 1, <16 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>)
   ret void
 }
 
-; Expanded ?
 define arm_aapcs_vfpcc void @unscaled_v8i8_i8(i8* %base, <8 x i8>* %offptr, <8 x i8> %input) {
 ; CHECK-LABEL: unscaled_v8i8_i8:
 ; CHECK:       @ %bb.0: @ %entry
-; CHECK-NEXT:    .save {r4, r5, r6, lr}
-; CHECK-NEXT:    push {r4, r5, r6, lr}
-; CHECK-NEXT:    vldrb.u32 q1, [r1]
-; CHECK-NEXT:    vmov.u16 r6, q0[0]
-; CHECK-NEXT:    vadd.i32 q1, q1, r0
-; CHECK-NEXT:    vmov r2, r3, d2
-; CHECK-NEXT:    vmov r12, lr, d3
-; CHECK-NEXT:    vldrb.u32 q1, [r1, #4]
-; CHECK-NEXT:    vadd.i32 q1, q1, r0
-; CHECK-NEXT:    vmov r0, r1, d2
-; CHECK-NEXT:    vmov r4, r5, d3
-; CHECK-NEXT:    strb r6, [r2]
-; CHECK-NEXT:    vmov.u16 r2, q0[1]
-; CHECK-NEXT:    strb r2, [r3]
-; CHECK-NEXT:    vmov.u16 r2, q0[2]
-; CHECK-NEXT:    strb.w r2, [r12]
-; CHECK-NEXT:    vmov.u16 r2, q0[3]
-; CHECK-NEXT:    strb.w r2, [lr]
-; CHECK-NEXT:    vmov.u16 r2, q0[4]
-; CHECK-NEXT:    strb r2, [r0]
-; CHECK-NEXT:    vmov.u16 r0, q0[5]
-; CHECK-NEXT:    strb r0, [r1]
-; CHECK-NEXT:    vmov.u16 r0, q0[6]
-; CHECK-NEXT:    strb r0, [r4]
-; CHECK-NEXT:    vmov.u16 r0, q0[7]
-; CHECK-NEXT:    strb r0, [r5]
-; CHECK-NEXT:    pop {r4, r5, r6, pc}
+; CHECK-NEXT:    vldrb.u16 q1, [r1]
+; CHECK-NEXT:    vmovlb.u8 q0, q0
+; CHECK-NEXT:    vstrb.16 q0, [r0, q1]
+; CHECK-NEXT:    bx lr
 entry:
   %offs = load <8 x i8>, <8 x i8>* %offptr, align 1
   %offs.zext = zext <8 x i8> %offs to <8 x i32>
   %ptrs = getelementptr inbounds i8, i8* %base, <8 x i32> %offs.zext
   call void @llvm.masked.scatter.v8i8.v8p0i8(<8 x i8> %input, <8 x i8*> %ptrs, i32 1, <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>)
   ret void
 }
 

llvm/test/CodeGen/Thumb2/mve-scatter-ptrs.ll

entry:
   %ext = trunc <4 x i32> %v to <4 x i16>
   call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %ext, <4 x i16*> %offs, i32 2, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
   ret void
 }
 
 define arm_aapcs_vfpcc void @ptr_v4i16_dup(i32 %v, <4 x i16*> %offs) {
 ; CHECK-LABEL: ptr_v4i16_dup:
 ; CHECK:       @ %bb.0: @ %entry
-; CHECK-NEXT:    vmov r1, r2, d0
-; CHECK-NEXT:    vmov r3, r12, d1
-; CHECK-NEXT:    strh r0, [r1]
-; CHECK-NEXT:    strh r0, [r2]
-; CHECK-NEXT:    strh r0, [r3]
-; CHECK-NEXT:    strh.w r0, [r12]
+; CHECK-NEXT:    vdup.32 q1, r0
+; CHECK-NEXT:    movs r1, #0
+; CHECK-NEXT:    vmovlb.u16 q1, q1
+; CHECK-NEXT:    vstrh.32 q1, [r1, q0]
 ; CHECK-NEXT:    bx lr
 entry:
   %ext = trunc i32 %v to i16
   %splatinsert = insertelement <4 x i16> poison, i16 %ext, i32 0
   %splat = shufflevector <4 x i16> %splatinsert, <4 x i16> poison, <4 x i32> zeroinitializer
   call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %splat, <4 x i16*> %offs, i32 2, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
   ret void
 }