[AArch64][SVE2p1] Allow more uses of mask in performActiveLaneMaskCombine #159360

kmclaughlin-arm · 2025-09-17T13:38:39Z

The combine replaces a get_active_lane_mask used by two extract subvectors with
a single paired whilelo intrinsic. When the instruction is used for control
flow in a vector loop, an additional extract of element 0 may introduce
other uses of the intrinsic such as ptest and reinterpret cast, which
is currently not supported.

This patch changes performActiveLaneMaskCombine to count the number of
extract subvectors using the mask instead of the total number of uses,
and returns the concatenated results of get_active_lane_mask.

llvmbot · 2025-09-17T13:39:10Z

@llvm/pr-subscribers-llvm-support

@llvm/pr-subscribers-backend-aarch64

Author: Kerry McLaughlin (kmclaughlin-arm)

Changes

The combine replaces a get_active_lane_mask used by two extract subvectors with
a single paired whilelo intrinsic. When the instruction is used for control
flow in a vector loop, an additional extract of element 0 may introduce
other uses of the intrinsic such as ptest and reinterpret cast, which
is currently not supported.

This patch changes performActiveLaneMaskCombine to count the number of
extract subvectors using the mask instead of the total number of uses,
and allows other uses by these additional operations.

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

3 Files Affected:

(modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+21-8)
(modified) llvm/lib/Target/AArch64/AArch64InstrInfo.cpp (+13-7)
(modified) llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll (+75)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index c9a756da0078d..9c7ecf944e763 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -18693,21 +18693,31 @@ performActiveLaneMaskCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
       (!ST->hasSVE2p1() && !(ST->hasSME2() && ST->isStreaming())))
     return SDValue();
 
-  unsigned NumUses = N->use_size();
+  // Count the number of users which are extract_vectors
+  // The only other valid users for this combine are ptest_first
+  // and reinterpret_cast.
+  unsigned NumExts = count_if(N->users(), [](SDNode *Use) {
+    return Use->getOpcode() == ISD::EXTRACT_SUBVECTOR;
+  });
+
   auto MaskEC = N->getValueType(0).getVectorElementCount();
-  if (!MaskEC.isKnownMultipleOf(NumUses))
+  if (!MaskEC.isKnownMultipleOf(NumExts))
     return SDValue();
 
-  ElementCount ExtMinEC = MaskEC.divideCoefficientBy(NumUses);
+  ElementCount ExtMinEC = MaskEC.divideCoefficientBy(NumExts);
   if (ExtMinEC.getKnownMinValue() < 2)
     return SDValue();
 
