- Add asserts for the result & operand types to ReplaceGetActiveLaneMaskResults

- Rename new test with partial extracts

Author: kmclaughlin-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -27335,12 +27335,19 @@ void AArch64TargetLowering::ReplaceGetActiveLaneMaskResults(
   if (!Subtarget->hasSVE2p1())
     return;
 
+  assert(N->getValueType(0) == MVT::nxv32i1 &&
+         "Unexpected result type for get.active.lane.mask");
+
   SDLoc DL(N);
   SDValue Idx = N->getOperand(0);
   SDValue TC = N->getOperand(1);
+
+  assert(Idx.getValueType().getFixedSizeInBits() <= 64 &&
+         "Unexpected operand type for get.active.lane.mask");
+
   if (Idx.getValueType() != MVT::i64) {
-    Idx = DAG.getZExtOrTrunc(Idx, DL, MVT::i64);
-    TC = DAG.getZExtOrTrunc(TC, DL, MVT::i64);
+    Idx = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, Idx);
+    TC = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i64, TC);
   }
 
   SDValue ID =

llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll

@@ -272,9 +272,8 @@ define void @test_2x16bit_mask_with_32bit_index_and_trip_count_ext8(i32 %i, i32
   ret void
 }
 
-; Negative test for when not extracting exactly two halves of the source vector
-define void @test_illegal_type_with_partial_extracts(i32 %i, i32 %n) #0 {
-; CHECK-SVE-LABEL: test_illegal_type_with_partial_extracts:
+define void @test_2x16bit_mask_with_32bit_index_and_trip_count_part_extracts(i32 %i, i32 %n) #0 {
+; CHECK-SVE-LABEL: test_2x16bit_mask_with_32bit_index_and_trip_count_part_extracts:
 ; CHECK-SVE:       // %bb.0:
 ; CHECK-SVE-NEXT:    rdvl x8, #1
 ; CHECK-SVE-NEXT:    adds w8, w0, w8
@@ -285,7 +284,7 @@ define void @test_illegal_type_with_partial_extracts(i32 %i, i32 %n) #0 {
 ; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
 ; CHECK-SVE-NEXT:    b use
 ;
-; CHECK-SVE2p1-LABEL: test_illegal_type_with_partial_extracts:
+; CHECK-SVE2p1-LABEL: test_2x16bit_mask_with_32bit_index_and_trip_count_part_extracts:
 ; CHECK-SVE2p1:       // %bb.0:
 ; CHECK-SVE2p1-NEXT:    mov w8, w1
 ; CHECK-SVE2p1-NEXT:    mov w9, w0