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

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

torch.manual_seed(0)
for groups, channels, shape in [(8, 64, (1, 64, 16)), (8, 512, (1, 512, 16)), (8, 1024, (1, 1024, 8)), (8, 2048, (1, 2048, 4)),
                                (8, 512, (1, 512, 4, 4)), (32, 512, (1, 512, 16))]:
    m = torch.nn.GroupNorm(groups, channels).eval()
    with torch.no_grad():
        m.weight.normal_()
        m.bias.normal_()
    x = torch.randn(*shape)
    et = to_edge_transform_and_lower(torch.export.export(m, (x,)), partitioner=[VulkanPartitioner()]).to_executorch()
    out = Runtime.get().load_program(et.buffer).load_method("forward").execute([x])[0]
    ref = m(x)
    print(f"GroupNorm({groups}, {channels}) on {tuple(shape)}: max abs diff {(out - ref).abs().max().item():.2e}")

Output on an H100:

GroupNorm(8, 64) on (1, 64, 16): max abs diff 9.20e+00
GroupNorm(8, 512) on (1, 512, 16): max abs diff 1.05e+01
GroupNorm(8, 1024) on (1, 1024, 8): max abs diff 7.85e+00
GroupNorm(8, 2048) on (1, 2048, 4): max abs diff 8.08e+00
GroupNorm(8, 512) on (1, 512, 4, 4): max abs diff 9.54e-07
GroupNorm(32, 512) on (1, 512, 16): max abs diff 9.56e+00

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 Index(torch.nn.Module):
    def forward(self, x, idx):
        return x[idx]


for dtype in (torch.float32, torch.int64, torch.int32, torch.bool):
    x = (torch.randn(8, 16) * 10).to(dtype)
    idx = torch.tensor([1, 3, 5])
    try:
        edge = to_edge_transform_and_lower(torch.export.export(Index(), (x, idx)), partitioner=[VulkanPartitioner()])
        names = [getattr(n.target, "__name__", "") for n in edge.exported_program().graph.nodes if n.op == "call_function"]
        method = Runtime.get().load_program(edge.to_executorch().buffer).load_method("forward")
        out = method.execute([x, idx])[0]
        print(f"{str(dtype):14s} delegated={names.count('executorch_call_delegate')} ok={torch.equal(out, x[idx])}")
    except Exception as error:
        print(f"{str(dtype):14s} {type(error).__name__}: {str(error).splitlines()[0].split('!')[-1].strip()[:120]}")

Output on an H100:

torch.float32  delegated=1 ok=True
torch.int64    RuntimeError: Could not find ShaderInfo with name index_tensor_buffer_buffer_int32
torch.int32    RuntimeError: Could not find ShaderInfo with name index_tensor_buffer_buffer_int32
torch.bool     delegated=0 ok=True

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 Mask(torch.nn.Module):
    def __init__(self, dtype):
        super().__init__()
        self.register_buffer("mask", (torch.arange(12) % 3 != 0).reshape(1, 12).to(dtype))

    def forward(self, x):
        return self.mask.clone(), x + 1


for dtype in (torch.bool, torch.float32, torch.int32):
    m = Mask(dtype)
    x = torch.randn(1, 12)
    edge = to_edge_transform_and_lower(torch.export.export(m, (x,)), partitioner=[VulkanPartitioner()])
    names = [getattr(n.target, "__name__", "") for n in edge.exported_program().graph.nodes if n.op == "call_function"]
    out = Runtime.get().load_program(edge.to_executorch().buffer).load_method("forward").execute([x])[0]
    print(f"{str(dtype):14s} delegated={names.count('executorch_call_delegate')} got={out.int().tolist()[0]} want={m.mask.int().tolist()[0]}")

Output on an H100:

torch.bool     delegated=1 got=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] want=[0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1]
torch.float32  delegated=1 got=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] want=[0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1]
torch.int32    delegated=1 got=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] want=[0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1]

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 Rms(torch.nn.Module):
    def __init__(self, scale):
        super().__init__()
        self.scale = scale
        self.weight = torch.nn.Parameter(torch.randn(64).to(torch.bfloat16 if scale == "bf16 weight" else torch.float32))

    def forward(self, x, s):
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6)
        return norm * (s if self.scale == "input" else self.weight)


x, s = torch.randn(1, 8, 64), torch.randn(1, 8, 64)
for scale in ("fp32 weight", "bf16 weight", "input"):
    m = Rms(scale).eval()
    try:
        et = to_edge_transform_and_lower(torch.export.export(m, (x, s)), partitioner=[VulkanPartitioner()]).to_executorch()
        out = Runtime.get().load_program(et.buffer).load_method("forward").execute([x, s])[0]
        print(f"scale is {scale}: loads, max abs diff {(out - m(x, s)).abs().max().item():.2e}")
    except Exception as error:
        print(f"scale is {scale}: fails to load, {str(error).split('!')[0].split('Staging.cpp:')[-1][4:]}")

Output on an H100:

scale is fp32 weight: loads, max abs diff 2.38e-07
scale is bf16 weight: fails to load,  (graph.val_is_tref(tensor_data)) is false
scale is input: fails to load,  (graph.val_is_tref(tensor_data)) is false