[ScalarizeMaskedMemIntrin] Scalarize constant mask expandload as shuffle(build_vector,pass_through)

[RKSimon]

Summary

As noticed on D66004, scalarization of an expandload with a constant mask as a chain of irregular loads+inserts makes it difficult to optimize before lowering, resulting in difficulties in merging loads etc.

This patch instead scalarizes the expansion to a build_vector(load0, load1, undef, load2,....) style pattern and then performs a blend shuffle with the pass through vector. This allows us to more easily make use of all the build_vector combines, merging of consecutive loads etc.

[spatel]

llvm/lib/CodeGen/ScalarizeMaskedMemIntrin.cpp

166

Would it make sense to do the same thing here? Ie, make a shared helper function.
It might be that later matching always gets this case, but it would still be an efficiency win to produce the blend shuffle here too?

630

-1 -> "UndefMaskElem"

[RKSimon]

llvm/lib/CodeGen/ScalarizeMaskedMemIntrin.cpp

166

Yes, I'm happy to do that as a followup if everyone agrees - I'd prefer to just get this one in first as its blocking D66004 and I'm keen to get that patch done (finally!),

I've been wondering how compress store of constant masks should be expressed as well.

[spatel]

I'm ok with this as-is (mod "-1" code nit) followed by cleanup, so LGTM.
Wait a bit to see if anyone else wants to comment though.

llvm/lib/CodeGen/ScalarizeMaskedMemIntrin.cpp

