[lslusarczyk]

Summary

This skill tunes a PyTorch module on an Intel GPU.

The default search is TorchAO, Torch Inductor, and eager.
The fastest one is kept.

On an Arc Pro B70, Torch Tweak reported these samples/s.

Module	Shape	Winner	TorchAO	Inductor	eager
Linear(16, 16)	[1, 16]	eager	21516	14340	31091
Linear(128, 128) plus 24 ReLUs	[32, 128]	Inductor	11379	13322	3107
Twelve ReLUs, length 4096	[8, 4096]	TorchAO	27501	24069	7207

Tuning is worth doing.
The winner on a GPU is not one backend.

OpenVINO is optional.
Add it only when the user asks.
On Linear(16, 16) and Linear(2048, 2048) it ran and lost to eager.

One HighestThroughputStrategy covers the run.
inspect walks the submodules.
You set min_depth and max_depth.

Install steps stay in the Torch Tweak install guide.
https://github.com/intel/torch-tweak/blob/main/docs/learn/install.md

Instruction to repeat tests on your own

Repeat the measurement with the script below.
Run one case per process.
A second tune in the same process hits the module that is already tuned.
Read the Selected line and the backend profiled lines.

python winners.py tiny-linear
python winners.py mix
python winners.py fused-elementwise
python winners.py tiny-linear-ov
python winners.py wide-linear-ov

"""One case per process. Pass the case name as argv[1]."""

import sys

import torch
import torch.nn as nn

import torch_tweak.torch as tt
from torch_tweak.torch.backend import (
    OpenVINOBackend,
    TorchAOBackend,
    TorchEagerBackend,
    TorchInductorBackend,
)


class Fused(nn.Module):
    def forward(self, x):
        y = x
        for _ in range(12):
            y = torch.relu(y * 1.01 + 0.1)
        return y


class Mix(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(128, 128)

    def forward(self, x):
        y = self.lin(x)
        for _ in range(24):
            y = torch.relu(y * 1.001 + 0.01)
        return y


CASES = {
    "tiny-linear": (nn.Linear(16, 16), torch.randn(16), [1], False),
    "mix": (Mix(), torch.randn(128), [32], False),
    "fused-elementwise": (Fused(), torch.randn(4096), [8], False),
    "tiny-linear-ov": (nn.Linear(16, 16), torch.randn(16), [1], True),
    "wide-linear-ov": (nn.Linear(2048, 2048), torch.randn(2048), [32], True),
}


def main():
    name = sys.argv[1]
    model, sample, batch_sizes, with_openvino = CASES[name]
    device = "xpu"
    model = model.eval().to(device)
    sample = sample.to(device)
    backends = [TorchAOBackend(), TorchInductorBackend(), TorchEagerBackend()]
    if with_openvino:
        backends.insert(2, OpenVINOBackend())
    strategy = tt.HighestThroughputStrategy(backends=backends).enable_find_max_batch_size(False)
    module = tt.Module(model, name=name, strategy=strategy)
    print(torch.xpu.get_device_name(0), flush=True)
    tt.tune(module, sample, batch_sizes=batch_sizes, device=device)


if __name__ == "__main__":
    main()

Prompt for an agent:

Torch Tweak is installed. An Intel GPU is available.
Save the script as winners.py.
Run these five commands, each in its own process:
python winners.py tiny-linear
python winners.py mix
python winners.py fused-elementwise
python winners.py tiny-linear-ov
python winners.py wide-linear-ov
From each log, read the Selected line and the backend profiled lines.
Make one table for tiny-linear, mix, and fused-elementwise.
Columns: module, shape, winner, TorchAO, Inductor, eager.
Under the table, say whether the winner stayed the same.
Keep OpenVINO out of that table.
For the two -ov runs, say whether OpenVINO beat eager.