#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 把 DepthART ONNX 中的 com.depthart::SelectiveScan 改写为 pulsar2 可编译的纯 ONNX 子图。 背景： 官方 relative_* ONNX 含自定义算子 SelectiveScan，pulsar2 前端直接报 "dont support SelectiveScan opr"。官方只提供 TensorRT 插件，没有 NPU 实现。 推理脚本里的 Loop 改写可以跑通 onnxruntime，但 AX 工具链不支持 Loop 控制流。 做法： 选择性扫描是线性递推 h_t = exp(dt_t*A) * h_{t-1} + dt_t*B_t*u_t，可写成 h_t = exp(L_t) * h_{-1} + sum_{s<=t} exp(L_t - L_s) * dt_s*B_s*u_s 其中 L_t = cumsum_{r<=t}(dt_r*A)。把时间轴切成固定长度 T（默认 14），每个 分块内用 CumSum + 下三角 mask + MatMul 一次算出，块间只传一个状态。所有 指数项的指数都 <= 0，不会溢出。所有算子（CumSum/Exp/Where/MatMul 等）都在 AX 前端支持列表内，不含 Loop/Scan/自定义算子。 校验： - 本工具用 Identity 替换 scan 跑一遍，拿到每个 scan 的 D/G/N/L，再按实际 尺寸生成静态图；时间轴不是 T 整数倍时自动补零并在输出处裁掉。 - 改写后的整模型输出与已验证的 Loop 版本最大绝对误差约 6e-6（FP32）。 - 同时把 4D `Einsum` 换成按 k 的批 MatMul（`b k d l, k c d -> b k c l` 和 `b k r l, k d r -> b k d l`）：pulsar2 的量化执行器不支持 4D Einsum， 会报 Op Execution Error。 用法（在 hs/DepthART 目录下执行）： # 默认：448 模型，块大小 14，输出 onnx/relative_s_448_default_axonnx.onnx ../../venv/bin/python src/tool/depthart_export_axonnx.py # 指定模型/输出/块大小 ../../venv/bin/python src/tool/depthart_export_axonnx.py \\ --model onnx/relative_b_448_default.onnx \\ --output onnx/relative_b_448_default_axonnx.onnx --chunk-size 14 # 改写后用推理脚本验证数值 ../../venv/bin/python src/infer/depthart_infer_onnx.py \\ --model onnx/relative_s_448_default_axonnx.onnx --image /path/a.jpg 注意： - 块内两两指数矩阵形状为 [D*N, T, T]，T 越大显存/内存占用越大； T=14 时 672 通道层约 1M 元素（4 MB），推荐保持默认。 - 输出是标准 ONNX，不再含自定义算子；原模型文件不会被修改。 """ import argparse import logging from pathlib import Path import numpy as np import onnx import onnxruntime as ort from onnx import helper, numpy_helper, TensorProto PROJECT_ROOT = Path(__file__).resolve().parents[2] DEFAULT_MODEL = PROJECT_ROOT / 'onnx' / 'relative_s_448_default.onnx' DOMAIN = 'com.depthart' OP = 'SelectiveScan' EINSUM_EQUATIONS = { 'b k d l, k c d -> b k c l', 'b k r l, k d r -> b k d l', } def collect_runtime_shapes(model): """用 ORT 实际跑一遍，读取 scan 与 Einsum 的静态形状。 scan 会被替换成 Identity（输出形状等于 u，不影响下游形状）；Einsum 正常执行。 返回 (scan_shapes, einsum_shapes)： scan_shapes : 按图顺序的 [(D, G, N, L), ...] einsum_shapes: {节点名: (x_shape, y_shape)} """ nodes = [] scan_outputs, b_aliases = [], {} einsum_x, einsum_y = {}, {} for index, node in enumerate(model.graph.node): if node.domain == DOMAIN and node.op_type == OP: nodes.append(helper.make_node('Identity', [node.input[0]], [node.output[0]], name=f'depthart_stub_scan_{index}')) alias = f'depthart_stub_b_{index}' nodes.append(helper.make_node('Identity', [node.input[3]], [alias], name=f'depthart_stub_bnode_{index}')) scan_outputs.append(node.output[0]) b_aliases[node.output[0]] = alias else: if node.op_type == 'Einsum': x_alias = f'depthart_stub_einsum_x_{index}' y_alias = f'depthart_stub_einsum_y_{index}' nodes.append(helper.make_node('Identity', [node.input[0]], [x_alias], name=f'{x_alias}_node')) nodes.append(helper.make_node('Identity', [node.output[0]], [y_alias], name=f'{y_alias}_node')) einsum_x[node.name] = x_alias einsum_y[node.name] = y_alias nodes.append(node) stub = onnx.ModelProto() stub.CopyFrom(model) del stub.graph.node[:] stub.graph.node.extend(nodes) for name in scan_outputs: stub.graph.output.append(helper.make_tensor_value_info(name, TensorProto.FLOAT, None)) stub.graph.output.append(helper.make_tensor_value_info(b_aliases[name], TensorProto.FLOAT, None)) for x_alias, y_alias in zip(einsum_x.values(), einsum_y.values()): stub.graph.output.append(helper.make_tensor_value_info(x_alias, TensorProto.FLOAT, None)) stub.graph.output.append(helper.make_tensor_value_info(y_alias, TensorProto.FLOAT, None)) meta = model.graph.input[0] shape = [dim.dim_value for dim in meta.type.tensor_type.shape.dim] if not all(shape): raise ValueError(f'输入必须是静态形状: {meta.name} {shape}') session = ort.InferenceSession(stub.SerializeToString()) values = dict(zip([out.name for out in session.get_outputs()], session.run(None, {meta.name: np.zeros(shape, dtype=np.float32)}))) scan_map = {} for name in scan_outputs: u_shape = values[name].shape b_shape = values[b_aliases[name]].shape if len(u_shape) != 3 or len(b_shape) != 4: raise ValueError(f'scan 形状不符合预期: u={u_shape} B={b_shape}') scan_map[name] = (int(u_shape[1]), int(b_shape[1]), int(b_shape[2]), int(u_shape[2])) scan_shapes = [scan_map[node.output[0]] for node in model.graph.node if node.domain == DOMAIN and node.op_type == OP] einsum_shapes = {name: (values[einsum_x[name]].shape, values[einsum_y[name]].shape) for name in einsum_x} return scan_shapes, einsum_shapes class Rewriter: def __init__(self, model, chunk_size, stub_scan=False, scan_method='pairwise'): self.model = model self.chunk_size = chunk_size self.stub_scan = stub_scan self.scan_method = scan_method self.initializers = list(model.graph.initializer) self.names = {tensor.name for tensor in self.initializers} self.initializer_shapes = {tensor.name: tuple(tensor.dims) for tensor in model.graph.initializer} self.shared = {} def build_einsum(self, prefix, node, x_shape, y_shape): """b k d l, k c d -> b k c l （以及 b k r l, k d r -> b k d l）。 两条式子都是按 k 的批 MatMul：W[k,out,in] @ x[b=1,k,in,l]。 替换掉 4D Einsum，避免量化执行器不支持。 """ out_name = node.output[0] w_shape = self.initializer_shapes.get(node.input[1]) if w_shape is None: raise ValueError(f'{node.name}: 第二个输入不是常量权重，暂不支持') if len(w_shape) != 3 or len(x_shape) != 4: raise ValueError(f'{node.name}: Einsum 形状不符合预期: {w_shape} {x_shape}') if int(x_shape[0]) != 1: raise ValueError(f'{node.name}: 仅支持 batch=1，实际 {x_shape}') K, out_dim, in_dim = (int(v) for v in w_shape) length = int(x_shape[3]) if int(x_shape[1]) != K: raise ValueError(f'{node.name}: K 不一致: {x_shape} vs {w_shape}') expected = (1, K, out_dim, length) if tuple(int(v) for v in y_shape) != expected: raise ValueError(f'{node.name}: 输出形状 {y_shape} != {expected}') nodes = [] # 用 Squeeze/Unsqueeze 而不是 Reshape 还原 batch 维：onnxslim 会把 # MatMul 前后的 Reshape 当成冗余节点合并/改名，导致下游张量断链。 nodes.append(helper.make_node( 'Squeeze', [node.input[0], self.shared_const('axes0', np.array([0], np.int64))], [f'{prefix}_x2'])) nodes.append(helper.make_node('MatMul', [node.input[1], f'{prefix}_x2'], [f'{prefix}_y3'])) nodes.append(helper.make_node( 'Unsqueeze', [f'{prefix}_y3', self.shared_const('axes0', np.array([0], np.int64))], [out_name])) return nodes def const(self, name, value): if name not in self.names: self.initializers.append(numpy_helper.from_array(np.asarray(value), name)) self.names.add(name) return name def shared_const(self, key, value): if key not in self.shared: self.shared[key] = self.const(f'depthart_{key}', value) return self.shared[key] def build_logspace(self, prefix, shapes, inputs): """对数空间分块递推：每块 O(T) 而非 O(T^2)，不构造 [D*N,T,T] 两两矩阵。 h_t = exp(L_t) * ( h_prev + cumsum_s( exp(-L_s) * dt_s * B_s * u_s ) ) 其中 L 为块内相对累计衰减（L_0=0，L_t<=0）。要求块内 exp(-L) 不超 U16 量程， 因此 T 要小（T=4 时最大约 e^9≈8e3）。 """ D, G, N, L = shapes u, delta, a, b, c = inputs[:5] d = inputs[5] if len(inputs) > 5 and inputs[5] else None bias = inputs[6] if len(inputs) > 6 and inputs[6] else None attrs = {item.name: item.i for item in self.node_attrs} T = self.chunk_size nodes = [] dt = delta if bias is not None: bias3 = self.const(f'{prefix}_bias3_shape', np.array([1, D, 1], np.int64)) nodes.append(helper.make_node('Reshape', [bias, bias3], [f'{prefix}_bias3'])) nodes.append(helper.make_node('Add', [delta, f'{prefix}_bias3'], [f'{prefix}_dtb'])) dt = f'{prefix}_dtb' if int(attrs.get('delta_softplus', 0)): nodes.append(helper.make_node('Softplus', [dt], [f'{prefix}_dt'])) dt = f'{prefix}_dt' # 补零到 T 的整数倍，避免最后一块长度不足 pad = (-L) % T Lp = L + pad if pad: nodes.append(helper.make_node( 'Concat', [dt, self.const(f'{prefix}_pad_dt', np.zeros((1, D, pad), np.float32))], [f'{prefix}_dt_pad'], axis=2)) nodes.append(helper.make_node( 'Concat', [u, self.const(f'{prefix}_pad_u', np.zeros((1, D, pad), np.float32))], [f'{prefix}_u_pad'], axis=2)) pad_bc = self.const(f'{prefix}_pad_bc', np.zeros((1, G, N, pad), np.float32)) nodes.append(helper.make_node('Concat', [b, pad_bc], [f'{prefix}_b_pad'], axis=3)) nodes.append(helper.make_node('Concat', [c, pad_bc], [f'{prefix}_c_pad'], axis=3)) dt, u, b, c = (f'{prefix}_dt_pad', f'{prefix}_u_pad', f'{prefix}_b_pad', f'{prefix}_c_pad') # dtA = dt * A -> [1,D,N,Lp] a4 = self.const(f'{prefix}_a4_shape', np.array([1, D, N, 1], np.int64)) nodes.append(helper.make_node('Reshape', [a, a4], [f'{prefix}_a4'])) dt3 = self.const(f'{prefix}_dt3_shape', np.array([1, D, 1, Lp], np.int64)) nodes.append(helper.make_node('Reshape', [dt, dt3], [f'{prefix}_dt4'])) nodes.append(helper.make_node('Mul', [f'{prefix}_dt4', f'{prefix}_a4'], [f'{prefix}_dtA'])) b_full = self.expand_groups(b, D, G, N, Lp, f'{prefix}_b', nodes) c_full = self.expand_groups(c, D, G, N, Lp, f'{prefix}_c', nodes) if d is not None: nodes.append(helper.make_node('Reshape', [d, self.const(f'{prefix}_d_shape', np.array([1, D, 1], np.int64))], [f'{prefix}_d3'])) d = f'{prefix}_d3' axes_u2 = self.shared_const('axes_u2', np.array([2], np.int64)) axes_rs = self.shared_const('axes_rs2', np.array([2], np.int64)) axis_t = self.shared_const('axis2', np.array([2], np.int64)) axis_bc = self.shared_const('axis3', np.array([3], np.int64)) h_prev = self.const(f'{prefix}_h0', np.zeros((1, D, N, 1), np.float32)) outputs = [] for k in range(Lp // T): start, end = k * T, (k + 1) * T tag = f'{prefix}_c{k}' starts = self.const(f'{tag}_starts', np.array([start], np.int64)) ends = self.const(f'{tag}_ends', np.array([end], np.int64)) for source, suffix, axis in ((f'{prefix}_dtA', 'dta', axis_bc), (dt, 'dt', axis_t), (u, 'u', axis_t), (b_full, 'b', axis_bc), (c_full, 'c', axis_bc)): nodes.append(helper.make_node('Slice', [source, starts, ends, axis], [f'{tag}_{suffix}'])) # L = cumsum(dtA)（块内，含首步）；h = exp(L) * (h_prev + cumsum(exp(-L)*b)) nodes.append(helper.make_node('CumSum', [f'{tag}_dta', self.shared_const('axis3s', np.array(3, np.int64))], [f'{tag}_L'])) nodes.append(helper.make_node('Exp', [f'{tag}_L'], [f'{tag}_eL'])) nodes.append(helper.make_node('Neg', [f'{tag}_L'], [f'{tag}_negL'])) nodes.append(helper.make_node('Exp', [f'{tag}_negL'], [f'{tag}_eneg'])) # b = dt * u * B；G = cumsum(exp(-Lp) * b) nodes.append(helper.make_node('Unsqueeze', [f'{tag}_dt', axes_u2], [f'{tag}_dt3'])) nodes.append(helper.make_node('Unsqueeze', [f'{tag}_u', axes_u2], [f'{tag}_u3'])) nodes.append(helper.make_node('Mul', [f'{tag}_dt3', f'{tag}_u3'], [f'{tag}_dtu'])) nodes.append(helper.make_node('Mul', [f'{tag}_dtu', f'{tag}_b'], [f'{tag}_bt'])) nodes.append(helper.make_node('Mul', [f'{tag}_eneg', f'{tag}_bt'], [f'{tag}_g'])) nodes.append(helper.make_node('CumSum', [f'{tag}_g', self.shared_const('axis3s', np.array(3, np.int64))], [f'{tag}_G'])) nodes.append(helper.make_node('Add', [f'{tag}_G', h_prev], [f'{tag}_hp'])) nodes.append(helper.make_node('Mul', [f'{tag}_eL', f'{tag}_hp'], [f'{tag}_h'])) nodes.append(helper.make_node('Mul', [f'{tag}_h', f'{tag}_c'], [f'{tag}_hc'])) nodes.append(helper.make_node('ReduceSum', [f'{tag}_hc', axes_rs], [f'{tag}_ysum'], keepdims=0)) if d is not None: nodes.append(helper.make_node('Mul', [d, f'{tag}_u'], [f'{tag}_du'])) nodes.append(helper.make_node('Add', [f'{tag}_ysum', f'{tag}_du'], [f'{tag}_y'])) else: nodes.append(helper.make_node('Identity', [f'{tag}_ysum'], [f'{tag}_y'])) outputs.append(f'{tag}_y') nodes.append(helper.make_node('Slice', [f'{tag}_h', self.const(f'{tag}_hl', np.array([T - 1], np.int64)), self.const(f'{tag}_he', np.array([T], np.int64)), axis_bc], [f'{tag}_hnext'])) h_prev = f'{tag}_hnext' nodes.append(helper.make_node('Concat', outputs, [f'{prefix}_ycat'], axis=2)) if pad: nodes.append(helper.make_node('Slice', [ f'{prefix}_ycat', self.const(f'{prefix}_out_start', np.array([0], np.int64)), self.const(f'{prefix}_out_end', np.array([L], np.int64)), self.shared_const('axis2', np.array([2], np.int64))], [self.out_name])) else: nodes.append(helper.make_node('Identity', [f'{prefix}_ycat'], [self.out_name])) return nodes def build(self, prefix, shapes, inputs): D, G, N, L = shapes u, delta, a, b, c = inputs[:5] d = inputs[5] if len(inputs) > 5 and inputs[5] else None bias = inputs[6] if len(inputs) > 6 and inputs[6] else None attrs = {item.name: item.i for item in self.node_attrs} T = self.chunk_size nodes = [] # dt = softplus(delta + bias) dt = delta if bias is not None: bias3 = self.const(f'{prefix}_bias3_shape', np.array([1, D, 1], np.int64)) nodes.append(helper.make_node('Reshape', [bias, bias3], [f'{prefix}_bias3'])) nodes.append(helper.make_node('Add', [delta, f'{prefix}_bias3'], [f'{prefix}_dtb'])) dt = f'{prefix}_dtb' if int(attrs.get('delta_softplus', 0)): nodes.append(helper.make_node('Softplus', [dt], [f'{prefix}_dt'])) dt = f'{prefix}_dt' # 时间轴补零到 T 的整数倍 pad = (-L) % T Lp = L + pad if pad: nodes.append(helper.make_node( 'Concat', [dt, self.const(f'{prefix}_pad_dt', np.zeros((1, D, pad), np.float32))], [f'{prefix}_dt_pad'], axis=2)) nodes.append(helper.make_node( 'Concat', [u, self.const(f'{prefix}_pad_u', np.zeros((1, D, pad), np.float32))], [f'{prefix}_u_pad'], axis=2)) pad_bc = self.const(f'{prefix}_pad_bc', np.zeros((1, G, N, pad), np.float32)) nodes.append(helper.make_node('Concat', [b, pad_bc], [f'{prefix}_b_pad'], axis=3)) nodes.append(helper.make_node('Concat', [c, pad_bc], [f'{prefix}_c_pad'], axis=3)) dt, u, b, c = (f'{prefix}_dt_pad', f'{prefix}_u_pad', f'{prefix}_b_pad', f'{prefix}_c_pad') # B/C 按组连续复制：[1,G,N,Lp] -> [1,D,N,Lp] b_full = self.expand_groups(b, D, G, N, Lp, f'{prefix}_b', nodes) c_full = self.expand_groups(c, D, G, N, Lp, f'{prefix}_c', nodes) # A -> [1,D,N,1]，dt -> [1,D,N,T]，再摊平成 [D*N,T] 做 CumSum a4 = self.const(f'{prefix}_a_shape', np.array([1, D, N, 1], np.int64)) nodes.append(helper.make_node('Reshape', [a, a4], [f'{prefix}_a4'])) if d is not None: d3 = self.const(f'{prefix}_d_shape', np.array([1, D, 1], np.int64)) nodes.append(helper.make_node('Reshape', [d, d3], [f'{prefix}_d3'])) d = f'{prefix}_d3' mask = np.tril(np.ones((T, T), dtype=bool)).reshape(1, T, T) mask_name = self.shared_const(f'mask_{T}', mask) neg_name = self.shared_const('mask_fill', np.array(-1e4, np.float32)) flat_axis = self.shared_const('axis1', np.array(1, np.int64)) axes_u2 = self.shared_const('axes_u2', np.array([2], np.int64)) axes_u3 = self.shared_const('axes_u3', np.array([3], np.int64)) axes_rs = self.shared_const('axes_rs2', np.array([2], np.int64)) h_prev = self.const(f'{prefix}_h0', np.zeros((1, D, N, 1), np.float32)) outputs = [] for k in range(Lp // T): start, end = k * T, (k + 1) * T tag = f'{prefix}_c{k}' starts = self.const(f'{tag}_starts', np.array([start], np.int64)) ends = self.const(f'{tag}_ends', np.array([end], np.int64)) axis_t = self.shared_const('axis2', np.array([2], np.int64)) axis_bc = self.shared_const('axis3', np.array([3], np.int64)) for source, suffix, axis in ((dt, 'dt', axis_t), (u, 'u', axis_t), (b_full, 'b', axis_bc), (c_full, 'c', axis_bc)): nodes.append(helper.make_node('Slice', [source, starts, ends, axis], [f'{tag}_{suffix}'])) # dtA = dt * A -> [1,D,N,T] -> [D*N,T]，L = cumsum_t(dtA) nodes.append(helper.make_node('Unsqueeze', [f'{tag}_dt', axes_u2], [f'{tag}_dt3'])) nodes.append(helper.make_node('Mul', [f'{tag}_dt3', f'{prefix}_a4'], [f'{tag}_dtA'])) nodes.append(helper.make_node('Reshape', [f'{tag}_dtA', self.const(f'{tag}_dtA2_shape', np.array([D * N, T], np.int64))], [f'{tag}_dtA2'])) nodes.append(helper.make_node('CumSum', [f'{tag}_dtA2', flat_axis], [f'{tag}_L2'])) nodes.append(helper.make_node('Reshape', [f'{tag}_L2', self.const(f'{tag}_L_shape', np.array([1, D, N, T], np.int64))], [f'{tag}_L'])) # 两两指数矩阵 E[t,s] = exp(L_t - L_s)（s>t 置为 -1e4 后 exp 为 0） nodes.append(helper.make_node('Reshape', [f'{tag}_L2', self.const(f'{tag}_Lt_shape', np.array([D * N, T, 1], np.int64))], [f'{tag}_Lt'])) nodes.append(helper.make_node('Reshape', [f'{tag}_L2', self.const(f'{tag}_Ls_shape', np.array([D * N, 1, T], np.int64))], [f'{tag}_Ls'])) nodes.append(helper.make_node('Sub', [f'{tag}_Lt', f'{tag}_Ls'], [f'{tag}_diff'])) nodes.append(helper.make_node('Where', [mask_name, f'{tag}_diff', neg_name], [f'{tag}_diffm'])) nodes.append(helper.make_node('Exp', [f'{tag}_diffm'], [f'{tag}_E'])) # b = dt * u * B，形状 [D*N,1,T]；contrib[t] = sum_{s<=t} E[t,s]*b[s] nodes.append(helper.make_node('Unsqueeze', [f'{tag}_u', axes_u2], [f'{tag}_u3'])) nodes.append(helper.make_node('Mul', [f'{tag}_dt3', f'{tag}_u3'], [f'{tag}_dtu'])) nodes.append(helper.make_node('Mul', [f'{tag}_dtu', f'{tag}_b'], [f'{tag}_bt'])) nodes.append(helper.make_node('Reshape', [f'{tag}_bt', self.const(f'{tag}_bt2_shape', np.array([D * N, 1, T], np.int64))], [f'{tag}_bt2'])) nodes.append(helper.make_node('Mul', [f'{tag}_E', f'{tag}_bt2'], [f'{tag}_eb'])) nodes.append(helper.make_node('ReduceSum', [f'{tag}_eb', self.shared_const('axis2s', np.array([2], np.int64))], [f'{tag}_contrib2'], keepdims=0)) nodes.append(helper.make_node('Reshape', [f'{tag}_contrib2', self.const(f'{tag}_contrib_shape', np.array([1, D, N, T], np.int64))], [f'{tag}_contrib'])) # 块间状态：h_t = exp(L_t) * h_prev + contrib nodes.append(helper.make_node('Exp', [f'{tag}_L'], [f'{tag}_expL'])) nodes.append(helper.make_node('Mul', [f'{tag}_expL', h_prev], [f'{tag}_carry'])) nodes.append(helper.make_node('Add', [f'{tag}_carry', f'{tag}_contrib'], [f'{tag}_h'])) # y = sum_n(h*C) + D*u nodes.append(helper.make_node('Mul', [f'{tag}_h', f'{tag}_c'], [f'{tag}_hc'])) nodes.append(helper.make_node('ReduceSum', [f'{tag}_hc', axes_rs], [f'{tag}_ysum'], keepdims=0)) if d is not None: nodes.append(helper.make_node('Mul', [d, f'{tag}_u'], [f'{tag}_du'])) nodes.append(helper.make_node('Add', [f'{tag}_ysum', f'{tag}_du'], [f'{tag}_y'])) else: nodes.append(helper.make_node('Identity', [f'{tag}_ysum'], [f'{tag}_y'])) outputs.append(f'{tag}_y') nodes.append(helper.make_node('Gather', [f'{tag}_h', self.const(f'{tag}_last', np.array(T - 1, np.int64))], [f'{tag}_hlast'], axis=3)) nodes.append(helper.make_node('Unsqueeze', [f'{tag}_hlast', axes_u3], [f'{tag}_hnext'])) h_prev = f'{tag}_hnext' nodes.append(helper.make_node('Concat', outputs, [f'{prefix}_ycat'], axis=2)) if pad: nodes.append(helper.make_node('Slice', [ f'{prefix}_ycat', self.const(f'{prefix}_out_start', np.array([0], np.int64)), self.const(f'{prefix}_out_end', np.array([L], np.int64)), self.shared_const('axis2', np.array([2], np.int64))], [self.out_name])) else: nodes.append(helper.make_node('Identity', [f'{prefix}_ycat'], [self.out_name])) return nodes def expand_groups(self, x, D, G, N, L, prefix, nodes): gs = D // G r1_shape = self.const(f'{prefix}_r1_shape', np.array([1, G, 1, N, L], np.int64)) nodes.append(helper.make_node('Reshape', [x, r1_shape], [f'{prefix}_r1'])) target = self.const(f'{prefix}_tgt', np.array([1, G, gs, N, L], np.int64)) nodes.append(helper.make_node('Expand', [f'{prefix}_r1', target], [f'{prefix}_e'])) r2_shape = self.const(f'{prefix}_r2_shape', np.array([1, D, N, L], np.int64)) nodes.append(helper.make_node('Reshape', [f'{prefix}_e', r2_shape], [f'{prefix}_bc4'])) return f'{prefix}_bc4' def rewrite(self): shapes, einsum_shapes = collect_runtime_shapes(self.model) if not shapes: raise ValueError('模型中没有 SelectiveScan 节点，无需改写') new_nodes, index, einsum_index = [], 0, 0 for node in self.model.graph.node: if node.domain == DOMAIN and node.op_type == OP: if self.stub_scan: # 仅用于测 NPU 耗时的对照模型：扫描直接透传 u，精度无意义。 stub_node = helper.make_node('Identity', [node.input[0]], [node.output[0]], name=f'depthart_stub_scan_{index}') new_nodes.append(stub_node) index += 1 continue self.node_attrs = node.attribute self.out_name = node.output[0] builder = self.build_logspace if self.scan_method == 'log' else self.build nodes = builder(f'depthart_pscan_{index}', shapes[index], node.input) new_nodes.extend(nodes) index += 1 elif node.op_type == 'Einsum': equation = None for attr in node.attribute: if attr.name == 'equation': equation = attr.s.decode() if equation not in EINSUM_EQUATIONS: raise ValueError(f'{node.name}: 不支持的 Einsum 方程 {equation!r}') x_shape, y_shape = einsum_shapes[node.name] nodes = self.build_einsum(f'depthart_einsum_{einsum_index}', node, x_shape, y_shape) new_nodes.extend(nodes) einsum_index += 1 else: new_nodes.append(node) out = onnx.ModelProto() out.CopyFrom(self.model) del out.graph.node[:] for position, node in enumerate(new_nodes): if not node.name: node.name = f'depthart_gen_{position}' out.graph.node.extend(new_nodes) out.graph.initializer.extend(self.initializers[len(self.model.graph.initializer):]) return out def main(): parser = argparse.ArgumentParser( description='把 DepthART SelectiveScan 改写为 pulsar2 可编译的纯 ONNX 子图') parser.add_argument('--model', type=Path, default=DEFAULT_MODEL, help=f'输入 ONNX，默认 {DEFAULT_MODEL.relative_to(PROJECT_ROOT)}') parser.add_argument('--output', type=Path, default=PROJECT_ROOT / 'onnx' / 'relative_s_448_default_axonnx.onnx', help='输出 ONNX（不会覆盖输入模型）') parser.add_argument('--chunk-size', type=int, default=14, help='分块长度 T，默认 14；越大显存占用越高') parser.add_argument('--stub-scan', action='store_true', help='把 SelectiveScan 替换为 Identity，生成仅用于测 NPU 耗时的对照模型（精度无意义）') parser.add_argument('--scan-method', choices=('pairwise', 'log'), default='pairwise', help='pairwise=两两指数矩阵(默认)；log=对数空间递推，无 T^2 矩阵，T 需小') args = parser.parse_args() logging.basicConfig(level=logging.INFO, format='%(message)s') if args.chunk_size < 1: parser.error('--chunk-size 必须为正数') if not args.model.is_file(): parser.error(f'--model 不存在: {args.model}') if args.output.resolve() == args.model.resolve(): parser.error('--output 不能与 --model 相同') model = onnx.load(str(args.model)) if not any(node.domain == DOMAIN and node.op_type == OP for node in model.graph.node): raise ValueError(f'{args.model} 中没有 {DOMAIN}::{OP} 节点') rewriter = Rewriter(model, args.chunk_size, args.stub_scan, args.scan_method) rewritten = rewriter.rewrite() args.output.parent.mkdir(parents=True, exist_ok=True) onnx.save(rewritten, str(args.output)) logging.info('已写出标准 ONNX: %s（节点数 %d）', args.output, len(rewritten.graph.node)) logging.info('下一步：把 pulsar2 配置的 input 指向该文件；' '可先用 src/infer/depthart_infer_onnx.py 验证数值。') if __name__ == '__main__': main()