Arm: how much of each LeRobot policy the Arm backend (VGF, Ethos-U85) takes, and what blocks the rest

import time

import torch
from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner
from executorch.backends.arm.quantizer import EthosUQuantizer, get_symmetric_quantization_config
from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner
from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e

CONFIG = EdgeCompileConfig(_check_ir_validity=False)
SKIP = ("executorch_call_delegate", "getitem", "quantize_per", "dequantize_per")


class Fn(torch.nn.Module):
    def __init__(self, fn, *modules):
        super().__init__()
        self.fn = fn
        self.layers = torch.nn.ModuleList(modules)

    def forward(self, x):
        return self.fn(self, x)


def summary(edge):
    """How many Arm programs the model lowers to, and which ops stay on the CPU."""
    names = [str(n.target).split("::")[-1].split("(")[0]
             for n in edge.exported_program().graph.nodes if n.op == "call_function"]
    return names.count("executorch_call_delegate"), sorted({n for n in names if not any(s in n for s in SKIP)})


def root(error):
    while error.__cause__ is not None:
        error = error.__cause__
    return f"{type(error).__name__}: {str(error).splitlines()[0][:140]}"


def vgf(model, x, spec=None):
    try:
        ep = torch.export.export(model.eval(), (x,))
        edge = to_edge_transform_and_lower(ep, partitioner=[VgfPartitioner(VgfCompileSpec(spec))], compile_config=CONFIG)
        return summary(edge)
    except Exception as error:
        return root(error)


def quantizer():
    q = EthosUQuantizer(EthosUCompileSpec("ethos-u85-256"))
    q.set_global(get_symmetric_quantization_config(is_per_channel=True))
    return q


def u85(model, x):
    m = prepare_pt2e(torch.export.export(model.eval(), (x,)).module(check_guards=False), quantizer())
    m(x)
    ep = torch.export.export(convert_pt2e(m), (x,))
    edge = to_edge_transform_and_lower(ep, partitioner=[EthosUPartitioner(EthosUCompileSpec("ethos-u85-256"))],
                                       compile_config=CONFIG)
    return summary(edge)


def case(label, result):
    print(f"{label:46s} -> {result}", flush=True)


# 1. The Arm quantizer fails on the module that ExportedProgram.module() returns by default.
conv_mean = Fn(lambda s, t: s.layers[0](t).mean(dim=(2, 3)), torch.nn.Conv2d(3, 8, 3))
ep = torch.export.export(conv_mean.eval(), (torch.randn(1, 3, 16, 16),))
for check_guards in (True, False):
    try:
        prepare_pt2e(ep.module(check_guards=check_guards), quantizer())
        result = "OK"
    except Exception as error:
        result = root(error)
    case(f"1. module(check_guards={check_guards}) -> prepare_pt2e", result)


# 2. An int64 index that is a tensor constant (not a buffer) fails the whole lowering.
class Buffer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("idx", torch.tensor([0, 2]))

    def forward(self, x):
        return x[:, self.idx]


