[LLVM][InstCombine] Extend masked_gather's demanded elt analysis.

[paulwalker-arm]

Add support for other Constant types for the mask operand.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Paul Walker (paulwalker-arm)

Changes

Add support for other Constant types for the mask operand.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp (+15-7)
• (modified) llvm/test/Transforms/InstCombine/masked_intrinsics.ll (+1)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
index 0e3436d12702d..c82dae7ac6e65 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
@@ -1834,14 +1834,22 @@ Value *InstCombinerImpl::SimplifyDemandedVectorElts(Value *V,
       // segfaults which didn't exist in the original program.
       APInt DemandedPtrs(APInt::getAllOnes(VWidth)),
           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 (auto *CMask = dyn_cast<Constant>(II->getOperand(2))) {
+        if (CMask->isNullValue())
+          DemandedPtrs.clearAllBits();
+        else if (CMask->isAllOnesValue())
+          DemandedPassThrough.clearAllBits();
+        else if (auto *CV = dyn_cast<ConstantVector>(CMask)) {
+          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, PoisonElts2);
       simplifyAndSetOp(II, 3, DemandedPassThrough, PoisonElts3);
diff --git a/llvm/test/Transforms/InstCombine/masked_intrinsics.ll b/llvm/test/Transforms/InstCombine/masked_intrinsics.ll
index d9f022442a02e..8f7683419a82a 100644
--- a/llvm/test/Transforms/InstCombine/masked_intrinsics.ll
+++ b/llvm/test/Transforms/InstCombine/masked_intrinsics.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+; RUN: opt -passes=instcombine -use-constant-int-for-fixed-length-splat -S < %s | FileCheck %s
 
 declare <2 x double> @llvm.masked.load.v2f64.p0(ptr %ptrs, i32, <2 x i1> %mask, <2 x double> %src0)
 declare void @llvm.masked.store.v2f64.p0(<2 x double> %val, ptr %ptrs, i32, <2 x i1> %mask)

[nikic]

I think an alternative would be to replace dyn_cast<ConstantVector> with dyn_cast<Constant>? Less efficient, but simpler.

[paulwalker-arm]

Works for me. Thanks.

[nikic]

LGTM