M10-3 | Fallback 報告：哪些 op 跑在 VPU、哪些切回 CPU

Open In Colab

VPU 只有 VPU_OPS 裡的 op 有 kernel。MyVPUQuantizer 不標其他 op，所以它們在 Q/DQ 圖上用 fp32 跑；部署時這些 op 得交給 CPU，圖因此被切成好幾段。

這本 notebook 寫一個極簡的 partitioner partition，讀 convert_pt2e 的 Q/DQ 圖，回答部署前最想知道的三件事：

問題	規則
每個 op 在哪裡跑？	在 `VPU_OPS` 裡、且輸入來自 dequantize 或另一個 VPU op → VPU；其餘 → CPU
切成幾個 VPU 子圖？	VPU op 之間經過 q → dq（int8 留在 VPU 上）或直接相連，就算同一個子圖
邊界在哪、搬多少？	q/dq 一邊是 VPU、另一邊是 CPU（或模型的輸入輸出）的地方

A. TinyConvNet：只有 adaptive_avg_pool2d 不在清單上。
B. OpZooNet（M09-3 的多 op 模型）：一半以上的 op 不在清單上，圖被切碎。

from dataclasses import dataclass

import graphviz
import torch
from torch.fx import Node

from common.pt2e import quiet
from common.graph_utils import short
from common.models import OpZooNet, TinyConvNet, example_inputs
from common.viz import FONT, table
from vpu_quantizer import VPU_OPS, vpu_ptq

quiet()
_ = torch.manual_seed(0)
QD = torch.ops.quantized_decomposed
Q_T, DQ_T, DQ_C = QD.quantize_per_tensor.default, QD.dequantize_per_tensor.default, QD.dequantize_per_channel.default

partitioner

先寫一個越過 q/dq 找 activation 來源的小函式，再寫 partition。同一個子圖的 op 用 union-find 合併（parent 指向同一組的根）：

def source(a: Node) -> Node | None:
    """越過「q → dq」找到 activation 真正的來源；weight／bias（get_attr、dequantize_per_channel）回傳 None。"""
    if a.op == "get_attr" or (a.op == "call_function" and a.target == DQ_C):
        return None
    if a.op == "call_function" and a.target == DQ_T:
        q = a.args[0]
        assert isinstance(q, Node) and isinstance(q.args[0], Node)
        return q.args[0]
    return a


def partition(gm):
    """回傳 (ops, device, group)：運算 node 依執行順序、每個 node 的裝置、VPU node 屬於第幾個子圖。"""
    ops = [n for n in gm.graph.nodes if n.op == "call_function" and n.target not in (Q_T, DQ_T, DQ_C)]
    device, parent = {}, {}

    def find(n):
        while parent[n] is not n:
            n = parent[n]
        return n

    for n in ops:
        ins = [a for a in n.all_input_nodes if source(a) is not None]
        # 輸入都經過 dequantize（int8 進來）或來自另一個 VPU op（融合的 relu）→ VPU
        quantized = all(a.target == DQ_T or device.get(a) == "VPU" for a in ins)
        device[n] = "VPU" if n.target in VPU_OPS and quantized else "CPU"
        if device[n] == "VPU":
            parent[n] = n
            for a in ins:
                src = source(a)
                if device.get(src) == "VPU":  # 前一個 VPU op 的 int8 輸出直接交給這個 op：同一個子圖
                    parent[find(n)] = find(src)
    roots = list(dict.fromkeys(find(n) for n in ops if device[n] == "VPU"))
    group = {n: roots.index(find(n)) for n in ops if device[n] == "VPU"}
    return ops, device, group


