What does torch.compile do? Find out in Python · torch.compile ทำอะไร? มาดูด้วย Python · torch.compile 做了什么？用 Python 看个明白

1 check_setup.py: what it needs · check_setup.py: สิ่งที่ต้องมี · check_setup.py：它需要什么

In the Ubuntu terminalในเทอร์มินัลของ Ubuntu在 Ubuntu 终端里

sudo apt update
sudo apt install g++ python3-dev python3-venv
mkdir ~/compile_lab
cd ~/compile_lab
python3 -m venv venv
source venv/bin/activate
pip install --index-url https://download.pytorch.org/whl/cpu torch
pip install numpy

not run by us: on a Mac, instead of the lines aboveเราไม่ได้รันเอง: บนแมค ใช้แทนบรรทัดข้างบน我们没有运行过：在 Mac 上，代替上面几行

xcode-select --install
mkdir ~/compile_lab
cd ~/compile_lab
python3 -m venv venv
source venv/bin/activate
pip install torch numpy

check_setup.py

# Is everything here that torch.compile needs?
import shutil
import sys
import sysconfig
from pathlib import Path

import torch

print("Python ", sys.version.split()[0])
print("PyTorch", torch.__version__)
print("GPU    ", torch.cuda.is_available())

# Inductor writes C++ for a CPU, so a C++ compiler has to do the last step.
for compiler in ["g++", "clang++", "cl"]:
    print(f"{compiler:8}", "found" if shutil.which(compiler) else "not found")

# The C++ that Inductor writes has #include <Python.h> in it.
header = Path(sysconfig.get_paths()["include"]) / "Python.h"
print("Python.h", "found" if header.exists() else "not found")

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

Python  3.14.4
PyTorch 2.15.0+cpu
GPU     False
g++      found
clang++  not found
cl       not found
Python.h found

torch._inductor.exc.InductorError: RuntimeError: Compiler: cl is not found.

2 A Python class that wraps a function · คลาส Python ที่ห่อฟังก์ชัน · 一个包住函数的 Python 类

toy_compile.py

# A pretend torch.compile in plain Python: a wrapper, a guard and a cache.
# Nothing here is PyTorch. It only shows the shape of the idea.
import time


class LittleCompile:
    def __init__(self, program):
        self.program = program
        self.cache = {}  # one finished version for each guard

    # __call__ lets us use the object as if it were a function.
    def __call__(self, numbers):
        # The guard: what this version was made for.
        guard = (len(numbers), type(numbers[0]).__name__)
        if guard in self.cache:
            print(guard, "guard matches: reuse")
        else:
            print(guard, "new: compile (slow)")
            time.sleep(0.5)  # stands in for the real work of compiling
            self.cache[guard] = self.program
        return self.cache[guard](numbers)


def double_all(numbers):
    return [n * 2 for n in numbers]


def main():
    compiled = LittleCompile(double_all)
    print("wrapped, and nothing is compiled yet:", compiled.cache)
    print(compiled([1, 2, 3]))
    print(compiled([4, 5, 6]))
    print(compiled([1, 2, 3, 4]))
    print(compiled([1.5, 2.5, 3.5]))
    print(compiled([7, 8, 9]))
    print("versions kept:", len(compiled.cache))


if __name__ == "__main__":
    main()

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

wrapped, and nothing is compiled yet: {}
(3, 'int') new: compile (slow)
[2, 4, 6]
(3, 'int') guard matches: reuse
[8, 10, 12]
(4, 'int') new: compile (slow)
[2, 4, 6, 8]
(3, 'float') new: compile (slow)
[3.0, 5.0, 7.0]
(3, 'int') guard matches: reuse
[14, 16, 18]
versions kept: 3

3 first_call.py: lazy compiling · first_call.py: คอมไพล์เมื่อถึงเวลา · first_call.py：懒惰的编译

first_call.py

