Make kernels extern, use stdlib for fns

EricLBuehler · Mar 12, 2024 · f31db99 · f31db99
1 parent cb5f4fa
commit f31db99
Showing 1 changed file with 5 additions and 16 deletions.
diff --git a/candle-kernels/src/layernorm.cu b/candle-kernels/src/layernorm.cu
@@ -1,18 +1,7 @@
-#define C10_WARP_SIZE 32
+#include <stdint.h>
 
-float max(float x, float y) {
-  return ::fmaxf(x, y);
-}
-double max(double x, double y) {
-  return ::fmax(x, y);
-}
+#define C10_WARP_SIZE 32
 
-float rsqrt(float x) {
-  return ::rsqrtf(x);
-}
-double rsqrt(double x) {
-  return ::rsqrt(x);
-}
 
 struct Block1D {
     static __forceinline__ __device__ int Tid() { return threadIdx.x; }
@@ -67,7 +56,7 @@ __inline__ __device__ T BlockReduceSum(T val, T* shared) {
 }
 
 #define ROWWISEMOMENTS(TYPENAME, FN_NAME) \
-__global__ void FN_NAME( \
+extern "C"  __global__ void FN_NAME( \
     int64_t N, \
     double eps, \
     const TYPENAME* X, \
@@ -95,7 +84,7 @@ __global__ void FN_NAME( \
 }
 
 #define LAYERNORM_BIAS(TYPENAME, FN_NAME) \
-__global__ void FN_NAME( \
+extern "C" __global__ void FN_NAME( \
     int64_t N, \
     const TYPENAME* X, \
     const TYPENAME* mean, \
@@ -116,7 +105,7 @@ __global__ void FN_NAME( \
 }
 
 #define LAYERNORM(TYPENAME, FN_NAME) \
-__global__ void FN_NAME( \
+extern "C" __global__ void FN_NAME( \
     int64_t N, \
     const TYPENAME* X, \
     const TYPENAME* mean, \