[InstCombine] simplify masked load intrinsics with constant masks

[spatel]

Summary

A masked load with a zero mask means there's no load.
A masked load with an allOnes mask means it's a normal vector load.

I think something similar may be happening in CodeGenPrepare with D13855, but it doesn't trigger for a target that actually supports these ops (an x86 AVX target for example). We may be able to remove some of that logic. Doing these transforms in InstCombine is a better solution because it will trigger sooner and allow more optimizations from other passes.

Eventually, I think we should be able to replace the x86 intrinsics with the llvm IR intrinsics.

[RKSimon]

Some minor thoughts - I'd like to see what others have to say as well though.

lib/Transforms/InstCombine/InstCombineCalls.cpp

690

Minor: You can relax the type requirements by using Constant::isNullValue() and Constant::isAllOnesValue()

But this is true for a lot of the code in InstCombineCalls.cpp .....

708

I'm really not sure whether we should try to do this here or leaving it until lowering. Its making a big assumption that the target is good at scalar loads + insertions.

[spatel]

lib/Transforms/InstCombine/InstCombineCalls.cpp

690

Yes, that makes the code simpler.

I think we still want all of the asserts for sanity checking the intrinsic signature, but let me know if I'm missing anything.

708

I wasn't sure about that when I wrote it, so I threw it in to see what the consensus might be.

But I agree with you now. Although I would only expect this intrinsic to be formed from C a intrinsic or the vectorizers, and therefore, implying that the masked op is supported by the target, we shouldn't assume that also means that other vector ops are supported equally. It could also be wrong if someone's trying to minimize size (-Oz). Let's handle other constant values in the DAG.

[spatel]

Patch updated based on feedback from Simon:

1. Cast the mask to plain 'Constant' for simpler predicates.
2. Remove the 'TODO' and associated test cases for constant masks with a single set bit.

[delena]

lib/Transforms/InstCombine/InstCombineCalls.cpp

678

Verifier checks intrinsic signature. You don't need these checks at all.

843

Not sure that you can handle masked gather/scatter this way.

[spatel]

Patch updated: remove assertions because these are or should be handled by the IR verifier.

[spatel]

lib/Transforms/InstCombine/InstCombineCalls.cpp

677–712

Thanks. I think the verifier has a hole - it doesn't confirm that the alignment operand is constant. But I've removed the assertions here because it does confirm the other things.

843

Please correct me if I've misunderstood, but I think scatter and gather each have one degenerate folding opportunity that we can handle here:

1. A scatter with zero mask is a nop.
2. A gather with zero mask will return the passthru arg.

[delena]

LGTM

lib/Transforms/InstCombine/InstCombineCalls.cpp

843

yes, for zero mask it can be optimized.

lib/Transforms/InstCombine/InstCombineCalls.cpp

case 0x7:
       return ConstantInt::getSigned(VecTy, -1); // TRUE
     }
 
     if (Value *Cmp = Builder.CreateICmp(Pred, II.getArgOperand(0), II.getArgOperand(1)))
       return Builder.CreateSExtOrTrunc(Cmp, VecTy);
   }
   return nullptr;
 }
 