-  SmallVector<SDNode *> Extracts(NumUses, nullptr);
+  SmallVector<SDNode *> Extracts(NumExts, nullptr);
   for (SDNode *Use : N->users()) {
+    if (Use->getOpcode() == AArch64ISD::PTEST_FIRST ||
+        Use->getOpcode() == AArch64ISD::REINTERPRET_CAST)
+      continue;
+
     if (Use->getOpcode() != ISD::EXTRACT_SUBVECTOR)
       return SDValue();
 
-    // Ensure the extract type is correct (e.g. if NumUses is 4 and
+    // Ensure the extract type is correct (e.g. if NumExts is 4 and
     // the mask return type is nxv8i1, each extract should be nxv2i1.
     if (Use->getValueType(0).getVectorElementCount() != ExtMinEC)
       return SDValue();
@@ -18741,11 +18751,13 @@ performActiveLaneMaskCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
   DCI.CombineTo(Extracts[0], R.getValue(0));
   DCI.CombineTo(Extracts[1], R.getValue(1));
 
-  if (NumUses == 2)
-    return SDValue(N, 0);
+  if (NumExts == 2) {
+    DAG.ReplaceAllUsesOfValueWith(SDValue(N, 0), R.getValue(0));
+    return SDValue(SDValue(N, 0));
+  }
 
   auto Elts = DAG.getElementCount(DL, OpVT, ExtVT.getVectorElementCount() * 2);
-  for (unsigned I = 2; I < NumUses; I += 2) {
+  for (unsigned I = 2; I < NumExts; I += 2) {
     // After the first whilelo_x2, we need to increment the starting value.
     Idx = DAG.getNode(ISD::UADDSAT, DL, OpVT, Idx, Elts);
     R = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, {ExtVT, ExtVT}, {ID, Idx, TC});
@@ -18753,6 +18765,7 @@ performActiveLaneMaskCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
     DCI.CombineTo(Extracts[I + 1], R.getValue(1));
   }
 
+  DAG.ReplaceAllUsesOfValueWith(SDValue(N, 0), R.getValue(0));
   return SDValue(N, 0);
 }
 
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index bf3d47ac43607..069d08663fdea 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -1495,13 +1495,19 @@ AArch64InstrInfo::canRemovePTestInstr(MachineInstr *PTest, MachineInstr *Mask,
     if ((Mask == Pred) && PTest->getOpcode() == AArch64::PTEST_PP_ANY)
       return PredOpcode;
 
-    // For PTEST(PTRUE_ALL, WHILE), if the element size matches, the PTEST is
-    // redundant since WHILE performs an implicit PTEST with an all active
-    // mask.
-    if (isPTrueOpcode(MaskOpcode) && Mask->getOperand(1).getImm() == 31 &&
-        getElementSizeForOpcode(MaskOpcode) ==
-            getElementSizeForOpcode(PredOpcode))
-      return PredOpcode;
+    if (isPTrueOpcode(MaskOpcode) && Mask->getOperand(1).getImm() == 31) {
+      auto PTestOp = MRI->getUniqueVRegDef(PTest->getOperand(1).getReg());
+      if (PTest->getOpcode() == AArch64::PTEST_PP_FIRST && PTestOp->isCopy() &&
+          PTestOp->getOperand(1).getSubReg() == AArch64::psub0)
+        return PredOpcode;
+
+      // For PTEST(PTRUE_ALL, WHILE), if the element size matches, the PTEST is
+      // redundant since WHILE performs an implicit PTEST with an all active
+      // mask.
+      if (getElementSizeForOpcode(MaskOpcode) ==
+          getElementSizeForOpcode(PredOpcode))
+        return PredOpcode;
+    }
 
     return {};
   }
diff --git a/llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll b/llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll
index 5e01612e3881a..3b18008605413 100644
--- a/llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll
+++ b/llvm/test/CodeGen/AArch64/get-active-lane-mask-extract.ll
@@ -310,6 +310,81 @@ define void @test_2x32bit_mask_with_32bit_index_and_trip_count(i32 %i, i32 %n) #
   ret void
 }
 