# torch.compile only wraps. The first call pays for compiling. Later calls do not.
import time

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(64, 64)

    def forward(self, x):
        return torch.relu(self.layer(x) + 1.0)


def main():
    net = TinyNet()
    x = torch.randn(8, 64)

    # The very first torch.compile also loads the compiler's own Python files.
    # So we wrap once to get that out of the way, and time the second wrap.
    torch.compile(net)
    start = time.perf_counter()
    model = torch.compile(net)
    print(f"wrapping took {time.perf_counter() - start:.4f} s")
    print("model is a", type(model).__name__, "around a", type(net).__name__)

    with torch.no_grad():
        for call in [1, 2, 3]:
            start = time.perf_counter()
            model(x)
            print(f"call {call}: {time.perf_counter() - start:.4f} s")


if __name__ == "__main__":
    main()

The first time we ever ran it, with an empty compile cacheครั้งแรกสุดที่เรารันมัน โดยแคชการคอมไพล์ว่างเปล่า我们第一次运行它，编译缓存是空的

wrapping took 0.0007 s
model is a OptimizedModule around a TinyNet
call 1: 8.5745 s
call 2: 0.0002 s
call 3: 0.0000 s

The same file, run again straight afterไฟล์เดิม รันอีกครั้งทันทีหลังจากนั้น同一个文件，紧接着再运行一次

wrapping took 0.0007 s
model is a OptimizedModule around a TinyNet
call 1: 0.5145 s
call 2: 0.0001 s
call 3: 0.0001 s

4 see_graph.py: Dynamo · see_graph.py: Dynamo · see_graph.py：Dynamo

see_graph.py

# Stage 1, Dynamo: what did it capture from our Python, and on what conditions?
import contextlib
import io

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(64, 64)

    def forward(self, x):
        return torch.relu(self.layer(x) + 1.0)


def main():
    net = TinyNet()
    x = torch.randn(8, 64)

    # explain runs Dynamo only. It stops before Inductor, so it needs no compiler.
    report = torch._dynamo.explain(net)(x)
    print("graphs:      ", report.graph_count)
    print("graph breaks:", report.graph_break_count)
    print("operations:  ", report.op_count)

    print("\nThe graph, one line for each step:")
    for node in report.graphs[0].graph.nodes:
        print(f"  {node.op:14} {node.name}")

    # backend="eager" also stops after Dynamo. The guard log goes to stderr;
    # here it goes into a string, so that we choose what to print.
    log = io.StringIO()
    with contextlib.redirect_stderr(log):
        torch._logging.set_logs(guards=True)
        torch.compile(net, backend="eager")(x)
    torch._logging.set_logs()

    lines = log.getvalue().splitlines()
    print("\nThe guard on our input x:")
    for line in lines:
        if "check_tensor(L['x']" in line:
            print(" ", line.split("TENSOR_MATCH: ")[1].split("  #")[0])
    print("Guards in all:", sum("_MATCH:" in line for line in lines))


if __name__ == "__main__":
    main()

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

graphs:       1
graph breaks: 0
operations:   3

The graph, one line for each step:
  placeholder    l_self_modules_layer_parameters_weight_
  placeholder    l_self_modules_layer_parameters_bias_
  placeholder    l_x_
  call_function  linear
  call_function  add
  call_function  relu
  output         output

The guard on our input x:
  check_tensor(L['x'], Tensor, DispatchKeySet(CPU, BackendSelect, ADInplaceOrView, AutogradCPU), torch.float32, device=None, requires_grad=False, size=[8, 64], stride=[64, 1])
Guards in all: 18

5 recompile.py: when a guard fails · recompile.py: เมื่อ guard ไม่ผ่าน · recompile.py：当 guard 不成立

recompile.py

# When a guard fails, PyTorch compiles a new version.
import time

import torch


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(64, 64)

    def forward(self, x):
        return torch.relu(self.layer(x) + 1.0)


