[InstCombine] Fold for masked gather when loading the same value each time.

[CarolineConcatto]

Summary

This patch checks in the masked gather when the first operand value is a
splat and the mask is all one, because the masked gather is reloading the
same value each time. This patch replaces this pattern of masked gather by
a scalar load of the value and splats it in a vector.

[david-arm]

This looks like a nice change @CarolineConcatto! I just have a few comments about the tests.

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

4 ↗

(On Diff #394226)

nit: Could you rename this to @valid_invariant_load_i64 just because the shortened form inv could also mean invalid?

18 ↗

(On Diff #394226)

nit: Same here - perhaps rename to @invalid_value_invariant_load_i64?

26 ↗

(On Diff #394226)

This is actually just creating a fully poison vector here because you're splatting the first element, which is poison. Perhaps you can just delete the %broadcast.splat line and pass %broadcast.splatinsert directly to the masked.gather?

32 ↗

(On Diff #394226)

nit: Rename to something like @invalid_mask_invariant_load_i64?

41 ↗

(On Diff #394226)

This is equivalent to a fully poison vector because you're splatting the first element, which is poison. Perhaps you can just pass a predicate to the function:

define <vscale x 2 x i64> @invalid_mask_invariant_load_i64(i64* %src, <vscale x 2 x i1> %mask) #0 {
  ...
  %res = call <vscale x 2 x i64> @llvm.masked.gather.nxv2i64(<vscale x 2 x i64*> %broadcast.splat, i32 8, <vscale x 2 x i1> %mask, <vscale x 2 x i64> undef)

?

Alternatively, if you want to test for a constant mask that isn't all zeroes you can do something like:

%res = call <vscale x 2 x i64> @llvm.masked.gather.nxv2i64(<vscale x 2 x i64*> %broadcast.splat, i32 8, <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 false, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer), <vscale x 2 x i64> undef)

which is an all-false constant mask?

[CarolineConcatto]

-address comments in the test

[CarolineConcatto]

Thank you @david-arm for the suggestions.
I believe I have addressed all of them.

[david-arm]

LGTM! Thanks for making the changes @CarolineConcatto. :)

[sdesmalen]

Hi @CarolineConcatto, the changes looks good to me, just left a few more nits on the tests.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

370

nit: load.scalar ?

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

1 ↗

(On Diff #395115)

nit: From what I can see, there is nothing really scalable-specific about this transformation, so maybe just add it to llvm/test/Transforms/InstCombine/masked_intrinsics.ll?
(and also add a test for fixed-width vectors)

5 ↗

(On Diff #395115)

nit: Did you mean uniform?

I'd also avoid using 'valid' and 'invalid', because all IR seems valid, it's just not necessarily the right input to trigger your transformation.

How about:

@gather_nxv2i64_uniform_ptrs_all_active_mask
@gather_nxv2i64_non_uniform_ptrs_all_active_mask
@gather_nxv2i64_uniform_ptrs_no_all_active_mask

18 ↗

(On Diff #395115)

nit: Vector of pointers is not a splat value.

25 ↗

(On Diff #395115)

nit: maybe just pass a <vscale x 2 x i64*> as argument?

30 ↗

(On Diff #395115)

nit: Do you just mean 'Not all active mask' ?

[CarolineConcatto]

• Move tests from vscale_masked_intrinsics.ll to masked_intrinsics.ll .
• Add test when mask is all zero

[CarolineConcatto]

Thank you Sander!

llvm/test/Transforms/InstCombine/vscale_masked_intrinsics.ll

1 ↗

(On Diff #395115)

Ok, I hope that there is no problem that the file you are pointing has no vscale tests in there yet.

30 ↗

(On Diff #395115)

It is both, because it cannot be all zero, this returns Null

if (!ConstMask)
  return nullptr;

[david-arm]

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

375

Do we actually need this attribute? This test file is in a target-independent location so we should avoid using target-specific attributes here.

[CarolineConcatto]

• remove sve attribute from masked_intrinsics.ll
• remove negative test with vscale

[sdesmalen]

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

273

nit: add Test gathers that can be simplified to scalar load + splat

274

Splat pointer?

297

just write: <2 x i64*> <i64* %src, i64* %src>

300

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

328

I think you can remove this test, because it's unaffected by your patch.

[CarolineConcatto]

• Improve comments in the test file

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

274

Yes, you are correct. It is pointer. Missed that

297

I believe this does not work, because src needs to be a constant.
I need to use gep and inserts in each position of the vector. That is the same as doing a splat

328

That is not true. I added these lines:

auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(2));
if (!ConstMask)
  return nullptr;

[sdesmalen]

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

297

You're right, I didn't realise that this notation doesn't work for constants!

328

Returning nullptr means that InstCombine is not combining this case, i.e. that your combine function makes no changes.
You can see that in InstCombinerImpl::run where it calls the visit function.

In negative_gather_v2i64_uniform_ptrs_all_inactive_mask you're testing that the code is optimized away entirely if the mask is all-zero, but that's not caused by your code.

[CarolineConcatto]

• Remove test negative_gather_v2i64_uniform_ptrs_no_all_active_mask

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

Instruction *InstCombinerImpl::simplifyMaskedStore(IntrinsicInst &II) {
   return nullptr;
 }
 
 // TODO, Obvious Missing Transforms:
 // * Single constant active lane load -> load
 // * Dereferenceable address & few lanes -> scalarize speculative load/selects
 // * Adjacent vector addresses -> masked.load
 // * Narrow width by halfs excluding zero/undef lanes
-// * Vector splat address w/known mask -> scalar load
 // * Vector incrementing address -> vector masked load
 Instruction *InstCombinerImpl::simplifyMaskedGather(IntrinsicInst &II) {
+  auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(2));
+  if (!ConstMask)
+    return nullptr;
+
+  // Vector splat address w/known mask -> scalar load
+  // Fold the gather to load the source vector first lane
+  // because it is reloading the same value each time
+  if (ConstMask->isAllOnesValue())
+    if (auto *SplatPtr = getSplatValue(II.getArgOperand(0))) {
+      auto *VecTy = cast<VectorType>(II.getType());
+      const Align Alignment =
+          cast<ConstantInt>(II.getArgOperand(1))->getAlignValue();
+      LoadInst *L = Builder.CreateAlignedLoad(VecTy->getElementType(), SplatPtr,
+                                              Alignment, "load.scalar");
+      Value *Shuf =
+          Builder.CreateVectorSplat(VecTy->getElementCount(), L, "broadcast");
+      return replaceInstUsesWith(II, cast<Instruction>(Shuf));
+    }
+
   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 incrementing address -> vector masked store

llvm/test/Transforms/InstCombine/masked_intrinsics.ll

 ;
   %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)  {
entry:
   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

 ; 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>)
+
+; Test gathers that can be simplified to scalar load + splat
+
+;; Splat address and all active mask
+define <vscale x 2 x i64> @gather_nxv2i64_uniform_ptrs_all_active_mask(i64* %src) {
+; CHECK-LABEL: @gather_nxv2i64_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:    [[LOAD_SCALAR:%.*]] = load i64, i64* [[SRC:%.*]], align 8
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT1:%.*]] = insertelement <vscale x 2 x i64> poison, i64 [[LOAD_SCALAR]], i64 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT2:%.*]] = shufflevector <vscale x 2 x i64> [[BROADCAST_SPLATINSERT1]], <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 2 x i64> [[BROADCAST_SPLAT2]]
+;
+  %broadcast.splatinsert = insertelement <vscale x 2 x i64*> poison, i64 *%src, i32 0
+  %broadcast.splat = shufflevector <vscale x 2 x i64*> %broadcast.splatinsert, <vscale x 2 x i64*> poison, <vscale x 2 x i32> zeroinitializer
+  %res = call <vscale x 2 x i64> @llvm.masked.gather.nxv2i64(<vscale x 2 x i64*> %broadcast.splat, i32 8, <vscale x 2 x i1> shufflevector (<vscale x 2 x i1> insertelement (<vscale x 2 x i1> poison, i1 true, i32 0), <vscale x 2 x i1> poison, <vscale x 2 x i32> zeroinitializer), <vscale x 2 x i64> undef)
+  ret <vscale x 2 x i64> %res
+}
+
+define <2 x i64> @gather_v2i64_uniform_ptrs_all_active_mask(i64* %src) {
+; CHECK-LABEL: @gather_v2i64_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:    [[LOAD_SCALAR:%.*]] = load i64, i64* [[SRC:%.*]], align 8
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT1:%.*]] = insertelement <2 x i64> poison, i64 [[LOAD_SCALAR]], i64 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT2:%.*]] = shufflevector <2 x i64> [[BROADCAST_SPLATINSERT1]], <2 x i64> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <2 x i64> [[BROADCAST_SPLAT2]]
+;
+  %broadcast.splatinsert = insertelement <2 x i64*> poison, i64 *%src, i32 0
+  %broadcast.splat = shufflevector <2 x i64*> %broadcast.splatinsert, <2 x i64*> poison, <2 x i32> zeroinitializer
+  %res = call <2 x i64> @llvm.masked.gather.v2i64(<2 x i64*> %broadcast.splat, i32 8, <2 x i1> <i1 1, i1 1>, <2 x i64> undef)
+  ret <2 x i64> %res
+}
+
+; Negative gather tests
+
+;; Vector of pointers is not a splat.
+define <2 x i64> @negative_gather_v2i64_non_uniform_ptrs_all_active_mask(<2 x i64*> %inVal, i64* %src ) {
+; CHECK-LABEL: @negative_gather_v2i64_non_uniform_ptrs_all_active_mask(
+; CHECK-NEXT:    [[INSERT_VALUE:%.*]] = insertelement <2 x i64*> [[INVAL:%.*]], i64* [[SRC:%.*]], i64 1
+; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.masked.gather.v2i64.v2p0i64(<2 x i64*> [[INSERT_VALUE]], i32 8, <2 x i1> <i1 true, i1 true>, <2 x i64> undef)
+; CHECK-NEXT:    ret <2 x i64> [[RES]]
+;
+  %insert.value = insertelement <2 x i64*> %inVal, i64 *%src, i32 1
+  %res = call <2 x i64> @llvm.masked.gather.v2i64(<2 x i64*> %insert.value, i32 8, <2 x i1><i1 1, i1 1>, <2 x i64> undef)
+  ret <2 x i64> %res
+}
+
+;; Unknown mask value
+define <2 x i64> @negative_gather_v2i64_uniform_ptrs_no_all_active_mask(i64* %src, <2 x i1> %mask) {
+; CHECK-LABEL: @negative_gather_v2i64_uniform_ptrs_no_all_active_mask(
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT:%.*]] = insertelement <2 x i64*> poison, i64* [[SRC:%.*]], i64 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT:%.*]] = shufflevector <2 x i64*> [[BROADCAST_SPLATINSERT]], <2 x i64*> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[RES:%.*]] = call <2 x i64> @llvm.masked.gather.v2i64.v2p0i64(<2 x i64*> [[BROADCAST_SPLAT]], i32 8, <2 x i1> [[MASK:%.*]], <2 x i64> undef)
+; CHECK-NEXT:    ret <2 x i64> [[RES]]
+;
+  %broadcast.splatinsert = insertelement <2 x i64*> poison, i64 *%src, i32 0
+  %broadcast.splat = shufflevector <2 x i64*> %broadcast.splatinsert, <2 x i64*> poison, <2 x i32> zeroinitializer
+  %res = call <2 x i64> @llvm.masked.gather.v2i64(<2 x i64*> %broadcast.splat, i32 8, <2 x i1> %mask, <2 x i64> undef)
+  ret <2 x i64> %res
+}
+
+; Function Attrs:
+declare <vscale x 2 x i64> @llvm.masked.gather.nxv2i64(<vscale x 2 x i64*>, i32, <vscale x 2 x i1>, <vscale x 2 x i64>)
+declare <2 x i64> @llvm.masked.gather.v2i64(<2 x i64*>, i32, <2 x i1>, <2 x i64>)
+