[llvm] [AArch64][SVE] Match (add_like x (lsr y, c)) -> usra x, y, c (PR #197657)

[Harry Ramsey via llvm-commits]

https://github.com/Harry-Ramsey created https://github.com/llvm/llvm-project/pull/197657

Modify SVEE USRA pattern to accept add_like, so both add and or disjoint forms can select usra.

Add known-bits support for predicated SVE logical shifts, allowing or_disjoint matching to prove disjointness for plain ORs where possible.

>From 0a89000985556bfebe4071f0b91a96e3778f5eaf Mon Sep 17 00:00:00 2001
From: Harry Ramsey <harry.ramsey at arm.com>
Date: Wed, 13 May 2026 15:46:55 +0000
Subject: [PATCH] [AArch64][SVE] Match (add_like x (lsr y, c)) -> usra x, y, c

Modify SVEE USRA pattern to accept add_like, so both add and or disjoint
forms can select usra.

Add known-bits support for predicated SVE logical shifts, allowing
or_disjoint matching to prove disjointness for plain ORs where possible.
---
 .../lib/CodeGen/SelectionDAG/SelectionDAG.cpp |  5 ----
 .../Target/AArch64/AArch64ISelLowering.cpp    | 12 ++++++++++
 .../lib/Target/AArch64/AArch64SVEInstrInfo.td |  2 +-
 llvm/test/CodeGen/AArch64/sve2-sli-sri.ll     | 24 ++++++++++++-------
 4 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 9af378ec28868..29bad5c513c55 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -4525,11 +4525,6 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
   case ISD::INTRINSIC_WO_CHAIN:
   case ISD::INTRINSIC_W_CHAIN:
   case ISD::INTRINSIC_VOID:
-    // TODO: Probably okay to remove after audit; here to reduce change size
-    // in initial enablement patch for scalable vectors
-    if (Op.getValueType().isScalableVector())
-      break;
-
     // Allow the target to implement this method for its nodes.
     TLI->computeKnownBitsForTargetNode(Op, Known, DemandedElts, *this, Depth);
     break;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9de0b461a216c..e78fd18f6604a 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2902,6 +2902,18 @@ void AArch64TargetLowering::computeKnownBitsForTargetNode(
     }
     break;
   }
+  case AArch64ISD::SHL_PRED:
+  case AArch64ISD::SRL_PRED: {
+    KnownBits KnownVal =
+        DAG.computeKnownBits(Op->getOperand(1), DemandedElts, Depth + 1);
+    KnownBits KnownAmt =
+        DAG.computeKnownBits(Op->getOperand(2), DemandedElts, Depth + 1);
+
+    Known = Op.getOpcode() == AArch64ISD::SHL_PRED
+                ? KnownBits::shl(KnownVal, KnownAmt)
+                : KnownBits::lshr(KnownVal, KnownAmt);
+    break;
+  }
   case ISD::INTRINSIC_WO_CHAIN:
   case ISD::INTRINSIC_VOID: {
     unsigned IntNo = Op.getConstantOperandVal(0);
diff --git a/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td b/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td
index 9df77f8e93c64..e4c281a84d72b 100644
--- a/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td
@@ -362,7 +362,7 @@ def AArch64uaba : PatFrags<(ops node:$op1, node:$op2, node:$op3),
 
 def AArch64usra : PatFrags<(ops node:$op1, node:$op2, node:$op3),
                            [(int_aarch64_sve_usra node:$op1, node:$op2, node:$op3),
-                            (add node:$op1, (AArch64lsr_p (SVEAnyPredicate), node:$op2, (SVEShiftSplatImmR (i32 node:$op3))))]>;
+                            (add_like node:$op1, (AArch64lsr_p (SVEAnyPredicate), node:$op2, (SVEShiftSplatImmR (i32 node:$op3))))]>;
 
 def AArch64ssra : PatFrags<(ops node:$op1, node:$op2, node:$op3),
                            [(int_aarch64_sve_ssra node:$op1, node:$op2, node:$op3),
diff --git a/llvm/test/CodeGen/AArch64/sve2-sli-sri.ll b/llvm/test/CodeGen/AArch64/sve2-sli-sri.ll
index 80999fb1f4864..a7048731850d3 100644
--- a/llvm/test/CodeGen/AArch64/sve2-sli-sri.ll
+++ b/llvm/test/CodeGen/AArch64/sve2-sli-sri.ll
@@ -118,14 +118,22 @@ define <vscale x 8 x i16> @testRightGood8x16(<vscale x 8 x i16> %src1, <vscale x
 }
 
 define <vscale x 8 x i16> @testRightBad8x16(<vscale x 8 x i16> %src1, <vscale x 8 x i16> %src2) {
-; CHECK-LABEL: testRightBad8x16:
-; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #16500 // =0x4074
-; CHECK-NEXT:    lsr z1.h, z1.h, #14
-; CHECK-NEXT:    mov z2.h, w8
-; CHECK-NEXT:    and z0.d, z0.d, z2.d
-; CHECK-NEXT:    orr z0.d, z0.d, z1.d
-; CHECK-NEXT:    ret
+; SVE-LABEL: testRightBad8x16:
+; SVE:       // %bb.0:
+; SVE-NEXT:    mov w8, #16500 // =0x4074
+; SVE-NEXT:    lsr z1.h, z1.h, #14
+; SVE-NEXT:    mov z2.h, w8
+; SVE-NEXT:    and z0.d, z0.d, z2.d
+; SVE-NEXT:    orr z0.d, z0.d, z1.d
+; SVE-NEXT:    ret
+;
+; SVE2-LABEL: testRightBad8x16:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    mov w8, #16500 // =0x4074
+; SVE2-NEXT:    mov z2.h, w8
+; SVE2-NEXT:    and z0.d, z0.d, z2.d
+; SVE2-NEXT:    usra z0.h, z1.h, #14
+; SVE2-NEXT:    ret
   %and.i = and <vscale x 8 x i16> %src1, splat(i16 16500)
   %vshl_n = lshr <vscale x 8 x i16> %src2, splat(i16 14)
   %result = or <vscale x 8 x i16> %and.i, %vshl_n