Skip to main content

Crate taconite_adaface

Crate taconite_adaface 

Source
Expand description

Embeddable Rust runtime for the AdaFace IR face embedders: NPU convolutions + CPU glue/head.

IrEmbedder replays a bundle exported by iron/applications/adaface_ir18/export_ir18.py (or the IR-101 wrapper export_ir101.py): the shared IRON bundle format (read with taconite-bundle), whose manifest lists the network’s steps in order. The steps encode the architecture, so this one runtime serves every exported IR backbone; the run_ir18 / run_ir101 binaries are thin wrappers over cli, which is itself a consumer of the public API:

use taconite_adaface::IrEmbedder;

let mut emb = IrEmbedder::load("/path/to/bundle")?;
// An aligned 112x112 face crop, CHW f32, cvlface preprocessing
// ((rgb/255) - 0.5) / 0.5. Channels beyond the image's (the bundle pads
// the stem to a multiple of 8) are zero-filled internally.
let chw = vec![0f32; 3 * 112 * 112];
let embedding = emb.embed(&chw)?; // 512-d

Loading reads every weight into RAM and checks every step against them. The first forward creates one hardware context per distinct conv kernel (16 for IR-18/IR-101) and the session keeps them ALL resident (an LRU cache sized to the NPU2’s 16-concurrent-context cap), so every later forward reloads nothing. Keep ONE IrEmbedder alive for the process and reuse it; a second simultaneous session would fight over the context cap.

Kernels run through taconite: over XRT (feature xrt, the default), or straight through the amdxdna driver’s ioctls with no XRT at all (feature direct, which wins when both are on; build with --no-default-features --features direct). BACKEND says which.

Pipeline per forward (all in Rust):

  • every convolution on the NPU, in the Conv2d channel-tiled [H, C/8, W, 8] bf16 layout; wide convs are summed over input-channel groups sized to fit L1, and the width%4 stages (14x14 / 7x7) zero-pad the input width to a multiple of 4 and crop the extra output columns;
  • the per-channel glue – leading BN0 affine, fused conv-bias + PReLU (prelu(x + bias)), plain conv bias – and the residual add and stride-2 subsample on the CPU (these are tiny and dispatch-bound, and float32 on the host is if anything more accurate);
  • the embedding head (both head BatchNorms folded into the Linear) as an exact f32 mat-vec on the CPU (M=1) – keeping it off the NPU lets all 16 conv kernels stay resident within the concurrent-context cap.

Threading: an IrEmbedder is Send (move it to a worker thread) but deliberately not Sync – all methods take &mut self, one forward at a time. Share it behind a Mutex if several threads need embeddings.

Structs§

IrEmbedder
A resident NPU session replaying one exported IR bundle.

Enums§

Error
Everything that can go wrong loading a bundle or running a forward. The library never prints or exits; the host application decides.

Constants§

BACKEND
Which NPU path this build runs kernels through: "xrt" or "direct".
BUNDLE_VERSION
The bundle format version this runtime reads (export_ir18.VERSION).

Functions§

cli
Full CLI entry point, shared by the per-network binaries (run_ir18, run_ir101). net_name labels the output when the bundle has no model param; the replayed network is whatever the bundle contains. Prints per-forward latency (cold vs warm) and the embedding’s cosine to the CPU reference (ref.cpu_embedding) and the recorder’s embedding (ref.npu_embedding); when the bundle has the raw f32 input (ref.input_chw), also exercises the public IrEmbedder::embed API and checks it reproduces the recorded replay.
cosine
Cosine similarity between two embeddings (0.0 if either is all-zero).