[llvm] [SelectionDAG] Add GET_ACTIVE_LANE_MASK to ComputeNumSignBits (PR #208977)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Usman Nadeem (UsmanNadeem)

<details>
<summary>Changes</summary>

`GET_ACTIVE_LANE_MASK ` acts like `icmp ult` so code is similar to `SETCC`.

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


3 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp (+6) 
- (modified) llvm/test/CodeGen/AArch64/fold-sext-in-reg-predicate-fixed-length.ll (+31) 
- (modified) llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp (+25) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 626803ed92a40..bf720addf5b8a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -5225,6 +5225,12 @@ unsigned SelectionDAG::ComputeNumSignBits(SDValue Op, const APInt &DemandedElts,
       return VTBits;
     break;
   }
+  case ISD::GET_ACTIVE_LANE_MASK:
+    // Semantically similar to icmp ult.
+    if (TLI->getBooleanContents(VT.isVector(), /*isFloat=*/false) ==
+        TargetLowering::ZeroOrNegativeOneBooleanContent)
+      return VTBits;
+    break;
   case ISD::ROTL:
   case ISD::ROTR:
     Tmp = ComputeNumSignBits(Op.getOperand(0), DemandedElts, Depth + 1);
diff --git a/llvm/test/CodeGen/AArch64/fold-sext-in-reg-predicate-fixed-length.ll b/llvm/test/CodeGen/AArch64/fold-sext-in-reg-predicate-fixed-length.ll
index 783cd6f693ed8..747c1619187f4 100644
--- a/llvm/test/CodeGen/AArch64/fold-sext-in-reg-predicate-fixed-length.ll
+++ b/llvm/test/CodeGen/AArch64/fold-sext-in-reg-predicate-fixed-length.ll
@@ -17,3 +17,34 @@ entry:
   %data = call <16 x i8> @llvm.masked.load.v16i8.p0(ptr %p, i32 1, <16 x i1> %mask, <16 x i8> zeroinitializer)
   ret <16 x i8> %data
 }
+
+define void @active_lane_mask_mstore_vscaleX2(ptr %p, i64 %n) #0 {
+; CHECK-LABEL: active_lane_mask_mstore_vscaleX2:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    whilelo p0.b, xzr, x1
+; CHECK-NEXT:    mov z0.h, #123 // =0x7b
+; CHECK-NEXT:    punpklo p0.h, p0.b
+; CHECK-NEXT:    st1h { z0.h }, p0, [x0]
+; CHECK-NEXT:    ret
+entry:
+  %mask = call <16 x i1> @llvm.get.active.lane.mask.v16i1(i64 0, i64 %n)
+  call void @llvm.masked.store.v16i16.p0(<16 x i16> splat(i16 123), ptr %p, <16 x i1> %mask)
+  ret void
+}
+
+define void @active_lane_mask_mstore_vscaleX4(ptr %p, i64 %n) #1 {
+; CHECK-LABEL: active_lane_mask_mstore_vscaleX4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    whilelo p0.b, xzr, x1
+; CHECK-NEXT:    mov z0.s, #123 // =0x7b
+; CHECK-NEXT:    punpklo p0.h, p0.b
+; CHECK-NEXT:    punpklo p0.h, p0.b
+; CHECK-NEXT:    st1w { z0.s }, p0, [x0]
+; CHECK-NEXT:    ret
+  %mask =  call <16 x i1> @llvm.get.active.lane.mask.v16i1.i64(i64 0, i64 %n)
+  call void @llvm.masked.store.v16i32.p0(<16 x i32> splat(i32 123), ptr %p, <16 x i1> %mask)
+  ret void
+}
+
+attributes #0 = { vscale_range(2,2) }
+attributes #1 = { vscale_range(4,4) }
diff --git a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
index b144060afe925..2362500254dc5 100644
--- a/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
+++ b/llvm/unittests/Target/AArch64/AArch64SelectionDAGTest.cpp
@@ -128,6 +128,31 @@ TEST_F(AArch64SelectionDAGTest, computeKnownBits_EXTRACT_SUBVECTOR) {
   EXPECT_TRUE(Known.isZero());
 }
 
+TEST_F(AArch64SelectionDAGTest, ComputeNumSignBits_GET_ACTIVE_LANE_MASK) {
+  // GET_ACTIVE_LANE_MASK promoted/widened to a vector integer type wider
+  // than i1 (e.g. v8i8) should report that all bits of each lane are sign
+  // bits.
+  SDLoc Loc;
+  auto Int64VT = EVT::getIntegerVT(Context, 64);
+  auto MaskVT = MVT::v8i8;
+  auto Base = DAG->getConstant(0, Loc, Int64VT);
+  auto TripCount = DAG->getConstant(8, Loc, Int64VT);
+  auto Op =
+      DAG->getNode(ISD::GET_ACTIVE_LANE_MASK, Loc, MaskVT, Base, TripCount);
+  EXPECT_EQ(DAG->ComputeNumSignBits(Op), 8u);
+}
+
+TEST_F(AArch64SelectionDAGTest, ComputeNumSignBitsSVE_GET_ACTIVE_LANE_MASK) {
+  SDLoc Loc;
+  auto Int64VT = EVT::getIntegerVT(Context, 64);
+  auto MaskVT = MVT::nxv8i16;
+  auto Base = DAG->getConstant(0, Loc, Int64VT);
+  auto TripCount = DAG->getConstant(8, Loc, Int64VT);
+  auto Op =
+      DAG->getNode(ISD::GET_ACTIVE_LANE_MASK, Loc, MaskVT, Base, TripCount);
+  EXPECT_EQ(DAG->ComputeNumSignBits(Op), 16u);
+}
+
 TEST_F(AArch64SelectionDAGTest, ComputeNumSignBits_SIGN_EXTEND_VECTOR_INREG) {
   SDLoc Loc;
   auto Int8VT = EVT::getIntegerVT(Context, 8);

``````````

</details>


https://github.com/llvm/llvm-project/pull/208977