Skip to content

Commit 6dbd889

Browse files
michalharakalclaude
michalharakal
and
claude
committed
Add Accelerate-backed CPU ops for native macOS/iOS
Replace scalar DefaultCpuOps on Apple platforms with AccelerateCpuOps that dispatches hot-path operations to Apple's Accelerate framework: - matmul: cblas_sgemm (NEON + AMX hardware acceleration) - add/subtract/multiply/divide: vDSP_vadd/vsub/vmul/vdiv - sum/mean: vDSP_sve/meanv - relu: vDSP_vthres - silu: optimized scalar loop on contiguous buffer - transpose: vDSP_mtrans All operations fall through to DefaultCpuOpsBase for non-FP32, non-contiguous, or complex broadcasting cases. Split PlatformCpuOpsFactory from single nativeMain actual into: - appleMain (macOS + iOS) → AccelerateCpuOps - linuxMain → DefaultCpuOps (scalar fallback) All 23 existing test suites pass on macosArm64 with zero failures. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 6e7422c commit 6dbd889

4 files changed

Lines changed: 359 additions & 2 deletions

File tree

‎skainet-backends/skainet-backend-cpu/build.gradle.kts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,16 +70,20 @@ kotlin {
7070
dependsOn(commonMain)
7171
}
7272

73+
val appleMain by creating {
74+
dependsOn(nativeMain)
75+
}
76+
7377
val linuxMain by creating {
7478
dependsOn(nativeMain)
7579
}
7680

7781
val iosMain by creating {
78-
dependsOn(nativeMain)
82+
dependsOn(appleMain)
7983
}
8084

8185
val macosMain by creating {
82-
dependsOn(nativeMain)
86+
dependsOn(appleMain)
8387
}
8488