def timed_call(model, x, label):
    start = time.perf_counter()
    model(x)
    print(f"{label:28} {time.perf_counter() - start:8.4f} s")


def main():
    model = torch.compile(TinyNet())

    # These logs say out loud why PyTorch compiled again.
    torch._logging.set_logs(recompiles=True)

    with torch.no_grad():
        timed_call(model, torch.randn(8, 64), "8 rows, first time")
        timed_call(model, torch.randn(8, 64), "8 rows again")
        timed_call(model, torch.randn(16, 64), "16 rows, first time")
        timed_call(model, torch.randn(8, 64), "8 rows, back again")
        timed_call(model, torch.randn(16, 64), "16 rows again")
        timed_call(model, torch.randn(32, 64), "32 rows, first time")


if __name__ == "__main__":
    main()

What it printed for us. The log lines came out last, because they go to a different stream from the times.สิ่งที่มันพิมพ์ให้เรา บรรทัดบันทึกออกมาท้ายสุด เพราะอยู่คนละสตรีมกับเวลา它给我们打印的内容。日志行排在最后，因为它们和时间走的是不同的流。

8 rows, first time             0.5358 s
8 rows again                   0.0001 s
16 rows, first time            1.1396 s
8 rows, back again             0.0001 s
16 rows again                  0.0000 s
32 rows, first time            0.0226 s
V1101 20:21:41.637000 564 torch/_dynamo/guards.py:6374] [0/1] [__recompiles] Recompiling function forward in ~/compile_lab/recompile.py:12
V1101 20:21:41.637000 564 torch/_dynamo/guards.py:6374] [0/1] [__recompiles]     triggered by the following guard failure(s):
V1101 20:21:41.637000 564 torch/_dynamo/guards.py:6374] [0/1] [__recompiles]     - 0/0: tensor 'x' size mismatch at index 0. expected 8, actual 16

6 train_graphs.py: AOT Autograd · train_graphs.py: AOT Autograd · train_graphs.py：AOT Autograd

train_graphs.py

# Stage 2, AOT Autograd: look at the graphs it prepares for the compiler.
# Our own small backend stands where Inductor would stand, and only prints.
import torch
from functorch.compile import make_boxed_func
from torch._dynamo.backends.common import aot_autograd


class TinyNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(4, 4)

    def forward(self, x):
        return torch.relu(self.layer(x) + 1.0)


def print_steps(title, graph_module):
    print(title)
    for node in graph_module.graph.nodes:
        if node.op == "call_function":
            print("  ", node.target)
    return make_boxed_func(graph_module.forward)


def forward_compiler(graph_module, example_inputs):
    return print_steps("FORWARD graph:", graph_module)


def backward_compiler(graph_module, example_inputs):
    return print_steps("BACKWARD graph:", graph_module)


def main():
    backend = aot_autograd(fw_compiler=forward_compiler, bw_compiler=backward_compiler)
    model = torch.compile(TinyNet(), backend=backend)
    x = torch.randn(2, 4)

    print("--- training: the gradients are needed ---")
    model(x).sum().backward()

    print("\n--- inference: a new batch size, and no gradients ---")
    with torch.no_grad():
        model(torch.randn(3, 4))


if __name__ == "__main__":
    main()

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

--- training: the gradients are needed ---
FORWARD graph:
   aten.t.default
   aten.addmm.default
   aten.add.Tensor
   aten.relu.default
   aten.detach.default
BACKWARD graph:
   aten.detach.default
   aten.threshold_backward.default
   aten.t.default
   aten.mm.default
   aten.t.default
   aten.sum.dim_IntList
   aten.view.default
   aten.t.default

--- inference: a new batch size, and no gradients ---
FORWARD graph:
   aten.t.default
   aten.addmm.default
   aten.add.Tensor
   aten.relu.default

7 see_kernels.py: Inductor and fusion · see_kernels.py: Inductor และการรวมงาน · see_kernels.py：Inductor 与融合

toy_fusion.py

