Fix compilation errors on aarch64

Author: turboderp

exllamav2/exllamav2_ext/cpp/sampling_avx2.cpp

@@ -12,8 +12,7 @@
 #include "profiling.h"
 
 #ifdef USE_AVX2
-    #include "avx_mathfun.h"
-#endif
+#include "avx_mathfun.h"
 
 AVX2_TARGET
 int softmax_cpu_avx2
@@ -26,10 +25,6 @@ int softmax_cpu_avx2
     float* output
 )
 {
-    #ifndef USE_AVX2
-    return 0;
-    #else
-
     profile_start("softmax_cpu (AVX2)");
 
     int vocab_size_aligned = ((vocab_size + 31) / 32) * 32;
@@ -148,5 +143,24 @@ int softmax_cpu_avx2
 //    }
 //    printf("sum: %f\n\n", summ);
     return maxi;
-    #endif
+
 }
+
+#else
+
+// Dummy function to avoid compilation errors on aarch64
+
+int softmax_cpu_avx2
+(
+    const int vocab_size,
+    const float temperature,
+    const float* logits,
+    const bool* logits_filter,
+    const float exponent,
+    float* output
+)
+{
+    return 0;
+}
+
+#endif

exllamav2/exllamav2_ext/ext_sampling.cpp

@@ -8,14 +8,22 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
-#ifdef __linux__
-#include <mm_malloc.h>
+#include "config.h"
+
+#ifdef USE_AVX2
+    #ifdef __linux__
+        #include <mm_malloc.h>
+    #else
+        #define _mm_malloc(a, b) _aligned_malloc(a, b)
+        #define _mm_free(a) _aligned_free(a)
+    #endif
 #else
-#define _mm_malloc(a, b) _aligned_malloc(a, b)
-#define _mm_free(a) _aligned_free(a)
+    // aarch64-linux apparently has no mm_malloc, but it would only be needed for AVX2 anyway, which is also not
+    // available on aarch64
+    #define _mm_malloc(a, b) malloc(a)
+    #define _mm_free(a) free(a)
 #endif
 
-#include "config.h"
 #include "ext_sampling.h"
 
 #include "cpp/generator.h"