static void scalarizeMaskedLoad(CallInst *CI, bool &ModifiedDT) {
   Type *NewPtrType =
       EltTy->getPointerTo(Ptr->getType()->getPointerAddressSpace());
   Value *FirstEltPtr = Builder.CreateBitCast(Ptr, NewPtrType);
   unsigned VectorWidth = cast<FixedVectorType>(VecType)->getNumElements();
 
   // The result vector
   Value *VResult = Src0;
 
   if (isConstantIntVector(Mask)) {
     for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
       if (cast<Constant>(Mask)->getAggregateElement(Idx)->isNullValue())
         continue;
       Value *Gep = Builder.CreateConstInBoundsGEP1_32(EltTy, FirstEltPtr, Idx);
       LoadInst *Load = Builder.CreateAlignedLoad(EltTy, Gep, AdjustedAlignVal);
       VResult = Builder.CreateInsertElement(VResult, Load, Idx);
     }
     CI->replaceAllUsesWith(VResult);
static void scalarizeMaskedExpandLoad(CallInst *CI, bool &ModifiedDT) {
   Builder.SetCurrentDebugLocation(CI->getDebugLoc());
 
   unsigned VectorWidth = VecType->getNumElements();
 
   // The result vector
   Value *VResult = PassThru;
 
   // Shorten the way if the mask is a vector of constants.
+  // Create a build_vector pattern, with loads/undefs as necessary and then
+  // shuffle blend with the pass through value.
   if (isConstantIntVector(Mask)) {
     unsigned MemIndex = 0;
+    VResult = UndefValue::get(VecType);
+    SmallVector<int, 16> ShuffleMask(VectorWidth, -1);
     for (unsigned Idx = 0; Idx < VectorWidth; ++Idx) {
-      if (cast<Constant>(Mask)->getAggregateElement(Idx)->isNullValue())
-        continue;
-      Value *NewPtr = Builder.CreateConstInBoundsGEP1_32(EltTy, Ptr, MemIndex);
-      LoadInst *Load = Builder.CreateAlignedLoad(EltTy, NewPtr, Align(1),
-                                                 "Load" + Twine(Idx));
-      VResult =
-          Builder.CreateInsertElement(VResult, Load, Idx, "Res" + Twine(Idx));
+      Value *InsertElt;
+      if (cast<Constant>(Mask)->getAggregateElement(Idx)->isNullValue()) {
+        InsertElt = UndefValue::get(EltTy);
+        ShuffleMask[Idx] = Idx + VectorWidth;
+      } else {
+        Value *NewPtr =
+            Builder.CreateConstInBoundsGEP1_32(EltTy, Ptr, MemIndex);
+        InsertElt = Builder.CreateAlignedLoad(EltTy, NewPtr, Align(1),
+        ShuffleMask[Idx] = Idx;
         ++MemIndex;
       }
+      VResult = Builder.CreateInsertElement(VResult, InsertElt, Idx,
+                                            "Res" + Twine(Idx));
+    }
+    VResult = Builder.CreateShuffleVector(VResult, PassThru, ShuffleMask);
     CI->replaceAllUsesWith(VResult);
     CI->eraseFromParent();
     return;
   }
 
   // If the mask is not v1i1, use scalar bit test operations. This generates
   // better results on X86 at least.
   Value *SclrMask;

llvm/test/CodeGen/X86/masked_expandload.ll

; AVX512VLBW-NEXT:    retq
   %mask = icmp eq <2 x i32> %trigger, zeroinitializer
   %res = call <2 x float> @llvm.masked.expandload.v2f32(float* %base, <2 x i1> %mask, <2 x float> %src0)
   ret <2 x float> %res
 }
 
 define <4 x float> @expandload_v4f32_const(float* %base, <4 x float> %src0) {
 ; SSE2-LABEL: expandload_v4f32_const:
 ; SSE2:       ## %bb.0:
-; SSE2-NEXT:    movss (%rdi), %xmm2 ## xmm2 = mem[0],zero,zero,zero
-; SSE2-NEXT:    movss 4(%rdi), %xmm1 ## xmm1 = mem[0],zero,zero,zero
-; SSE2-NEXT:    shufps {{.*#+}} xmm1 = xmm1[0,0],xmm2[0,0]
 ; SSE2-NEXT:    movss 8(%rdi), %xmm2 ## xmm2 = mem[0],zero,zero,zero
-; SSE2-NEXT:    shufps {{.*#+}} xmm2 = xmm2[0,0],xmm0[3,0]
-; SSE2-NEXT:    shufps {{.*#+}} xmm1 = xmm1[2,0],xmm2[0,2]
+; SSE2-NEXT:    movsd (%rdi), %xmm1 ## xmm1 = mem[0],zero
+; SSE2-NEXT:    movaps %xmm1, %xmm3
+; SSE2-NEXT:    movlhps {{.*#+}} xmm3 = xmm3[0],xmm2[0]
+; SSE2-NEXT:    shufps {{.*#+}} xmm0 = xmm0[3,0],xmm3[2,0]
+; SSE2-NEXT:    shufps {{.*#+}} xmm1 = xmm1[0,1],xmm0[2,0]
 ; SSE2-NEXT:    movaps %xmm1, %xmm0
 ; SSE2-NEXT:    retq
 ;
 ; SSE42-LABEL: expandload_v4f32_const:
 ; SSE42:       ## %bb.0:
-; SSE42-NEXT:    movss (%rdi), %xmm1 ## xmm1 = mem[0],zero,zero,zero
-; SSE42-NEXT:    blendps {{.*#+}} xmm0 = xmm1[0],xmm0[1,2,3]
-; SSE42-NEXT:    insertps $16, 4(%rdi), %xmm0 ## xmm0 = xmm0[0],mem[0],xmm0[2,3]
-; SSE42-NEXT:    insertps $32, 8(%rdi), %xmm0 ## xmm0 = xmm0[0,1],mem[0],xmm0[3]
+; SSE42-NEXT:    movsd (%rdi), %xmm1 ## xmm1 = mem[0],zero
+; SSE42-NEXT:    insertps $32, 8(%rdi), %xmm1 ## xmm1 = xmm1[0,1],mem[0],xmm1[3]
+; SSE42-NEXT:    blendps {{.*#+}} xmm0 = xmm1[0,1,2],xmm0[3]
 ; SSE42-NEXT:    retq
 ;
 ; AVX1OR2-LABEL: expandload_v4f32_const:
 ; AVX1OR2:       ## %bb.0:
-; AVX1OR2-NEXT:    vmovss (%rdi), %xmm1 ## xmm1 = mem[0],zero,zero,zero
-; AVX1OR2-NEXT:    vblendps {{.*#+}} xmm0 = xmm1[0],xmm0[1,2,3]
-; AVX1OR2-NEXT:    vinsertps $16, 4(%rdi), %xmm0, %xmm0 ## xmm0 = xmm0[0],mem[0],xmm0[2,3]
-; AVX1OR2-NEXT:    vinsertps $32, 8(%rdi), %xmm0, %xmm0 ## xmm0 = xmm0[0,1],mem[0],xmm0[3]
+; AVX1OR2-NEXT:    vmovsd (%rdi), %xmm1 ## xmm1 = mem[0],zero
+; AVX1OR2-NEXT:    vinsertps $32, 8(%rdi), %xmm1, %xmm1 ## xmm1 = xmm1[0,1],mem[0],xmm1[3]
+; AVX1OR2-NEXT:    vblendps {{.*#+}} xmm0 = xmm1[0,1,2],xmm0[3]
 ; AVX1OR2-NEXT:    retq
 ;
 ; AVX512F-LABEL: expandload_v4f32_const:
 ; AVX512F:       ## %bb.0:
 ; AVX512F-NEXT:    ## kill: def $xmm0 killed $xmm0 def $zmm0
 ; AVX512F-NEXT:    movw $7, %ax
 ; AVX512F-NEXT:    kmovw %eax, %k1
 ; AVX512F-NEXT:    vexpandps (%rdi), %zmm0 {%k1}

 ; AVX512VLBW-NEXT:    retq
   %res = call <4 x float> @llvm.masked.expandload.v4f32(float* %base, <4 x i1> <i1 true, i1 true, i1 true, i1 false>, <4 x float> %src0)
   ret <4 x float>%res
 }
 
 define <16 x float> @expandload_v16f32_const(float* %base, <16 x float> %src0) {
 ; SSE2-LABEL: expandload_v16f32_const:
 ; SSE2:       ## %bb.0:
+; SSE2-NEXT:    movss 52(%rdi), %xmm0 ## xmm0 = mem[0],zero,zero,zero
+; SSE2-NEXT:    movsd 44(%rdi), %xmm4 ## xmm4 = mem[0],zero
+; SSE2-NEXT:    movaps %xmm4, %xmm6
+; SSE2-NEXT:    movlhps {{.*#+}} xmm6 = xmm6[0],xmm0[0]
+; SSE2-NEXT:    movss 40(%rdi), %xmm0 ## xmm0 = mem[0],zero,zero,zero
+; SSE2-NEXT:    movsd 32(%rdi), %xmm5 ## xmm5 = mem[0],zero
+; SSE2-NEXT:    movaps %xmm5, %xmm7
+; SSE2-NEXT:    movlhps {{.*#+}} xmm7 = xmm7[0],xmm0[0]
 ; SSE2-NEXT:    movups (%rdi), %xmm0
 ; SSE2-NEXT:    movups 16(%rdi), %xmm1
-; SSE2-NEXT:    movss 32(%rdi), %xmm5 ## xmm5 = mem[0],zero,zero,zero
-; SSE2-NEXT:    movss 36(%rdi), %xmm4 ## xmm4 = mem[0],zero,zero,zero
-; SSE2-NEXT:    shufps {{.*#+}} xmm4 = xmm4[0,0],xmm5[0,0]
-; SSE2-NEXT:    movss 40(%rdi), %xmm5 ## xmm5 = mem[0],zero,zero,zero
-; SSE2-NEXT:    shufps {{.*#+}} xmm5 = xmm5[0,0],xmm2[3,0]
-; SSE2-NEXT:    shufps {{.*#+}} xmm4 = xmm4[2,0],xmm5[0,2]
-; SSE2-NEXT:    movss 44(%rdi), %xmm2 ## xmm2 = mem[0],zero,zero,zero
-; SSE2-NEXT:    movss 48(%rdi), %xmm5 ## xmm5 = mem[0],zero,zero,zero
-; SSE2-NEXT:    shufps {{.*#+}} xmm5 = xmm5[0,0],xmm2[0,0]
-; SSE2-NEXT:    movss 52(%rdi), %xmm2 ## xmm2 = mem[0],zero,zero,zero
-; SSE2-NEXT:    shufps {{.*#+}} xmm2 = xmm2[0,0],xmm3[3,0]
-; SSE2-NEXT:    shufps {{.*#+}} xmm5 = xmm5[2,0],xmm2[0,2]
-; SSE2-NEXT:    movaps %xmm4, %xmm2
-; SSE2-NEXT:    movaps %xmm5, %xmm3
+; SSE2-NEXT:    shufps {{.*#+}} xmm2 = xmm2[3,0],xmm7[2,0]
+; SSE2-NEXT:    shufps {{.*#+}} xmm5 = xmm5[0,1],xmm2[2,0]
+; SSE2-NEXT:    shufps {{.*#+}} xmm3 = xmm3[3,0],xmm6[2,0]
+; SSE2-NEXT:    shufps {{.*#+}} xmm4 = xmm4[0,1],xmm3[2,0]
+; SSE2-NEXT:    movaps %xmm5, %xmm2
+; SSE2-NEXT:    movaps %xmm4, %xmm3
 ; SSE2-NEXT:    retq
 ;
 ; SSE42-LABEL: expandload_v16f32_const:
 ; SSE42:       ## %bb.0:
+; SSE42-NEXT:    movsd 44(%rdi), %xmm4 ## xmm4 = mem[0],zero
+; SSE42-NEXT:    insertps $32, 52(%rdi), %xmm4 ## xmm4 = xmm4[0,1],mem[0],xmm4[3]
+; SSE42-NEXT:    movsd 32(%rdi), %xmm5 ## xmm5 = mem[0],zero
+; SSE42-NEXT:    insertps $32, 40(%rdi), %xmm5 ## xmm5 = xmm5[0,1],mem[0],xmm5[3]
 ; SSE42-NEXT:    movups (%rdi), %xmm0
 ; SSE42-NEXT:    movups 16(%rdi), %xmm1
-; SSE42-NEXT:    movss 32(%rdi), %xmm4 ## xmm4 = mem[0],zero,zero,zero
-; SSE42-NEXT:    blendps {{.*#+}} xmm2 = xmm4[0],xmm2[1,2,3]
-; SSE42-NEXT:    insertps $16, 36(%rdi), %xmm2 ## xmm2 = xmm2[0],mem[0],xmm2[2,3]
-; SSE42-NEXT:    insertps $32, 40(%rdi), %xmm2 ## xmm2 = xmm2[0,1],mem[0],xmm2[3]
-; SSE42-NEXT:    movss 44(%rdi), %xmm4 ## xmm4 = mem[0],zero,zero,zero
-; SSE42-NEXT:    blendps {{.*#+}} xmm3 = xmm4[0],xmm3[1,2,3]
-; SSE42-NEXT:    insertps $16, 48(%rdi), %xmm3 ## xmm3 = xmm3[0],mem[0],xmm3[2,3]
-; SSE42-NEXT:    insertps $32, 52(%rdi), %xmm3 ## xmm3 = xmm3[0,1],mem[0],xmm3[3]
+; SSE42-NEXT:    blendps {{.*#+}} xmm2 = xmm5[0,1,2],xmm2[3]
+; SSE42-NEXT:    blendps {{.*#+}} xmm3 = xmm4[0,1,2],xmm3[3]
 ; SSE42-NEXT:    retq
 ;
 ; AVX1OR2-LABEL: expandload_v16f32_const:
 ; AVX1OR2:       ## %bb.0:
-; AVX1OR2-NEXT:    vmovups (%rdi), %xmm0
-; AVX1OR2-NEXT:    vmovsd 16(%rdi), %xmm2 ## xmm2 = mem[0],zero
-; AVX1OR2-NEXT:    vinsertps $32, 24(%rdi), %xmm2, %xmm2 ## xmm2 = xmm2[0,1],mem[0],xmm2[3]
-; AVX1OR2-NEXT:    vinsertps $48, 28(%rdi), %xmm2, %xmm2 ## xmm2 = xmm2[0,1,2],mem[0]
-; AVX1OR2-NEXT:    vinsertf128 $1, %xmm2, %ymm0, %ymm0
-; AVX1OR2-NEXT:    vmovss 32(%rdi), %xmm2 ## xmm2 = mem[0],zero,zero,zero
-; AVX1OR2-NEXT:    vblendps {{.*#+}} xmm2 = xmm2[0],xmm1[1,2,3]
-; AVX1OR2-NEXT:    vinsertps $16, 36(%rdi), %xmm2, %xmm2 ## xmm2 = xmm2[0],mem[0],xmm2[2,3]
+; AVX1OR2-NEXT:    vmovsd 44(%rdi), %xmm0 ## xmm0 = mem[0],zero
+; AVX1OR2-NEXT:    vinsertps $32, 52(%rdi), %xmm0, %xmm0 ## xmm0 = xmm0[0,1],mem[0],xmm0[3]
+; AVX1OR2-NEXT:    vmovsd 32(%rdi), %xmm2 ## xmm2 = mem[0],zero
 ; AVX1OR2-NEXT:    vinsertps $32, 40(%rdi), %xmm2, %xmm2 ## xmm2 = xmm2[0,1],mem[0],xmm2[3]
-; AVX1OR2-NEXT:    vmovss 44(%rdi), %xmm3 ## xmm3 = mem[0],zero,zero,zero
-; AVX1OR2-NEXT:    vextractf128 $1, %ymm1, %xmm1
-; AVX1OR2-NEXT:    vblendps {{.*#+}} xmm1 = xmm3[0],xmm1[1,2,3]
-; AVX1OR2-NEXT:    vinsertps $16, 48(%rdi), %xmm1, %xmm1 ## xmm1 = xmm1[0],mem[0],xmm1[2,3]
-; AVX1OR2-NEXT:    vinsertps $32, 52(%rdi), %xmm1, %xmm1 ## xmm1 = xmm1[0,1],mem[0],xmm1[3]
-; AVX1OR2-NEXT:    vinsertf128 $1, %xmm1, %ymm2, %ymm1
+; AVX1OR2-NEXT:    vinsertf128 $1, %xmm0, %ymm2, %ymm2
+; AVX1OR2-NEXT:    vmovups (%rdi), %ymm0
+; AVX1OR2-NEXT:    vblendps {{.*#+}} ymm1 = ymm2[0,1,2],ymm1[3],ymm2[4,5,6],ymm1[7]
 ; AVX1OR2-NEXT:    retq
 ;
 ; AVX512F-LABEL: expandload_v16f32_const:
 ; AVX512F:       ## %bb.0:
 ; AVX512F-NEXT:    movw $30719, %ax ## imm = 0x77FF
 ; AVX512F-NEXT:    kmovw %eax, %k1
 ; AVX512F-NEXT:    vexpandps (%rdi), %zmm0 {%k1}
 ; AVX512F-NEXT:    retq