def crossings(gm, ops, device):
    """跨裝置的邊界：[(方向, 來源 node, 交給誰, shape, int8 bytes)]。模型的輸入輸出算在 CPU（host）那一側。"""
    where = lambda n: device.get(n, "CPU")  # noqa: E731
    out = []
    for n in ops + [next(n for n in gm.graph.nodes if n.op == "output")]:
        for a in n.all_input_nodes:
            src = source(a)
            if src is not None and where(src) != where(n):
                val = src.meta.get("val")
                shape = tuple(val.shape) if isinstance(val, torch.Tensor) else ()
                out.append((f"{where(src)}→{where(n)}", src, n, shape, int(torch.Size(shape).numel())))
    return out

分區圖：只畫運算 op（不畫 q/dq），每個 VPU 子圖一個框，CPU op 在框外，跨裝置的邊標出 int8 bytes。畫圖的程式只服務這本 notebook，摺起來：

畫分區圖的程式

def draw_partition(gm, title):
    ops, device, group = partition(gm)
    cross = {(c[1], c[2]): c[4] for c in crossings(gm, ops, device)}
    g = graphviz.Digraph(format="svg")
    g.attr(rankdir="TB", bgcolor="white", fontname=FONT, label=title, labelloc="t", fontsize="14", nodesep="0.3", ranksep="0.35")
    g.attr("node", fontname=FONT, fontsize="11", shape="box", style="filled,rounded", margin="0.12,0.05", height="0.3")
    g.attr("edge", fontname=FONT, fontsize="9", color="#495057", arrowsize="0.6")
    io = dict(shape="ellipse", style="filled", fillcolor="#E7F5FF", color="#1C7ED6")
    g.node("x", "x", **io)
    g.node("output", "output", **io)
    for k in sorted(set(group.values())):
        c = graphviz.Digraph(name=f"cluster_{k}")  # graphviz 以 cluster_ 開頭的子圖畫成一個框
        c.attr(label=f"VPU 子圖 #{k}", style="rounded", color="#2B8A3E", fontcolor="#2B8A3E", fontsize="11")
        for n in ops:
            if group.get(n) == k:
                c.node(n.name, short(n.target), fillcolor="#EBFBEE", color="#2B8A3E")
        g.subgraph(c)
    for n in ops:
        if device[n] == "CPU":
            g.node(n.name, f"{short(n.target)}\n（CPU，fp32）", fillcolor="#FFE3E3", color="#C92A2A")
    out_node = next(n for n in gm.graph.nodes if n.op == "output")
    for n in ops + [out_node]:
        for a in n.all_input_nodes:
            src = source(a)
            if src is None:
                continue
            dst = "output" if n is out_node else n.name
            if (src, n) in cross:
                g.edge(src.name if src.op != "placeholder" else "x", dst, label=f"{cross[(src, n)]} B",
                       color="#C92A2A", fontcolor="#C92A2A", style="dashed")
            else:
                g.edge(src.name, dst)
    return g

A. TinyConvNet

Q/DQ 圖來自 vpu_ptq（M10-2）：torch.export.export（batch 1，static shape）→ prepare_pt2e(…, MyVPUQuantizer()) → 校準 → convert_pt2e。

x = example_inputs(1)
tiny = vpu_ptq(TinyConvNet().eval(), x)
draw_partition(tiny, "TinyConvNet：綠框是 VPU 子圖，紅色虛線是跨裝置的邊界")

@dataclass
class Summary:
    vpu: int          # VPU op 數
    cpu: int          # CPU op 數
    sizes: list[int]  # 每個 VPU 子圖的 op 數
    n_cross: int      # 邊界數
    bytes: int        # 每次推論跨裝置搬移的 int8 bytes


def boundary_table(gm):
    ops, device, group = partition(gm)
    cross = crossings(gm, ops, device)
    rows = [[d, s.name, n.name if n.op != "output" else "output", str(shape), b] for d, s, n, shape, b in cross]
    sizes = [sum(g == k for g in group.values()) for k in range(len(set(group.values())))]
    n_vpu = sum(d == "VPU" for d in device.values())
    summary = Summary(vpu=n_vpu, cpu=len(ops) - n_vpu, sizes=sizes, n_cross=len(cross), bytes=sum(c[4] for c in cross))
    return rows, summary


