modified sve_utils, added Copyright info, rebased

Author: ashwins990

src/plugins/intel_cpu/src/nodes/kernels/aarch64/pa_kernels.hpp

@@ -0,0 +1,53 @@
+// Copyright (C) 2024 FUJITSU LIMITED
+// SPDX-License-Identifier: Apache-2.0
+//
+#include <arm_sve.h>
+
+#include "openvino/core/type/float16.hpp"
+
+template <typename T, typename... Args>
+constexpr bool one_of(T val, Args... args) {
+    return ((val == args) || ...);
+}
+
+template <size_t T_SIZE>
+svbool_t sve_predicate() {
+    static_assert(one_of(T_SIZE, 8, 16, 32, 64), "Unexpected parameter size");
+    if constexpr (8 == T_SIZE) {
+        return svptrue_b8();
+    } else if (16 == T_SIZE) {
+        return svptrue_b16();
+    } else if (32 == T_SIZE) {
+        return svptrue_b32();
+    } else if (64 == T_SIZE) {
+        return svptrue_b64();
+    }
+}
+
+template <typename T_TYPE, size_t T_SIZE>
+svbool_t sve_predicate(T_TYPE lower, T_TYPE higher) {
+    static_assert(one_of(T_SIZE, 8, 16, 32, 64), "Unexpected parameter size");
+    if constexpr (8 == T_SIZE) {
+        return svwhilelt_b8(lower, higher);
+    } else if (16 == T_SIZE) {
+        return svwhilelt_b16(lower, higher);
+    } else if (32 == T_SIZE) {
+        return svwhilelt_b32(lower, higher);
+    } else if (64 == T_SIZE) {
+        return svwhilelt_b64(lower, higher);
+    }
+}
+
+template <size_t T_SIZE>
+size_t sve_vlen() {
+    static_assert(one_of(T_SIZE, 8, 16, 32, 64), "Unexpected parameter size");
+    if constexpr (8 == T_SIZE) {
+        return svcntb();
+    } else if (16 == T_SIZE) {
+        return svcnth();
+    } else if (32 == T_SIZE) {
+        return svcntw();
+    } else if (64 == T_SIZE) {
+        return svcntd();
+    }
+}

src/plugins/intel_cpu/src/nodes/kernels/aarch64/sve_utils.hpp

@@ -1,3 +1,9 @@
+// Copyright (C) 2024 FUJITSU LIMITED
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
 #include <arm_neon.h>
 #include <kai/ukernels/matmul/matmul_clamp_f16_f16_f16p/kai_matmul_clamp_f16_f16_f16p16x1biasf16_6x16x8_neon_mla.h>
 #include <kai/ukernels/matmul/matmul_clamp_f16_f16_f16p/kai_matmul_clamp_f16_f16_f16p_interface.h>

src/plugins/intel_cpu/src/nodes/kernels/kai/kleidi_kernel.hpp

@@ -30,7 +30,7 @@
 #elif defined(OPENVINO_ARCH_ARM64) && defined(HAVE_SVE)
 #    include "arm_sve.h"
 #    include "nodes/kernels/aarch64/brgemm_kernel.hpp"
-#    include "nodes/kernels/aarch64/pa_kernels.hpp"
+#    include "nodes/kernels/aarch64/sve_utils.hpp"
 #    include "nodes/kernels/kai/kleidi_kernel.hpp"
 #endif
 