+; Extra use of the get_active_lane_mask from an extractelement, which is replaced with ptest_first.
+
+define void @test_2x8bit_mask_with_extracts_and_ptest(i64 %i, i64 %n) {
+; CHECK-SVE-LABEL: test_2x8bit_mask_with_extracts_and_ptest:
+; CHECK-SVE:       // %bb.0: // %entry
+; CHECK-SVE-NEXT:    whilelo p1.b, x0, x1
+; CHECK-SVE-NEXT:    b.pl .LBB11_2
+; CHECK-SVE-NEXT:  // %bb.1: // %if.then
+; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
+; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SVE-NEXT:    b use
+; CHECK-SVE-NEXT:  .LBB11_2: // %if.end
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-SME2-LABEL: test_2x8bit_mask_with_extracts_and_ptest:
+; CHECK-SVE2p1-SME2:       // %bb.0: // %entry
+; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.h, p1.h }, x0, x1
+; CHECK-SVE2p1-SME2-NEXT:    b.pl .LBB11_2
+; CHECK-SVE2p1-SME2-NEXT:  // %bb.1: // %if.then
+; CHECK-SVE2p1-SME2-NEXT:    b use
+; CHECK-SVE2p1-SME2-NEXT:  .LBB11_2: // %if.end
+; CHECK-SVE2p1-SME2-NEXT:    ret
+entry:
+    %r = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i32(i64 %i, i64 %n)
+    %v0 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 0)
+    %v1 = call <vscale x 8 x i1> @llvm.vector.extract.nxv8i1.nxv16i1.i64(<vscale x 16 x i1> %r, i64 8)
+    %elt0 = extractelement <vscale x 16 x i1> %r, i32 0
+    br i1 %elt0, label %if.then, label %if.end
+
+if.then:
+    tail call void @use(<vscale x 8 x i1> %v0, <vscale x 8 x i1> %v1)
+    br label %if.end
+
+if.end:
+    ret void
+}
+
+; Extra use of the get_active_lane_mask from an extractelement, which is
+; replaced with ptest_first and reinterpret_casts because the extract is not nxv16i1.
+
+define void @test_2x8bit_mask_with_extracts_and_reinterpret_casts(i64 %i, i64 %n) {
+; CHECK-SVE-LABEL: test_2x8bit_mask_with_extracts_and_reinterpret_casts:
+; CHECK-SVE:       // %bb.0: // %entry
+; CHECK-SVE-NEXT:    whilelo p1.h, x0, x1
+; CHECK-SVE-NEXT:    b.pl .LBB12_2
+; CHECK-SVE-NEXT:  // %bb.1: // %if.then
+; CHECK-SVE-NEXT:    punpklo p0.h, p1.b
+; CHECK-SVE-NEXT:    punpkhi p1.h, p1.b
+; CHECK-SVE-NEXT:    b use
+; CHECK-SVE-NEXT:  .LBB12_2: // %if.end
+; CHECK-SVE-NEXT:    ret
+;
+; CHECK-SVE2p1-SME2-LABEL: test_2x8bit_mask_with_extracts_and_reinterpret_casts:
+; CHECK-SVE2p1-SME2:       // %bb.0: // %entry
+; CHECK-SVE2p1-SME2-NEXT:    whilelo { p0.s, p1.s }, x0, x1
+; CHECK-SVE2p1-SME2-NEXT:    b.pl .LBB12_2
+; CHECK-SVE2p1-SME2-NEXT:  // %bb.1: // %if.then
+; CHECK-SVE2p1-SME2-NEXT:    b use
+; CHECK-SVE2p1-SME2-NEXT:  .LBB12_2: // %if.end
+; CHECK-SVE2p1-SME2-NEXT:    ret
+entry:
+    %r = call <vscale x 8 x i1> @llvm.get.active.lane.mask.nxv8i1.i64(i64 %i, i64 %n)
+    %v0 = tail call <vscale x 4 x i1> @llvm.vector.extract.nxv4i1.nxv8i1(<vscale x 8 x i1> %r, i64 0)
+    %v1 = tail call <vscale x 4 x i1> @llvm.vector.extract.nxv4i1.nxv8i1(<vscale x 8 x i1> %r, i64 4)
+    %elt0 = extractelement <vscale x 8 x i1> %r, i64 0
+    br i1 %elt0, label %if.then, label %if.end
+
+if.then:
+    tail call void @use(<vscale x 4 x i1> %v0, <vscale x 4 x i1> %v1)
+    br label %if.end
+
+if.end:
+    ret void
+}
+
 declare void @use(...)
 
 attributes #0 = { nounwind }

paulwalker-arm · 2025-09-17T16:50:24Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

    // After the first whilelo_x2, we need to increment the starting value.
    Idx = DAG.getNode(ISD::UADDSAT, DL, OpVT, Idx, Elts);
    R = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL, {ExtVT, ExtVT}, {ID, Idx, TC});
    DCI.CombineTo(Extracts[I], R.getValue(0));
    DCI.CombineTo(Extracts[I + 1], R.getValue(1));
  }

+  DAG.ReplaceAllUsesOfValueWith(SDValue(N, 0), R.getValue(0));


Similar to the above but in this case there's the extra problem that you're replacing N with the first result of the last instance of emitted while_pair, which is even more wrong if the uses happened to be a PTEST_FIRST.

paulwalker-arm · 2025-09-17T16:56:15Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+    if (Use->getOpcode() == AArch64ISD::PTEST_FIRST ||
+        Use->getOpcode() == AArch64ISD::REINTERPRET_CAST)
+      continue;


