iris-3b Turbo on LiteRT: 20 fp16-weight graphs for the phone GPU at FP32, a Python reference pipeline and the Android engine
743fd20 verified | package com.iris3b | |
| import java.io.File | |
| import java.nio.file.Files | |
| import java.security.MessageDigest | |
| /** | |
| * The files of one run in the app's files dir, as the rows file names them (no Android API, so the | |
| * JVM check reads the very files the phone gets): | |
| * | |
| * - rows file: `{"kind": "generic", "config": <config file>, "inputs": [{"file": ...}, ...] | |
| * (every file the run needs; devicebench pushes them), "rows": [{"key", "prompt", "noise"}]}`; | |
| * - config file: the residency policy (P1 / P2 / P4, with P4's `prefetch` rule), the dumps, the | |
| * tokenizer and embedding files, the host tensor pack and its index, and per graph its file, | |
| * accel, precision, signature, input names in feed order and output names with element counts; | |
| * - tokenizer file: `{"files": {name: text}}`, written out to `tok/` for [PromptTokens], with | |
| * `tok/embed_bf16.bin` a symbolic link to the embedding table. | |
| */ | |
| class RunFiles(val dir: File, rowsName: String) { | |
| class Row(val key: String, val prompt: String, val noise: String) | |
| class Tensor(val offset: Long, val count: Int, val shape: List<Int>) | |
| val rowsDoc: Map<String, Any?> = obj(Json.parse(File(dir, rowsName).readText())) | |
| val configName: String = rowsDoc.str("config") | |
| val config: Map<String, Any?> = obj(Json.parse(File(dir, configName).readText())) | |
| val rows: List<Row> = rowsDoc.list("rows").map { | |
| val r = obj(it) | |
| Row(r.str("key"), r.str("prompt"), (r["noise"] as? String) ?: "noise_pix") | |
| } | |
| val residency: Residency = Residency.valueOf(config.str("residency")) | |
| /** | |
| * Graph roles whose GPU program cache is serialized, as role prefixes: `"program_cache": true` | |
| * = every graph, a list (e.g. `["dit_"]`) = the roles starting with one of them, absent = none. | |
| */ | |
| val programCacheRoles: List<String> = when (val pc = config["program_cache"]) { | |
| true -> listOf("") | |
| is List<*> -> pc.map { it as String } | |
| else -> emptyList() | |
| } | |
| val programCache: Boolean = programCacheRoles.isNotEmpty() | |
| val dumps: List<String> = (config["dumps"] as? List<*>)?.map { it as String } ?: emptyList() | |
| private val hostFile = File(dir, obj(config["host"]).str("file")) | |
| private val hostIndex: Map<String, Tensor> = obj(obj(config["host"])["tensors"]).mapValues { (_, v) -> | |
| val t = obj(v) | |
| Tensor(t.num("offset").toLong(), t.num("count").toInt(), (t["shape"] as List<*>).map { (it as Double).toInt() }) | |
| } | |
| private val hostCache = HashMap<String, FloatArray>() | |
| /** Host tensors from the pack, each read once. */ | |
| val host = HostTensors { name -> | |
| hostCache.getOrPut(name) { | |
| val t = requireNotNull(hostIndex[name]) { "host tensor $name is not in ${hostFile.name}" } | |
| HostOps.readFloats(hostFile, t.offset, t.count) | |
| } | |
| } | |
| fun hostShape(name: String): List<Int> = requireNotNull(hostIndex[name]) { "no host tensor $name" }.shape | |
| /** Every file the rows file declares, with a check that each is a regular file. */ | |
| fun missingInputs(): List<String> = | |
| rowsDoc.list("inputs").map { obj(it).str("file") }.filter { !File(dir, it).isFile } | |
| /** The [EngineConfig] the config file describes. */ | |
| fun engineConfig(): EngineConfig { | |
| val g = obj(config["graphs"]) | |
| return EngineConfig( | |
| residency = residency, | |
| textEncoder = g.list("te").map { spec(it) }, | |
| adapter = spec(g["adapter"]), | |
| dual = g.list("dual").map { spec(it) }, | |
| dualModRows = g.list("dual_mod_rows").map { range(it) }, | |
| single = g.list("single").map { spec(it) }, | |
| singleModRows = g.list("single_mod_rows").map { range(it) }, | |
| pixel = spec(g["pixel"]), | |
| dumps = dumps.toSet(), | |
| prefetch = prefetchRule(), | |
| ) | |
| } | |
| /** P4's memory rule: `"prefetch": {"factor": x, "margin_bytes": n}`, each defaulting to [PrefetchRule]'s. */ | |
| private fun prefetchRule(): PrefetchRule { | |
| val p = config["prefetch"] as? Map<*, *> ?: return PrefetchRule() | |
| val default = PrefetchRule() | |
| return PrefetchRule( | |
| factor = (p["factor"] as? Double) ?: default.factor, | |
| marginBytes = (p["margin_bytes"] as? Double)?.toLong() ?: default.marginBytes, | |
| ) | |
| } | |
| /** Every graph file the config names, in pipeline order. */ | |
| fun graphFiles(): List<String> { | |
| val c = engineConfig() | |
| return (c.textEncoder + c.adapter + c.dual + c.single + c.pixel).map { it.file } | |
| } | |
| /** | |
| * Writes the tokenizer files into `tok/` and links `tok/embed_bf16.bin` to the embedding | |
| * table; returns the sha256 of every file written. | |
| */ | |
| fun materializeTokenizer(): Map<String, String> { | |
| val tok = File(dir, "tok") | |
| tok.mkdirs() | |
| val files = obj(obj(Json.parse(File(dir, config.str("tokenizer")).readText()))["files"]) | |
| val sums = LinkedHashMap<String, String>() | |
| for ((name, text) in files) { | |
| val bytes = (text as String).toByteArray(Charsets.UTF_8) | |
| File(tok, name).writeBytes(bytes) | |
| sums[name] = sha256(bytes) | |
| } | |
| val link = File(tok, PromptTokens.EMBEDDING_FILE).toPath() | |
| Files.deleteIfExists(link) | |
| Files.createSymbolicLink(link, File(dir, config.str("embed")).absoluteFile.toPath()) | |
| return sums | |
| } | |
| /** The output tensors of one row in out-file order: the image, then the dumps. */ | |
| fun outputs(): List<Pair<String, Int>> { | |
| val out = arrayListOf("image" to HostOps.IMAGE_FLOATS) | |
| for (d in dumps) { | |
| when (d) { | |
| "te" -> out.add("te" to HostOps.TE_FLOATS) | |
| "y0" -> out.add("y0" to HostOps.WINDOW * HostOps.DIT_HIDDEN) | |
| "v" -> for (s in 0 until 4) out.add("v_s$s" to HostOps.PATCHES * HostOps.PIXELS * HostOps.CHANNELS) | |
| else -> throw IllegalArgumentException("unknown dump $d") | |
| } | |
| } | |
| return out | |
| } | |
| private fun spec(any: Any?): GraphSpec { | |
| val m = obj(any) | |
| val accel = (m["accel"] as? String) ?: "gpu" | |
| require(accel == "gpu" || accel == "cpu") { "${m["role"]}: accel $accel (gpu or cpu)" } | |
| val file = m.str("file") | |
| return GraphSpec( | |
| role = m.str("role"), | |
| file = file, | |
| precision = m.str("precision"), | |
| signature = m.str("signature"), | |
| inputs = m.list("inputs").map { it as String }, | |
| outputs = m.list("outputs").map { | |
| val pair = it as List<*> | |
| (pair[0] as String) to (pair[1] as Double).toInt() | |
| }, | |
| accel = accel, | |
| threads = ((m["threads"] ?: config["cpu_threads"]) as? Double)?.toInt() ?: 4, | |
| xnnpackCache = m["xnnpack_cache"] as? Boolean ?: false, | |
| bytes = File(dir, file).length(), | |
| ) | |
| } | |
| private fun range(any: Any?): IntArray = (any as List<*>).map { (it as Double).toInt() }.toIntArray() | |
| companion object { | |
| fun obj(any: Any?): Map<String, Any?> = requireNotNull(any as? Map<String, Any?>) { "expected a JSON object" } | |
| fun Map<String, Any?>.str(key: String): String = requireNotNull(this[key] as? String) { "missing string $key" } | |
| fun Map<String, Any?>.num(key: String): Double = requireNotNull(this[key] as? Double) { "missing number $key" } | |
| fun Map<String, Any?>.list(key: String): List<Any?> = requireNotNull(this[key] as? List<*>) { "missing list $key" } | |
| fun sha256(bytes: ByteArray): String = | |
| MessageDigest.getInstance("SHA-256").digest(bytes).joinToString("") { "%02x".format(it) } | |
| } | |
| } | |