tools/maya_export/quant_depth_frame.py

"""Weight-only block quantization (MatMulNBits) of the fp16-storage depth graphs, streamed one tensor at a time.

Input: an fp16-storage depth graph from export_depth_all (depth_frame_fp16.onnx, the sampled runtime graph, or
depth_all.onnx, the teacher-forced graph used by the quality gate). In both, each Linear weight W (N, K) is an fp16
initializer used as Cast(W -> fp32) followed by either Transpose(perm 1,0) -> MatMul(s) or Gemm(transB=1).
Each such use becomes MatMulNBits(A, Wq, scales) with Wq quantized from the fp32 upcast of W (lossless upcast of
the fp16 source). Shared weights (group 8 for levels 9..15, the cross blocks for all levels) are quantized once
and referenced by every MatMulNBits node that used them.

Weights matching --keep (regex on the initializer name) stay fp16-stored behind their Cast. Embeddings (Gather)
are left as they are. Graph I/O is unchanged.

Never holds more than one weight (plus its quantized copy) in memory: initializers are read from the source
external-data file by offset and written sequentially to the destination data file.

    PYTHONPATH=. .venv-tts/bin/python -m tools.maya_export.quant_depth_frame \
        --src ~/.cache/maya_fast/depth_build/fp16/depth_frame_fp16.onnx \
        --dst ~/.cache/maya_fast/depth_build/q8a1/depth_frame_q8a1.onnx --bits 8 --block 32 --accuracy-level 1

Note (ORT MLAS, checked on main): 8-bit MatMulNBits has native kernels only for accuracy_level 4 (int8 compute);
accuracy_level 1 runs the generic path on every CPU: dequantize B to an fp32 scratch buffer on each call, then
SGEMM. That path is the same C++ code on x86 and ARM.
"""
from __future__ import annotations

import argparse
import collections
import os
import re
from pathlib import Path

import numpy as np
import onnx
from onnx import TensorProto, helper

NP = {TensorProto.FLOAT: np.float32, TensorProto.FLOAT16: np.float16, TensorProto.INT64: np.int64,
      TensorProto.INT32: np.int32, TensorProto.INT8: np.int8, TensorProto.UINT8: np.uint8,
      TensorProto.BOOL: np.bool_, TensorProto.DOUBLE: np.float64}


def _attr(node, name, default=None):
    for a in node.attribute:
        if a.name == name:
            return helper.get_attribute_value(a)
    return default


def _read(t: TensorProto, src_dir: str) -> np.ndarray:
    if t.data_location == TensorProto.EXTERNAL:
        info = {e.key: e.value for e in t.external_data}
        off = int(info.get("offset", 0))
        n = int(np.prod(list(t.dims), dtype=np.int64)) * np.dtype(NP[t.data_type]).itemsize
        length = int(info.get("length", n))
        with open(os.path.join(src_dir, info["location"]), "rb") as fh:
            fh.seek(off)
            raw = fh.read(length)
        return np.frombuffer(raw, dtype=NP[t.data_type]).reshape(tuple(t.dims))
    return onnx.numpy_helper.to_array(t)


