[llvm] [AArch64] Improve fixed vector lowering for cttz/ctlz when sve (PR #192427)

[Yashwant Singh via llvm-commits]

https://github.com/yashssh updated https://github.com/llvm/llvm-project/pull/192427

>From 59b491b56150318d42dd575ad616c06a2d83ec0c Mon Sep 17 00:00:00 2001
From: Yashwant Singh <yashwants at nvidia.com>
Date: Thu, 16 Apr 2026 04:34:35 -0700
Subject: [PATCH] [AArch64] Improve fixed vector lowering for cttz/ctlz when
 sve

When available we should be able to lower fixed size 64/128 bit cttz/ctlz vector instructions using sve. Below are the performance uplifts

+--------------+--------------+--------------+
| Function | CTTZ Speedup | CTLZ Speedup |
+--------------+--------------+--------------+
| u8x8 | 1.33x | 1.00x |
| u16x4 | 1.26x | 1.00x |
| u32x2 | 1.26x | 1.00x |
| u64x1 | 1.00x | 7.00x |
+--------------+--------------+--------------+
| u8x16 | 1.67x | 1.00x |
| u16x8 | 1.51x | 1.00x |
| u32x4 | 1.51x | 1.00x |
| u64x2 | 1.26x | 7.00x |
+--------------+--------------+--------------+
---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp     | 13 +++++++++----
 .../AArch64/sve-fixed-length-bit-counting.ll        |  8 ++++----
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 8c53ea0dd94eb..73c0e39e98f82 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1943,6 +1943,11 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       setOperationAction(ISD::VECREDUCE_UMIN, MVT::v2i64, Custom);
     }
 
+    for (auto VT : {MVT::v1i64, MVT::v2i64}) {
+      setOperationAction(ISD::CTLZ, VT, Custom);
+      setOperationAction(ISD::CTTZ, VT, Custom);
+    }
+
     // NOTE: Currently this has to happen after computeRegisterProperties rather
     // than the preferred option of combining it with the addRegisterClass call.
     if (Subtarget->useSVEForFixedLengthVectors()) {
@@ -1970,9 +1975,6 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
 
       // These operations are not supported on NEON but SVE can do them.
       setOperationAction(ISD::BITREVERSE, MVT::v1i64, Custom);
-      setOperationAction(ISD::CTLZ, MVT::v1i64, Custom);
-      setOperationAction(ISD::CTLZ, MVT::v2i64, Custom);
-      setOperationAction(ISD::CTTZ, MVT::v1i64, Custom);
       setOperationAction(ISD::SMAX, MVT::v1i64, Custom);
       setOperationAction(ISD::SMAX, MVT::v2i64, Custom);
       setOperationAction(ISD::SMIN, MVT::v1i64, Custom);
@@ -1990,11 +1992,14 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       for (auto VT : {MVT::v8i8, MVT::v16i8, MVT::v4i16, MVT::v8i16,
                       MVT::v2i32, MVT::v4i32, MVT::v2i64}) {
         setOperationAction(ISD::BITREVERSE, VT, Custom);
-        setOperationAction(ISD::CTTZ, VT, Custom);
         setOperationAction(ISD::VECREDUCE_AND, VT, Custom);
         setOperationAction(ISD::VECREDUCE_OR, VT, Custom);
         setOperationAction(ISD::VECREDUCE_XOR, VT, Custom);
       }
+      for (auto VT : {MVT::v8i8, MVT::v16i8, MVT::v4i16, MVT::v8i16, MVT::v2i32,
+                      MVT::v4i32}) {
+        setOperationAction(ISD::CTTZ, VT, Custom);
+      }
 
       // Use SVE for vectors with more than 2 elements.
       for (auto VT : {MVT::v4f16, MVT::v8f16, MVT::v4f32})
diff --git a/llvm/test/CodeGen/AArch64/sve-fixed-length-bit-counting.ll b/llvm/test/CodeGen/AArch64/sve-fixed-length-bit-counting.ll
index b62b850434469..57797d04453bc 100644
--- a/llvm/test/CodeGen/AArch64/sve-fixed-length-bit-counting.ll
+++ b/llvm/test/CodeGen/AArch64/sve-fixed-length-bit-counting.ll
@@ -273,7 +273,7 @@ define void @ctlz_v64i32(ptr %a)  vscale_range(16,0) #0 {
   ret void
 }
 
-define <1 x i64> @ctlz_v1i64(<1 x i64> %op) vscale_range(2,0) #0 {
+define <1 x i64> @ctlz_v1i64(<1 x i64> %op) #0 {
 ; CHECK-LABEL: ctlz_v1i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d, vl1
@@ -285,7 +285,7 @@ define <1 x i64> @ctlz_v1i64(<1 x i64> %op) vscale_range(2,0) #0 {
   ret <1 x i64> %res
 }
 
-define <2 x i64> @ctlz_v2i64(<2 x i64> %op) vscale_range(2,0) #0 {
+define <2 x i64> @ctlz_v2i64(<2 x i64> %op) #0 {
 ; CHECK-LABEL: ctlz_v2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d, vl2
@@ -1033,7 +1033,7 @@ define void @cttz_v64i32(ptr %a) vscale_range(16,0) #0 {
   ret void
 }
 
-define <1 x i64> @cttz_v1i64(<1 x i64> %op) vscale_range(2,0) #0 {
+define <1 x i64> @cttz_v1i64(<1 x i64> %op) #0 {
 ; CHECK-LABEL: cttz_v1i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d, vl1
@@ -1046,7 +1046,7 @@ define <1 x i64> @cttz_v1i64(<1 x i64> %op) vscale_range(2,0) #0 {
   ret <1 x i64> %res
 }
 
-define <2 x i64> @cttz_v2i64(<2 x i64> %op) vscale_range(2,0) #0 {
+define <2 x i64> @cttz_v2i64(<2 x i64> %op) #0 {
 ; CHECK-LABEL: cttz_v2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d, vl2