[InstCombine] Fold for masked scatters to a uniform address

[CarolineConcatto]

Summary

When masked scatter intrinsic does a uniform store to a destination
address from a source vector, and in this case, the mask is all one value.
This patch replaces the masked scatter with an extracted element of the
last lane of the source vector and stores it in the destination vector.

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

378

I think you can distinguish two cases here:

1. The pointer is a splat(), the value is not a splat, the mask is an all-active mask.
2. The pointer is a splat(), the value is a splat(), the mask is not all-inactive (but not necessarily all active).

For 2, you can extract any lane instead of the last lane.

379

What is operand 1? Can you change Splat to something more descriptive? e.g. SplatPtr ?

381

nit: WideLoadTy ?

382

this must be cast<VectorType>

386

This is only valid if the the type is a ScalableVectorType

[david-arm]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

378

For 2, I don't think you can extract any lane - it has to be the last active lane otherwise it's not functionally correct since that's the only thing that matches the scalar equivalent? The trouble with extracting the last active lane is that the resulting code is potentially going to be worse than a scatter I think? I don't think there is a trivial way in LLVM IR to ask for the last active lane.

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

378

If all lanes contain the same value (the value is a splat()), and at least one lane is active, why couldn't you extract any lane?

[david-arm]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

378

Ah sorry, I misunderstood your original comment. I see what you mean now - if the values are all the same yeah you're right.

[CarolineConcatto]

• Add fold for the value when it is a splat.
• Add tests for fixed vector

[CarolineConcatto]

Thank you @sdesmalen and @david-arm.
I believe both folds:

1. The pointer is a splat(), the value is not a splat, the mask is an all-active mask.
2. The pointer is a splat(), the value is a splat(), the mask is not all-inactive (but not necessarily all active). are now implemented.

I noticed that for option 2 (when the value is a splat) in the extract element the value is the same as the original value that was splatted.
So in case 2, the fold is to a scalar store of value in the destination pointer. I don't need to extract the element from any line.
So I only have:
StoreInst *S = new StoreInst(SplatValue, SplatPtr, false, Alignment);
I also added the same tests for fixed and scalable vector.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

386

I believe this fold should also work for fixed vector.
So I've update it to work for fixed vector too.
And I added tests for it as well.

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

379

This case still needs a check to ensure that !ConstMask->isNullValue(), because at least one lane must be active for this to be valid. (see my comment on your test to explain why this is currently untested)

385

nit: Vector splat address with all-active mask -> scalar store of last lane.

386–388

These lines of comments can be dropped if you take the previous suggestion.

396

nit: remove comment, the line of code below says as much.

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

285

This is still an all-active mask though (contrary to the comment on this test), so this is not testing the code you're aiming to test.

