From 721ad4aa180d5b372323a0a5d06b799988d708be Mon Sep 17 00:00:00 2001
From: George Steed <george.steed@arm.com>
Date: Fri, 4 Oct 2024 14:20:26 +0100
Subject: [PATCH] [AArch64] Add SME implementation of ScaleUVRowDown2Box

There is no benefit from an SVE version of this kernel for devices with
an SVE vector length of 128-bits, so skip directly to SME instead. We do
not use the ZA tile here, so this is a purely streaming-SVE (SSVE)
implementation.

Change-Id: Ie15bb4e7484b61e78f405ad4e8a8a7bbb66b7edb
Reviewed-on: https://chromium-review.googlesource.com/c/libyuv/libyuv/+/5979727
Reviewed-by: Justin Green <greenjustin@google.com>
Reviewed-by: Frank Barchard <fbarchard@chromium.org>
---
 include/libyuv/scale_row.h |  5 ++++
 source/scale_sme.cc        | 54 ++++++++++++++++++++++++++++++++++++++
 source/scale_uv.cc         | 13 ++++++---
 3 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/include/libyuv/scale_row.h b/include/libyuv/scale_row.h
index c4df7348a..101ccbf84 100644
--- a/include/libyuv/scale_row.h
+++ b/include/libyuv/scale_row.h
@@ -122,6 +122,7 @@ extern "C" {
 #define HAS_SCALEROWDOWN2_SME
 #define HAS_SCALEUVROWDOWN2_SME
 #define HAS_SCALEUVROWDOWN2LINEAR_SME
+#define HAS_SCALEUVROWDOWN2BOX_SME
 #endif
 
 #if !defined(LIBYUV_DISABLE_MSA) && defined(__mips_msa)
@@ -1172,6 +1173,10 @@ void ScaleUVRowDown2Box_NEON(const uint8_t* src_ptr,
                              ptrdiff_t src_stride,
                              uint8_t* dst,
                              int dst_width);
+void ScaleUVRowDown2Box_SME(const uint8_t* src_ptr,
+                            ptrdiff_t src_stride,
+                            uint8_t* dst,
+                            int dst_width);
 void ScaleUVRowDown2_MSA(const uint8_t* src_ptr,
                          ptrdiff_t src_stride,
                          uint8_t* dst_uv,
diff --git a/source/scale_sme.cc b/source/scale_sme.cc
index ee63c090b..f4166a5b1 100644
--- a/source/scale_sme.cc
+++ b/source/scale_sme.cc
@@ -228,6 +228,60 @@ __arm_locally_streaming void ScaleUVRowDown2Linear_SME(const uint8_t* src_uv,
       : "z0", "z1", "p0", "p1");
 }
 
+#define SCALEUVROWDOWN2BOX_SVE                               \
+  "ld2h     {z0.h, z1.h}, p0/z, [%[src_uv]]              \n" \
+  "ld2h     {z2.h, z3.h}, p0/z, [%[src2_uv]]             \n" \
+  "incb     %[src_uv], all, mul #2                       \n" \
+  "incb     %[src2_uv], all, mul #2                      \n" \
+  "uaddlb   z4.h, z0.b, z1.b                             \n" \
+  "uaddlt   z5.h, z0.b, z1.b                             \n" \
+  "uaddlb   z6.h, z2.b, z3.b                             \n" \
+  "uaddlt   z7.h, z2.b, z3.b                             \n" \
+  "add      z4.h, z4.h, z6.h                             \n" \
+  "add      z5.h, z5.h, z7.h                             \n" \
+  "rshrnb   z0.b, z4.h, #2                               \n" \
+  "rshrnt   z0.b, z5.h, #2                               \n" \
+  "st1h     {z0.h}, p0, [%[dst_uv]]                      \n" \
+  "incb     %[dst_uv], all, mul #1                       \n"
+
+__arm_locally_streaming void ScaleUVRowDown2Box_SME(const uint8_t* src_uv,
+                                                    ptrdiff_t src_stride,
+                                                    uint8_t* dst_uv,
+                                                    int dst_width) {
+  // Streaming-SVE only, no use of ZA tile.
+  const uint8_t* src2_uv = src_uv + src_stride;
+  int vl;
+  asm volatile(
+      "cnth     %x[vl]                                       \n"
+      "ptrue    p1.b                                         \n"
+      "subs     %w[dst_width], %w[dst_width], %w[vl]         \n"
+      "b.lt     2f                                           \n"
+
+      "ptrue    p0.h                                         \n"
+      "1:                                                    \n"  //
+      SCALEUVROWDOWN2BOX_SVE
+      "subs     %w[dst_width], %w[dst_width], %w[vl]         \n"
+      "b.ge     1b                                           \n"
+
+      "2:                                                    \n"
+      "adds     %w[dst_width], %w[dst_width], %w[vl]         \n"
+      "b.eq     99f                                          \n"
+
+      "whilelt  p0.h, wzr, %w[dst_width]                     \n"  //
+      SCALEUVROWDOWN2BOX_SVE
+
+      "99:                                                   \n"
+      : [src_uv] "+r"(src_uv),        // %[src_uv]
+        [src2_uv] "+r"(src2_uv),      // %[src2_uv]
+        [dst_uv] "+r"(dst_uv),        // %[dst_uv]
+        [dst_width] "+r"(dst_width),  // %[dst_width]
+        [vl] "=r"(vl)                 // %[vl]
+      :
+      : "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", "p0", "p1");
+}
+
+#undef SCALEUVROWDOWN2BOX_SVE
+
 #endif  // !defined(LIBYUV_DISABLE_SME) && defined(CLANG_HAS_SME) &&
         // defined(__aarch64__)
 
diff --git a/source/scale_uv.cc b/source/scale_uv.cc
index 4498027ea..9ef2e1387 100644
--- a/source/scale_uv.cc
+++ b/source/scale_uv.cc
@@ -121,10 +121,10 @@ static void ScaleUVDown2(int src_width,
   }
 #endif
 #if defined(HAS_SCALEUVROWDOWN2_SME)
-  if (TestCpuFlag(kCpuHasSME) &&
-      (filtering == kFilterNone || filtering == kFilterLinear)) {
-    ScaleUVRowDown2 = filtering == kFilterNone ? ScaleUVRowDown2_SME
-                                               : ScaleUVRowDown2Linear_SME;
+  if (TestCpuFlag(kCpuHasSME)) {
+    ScaleUVRowDown2 = filtering == kFilterNone     ? ScaleUVRowDown2_SME
+                      : filtering == kFilterLinear ? ScaleUVRowDown2Linear_SME
+                                                   : ScaleUVRowDown2Box_SME;
   }
 #endif
 #if defined(HAS_SCALEUVROWDOWN2_RVV)
@@ -241,6 +241,11 @@ static int ScaleUVDown4Box(int src_width,
     }
   }
 #endif
+#if defined(HAS_SCALEUVROWDOWN2BOX_SME)
+  if (TestCpuFlag(kCpuHasSME)) {
+    ScaleUVRowDown2 = ScaleUVRowDown2Box_SME;
+  }
+#endif
 #if defined(HAS_SCALEUVROWDOWN2BOX_RVV)
   if (TestCpuFlag(kCpuHasRVV)) {
     ScaleUVRowDown2 = ScaleUVRowDown2Box_RVV;