[EarlyCSE] Handle masked loads and stores

[kparzysz]

Summary

Extend the handling of memory intrinsics to also include non-target-specific intrinsics, in particular masked loads and stores.

Invent isHandledNonTargetIntrinsic to distinguish between intrinsics that should be handled natively from intrinsics that can be passed to TTI.

Add code that handles masked loads and stores and update the testcase to reflect the results.

[kparzysz]

Ping?

This should be NFC for everything that doesn't involve masked loads/stores.

[jdoerfert]

In D87340#2271039, @kparzysz wrote:

This should be NFC for everything that doesn't involve masked loads/stores.

Can us split that part out then? Makes reviews easier.

[kparzysz]

Created D87691 with the refactoring. Will rebase this patch soon.

[kparzysz]

Rebased on top of D87691.

[efriedma]

Overall this seems like a natural extension of the existing EarlyCSE handling of load/store. I don't see any major issues.

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

701

Info.MatchingID is getting implicitly set to zero? Could that conflict with a target-specific ID? I guess you explicitly check the identity of the intrinsic later, but it seems a little confusing that you aren't doing anything explicit with getMatchingId().

873

Is there any reason to expect !isa<UndefValue>(Mask0) to hold here? I mean, it would be weird to see a load with an undef mask, but I don't see any code that would enforce it, and the assertion itself doesn't explain.

876

The Vec0->getType() != Vec1->getType() shouldn't be necessary: if the values are the same type, that implies the masks are the same type.

884

There should be some test coverage for the unequal-masks case.

If the unequal-masks cases are important, you might run into cases where a masked operation with an all-true mask got folded to a plain load/store. Not sure how likely that is.

[kparzysz]

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

701

Setting it to masked_load (with an explanation in a comment).

873

You're right, I changed it to an explicit check.

876

There is a check if the pointer operands are equal, so you're right, the types here are implied to be the same. I added it as a safeguard for when that check (which is at another place) is relaxed to checking that the pointers point to the same location, but are not necessarily equal as Value's. I changed the type check here to an assertion, so this will still fail when this condition is accidentally violated.

[kparzysz]

Add testcases with unequal masks.

Fix detecting dead masked stores.

[lebedev.ri]

llvm/test/Transforms/EarlyCSE/masked-intrinsics-unequal-masks.ll

1

Precommit new tests please

[kparzysz]

Pre-committed new test file.

[efriedma]

LGTM

[efriedma]

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

883

Just spotted an issue here.

This cast<ConstantInt> could fail. The operand of a ConstantVector has to be a Constant, but it doesn't have to be a ConstantInt. It could be an UndefValue, or a ConstantExpr.

[kparzysz]

Fixed submask detection.

[efriedma]

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

898

!Int0->isZero() && Int1->isZero() is always true.

[kparzysz]

Removed unnecessary checks in isSubmask.

[efriedma]

LGTM

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

private:
 
   /// Wrapper class to handle memory instructions, including loads,
   /// stores and intrinsic loads and stores defined by the target.
   class ParseMemoryInst {
   public:
     ParseMemoryInst(Instruction *Inst, const TargetTransformInfo &TTI)
       : Inst(Inst) {
       if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(Inst)) {
-        if (TTI.getTgtMemIntrinsic(II, Info))
-          IntrID = II->getIntrinsicID();
+        if (TTI.getTgtMemIntrinsic(II, Info))
+          return;
+        if (isHandledNonTargetIntrinsic(IntrID)) {
+          switch (IntrID) {
+          case Intrinsic::masked_load:
+            Info.PtrVal = Inst->getOperand(0);
+            Info.MatchingId = Intrinsic::masked_load;
+            Info.ReadMem = true;
+            Info.WriteMem = false;
+            Info.IsVolatile = false;
+            break;
+          case Intrinsic::masked_store:
+            Info.PtrVal = Inst->getOperand(1);
+            // Use the ID of masked load as the "matching id". This will
+            // prevent matching non-masked loads/stores with masked ones
+            // (which could be done), but at the moment, the code here
+            // does not support matching intrinsics with non-intrinsics,
+            // so keep the MatchingIds specific to masked instructions
+            // for now (TODO).
+            Info.MatchingId = Intrinsic::masked_load;
+            Info.ReadMem = false;
+            Info.WriteMem = true;
+            Info.IsVolatile = false;
+            break;
+          }
+        }
       }
     }
 
     Instruction *get() { return Inst; }
     const Instruction *get() const { return Inst; }
 
     bool isLoad() const {
       if (IntrID != 0)
public:
     }
 
     bool isInvariantLoad() const {
       if (auto *LI = dyn_cast<LoadInst>(Inst))
         return LI->hasMetadata(LLVMContext::MD_invariant_load);
       return false;
     }
 