Also, fixed-width vectors can be written like this: <4 x i1> <i32 1, i32 1, i32 1, i32 0> (for a mask that is not all-active, but also not all-inactive, which I think is what you're after)

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

16 ↗

(On Diff #394832)

instead of passing an all true mask, you could simply insert 1 i1 1 element into a zeroinitializer

[CarolineConcatto]

• Address comments in the test
• Remove redundant comment in the function

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

379

As discussed internally this is done in line 371:

// If the mask is all zeros, a scatter does nothing.
if (ConstMask->isNullValue())
  return eraseInstFromFunction(II);

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

285

I have changed all the fixed vector to be as you said.
<i32 X, i32 X, i32 X, i32 X>, where X can be 0 or 1 depending on what is being tested.

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

16 ↗

(On Diff #394832)

I have updated the all zero mask to be:
<vscale x 4 x i1> zeroinitializer,
I don't know if I follow what you want here.

[sdesmalen]

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

16 ↗

(On Diff #394832)

The comment says:

;; Value splat and  mask is not not all active/one

Yet in the test it uses an all-poison mask:

<vscale x 4 x i1> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> poison, i1 true, i32 1), <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer

this inserts i1 1 into element 1 (not zero) of a poison vector, so you get:

<poison, 1, poison, poison, ... >

and then does a splat of element 0, leading to

<poison, poison, poison, poison, ...>

If you want to pass a mask that is not all active *and* not all ones *and* not all poison, I think you want something like this:

<vscale x 4 x i1> insertelement (<vscale x 4 x i1> zeroinitializer, i1 true, i32 1)
<=>
<0, 1, 0, 0, ...>

[CarolineConcatto]

-Changed the mask to be an input for vscale test when it is
not all active or inactive

[CarolineConcatto]

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

16 ↗

(On Diff #394832)

I have done something different, suggested by David.
Mask is an input, so we don't know if it is all active or inactive.

[sdesmalen]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

377–378

nit: // scatter(splat(value), splat(ptr), non-zero-mask) -> store value, ptr

381

nit: /*IsVolatile=*/ false

385

nit: // scatter(vector, splat(ptr), splat(true)) -> store extract(vector, lastlane), ptr

396

nit: /*IsVolatile=*/ false

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

274

nit: The scatters in the tests below take a splatted pointer

276

can you stick with only 1 terminology? i.e. just 'all active'.

277

Like mentioned on D115726, I'd avoid using 'valid' and 'invalid', but instead do:

@scatter_v4i16_uniform_vals_uniform_ptrs_no_all_active_mask

304

This case isn't affected by your patch, so I would just remove it.

309

The mask is <0, 0, 1, 1> which is not all active.

Maybe place these under 'negative tests', since they cannot be optimised.

350–351

Just pass a <4 x i16*> as argument, it's a bit opaque that this is not a splat. (i.e. you're splatting element zero, which is not where %dst is inserted, but that's not an obvious pattern to spot)

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

8 ↗

(On Diff #395106)

I think the tests for fixed-width vectors are largely sufficient, only for the two positive tests/cases I think you should add a scalable-vector test.

[CarolineConcatto]

• Move positive tests in vscale_masked_intrinsics.ll to masked_intrinsics.ll
• Rename tests as suggested in D115726
• Fix test scatter_nxv4i16_uniform_vals_uniform_ptrs_all_active_mask

[CarolineConcatto]

• Remove sve attribute from masked_intrinsics.ll

[sdesmalen]

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

274

nit: ; Test scatters that can be simplified to scalar stores.

292

this test seems is redundant, because it's covered by the test above.

307

nit: remove one of the newlines here

346

nit: can you add a comment here that these are negative tests?
(and change the name of the functions to @negative_scatter_...

[CarolineConcatto]

• Improve comments in the test file

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

292

Yes, it is, but this one is with vscale. It is the same on gather.
This is a positive test, and not negative!

[sdesmalen]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

 Instruction *InstCombinerImpl::simplifyMaskedGather(IntrinsicInst &II) {
   return nullptr;
 }
 
 // TODO, Obvious Missing Transforms:
 // * Single constant active lane -> store
 // * Adjacent vector addresses -> masked.store
 // * Narrow store width by halfs excluding zero/undef lanes
-// * Vector splat address w/known mask -> scalar store
 // * Vector incrementing address -> vector masked store
 Instruction *InstCombinerImpl::simplifyMaskedScatter(IntrinsicInst &II) {
   auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(3));
   if (!ConstMask)
     return nullptr;
 
   // If the mask is all zeros, a scatter does nothing.
   if (ConstMask->isNullValue())
     return eraseInstFromFunction(II);
 
+  // Vector splat address -> scalar store
+  if (auto *SplatPtr = getSplatValue(II.getArgOperand(1))) {
+    // scatter(splat(value), splat(ptr), non-zero-mask) -> store value, ptr
+    if (auto *SplatValue = getSplatValue(II.getArgOperand(0))) {
+      Align Alignment = cast<ConstantInt>(II.getArgOperand(2))->getAlignValue();
+      StoreInst *S =
+          new StoreInst(SplatValue, SplatPtr, /*IsVolatile=*/false, Alignment);
+      S->copyMetadata(II);
+      return S;
+    }
+    // scatter(vector, splat(ptr), splat(true)) -> store extract(vector,
+    // lastlane), ptr
+    if (ConstMask->isAllOnesValue()) {
+      Align Alignment = cast<ConstantInt>(II.getArgOperand(2))->getAlignValue();
+      VectorType *WideLoadTy = cast<VectorType>(II.getArgOperand(1)->getType());
+      ElementCount VF = WideLoadTy->getElementCount();
+      Constant *EC =
+          ConstantInt::get(Builder.getInt32Ty(), VF.getKnownMinValue());
+      Value *RunTimeVF = VF.isScalable() ? Builder.CreateVScale(EC) : EC;
+      Value *LastLane = Builder.CreateSub(RunTimeVF, Builder.getInt32(1));
+      Value *Extract =
+          Builder.CreateExtractElement(II.getArgOperand(0), LastLane);
+      StoreInst *S =
+          new StoreInst(Extract, SplatPtr, /*IsVolatile=*/false, Alignment);
+      S->copyMetadata(II);
+      return S;
+    }
+  }
   if (isa<ScalableVectorType>(ConstMask->getType()))
     return nullptr;
 
   // Use masked off lanes to simplify operands via SimplifyDemandedVectorElts
   APInt DemandedElts = possiblyDemandedEltsInMask(ConstMask);
   APInt UndefElts(DemandedElts.getBitWidth(), 0);
   if (Value *V =
           SimplifyDemandedVectorElts(II.getOperand(0), DemandedElts, UndefElts))

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

 ; CHECK-NEXT:    ret void
 ;
   %ptrs = getelementptr double, double* %ptr, <2 x i64> <i64 0, i64 1>
   %valvec1 = insertelement <2 x double> undef, double %val, i32 0
   %valvec2 = insertelement <2 x double> %valvec1, double %val, i32 1
   call void @llvm.masked.scatter.v2f64.v2p0f64(<2 x double> %valvec2, <2 x double*> %ptrs, i32 8, <2 x i1> <i1 true, i1 false>)
   ret void
 }
+
+
+; Test scatters that can be simplified to scalar stores.
+
+;; Value splat (mask is not used)
+define void @scatter_v4i16_uniform_vals_uniform_ptrs_no_all_active_mask(i16* %dst, i16 %val) {
+; CHECK-LABEL: @scatter_v4i16_uniform_vals_uniform_ptrs_no_all_active_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i16 [[VAL:%.*]], i16* [[DST:%.*]], align 2
+; CHECK-NEXT:    ret void
+;
+entry:
+  %broadcast.splatinsert = insertelement <4 x i16*> poison, i16* %dst, i32 0
+  %broadcast.splat = shufflevector <4 x i16*> %broadcast.splatinsert, <4 x i16*> poison, <4 x i32> zeroinitializer
+  %broadcast.value = insertelement <4 x i16> poison, i16 %val, i32 0
+  %broadcast.splatvalue = shufflevector <4 x i16> %broadcast.value, <4 x i16> poison, <4 x i32> zeroinitializer
+  call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %broadcast.splatvalue, <4 x i16*> %broadcast.splat, i32 2, <4 x i1> <i1 0, i1 0, i1 1, i1 1>)
+  ret void
+}
+
+define void @scatter_nxv4i16_uniform_vals_uniform_ptrs_all_active_mask(i16* %dst, i16 %val) {
+; CHECK-LABEL: @scatter_nxv4i16_uniform_vals_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i16 [[VAL:%.*]], i16* [[DST:%.*]], align 2
+; CHECK-NEXT:    ret void
+;
+entry:
+  %broadcast.splatinsert = insertelement <vscale x 4 x i16*> poison, i16* %dst, i32 0
+  %broadcast.splat = shufflevector <vscale x 4 x i16*> %broadcast.splatinsert, <vscale x 4 x i16*> poison, <vscale x 4 x i32> zeroinitializer
+  %broadcast.value = insertelement <vscale x 4 x i16> poison, i16 %val, i32 0
+  %broadcast.splatvalue = shufflevector <vscale x 4 x i16> %broadcast.value, <vscale x 4 x i16> poison, <vscale x 4 x i32> zeroinitializer
+  call void @llvm.masked.scatter.nxv4i16.nxv4p0i16(<vscale x 4 x i16> %broadcast.splatvalue, <vscale x 4 x i16*> %broadcast.splat, i32 2, <vscale x 4 x i1> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> zeroinitializer , i1 true, i32 0), <vscale x 4 x i1> zeroinitializer, <vscale x 4 x i32> zeroinitializer))
+  ret void
+}
+
+;; The pointer is splat and mask is all active, but value is not a splat
+define void @scatter_v4i16_no_uniform_vals_uniform_ptrs_all_active_mask(i16* %dst, <4 x i16>*  %src)  {
+; CHECK-LABEL: @scatter_v4i16_no_uniform_vals_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x i16>, <4 x i16>* [[SRC:%.*]], align 2
+; CHECK-NEXT:    [[TMP0:%.*]] = extractelement <4 x i16> [[WIDE_LOAD]], i64 3
+; CHECK-NEXT:    store i16 [[TMP0]], i16* [[DST:%.*]], align 2
+; CHECK-NEXT:    ret void
+;
+entry:
+  %broadcast.splatinsert = insertelement <4 x i16*> poison, i16* %dst, i32 0
+  %broadcast.splat = shufflevector <4 x i16*> %broadcast.splatinsert, <4 x i16*> poison, <4 x i32> zeroinitializer
+  %wide.load = load <4 x i16>, <4 x i16>* %src, align 2
+  call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %wide.load, <4 x i16*> %broadcast.splat, i32 2, <4 x i1> <i1 1, i1 1, i1 1, i1 1>)
+  ret void
+}
+
+define void @scatter_nxv4i16_no_uniform_vals_uniform_ptrs_all_active_mask(i16* %dst, <vscale x 4 x i16>* %src) {
+; CHECK-LABEL: @scatter_nxv4i16_no_uniform_vals_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <vscale x 4 x i16>, <vscale x 4 x i16>* [[SRC:%.*]], align 2
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[TMP0]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[TMP1]], -1
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <vscale x 4 x i16> [[WIDE_LOAD]], i32 [[TMP2]]
+; CHECK-NEXT:    store i16 [[TMP3]], i16* [[DST:%.*]], align 2
+; CHECK-NEXT:    ret void
+;
+entry:
+  %broadcast.splatinsert = insertelement <vscale x 4 x i16*> poison, i16* %dst, i32 0
+  %broadcast.splat = shufflevector <vscale x 4 x i16*> %broadcast.splatinsert, <vscale x 4 x i16*> poison, <vscale x 4 x i32> zeroinitializer
+  %wide.load = load <vscale x 4 x i16>, <vscale x 4 x i16>* %src, align 2
+  call void @llvm.masked.scatter.nxv4i16.nxv4p0i16(<vscale x 4 x i16> %wide.load, <vscale x 4 x i16*> %broadcast.splat, i32 2, <vscale x 4 x i1> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> poison, i1 true, i32 0), <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer))
+  ret void
+}
+
+; Negative scatter tests
+
+;; Pointer is splat, but mask is not all active and  value is not a splat
+define void @negative_scatter_v4i16_no_uniform_vals_uniform_ptrs_all_inactive_mask(i16* %dst, <4 x i16>* %src) {
+; CHECK-LABEL: @negative_scatter_v4i16_no_uniform_vals_uniform_ptrs_all_inactive_mask(
+; CHECK-NEXT:    [[INSERT_ELT:%.*]] = insertelement <4 x i16*> poison, i16* [[DST:%.*]], i64 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT:%.*]] = shufflevector <4 x i16*> [[INSERT_ELT]], <4 x i16*> poison, <4 x i32> <i32 undef, i32 undef, i32 0, i32 0>
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x i16>, <4 x i16>* [[SRC:%.*]], align 2
+; CHECK-NEXT:    call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> [[WIDE_LOAD]], <4 x i16*> [[BROADCAST_SPLAT]], i32 2, <4 x i1> <i1 false, i1 false, i1 true, i1 true>)
+; CHECK-NEXT:    ret void
+;
+  %insert.elt = insertelement <4 x i16*> poison, i16* %dst, i32 0
+  %broadcast.splat = shufflevector <4 x i16*> %insert.elt, <4 x i16*> poison, <4 x i32> zeroinitializer
+  %wide.load = load <4 x i16>, <4 x i16>* %src, align 2
+  call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %wide.load, <4 x i16*> %broadcast.splat, i32 2, <4 x i1> <i1 0, i1 0, i1 1, i1 1>)
+  ret void
+}
+
+;; The pointer in NOT a splat
+define void @negative_scatter_v4i16_no_uniform_vals_no_uniform_ptrs_all_active_mask(<4 x i16*> %inPtr, <4 x i16>* %src) {
+; CHECK-LABEL: @negative_scatter_v4i16_no_uniform_vals_no_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:    [[BROADCAST:%.*]] = shufflevector <4 x i16*> [[INPTR:%.*]], <4 x i16*> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x i16>, <4 x i16>* [[SRC:%.*]], align 2
+; CHECK-NEXT:    call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> [[WIDE_LOAD]], <4 x i16*> [[BROADCAST]], i32 2, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
+; CHECK-NEXT:    ret void
+;
+  %broadcast= shufflevector <4 x i16*> %inPtr, <4 x i16*> poison, <4 x i32> zeroinitializer
+  %wide.load = load <4 x i16>, <4 x i16>* %src, align 2
+  call void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16> %wide.load, <4 x i16*> %broadcast, i32 2, <4 x i1> <i1 1, i1 1, i1 1, i1 1> )
+  ret void
+}
+
+
+; Function Attrs:
+declare void @llvm.masked.scatter.v4i16.v4p0i16(<4 x i16>, <4 x i16*>, i32 immarg, <4 x i1>)
+declare void @llvm.masked.scatter.nxv4i16.nxv4p0i16(<vscale x 4 x i16>, <vscale x 4 x i16*>, i32 immarg, <vscale x 4 x i1>)