[InstComb] Handle undef in simplifyMasked(Store|Scatter)

[artagnon]

No description provided.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Ramkumar Ramachandra (artagnon)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/161825.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+19-15)
• (modified) llvm/test/Transforms/InstCombine/masked_intrinsics.ll (+25)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index cf6d0ecab4f69..5b7fa8458c6f4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -292,25 +292,29 @@ Value *InstCombinerImpl::simplifyMaskedLoad(IntrinsicInst &II) {
   const Align Alignment =
       cast<ConstantInt>(II.getArgOperand(1))->getAlignValue();
 
+  // If we can unconditionally load from this address, replace with a
+  // load/select idiom.
+  if (isDereferenceablePointer(LoadPtr, II.getType(), II.getDataLayout(), &II,
+                               &AC, &DT)) {
+    LoadInst *LI = Builder.CreateAlignedLoad(II.getType(), LoadPtr, Alignment,
+                                             "unmaskedload");
+    LI->copyMetadata(II);
+    return Builder.CreateSelect(II.getArgOperand(2), LI, II.getArgOperand(3));
+  }
+
+  auto *ConstMask = dyn_cast<Constant>(II.getArgOperand(2));
+  if (!ConstMask)
+    return nullptr;
+
   // If the mask is all ones or undefs, this is a plain vector load of the 1st
   // argument.
-  if (maskIsAllOneOrUndef(II.getArgOperand(2))) {
+  if (maskIsAllOneOrUndef(ConstMask)) {
     LoadInst *L = Builder.CreateAlignedLoad(II.getType(), LoadPtr, Alignment,
                                             "unmaskedload");
     L->copyMetadata(II);
     return L;
   }
 
-  // If we can unconditionally load from this address, replace with a
-  // load/select idiom. TODO: use DT for context sensitive query
-  if (isDereferenceablePointer(LoadPtr, II.getType(),
-                               II.getDataLayout(), &II, &AC)) {
-    LoadInst *LI = Builder.CreateAlignedLoad(II.getType(), LoadPtr, Alignment,
-                                             "unmaskedload");
-    LI->copyMetadata(II);
-    return Builder.CreateSelect(II.getArgOperand(2), LI, II.getArgOperand(3));
-  }
-
   return nullptr;
 }
 
@@ -318,18 +322,18 @@ Value *InstCombinerImpl::simplifyMaskedLoad(IntrinsicInst &II) {
 // * Single constant active lane -> store
 // * Narrow width by halfs excluding zero/undef lanes
 Instruction *InstCombinerImpl::simplifyMaskedStore(IntrinsicInst &II) {
+  Value *StorePtr = II.getArgOperand(1);
+  Align Alignment = cast<ConstantInt>(II.getArgOperand(2))->getAlignValue();
   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())
+  if (maskIsAllZeroOrUndef(ConstMask))
     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);
-    Align Alignment = cast<ConstantInt>(II.getArgOperand(2))->getAlignValue();
+  if (maskIsAllOneOrUndef(ConstMask)) {
     StoreInst *S =
         new StoreInst(II.getArgOperand(0), StorePtr, false, Alignment);
     S->copyMetadata(II);
diff --git a/llvm/test/Transforms/InstCombine/masked_intrinsics.ll b/llvm/test/Transforms/InstCombine/masked_intrinsics.ll
index 8f7683419a82a..55be926230c45 100644
--- a/llvm/test/Transforms/InstCombine/masked_intrinsics.ll
+++ b/llvm/test/Transforms/InstCombine/masked_intrinsics.ll
@@ -16,6 +16,14 @@ define <2 x double> @load_zeromask(ptr %ptr, <2 x double> %passthru)  {
   ret <2 x double> %res
 }
 
+define <2 x double> @load_poisonmask(ptr %ptr, <2 x double> %passthru)  {
+; CHECK-LABEL: @load_poisonmask(
+; CHECK-NEXT:    ret <2 x double> [[PASSTHRU:%.*]]
+;
+  %res = call <2 x double> @llvm.masked.load.v2f64.p0(ptr %ptr, i32 1, <2 x i1> splat(i1 poison), <2 x double> %passthru)
+  ret <2 x double> %res
+}
+
 define <2 x double> @load_onemask(ptr %ptr, <2 x double> %passthru)  {
 ; CHECK-LABEL: @load_onemask(
 ; CHECK-NEXT:    [[UNMASKEDLOAD:%.*]] = load <2 x double>, ptr [[PTR:%.*]], align 2
@@ -150,6 +158,14 @@ define void @store_zeromask(ptr %ptr, <2 x double> %val)  {
   ret void
 }
 
+define void @store_poisonmask(ptr %ptr, <2 x double> %val)  {
+; CHECK-LABEL: @store_poisonmask(
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.masked.store.v2f64.p0(<2 x double> %val, ptr %ptr, i32 4, <2 x i1> splat(i1 poison))
+  ret void
+}
+
 define void @store_onemask(ptr %ptr, <2 x double> %val)  {
 ; CHECK-LABEL: @store_onemask(
 ; CHECK-NEXT:    store <2 x double> [[VAL:%.*]], ptr [[PTR:%.*]], align 4
@@ -159,6 +175,15 @@ define void @store_onemask(ptr %ptr, <2 x double> %val)  {
   ret void
 }
 
+define void @store_one_withpoison_mask(ptr %ptr, <2 x double> %val)  {
+; CHECK-LABEL: @store_one_withpoison_mask(
+; CHECK-NEXT:    store <2 x double> [[VAL:%.*]], ptr [[PTR:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.masked.store.v2f64.p0(<2 x double> %val, ptr %ptr, i32 4, <2 x i1> <i1 1, i1 poison>)
+  ret void
+}
+
 define void @store_demandedelts(ptr %ptr, double %val)  {
 ; CHECK-LABEL: @store_demandedelts(
 ; CHECK-NEXT:    [[VALVEC1:%.*]] = insertelement <2 x double> poison, double [[VAL:%.*]], i64 0

[nikic]

LGTM