rows, s = boundary_table(tiny)
print(f"VPU op {s.vpu} 個、CPU op {s.cpu} 個；VPU 子圖 {len(s.sizes)} 個（各含 {s.sizes} 個 op）；"
      f"邊界 {s.n_cross} 個，每次推論搬 {s.bytes} bytes（以 int8 計）")
table(rows, ["方向", "tensor（產生它的 node）", "交給", "shape", "int8 bytes"])

VPU op 13 個、CPU op 1 個；VPU 子圖 2 個（各含 [11, 2] 個 op）；邊界 4 個，每次推論搬 1316 bytes（以 int8 計）

方向	tensor（產生它的 node）	交給	shape	int8 bytes
CPU→VPU	x	conv2d	(1, 3, 16, 16)	768
VPU→CPU	relu_3	adaptive_avg_pool2d	(1, 32, 4, 4)	512
CPU→VPU	adaptive_avg_pool2d	flatten	(1, 32, 1, 1)	32
VPU→CPU	linear	output	(1, 4)	4

assert s.cpu == 1 and len(s.sizes) == 2 and s.n_cross == 4

adaptive_avg_pool2d 不在 VPU_OPS，MyVPUQuantizer 沒有標它（M10-1 D）。它在 Q/DQ 圖上夾在 dq 與 q 之間用 fp32 跑，partition 把它歸給 CPU，前後的 VPU op 因此分成兩個子圖。

邊界上的 q/dq 都是 quantizer 的標籤留下的：它前面的 dq 是 relu_3 的輸出標籤（VPU→CPU），後面的 q 是 flatten 的輸入邊標籤（CPU→VPU）。每個邊界是一次 int8 ↔︎ fp32 轉換加一次裝置間搬移；模型的輸入、輸出也各算一個。

B. OpZooNet：一半以上的 op 不在清單上

sigmoid、leaky_relu、mul、avg_pool2d、mean、matmul、softmax 都不在 VPU_OPS：

zoo = vpu_ptq(OpZooNet().eval(), x)
draw_partition(zoo, "OpZooNet：圖被切成好幾個單一 op 的子圖")

rows, s = boundary_table(zoo)
lone = sum(k == 1 for k in s.sizes)
print(f"VPU op {s.vpu} 個、CPU op {s.cpu} 個；VPU 子圖 {len(s.sizes)} 個，其中 {lone} 個只含 1 個 op；"
      f"邊界 {s.n_cross} 個，每次推論搬 {s.bytes} bytes（以 int8 計）")
assert s.cpu > s.vpu and lone == len(s.sizes)

VPU op 5 個、CPU op 7 個；VPU 子圖 5 個，其中 5 個只含 1 個 op；邊界 11 個，每次推論搬 18970 bytes（以 int8 計）

每個 VPU 子圖都只有 1 個 op。例如 max_pool2d 前面是 mul、後面是 avg_pool2d，兩邊都在 CPU，它進出 VPU 各一次只做一個 pooling。

這種單一 op 的子圖，搬移與轉換的成本可能比在 CPU 上直接算還高（推論，本專案沒有實機量測）；partitioner 可以再加一條「子圖太小就留在 CPU」的規則。

小結

quantizer 與 partitioner 讀同一份 VPU_OPS：quantizer 不標的 op，就是 partitioner 切回 CPU 的 op。兩邊一致，CPU 那段才會拿到 fp32、VPU 那段才會拿到 int8。
這和 M08-3（XnnpackPartitioner 不收 cumsum）、M09-3（ORT 的 QDQ fusion 清單）是同一件事：模型能有多少比例跑在加速器上，要看 op 清單，而不是只看 quantizer 標了什麼。
部署前先看這張分區圖：子圖數、邊界數、搬移量，都是在選模型架構時就能調整的（例如把 sigmoid 換成清單內的 relu）。