FFHT: ARM NEON port (#5289)

Summary:
Pull Request resolved: #5289

Patch the code generator to be capable of generating NEON
code and leave it configured to do that since we already have the
checked-in generated AVX and SSE code. Generated code size was
 a potential issue so I also patched the generator to 1) reuse
generated code for previous smaller sizes whereapplicable and 2)
 choose the smallest code that isn't more than 10%
slower than the very fastest code.
ghstack-source-id: 242230777
exported-using-ghexport

Reviewed By: kimishpatel

Differential Revision: D60194970

fbshipit-source-id: 37aab6813222c5a965c060286b5d5453ced22a0c

Author: swolchok

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"