Demanded elements support for masked.load and masked.gather

[reames]

Summary

Teach instcombine to propagate demanded elements through a masked load or masked gather instruction.

This is in the broader context of improving vector pointer instcombine under https://reviews.llvm.org/D57140.

[RKSimon]

rebase after rL352653 ?

[reames]

In D57372#1380603, @RKSimon wrote:

rebase after rL352653 ?

[reames]

Rebase

[craig.topper]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1435

I think DemandedLoad should maybe be called DemandedPtrVec.

[reames]

Address review comment re: naming.

I chose DemandedPtrs instead of DemandedPtrVec for consistency w/other variables in the same file.

[RKSimon]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1441

We're not 100% consistent with this but for selects we try to use:

if (CElt->isNullValue())
  DemandedPtrs.clearBit(i);
else
  DemandedPassThrough.clearBit(i);

[reames]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1441

This would be wrong. Consider any constant expr or other constant of indeterminate value.

[RKSimon]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1441

@spatel may be able to advise, but Instruction::Select does this:

Constant *CElt = CV->getAggregateElement(i);
if (isa<ConstantExpr>(CElt))
  continue;
// TODO: If a select condition element is undef, we can demand from
// either side. If one side is known undef, choosing that side would
// propagate undef.
if (CElt->isNullValue())
  DemandedLHS.clearBit(i);
else
  DemandedRHS.clearBit(i);

[spatel]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1441

I haven't been following the details of this review, but that code specifically guards against ConstantExpr. Is there some other constant-derived possibility that we failed to account for?

The code here looks safe since it is explicitly checking for constant 0 or -1, but it doesn't account for the possibility of an undef element in the mask.

I'm not sure what we want the result to be in that case. Either way, we probably want a regression test to account for that (and a test with a constant expression, so we're sure we don't crash on that?).

[reames]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

1441

I'm happy to add a test case for both constexpr, and undef. I specifically *do not* want to be aggressive about undef at this time though. That opens up a bunch of complexity since we need to be *consistent* about how we interpret an undef mask element, and I'd really rather separate that into it's own patch.

p.s. That TODO comment in Select looks wrong. We already interpret undef as falling into the else. I'm not sure that's correct, but that's a separate patch and a separate discussion.

p.p.s. In a future patch, abstracting out a helper function for this masked demanded bit splitting operation is planned.

[reames]

Rebase on top of requested test change.

[spatel]

LGTM

test/Transforms/InstCombine/masked_intrinsics.ll

38–45

This constant expression is too easy. :)
As we can see from the optimized check lines, we fold the ConstantExpr into a regular constant, so it's not providing the hoped-for code coverage.
I'd change this to something like:

@g = external global i8

define <2 x double> @load_cemask(<2 x double>* %ptr, <2 x double> %passthru)  {
  %res = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* %ptr, i32 2, <2 x i1> <i1 1, i1 ptrtoint (i8* @g to i1)>, <2 x double> %passthru)
  ret <2 x double> %res
}

[reames]

Be aware that this patch turned out to contain a nasty miscompile. It's now been fixed in https://reviews.llvm.org/rL358299. I recommend anyone who has branched with this change to consider cherry-picking the fix.

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