I don't think it's worth trying to special case based on the users beyond verifying the presence of the relevant ISD::EXTRACT_SUBVECTOR to prove the value of using the while_pair instructions. Even if the original while remains, the resulting code might be better because multiple extracts have been replaced by a single while_pair?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

paulwalker-arm · 2025-09-17T16:59:40Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

-  if (NumUses == 2)
-    return SDValue(N, 0);
+  if (NumExts == 2) {
+    DAG.ReplaceAllUsesOfValueWith(SDValue(N, 0), R.getValue(0));


Not sure if I'm misunderstanding something but I don't know how this works because SDValue(N, 0) and R.getValue(0) are going to have different result types? so the post combine DAG is likely to be broken.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

paulwalker-arm · 2025-09-17T17:24:41Z

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

+    if (isPTrueOpcode(MaskOpcode) && Mask->getOperand(1).getImm() == 31) {
+      auto PTestOp = MRI->getUniqueVRegDef(PTest->getOperand(1).getReg());
+      if (PTest->getOpcode() == AArch64::PTEST_PP_FIRST && PTestOp->isCopy() &&
+          PTestOp->getOperand(1).getSubReg() == AArch64::psub0)
+        return PredOpcode;
+
+      // For PTEST(PTRUE_ALL, WHILE), if the element size matches, the PTEST is
+      // redundant since WHILE performs an implicit PTEST with an all active
+      // mask.
+      if (getElementSizeForOpcode(MaskOpcode) ==
+          getElementSizeForOpcode(PredOpcode))
+        return PredOpcode;
+    }


As discussed, I'm kind of hoping this is just fallout from the DAG being broken and will not be necessary once fixed.

paulwalker-arm · 2025-09-17T17:31:31Z

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

+; Extra use of the get_active_lane_mask from an extractelement, which is
+; replaced with ptest_first and reinterpret_casts because the extract is not nxv16i1.
+
+define void @test_2x8bit_mask_with_extracts_and_reinterpret_casts(i64 %i, i64 %n) {


It's worth adding a similar test for the NumExts != 2 case, if only to see if that better exposes the issues I believe exist in the PR as it stands today.

…bine The combine replaces a get_active_lane_mask used by two extract subvectors with a single paired whilelo intrinsic. When the instruction is used for control flow in a vector loop, an additional extract of element 0 may introduce other uses of the intrinsic such as ptest and reinterpret cast, which is currently not supported. This patch changes performActiveLaneMaskCombine to count the number of extract subvectors using the mask instead of the total number of uses, and allows other uses by these additional operations.

…performActiveLaneMaskCombine - Add tests for the 4 extracts case which will use ptest & reinterpret_cast - Remove changes to canRemovePTestInstr

david-arm · 2025-09-18T15:58:48Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 performActiveLaneMaskCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
                             const AArch64Subtarget *ST) {
-  if (DCI.isBeforeLegalize())
+  if (DCI.isBeforeLegalize() && !!DCI.isBeforeLegalizeOps())


nit: The !! looks a little odd. Is it possible to just use DCI.isBeforeLegalizeOps()?

This was a mistake, it should be !DCI.isBeforeLegalizeOps()

I don't think this is doing what you intend. Given !DCI.isBeforeLegalizeOps()) means AfterLegalizeVectorOps, you've effectively written if "before type legalisation" and "after vector ops legalisation", which is always going to be false because they are are opposite ends of selection.

I've now removed this altogether, since finding that a concat with 4+ inputs will be split after this combine without changing the stage at which it applies (although I've also made changes to create multiple concat_vectors in the latest commit too).

- Create multiple concats in performActiveLaneMaskCombine when there are more than 2 extracts

paulwalker-arm · 2025-09-26T14:06:08Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+
  auto MaskEC = N->getValueType(0).getVectorElementCount();
-  if (!MaskEC.isKnownMultipleOf(NumUses))
+  if (NumExts == 0 || !MaskEC.isKnownMultipleOf(NumExts))


Is it possible to move the zero check into isKnownMultipleOf?

paulwalker-arm · 2025-09-26T14:15:23Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

  DCI.CombineTo(Extracts[0], R.getValue(0));
  DCI.CombineTo(Extracts[1], R.getValue(1));
+  SmallVector<SDValue> Concats = {DAG.getNode(
+      ISD::CONCAT_VECTORS, DL, DoubleExtVT, {R.getValue(0), R.getValue(1)})};


Is wrapping the operands in {} necessary?

No, I've removed them in both places

paulwalker-arm · 2025-09-26T14:15:42Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+    Concats.push_back(DAG.getNode(ISD::CONCAT_VECTORS, DL, DoubleExtVT,
+                                  {R.getValue(0), R.getValue(1)}));


Is wrapping the operands in {} necessary?

…t.cpp - Remove unnecessary {} from getNode in performActiveLaneMaskCombine

paulwalker-arm · 2025-09-30T10:51:52Z

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

+  // The only other valid users for this combine are ptest_first
+  // and reinterpret_cast.


This last bit can be removed because the ptest_first restriction no longer applies.

…bine (llvm#159360) The combine replaces a get_active_lane_mask used by two extract subvectors with a single paired whilelo intrinsic. When the instruction is used for control flow in a vector loop, an additional extract of element 0 may introduce other uses of the intrinsic such as ptest and reinterpret cast, which is currently not supported. This patch changes performActiveLaneMaskCombine to count the number of extract subvectors using the mask instead of the total number of uses, and returns the concatenated results of get_active_lane_mask.

llvmbot added the backend:AArch64 label Sep 17, 2025

kmclaughlin-arm requested review from CarolineConcatto, david-arm and paulwalker-arm September 17, 2025 13:48

kmclaughlin-arm force-pushed the alm-combine-ptest branch from ea8a053 to 02a75b7 Compare September 17, 2025 16:13

paulwalker-arm reviewed Sep 17, 2025

View reviewed changes

kmclaughlin-arm added 3 commits September 18, 2025 14:17

- Tests for get_active_lane_mask with more uses than extract_vector

7f042ec

- Return a concat_vector created from the results of whilelo_x2 from …

70adaf7

…performActiveLaneMaskCombine - Add tests for the 4 extracts case which will use ptest & reinterpret_cast - Remove changes to canRemovePTestInstr

kmclaughlin-arm force-pushed the alm-combine-ptest branch from 02a75b7 to 70adaf7 Compare September 18, 2025 15:35

david-arm reviewed Sep 18, 2025

View reviewed changes

kmclaughlin-arm added 2 commits September 18, 2025 16:32

- Fix !!DCI.isBeforeLegalizeOps -> !DCI.isBeforeLegalizeOps

083701d

- Remove !DCI.isBeforeLegalizeOps

d64e2d8

- Create multiple concats in performActiveLaneMaskCombine when there are more than 2 extracts

paulwalker-arm reviewed Sep 26, 2025

View reviewed changes

- Moved zero check into isKnownMultipleOf & added test to TypeSizeTes…

c077d34

…t.cpp - Remove unnecessary {} from getNode in performActiveLaneMaskCombine

llvmbot added the llvm:support label Sep 26, 2025

paulwalker-arm approved these changes Sep 30, 2025

View reviewed changes

- Fix comment in performActiveLaneMaskCombine

a815152

kmclaughlin-arm merged commit cf50bbf into llvm:main Sep 30, 2025
9 checks passed

kmclaughlin-arm deleted the alm-combine-ptest branch October 7, 2025 08:39

aemerson mentioned this pull request Oct 16, 2025

[clang-reorder-fields] Check for flexible array member (#160262) #163697

Closed

		Concats.push_back(DAG.getNode(ISD::CONCAT_VECTORS, DL, DoubleExtVT,
		{R.getValue(0), R.getValue(1)}));

		// The only other valid users for this combine are ptest_first
		// and reinterpret_cast.

[AArch64][SVE2p1] Allow more uses of mask in performActiveLaneMaskCombine #159360

[AArch64][SVE2p1] Allow more uses of mask in performActiveLaneMaskCombine #159360

Uh oh!

Conversation

kmclaughlin-arm commented Sep 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

kmclaughlin-arm commented Sep 17, 2025 •

edited

Loading

llvmbot commented Sep 17, 2025 •

edited

Loading