-    bool isMatchingMemLoc(const ParseMemoryInst &Inst) const {
-      return (getPointerOperand() == Inst.getPointerOperand() &&
-              getMatchingId() == Inst.getMatchingId());
-    }
-
     bool isValid() const { return getPointerOperand() != nullptr; }
 
     // For regular (non-intrinsic) loads/stores, this is set to -1. For
     // intrinsic loads/stores, the id is retrieved from the corresponding
     // field in the MemIntrinsicInfo structure.  That field contains
     // non-negative values only.
     int getMatchingId() const {
       if (IntrID != 0)
public:
     }
 
   private:
     Intrinsic::ID IntrID = 0;
     MemIntrinsicInfo Info;
     Instruction *Inst;
   };
 
+  // This function is to prevent accidentally passing a non-target
+  // intrinsic ID to TargetTransformInfo.
+  static bool isHandledNonTargetIntrinsic(Intrinsic::ID ID) {
+    switch (ID) {
+    case Intrinsic::masked_load:
+    case Intrinsic::masked_store:
+      return true;
+    }
+    return false;
+  }
+  static bool isHandledNonTargetIntrinsic(const Value *V) {
+    if (auto *II = dyn_cast<IntrinsicInst>(V))
+      return isHandledNonTargetIntrinsic(II->getIntrinsicID());
+    return false;
+  }
+
   bool processNode(DomTreeNode *Node);
 
   bool handleBranchCondition(Instruction *CondInst, const BranchInst *BI,
                              const BasicBlock *BB, const BasicBlock *Pred);
 
   Value *getMatchingValue(LoadValue &InVal, ParseMemoryInst &MemInst,
                           unsigned CurrentGeneration);
 
+  bool overridingStores(const ParseMemoryInst &Earlier,
+                        const ParseMemoryInst &Later);
+
   Value *getOrCreateResult(Value *Inst, Type *ExpectedType) const {
     if (auto *LI = dyn_cast<LoadInst>(Inst))
       return LI;
     if (auto *SI = dyn_cast<StoreInst>(Inst))
       return SI->getValueOperand();
     assert(isa<IntrinsicInst>(Inst) && "Instruction not supported");
-    return TTI.getOrCreateResultFromMemIntrinsic(cast<IntrinsicInst>(Inst),
-                                                 ExpectedType);
+    auto *II = cast<IntrinsicInst>(Inst);
+    if (isHandledNonTargetIntrinsic(II->getIntrinsicID()))
+      return getOrCreateResultNonTargetMemIntrinsic(II, ExpectedType);
+    return TTI.getOrCreateResultFromMemIntrinsic(II, ExpectedType);
+  }
+
+  Value *getOrCreateResultNonTargetMemIntrinsic(IntrinsicInst *II,
+                                                Type *ExpectedType) const {
+    switch (II->getIntrinsicID()) {
+    case Intrinsic::masked_load:
+      return II;
+    case Intrinsic::masked_store:
+      return II->getOperand(0);
+    }
+    return nullptr;
   }
 
   /// Return true if the instruction is known to only operate on memory
   /// provably invariant in the given "generation".
   bool isOperatingOnInvariantMemAt(Instruction *I, unsigned GenAt);
 
   bool isSameMemGeneration(unsigned EarlierGeneration, unsigned LaterGeneration,
                            Instruction *EarlierInst, Instruction *LaterInst);
 
