Skip to content

Commit 961d68f

Browse files
michalharakalmichalharakal
authored
Merge pull request #571 from SKaiNET-developers/feature/native-ffm-kernel-provider
feat(native-cpu): scaffold FFM kernel provider module (PR 1 of 5)
2 parents 90bcf1f + b23bd5a commit 961d68f

11 files changed

Lines changed: 454 additions & 0 deletions

File tree

‎settings.gradle.kts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ include("skainet-compile:skainet-compile-c")
3939
// ====== BACKENDS
4040
include("skainet-backends:skainet-backend-api")
4141
include("skainet-backends:skainet-backend-cpu")
42+
include("skainet-backends:skainet-backend-native-cpu")
4243

4344
// ====== BENCHMARKS
4445
include("skainet-backends:benchmarks:jvm-cpu-jmh")

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

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
plugins {
2+
alias(libs.plugins.kotlinMultiplatform)
3+
}
4+
5+
kotlin {
6+
explicitApi()
7+
jvm()
8+
9+
sourceSets {
10+
val jvmMain by getting {
11+
dependencies {
12+
implementation(project(":skainet-backends:skainet-backend-api"))
13+
}
14+
}
15+
val jvmTest by getting {
16+
dependencies {
17+
implementation(libs.kotlin.test)
18+
}
19+
}
20+
}
21+
}
22+
23+
// --- Native (CMake) wiring -------------------------------------------------
24+
//
25+
// PR 1 builds for the host arch only. Cross-arch CI matrix is deferred per
26+
// the native-ffm-plan asciidoc. The artifact lands at
27+
// build/native/resources/native/<os>-<arch>/libskainet_kernels.{so|dylib|dll}
28+
// and is bundled into the JAR via an extra resources srcDir on jvmMain.
29+
30+
val nativeOsArch: String = run {
31+
val os = System.getProperty("os.name").lowercase()
32+
val osTag = when {
33+
os.contains("linux") -> "linux"
34+
os.contains("mac") || os.contains("darwin") -> "macos"
35+
os.contains("windows") -> "windows"
36+
else -> error("Unsupported OS for skainet-backend-native-cpu: $os")
37+
}
38+
val archRaw = System.getProperty("os.arch").lowercase()
39+
val archTag = when (archRaw) {
40+
"x86_64", "amd64" -> "x86_64"
41+
"aarch64", "arm64" -> "arm64"
42+
else -> error("Unsupported arch for skainet-backend-native-cpu: $archRaw")
43+
}
44+
"$osTag-$archTag"
45+
}
46+
47+
// Capture as plain Strings up front so configuration-cache serialization
48+
// doesn't have to walk back to the build script for Provider resolution.
49+
val nativeSourcePath: String = layout.projectDirectory.dir("native").asFile.absolutePath
50+
val cmakeBuildPath: String = layout.buildDirectory.dir("native/cmake-build").get().asFile.absolutePath
51+
val nativeResourcesRoot = layout.buildDirectory.dir("native/resources")
52+
val nativeResourceTargetDir = nativeResourcesRoot.map { it.dir("native/$nativeOsArch") }
53+
54+
val configureNativeKernels by tasks.registering(Exec::class) {
55+
group = "build"
56+
description = "Run CMake configure step for the native kernels library."
57+
inputs.file("$nativeSourcePath/CMakeLists.txt")
58+
inputs.dir("$nativeSourcePath/src")
59+
inputs.dir("$nativeSourcePath/include")
60+
outputs.dir(cmakeBuildPath)
61+
// CMake auto-creates the -B directory; no doFirst mkdirs needed.
62+
commandLine = listOf(
63+
"cmake",
64+
"-S", nativeSourcePath,
65+
"-B", cmakeBuildPath,
66+
"-DCMAKE_BUILD_TYPE=Release",
67+
)
68+
}
69+
70+
val buildNativeKernels by tasks.registering(Exec::class) {
71+
group = "build"
72+
description = "Build the native kernels shared library via CMake."
73+
dependsOn(configureNativeKernels)
74+
inputs.file("$nativeSourcePath/CMakeLists.txt")
75+
inputs.dir("$nativeSourcePath/src")
76+
inputs.dir("$nativeSourcePath/include")
77+
outputs.dir(cmakeBuildPath)
78+
commandLine = listOf(
79+
"cmake",
80+
"--build", cmakeBuildPath,
81+
"--config", "Release",
82+
)
83+
}
84+
85+
val packageNativeKernels by tasks.registering(Copy::class) {
86+
group = "build"
87+
description = "Stage the built native kernels library into JVM resources."
88+
dependsOn(buildNativeKernels)
89+
from(cmakeBuildPath) {
90+
include(
91+
"libskainet_kernels.so",
92+
"libskainet_kernels.dylib",
93+
"skainet_kernels.dll",
94+
"Release/skainet_kernels.dll",
95+
)
96+
eachFile { path = name }
97+
}
98+
into(nativeResourceTargetDir)
99+
}
100+
101+
kotlin.sourceSets.named("jvmMain") {
102+
resources.srcDir(nativeResourcesRoot)
103+
}
104+
105+
tasks.named("jvmProcessResources") {
106+
dependsOn(packageNativeKernels)
107+
}
108+
109+
tasks.withType<Test>().configureEach {
110+
jvmArgs("--enable-preview", "--enable-native-access=ALL-UNNAMED")
111+
}
112+
113+
tasks.withType<JavaExec>().configureEach {
114+
jvmArgs("--enable-preview", "--enable-native-access=ALL-UNNAMED")
115+
}

