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

Add support for other Constant types for the mask operand.

Author: paulwalker-arm

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);

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)