diff --git a/CHANGELOG.md b/CHANGELOG.md index e1eedecc..1c734451 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- **`IreeModuleRunner` + `libskainet_iree_module.so`** (`runtime-iree-android`, `arm64-v8a` / `armeabi-v7a`, + CPU and Vulkan drivers): a generic single-shot runner for one compiled IREE module — one float32 tensor in, + one float32 tensor out, weights from an optional `.irpa` (scope `model`). For batch cartridges such as the + `detection-yolov8-iree` blueprint, which extracts the `.so` from this AAR. `FloatTensor` carries data + shape. + ## [0.63.0] — 2026-10-21 Against engine **0.62.0** (the JVM ↔ Kotlin/Wasm parity release: every kernel on the wasm tier, resident diff --git a/llm-runtime/iree-android/README.md b/llm-runtime/iree-android/README.md index f385deb9..c9b64ac5 100644 --- a/llm-runtime/iree-android/README.md +++ b/llm-runtime/iree-android/README.md @@ -205,3 +205,19 @@ the M = 1 matrix-vector products); encoder and prefill do — hence the CPU deco `native/build-whisper-session.sh --vulkan` builds the library like the others. `native/whisper_session_check.c` drives it on a device without an app (a minimal `JNIEnv` over `dlopen`) — how the numbers above were taken. + +## Single-shot modules (`IreeModuleRunner`, `libskainet_iree_module.so`) + +For batch graphs that are a plain tensor function — one float32 tensor in, one float32 tensor out, static +shapes: a detector with its decoded head, a classifier, an embedding graph. Nothing model-specific: the vmfb, +the optional `.irpa` (scope `model`; `null` when the vmfb carries its weights) and the function name are the +caller's. + +```kotlin +IreeModuleRunner("$dir/yolov8n.vmfb", "$dir/params.irpa", "yolov8_detect", IreeModuleRunner.VULKAN_DEVICE).use { runner -> + val head: FloatTensor = runner.run(FloatTensor(pixels, intArrayOf(1, 3, 640, 640))) // head.shape = [1, 84, 8400] +} +``` + +Failures throw, with the IREE status in logcat under `skainet_iree`. One `run` at a time per instance. +`native/build-iree-module.sh --vulkan` builds it like the redecode `.so` (same image and links). diff --git a/llm-runtime/iree-android/build.gradle.kts b/llm-runtime/iree-android/build.gradle.kts index 815c2310..15dbef98 100644 --- a/llm-runtime/iree-android/build.gradle.kts +++ b/llm-runtime/iree-android/build.gradle.kts @@ -23,6 +23,8 @@ android { defaultConfig { minSdk = libs.versions.android.minSdk.get().toInt() + // On-device check over served module cases (src/androidTest, IreeModuleRunnerDeviceTest). + testInstrumentationRunner = "androidx.test.runner.AndroidJUnitRunner" } androidResources { @@ -51,4 +53,6 @@ dependencies { testImplementation(libs.kotlin.test) testImplementation(libs.junit.jupiter) testRuntimeOnly(libs.junit.platform.launcher) + androidTestImplementation(libs.androidx.test.junit) + androidTestImplementation(libs.androidx.test.runner) } diff --git a/llm-runtime/iree-android/native/build-iree-module.sh b/llm-runtime/iree-android/native/build-iree-module.sh new file mode 100755 index 00000000..90562683 --- /dev/null +++ b/llm-runtime/iree-android/native/build-iree-module.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# Cross-build libskainet_iree_module.so for an Android ABI using the consolidated +# skainet/iree-android image from SKaiNET-iree-toolchain (IREE v3.11 runtime tree + NDK +# r27c; build it there with `make android`). The build-so subcommand injects the CMake +# target, configures the local-task/local-sync CPU drivers (+ Vulkan when --vulkan is +# passed), builds and strips. +# +# This is the GENERIC single-shot IREE module runner (llm-runtime/iree-android, IreeModuleRunner): +# one float32 tensor in, one out. It knows nothing +# about any specific model. Weights for whatever vmfb it's pointed at are EXTERNAL (an +# `.irpa` loaded at session-create time via the io_parameters VM module), so the `.so` +# needs the io_parameters VM module and its iree/io dependencies explicitly linked — the +# default `iree_runtime_unified` target alone does not pull them in. This requirement is a +# property of the runtime (any external-weights vmfb needs it), not of any one model. +# +# ./build-iree-module.sh [armeabi-v7a|arm64-v8a] [--vulkan] [cache-dir] +# +# Cache dir (default $IREE_CACHE, or ~/.cache/skainet-iree/) warms /iree/build-host +# and /iree/build- across runs and ABIs. +# Output: out/libskainet_iree_module.so -> copy to +# ../src/main/jniLibs//libskainet_iree_module.so +# +# The flags + task-api links back nativeCreateWithTopology (SKEEP-005 phase 2): the JNI sets +# --task_topology_group_count through IREE's flag parser before the local-task device is +# created. A `.so` built without this source predates the symbol; the Kotlin side fails +# loudly (UnsatisfiedLinkError → IllegalStateException) rather than ignoring the knob. +set -euo pipefail +ABI="${1:-arm64-v8a}" +VULKAN="" +if [ "${2:-}" = "--vulkan" ]; then VULKAN="--vulkan"; shift; fi +CACHE="${2:-${IREE_CACHE:-$HOME/.cache/skainet-iree/3.11.0}}" +HERE="$(cd "$(dirname "$0")" && pwd)" +mkdir -p "$HERE/out" "$CACHE/build-host" "$CACHE/build-$ABI" + +docker run --rm \ + -v "$HERE":/src:ro -v "$HERE/out":/out \ + -v "$CACHE/build-host":/iree/build-host \ + -v "$CACHE/build-$ABI":/iree/build-$ABI \ + skainet/iree-android:3.11.0 \ + build-so "$ABI" --name skainet_iree_module --src /src/iree_module_jni.c $VULKAN \ + --link iree_modules_io_parameters_parameters \ + --link iree_io_parameter_index \ + --link iree_io_parameter_index_provider \ + --link iree_io_formats_irpa_irpa \ + --link iree_base_tooling_flags \ + --link iree_task_api diff --git a/llm-runtime/iree-android/native/iree_module_jni.c b/llm-runtime/iree-android/native/iree_module_jni.c new file mode 100644 index 00000000..14340a44 --- /dev/null +++ b/llm-runtime/iree-android/native/iree_module_jni.c @@ -0,0 +1,228 @@ +/* iree_module_jni.c — generic single-shot runner for one compiled IREE module. + * + * Drives ANY vmfb whose entry function takes one float32 tensor and returns one float32 tensor + * (`tensor -> tensor`, static shapes): a vision backbone with its head, a + * classifier, an embedding graph. Weights are either baked into the vmfb or EXTERNAL (scope + * "model") — bound at create time from a `.irpa` parameter archive via the io_parameters VM + * module, exactly as in iree_redecode_jni.c. The vmfb path, the optional `.irpa` path and the + * function name are caller-supplied; this file knows nothing about any specific model. + * + * Exposes to sk.ainet.transformers.iree.android.IreeModuleRunner: + * long nativeCreate(String device, String vmfbPath, String irpaPath, String functionName, + * int taskTopologyGroupCount) + * irpaPath "" = weights baked into the vmfb; group count <= 0 = IREE's own topology + * float[] nativeRun(long h, float[] input, int[] inputShape, int[] outputShape) + * outputShape (length >= 1 + MAX_RANK) receives [rank, d0, d1, …] of the result + * void nativeDestroy(long h) + * + * Every failure is logged under the "skainet_iree" tag and returns 0 / NULL (see #404). + */ +#include +#include +#include +#include +#include +#include "iree/base/tooling/flags.h" +#include "iree/runtime/api.h" +#include "iree/io/file_handle.h" +#include "iree/io/parameter_index.h" +#include "iree/io/parameter_index_provider.h" +#include "iree/io/formats/irpa/irpa_parser.h" +#include "iree/modules/io/parameters/module.h" +#include + +#define TAG "skainet_iree" +#define PARAM_SCOPE "model" +#define MAX_CONCURRENT_PARAM_OPS 16 +#define MAX_RANK 8 + +static void log_and_ignore(const char* what, iree_status_t st) { + if (iree_status_is_ok(st)) return; + char* buf = NULL; iree_host_size_t len = 0; + iree_allocator_t alloc = iree_allocator_system(); + if (iree_status_to_string(st, &alloc, &buf, &len) && buf) { + __android_log_print(ANDROID_LOG_ERROR, TAG, "%s: %.*s", what, (int)len, buf); + iree_allocator_free(alloc, buf); + } else { + __android_log_print(ANDROID_LOG_ERROR, TAG, "%s: status code %d", what, (int)iree_status_code(st)); + } + iree_status_ignore(st); +} + +typedef struct { + iree_runtime_instance_t* inst; + iree_hal_device_t* dev; + iree_runtime_session_t* sess; + iree_io_parameter_provider_t* param_provider; + char* fn_name; +} Runner; + +#define JNIFN(name) Java_sk_ainet_transformers_iree_android_IreeModuleRunner_##name + +/* The `.irpa` as a "model"-scoped parameter provider, appended BEFORE the bytecode module so its + * util.global initializers resolve at link time. Same as iree_redecode_jni.c. */ +static iree_status_t append_parameters_module(Runner* r, const char* irpa) { + iree_allocator_t alloc = iree_allocator_system(); + iree_io_file_handle_t* file = NULL; + IREE_RETURN_IF_ERROR(iree_io_file_handle_open( + IREE_IO_FILE_MODE_READ, iree_make_cstring_view(irpa), alloc, &file)); + + iree_io_parameter_index_t* index = NULL; + iree_status_t st = iree_io_parameter_index_create(alloc, &index); + if (iree_status_is_ok(st)) st = iree_io_parse_irpa_index(file, index, alloc); + iree_io_file_handle_release(file); + if (!iree_status_is_ok(st)) { if (index) iree_io_parameter_index_release(index); return st; } + + iree_io_parameter_provider_t* provider = NULL; + st = iree_io_parameter_index_provider_create( + iree_make_cstring_view(PARAM_SCOPE), index, MAX_CONCURRENT_PARAM_OPS, alloc, &provider); + iree_io_parameter_index_release(index); + if (!iree_status_is_ok(st)) return st; + + iree_vm_module_t* params_module = NULL; + st = iree_io_parameters_module_create( + iree_runtime_instance_vm_instance(r->inst), 1, &provider, alloc, ¶ms_module); + if (iree_status_is_ok(st)) st = iree_runtime_session_append_module(r->sess, params_module); + iree_vm_module_release(params_module); + if (!iree_status_is_ok(st)) { iree_io_parameter_provider_release(provider); return st; } + r->param_provider = provider; + return st; +} + +/* --task_topology_group_count for the next local-task device; <= 0 leaves IREE's detection. */ +static iree_status_t apply_task_topology(int group_count) { + if (group_count <= 0) return iree_ok_status(); + char flag[64]; + snprintf(flag, sizeof(flag), "--task_topology_group_count=%d", group_count); + char* argv_storage[] = { "skainet_iree_module", flag }; + char** argv = argv_storage; + int argc = 2; + return iree_flags_parse(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK, &argc, &argv); +} + +static void runner_free(Runner* r) { + if (!r) return; + if (r->param_provider) iree_io_parameter_provider_release(r->param_provider); + if (r->sess) iree_runtime_session_release(r->sess); + if (r->dev) iree_hal_device_release(r->dev); + if (r->inst) iree_runtime_instance_release(r->inst); + free(r->fn_name); + free(r); +} + +JNIEXPORT jlong JNICALL JNIFN(nativeCreate)(JNIEnv* env, jobject thiz, + jstring jdev, jstring jvmfb, jstring jirpa, jstring jfn, jint group_count) { + const char* dev = (*env)->GetStringUTFChars(env, jdev, 0); + const char* vmfb = (*env)->GetStringUTFChars(env, jvmfb, 0); + const char* irpa = (*env)->GetStringUTFChars(env, jirpa, 0); + const char* fn = (*env)->GetStringUTFChars(env, jfn, 0); + Runner* r = calloc(1, sizeof(Runner)); + if (r) r->fn_name = strdup(fn); + + iree_runtime_instance_options_t io; iree_runtime_instance_options_initialize(&io); + iree_runtime_instance_options_use_all_available_drivers(&io); + iree_status_t st = (r && r->fn_name) + ? iree_runtime_instance_create(&io, iree_allocator_system(), &r->inst) + : iree_status_from_code(IREE_STATUS_RESOURCE_EXHAUSTED); + if (iree_status_is_ok(st)) st = apply_task_topology((int)group_count); + if (iree_status_is_ok(st)) { + st = iree_runtime_instance_try_create_default_device(r->inst, iree_make_cstring_view(dev), &r->dev); + } + if (iree_status_is_ok(st)) { + iree_runtime_session_options_t o; iree_runtime_session_options_initialize(&o); + st = iree_runtime_session_create_with_device( + r->inst, &o, r->dev, iree_runtime_instance_host_allocator(r->inst), &r->sess); + } + if (iree_status_is_ok(st) && irpa[0] != '\0') st = append_parameters_module(r, irpa); + if (iree_status_is_ok(st)) st = iree_runtime_session_append_bytecode_module_from_file(r->sess, vmfb); + + (*env)->ReleaseStringUTFChars(env, jdev, dev); + (*env)->ReleaseStringUTFChars(env, jvmfb, vmfb); + (*env)->ReleaseStringUTFChars(env, jirpa, irpa); + (*env)->ReleaseStringUTFChars(env, jfn, fn); + if (!iree_status_is_ok(st)) { + log_and_ignore("IreeModuleRunner.nativeCreate (instance/device/session/parameters/module)", st); + runner_free(r); + return 0; + } + return (jlong)(intptr_t)r; +} + +JNIEXPORT jfloatArray JNICALL JNIFN(nativeRun)(JNIEnv* env, jobject thiz, + jlong handle, jfloatArray jin, jintArray jshape, jintArray joutShape) { + Runner* r = (Runner*)(intptr_t)handle; if (!r) return NULL; + jsize rank = (*env)->GetArrayLength(env, jshape); + if (rank <= 0 || rank > MAX_RANK || (*env)->GetArrayLength(env, joutShape) < 1 + MAX_RANK) return NULL; + + iree_hal_dim_t dims[MAX_RANK]; + jint* sh = (*env)->GetIntArrayElements(env, jshape, 0); + iree_host_size_t count = 1; + for (int i = 0; i < rank; ++i) { dims[i] = (iree_hal_dim_t)sh[i]; count *= (iree_host_size_t)sh[i]; } + (*env)->ReleaseIntArrayElements(env, jshape, sh, JNI_ABORT); + if (count != (iree_host_size_t)(*env)->GetArrayLength(env, jin)) { + __android_log_print(ANDROID_LOG_ERROR, TAG, "nativeRun: input has %d values, its shape needs %zu", + (int)(*env)->GetArrayLength(env, jin), (size_t)count); + return NULL; + } + + /* Copy straight from the pinned Java array into a device buffer. */ + jfloat* src = (*env)->GetPrimitiveArrayCritical(env, jin, 0); + if (!src) return NULL; + iree_hal_buffer_view_t* in = NULL; + iree_status_t st = iree_hal_buffer_view_allocate_buffer_copy( + r->dev, iree_runtime_session_device_allocator(r->sess), (iree_host_size_t)rank, dims, + IREE_HAL_ELEMENT_TYPE_FLOAT_32, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, + (iree_hal_buffer_params_t){ .type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL, + .access = IREE_HAL_MEMORY_ACCESS_ALL, .usage = IREE_HAL_BUFFER_USAGE_DEFAULT }, + iree_make_const_byte_span(src, count * sizeof(float)), &in); + (*env)->ReleasePrimitiveArrayCritical(env, jin, src, JNI_ABORT); + if (!iree_status_is_ok(st)) { log_and_ignore("nativeRun: input buffer_view", st); return NULL; } + + iree_runtime_call_t c; + st = iree_runtime_call_initialize_by_name(r->sess, iree_make_cstring_view(r->fn_name), &c); + if (!iree_status_is_ok(st)) { + __android_log_print(ANDROID_LOG_ERROR, TAG, "nativeRun: function '%s' not found (names are module-qualified, e.g. module.main)", r->fn_name); + log_and_ignore("nativeRun: call_initialize_by_name", st); + iree_hal_buffer_view_release(in); + return NULL; + } + st = iree_runtime_call_inputs_push_back_buffer_view(&c, in); + iree_hal_buffer_view_release(in); + if (iree_status_is_ok(st)) st = iree_runtime_call_invoke(&c, 0); + iree_hal_buffer_view_t* out = NULL; + if (iree_status_is_ok(st)) st = iree_runtime_call_outputs_pop_front_buffer_view(&c, &out); + iree_runtime_call_deinitialize(&c); + if (!iree_status_is_ok(st)) { log_and_ignore("nativeRun: invoke / output", st); if (out) iree_hal_buffer_view_release(out); return NULL; } + + iree_host_size_t outRank = iree_hal_buffer_view_shape_rank(out); + if (iree_hal_buffer_view_element_type(out) != IREE_HAL_ELEMENT_TYPE_FLOAT_32 || outRank > MAX_RANK) { + __android_log_print(ANDROID_LOG_ERROR, TAG, "nativeRun: '%s' must return one float32 tensor of rank <= %d (rank %u)", + r->fn_name, MAX_RANK, (unsigned)outRank); + iree_hal_buffer_view_release(out); + return NULL; + } + jint outShape[1 + MAX_RANK] = { (jint)outRank }; + iree_host_size_t outCount = 1; + for (iree_host_size_t i = 0; i < outRank; ++i) { + iree_hal_dim_t d = iree_hal_buffer_view_shape_dim(out, i); + outShape[1 + i] = (jint)d; outCount *= (iree_host_size_t)d; + } + + jfloatArray result = (*env)->NewFloatArray(env, (jsize)outCount); + if (result) { + jfloat* dst = (*env)->GetPrimitiveArrayCritical(env, result, 0); + st = dst ? iree_hal_device_transfer_d2h(r->dev, iree_hal_buffer_view_buffer(out), 0, dst, + outCount * sizeof(float), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, + iree_infinite_timeout()) + : iree_status_from_code(IREE_STATUS_RESOURCE_EXHAUSTED); + if (dst) (*env)->ReleasePrimitiveArrayCritical(env, result, dst, 0); + if (!iree_status_is_ok(st)) { log_and_ignore("nativeRun: transfer_d2h", st); result = NULL; } + } + iree_hal_buffer_view_release(out); + if (result) (*env)->SetIntArrayRegion(env, joutShape, 0, (jsize)(1 + outRank), outShape); + return result; +} + +JNIEXPORT void JNICALL JNIFN(nativeDestroy)(JNIEnv* env, jobject thiz, jlong handle) { + runner_free((Runner*)(intptr_t)handle); +} diff --git a/llm-runtime/iree-android/src/androidTest/AndroidManifest.xml b/llm-runtime/iree-android/src/androidTest/AndroidManifest.xml new file mode 100644 index 00000000..f0170e43 --- /dev/null +++ b/llm-runtime/iree-android/src/androidTest/AndroidManifest.xml @@ -0,0 +1,7 @@ + + + + + + diff --git a/llm-runtime/iree-android/src/androidTest/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunnerDeviceTest.kt b/llm-runtime/iree-android/src/androidTest/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunnerDeviceTest.kt new file mode 100644 index 00000000..51cb2e4c --- /dev/null +++ b/llm-runtime/iree-android/src/androidTest/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunnerDeviceTest.kt @@ -0,0 +1,127 @@ +package sk.ainet.transformers.iree.android + +import android.util.Log +import androidx.test.ext.junit.runners.AndroidJUnit4 +import androidx.test.platform.app.InstrumentationRegistry +import org.json.JSONArray +import org.json.JSONObject +import org.junit.Assert.assertTrue +import org.junit.Assume.assumeTrue +import org.junit.Test +import org.junit.runner.RunWith +import java.io.File +import java.net.HttpURLConnection +import java.net.URL +import java.nio.ByteBuffer +import java.nio.ByteOrder + +/** + * On-device check of [IreeModuleRunner]: runs every case of a `cases.json` served from the host and compares the + * output with the case's reference. Nothing model-specific lives here — the cases (vmfb, optional irpa, function, + * device, raw little-endian float32 input and expected output) come from whoever serves them: + * + * ``` + * adb reverse tcp:8766 tcp:8766 + * (cd /path/to/cases && python3 -m http.server 8766) & + * ./gradlew :llm-runtime:iree-android:connectedAndroidTest + * ``` + * + * `cases.json`: `[{"name", "vmfb", "irpa"?, "function", "device", "input", "input_shape", "expected", "rtol", "atol", + * "runs"?}]`, file names relative to the served directory. Skips (does not fail) when no server is reachable. + * Instrumentation arg `moduleCasesUrl` overrides the base URL. Per case it logs max |Δ| and p50/p95 latency under + * tag `IreeModuleDevice`. + */ +@RunWith(AndroidJUnit4::class) +class IreeModuleRunnerDeviceTest { + + private companion object { + const val TAG = "IreeModuleDevice" + const val DEFAULT_URL = "http://localhost:8766" + } + + private val instrumentation = InstrumentationRegistry.getInstrumentation() + private val baseUrl = InstrumentationRegistry.getArguments().getString("moduleCasesUrl") ?: DEFAULT_URL + private val cacheDir = File(instrumentation.targetContext.cacheDir, "iree-module-cases").apply { mkdirs() } + + private fun log(msg: String) { + Log.i(TAG, msg) + println(msg) + } + + private fun fetch(name: String): File? { + val dest = File(cacheDir, name) + return try { + val conn = URL("$baseUrl/$name").openConnection() as HttpURLConnection + conn.connectTimeout = 3_000 + conn.readTimeout = 60_000 + conn.inputStream.use { input -> dest.outputStream().use { input.copyTo(it) } } + dest + } catch (e: Exception) { + log("cannot fetch $baseUrl/$name: ${e.message}") + null + } + } + + private fun floats(file: File): FloatArray { + val bytes = file.readBytes() + return FloatArray(bytes.size / 4).also { ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer().get(it) } + } + + private fun JSONArray.ints() = IntArray(length()) { getInt(it) } + + @Test + fun everyServedCaseMatchesItsReference() { + val index = fetch("cases.json") + assumeTrue("no case server at $baseUrl (see the class doc)", index != null) + val cases = JSONArray(index!!.readText()) + assertTrue("cases.json is empty", cases.length() > 0) + val failures = mutableListOf() + for (i in 0 until cases.length()) runCase(cases.getJSONObject(i))?.let { failures += it } + assertTrue("failed cases:\n" + failures.joinToString("\n"), failures.isEmpty()) + } + + /** Returns a failure line, or null when the case passes. */ + private fun runCase(c: JSONObject): String? { + val name = c.getString("name") + val vmfb = fetch(c.getString("vmfb")) ?: return "$name: vmfb not served" + val irpa = if (c.has("irpa")) fetch(c.getString("irpa")) ?: return "$name: irpa not served" else null + val input = FloatTensor(floats(fetch(c.getString("input")) ?: return "$name: input not served"), c.getJSONArray("input_shape").ints()) + val expected = floats(fetch(c.getString("expected")) ?: return "$name: expected not served") + val rtol = c.getDouble("rtol").toFloat() + val atol = c.getDouble("atol").toFloat() + val runs = c.optInt("runs", 1).coerceAtLeast(1) + + val createStart = System.nanoTime() + val runner = try { + IreeModuleRunner(vmfb.path, irpa?.path, c.getString("function"), c.getString("device")) + } catch (e: IllegalStateException) { + return "$name: ${e.message}" + } + val createMs = (System.nanoTime() - createStart) / 1e6 + runner.use { r -> + val times = DoubleArray(runs) + var out: FloatTensor? = null + for (k in 0 until runs) { + val t0 = System.nanoTime() + out = r.run(input) + times[k] = (System.nanoTime() - t0) / 1e6 + } + val actual = out!!.data + if (actual.size != expected.size) return "$name: output $out has ${actual.size} values, reference ${expected.size}" + var maxDiff = 0f + var bad = 0 + for (j in actual.indices) { + val d = kotlin.math.abs(actual[j] - expected[j]) + if (d > maxDiff) maxDiff = d + if (!(d <= atol + rtol * kotlin.math.abs(expected[j]))) bad++ + } + times.sort() + fun pct(p: Double) = times[((times.size - 1) * p).toInt()] + log( + "$name: device=${c.getString("device")} output $out max|Δ|=$maxDiff outside tolerance=$bad " + + "create=%.0f ms first=%.1f ms p50=%.1f ms p95=%.1f ms (runs=$runs)".format(createMs, times.first(), pct(0.5), pct(0.95)), + ) + return if (bad == 0) null else "$name: $bad of ${actual.size} values outside rtol $rtol / atol $atol (max |Δ| $maxDiff)" + } + } +} diff --git a/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_module.so b/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_module.so new file mode 100755 index 00000000..0df01694 Binary files /dev/null and b/llm-runtime/iree-android/src/main/jniLibs/arm64-v8a/libskainet_iree_module.so differ diff --git a/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_module.so b/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_module.so new file mode 100755 index 00000000..a683bf21 Binary files /dev/null and b/llm-runtime/iree-android/src/main/jniLibs/armeabi-v7a/libskainet_iree_module.so differ diff --git a/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunner.kt b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunner.kt new file mode 100644 index 00000000..7f1eb0fe --- /dev/null +++ b/llm-runtime/iree-android/src/main/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunner.kt @@ -0,0 +1,109 @@ +package sk.ainet.transformers.iree.android + +/** A dense float32 tensor, row-major: what [IreeModuleRunner] takes and returns. */ +public class FloatTensor(public val data: FloatArray, public val shape: IntArray) { + init { + require(shape.isNotEmpty() && shape.size <= MAX_RANK) { "rank must be 1..$MAX_RANK, got ${shape.size}" } + require(shape.all { it > 0 }) { "dimensions must be positive: ${shape.toList()}" } + val n = shape.fold(1L) { a, d -> a * d } + require(n == data.size.toLong()) { "shape ${shape.toList()} needs $n values, got ${data.size}" } + } + + override fun toString(): String = "FloatTensor(${shape.joinToString("x")})" + + internal companion object { + const val MAX_RANK = 8 + + /** The native side's `[rank, d0, d1, …]` shape buffer → the result tensor. */ + fun fromNative(data: FloatArray, shapeBuffer: IntArray): FloatTensor { + val rank = shapeBuffer[0] + require(rank in 1..MAX_RANK) { "native output rank $rank" } + return FloatTensor(data, shapeBuffer.copyOfRange(1, 1 + rank)) + } + } +} + +/** + * JNI wrapper for `libskainet_iree_module.so` — a generic single-shot runner over the IREE C runtime: one + * compiled module, one entry function taking one float32 tensor and returning one float32 tensor (static + * shapes). For batch cartridges — a detector, a classifier, an embedding graph — that are a plain tensor function. + * + * This class knows nothing about any specific model: [vmfbPath], [irpaPath] and [functionName] are all + * caller-supplied. With an [irpaPath], weights are EXTERNAL, bound at create time from the `.irpa` (scope `model`) + * via the IREE `io_parameters` VM module; with `null`, the vmfb carries them. + * + * Not thread-safe: one [run] at a time per instance. The package/class name is the JNI symbol contract with the + * `.so` — do not move/rename. + * + * @param functionName the exported function, bare (`yolov8_detect`) or module-qualified (`module.yolov8_detect`). + * @param taskTopologyGroupCount local-task worker groups (see [IreeTaskTopology]); `null` lets IREE decide. + */ +public class IreeModuleRunner( + vmfbPath: String, + irpaPath: String?, + functionName: String, + device: String = DEFAULT_DEVICE, + taskTopologyGroupCount: Int? = null, +) : AutoCloseable { + + public val functionName: String = qualify(functionName) + private var handle: Long = 0 + + init { + require(taskTopologyGroupCount == null || taskTopologyGroupCount > 0) { + "taskTopologyGroupCount must be >= 1, got $taskTopologyGroupCount" + } + loadNativeLibrary() + handle = nativeCreate(device, vmfbPath, irpaPath.orEmpty(), this.functionName, taskTopologyGroupCount ?: 0) + check(handle != 0L) { + "IREE module runner create failed (device='$device', function='${this.functionName}', vmfb='$vmfbPath', " + + "irpa=${irpaPath ?: "none"}); the cause is in logcat under tag skainet_iree." + } + } + + /** Invokes the function on [input]; throws with the IREE status logged under `skainet_iree` on failure. */ + public fun run(input: FloatTensor): FloatTensor { + check(handle != 0L) { "runner is closed" } + val shapeBuffer = IntArray(1 + FloatTensor.MAX_RANK) + val out = nativeRun(handle, input.data, input.shape, shapeBuffer) + ?: throw IllegalStateException( + "IREE module run failed (function='$functionName', input $input); the cause is in logcat under tag skainet_iree.", + ) + return FloatTensor.fromNative(out, shapeBuffer) + } + + override fun close() { + if (handle != 0L) { + nativeDestroy(handle) + handle = 0 + } + } + + private external fun nativeCreate(device: String, vmfbPath: String, irpaPath: String, functionName: String, taskTopologyGroupCount: Int): Long + private external fun nativeRun(handle: Long, input: FloatArray, inputShape: IntArray, outputShape: IntArray): FloatArray? + private external fun nativeDestroy(handle: Long) + + public companion object { + /** CPU driver. */ + public const val DEFAULT_DEVICE: String = "local-task" + + /** GPU driver (Vulkan); needs a `vulkan-spirv`-compiled vmfb. The `.so` ships with it compiled in. */ + public const val VULKAN_DEVICE: String = "vulkan" + + /** The runtime resolves functions by module-qualified name; a bare export name gets `module.`. */ + internal fun qualify(name: String): String { + require(name.isNotBlank()) { "functionName must not be blank" } + return if ('.' in name) name else "module.$name" + } + + @Volatile private var loaded: Boolean = false + private fun loadNativeLibrary() { + if (loaded) return + synchronized(this) { + if (loaded) return + System.loadLibrary("skainet_iree_module") + loaded = true + } + } + } +} diff --git a/llm-runtime/iree-android/src/test/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunnerTest.kt b/llm-runtime/iree-android/src/test/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunnerTest.kt new file mode 100644 index 00000000..21d294ba --- /dev/null +++ b/llm-runtime/iree-android/src/test/kotlin/sk/ainet/transformers/iree/android/IreeModuleRunnerTest.kt @@ -0,0 +1,31 @@ +package sk.ainet.transformers.iree.android + +import org.junit.jupiter.api.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith + +class IreeModuleRunnerTest { + @Test + fun bareFunctionNamesAreModuleQualified() { + assertEquals("module.yolov8_detect", IreeModuleRunner.qualify("yolov8_detect")) + assertEquals("other.main", IreeModuleRunner.qualify("other.main")) + assertFailsWith { IreeModuleRunner.qualify(" ") } + } + + @Test + fun aTensorMustHaveAsManyValuesAsItsShape() { + FloatTensor(FloatArray(6), intArrayOf(1, 2, 3)) + assertFailsWith { FloatTensor(FloatArray(5), intArrayOf(1, 2, 3)) } + assertFailsWith { FloatTensor(FloatArray(0), intArrayOf()) } + assertFailsWith { FloatTensor(FloatArray(0), intArrayOf(0, 3)) } + assertFailsWith { FloatTensor(FloatArray(1), IntArray(9) { 1 }) } + } + + @Test + fun theNativeShapeBufferIsRankThenDims() { + val t = FloatTensor.fromNative(FloatArray(84 * 8400), intArrayOf(3, 1, 84, 8400, 0, 0, 0, 0, 0)) + assertContentEquals(intArrayOf(1, 84, 8400), t.shape) + assertFailsWith { FloatTensor.fromNative(FloatArray(1), intArrayOf(0, 0, 0, 0, 0, 0, 0, 0, 0)) } + } +}