‎skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
cmake_minimum_required(VERSION 3.20)
2+
project(skainet_kernels C)
3+
4+
set(CMAKE_C_STANDARD 11)
5+
set(CMAKE_C_STANDARD_REQUIRED ON)
6+
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
7+
8+
if(NOT CMAKE_BUILD_TYPE)
9+
set(CMAKE_BUILD_TYPE Release)
10+
endif()
11+
12+
add_library(skainet_kernels SHARED
13+
src/skainet_smoke.c
14+
)
15+
16+
target_include_directories(skainet_kernels PUBLIC
17+
${CMAKE_CURRENT_SOURCE_DIR}/include
18+
)
19+
20+
# Strip the "lib" prefix on Windows so the artifact name is consistent
21+
# with the resource-bundle path skainet_kernels.{dll,so,dylib}.
22+
if(WIN32)
23+
set_target_properties(skainet_kernels PROPERTIES PREFIX "")
24+
endif()
25+
26+
# Hide non-exported symbols on ELF / Mach-O for a smaller surface area.
27+
if(CMAKE_C_COMPILER_ID MATCHES "Clang|GNU")
28+
target_compile_options(skainet_kernels PRIVATE -fvisibility=hidden -Wall -Wextra)
29+
set_target_properties(skainet_kernels PROPERTIES C_VISIBILITY_PRESET hidden)
30+
endif()

‎skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
#ifndef SKAINET_KERNELS_H
2+
#define SKAINET_KERNELS_H
3+
4+
#include <stdint.h>
5+
6+
#if defined(_WIN32) || defined(__CYGWIN__)
7+
# define SKAINET_API __declspec(dllexport)
8+
#elif defined(__GNUC__) || defined(__clang__)
9+
# define SKAINET_API __attribute__((visibility("default")))
10+
#else
11+
# define SKAINET_API
12+
#endif
13+
14+
#ifdef __cplusplus
15+
extern "C" {
16+
#endif
17+
18+
/*
19+
* Trivial smoke kernel proving the FFM downcall pipeline end-to-end.
20+
*
21+
* for (int i = 0; i < length; ++i) output[i] = 2.0f * input[i];
22+
*
23+
* The Kotlin caller owns the memory backing `input` and `output`; the
24+
* kernel must not retain pointers past return.
25+
*/
26+
SKAINET_API void skainet_smoke_double(const float* input, float* output, int32_t length);
27+
28+
#ifdef __cplusplus
29+
}
30+
#endif
31+
32+
#endif /* SKAINET_KERNELS_H */

‎skainet-backends/skainet-backend-native-cpu/native/src/skainet_smoke.c‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#include "skainet_kernels.h"
2+
3+
void skainet_smoke_double(const float* input, float* output, int32_t length) {
4+
for (int32_t i = 0; i < length; ++i) {
5+
output[i] = 2.0f * input[i];
6+
}
7+
}