+  bool isNonTargetIntrinsicMatch(const IntrinsicInst *Earlier,
+                                 const IntrinsicInst *Later) {
+    auto IsSubmask = [](const Value *Mask0, const Value *Mask1) {
+      // Is Mask0 a submask of Mask1?
+      if (Mask0 == Mask1)
+        return true;
+      if (isa<UndefValue>(Mask0) || isa<UndefValue>(Mask1))
+        return false;
+      auto *Vec0 = dyn_cast<ConstantVector>(Mask0);
+      auto *Vec1 = dyn_cast<ConstantVector>(Mask1);
+      if (!Vec0 || !Vec1)
+        return false;
+      assert(Vec0->getType() == Vec1->getType() &&
+             "Masks should have the same type");
+      for (int i = 0, e = Vec0->getNumOperands(); i != e; ++i) {
+        Constant *Elem0 = Vec0->getOperand(i);
+        Constant *Elem1 = Vec1->getOperand(i);
+        auto *Int0 = dyn_cast<ConstantInt>(Elem0);
+        if (Int0 && Int0->isZero())
+          continue;
+        auto *Int1 = dyn_cast<ConstantInt>(Elem1);
+        if (Int1 && !Int1->isZero())
+          continue;
+        if (isa<UndefValue>(Elem0) || isa<UndefValue>(Elem1))
+          return false;
+        if (Elem0 == Elem1)
+          continue;
+        return false;
+      }
+      return true;
+    };
+    auto PtrOp = [](const IntrinsicInst *II) {
+      if (II->getIntrinsicID() == Intrinsic::masked_load)
+        return II->getOperand(0);
+      if (II->getIntrinsicID() == Intrinsic::masked_store)
+        return II->getOperand(1);
+      llvm_unreachable("Unexpected IntrinsicInst");
+    };
+    auto MaskOp = [](const IntrinsicInst *II) {
+      if (II->getIntrinsicID() == Intrinsic::masked_load)
+        return II->getOperand(2);
+      if (II->getIntrinsicID() == Intrinsic::masked_store)
+        return II->getOperand(3);
+      llvm_unreachable("Unexpected IntrinsicInst");
+    };
+    auto ThruOp = [](const IntrinsicInst *II) {
+      if (II->getIntrinsicID() == Intrinsic::masked_load)
+        return II->getOperand(3);
+      llvm_unreachable("Unexpected IntrinsicInst");
+    };
+
+    if (PtrOp(Earlier) != PtrOp(Later))
+      return false;
+
+    Intrinsic::ID IDE = Earlier->getIntrinsicID();
+    Intrinsic::ID IDL = Later->getIntrinsicID();
+    // We could really use specific intrinsic classes for masked loads
+    // and stores in IntrinsicInst.h.
+    if (IDE == Intrinsic::masked_load && IDL == Intrinsic::masked_load) {
+      // Trying to replace later masked load with the earlier one.
+      // Check that the pointers are the same, and
+      // - masks and pass-throughs are the same, or
+      // - replacee's pass-through is "undef" and replacer's mask is a
+      //   super-set of the replacee's mask.
+      if (MaskOp(Earlier) == MaskOp(Later) && ThruOp(Earlier) == ThruOp(Later))
+        return true;
+      if (!isa<UndefValue>(ThruOp(Later)))
+        return false;
+      return IsSubmask(MaskOp(Later), MaskOp(Earlier));
+    }
+    if (IDE == Intrinsic::masked_store && IDL == Intrinsic::masked_load) {
+      // Trying to replace a load of a stored value with the store's value.
+      // Check that the pointers are the same, and
+      // - load's mask is a subset of store's mask, and
+      // - load's pass-through is "undef".
+      if (!IsSubmask(MaskOp(Later), MaskOp(Earlier)))
+        return false;
+      return isa<UndefValue>(ThruOp(Later));
+    }
+    if (IDE == Intrinsic::masked_load && IDL == Intrinsic::masked_store) {
+      // Trying to remove a store of the loaded value.
+      // Check that the pointers are the same, and
+      // - store's mask is a subset of the load's mask.
+      return IsSubmask(MaskOp(Later), MaskOp(Earlier));
+    }
+    if (IDE == Intrinsic::masked_store && IDL == Intrinsic::masked_store) {
+      // Trying to remove a dead store (earlier).
+      // Check that the pointers are the same,
+      // - the to-be-removed store's mask is a subset of the other store's
+      //   mask.
+      return IsSubmask(MaskOp(Earlier), MaskOp(Later));
+    }
+    return false;
+  }
+
   void removeMSSA(Instruction &Inst) {
     if (!MSSA)
       return;
     if (VerifyMemorySSA)
       MSSA->verifyMemorySSA();
     // Removing a store here can leave MemorySSA in an unoptimized state by
     // creating MemoryPhis that have identical arguments and by creating
     // MemoryUses whose defining access is not an actual clobber. The phi case
Value *EarlyCSE::getMatchingValue(LoadValue &InVal, ParseMemoryInst &MemInst,
   // The value V returned from this function is used differently depending
   // on whether MemInst is a load or a store. If it's a load, we will replace
   // MemInst with V, if it's a store, we will check if V is the same as the
   // available value.
   bool MemInstMatching = !MemInst.isLoad();
   Instruction *Matching = MemInstMatching ? MemInst.get() : InVal.DefInst;
   Instruction *Other = MemInstMatching ? InVal.DefInst : MemInst.get();
 
+  // Deal with non-target memory intrinsics.
+  bool MatchingNTI = isHandledNonTargetIntrinsic(Matching);
+  bool OtherNTI = isHandledNonTargetIntrinsic(Other);
+  if (OtherNTI != MatchingNTI)
+    return nullptr;
+  if (OtherNTI && MatchingNTI) {
+    if (!isNonTargetIntrinsicMatch(cast<IntrinsicInst>(InVal.DefInst),
+                                   cast<IntrinsicInst>(MemInst.get())))
+      return nullptr;
+  }
+
   if (!isOperatingOnInvariantMemAt(MemInst.get(), InVal.Generation) &&
       !isSameMemGeneration(InVal.Generation, CurrentGeneration, InVal.DefInst,
                            MemInst.get()))
     return nullptr;
   return getOrCreateResult(Matching, Other->getType());
 }
 
+bool EarlyCSE::overridingStores(const ParseMemoryInst &Earlier,
+                                const ParseMemoryInst &Later) {
+  // Can we remove Earlier store because of Later store?
+
+  assert(Earlier.isUnordered() && !Earlier.isVolatile() &&
+         "Violated invariant");
+  if (Earlier.getPointerOperand() != Later.getPointerOperand())
+    return false;
+  if (Earlier.getMatchingId() != Later.getMatchingId())
+    return false;
+  // At the moment, we don't remove ordered stores, but do remove
+  // unordered atomic stores.  There's no special requirement (for
+  // unordered atomics) about removing atomic stores only in favor of
+  // other atomic stores since we were going to execute the non-atomic
+  // one anyway and the atomic one might never have become visible.
+  if (!Earlier.isUnordered() || !Later.isUnordered())
+    return false;
+
+  // Deal with non-target memory intrinsics.
+  bool ENTI = isHandledNonTargetIntrinsic(Earlier.get());
+  bool LNTI = isHandledNonTargetIntrinsic(Later.get());
+  if (ENTI && LNTI)
+    return isNonTargetIntrinsicMatch(cast<IntrinsicInst>(Earlier.get()),
+                                     cast<IntrinsicInst>(Later.get()));
+
+  // Because of the check above, at least one of them is false.
+  // For now disallow matching intrinsics with non-intrinsics,
+  // so assume that the stores match if neither is an intrinsic.
+  return ENTI == LNTI;
+}
+
 bool EarlyCSE::processNode(DomTreeNode *Node) {
   bool Changed = false;
   BasicBlock *BB = Node->getBlock();
 
   // If this block has a single predecessor, then the predecessor is the parent
   // of the domtree node and all of the live out memory values are still current
   // in this block.  If this block has multiple predecessors, then they could
   // have invalidated the live-out memory values of our parent value.  For now,
for (Instruction &Inst : make_early_inc_range(BB->getInstList())) {
     // something that could modify memory.  If so, our available memory values
     // cannot be used so bump the generation count.
     if (Inst.mayWriteToMemory()) {
       ++CurrentGeneration;
 
       if (MemInst.isValid() && MemInst.isStore()) {
         // We do a trivial form of DSE if there are two stores to the same
         // location with no intervening loads.  Delete the earlier store.
-        // At the moment, we don't remove ordered stores, but do remove
-        // unordered atomic stores.  There's no special requirement (for
-        // unordered atomics) about removing atomic stores only in favor of
-        // other atomic stores since we were going to execute the non-atomic
-        // one anyway and the atomic one might never have become visible.
         if (LastStore) {
-          ParseMemoryInst LastStoreMemInst(LastStore, TTI);
-          assert(LastStoreMemInst.isUnordered() &&
-                 !LastStoreMemInst.isVolatile() &&
-                 "Violated invariant");
-          if (LastStoreMemInst.isMatchingMemLoc(MemInst)) {
+          if (overridingStores(ParseMemoryInst(LastStore, TTI), MemInst)) {
             LLVM_DEBUG(dbgs() << "EarlyCSE DEAD STORE: " << *LastStore
                               << "  due to: " << Inst << '\n');
             if (!DebugCounter::shouldExecute(CSECounter)) {
               LLVM_DEBUG(dbgs() << "Skipping due to debug counter\n");
             } else {
               salvageKnowledge(&Inst, &AC);
               removeMSSA(*LastStore);
               LastStore->eraseFromParent();

llvm/test/Transforms/EarlyCSE/masked-intrinsics-unequal-masks.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -early-cse < %s | FileCheck %s
 
 ; Unequal mask check.
 
 ; Load-load: the second load can be removed if (assuming unequal masks) the
 ; second loaded value is a subset of the first loaded value considering the
 ; non-undef vector elements. In other words, if the second mask is a submask
 ; of the first one, and the through value of the second load is undef.
 
 ; Load-load, second mask is a submask of the first, second through is undef.
 ; Expect the second load to be removed.
 define <4 x i32> @f3(<4 x i32>* %a0, <4 x i32> %a1) {
 ; CHECK-LABEL: @f3(
 ; CHECK-NEXT:    [[V0:%.*]] = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* [[A0:%.*]], i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>, <4 x i32> [[A1:%.*]])
-; CHECK-NEXT:    [[V1:%.*]] = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* [[A0]], i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> undef)
-; CHECK-NEXT:    [[V2:%.*]] = add <4 x i32> [[V0]], [[V1]]
+; CHECK-NEXT:    [[V2:%.*]] = add <4 x i32> [[V0]], [[V0]]
 ; CHECK-NEXT:    ret <4 x i32> [[V2]]
 ;
   %v0 = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %a0, i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>, <4 x i32> %a1)
   %v1 = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %a0, i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> undef)
   %v2 = add <4 x i32> %v0, %v1
   ret <4 x i32> %v2
 }
 

 
 ; Store-store: the first store can be removed if the first; mask is a submask
 ; of the second mask.
 
 ; Store-store, first mask is a submask of the second.
 ; Expect the first store to be removed.
 define void @f6(<4 x i32> %a0, <4 x i32>* %a1) {
 ; CHECK-LABEL: @f6(
-; CHECK-NEXT:    call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> [[A0:%.*]], <4 x i32>* [[A1:%.*]], i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>)
-; CHECK-NEXT:    call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> [[A0]], <4 x i32>* [[A1]], i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>)
+; CHECK-NEXT:    call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> [[A0:%.*]], <4 x i32>* [[A1:%.*]], i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>)
 ; CHECK-NEXT:    ret void
 ;
   call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> %a0, <4 x i32>* %a1, i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>)
   call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> %a0, <4 x i32>* %a1, i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>)
   ret void
 }
 
 ; Store-store, first mask is not a submask of the second.

 ; Load-store: the store can be removed if the store's mask is a submask of the
 ; load's mask.
 
 ; Load-store, second mask is a submask of the first.
 ; Expect the store to be removed.
 define <4 x i32> @f8(<4 x i32>* %a0, <4 x i32> %a1) {
 ; CHECK-LABEL: @f8(
 ; CHECK-NEXT:    [[V0:%.*]] = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* [[A0:%.*]], i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>, <4 x i32> [[A1:%.*]])
-; CHECK-NEXT:    call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> [[V0]], <4 x i32>* [[A0]], i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>)
 ; CHECK-NEXT:    ret <4 x i32> [[V0]]
 ;
   %v0 = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %a0, i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>, <4 x i32> %a1)
   call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> %v0, <4 x i32>* %a0, i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>)
   ret <4 x i32> %v0
 }
 
 ; Load-store, second mask is not a submask of the first.

 ; Store-load: the load can be removed if load's mask is a submask of the
 ; store's mask, and the load's through value is undef.
 
 ; Store-load, load's mask is a submask of store's mask, thru is undef.
 ; Expect the load to be removed.
 define <4 x i32> @fa(<4 x i32> %a0, <4 x i32>* %a1) {
 ; CHECK-LABEL: @fa(
 ; CHECK-NEXT:    call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> [[A0:%.*]], <4 x i32>* [[A1:%.*]], i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>)
-; CHECK-NEXT:    [[V0:%.*]] = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* [[A1]], i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> undef)
-; CHECK-NEXT:    ret <4 x i32> [[V0]]
+; CHECK-NEXT:    ret <4 x i32> [[A0]]
 ;
   call void @llvm.masked.store.v4i32.p0v4i32(<4 x i32> %a0, <4 x i32>* %a1, i32 4, <4 x i1> <i1 true, i1 true, i1 false, i1 true>)
   %v0 = call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* %a1, i32 4, <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> undef)
   ret <4 x i32> %v0
 }
 
 ; Store-load, load's mask is a submask of store's mask, thru is not undef.
 ; Expect the load to remain.