def quant_block(w_kn: np.ndarray, bits: int, block: int, symmetric: bool):
    """w_kn (K, N) fp32 (MatMul B layout). Returns packed (N, k_blocks, blob) uint8, scales (N, k_blocks) fp32,
    zero points (N, ceil(k_blocks / kpack)) uint8. Same C++ kernels as onnxruntime's MatMulNBitsQuantizer."""
    from onnxruntime.capi._pybind_state import quantize_matmul_4bits, quantize_matmul_8bits

    K, N = w_kn.shape
    kpack = 8 // bits
    k_blocks = (K + block - 1) // block
    blob = (block + kpack - 1) // kpack
    pad = k_blocks * block - K
    if pad:
        w_kn = np.pad(w_kn, ((0, pad), (0, 0)))
    w_kn = np.ascontiguousarray(w_kn, dtype=np.float32)
    packed = np.zeros((N, k_blocks, blob), dtype=np.uint8)
    zp = np.zeros((N, (k_blocks + kpack - 1) // kpack), dtype=np.uint8)
    scales = np.zeros((N, k_blocks), dtype=np.float32)
    fn = quantize_matmul_8bits if bits == 8 else quantize_matmul_4bits
    fn(packed, w_kn, scales, zp, block, N, K, symmetric)
    return packed, scales, zp


def plan(model: onnx.ModelProto, keep: re.Pattern | None):
    """Find the quantizable weights. Returns {init_name: [(kind, node, extra)]} and the set of node names to drop."""
    g = model.graph
    inits = {t.name: t for t in g.initializer}
    cons = collections.defaultdict(list)
    for n in g.node:
        for i in n.input:
            cons[i].append(n)
    out = {}
    unsupported = set()  # weights with at least one Cast use that cannot be converted
    drop = set()
    for n in g.node:
        if n.op_type != "Cast" or n.input[0] not in inits:
            continue
        t = inits[n.input[0]]
        if t.data_type != TensorProto.FLOAT16 or len(t.dims) != 2 or _attr(n, "to") != TensorProto.FLOAT:
            continue
        if keep is not None and keep.search(t.name):
            continue
        uses, ok, local_drop = [], True, {n.name}
        for c in cons[n.output[0]]:
            if c.op_type == "Transpose" and list(_attr(c, "perm", [])) == [1, 0]:
                local_drop.add(c.name)
                for mm in cons[c.output[0]]:
                    if mm.op_type == "MatMul" and mm.input[1] == c.output[0]:
                        uses.append(("matmul", mm))
                    else:
                        ok = False
            elif (c.op_type == "Gemm" and c.input[1] == n.output[0] and _attr(c, "transB", 0) == 1
                  and _attr(c, "transA", 0) == 0 and _attr(c, "alpha", 1.0) == 1.0 and _attr(c, "beta", 1.0) == 1.0):
                uses.append(("gemm", c))
            else:
                ok = False
        if ok and uses:
            out.setdefault(t.name, []).extend(uses)  # one weight may feed several Casts
            drop |= local_drop
        else:
            unsupported.add(t.name)
    mixed = sorted(set(out) & unsupported)
    if mixed:
        raise ValueError(f"weights with both convertible and unconvertible Cast uses: {mixed}")
    return out, drop


def convert(src: Path, dst: Path, bits: int, block: int, accuracy_level: int, symmetric: bool,
            keep: str | None) -> dict:
    model = onnx.load(str(src), load_external_data=False)
    g = model.graph
    keep_re = re.compile(keep) if keep else None
    targets, drop_nodes = plan(model, keep_re)
    replaced = {}  # node name -> list of new nodes
    dst.parent.mkdir(parents=True, exist_ok=True)
    loc = dst.name + ".data"
    data_path = dst.with_name(loc)
    data_path.unlink(missing_ok=True)
    src_dir = str(src.parent)
    new_inits, pos = [], 0
    stats = collections.Counter()
    used = set()
    for n in g.node:
        if n.name not in drop_nodes:
            used.update(n.input)

    def add_tensor(fh, name, arr, dtype):
        nonlocal pos
        raw = np.ascontiguousarray(arr).tobytes()
        t = TensorProto()
        t.name, t.data_type = name, dtype
        t.dims.extend(arr.shape)
        if len(raw) < 1024:  # small tensors inline (ORT shape inference needs them)
            t.raw_data = raw
        else:
            fh.write(raw)
            for k, v in (("location", loc), ("offset", str(pos)), ("length", str(len(raw)))):
                e = t.external_data.add()
                e.key, e.value = k, v
            t.data_location = TensorProto.EXTERNAL
            pos += len(raw)
        new_inits.append(t)

    with open(data_path, "wb") as fh:
        for t in g.initializer:
            if t.name in targets:
                w = _read(t, src_dir).astype(np.float32)  # (N, K)
                N, K = w.shape
                packed, scales, zp = quant_block(w.T, bits, block, symmetric)
                del w
                qn, sn, zn = t.name + f"_Q{bits}", t.name + "_scales", t.name + "_zp"
                add_tensor(fh, qn, packed, TensorProto.UINT8)
                add_tensor(fh, sn, scales, TensorProto.FLOAT)
                ins = [qn, sn]
                if not symmetric:
                    add_tensor(fh, zn, zp, TensorProto.UINT8)
                    ins.append(zn)
                stats["quantized_weights"] += 1
                stats["quantized_params"] += N * K
                for kind, node in targets[t.name]:
                    attrs = dict(K=K, N=N, bits=bits, block_size=block)
                    if accuracy_level:
                        attrs["accuracy_level"] = accuracy_level
                    nodes = []
                    if kind == "gemm" and len(node.input) > 2 and node.input[2]:
                        mid = node.output[0] + "_nbits"
                        nodes.append(helper.make_node("MatMulNBits", [node.input[0]] + ins, [mid],
                                                      name=node.name + "_nbits", domain="com.microsoft", **attrs))
                        nodes.append(helper.make_node("Add", [mid, node.input[2]], [node.output[0]],
                                                      name=node.name + "_bias"))
                    else:
                        nodes.append(helper.make_node("MatMulNBits", [node.input[0]] + ins, [node.output[0]],
                                                      name=node.name + "_nbits", domain="com.microsoft", **attrs))
                    replaced[node.name] = nodes
                    stats["matmulnbits_nodes"] += 1
                continue
            if t.name not in used:
                stats["dropped_inits"] += 1
                continue
            arr = _read(t, src_dir)
            if t.data_type == TensorProto.FLOAT16 and arr.ndim == 2:
                stats["kept_fp16_params"] += arr.size
            add_tensor(fh, t.name, arr, t.data_type)
    nodes = []
    for n in g.node:
        if n.name in drop_nodes:
            continue
        nodes.extend(replaced.get(n.name, [n]))
    del g.node[:]
    g.node.extend(nodes)
    del g.initializer[:]
    g.initializer.extend(new_inits)
    names = {t.name for t in new_inits}
    keep_inputs = [i for i in g.input if i.name not in names or i.name in {"h", "cb0", "mem", "noise", "codes"}]
    del g.input[:]
    g.input.extend(keep_inputs)
    if not any(o.domain == "com.microsoft" for o in model.opset_import):
        model.opset_import.append(helper.make_opsetid("com.microsoft", 1))
    onnx.save_model(model, str(dst))
    stats["data_bytes"] = pos
    return dict(stats)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", type=Path, required=True)
    ap.add_argument("--dst", type=Path, required=True)
    ap.add_argument("--bits", type=int, choices=(4, 8), default=8)
    ap.add_argument("--block", type=int, default=32)
    ap.add_argument("--accuracy-level", type=int, default=1)
    ap.add_argument("--asym", action="store_true")
    ap.add_argument("--keep", default=None, help="regex: initializers kept fp16-stored (e.g. 'heads\\.|cross\\.')")
    a = ap.parse_args()
    res = convert(a.src.expanduser(), a.dst.expanduser(), a.bits, a.block, a.accuracy_level, not a.asym, a.keep)
    print(res, flush=True)


if __name__ == "__main__":
    main()