Simply operands of masked stores and scatters based on demanded elements

[reames]

Summary

If we know we're not storing a lane, we don't need to compute the lane. This could be improved by using the undef element result to further prune the mask, but I want to separate that into its own change since it's relatively likely to expose other problems.

Question for reviewer: Surely there's a better way to go from an <N x i1> mask to the demanded elements bits? One that works for more than just constant masks?

In the same spirit as https://reviews.llvm.org/D57177. Once these two are in, will extend to gather and masked.load. This is in the broader context of improving vector pointer instcombine under https://reviews.llvm.org/D57140.

[RKSimon]

Would it be better to put this inside InstCombiner::SimplifyDemandedVectorElts so we can also simplify the mask based on the demanded elts from its user(s)?

[reames]

In D57247#1371550, @RKSimon wrote:

Would it be better to put this inside InstCombiner::SimplifyDemandedVectorElts so we can also simplify the mask based on the demanded elts from its user(s)?

Is that idiomatic for a void result instruction? I thought SimplifyDemandedVectorElts expected to be given a value of vector type?

[RKSimon]

In D57247#1371663, @reames wrote:

In D57247#1371550, @RKSimon wrote:

Would it be better to put this inside InstCombiner::SimplifyDemandedVectorElts so we can also simplify the mask based on the demanded elts from its user(s)?

Is that idiomatic for a void result instruction? I thought SimplifyDemandedVectorElts expected to be given a value of vector type?

*facepalm* sorry, my brain was in Friday mode and I was thinking gathers/masked-loads.....