llvm/test/Transforms/EarlyCSE/masked-intrinsics.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -early-cse < %s | FileCheck %s
 
 define <128 x i8> @f0(<128 x i8>* %a0, <128 x i8> %a1, <128 x i8> %a2) {
 ; CHECK-LABEL: @f0(
 ; CHECK-NEXT:    [[V0:%.*]] = icmp eq <128 x i8> [[A1:%.*]], [[A2:%.*]]
 ; CHECK-NEXT:    call void @llvm.masked.store.v128i8.p0v128i8(<128 x i8> [[A1]], <128 x i8>* [[A0:%.*]], i32 4, <128 x i1> [[V0]])
-; CHECK-NEXT:    [[V1:%.*]] = call <128 x i8> @llvm.masked.load.v128i8.p0v128i8(<128 x i8>* [[A0]], i32 4, <128 x i1> [[V0]], <128 x i8> undef)
-; CHECK-NEXT:    ret <128 x i8> [[V1]]
+; CHECK-NEXT:    ret <128 x i8> [[A1]]
 ;
   %v0 = icmp eq <128 x i8> %a1, %a2
   call void @llvm.masked.store.v128i8.p0v128i8(<128 x i8> %a1, <128 x i8>* %a0, i32 4, <128 x i1> %v0)
   %v1 = call <128 x i8> @llvm.masked.load.v128i8.p0v128i8(<128 x i8>* %a0, i32 4, <128 x i1> %v0, <128 x i8> undef)
   ret <128 x i8> %v1
 }
 
 define <128 x i8> @f1(<128 x i8>* %a0, <128 x i8> %a1, <128 x i8> %a2) {
 ; CHECK-LABEL: @f1(
 ; CHECK-NEXT:    [[V0:%.*]] = icmp eq <128 x i8> [[A1:%.*]], [[A2:%.*]]
 ; CHECK-NEXT:    [[V1:%.*]] = call <128 x i8> @llvm.masked.load.v128i8.p0v128i8(<128 x i8>* [[A0:%.*]], i32 4, <128 x i1> [[V0]], <128 x i8> undef)
-; CHECK-NEXT:    call void @llvm.masked.store.v128i8.p0v128i8(<128 x i8> [[V1]], <128 x i8>* [[A0]], i32 4, <128 x i1> [[V0]])
 ; CHECK-NEXT:    ret <128 x i8> [[V1]]
 ;
   %v0 = icmp eq <128 x i8> %a1, %a2
   %v1 = call <128 x i8> @llvm.masked.load.v128i8.p0v128i8(<128 x i8>* %a0, i32 4, <128 x i1> %v0, <128 x i8> undef)
   call void @llvm.masked.store.v128i8.p0v128i8(<128 x i8> %v1, <128 x i8>* %a0, i32 4, <128 x i1> %v0)
   ret <128 x i8> %v1
 }