to_dense_batch.py

import torch
from torch.utils.benchmark import Compare, Timer

from torch_geometric.utils import to_dense_batch

NUM_STEPS = 5
torch.set_num_threads(8)

configs = [
    {"batch_size": 1024, "max_num_nodes": 64},
    {"batch_size": 1024, "max_num_nodes": 512},
    {"batch_size": 1024, "max_num_nodes": 2048},
]

devices = ["cpu", "cuda"]

results = []

for cfg in configs:
    bs = cfg["batch_size"]
    mn = cfg["max_num_nodes"]
    label = f"bs={bs} mn={mn}"

    cpu_data = []
    for _ in range(NUM_STEPS):
        batch = torch.repeat_interleave(
            torch.arange(bs),
            torch.randint(1, 2 * mn + 1, (bs,)),
        )
        x = torch.randn(batch.size(0), 128)
        cpu_data.append((x, batch))

    for device in devices:
        data = [(x.to(device), b.to(device)) for x, b in cpu_data]

        stmt = """\
for x, batch in data:
    fn(x, batch, batch_size=bs, max_num_nodes=mn)
"""
        globs = {"fn": to_dense_batch, "data": data, "bs": bs, "mn": mn}

        overflow_mn = max(1, mn // 2)
        cases = [
            ("eager", stmt, globs),
            ("overflow", stmt, {**globs, "mn": overflow_mn}),
        ]

        compiled = torch.compile(to_dense_batch)
        for x, batch in data[:3]:
            compiled(x, batch, batch_size=bs, max_num_nodes=mn)
        cases.append(("torch.compile", stmt, {**globs, "fn": compiled}))

        for sub_label, s, g in cases:
            t = Timer(
                stmt=s,
                globals=g,
                label=label,
                sub_label=sub_label,
                description=device,
                num_threads=8,
            )
            results.append(t.blocked_autorange(min_run_time=1))

        del data

compare = Compare(results)
compare.colorize()
compare.print()

to_dense_batch.txt

[---------- bs=1024 mn=64 ----------]
                     |  cpu   |  cuda
8 threads: --------------------------
      eager          |  34.7  |  2.8
      overflow       |  16.3  |  2.5
      torch.compile  |  33.3  |  1.4

Times are in milliseconds (ms).

[---------- bs=1024 mn=512 ----------]
                     |   cpu   |  cuda
8 threads: ---------------------------
      eager          |  290.3  |  16.8
      overflow       |  214.0  |  12.4
      torch.compile  |  294.4  |  15.5

Times are in milliseconds (ms).

[---------- bs=1024 mn=2048 ----------]
                     |   cpu    |  cuda
8 threads: ----------------------------
      eager          |  1096.5  |  67.1
      overflow       |   827.8  |  49.6
      torch.compile  |  1062.9  |  63.0

Times are in milliseconds (ms).

to_dense_batch_master.txt

[---------- bs=1024 mn=64 ----------]
                     |  cpu   |  cuda
8 threads: --------------------------
      eager          |  44.2  |  4.3
      overflow       |   9.2  |  2.9
      torch.compile  |  41.9  |  3.7

Times are in milliseconds (ms).

[---------- bs=1024 mn=512 ----------]
                     |   cpu   |  cuda
8 threads: ---------------------------
      eager          |  465.1  |  24.4
      overflow       |  252.9  |  14.7
      torch.compile  |  455.5  |  23.8

Times are in milliseconds (ms).

[---------- bs=1024 mn=2048 ----------]
                     |   cpu    |  cuda
8 threads: ----------------------------
      eager          |  1670.7  |  95.9
      overflow       |   941.0  |  56.2
      torch.compile  |  1651.8  |  93.2

Times are in milliseconds (ms).