class Constant(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.idx = torch.tensor([0, 2])  # a plain tensor attribute, which export lifts as a constant

    def forward(self, x):
        return x[:, self.idx]


x = torch.randn(1, 4, 8)
case("2. VGF x[:, idx], idx an int64 buffer", vgf(Buffer(), x))
case("2. VGF x[:, idx], idx an int64 constant", vgf(Constant(), x))


# 3. The Ethos-U quantizer does not annotate Mish, so it stays in float on the CPU.
def block(act):
    return torch.nn.Sequential(torch.nn.Conv1d(16, 16, 3, padding=1), act, torch.nn.Conv1d(16, 16, 3, padding=1))


x = torch.randn(1, 16, 32)
for label, act in (("SiLU", torch.nn.SiLU()), ("ELU", torch.nn.ELU()), ("Mish", torch.nn.Mish())):
    case(f"3. U85 Conv1d -> {label} -> Conv1d", u85(block(act), x))

# 4. Ops with no TOSA lowering stay on the CPU and split the graph.
case("4. VGF ConvTranspose1d", vgf(torch.nn.ConvTranspose1d(16, 8, 4, stride=2, padding=1), x))
case("4. VGF ConvTranspose2d (for comparison)", vgf(torch.nn.ConvTranspose2d(16, 8, 4, stride=2, padding=1), x.unsqueeze(-1)))
s = torch.randn(64, 16)
case("4. VGF argmax (for comparison)", vgf(Fn(lambda _, t: t.argmax(-1).to(torch.int32)), s))
case("4. VGF argmin", vgf(Fn(lambda _, t: t.argmin(-1).to(torch.int32)), s))
case("4. VGF topk(4) of a [64, 16] score", vgf(Fn(lambda _, t: t.topk(4).values), s))
case("4. VGF topk(8) of a [64, 16] score", vgf(Fn(lambda _, t: t.topk(8).values), s))
case("4. VGF rows picked by topk(8) indices", vgf(Fn(lambda _, t: t[t.sum(-1).topk(8).indices]), s))
case("4. VGF index % 7 (int64 math)", vgf(Fn(lambda _, t: (t.argmax(-1) % 7).to(torch.float32)), s))

# 5. bf16 is not in the default VGF TOSA spec, so bf16 ops stay on the CPU without a warning.
lin16 = torch.nn.Linear(64, 64).to(torch.bfloat16)
x16 = torch.randn(4, 64, dtype=torch.bfloat16)
case("5. VGF bf16 Linear, default VgfCompileSpec()", vgf(lin16, x16))
case("5. VGF bf16 Linear, TOSA spec with +bf16", vgf(lin16, x16, "TOSA-1.0+FP+INT+int4+int16+bf16"))


# 6. Partitioning time grows fast with graph size when CPU ops sit on side branches.
class Blocks(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.layers = torch.nn.ModuleList(torch.nn.Linear(32, 32) for _ in range(n))

    def forward(self, x):
        for layer in self.layers:
            a = layer(x)
            x = a * torch.sigmoid(a) + torch.atan2(a, a + 2)  # atan2 stays on the CPU
        return x


for n in (50, 100, 200, 400):
    ep = to_edge(torch.export.export(Blocks(n).eval(), (torch.randn(1, 32),)), compile_config=CONFIG).exported_program()
    start = time.time()
    tags = VgfPartitioner(VgfCompileSpec()).partition(ep).partition_tags
    print(f"6. {n:4d} blocks: {len(tags):4d} VGF partitions, partition() took {time.time() - start:6.1f} s", flush=True)

Output:

1. module(check_guards=True) -> prepare_pt2e   -> ExportPassBaseError: call_module is not supported.
1. module(check_guards=False) -> prepare_pt2e  -> OK
2. VGF x[:, idx], idx an int64 buffer          -> (1, [])
2. VGF x[:, idx], idx an int64 constant        -> AssertionError: index.Tensor requires index dtype must be int32
3. U85 Conv1d -> SiLU -> Conv1d                -> (1, [])
3. U85 Conv1d -> ELU -> Conv1d                 -> (1, [])
3. U85 Conv1d -> Mish -> Conv1d                -> (2, ['exp', 'gt.Scalar', 'log1p', 'mul.Tensor', 'tanh', 'where.self'])
4. VGF ConvTranspose1d                         -> (0, ['convolution'])
4. VGF ConvTranspose2d (for comparison)        -> (1, [])
4. VGF argmax (for comparison)                 -> (1, [])
4. VGF argmin                                  -> (0, ['_to_dim_order_copy', 'argmin'])
4. VGF topk(4) of a [64, 16] score             -> (1, [])
4. VGF topk(8) of a [64, 16] score             -> (0, ['topk'])
4. VGF rows picked by topk(8) indices          -> (2, ['_to_dim_order_copy', 'topk'])
4. VGF index % 7 (int64 math)                  -> (0, ['_to_dim_order_copy', 'argmax', 'remainder.Scalar'])
5. VGF bf16 Linear, default VgfCompileSpec()   -> (0, ['linear'])
5. VGF bf16 Linear, TOSA spec with +bf16       -> (1, [])
6.   50 blocks:   51 VGF partitions, partition() took    0.4 s
6.  100 blocks:  101 VGF partitions, partition() took    2.2 s
6.  200 blocks:  201 VGF partitions, partition() took   18.9 s
6.  400 blocks:  401 VGF partitions, partition() took  133.6 s