[ExecuTorch] FFHT: ARM NEON port

extension/llm/custom_ops/spinquant/third-party/FFHT/fht.h

@@ -1,37 +1,44 @@
 #ifndef _FHT_H_
 #define _FHT_H_
-#include <string.h>
 #include <stdlib.h>
+#include <string.h>
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-int fht_float(float *buf, int log_n);
-int fht_double(double *buf, int log_n);
-int fht_float_oop(float *in, float *out, int log_n);
-int fht_double_oop(double *in, double *out, int log_n);
-
+int fht_float(float* buf, int log_n);
+#ifndef __aarch64__
+int fht_double(double* buf, int log_n);
+#endif
+int fht_float_oop(float* in, float* out, int log_n);
+#ifndef __aarch64__
+int fht_double_oop(double* in, double* out, int log_n);
+#endif
 
 #ifdef __cplusplus
 
 } // extern "C"
 
-static inline int fht(float *buf, int log_n) {
-    return fht_float(buf, log_n);
+static inline int fht(float* buf, int log_n) {
+  return fht_float(buf, log_n);
 }
 
-static inline int fht(double *buf, int log_n) {
-    return fht_double(buf, log_n);
+#ifndef __aarch64__
+static inline int fht(double* buf, int log_n) {
+  return fht_double(buf, log_n);
 }
+#endif
 
-static inline int fht(float *buf, float *out, int log_n) {
-    return fht_float_oop(buf, out, log_n);
+static inline int fht(float* buf, float* out, int log_n) {
+  return fht_float_oop(buf, out, log_n);
 }
 
-static inline int fht(double *buf, double *out, int log_n) {
-    return fht_double_oop(buf, out, log_n);
+#ifndef __aarch64__
+static inline int fht(double* buf, double* out, int log_n) {
+  return fht_double_oop(buf, out, log_n);
 }
+#endif
 
 #endif
 

extension/llm/custom_ops/spinquant/third-party/FFHT/fht_impl.h

@@ -7,23 +7,30 @@
 extern "C" {
 #endif
 
+#ifdef __aarch64__
+#include "fht_neon.c"
+#define VECTOR_WIDTH (16u)
+#else
 #ifdef __AVX__
 #include "fht_avx.c"
 #define VECTOR_WIDTH (32u)
 #else
 #include "fht_sse.c"
 #define VECTOR_WIDTH (16u)
 #endif
+#endif
 
-int fht_float_oop(float *in, float *out, int log_n) {
-    fast_copy(out, in, sizeof(float) << log_n);
-    return fht_float(out, log_n);
+int fht_float_oop(float* in, float* out, int log_n) {
+  fast_copy(out, in, sizeof(float) << log_n);
+  return fht_float(out, log_n);
 }
 
-int fht_double_oop(double *in, double *out, int log_n) {
-    fast_copy(out, in, sizeof(double) << log_n);
-    return fht_double(out, log_n);
+#ifndef __aarch64__
+int fht_double_oop(double* in, double* out, int log_n) {
+  fast_copy(out, in, sizeof(double) << log_n);
+  return fht_double(out, log_n);
 }
+#endif
 
 #ifdef __cplusplus
 } // extern "C"