I agree some kind of helper function that can extract vector values from various Constant vector forms would be useful - they do keep turning up in code in various guises (and if you do gathers/maskedloads you'll be duplicating the code yet again!).

[reames]

In D57247#1372418, @RKSimon wrote:

I agree some kind of helper function that can extract vector values from various Constant vector forms would be useful - they do keep turning up in code in various guises (and if you do gathers/maskedloads you'll be duplicating the code yet again!).

Mind if I do this in a follow up commit? And can you point me to a few places that duplicate this? I'm just getting up to speed on how we handle masks and don't have the code indexed in my head yet. :)

[reames]

ping

[reames]

ping

[spatel]

I added the tests with baseline checks; please rebase after rL356283.
This seems ok to me, but let's double-check with @RKSimon about the refactoring options.

include/llvm/Analysis/VectorUtils.h

78

Why is this changing in this patch?

[reames]

Rebase, and factor out some common code.

[RKSimon]

LGTM with a couple of minor comments.

lib/Transforms/InstCombine/InstCombineCalls.cpp

1185

Not sure if its any better but you can get this from DemandedElts.getBitWidth():

APInt UndefElts(DemandedElts.getBitWidth(), 0);

1258

DemandedElts.getBitWidth()

include/llvm/Analysis/VectorUtils.h

 /// Given a vector and an element number, see if the scalar value is
 /// already around as a register, for example if it were inserted then extracted
 /// from the vector.
 Value *findScalarElement(Value *V, unsigned EltNo);
 
 /// Get splat value if the input is a splat vector or return nullptr.
 /// The value may be extracted from a splat constants vector or from
 /// a sequence of instructions that broadcast a single value into a vector.
-const Value *getSplatValue(const Value *V);
+Value *getSplatValue(const Value *V);
 
 /// Compute a map of integer instructions to their minimum legal type
 /// size.
 ///
 /// C semantics force sub-int-sized values (e.g. i8, i16) to be promoted to int
 /// type (e.g. i32) whenever arithmetic is performed on them.
 ///
 /// For targets with native i8 or i16 operations, usually InstCombine can shrink

lib/Analysis/VectorUtils.cpp

Value *llvm::findScalarElement(Value *V, unsigned EltNo) {
   return nullptr;
 }
 
 /// Get splat value if the input is a splat vector or return nullptr.
 /// This function is not fully general. It checks only 2 cases:
 /// the input value is (1) a splat constants vector or (2) a sequence
 /// of instructions that broadcast a single value into a vector.
 ///
-const llvm::Value *llvm::getSplatValue(const Value *V) {
+llvm::Value *llvm::getSplatValue(const Value *V) {
 
   if (auto *C = dyn_cast<Constant>(V))
     if (isa<VectorType>(V->getType()))
       return C->getSplatValue();
 
   auto *ShuffleInst = dyn_cast<ShuffleVectorInst>(V);
   if (!ShuffleInst)
     return nullptr;

lib/Transforms/InstCombine/InstCombineCalls.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
if (maskIsAllOneOrUndef(II.getArgOperand(2))) {
     Value *LoadPtr = II.getArgOperand(0);
     unsigned Alignment = cast<ConstantInt>(II.getArgOperand(1))->getZExtValue();
     return Builder.CreateAlignedLoad(LoadPtr, Alignment, "unmaskedload");
   }
 
   return nullptr;
 }
 
-static Instruction *simplifyMaskedStore(IntrinsicInst &II, InstCombiner &IC) {
+Instruction *InstCombiner::simplifyMaskedStore(IntrinsicInst &II) {
   auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(3));
   if (!ConstMask)
     return nullptr;
 
   // If the mask is all zeros, this instruction does nothing.
   if (ConstMask->isNullValue())
-    return IC.eraseInstFromFunction(II);
+    return eraseInstFromFunction(II);
 
   // If the mask is all ones, this is a plain vector store of the 1st argument.
   if (ConstMask->isAllOnesValue()) {
     Value *StorePtr = II.getArgOperand(1);
     unsigned Alignment = cast<ConstantInt>(II.getArgOperand(2))->getZExtValue();
     return new StoreInst(II.getArgOperand(0), StorePtr, false, Alignment);
   }
 
+  // Use masked off lanes to simplify operands via SimplifyDemandedVectorElts
+  // TODO: Surely we have something more general for understanding masks?
+  auto *CV = dyn_cast<ConstantVector>(ConstMask);
+  if (!CV)
+    return nullptr;
+  
+  const unsigned VWidth = CV->getType()->getNumElements();
+  APInt DemandedElts(VWidth, 0);
+  for (unsigned i = 0; i < VWidth; i++)
+    if (!CV->getAggregateElement(i)->isNullValue())
+      DemandedElts.setBit(i); 
+  APInt UndefElts(VWidth, 0);
+  if (Value *V = SimplifyDemandedVectorElts(II.getOperand(0),
+                                            DemandedElts, UndefElts)) {
+    II.setOperand(0, V);
+    return &II;
+  }
+
   return nullptr;
 }
 
 static Instruction *simplifyMaskedGather(IntrinsicInst &II, InstCombiner &IC) {
   // If the mask is all zeros, return the "passthru" argument of the gather.
   auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(2));
   if (ConstMask && ConstMask->isNullValue())
     return IC.replaceInstUsesWith(II, II.getArgOperand(3));
if (Result->getType()->getPointerAddressSpace() !=
       II.getType()->getPointerAddressSpace())
     Result = IC.Builder.CreateAddrSpaceCast(Result, II.getType());
   if (Result->getType() != II.getType())
     Result = IC.Builder.CreateBitCast(Result, II.getType());
 
   return cast<Instruction>(Result);
 }
 
-static Instruction *simplifyMaskedScatter(IntrinsicInst &II, InstCombiner &IC) {
+Instruction *InstCombiner::simplifyMaskedScatter(IntrinsicInst &II) {
   // If the mask is all zeros, a scatter does nothing.
   auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(3));
   if (ConstMask && ConstMask->isNullValue())
-    return IC.eraseInstFromFunction(II);
+    return eraseInstFromFunction(II);
+
+  auto *CV = dyn_cast_or_null<ConstantVector>(ConstMask);
+  if (!CV)
+    return nullptr;
+  
+  const unsigned VWidth = CV->getType()->getNumElements();
+  APInt DemandedElts(VWidth, 0);
+  for (unsigned i = 0; i < VWidth; i++)
+    if (!CV->getAggregateElement(i)->isNullValue())
+      DemandedElts.setBit(i); 
+  APInt UndefElts(VWidth, 0);
+  if (Value *V = SimplifyDemandedVectorElts(II.getOperand(0),
+                                            DemandedElts, UndefElts)) {
+    II.setOperand(0, V);
+    return &II;
+  }
+  if (Value *V = SimplifyDemandedVectorElts(II.getOperand(1),
+                                            DemandedElts, UndefElts)) {
+    II.setOperand(1, V);
+    return &II;
+  }
 
   return nullptr;
 }
 
 static Instruction *foldCttzCtlz(IntrinsicInst &II, InstCombiner &IC) {
   assert((II.getIntrinsicID() == Intrinsic::cttz ||
           II.getIntrinsicID() == Intrinsic::ctlz) &&
          "Expected cttz or ctlz intrinsic");
case Intrinsic::bswap: {
     }
     break;
   }
   case Intrinsic::masked_load:
     if (Value *SimplifiedMaskedOp = simplifyMaskedLoad(*II, Builder))
       return replaceInstUsesWith(CI, SimplifiedMaskedOp);
     break;
   case Intrinsic::masked_store:
-    return simplifyMaskedStore(*II, *this);
+    return simplifyMaskedStore(*II);
   case Intrinsic::masked_gather:
     return simplifyMaskedGather(*II, *this);
   case Intrinsic::masked_scatter:
-    return simplifyMaskedScatter(*II, *this);
+    return simplifyMaskedScatter(*II);
   case Intrinsic::launder_invariant_group:
   case Intrinsic::strip_invariant_group:
     if (auto *SkippedBarrier = simplifyInvariantGroupIntrinsic(*II, *this))
       return replaceInstUsesWith(*II, SkippedBarrier);
     break;
   case Intrinsic::powi:
     if (ConstantInt *Power = dyn_cast<ConstantInt>(II->getArgOperand(1))) {
       // 0 and 1 are handled in instsimplify

lib/Transforms/InstCombine/InstCombineInternal.h

bool OptimizeOverflowCheck(OverflowCheckFlavor OCF, Value *LHS, Value *RHS,
                              Constant *&OverflowResult);
 
   Instruction *visitCallSite(CallSite CS);
   Instruction *tryOptimizeCall(CallInst *CI);
   bool transformConstExprCastCall(CallSite CS);
   Instruction *transformCallThroughTrampoline(CallSite CS,
                                               IntrinsicInst *Tramp);
 
+  Instruction *simplifyMaskedStore(IntrinsicInst &II);
+  Instruction *simplifyMaskedScatter(IntrinsicInst &II);
+  
   /// Transform (zext icmp) to bitwise / integer operations in order to
   /// eliminate it.
   ///
   /// \param ICI The icmp of the (zext icmp) pair we are interested in.
   /// \parem CI The zext of the (zext icmp) pair we are interested in.
   /// \param DoTransform Pass false to just test whether the given (zext icmp)
   /// would be transformed. Pass true to actually perform the transformation.
   ///

test/Transforms/InstCombine/masked_intrinsics.ll

define void @store_onemask(<2 x double>* %ptr, <2 x double> %val)  {
   call void @llvm.masked.store.v2f64.p0v2f64(<2 x double> %val, <2 x double>* %ptr, i32 4, <2 x i1> <i1 1, i1 1>)
   ret void
 
 ; CHECK-LABEL: @store_onemask(
 ; CHECK-NEXT:  store <2 x double> %val, <2 x double>* %ptr, align 4
 ; CHECK-NEXT:  ret void
 }
 
+define void @store_demandedelts(<2 x double>* %ptr, double %val)  {
+  %valvec1 = insertelement <2 x double> undef, double %val, i32 0
+  %valvec2 = insertelement <2 x double> %valvec1, double %val, i32 1
+  call void @llvm.masked.store.v2f64.p0v2f64(<2 x double> %valvec2, <2 x double>* %ptr, i32 4, <2 x i1> <i1 true, i1 false>)
+  ret void
+
+; CHECK-LABEL: @store_demandedelts(
+; CHECK-NEXT:  %valvec2 = insertelement <2 x double> undef, double %val, i32 0
+; CHECK-NEXT:  call void @llvm.masked.store.v2f64.p0v2f64(<2 x double> %valvec2, <2 x double>* %ptr, i32 4, <2 x i1> <i1 true, i1 false>)
+; CHECK-NEXT:  ret void
+}
+
+
 define <2 x double> @gather_zeromask(<2 x double*> %ptrs, <2 x double> %passthru)  {
   %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
 
 ; CHECK-LABEL: @gather_zeromask(
 ; CHECK-NEXT:  ret <2 x double> %passthru
 }
 
 define void @scatter_zeromask(<2 x double*> %ptrs, <2 x double> %val)  {
-  call void @llvm.masked.scatter.v2f64.v2p0f64(<2 x double> %val, <2 x double*> %ptrs, i32 6, <2 x i1> zeroinitializer)
+  call void @llvm.masked.scatter.v2f64.v2p0f64(<2 x double> %val, <2 x double*> %ptrs, i32 8, <2 x i1> zeroinitializer)
   ret void
 
 ; CHECK-LABEL: @scatter_zeromask(
 ; CHECK-NEXT:  ret void
 }
 
+define void @scatter_demandedelts(double* %ptr, double %val)  {
+  %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
+
+; CHECK-LABEL: @scatter_demandedelts(
+; CHECK-NEXT:  %ptrs = getelementptr double, double* %ptr, <2 x i64> <i64 0, i64 1>
+; CHECK-NEXT:  %valvec2 = insertelement <2 x double> undef, double %val, i32 0
+; CHECK-NEXT:  call void @llvm.masked.scatter.v2f64.v2p0f64(<2 x double> %valvec2, <2 x double*> %ptrs, i32 8, <2 x i1> <i1 true, i1 false>)
+; CHECK-NEXT:  ret void
+}
+