Add fast module

frost-beta · Apr 15, 2024 · 369cf27 · 369cf27
1 parent 337afe7
commit 369cf27
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/src/bindings.cc b/src/bindings.cc
@@ -12,6 +12,7 @@ napi_value Init(napi_env env, napi_value exports) {
   InitRandom(env, exports);
   InitFFT(env, exports);
   InitLinalg(env, exports);
+  InitFast(env, exports);
   return exports;
 }
 

diff --git a/src/bindings.h b/src/bindings.h
@@ -15,5 +15,6 @@ void InitIO(napi_env env, napi_value exports);
 void InitRandom(napi_env env, napi_value exports);
 void InitFFT(napi_env env, napi_value exports);
 void InitLinalg(napi_env env, napi_value exports);
+void InitFast(napi_env env, napi_value exports);
 
 #endif  // SRC_BINDINGS_H_
diff --git a/src/fast.cc b/src/fast.cc
@@ -0,0 +1,13 @@
+#include "src/array.h"
+#include "src/stream.h"
+
+void InitFast(napi_env env, napi_value exports) {
+  napi_value fast = ki::CreateObject(env);
+  ki::Set(env, exports, "fast", fast);
+
+  ki::Set(env, fast,
+          "rmsNorm", &mx::fast::rms_norm,
+          "layerNorm", &mx::fast::layer_norm,
+          "rope", &mx::fast::rope,
+          "scaledDotProductAttention", &mx::fast::scaled_dot_product_attention);
+}