Add sve implementation for float matrix transpose (#3421)

Summary:

X-link: facebookresearch/FBGEMM#509

Adding sve-based function for transposing float matrixes

Differential Revision: D66528598

Author: Nicoshev

defs.bzl

@@ -143,10 +143,10 @@ def get_fbgemm_inline_avx512_srcs(msvc = False, buck = False):
     return asm_srcs if not msvc else intrinsics_srcs
 
 def get_fbgemm_inline_sve_srcs(msvc = False, buck = False):
-    intrinsics_srcs = ["src/FbgemmFP16UKernelsSve128.cc"]
+    intrinsics_srcs = ["src/FbgemmFP16UKernelsSve128.cc", "src/UtilsSve.cc"]
 
     #FP16 kernels contain inline assembly and inline assembly syntax for MSVC is different.
-    asm_srcs = ["src/FbgemmFP16UKernelsSve128.cc"]
+    asm_srcs = ["src/FbgemmFP16UKernelsSve128.cc", "src/UtilsSve.cc"]
     if buck:
         return select({
             "DEFAULT": asm_srcs,

src/TransposeUtils.cc

@@ -46,6 +46,14 @@ void transpose_simd(
     }
     return;
   }
+
+#ifdef __aarch64__
+  if constexpr (std::is_same<T, float>::value) {
+    internal::transpose_sve<T>(M, N, src, ld_src, dst, ld_dst);
+  } else {
+    transpose_ref<T>(M, N, src, ld_src, dst, ld_dst);
+  }
+#else
   static const auto iset = fbgemmInstructionSet();
   // Run time CPU detection
   if (isZmm(iset)) {
@@ -55,6 +63,8 @@ void transpose_simd(
   } else {
     transpose_ref<T>(M, N, src, ld_src, dst, ld_dst);
   }
+
+#endif
 }
 
 template void transpose_ref<float>(

src/TransposeUtils.h

@@ -62,6 +62,22 @@ void transpose_avx512(
     T* dst,
     int64_t ld_dst);
 
+#ifdef __aarch64__
+/**
+ * @brief Transpose a matrix using Intel AVX2.
+ *
+ * This is called if the code is running on a CPU with Intel AVX2 support.
+ */
+template <typename T>
+void transpose_sve(
+    int64_t M,
+    int64_t N,
+    const T* src,
+    int64_t ld_src,
+    T* dst,
+    int64_t ld_dst);
+#endif // __aarch64__
+
 } // namespace internal
 
 } // namespace fbgemm