8589
val iosArm64Main by getting {

‎skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/AccelerateCpuOps.kt‎

Lines changed: 344 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,344 @@
1+
@file:OptIn(kotlinx.cinterop.ExperimentalForeignApi::class)
2+
3+
package sk.ainet.exec.tensor.ops
4+
5+
import kotlinx.cinterop.addressOf
6+
import kotlinx.cinterop.usePinned
7+
import platform.Accelerate.CblasNoTrans
8+
import platform.Accelerate.CblasRowMajor
9+
import platform.Accelerate.cblas_sgemm
10+
import platform.Accelerate.vDSP_vadd
11+
import platform.Accelerate.vDSP_vsub
12+
import platform.Accelerate.vDSP_vmul
13+
import platform.Accelerate.vDSP_vdiv
14+
import platform.Accelerate.vDSP_sve
15+
import platform.Accelerate.vDSP_meanv
16+
import platform.Accelerate.vDSP_mtrans
17+
import platform.Accelerate.vDSP_vthres
18+
import sk.ainet.lang.tensor.Shape
19+
import sk.ainet.lang.tensor.Tensor
20+
import sk.ainet.lang.tensor.data.FloatArrayTensorData
21+
import sk.ainet.lang.tensor.data.TensorDataFactory
22+
import sk.ainet.lang.types.DType
23+
import sk.ainet.lang.types.FP32
24+
25+
/**
26+
* CPU operations accelerated by Apple's Accelerate framework.
27+
* Overrides hot-path operations (matmul, elementwise, reductions) with
28+
* hardware-optimized routines that leverage ARM NEON and AMX.
29+
*
30+
* Falls through to [DefaultCpuOpsBase] for non-FP32, non-contiguous,
31+
* or complex broadcasting cases.
32+
*/
33+
public class AccelerateCpuOps(
34+
dataFactory: TensorDataFactory,
35+
) : DefaultCpuOpsBase(dataFactory) {
36+
37+
// ── matmul ──────────────────────────────────────────────────────────
38+
39+
override fun <T : DType, V> matmul(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
40+
if (a.rank == 2 && b.rank == 2
41+
&& a.dtype == FP32::class
42+
&& a.data is FloatArrayTensorData<*>
43+
&& b.data is FloatArrayTensorData<*>
44+
) {
45+
val aBuf = (a.data as FloatArrayTensorData<*>).buffer
46+
val bBuf = (b.data as FloatArrayTensorData<*>).buffer
47+
val m = a.shape[0]
48+
val k = a.shape[1]
49+
val n = b.shape[1]
50+
require(k == b.shape[0]) { "matmul shape mismatch: ${a.shape} vs ${b.shape}" }
51+
52+
val out = FloatArray(m * n)
53+
// cblas_sgemm: C = alpha * A * B + beta * C
54+
aBuf.usePinned { aPin ->
55+
bBuf.usePinned { bPin ->
56+
out.usePinned { cPin ->
57+
cblas_sgemm(
58+
CblasRowMajor,
59+
CblasNoTrans, CblasNoTrans,
60+
m, n, k,
61+
1.0f, // alpha
62+
aPin.addressOf(0), k, // A, lda
63+
bPin.addressOf(0), n, // B, ldb
64+
0.0f, // beta
65+
cPin.addressOf(0), n, // C, ldc
66+
)
67+
}
68+
}
69+
}
70+
71+
@Suppress("UNCHECKED_CAST")
72+
val outData = dataFactory.fromFloatArray<T, Float>(Shape(m, n), a.dtype, out)
73+
as sk.ainet.lang.tensor.data.TensorData<T, V>
74+
return newTensor(outData, a.dtype, a, b)
75+
}
76+
77+
return super.matmul(a, b)
78+
}
79+
80+
// ── elementwise binary ops ──────────────────────────────────────────
81+
82+
override fun <T : DType, V> add(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
83+
val result = tryVdspBinary(a, b, ::vdspAdd)
84+
return result ?: super.add(a, b)
85+
}
86+
87+
override fun <T : DType, V> subtract(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
88+
val result = tryVdspBinary(a, b, ::vdspSub)
89+
return result ?: super.subtract(a, b)
90+
}
91+
92+
override fun <T : DType, V> multiply(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
93+
val result = tryVdspBinary(a, b, ::vdspMul)
94+
return result ?: super.multiply(a, b)
95+
}
96+
97+
override fun <T : DType, V> divide(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> {
98+
val result = tryVdspBinary(a, b, ::vdspDiv)
99+
return result ?: super.divide(a, b)
100+
}
101+
102+
// ── reductions ──────────────────────────────────────────────────────
103+
104+
override fun <T : DType, V> sum(tensor: Tensor<T, V>, dim: Int?): Tensor<T, V> {
105+
if (dim == null
106+
&& tensor.dtype == FP32::class
107+
&& tensor.data is FloatArrayTensorData<*>
108+
) {
109+
val buf = (tensor.data as FloatArrayTensorData<*>).buffer
110+
val n = buf.size
111+
if (n > 0) {
112+
val result = FloatArray(1)
113+
buf.usePinned { pin ->
114+
result.usePinned { rPin ->
115+
vDSP_sve(pin.addressOf(0), 1, rPin.addressOf(0), n.toULong())
116+
}
117+
}
118+
@Suppress("UNCHECKED_CAST")
119+
val outData = dataFactory.fromFloatArray<T, Float>(Shape(), tensor.dtype, floatArrayOf(result[0]))
120+
as sk.ainet.lang.tensor.data.TensorData<T, V>
121+
return newTensor(outData, tensor.dtype, tensor)
122+
}
123+
}
124+
return super.sum(tensor, dim)
125+
}
126+
127+
override fun <T : DType, V> mean(tensor: Tensor<T, V>, dim: Int?): Tensor<T, V> {
128+
if (dim == null
129+
&& tensor.dtype == FP32::class
130+
&& tensor.data is FloatArrayTensorData<*>
131+
) {
132+
val buf = (tensor.data as FloatArrayTensorData<*>).buffer
133+
val n = buf.size
134+
if (n > 0) {
135+
val result = FloatArray(1)
136+
buf.usePinned { pin ->
137+
result.usePinned { rPin ->
138+
vDSP_meanv(pin.addressOf(0), 1, rPin.addressOf(0), n.toULong())
139+
}
140+
}
141+
@Suppress("UNCHECKED_CAST")
142+
val outData = dataFactory.fromFloatArray<T, Float>(Shape(), tensor.dtype, floatArrayOf(result[0]))
143+
as sk.ainet.lang.tensor.data.TensorData<T, V>
144+
return newTensor(outData, tensor.dtype, tensor)
145+
}
146+
}
147+
return super.mean(tensor, dim)
148+
}
149+
150+
// ── activations ─────────────────────────────────────────────────────
151+
152+
override fun <T : DType, V> relu(tensor: Tensor<T, V>): Tensor<T, V> {
153+
if (tensor.dtype == FP32::class && tensor.data is FloatArrayTensorData<*>) {
154+
val buf = (tensor.data as FloatArrayTensorData<*>).buffer
155+
val n = buf.size
156+
val out = FloatArray(n)
157+
buf.usePinned { pin ->
158+
out.usePinned { oPin ->
159+
val threshold = FloatArray(1) { 0.0f }
160+
threshold.usePinned { tPin ->
161+
vDSP_vthres(pin.addressOf(0), 1, tPin.addressOf(0), oPin.addressOf(0), 1, n.toULong())
162+
}
163+
}
164+
}
165+
@Suppress("UNCHECKED_CAST")
166+
val outData = dataFactory.fromFloatArray<T, Float>(tensor.shape, tensor.dtype, out)
167+
as sk.ainet.lang.tensor.data.TensorData<T, V>
168+
return newTensor(outData, tensor.dtype, tensor)
169+
}
170+
return super.relu(tensor)
171+
}
172+
173+
override fun <T : DType, V> silu(tensor: Tensor<T, V>): Tensor<T, V> {
174+
if (tensor.dtype == FP32::class && tensor.data is FloatArrayTensorData<*>) {
175+
val buf = (tensor.data as FloatArrayTensorData<*>).buffer
176+
val n = buf.size
177+
val out = FloatArray(n)
178+
for (i in 0 until n) {
179+
val x = buf[i]
180+
out[i] = x / (1.0f + kotlin.math.exp(-x))
181+
}
182+
@Suppress("UNCHECKED_CAST")
183+
val outData = dataFactory.fromFloatArray<T, Float>(tensor.shape, tensor.dtype, out)
184+
as sk.ainet.lang.tensor.data.TensorData<T, V>
185+
return newTensor(outData, tensor.dtype, tensor)
186+
}
187+
return super.silu(tensor)
188+
}
189+
190+
// ── transpose ───────────────────────────────────────────────────────
191+
192+
override fun <T : DType, V> transpose(tensor: Tensor<T, V>): Tensor<T, V> {
193+
if (tensor.rank == 2
194+
&& tensor.dtype == FP32::class
195+
&& tensor.data is FloatArrayTensorData<*>
196+
) {
197+
val buf = (tensor.data as FloatArrayTensorData<*>).buffer
198+
val rows = tensor.shape[0]
199+
val cols = tensor.shape[1]
200+
val out = FloatArray(rows * cols)
201+
buf.usePinned { pin ->
202+
out.usePinned { oPin ->
203+
vDSP_mtrans(
204+
pin.addressOf(0), 1,
205+
oPin.addressOf(0), 1,
206+
cols.toULong(), rows.toULong(),
207+
)
208+
}
209+
}
210+
@Suppress("UNCHECKED_CAST")
211+
val outData = dataFactory.fromFloatArray<T, Float>(Shape(cols, rows), tensor.dtype, out)
212+
as sk.ainet.lang.tensor.data.TensorData<T, V>
213+
return newTensor(outData, tensor.dtype, tensor)
214+
}
215+
return super.transpose(tensor)
216+
}
217+
218+
// ── vDSP binary helpers ─────────────────────────────────────────────
219+
220+
/**
221+
* Attempt to dispatch a binary elementwise op to vDSP.
222+
* Returns null if the tensors are not eligible (non-FP32, non-contiguous,
223+
* complex broadcasting).
224+
*/
225+
private fun <T : DType, V> tryVdspBinary(
226+
a: Tensor<T, V>,
227+
b: Tensor<T, V>,
228+
op: (FloatArray, FloatArray, FloatArray, Int) -> Unit,
229+
): Tensor<T, V>? {
230+
if (a.dtype != FP32::class) return null
231+
if (a.data !is FloatArrayTensorData<*> || b.data !is FloatArrayTensorData<*>) return null
232+
233+
val aBuf = (a.data as FloatArrayTensorData<*>).buffer
234+
val bBuf = (b.data as FloatArrayTensorData<*>).buffer
235+
236+
// Same shape: straightforward vectorized op
237+
if (a.shape == b.shape) {
238+
val n = aBuf.size
239+
val out = FloatArray(n)
240+
op(aBuf, bBuf, out, n)
241+
@Suppress("UNCHECKED_CAST")
242+
val outData = dataFactory.fromFloatArray<T, Float>(a.shape, a.dtype, out)
243+
as sk.ainet.lang.tensor.data.TensorData<T, V>
244+
return newTensor(outData, a.dtype, a, b)
245+
}
246+
247+
// Scalar broadcast: b is a single element
248+
if (bBuf.size == 1) {
249+
val n = aBuf.size
250+
val expanded = FloatArray(n) { bBuf[0] }
251+
val out = FloatArray(n)
252+
op(aBuf, expanded, out, n)
253+
@Suppress("UNCHECKED_CAST")
254+
val outData = dataFactory.fromFloatArray<T, Float>(a.shape, a.dtype, out)
255+
as sk.ainet.lang.tensor.data.TensorData<T, V>
256+
return newTensor(outData, a.dtype, a, b)
257+
}
258+
259+
// Scalar broadcast: a is a single element
260+
if (aBuf.size == 1) {
261+
val n = bBuf.size
262+
val expanded = FloatArray(n) { aBuf[0] }
263+
val out = FloatArray(n)
264+
op(expanded, bBuf, out, n)
265+
@Suppress("UNCHECKED_CAST")
266+
val outData = dataFactory.fromFloatArray<T, Float>(b.shape, a.dtype, out)
267+
as sk.ainet.lang.tensor.data.TensorData<T, V>
268+
return newTensor(outData, a.dtype, a, b)
269+
}
270+
271+
// Last-dim broadcast: b has shape [1, ..., 1, N] matching a's last dim
272+
// Common for bias add: [batch, features] + [features]
273+
if (b.rank <= a.rank) {
274+
val bDims = b.shape.dimensions
275+
val aDims = a.shape.dimensions
276+
val offset = aDims.size - bDims.size
277+
var isBiasBroadcast = true
278+
for (i in bDims.indices) {
279+
if (i < bDims.size - 1 && bDims[i] != 1) { isBiasBroadcast = false; break }
280+
if (i == bDims.size - 1 && bDims[i] != aDims[offset + i]) { isBiasBroadcast = false; break }
281+
}
282+
if (isBiasBroadcast && bDims.last() > 1) {
283+
val lastDim = bDims.last()
284+
val batches = aBuf.size / lastDim
285+
val out = FloatArray(aBuf.size)
286+
for (batch in 0 until batches) {
287+
val aSlice = FloatArray(lastDim)
288+
aBuf.copyInto(aSlice, 0, batch * lastDim, (batch + 1) * lastDim)
289+
val oSlice = FloatArray(lastDim)
290+
op(aSlice, bBuf, oSlice, lastDim)
291+
oSlice.copyInto(out, batch * lastDim)
292+
}
293+
@Suppress("UNCHECKED_CAST")
294+
val outData = dataFactory.fromFloatArray<T, Float>(a.shape, a.dtype, out)
295+
as sk.ainet.lang.tensor.data.TensorData<T, V>
296+
return newTensor(outData, a.dtype, a, b)
297+
}
298+
}
299+
300+
return null // fall through to scalar
301+
}
302+
303+
private fun vdspAdd(a: FloatArray, b: FloatArray, out: FloatArray, n: Int) {
304+
a.usePinned { aPin ->
305+
b.usePinned { bPin ->
306+
out.usePinned { oPin ->
307+
vDSP_vadd(aPin.addressOf(0), 1, bPin.addressOf(0), 1, oPin.addressOf(0), 1, n.toULong())
308+
}
309+
}
310+
}
311+
}
312+
313+
private fun vdspSub(a: FloatArray, b: FloatArray, out: FloatArray, n: Int) {
314+
// vDSP_vsub computes out = B - A (reversed!), so swap args
315+
a.usePinned { aPin ->
316+
b.usePinned { bPin ->
317+
out.usePinned { oPin ->
318+
vDSP_vsub(bPin.addressOf(0), 1, aPin.addressOf(0), 1, oPin.addressOf(0), 1, n.toULong())
319+
}
320+
}
321+
}
322+
}
323+
324+
private fun vdspMul(a: FloatArray, b: FloatArray, out: FloatArray, n: Int) {
325+
a.usePinned { aPin ->
326+
b.usePinned { bPin ->
327+
out.usePinned { oPin ->
328+
vDSP_vmul(aPin.addressOf(0), 1, bPin.addressOf(0), 1, oPin.addressOf(0), 1, n.toULong())
329+
}
330+
}
331+
}
332+
}
333+
334+
private fun vdspDiv(a: FloatArray, b: FloatArray, out: FloatArray, n: Int) {
335+
// vDSP_vdiv computes out = B / A (reversed!), so swap args
336+
a.usePinned { aPin ->
337+
b.usePinned { bPin ->
338+
out.usePinned { oPin ->
339+
vDSP_vdiv(bPin.addressOf(0), 1, aPin.addressOf(0), 1, oPin.addressOf(0), 1, n.toULong())
340+
}
341+
}
342+
}
343+
}
344+
}

‎skainet-backends/skainet-backend-cpu/src/appleMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.apple.kt‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package sk.ainet.exec.tensor.ops
2+
3+
import sk.ainet.lang.tensor.data.TensorDataFactory
4+
import sk.ainet.lang.tensor.ops.TensorOps
5+
6+
internal actual fun platformDefaultCpuOpsFactory(): (TensorDataFactory) -> TensorOps {
7+
println("[SKaiNET] Using Accelerate-backed CPU operations (ARM NEON + AMX)")
8+
return { factory -> AccelerateCpuOps(factory) }
9+
}

skainet-backends/skainet-backend-cpu/src/nativeMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.native.kt renamed to skainet-backends/skainet-backend-cpu/src/linuxMain/kotlin/sk/ainet/exec/tensor/ops/PlatformCpuOpsFactory.linux.kt

File renamed without changes.

0 commit comments

Comments
 (0)