‎skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFfmSmoke.kt‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package sk.ainet.exec.kernel
2+
3+
import java.lang.foreign.Arena
4+
import java.lang.foreign.FunctionDescriptor
5+
import java.lang.foreign.Linker
6+
import java.lang.foreign.MemorySegment
7+
import java.lang.foreign.ValueLayout
8+
import java.lang.invoke.MethodHandle
9+
10+
/**
11+
* End-to-end smoke test of the FFM downcall pipeline used by the
12+
* native kernel provider. Calls the bundled native function
13+
*
14+
* void skainet_smoke_double(const float* input, float* output, int32_t length);
15+
*
16+
* which writes `output[i] = 2.0f * input[i]`. This object exists only
17+
* to validate the loader → Linker → MethodHandle path on real hardware
18+
* before any production kernel ships in PR 2.
19+
*
20+
* Not part of the public SPI: `internal` visibility, exposed to tests
21+
* via the same package.
22+
*/
23+
internal object NativeFfmSmoke {
24+
25+
fun isAvailable(): Boolean = handle != null
26+
27+
/**
28+
* Run the bundled smoke kernel on [input] and return a fresh
29+
* `FloatArray` with `output[i] = 2.0f * input[i]`. Returns `null`
30+
* when the native lib failed to load (callers fall back to a
31+
* pure-Kotlin reference).
32+
*/
33+
fun double(input: FloatArray): FloatArray? {
34+
val mh = handle ?: return null
35+
val output = FloatArray(input.size)
36+
val byteSize = input.size.toLong() * java.lang.Float.BYTES
37+
val byteAlign = ValueLayout.JAVA_FLOAT.byteAlignment()
38+
Arena.ofConfined().use { arena ->
39+
val inSeg = arena.allocate(byteSize, byteAlign)
40+
val outSeg = arena.allocate(byteSize, byteAlign)
41+
MemorySegment.copy(input, 0, inSeg, ValueLayout.JAVA_FLOAT, 0L, input.size)
42+
mh.invoke(inSeg, outSeg, input.size)
43+
MemorySegment.copy(outSeg, ValueLayout.JAVA_FLOAT, 0L, output, 0, output.size)
44+
}
45+
return output
46+
}
47+
48+
private val handle: MethodHandle? by lazy {
49+
val lookup = NativeLibraryLoader.lookup() ?: return@lazy null
50+
val symbol = lookup.find("skainet_smoke_double").orElse(null) ?: return@lazy null
51+
val descriptor = FunctionDescriptor.ofVoid(
52+
ValueLayout.ADDRESS,
53+
ValueLayout.ADDRESS,
54+
ValueLayout.JAVA_INT,
55+
)
56+
runCatching { Linker.nativeLinker().downcallHandle(symbol, descriptor) }.getOrNull()
57+
}
58+
}

‎skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package sk.ainet.exec.kernel
2+
3+
import sk.ainet.backend.api.kernel.Fp32MatmulKernel
4+
import sk.ainet.backend.api.kernel.KernelProvider
5+
import sk.ainet.backend.api.kernel.Q4KMatmulKernel
6+
7+
/**
8+
* Native (FFM) [KernelProvider]. Sits at priority `100`, above
9+
* [PanamaVectorKernelProvider] (`50`) and the scalar reference (`0`).
10+
*
11+
* PR 1 of the staged native-FFM rollout (see the `native-ffm-plan`
12+
* asciidoc) only ships the module scaffolding: the Gradle ↔ CMake
13+
* pipeline that produces a host-arch shared library, its bundling into
14+
* JAR resources, and an end-to-end FFM smoke downcall test. No real
15+
* matmul kernel is wired into the public SPI yet.
16+
*
17+
* Until [NativeQ4KMatmulKernel] (or its `MemSegment`-input sibling)
18+
* lands in PR 2, this provider deliberately reports `isAvailable() =
19+
* false` and returns `null` from every kernel accessor. That keeps
20+
* `KernelRegistry.bestAvailable()` cleanly cascading down to the
21+
* Panama priority-50 provider on every shape we measure today, so
22+
* adding the new module to the classpath produces no behavior change.
23+
*/
24+
public object NativeKernelProvider : KernelProvider {
25+
override val name: String = "native-ffm"
26+
override val priority: Int = 100
27+
28+
override fun isAvailable(): Boolean = false
29+
30+
override fun matmulFp32(): Fp32MatmulKernel? = null
31+
32+
override fun matmulQ4K(): Q4KMatmulKernel? = null
33+
}

‎skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProviderFactory.kt‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package sk.ainet.exec.kernel
2+
3+
import sk.ainet.backend.api.kernel.KernelProvider
4+
5+
/**
6+
* `ServiceLoader`-friendly wrapper around [NativeKernelProvider]. The
7+
* platform `ServiceLoader` machinery requires a public no-arg
8+
* constructor, which a Kotlin `object` does not expose; this factory
9+
* delegates every [KernelProvider] member back to the singleton.
10+
*
11+
* Listed in
12+
* `META-INF/services/sk.ainet.backend.api.kernel.KernelProvider` so
13+
* `KernelServiceLoader.installAll()` discovers the provider on JVM
14+
* startup.
15+
*/
16+
public class NativeKernelProviderFactory : KernelProvider by NativeKernelProvider

0 commit comments

Comments
 (0)