Example Metal implementation for batch_inverse (wip)

jarnesino · Jan 3, 2025 · 2e73e8d · 2e73e8d
1 parent 36b17ac
commit 2e73e8d
Show file tree

Hide file tree

Showing 7 changed files with 395 additions and 0 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/prover/Cargo.toml b/crates/prover/Cargo.toml
@@ -17,6 +17,7 @@ cfg-if = "1.0.0"
 educe.workspace = true
 hex.workspace = true
 itertools.workspace = true
+metal = "0.30"
 num-traits.workspace = true
 rand = { version = "0.8.5", default-features = false, features = ["small_rng"] }
 starknet-crypto = "0.6.2"

diff --git a/crates/prover/src/core/backend/metal/batch_inverse.air b/crates/prover/src/core/backend/metal/batch_inverse.air
diff --git a/crates/prover/src/core/backend/metal/batch_inverse.metal b/crates/prover/src/core/backend/metal/batch_inverse.metal
@@ -0,0 +1,10 @@
+#include <metal_stdlib>
+
+using namespace metal;
+
+kernel void add_arrays(const device float* in1 [[buffer(0)]],
+                       const device float* in2 [[buffer(1)]],
+                       device float* result [[buffer(2)]],
+                       uint id [[thread_position_in_grid]]) {
+    result[id] = in1[id] + in2[id];
+}
diff --git a/crates/prover/src/core/backend/metal/batch_inverse.metallib b/crates/prover/src/core/backend/metal/batch_inverse.metallib