Value *InstCombiner::SimplifyDemandedVectorElts(Value *V, APInt DemandedElts,
   case Instruction::FPExt:
     simplifyAndSetOp(I, 0, DemandedElts, UndefElts);
     break;
 
   case Instruction::Call: {
     IntrinsicInst *II = dyn_cast<IntrinsicInst>(I);
     if (!II) break;
     switch (II->getIntrinsicID()) {
+    case Intrinsic::masked_gather: // fallthrough
+    case Intrinsic::masked_load: {
+      // Subtle - We *can't* simplify the zero bits of the mask just because
+      // they're unused because that might introduce a fault.
+      // TODO: but we can fold them to zero in a constant mask
+      APInt DemandedPtrs(DemandedElts), DemandedPassThrough(DemandedElts);
+      if (auto *CV = dyn_cast<ConstantVector>(II->getOperand(2)))
+        for (unsigned i = 0; i < VWidth; i++) {
+          Constant *CElt = CV->getAggregateElement(i);
+          if (CElt->isNullValue())
+            DemandedPtrs.clearBit(i);
+          else if (CElt->isAllOnesValue())
+            DemandedPassThrough.clearBit(i);
+        }
+      if (II->getIntrinsicID() == Intrinsic::masked_gather)
+        simplifyAndSetOp(II, 0, DemandedPtrs, UndefElts2);
+      simplifyAndSetOp(II, 3, DemandedPassThrough, UndefElts3);
+      
+      // Output elements are undefined if the element from both sources are.
+      // TODO: can strengthen via mask as well.
+      UndefElts = UndefElts2 & UndefElts3;
+      break;
+    }
     case Intrinsic::x86_xop_vfrcz_ss:
     case Intrinsic::x86_xop_vfrcz_sd:
       // The instructions for these intrinsics are speced to zero upper bits not
       // pass them through like other scalar intrinsics. So we shouldn't just
       // use Arg0 if DemandedElts[0] is clear like we do for other intrinsics.
       // Instead we should return a zero vector.
       if (!DemandedElts[0]) {
         Worklist.Add(II);

test/Transforms/InstCombine/X86/x86-masked-memops.ll

 
 }
 
 ; It's a constant mask, so convert to an LLVM intrinsic. The backend should optimize further.
 
 define <4 x float> @mload_one_one(i8* %f) {
 ; CHECK-LABEL: @mload_one_one(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <4 x float>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>* [[CASTVEC]], i32 1, <4 x i1> <i1 false, i1 false, i1 false, i1 true>, <4 x float> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>* [[CASTVEC]], i32 1, <4 x i1> <i1 false, i1 false, i1 false, i1 true>, <4 x float> <float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float undef>)
 ; CHECK-NEXT:    ret <4 x float> [[TMP1]]
 ;
   %ld = tail call <4 x float> @llvm.x86.avx.maskload.ps(i8* %f, <4 x i32> <i32 0, i32 0, i32 0, i32 -1>)
   ret <4 x float> %ld
 
 }
 
 ; Try doubles.
 
 define <2 x double> @mload_one_one_double(i8* %f) {
 ; CHECK-LABEL: @mload_one_one_double(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <2 x double>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* [[CASTVEC]], i32 1, <2 x i1> <i1 true, i1 false>, <2 x double> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* [[CASTVEC]], i32 1, <2 x i1> <i1 true, i1 false>, <2 x double> <double undef, double 0.000000e+00>)
 ; CHECK-NEXT:    ret <2 x double> [[TMP1]]
 ;
   %ld = tail call <2 x double> @llvm.x86.avx.maskload.pd(i8* %f, <2 x i64> <i64 -1, i64 0>)
   ret <2 x double> %ld
 
 }
 
 ; Try 256-bit FP ops.
 
 define <8 x float> @mload_v8f32(i8* %f) {
 ; CHECK-LABEL: @mload_v8f32(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <8 x float>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>* [[CASTVEC]], i32 1, <8 x i1> <i1 false, i1 false, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false>, <8 x float> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>* [[CASTVEC]], i32 1, <8 x i1> <i1 false, i1 false, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false>, <8 x float> <float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float undef, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float 0.000000e+00>)
 ; CHECK-NEXT:    ret <8 x float> [[TMP1]]
 ;
   %ld = tail call <8 x float> @llvm.x86.avx.maskload.ps.256(i8* %f, <8 x i32> <i32 0, i32 0, i32 0, i32 -1, i32 0, i32 0, i32 0, i32 0>)
   ret <8 x float> %ld
 
 }
 
 define <4 x double> @mload_v4f64(i8* %f) {
 ; CHECK-LABEL: @mload_v4f64(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <4 x double>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x double> @llvm.masked.load.v4f64.p0v4f64(<4 x double>* [[CASTVEC]], i32 1, <4 x i1> <i1 true, i1 false, i1 false, i1 false>, <4 x double> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x double> @llvm.masked.load.v4f64.p0v4f64(<4 x double>* [[CASTVEC]], i32 1, <4 x i1> <i1 true, i1 false, i1 false, i1 false>, <4 x double> <double undef, double 0.000000e+00, double 0.000000e+00, double 0.000000e+00>)
 ; CHECK-NEXT:    ret <4 x double> [[TMP1]]
 ;
   %ld = tail call <4 x double> @llvm.x86.avx.maskload.pd.256(i8* %f, <4 x i64> <i64 -1, i64 0, i64 0, i64 0>)
   ret <4 x double> %ld
 
 }
 
 ; Try the AVX2 variants.
 
 define <4 x i32> @mload_v4i32(i8* %f) {
 ; CHECK-LABEL: @mload_v4i32(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <4 x i32>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* [[CASTVEC]], i32 1, <4 x i1> <i1 false, i1 false, i1 false, i1 true>, <4 x i32> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* [[CASTVEC]], i32 1, <4 x i1> <i1 false, i1 false, i1 false, i1 true>, <4 x i32> <i32 0, i32 0, i32 0, i32 undef>)
 ; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %ld = tail call <4 x i32> @llvm.x86.avx2.maskload.d(i8* %f, <4 x i32> <i32 0, i32 0, i32 0, i32 -1>)
   ret <4 x i32> %ld
 
 }
 
 define <2 x i64> @mload_v2i64(i8* %f) {
 ; CHECK-LABEL: @mload_v2i64(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <2 x i64>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i64> @llvm.masked.load.v2i64.p0v2i64(<2 x i64>* [[CASTVEC]], i32 1, <2 x i1> <i1 true, i1 false>, <2 x i64> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i64> @llvm.masked.load.v2i64.p0v2i64(<2 x i64>* [[CASTVEC]], i32 1, <2 x i1> <i1 true, i1 false>, <2 x i64> <i64 undef, i64 0>)
 ; CHECK-NEXT:    ret <2 x i64> [[TMP1]]
 ;
   %ld = tail call <2 x i64> @llvm.x86.avx2.maskload.q(i8* %f, <2 x i64> <i64 -1, i64 0>)
   ret <2 x i64> %ld
 
 }
 
 define <8 x i32> @mload_v8i32(i8* %f) {
 ; CHECK-LABEL: @mload_v8i32(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <8 x i32>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <8 x i32> @llvm.masked.load.v8i32.p0v8i32(<8 x i32>* [[CASTVEC]], i32 1, <8 x i1> <i1 false, i1 false, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false>, <8 x i32> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <8 x i32> @llvm.masked.load.v8i32.p0v8i32(<8 x i32>* [[CASTVEC]], i32 1, <8 x i1> <i1 false, i1 false, i1 false, i1 true, i1 false, i1 false, i1 false, i1 false>, <8 x i32> <i32 0, i32 0, i32 0, i32 undef, i32 0, i32 0, i32 0, i32 0>)
 ; CHECK-NEXT:    ret <8 x i32> [[TMP1]]
 ;
   %ld = tail call <8 x i32> @llvm.x86.avx2.maskload.d.256(i8* %f, <8 x i32> <i32 0, i32 0, i32 0, i32 -1, i32 0, i32 0, i32 0, i32 0>)
   ret <8 x i32> %ld
 
 }
 
 define <4 x i64> @mload_v4i64(i8* %f) {
 ; CHECK-LABEL: @mload_v4i64(
 ; CHECK-NEXT:    [[CASTVEC:%.*]] = bitcast i8* [[F:%.*]] to <4 x i64>*
-; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i64> @llvm.masked.load.v4i64.p0v4i64(<4 x i64>* [[CASTVEC]], i32 1, <4 x i1> <i1 true, i1 false, i1 false, i1 false>, <4 x i64> zeroinitializer)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i64> @llvm.masked.load.v4i64.p0v4i64(<4 x i64>* [[CASTVEC]], i32 1, <4 x i1> <i1 true, i1 false, i1 false, i1 false>, <4 x i64> <i64 undef, i64 0, i64 0, i64 0>)
 ; CHECK-NEXT:    ret <4 x i64> [[TMP1]]
 ;
   %ld = tail call <4 x i64> @llvm.x86.avx2.maskload.q.256(i8* %f, <4 x i64> <i64 -1, i64 0, i64 0, i64 0>)
   ret <4 x i64> %ld
 
 }
 
 

test/Transforms/InstCombine/masked_intrinsics.ll

 ; CHECK-NEXT:    [[UNMASKEDLOAD:%.*]] = load <2 x double>, <2 x double>* [[PTR:%.*]], align 2
 ; CHECK-NEXT:    ret <2 x double> [[UNMASKEDLOAD]]
 ;
   %res = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* %ptr, i32 2, <2 x i1> <i1 1, i1 undef>, <2 x double> %passthru)
   ret <2 x double> %res
 
 }
 
 define <2 x double> @load_cemask(<2 x double>* %ptr, <2 x double> %passthru)  {
 ; CHECK-LABEL: @load_cemask(
 ; CHECK-NEXT:    [[RES:%.*]] = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* [[PTR:%.*]], i32 2, <2 x i1> <i1 true, i1 false>, <2 x double> [[PASSTHRU:%.*]])
 ; CHECK-NEXT:    ret <2 x double> [[RES]]
 ;
   %res = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* %ptr, i32 2, <2 x i1> <i1 1, i1 trunc (i32 0 to i1)>, <2 x double> %passthru)
   ret <2 x double> %res
 }
 
 define <2 x double> @load_lane0(<2 x double>* %ptr, double %pt)  {
 ; CHECK-LABEL: @load_lane0(
-; CHECK-NEXT:    [[PTV1:%.*]] = insertelement <2 x double> undef, double [[PT:%.*]], i64 0
-; CHECK-NEXT:    [[PTV2:%.*]] = shufflevector <2 x double> [[PTV1]], <2 x double> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[PTV2:%.*]] = insertelement <2 x double> undef, double [[PT:%.*]], i64 1
 ; CHECK-NEXT:    [[RES:%.*]] = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* [[PTR:%.*]], i32 2, <2 x i1> <i1 true, i1 false>, <2 x double> [[PTV2]])
 ; CHECK-NEXT:    ret <2 x double> [[RES]]
 ;
   %ptv1 = insertelement <2 x double> undef, double %pt, i64 0
   %ptv2 = insertelement <2 x double> %ptv1, double %pt, i64 1
   %res = call <2 x double> @llvm.masked.load.v2f64.p0v2f64(<2 x double>* %ptr, i32 2, <2 x i1> <i1 true, i1 false>, <2 x double> %ptv2)
   ret <2 x double> %res
 
;
   %res = call <2 x double> @llvm.masked.gather.v2f64.v2p0f64(<2 x double*> %ptrs, i32 5, <2 x i1> zeroinitializer, <2 x double> %passthru)
   ret <2 x double> %res
 
 }
 
 
 define <2 x double> @gather_onemask(<2 x double*> %ptrs, <2 x double> %passthru)  {
 ; CHECK-LABEL: @gather_onemask(
-; CHECK-NEXT:    [[RES:%.*]] = call <2 x double> @llvm.masked.gather.v2f64.v2p0f64(<2 x double*> [[PTRS:%.*]], i32 5, <2 x i1> <i1 true, i1 true>, <2 x double> [[PASSTHRU:%.*]])
+; CHECK-NEXT:    [[RES:%.*]] = call <2 x double> @llvm.masked.gather.v2f64.v2p0f64(<2 x double*> [[PTRS:%.*]], i32 5, <2 x i1> <i1 true, i1 true>, <2 x double> undef)
 ; CHECK-NEXT:    ret <2 x double> [[RES]]
 ;
   %res = call <2 x double> @llvm.masked.gather.v2f64.v2p0f64(<2 x double*> %ptrs, i32 5, <2 x i1> <i1 true, i1 true>, <2 x double> %passthru)
   ret <2 x double> %res
 
 }
 
 define <2 x double> @gather_lane0(double* %base, double %pt)  {
 ; CHECK-LABEL: @gather_lane0(
-; CHECK-NEXT:    [[PTRS:%.*]] = getelementptr double, double* [[BASE:%.*]], <2 x i64> <i64 0, i64 1>
-; CHECK-NEXT:    [[PT_V1:%.*]] = insertelement <2 x double> undef, double [[PT:%.*]], i64 0
-; CHECK-NEXT:    [[PT_V2:%.*]] = shufflevector <2 x double> [[PT_V1]], <2 x double> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[PTRS:%.*]] = getelementptr double, double* [[BASE:%.*]], <2 x i64> <i64 0, i64 undef>
+; CHECK-NEXT:    [[PT_V2:%.*]] = insertelement <2 x double> undef, double [[PT:%.*]], i64 1
 ; CHECK-NEXT:    [[RES:%.*]] = call <2 x double> @llvm.masked.gather.v2f64.v2p0f64(<2 x double*> [[PTRS]], i32 5, <2 x i1> <i1 true, i1 false>, <2 x double> [[PT_V2]])
 ; CHECK-NEXT:    ret <2 x double> [[RES]]
 ;
   %ptrs = getelementptr double, double *%base, <2 x i64> <i64 0, i64 1>
   %pt_v1 = insertelement <2 x double> undef, double %pt, i64 0
   %pt_v2 = insertelement <2 x double> %pt_v1, double %pt, i64 1
   %res = call <2 x double> @llvm.masked.gather.v2f64.v2p0f64(<2 x double*> %ptrs, i32 5, <2 x i1> <i1 true, i1 false>, <2 x double> %pt_v2)
   ret <2 x double> %res