Vulkan: how much of each LeRobot policy the Vulkan backend takes, and what blocks the rest

import time

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge, to_edge_transform_and_lower


def lowered(model, *inputs):
    """How many Vulkan programs the model lowers to, and which ops stay on the CPU."""
    edge = to_edge_transform_and_lower(torch.export.export(model.eval(), inputs), partitioner=[VulkanPartitioner()])
    nodes = [n for n in edge.exported_program().graph.nodes if n.op == "call_function"]
    names = [getattr(n.target, "__name__", str(n.target)) for n in nodes]
    vulkan = names.count("executorch_call_delegate")
    cpu = sorted({name for name in names if name not in ("executorch_call_delegate", "getitem")})
    return vulkan, cpu


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

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


def case(label, model, *inputs):
    print(f"{label:44s} -> Vulkan programs, CPU ops: {lowered(model, *inputs)}")


# 1. bf16: no op is delegated.
case("Linear float32", torch.nn.Linear(64, 64), torch.randn(4, 64))
case("Linear bfloat16", torch.nn.Linear(64, 64).to(torch.bfloat16), torch.randn(4, 64, dtype=torch.bfloat16))

# 2. A 4D convolution is taken only with batch size 1.
case("Conv2d batch 1", torch.nn.Conv2d(3, 8, 3), torch.randn(1, 3, 32, 32))
case("Conv2d batch 2", torch.nn.Conv2d(3, 8, 3), torch.randn(2, 3, 32, 32))

# 3. Ops with no Vulkan kernel.
x = torch.randn(1, 16, 32)
case("ELU", torch.nn.ELU(), x)
case("Mish (softplus -> log1p)", torch.nn.Mish(), x)
case("var", Fn(lambda t: t.var(dim=-1, keepdim=True)), x)
case("ConvTranspose1d", torch.nn.ConvTranspose1d(16, 8, 4, stride=2, padding=1), x)


class Put(torch.nn.Module):
    def forward(self, t, idx):
        y = t.clone()
        y[idx] = 0.0
        return y


case("index_put (y[idx] = 0)", Put(), torch.randn(8, 32), torch.tensor([1, 3]))


# 4. 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.nn.functional.elu(a)  # elu 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),))).exported_program()
    start = time.time()
    tags = VulkanPartitioner().partition(ep).partition_tags
    print(f"{n:4d} blocks: {len(tags):4d} Vulkan partitions, partition() took {time.time() - start:6.1f} s")

Output:

Linear float32                               -> Vulkan programs, CPU ops: (1, [])
Linear bfloat16                              -> Vulkan programs, CPU ops: (0, ['aten.addmm.default', 'aten.permute_copy.default'])
Conv2d batch 1                               -> Vulkan programs, CPU ops: (1, [])
Conv2d batch 2                               -> Vulkan programs, CPU ops: (0, ['aten.convolution.default'])
ELU                                          -> Vulkan programs, CPU ops: (0, ['aten.elu.default'])
Mish (softplus -> log1p)                     -> Vulkan programs, CPU ops: (2, ['aten.log1p.default'])
var                                          -> Vulkan programs, CPU ops: (0, ['aten.var.correction'])
ConvTranspose1d                              -> Vulkan programs, CPU ops: (0, ['aten.convolution.default'])
index_put (y[idx] = 0)                       -> Vulkan programs, CPU ops: (1, ['aten.index_put.default'])
  50 blocks:   51 Vulkan partitions, partition() took    2.6 s
 100 blocks:  101 Vulkan partitions, partition() took    3.8 s
 200 blocks:  201 Vulkan partitions, partition() took   13.1 s
 400 blocks:  401 Vulkan partitions, partition() took   94.7 s

import sys

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower
from executorch.runtime import Runtime


class Many(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.n = n

    def forward(self, x):
        for i in range(self.n):
            x = x * 1.0001 + 0.001  # 2 dispatches per step
        return x


for n in [int(a) for a in sys.argv[1:]] or [500, 1000, 2000]:
    x = torch.randn(1, 64)
    et = to_edge_transform_and_lower(torch.export.export(Many(n), (x,)), partitioner=[VulkanPartitioner()]).to_executorch()
    method = Runtime.get().load_program(et.buffer).load_method("forward")
    try:
        out = method.execute([x])[0]
        print(n, "steps: ok, max abs diff", (out - Many(n)(x)).abs().max().item())
    except RuntimeError as error:
        print(n, "steps: RuntimeError:", str(error).split("\n")[0][:200])

Output on an H100:

# built with the event tracer on (the default of the Python wheel's build preset)
500 steps: ok, max abs diff 0.0
1000 steps: ok, max abs diff 0.0
1500 steps: RuntimeError: Exception raised from write_timestamp at .../backends/vulkan/runtime/vk_api/QueryPool.cpp:79: (num_queries_ < config_.max_query_count) is false! ...
2000 steps: RuntimeError: (same)
# built with -DEXECUTORCH_ENABLE_EVENT_TRACER=OFF
1500 steps: ok, max abs diff 0.0
4000 steps: ok, max abs diff 0.0