[InstSimplify] Consider vscale_range for get active lane mask (#160073)

Scalable get_active_lane_mask intrinsic calls can be simplified to i1
splat (ptrue) when its constant range is larger than or equal to the
maximum possible number of elements, which can be inferred from
vscale_range(x, y)

Author: MDevereau

llvm/lib/Analysis/InstructionSimplify.cpp

@@ -6514,10 +6514,27 @@ Value *llvm::simplifyBinaryIntrinsic(Intrinsic::ID IID, Type *ReturnType,
                                      const CallBase *Call) {
   unsigned BitWidth = ReturnType->getScalarSizeInBits();
   switch (IID) {
-  case Intrinsic::get_active_lane_mask:
+  case Intrinsic::get_active_lane_mask: {
     if (match(Op1, m_Zero()))
       return ConstantInt::getFalse(ReturnType);
+
+    const Function *F = Call->getFunction();
+    auto *ScalableTy = dyn_cast<ScalableVectorType>(ReturnType);
+    Attribute Attr = F->getFnAttribute(Attribute::VScaleRange);
+    if (ScalableTy && Attr.isValid()) {
+      std::optional<unsigned> VScaleMax = Attr.getVScaleRangeMax();
+      if (!VScaleMax)
+        break;
+      uint64_t MaxPossibleMaskElements =
+          (uint64_t)ScalableTy->getMinNumElements() * (*VScaleMax);
+
+      const APInt *Op1Val;
+      if (match(Op0, m_Zero()) && match(Op1, m_APInt(Op1Val)) &&
+          Op1Val->uge(MaxPossibleMaskElements))
+        return ConstantInt::getAllOnesValue(ReturnType);
+    }
     break;
+  }
   case Intrinsic::abs:
     // abs(abs(x)) -> abs(x). We don't need to worry about the nsw arg here.
     // It is always ok to pick the earlier abs. We'll just lose nsw if its only

llvm/test/Transforms/InstSimplify/get_active_lane_mask.ll

@@ -18,3 +18,51 @@ define <vscale x 8 x i1> @foo_nxv8i1(i32 %a) {
   %mask = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1(i32 %a, i32 0)
   ret <vscale x 8 x i1> %mask
 }
+
+define <vscale x 16 x i1> @foo_vscale_max_255() vscale_range(1,16) {
+; CHECK-LABEL: define <vscale x 16 x i1> @foo_vscale_max_255(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[MASK:%.*]] = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i32(i32 0, i32 255)
+; CHECK-NEXT:    ret <vscale x 16 x i1> [[MASK]]
+;
+  %mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1(i32 0, i32 255)
+  ret <vscale x 16 x i1> %mask
+}
+
+define <vscale x 16 x i1> @foo_vscale_max_256() vscale_range(1,16) {
+; CHECK-LABEL: define <vscale x 16 x i1> @foo_vscale_max_256(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:    ret <vscale x 16 x i1> splat (i1 true)
+;
+  %mask = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1(i32 0, i32 256)
+  ret <vscale x 16 x i1> %mask
+}
+
+define <vscale x 2 x i1> @foo_vscale_max_nxv2i1_1_1_2() vscale_range(1,1) {
+; CHECK-LABEL: define <vscale x 2 x i1> @foo_vscale_max_nxv2i1_1_1_2(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:    ret <vscale x 2 x i1> splat (i1 true)
+;
+  %mask = call <vscale x 2 x i1> @llvm.get.active.lane.mask.nxv2i1(i32 0, i32 2)
+  ret <vscale x 2 x i1> %mask
+}
+
+define <vscale x 4 x i1> @foo_vscale_max_nxv4i1_2_4_16() vscale_range(2,4) {
+; CHECK-LABEL: define <vscale x 4 x i1> @foo_vscale_max_nxv4i1_2_4_16(
+; CHECK-SAME: ) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT:    ret <vscale x 4 x i1> splat (i1 true)
+;
+  %mask = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1(i128 0, i128 16)
+  ret <vscale x 4 x i1> %mask
+}
+
+define <vscale x 4 x i1> @foo_vscale_max_nxv4i1_2_4_1_16() vscale_range(2,4) {
+; CHECK-LABEL: define <vscale x 4 x i1> @foo_vscale_max_nxv4i1_2_4_1_16(
+; CHECK-SAME: ) #[[ATTR2]] {
+; CHECK-NEXT:    [[MASK:%.*]] = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1.i128(i128 1, i128 16)
+; CHECK-NEXT:    ret <vscale x 4 x i1> [[MASK]]
+;
+  %mask = call <vscale x 4 x i1> @llvm.get.active.lane.mask.nxv4i1(i128 1, i128 16)
+  ret <vscale x 4 x i1> %mask
+}
+