perf: RKNN half RAM model usage + no weight duplication in multi-core (#23958)
* perf(ml): share RKNN weights across NPU cores and halve model RAM * perf(ml): default rknn threads to npu cores and cache shape setup * perf(ml): serve rknn runs in arrival order (cherry picked from commit e659fb62f2bb3f6b1db08ed643d3d054a539c415)
This commit is contained in:
Horatiu George Todoran
committed by
Codex
Codex
parent
1f422ea4c4
commit
f418967d55
@@ -70,7 +70,7 @@ class Settings(BaseSettings):
|
||||
ann_fp16_turbo: bool = False
|
||||
ann_tuning_level: int = 2
|
||||
rknn: bool = True
|
||||
rknn_threads: int = 1
|
||||
rknn_threads: int | None = None
|
||||
preload: PreloadModelData | None = None
|
||||
max_batch_size: MaxBatchSize = MaxBatchSize()
|
||||
model_organization: ModelOrganization = ModelOrganization.APP
|
||||
|
||||
@@ -100,8 +100,7 @@ SUPPORTED_PROVIDERS = [
|
||||
"CPUExecutionProvider",
|
||||
]
|
||||
|
||||
RKNN_SUPPORTED_SOCS = ["rk3566", "rk3568", "rk3576", "rk3588"]
|
||||
RKNN_COREMASK_SUPPORTED_SOCS = ["rk3576", "rk3588"]
|
||||
RKNN_NPU_CORES = {"rk3566": 1, "rk3568": 1, "rk3576": 2, "rk3588": 3}
|
||||
|
||||
|
||||
WEBLATE_TO_FLORES200 = {
|
||||
|
||||
@@ -9,9 +9,10 @@ import orjson
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.config import log, settings
|
||||
from immich_ml.models.constants import RKNN_NPU_CORES
|
||||
from immich_ml.schemas import ModelInput, ModelTensor, SessionNode, Shape
|
||||
|
||||
from .rknnpool import RknnNode, RknnPoolExecutor, is_available, native_outputs, soc_name
|
||||
from .rknnpool import RknnNode, RknnPoolExecutor, is_available, soc_name
|
||||
|
||||
is_available = is_available and settings.rknn
|
||||
model_prefix = Path("rknpu") / soc_name if is_available and soc_name is not None else None
|
||||
@@ -21,12 +22,6 @@ def model_path(model_dir: Path, variant: str = "") -> Path:
|
||||
return (model_dir / model_prefix if model_prefix else model_dir) / variant / "model.rknn"
|
||||
|
||||
|
||||
def run_inference(rknn_lite: Any, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
|
||||
rknn_lite.rknn_runtime.set_inputs(inputs, None, data_format)
|
||||
rknn_lite.rknn_runtime.run(False)
|
||||
return native_outputs(rknn_lite)
|
||||
|
||||
|
||||
def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
|
||||
if array.ndim != 4:
|
||||
return None # token ids carry no layout
|
||||
@@ -39,11 +34,11 @@ def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
|
||||
|
||||
class RknnSession:
|
||||
def __init__(self, model_path: Path) -> None:
|
||||
self.tpe = settings.rknn_threads
|
||||
self.tpe = settings.rknn_threads or RKNN_NPU_CORES.get(soc_name or "", 1)
|
||||
|
||||
log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.")
|
||||
self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe, func=run_inference)
|
||||
log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads.")
|
||||
self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe)
|
||||
log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads sharing its weights.")
|
||||
# the shapes the binary was compiled for, among which the runtime routes itself
|
||||
batch = self.rknnpool.inputs[0].shape[0]
|
||||
declared = orjson.loads(self.rknnpool.custom_string) if self.rknnpool.custom_string else {}
|
||||
@@ -77,7 +72,7 @@ class RknnSession:
|
||||
input_feed: ModelInput,
|
||||
run_options: Any = None,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
inputs = [array if array.flags.c_contiguous else array.copy() for array in input_feed.values()]
|
||||
inputs = list(input_feed.values())
|
||||
return self.rknnpool.run(inputs, input_layout(self.rknnpool.inputs[0].shape, inputs[0]))
|
||||
|
||||
|
||||
|
||||
@@ -1,18 +1,15 @@
|
||||
# This code is from leafqycc/rknn-multi-threaded
|
||||
# Following Apache License 2.0
|
||||
|
||||
import ctypes
|
||||
import logging
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from ctypes import byref, sizeof
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, NamedTuple
|
||||
from queue import SimpleQueue
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from immich_ml.config import log
|
||||
from immich_ml.models.constants import RKNN_COREMASK_SUPPORTED_SOCS, RKNN_SUPPORTED_SOCS
|
||||
from immich_ml.models.constants import RKNN_NPU_CORES
|
||||
from immich_ml.schemas import ModelTensor
|
||||
|
||||
|
||||
@@ -20,7 +17,7 @@ def get_soc(device_tree_path: Path | str) -> str | None:
|
||||
try:
|
||||
with Path(device_tree_path).open() as f:
|
||||
device_compatible_str = f.read()
|
||||
for soc in RKNN_SUPPORTED_SOCS:
|
||||
for soc in RKNN_NPU_CORES:
|
||||
if soc in device_compatible_str:
|
||||
return soc
|
||||
log.warning("Device is not supported for RKNN")
|
||||
@@ -29,26 +26,34 @@ def get_soc(device_tree_path: Path | str) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
soc_name = None
|
||||
is_available = False
|
||||
try:
|
||||
from rknnlite.api import RKNNLite, rknn_runtime
|
||||
|
||||
soc_name = get_soc("/proc/device-tree/compatible")
|
||||
is_available = soc_name is not None
|
||||
except ImportError:
|
||||
log.debug("RKNN is not available")
|
||||
|
||||
|
||||
class RknnNode(NamedTuple):
|
||||
name: str
|
||||
shape: tuple[int, ...]
|
||||
|
||||
|
||||
_CUSTOM_STRING_QUERY = 7 # RKNN_QUERY_CUSTOM_STRING
|
||||
# RKNN_QUERY_OUTPUT_ATTR, and RKNN_QUERY_CURRENT_OUTPUT_ATTR for the shape a dynamic binary last produced
|
||||
_OUTPUT_QUERY, _CURRENT_OUTPUT_QUERY = 2, 15
|
||||
_DTYPES = {0: np.float32, 1: np.float16} # rknn_tensor_type
|
||||
# rknn_query_cmd
|
||||
_IN_OUT_NUM, _INPUT_ATTR, _OUTPUT_ATTR, _CUSTOM_STRING = 0, 1, 2, 7
|
||||
_INPUT_DYNAMIC_RANGE, _CURRENT_OUTPUT_ATTR = 13, 15
|
||||
# rknn_tensor_format
|
||||
_NCHW, _NHWC = 0, 1
|
||||
# rknn_tensor_type, in enum order
|
||||
_DTYPES: list[type[np.generic]] = [
|
||||
np.float32,
|
||||
np.float16,
|
||||
np.int8,
|
||||
np.uint8,
|
||||
np.int16,
|
||||
np.uint16,
|
||||
np.int32,
|
||||
np.uint32,
|
||||
np.int64,
|
||||
np.bool_,
|
||||
]
|
||||
_TYPES = {np.dtype(dtype): index for index, dtype in enumerate(_DTYPES)}
|
||||
|
||||
|
||||
class _IoNum(ctypes.Structure):
|
||||
_fields_ = [("n_input", ctypes.c_uint32), ("n_output", ctypes.c_uint32)]
|
||||
|
||||
|
||||
class _CustomString(ctypes.Structure):
|
||||
@@ -76,94 +81,194 @@ class _TensorAttr(ctypes.Structure): # rknn_tensor_attr
|
||||
]
|
||||
|
||||
|
||||
def native_outputs(rknn_lite: "RKNNLite") -> list[NDArray[Any]]:
|
||||
"""The outputs in the precision the NPU computed them in, which RKNNLite would widen to fp32 on the CPU."""
|
||||
runtime: Any = rknn_lite.rknn_runtime
|
||||
count = runtime.get_in_out_num()[1]
|
||||
buffers = (rknn_runtime.RKNNOutput * count)()
|
||||
for index, buffer in enumerate(buffers):
|
||||
buffer.index = index
|
||||
if runtime.lib.rknn_outputs_get(runtime.context, count, buffers, None) != 0:
|
||||
raise RuntimeError("RKNN inference failed!")
|
||||
query = _CURRENT_OUTPUT_QUERY if runtime.is_dynamic_shape else _OUTPUT_QUERY
|
||||
outputs = []
|
||||
for buffer in buffers:
|
||||
attr = _TensorAttr(index=buffer.index)
|
||||
runtime.lib.rknn_query(runtime.context, query, ctypes.byref(attr), ctypes.sizeof(attr))
|
||||
data = ctypes.cast(buffer.buf, ctypes.POINTER(ctypes.c_char * buffer.size)).contents
|
||||
outputs.append(np.frombuffer(data, _DTYPES[attr.type]).reshape(attr.dims[: attr.n_dims]).copy())
|
||||
runtime.lib.rknn_outputs_release(runtime.context, count, buffers)
|
||||
return outputs
|
||||
class _InputRange(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("index", ctypes.c_uint32),
|
||||
("shape_number", ctypes.c_uint32),
|
||||
("fmt", ctypes.c_int),
|
||||
("name", ctypes.c_char * 256),
|
||||
("dyn_range", (ctypes.c_uint32 * 16) * 512),
|
||||
("n_dims", ctypes.c_uint32),
|
||||
]
|
||||
|
||||
|
||||
def tensor_nodes(rknn_lite: "RKNNLite") -> tuple[list[RknnNode], list[RknnNode]]:
|
||||
runtime = rknn_lite.rknn_runtime
|
||||
n_inputs, n_outputs = runtime.get_in_out_num()
|
||||
|
||||
def node(index: int, is_output: bool) -> RknnNode:
|
||||
attr = runtime.get_tensor_attr(index, is_output=is_output)
|
||||
return RknnNode(attr.name.decode(), tuple(attr.dims[: attr.n_dims]))
|
||||
|
||||
return [node(i, False) for i in range(n_inputs)], [node(i, True) for i in range(n_outputs)]
|
||||
class _Input(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("index", ctypes.c_uint32),
|
||||
("buf", ctypes.c_void_p),
|
||||
("size", ctypes.c_uint32),
|
||||
("pass_through", ctypes.c_uint8),
|
||||
("type", ctypes.c_int),
|
||||
("fmt", ctypes.c_int),
|
||||
]
|
||||
|
||||
|
||||
def custom_string(rknn_lite: "RKNNLite") -> str:
|
||||
runtime: Any = rknn_lite.rknn_runtime
|
||||
buffer = _CustomString()
|
||||
# rknnlite declares argtypes on the shared handle, so pass plain values and let ctypes convert
|
||||
runtime.lib.rknn_query(runtime.context, _CUSTOM_STRING_QUERY, ctypes.byref(buffer), ctypes.sizeof(buffer))
|
||||
declared: str = buffer.string.decode()
|
||||
return declared
|
||||
class _Output(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("want_float", ctypes.c_uint8),
|
||||
("is_prealloc", ctypes.c_uint8),
|
||||
("index", ctypes.c_uint32),
|
||||
("buf", ctypes.c_void_p),
|
||||
("size", ctypes.c_uint32),
|
||||
]
|
||||
|
||||
|
||||
def init_rknn(model_path: str) -> "RKNNLite":
|
||||
if not is_available:
|
||||
raise RuntimeError("rknn is not available!")
|
||||
rknn_lite = RKNNLite()
|
||||
rknn_lite.rknn_log.logger.setLevel(logging.ERROR)
|
||||
ret = rknn_lite.load_rknn(model_path)
|
||||
_Key = tuple[tuple[int, ...], ...]
|
||||
_Spec = tuple[tuple[int, ...], np.dtype[Any]]
|
||||
_context = ctypes.c_uint64 # rknn_context on 64-bit
|
||||
_ctx_p = ctypes.POINTER(_context)
|
||||
_u32 = ctypes.c_uint32
|
||||
|
||||
soc_name = None
|
||||
is_available = False
|
||||
try:
|
||||
lib = ctypes.CDLL("librknnrt.so") # CDLL releases the GIL for each call, so contexts run in parallel
|
||||
lib.rknn_init.argtypes = _ctx_p, ctypes.c_char_p, _u32, _u32, ctypes.c_void_p
|
||||
lib.rknn_dup_context.argtypes = _ctx_p, _ctx_p
|
||||
lib.rknn_destroy.argtypes = (_context,)
|
||||
lib.rknn_query.argtypes = _context, ctypes.c_int, ctypes.c_void_p, _u32
|
||||
lib.rknn_set_input_shapes.argtypes = _context, _u32, ctypes.POINTER(_TensorAttr)
|
||||
lib.rknn_inputs_set.argtypes = _context, _u32, ctypes.POINTER(_Input)
|
||||
lib.rknn_run.argtypes = _context, ctypes.c_void_p
|
||||
lib.rknn_outputs_get.argtypes = _context, _u32, ctypes.POINTER(_Output), ctypes.c_void_p
|
||||
lib.rknn_outputs_release.argtypes = _context, _u32, ctypes.POINTER(_Output)
|
||||
|
||||
soc_name = get_soc("/proc/device-tree/compatible")
|
||||
is_available = soc_name is not None
|
||||
except OSError:
|
||||
log.debug("RKNN is not available")
|
||||
|
||||
|
||||
def _check(ret: int, call: str) -> None:
|
||||
if ret != 0:
|
||||
raise RuntimeError("Failed to load RKNN model")
|
||||
raise RuntimeError(f"{call} failed with RKNN error {ret}")
|
||||
|
||||
if soc_name in RKNN_COREMASK_SUPPORTED_SOCS:
|
||||
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
|
||||
else:
|
||||
ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms.
|
||||
|
||||
if ret != 0:
|
||||
raise RuntimeError("Failed to initialize RKNN runtime environment")
|
||||
def _query(ctx: int, cmd: int, info: ctypes.Structure) -> None:
|
||||
_check(lib.rknn_query(ctx, cmd, byref(info), sizeof(info)), "rknn_query")
|
||||
|
||||
return rknn_lite
|
||||
|
||||
def _attr(ctx: int, cmd: int, index: int) -> _TensorAttr:
|
||||
attr = _TensorAttr(index=index)
|
||||
_query(ctx, cmd, attr)
|
||||
return attr
|
||||
|
||||
|
||||
def _specs(attrs: list[_TensorAttr]) -> list[_Spec]:
|
||||
return [(tuple(attr.dims[: attr.n_dims]), np.dtype(_DTYPES[attr.type])) for attr in attrs]
|
||||
|
||||
|
||||
class RknnPoolExecutor:
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
tpes: int,
|
||||
func: Callable[["RKNNLite", list[ModelTensor], str | None], list[NDArray[np.float32]]],
|
||||
) -> None:
|
||||
self.tpes = tpes
|
||||
self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)]
|
||||
self.inputs, self.outputs = tensor_nodes(self.rknn_pool[0])
|
||||
self.custom_string = custom_string(self.rknn_pool[0])
|
||||
self.pool = ThreadPoolExecutor(max_workers=tpes)
|
||||
self.func = func
|
||||
self.num = 0
|
||||
self.lock = threading.Lock()
|
||||
"""One context per thread, all but the first duplicated from it so that they share its weights.
|
||||
|
||||
def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
|
||||
with self.lock:
|
||||
idx = self.num % self.tpes
|
||||
self.num += 1
|
||||
The model is loaded by path rather than from a buffer, so the runtime holds the only copy of it."""
|
||||
|
||||
fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs, data_format)
|
||||
return fut.result()
|
||||
def __init__(self, model_path: str, tpes: int) -> None:
|
||||
self.contexts: list[int] = []
|
||||
self.executor = ThreadPoolExecutor(tpes, thread_name_prefix="rknn")
|
||||
master = _context()
|
||||
_check(lib.rknn_init(byref(master), model_path.encode(), 0, 0, None), "rknn_init")
|
||||
self.contexts.append(master.value)
|
||||
for _ in range(tpes - 1):
|
||||
child = _context()
|
||||
_check(lib.rknn_dup_context(byref(master), byref(child)), "rknn_dup_context")
|
||||
self.contexts.append(child.value)
|
||||
self.idle: SimpleQueue[int] = SimpleQueue()
|
||||
for ctx in self.contexts:
|
||||
self.idle.put(ctx)
|
||||
|
||||
io = _IoNum()
|
||||
_query(master.value, _IN_OUT_NUM, io)
|
||||
self.input_attrs = [_attr(master.value, _INPUT_ATTR, i) for i in range(io.n_input)]
|
||||
self.output_attrs = [_attr(master.value, _OUTPUT_ATTR, i) for i in range(io.n_output)]
|
||||
self.inputs = [RknnNode(a.name.decode(), tuple(a.dims[: a.n_dims])) for a in self.input_attrs]
|
||||
self.outputs = [RknnNode(a.name.decode(), tuple(a.dims[: a.n_dims])) for a in self.output_attrs]
|
||||
custom = _CustomString()
|
||||
lib.rknn_query(master.value, _CUSTOM_STRING, byref(custom), sizeof(custom))
|
||||
self.custom_string = custom.string.decode()
|
||||
|
||||
ranges = [_InputRange(index=index) for index in range(io.n_input)]
|
||||
for rng in ranges:
|
||||
lib.rknn_query(master.value, _INPUT_DYNAMIC_RANGE, byref(rng), sizeof(rng))
|
||||
self.formats = [rng.fmt for rng in ranges]
|
||||
self.compiled = [{tuple(rng.dyn_range[i][: rng.n_dims]) for i in range(rng.shape_number)} for rng in ranges]
|
||||
self.dynamic = any(self.compiled)
|
||||
self.shape_attrs: dict[_Key, ctypes.Array[_TensorAttr]] = {}
|
||||
self.output_specs: dict[_Key, list[_Spec]] = {(): _specs(self.output_attrs)}
|
||||
self.current: dict[int, _Key] = {}
|
||||
|
||||
def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[Any]]:
|
||||
return self.executor.submit(self._next, inputs, data_format).result()
|
||||
|
||||
def _next(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[Any]]:
|
||||
ctx = self.idle.get()
|
||||
try:
|
||||
return self._run(ctx, inputs, data_format)
|
||||
finally:
|
||||
self.idle.put(ctx)
|
||||
|
||||
def _run(self, ctx: int, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[Any]]:
|
||||
fed: list[NDArray[Any]] = [
|
||||
np.ascontiguousarray(array.transpose(0, 2, 3, 1) if data_format == "nchw" and array.ndim == 4 else array)
|
||||
for array in inputs
|
||||
]
|
||||
key = self._set_shapes(ctx, fed) if self.dynamic else ()
|
||||
feed = (_Input * len(fed))(
|
||||
*(
|
||||
_Input(index=i, buf=array.ctypes.data, size=array.nbytes, type=_TYPES[array.dtype], fmt=_NHWC)
|
||||
for i, array in enumerate(fed)
|
||||
)
|
||||
)
|
||||
_check(lib.rknn_inputs_set(ctx, len(fed), feed), "rknn_inputs_set")
|
||||
_check(lib.rknn_run(ctx, None), "rknn_run")
|
||||
|
||||
specs = self.output_specs.get(key)
|
||||
if specs is None:
|
||||
attrs = [_attr(ctx, _CURRENT_OUTPUT_ATTR, i) for i in range(len(self.output_attrs))]
|
||||
specs = self.output_specs[key] = _specs(attrs)
|
||||
outputs = [np.empty(dims, dtype) for dims, dtype in specs]
|
||||
buffers = (_Output * len(outputs))(
|
||||
*(
|
||||
_Output(is_prealloc=1, index=i, buf=array.ctypes.data, size=array.nbytes)
|
||||
for i, array in enumerate(outputs)
|
||||
)
|
||||
)
|
||||
_check(lib.rknn_outputs_get(ctx, len(outputs), buffers, None), "rknn_outputs_get")
|
||||
lib.rknn_outputs_release(ctx, len(outputs), buffers)
|
||||
return outputs
|
||||
|
||||
def _set_shapes(self, ctx: int, inputs: list[NDArray[Any]]) -> _Key:
|
||||
key = tuple(
|
||||
self._compiled_shape(fmt, compiled, array.shape)
|
||||
for fmt, compiled, array in zip(self.formats, self.compiled, inputs)
|
||||
)
|
||||
if self.current.get(ctx) == key:
|
||||
return key
|
||||
attrs = self.shape_attrs.get(key)
|
||||
if attrs is None:
|
||||
attrs = self.shape_attrs[key] = (_TensorAttr * len(key))()
|
||||
for i, (source, compiled, fmt, shape) in enumerate(zip(self.input_attrs, self.compiled, self.formats, key)):
|
||||
attrs[i] = _TensorAttr.from_buffer_copy(source)
|
||||
if compiled:
|
||||
attrs[i].n_dims = len(shape)
|
||||
attrs[i].dims[: len(shape)] = shape
|
||||
attrs[i].fmt = fmt
|
||||
_check(lib.rknn_set_input_shapes(ctx, len(key), attrs), "rknn_set_input_shapes")
|
||||
self.current[ctx] = key
|
||||
return key
|
||||
|
||||
@staticmethod
|
||||
def _compiled_shape(fmt: int, compiled: set[tuple[int, ...]], fed: tuple[int, ...]) -> tuple[int, ...]:
|
||||
shape = (fed[0], fed[3], fed[1], fed[2]) if len(fed) == 4 and fmt == _NCHW else fed
|
||||
if compiled and shape not in compiled:
|
||||
raise ValueError(f"binary was compiled for {sorted(compiled)}, fed {list(fed)}")
|
||||
return shape
|
||||
|
||||
def release(self) -> None:
|
||||
self.pool.shutdown()
|
||||
for rknn_lite in self.rknn_pool:
|
||||
rknn_lite.release()
|
||||
self.executor.shutdown()
|
||||
for ctx in reversed(self.contexts): # duplicates before the context whose weights they share
|
||||
lib.rknn_destroy(ctx)
|
||||
self.contexts = []
|
||||
|
||||
def __del__(self) -> None:
|
||||
self.release()
|
||||
|
||||
Reference in New Issue
Block a user