# What fusion means, in plain Python: one trip through the data, not two.
import time

SIZE = 1_000_000


def two_trips(numbers):
    # Add a bias, and keep the whole in-between list in memory.
    added = [n + 1.0 for n in numbers]
    # Then make a second trip for the activation.
    return [max(n, 0.0) for n in added]


def one_trip(numbers):
    # The same two jobs in one trip, and no in-between list.
    return [max(n + 1.0, 0.0) for n in numbers]


def seconds_for(job, numbers):
    start = time.perf_counter()
    answer = job(numbers)
    return time.perf_counter() - start, answer


def main():
    numbers = [float(n % 7) - 3.0 for n in range(SIZE)]
    slow, answer_a = seconds_for(two_trips, numbers)
    fast, answer_b = seconds_for(one_trip, numbers)
    print("same answer:", answer_a == answer_b)
    print(f"two trips: {slow:.3f} s")
    print(f"one trip:  {fast:.3f} s")


if __name__ == "__main__":
    main()

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

same answer: True
two trips: 0.058 s
one trip:  0.041 s

see_kernels.py

# Stage 3, Inductor: read the code it wrote, for two small models.
import contextlib
import io

import torch


class BiasThenRelu(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.bias = torch.nn.Parameter(torch.randn(64))

    def forward(self, x):
        return torch.relu(x + self.bias)


class MatmulThenRelu(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = torch.nn.Parameter(torch.randn(64, 64))

    def forward(self, x):
        return torch.relu(x @ self.weights)


def show_code_for(net):
    # The log goes to stderr. Here it goes into a string, so we choose what to print.
    log = io.StringIO()
    with contextlib.redirect_stderr(log), torch.no_grad():
        torch._logging.set_logs(output_code=True)
        torch.compile(net)(torch.randn(8, 64))
    torch._logging.set_logs()

    lines = log.getvalue().splitlines()
    print(type(net).__name__, "-", len(lines), "lines of log")
    for line in lines:
        if "= async_compile" in line:
            print("   a kernel it wrote:", line.split(" = async_compile")[0].split()[-1])
        elif "extern_kernels." in line:
            print("   a library call:   ", line.split("extern_kernels.")[1].split("(")[0])
        elif "Original ATen" in line:
            print("     made from:", line.split("Original ATen: ")[1])


def main():
    show_code_for(BiasThenRelu())
    show_code_for(MatmulThenRelu())


if __name__ == "__main__":
    main()

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

BiasThenRelu - 115 lines of log
   a kernel it wrote: cpp_fused_add_relu_0
MatmulThenRelu - 110 lines of log
   a kernel it wrote: cpp_fused_relu_0
     made from: [aten.mm]
   a library call:    mm

8 graph_break.py: graph breaks · graph_break.py: กราฟขาด · graph_break.py：图中断

graph_break.py

# A graph break: Dynamo meets code it cannot capture, and cuts the graph there.
import torch


class WithIf(torch.nn.Module):
    def forward(self, x):
        x = x + 1.0
        # Python needs a real True or False, so this asks for a value in a tensor.
        if x.sum().item() > 0:
            x = x * 2.0
        return torch.relu(x)


class WithWhere(torch.nn.Module):
    def forward(self, x):
        x = x + 1.0
        # torch.where keeps the choice inside the graph.
        x = torch.where(x.sum() > 0, x * 2.0, x)
        return torch.relu(x)


def show(name, net, x):
    report = torch._dynamo.explain(net)(x)
    print(name)
    print("  graphs:", report.graph_count, " breaks:", report.graph_break_count)
    for reason in report.break_reasons:
        print("  why:", reason.reason.splitlines()[0])


def main():
    x = torch.ones(4, 4)
    show("WithIf", WithIf(), x)
    show("WithWhere", WithWhere(), x)

    print("\nfullgraph=True on each of them:")
    for net in [WithWhere(), WithIf()]:
        name = type(net).__name__
        try:
            torch.compile(net, fullgraph=True)(x)
            print(" ", name, "compiled as one graph")
        except Exception as error:
            print(" ", name, "stopped with", type(error).__name__)


if __name__ == "__main__":
    main()

What it printed for us. The warning is PyTorch's own, and it names the line.สิ่งที่มันพิมพ์ให้เรา คำเตือนเป็นของ PyTorch เอง และบอกบรรทัดด้วย它给我们打印的内容。警告是 PyTorch 自己发的，还指出了是哪一行。

WithIf
  graphs: 2  breaks: 1
  why: Unsupported Tensor.item() call with capture_scalar_outputs=False
WithWhere
  graphs: 1  breaks: 0

fullgraph=True on each of them:
  WithWhere compiled as one graph
  WithIf stopped with UserError
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] Graph break from `Tensor.item()`, consider setting:
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0]     torch._dynamo.config.capture_scalar_outputs = True
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] or:
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0]     env TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] to include these operations in the captured graph.
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] 
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] Graph break: from user code at:
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0]   File "~/compile_lab/graph_break.py", line 9, in forward
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0]     if x.sum().item() > 0:
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0]        ~~~~~~~~~~~~^^
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] 
W1101 20:21:53.837000 673 torch/_dynamo/variables/tensor.py:2043] [0/0] 