@@ -39,7 +39,7 @@ namespace ov::Extensions::Cpu::XARCH {
 using namespace ov;
 using namespace ov::intel_cpu;
 
-// currently depends on brgemm which only support x64
+// currently depends on brgemm which only support x64 or ARM SVE
 #if defined(OPENVINO_ARCH_X86_64) || (defined(OPENVINO_ARCH_ARM64) && defined(HAVE_SVE))
 
 #    if defined(HAVE_AVX2) || defined(HAVE_AVX512F)
@@ -72,16 +72,27 @@ void cvt_copy(TA* dst, TB* src, size_t n) {
     }
 #    elif defined(HAVE_SVE)
     if constexpr (std::is_same<TA, TB>::value) {
-        SVE_PREDICATE(pg_dst, TA)
-        SVE_VLEN(vlen, TA)
+        auto pg_dst = sve_predicate<sizeof(TA)>();
+        auto vlen = sve_vlen<sizeof(TA)>();
         for (; i + vlen <= n; i += vlen) {
             auto vb = svld1(pg_dst, src + i);
             svst1(pg_dst, dst + i, vb);
         }
-        SVE_PREDICATE_WHILELT(pgt, TA, i, n)
+        auto pgt = sve_predicate<TA, sizeof(TA)>(i, n);
         auto vb = svld1(pg_dst, src + i);
         svst1(pg_dst, dst + i, vb);
         return;
+    } else if constexpr (std::is_same<TA, float>::value && std::is_same<TB, ov::float16>::value) {
+        auto src_ptr = reinterpret_cast<float16_t*>(src);
+        auto pg_vl2 = svwhilelt_b16(svcnth() / 2, svcnth());
+        auto vlen = svcnth() / 2;
+        auto pg_dst = svptrue_b32();
+        for (; i + vlen <= n; i += vlen) {
+            auto load_src = svld1_f16(pg_vl2, src_ptr + i);
+            auto src_interleave = svzip1_f16(load_src, load_src);
+            auto cvt_dst = svcvt_f32_f16_z(pg_dst, src_interleave);
+            svst1(pg_dst, dst + i, cvt_dst);
+        }
     }
 #    endif
     for (; i < n; i++) {
@@ -1715,7 +1726,7 @@ struct MHAHelper {
     void init_reorder_buffers(size_t batch, size_t kv_len_in_blocks) {
         _qk_scratch_b.resize<DATA_TYPE>({batch, kv_len_in_blocks, Hk, _block_size * S});
         if (AarchF16) {
-            // Required to keep kv_cache continuous in mem, as kleidi do to support accumulation
+            // It is required to keep kv_cache continuous in mem, as kleidi do not support accumulation
             _wv_scratch_b.resize<DATA_TYPE>({batch, Hk, kv_len_in_blocks, _block_size * rnd_up(SV, _block_size)});
         } else {
             _wv_scratch_b.resize<DATA_TYPE>({batch, kv_len_in_blocks, Hk, _block_size * rnd_up(SV, _block_size)});
@@ -1918,7 +1929,7 @@ struct MHAHelper {
         auto _score_stride = _weight.stride_bytes(2) / 2;
         for (size_t h = hq_beg; h < hq_end; h++) {
             auto* q_ptr = query.ptr<DATA_TYPE>(h, q_start, 0);
-            float* c_ptr = _weight.ptr<float>(ithr, h, 0, 0);
+            float* c_ptr = _weight.ptr<float>(ithr, h - hq_beg, 0, 0);
             // for each query block, loop through all key block
             // for blocks:
             // 1 0 0 0 ...
@@ -1947,8 +1958,8 @@ struct MHAHelper {
             for (size_t m = q_start; m < q_end; m++) {
                 // apply softmax in f32 precision
                 auto ncausal = (cur_kv_len - q_cnt + (m - q_start) + 1);
-                auto soft_in = _weight.ptr<float>(ithr, h, m - q_start);
-                auto score = _weight.ptr<float>(ithr, h, m - q_start);
+                auto soft_in = _weight.ptr<float>(ithr, h - hq_beg, m - q_start);
+                auto score = _weight.ptr<float>(ithr, h - hq_beg, m - q_start);
                 PlainTensor f32_cvt;
                 if (q_is_xf16) {
                     f32_cvt.resize<float>({size_t{rnd_up(cur_kv_len, _block_size)}});
@@ -2007,7 +2018,7 @@ struct MHAHelper {
             }
 
             // reuse float buffer, need to use float to compute offset
-            auto* w_ptr = reinterpret_cast<DATA_TYPE*>(_weight.ptr<float>(ithr, h, 0, 0));
+            auto* w_ptr = reinterpret_cast<DATA_TYPE*>(_weight.ptr<float>(ithr, h - hq_beg, 0, 0));
             DATA_TYPE* out_ptr = output_emb.ptr<DATA_TYPE>(q_start, h * SV);
             DATA_TYPE* v_ptr;
             v_ptr = wv_scratch_b.ptr<DATA_TYPE>(hk, 0);