+static Value *simplifyMaskedLoad(const IntrinsicInst &II,
+                                 InstCombiner::BuilderTy &Builder) {
+  assert(II.getNumArgOperands() == 4 &&
+         "Wrong masked intrinsic format");
+  assert(isa<PointerType>(II.getArgOperand(0)->getType()) &&
+         "Wrong type for 1st arg");
+  assert(isa<ConstantInt>(II.getArgOperand(1)) &&
+         "Wrong type for 2nd arg");
+  assert(isa<VectorType>(II.getArgOperand(2)->getType()) &&
+         "Wrong type for 3rd arg");
+  assert(II.getArgOperand(3)->getType() == II.getType() &&
+         "Wrong type for 4th arg");
+
+  Value *LoadPtr = II.getArgOperand(0);
+  unsigned Alignment = cast<ConstantInt>(II.getArgOperand(1))->getZExtValue();
+  Value *Mask = II.getArgOperand(2);
+  Value *PassThru = II.getArgOperand(3);
+
+  // If the mask is all zeros, the "passthru" argument is the result.
+  if (isa<ConstantAggregateZero>(Mask))
+    return PassThru;
+
+  // If the mask is all ones, this is a plain vector load of the first argument.
+  // Note that the mask type is <N x i1> which is not a ConstantDataVector.
+  if (auto *ConstMask = dyn_cast<ConstantVector>(Mask))
+    if (ConstMask->isAllOnesValue())
+      return Builder.CreateAlignedLoad(LoadPtr, Alignment, "unmaskedload");
+
+  // TODO: A mask with only one set bit can be reduced to a scalar load and
+  // insertelement into the passthru vector.
+
+  return nullptr;
+}
+
 /// CallInst simplification. This mostly only handles folding of intrinsic
 /// instructions. For normal calls, it allows visitCallSite to do the heavy
 /// lifting.
 Instruction *InstCombiner::visitCallInst(CallInst &CI) {
   auto Args = CI.arg_operands();
   if (Value *V = SimplifyCall(CI.getCalledValue(), Args.begin(), Args.end(), DL,
                               TLI, DT, AC))
     return ReplaceInstUsesWith(CI, V);
 
case Intrinsic::bitreverse: {
     Value *X = nullptr;
 
     // bitreverse(bitreverse(x)) -> x
     if (match(IIOperand, m_Intrinsic<Intrinsic::bitreverse>(m_Value(X))))
       return ReplaceInstUsesWith(CI, X);
     break;
   }
 
+  case Intrinsic::masked_load:
+    if (Value *SimplifiedMaskedOp = simplifyMaskedLoad(*II, *Builder))
+      return ReplaceInstUsesWith(CI, SimplifiedMaskedOp);
+    break;
+
+  // TODO: Handle the other masked ops.
+  // case Intrinsic::masked_store:
+  // case Intrinsic::masked_gather:
+  // case Intrinsic::masked_scatter:
+
   case Intrinsic::powi:
     if (ConstantInt *Power = dyn_cast<ConstantInt>(II->getArgOperand(1))) {
       // powi(x, 0) -> 1.0
       if (Power->isZero())
         return ReplaceInstUsesWith(CI, ConstantFP::get(CI.getType(), 1.0));
       // powi(x, 1) -> x
       if (Power->isOne())
         return ReplaceInstUsesWith(CI, II->getArgOperand(0));

test/Transforms/InstCombine/masked_intrinsics.ll

 ; RUN: opt -instcombine -S < %s | FileCheck %s
 
 declare <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptrs, i32, <2 x i1> %mask, <2 x double> %src0)
 
-; FIXME: All of these could be simplified.
 
 define <2 x double> @load_zeromask(<2 x double>* %ptr, <2 x double> %passthru)  {
   %res = call <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptr, i32 1, <2 x i1> zeroinitializer, <2 x double> %passthru)
   ret <2 x double> %res
 
 ; CHECK-LABEL: @load_zeromask(
-; CHECK-NEXT:  %res = call <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptr, i32 1, <2 x i1> zeroinitializer, <2 x double> %passthru)
-; CHECK-NEXT   ret <2 x double> %res
+; CHECK-NEXT   ret <2 x double> %passthru
 }
 
 define <2 x double> @load_onemask(<2 x double>* %ptr, <2 x double> %passthru)  {
   %res = call <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptr, i32 2, <2 x i1> <i1 1, i1 1>, <2 x double> %passthru)
   ret <2 x double> %res
 
 ; CHECK-LABEL: @load_onemask(
-; CHECK-NEXT:  %res = call <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptr, i32 2, <2 x i1> <i1 true, i1 true>, <2 x double> %passthru)
+; CHECK-NEXT:  %unmaskedload = load <2 x double>, <2 x double>* %ptr, align 2
 ; CHECK-NEXT   ret <2 x double> %res
 }
 
+; FIXME: These could be simplified.
+
 define <2 x double> @load_onesetbitmask1(<2 x double>* %ptr, <2 x double> %passthru)  {
   %res = call <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptr, i32 3, <2 x i1> <i1 0, i1 1>, <2 x double> %passthru)
   ret <2 x double> %res
 
 ; CHECK-LABEL: @load_onesetbitmask1(
 ; CHECK-NEXT:  %res = call <2 x double> @llvm.masked.load.v2f64(<2 x double>* %ptr, i32 3, <2 x i1> <i1 false, i1 true>, <2 x double> %passthru)
 ; CHECK-NEXT   ret <2 x double> %res
 }