9 benchmark.py: is it faster? · benchmark.py: เร็วขึ้นจริงไหม · benchmark.py：真的更快吗

benchmark.py

# Is it faster? Keep compile time apart from the steady speed, and ask two models.
import statistics
import time

import torch

WARM_UP_CALLS = 3
TIMED_CALLS = 50


class ManySmallSteps(torch.nn.Module):
    # Lots of small steps, one after another, on every number.
    def forward(self, x):
        for _ in range(20):
            x = torch.relu(x + 0.1) * 0.99
        return x


class OneBigMatmul(torch.nn.Module):
    # One large matrix multiplication, already done by a tuned library.
    def forward(self, x):
        return x @ x


def finish_gpu_work():
    # A GPU works in the background, so wait for it before reading the clock.
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def seconds_for_one_call(model, x):
    start = time.perf_counter()
    model(x)
    finish_gpu_work()
    return time.perf_counter() - start


def compare(net, x):
    model = torch.compile(net)
    first = seconds_for_one_call(model, x)

    for _ in range(WARM_UP_CALLS):
        seconds_for_one_call(net, x)
        seconds_for_one_call(model, x)

    # They take turns, so that neither one gets a quieter moment of the machine.
    plain_times = []
    compiled_times = []
    for _ in range(TIMED_CALLS):
        plain_times.append(seconds_for_one_call(net, x))
        compiled_times.append(seconds_for_one_call(model, x))

    plain_ms = statistics.median(plain_times) * 1000
    compiled_ms = statistics.median(compiled_times) * 1000
    print(type(net).__name__)
    print(f"  first compiled call {first:8.2f} s")
    print(f"  plain, each call    {plain_ms:8.2f} ms")
    print(f"  compiled, each call {compiled_ms:8.2f} ms")
    print(f"  plain / compiled    {plain_ms / compiled_ms:8.2f} times")


def main():
    print(torch.get_num_threads(), "CPU threads")
    with torch.no_grad():
        compare(ManySmallSteps(), torch.randn(256, 256))
        compare(OneBigMatmul(), torch.randn(1500, 1500))


if __name__ == "__main__":
    main()

What it printed for usสิ่งที่มันพิมพ์ให้เรา它给我们打印的内容

11 CPU threads
ManySmallSteps
  first compiled call     1.60 s
  plain, each call        0.46 ms
  compiled, each call     0.06 ms
  plain / compiled        7.15 times
OneBigMatmul
  first compiled call     0.11 s
  plain, each call       17.25 ms
  compiled, each call    17.89 ms
  plain / compiled        0.96 times