评价此页

Inductor CPU 后端调试与性能分析#

创建日期:2023年7月1日 | 最后更新:2025年1月8日 | 最后验证:2024年11月5日

作者Xuan Liao, Haozhe Zhu, Jiong Gong, Weihan Wang

概述#

PyTorch 2.0 引入了名为 torch.compile 的编译 API。该新功能通过默认的 Inductor 后端进行图级优化,相比动态图(eager mode)执行提供了显著的加速。

本教程旨在深入探讨 torch.compile 的内部细节,并对 Inductor CPU 后端的调试与性能分析进行深入介绍。

同时,你还可以参考其他关于 torch.compile 的相关教程,包括基础用法、全面的故障排查以及 GPU 专项知识,如 GPU 性能分析

我们将从一个触发编译问题和精度问题的动机示例开始,演示调试过程并定位问题。

通过启用日志记录并探索底层生成的代码,你可以学习如何逐步缩小失败范围,并最终查明根本原因。

随后,我们将讨论如何分析编译后的代码,并通过与动态图模式的性能对比,阐述为何 torch.compile 相比于动态图能提供额外的性能提升。

调试#

以下是一个使用 Inductor 运行 torch.compile 并将其结果与动态图模式进行比较的简单示例。

import torch

def foo1(x1, x2):
    a = torch.neg(x1)
    b = torch.maximum(x2, a)
    y = torch.cat([b], dim=0)
    return y

x1 = torch.randint(256, (1, 8), dtype=torch.uint8)
x2 = torch.randint(256, (8390, 8), dtype=torch.uint8)

compiled_foo1 = torch.compile(foo1)
result = compiled_foo1(x1, x2)

cpp 代码生成器中,neg 的正确实现如下:

def neg1(x):
    return f"decltype({x})(-{x})"

为了演示调试过程,我们稍后会将其函数修改为一个错误的实现。

获取更多日志信息#

默认运行此简单示例时不会提供调试信息。为了获得更有用的调试和日志信息,我们通常添加 TORCH_COMPILE_DEBUG 环境变量,如下所示:

TORCH_COMPILE_DEBUG=1 python xx.py

这将在输出日志中打印更多调试信息,并转储代码生成过程中生成的中间 IR。你可以在日志中找到如下所示的转储文件路径:

torch._inductor.debug: [WARNING] model___20 debug trace: /tmp/torchinductor_root/rx/crxfi2ybd7yp5sbj2pnhw33wfhtdw7wumvrobyp5sjvdui5ktjc2.debug

在该目录中,为调试目的保存了以下文件:

File

描述

fx_graph_runnable.py

可执行的 FX 图,位于分解后、模式匹配前

fx_graph_transformed.py

转换后的 FX 图,位于模式匹配后

ir_pre_fusion.txt

融合前的 Inductor IR

ir_post_fusion.txt

融合后的 Inductor IR

output_code.py

为该图生成的 Python 代码,包含 C++/Triton 内核

请注意,fx_graph_runnable.pyoutput_code.py 均可运行且可编辑,以便于调试。以下是从文件中提取的代码主要部分,我们将 C++ 生成的行与 FX 代码行进行了关联。

fx_graph_runnable:

def forward1(self, arg0_1, arg1_1):
    neg = torch.ops.aten.neg.default(arg0_1);  arg0_1 = None
    maximum = torch.ops.aten.maximum.default(arg1_1, neg);  arg1_1 = neg = None
    clone = torch.ops.aten.clone.default(maximum);  maximum = None
    return (clone,)

output_code 中的 C++ 内核

import torch
from torch._inductor.async_compile import AsyncCompile
async_compile = AsyncCompile()

cpp_fused_cat_maximum_neg_0 = async_compile.cpp('''
#include "/tmp/torchinductor_root/gv/cgv6n5aotqjo5w4vknjibhengeycuattfto532hkxpozszcgxr3x.h"
extern "C" void kernel(const unsigned char* in_ptr0,
                       const unsigned char* in_ptr1,
                       unsigned char* out_ptr0)
{
    {
        #pragma GCC ivdep
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(8390L); i0+=static_cast<long>(1L))
        {
            #pragma GCC ivdep
            for(long i1=static_cast<long>(0L); i1<static_cast<long>(8L); i1+=static_cast<long>(1L))
            {
                auto tmp0 = in_ptr0[static_cast<long>(i1 + (8L*i0))];
                auto tmp1 = in_ptr1[static_cast<long>(i1)];
                // Corresponding FX code line: neg = torch.ops.aten.neg.default(arg0_1);  arg0_1 = None
                auto tmp2 = decltype(tmp1)(-tmp1);
                // Corresponding FX code line: maximum = torch.ops.aten.maximum.default(arg1_1, neg);  arg1_1 = neg = None
                auto tmp3 = max_propagate_nan(tmp0, tmp2);
                // Corresponding FX code line: clone = torch.ops.aten.clone.default(maximum);  maximum = None
                out_ptr0[static_cast<long>(i1 + (8L*i0))] = tmp3;
            }
        }
    }
}''')

确定错误的组成部分#

当遇到错误或精度问题时,定位错误的一个直接方法是缩小问题范围。首先要做的是确定错误发生的组件。幸运的是,这可以通过简单地更改 torch.compile 的后端来实现。

代码

描述

torch.compile(fn, backend="eager")

启用 Dynamo

torch.compile(fn, backend="aot_eager")

启用 Dynamo + AOT Autograd

torch.compile(fn, backend="inductor")

启用 Dynamo + AOT Autograd + Inductor

如果模型在后端设置为 eageraot_eager 时能成功运行,但在 inductor 下失败,我们就可以将失败范围缩小到 Inductor。

编译错误#

众所周知,图级优化的演进链条大致如下:

torch.neg (Python) -> torch.ops.aten.neg.default (within FX graph) -> ops.neg (within IR node) -> tmp2 = -tmp1 (within C++ kernel)

如果你遇到编译错误,说明在输出代码中编译 C++ 内核时出现了问题。这类错误表明在将 IR 节点降级(lowering)到输出代码时引入了 Bug。编译错误的根本原因通常显示在回溯日志中。

例如,neg 函数被修改如下:

def neg2(x):
    return f"-{x}"

日志给出了以下编译错误,原因非常明确。

 torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
 CppCompileError: C++ compile error
 /tmp/torchinductor_root/xg/cxga5tk3b4lkwoxyigrtocjp5s7vc5cg2ikuscf6bk6pjqip2bhx.cpp: In function ‘void kernel(const unsigned char*, const unsigned char*, unsigned char*)’:
 /tmp/torchinductor_root/xg/cxga5tk3b4lkwoxyigrtocjp5s7vc5cg2ikuscf6bk6pjqip2bhx.cpp:17:57: error: no matching function for call to ‘max_propagate_nan(unsigned char&, int&)’
   17 |                 auto tmp3 = max_propagate_nan(tmp0, tmp2);
        |                                                         ^
 In file included from /tmp/torchinductor_root/xg/cxga5tk3b4lkwoxyigrtocjp5s7vc5cg2ikuscf6bk6pjqip2bhx.cpp:2:
 /tmp/torchinductor_root/gv/cgv6n5aotqjo5w4vknjibhengeycuattfto532hkxpozszcgxr3x.h:27:17: note: candidate: ‘template<class scalar_t> scalar_t max_propagate_nan(scalar_t, scalar_t)’
 27 | inline scalar_t max_propagate_nan(scalar_t a, scalar_t b) {
      |                 ^~~~~~~~~~~~~~~~~
 /tmp/torchinductor_root/gv/cgv6n5aotqjo5w4vknjibhengeycuattfto532hkxpozszcgxr3x.h:27:17: note:   template argument deduction/substitution failed:
/tmp/torchinductor_root/xg/cxga5tk3b4lkwoxyigrtocjp5s7vc5cg2ikuscf6bk6pjqip2bhx.cpp:17:57: note:   deduced conflicting types for parameter ‘scalar_t’ (‘unsigned char’ and ‘int’)
 17 |                 auto tmp3 = max_propagate_nan(tmp0, tmp2);
      |                                                         ^

让我们看看输出代码中相应的 C++ 内核和 IR 节点。

C++ 内核

include "/tmp/torchinductor_root/gv/cgv6n5aotqjo5w4vknjibhengeycuattfto532hkxpozszcgxr3x.h"
extern "C" void kernel(const unsigned char* in_ptr0,
                    const unsigned char* in_ptr1,
                    unsigned char* out_ptr0)
{
    {
        #pragma GCC ivdep
        for(long i0=static_cast<long>(0L); i0<static_cast<long>(8390L); i0+=static_cast<long>(1L))
        {
            #pragma GCC ivdep
            for(long i1=static_cast<long>(0L); i1<static_cast<long>(8L); i1+=static_cast<long>(1L))
            {
                auto tmp0 = in_ptr0[static_cast<long>(i1 + (8L*i0))];
                auto tmp1 = in_ptr1[static_cast<long>(i1)];
                auto tmp2 = -tmp1;
                auto tmp3 = max_propagate_nan(tmp0, tmp2);
                out_ptr0[static_cast<long>(i1 + (8L*i0))] = tmp3;
            }
        }
    }
}

IR 节点

buf0: SchedulerNode(ComputedBuffer)
buf0.writes = [MemoryDep('buf0', c0, {c0: 67120})]
buf0.unmet_dependencies = []
buf0.met_dependencies =
    [   MemoryDep('arg0_1', c1, {c0: 8390, c1: 8}),
        MemoryDep('arg1_1', c0, {c0: 67120})]
buf0.users = [NodeUser(node=OUTPUT, can_inplace=False)]
buf0.group.device = cpu
buf0.group.iteration = ((8390, 8), ())
buf0.sizes = ([8390, 8], [])
class buf0_loop_body:
    var_ranges = {z0: 8390, z1: 8}
    index0 = 8*z0 + z1
    index1 = z1
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('arg1_1', get_index)
        get_index_1 = self.get_index('index1')
        load_1 = ops.load('arg0_1', get_index_1)
        neg = ops.neg(load_1)
        maximum = ops.maximum(load, neg)
        get_index_2 = self.get_index('index0')
        store = ops.store('buf0', get_index_2, maximum, None)
        return store

根据回溯日志,编译错误是由 max_propagate_nan 输入的数据类型不一致引起的。通过检查 C++ 内核,我们知道由于 tmp0long 类型,在执行 -tmp2 不再是 long。我们可以轻松地将 C++ 内核中的 -max_propagate_nan 与 IR 节点中的 ops.negops.maximum 对应起来。

现在我们成功发现根本原因是 cpp 代码生成器中 ops.neg 的实现,它在执行 neg 时悄悄改变了数据类型。

精度调试#

另一方面,如果模型运行出现其他错误或精度问题,你可以使用名为 Minifier 的 PyTorch 调试工具。

Minifier 的核心思想是不断移除图中的节点和输入,直到找到出现问题的最小图。它通过四种策略帮助自动生成最小化的有问题的图:截断后缀、Delta 调试、消除死代码和移除未使用输入。

我们现在将借助 Minifier 演示精度问题的调试过程。精度问题指的是后端 eager 和 inductor 的输出不一致的情况。

例如,我们将示例修改如下:

from torch._dynamo.utils import same

def foo2(x1, x2):
    a = torch.neg(x1)
    b = torch.maximum(x2, a)
    y = torch.cat([b], dim=0)
    return y

x1 = torch.randn((1, 8), dtype=torch.float32)
x2 = torch.randn((8390, 8), dtype=torch.float32)

expected_result = foo2(x1, x2)

compiled_foo2 = torch.compile(foo2)
actual_result = compiled_foo2(x1, x2)

assert same(expected_result, actual_result) == True

并同时修改 neg 函数。

def neg3(x):
    return f"decltype({x})(2 * {x})"

会出现如下精度问题:

torch._dynamo.utils: [ERROR] Accuracy failed: allclose not within tol=0.0001
Traceback (most recent call last):
  File "test_script.py", line 18, in <module>
    assert same(expected_result, actual_result) == True
AssertionError

要使用 Minifier 调试精度问题,需要两个环境变量:

TORCHDYNAMO_REPRO_AFTER="aot" TORCHDYNAMO_REPRO_LEVEL=4 python xx.py

这为我们提供了展示最小化过程的日志信息。

Started off with 6 nodes

Trying granularity 2
Strategy: Truncate suffix (G: 2) (6 nodes, 2 inputs)
SUCCESS: Went from 6 to 4 nodes

Trying granularity 4
Strategy: Remove unused inputs (G: 4) (4 nodes, 2 inputs)
SUCCESS: Went from 4 to 3 nodes

运行后,我们得到包含目标节点 neg 的最终最小化图。

def forward2(self, arg0_1):
    neg = torch.ops.aten.neg.default(arg0_1);  arg0_1 = None
    return (neg,)

有关 Minifier 的更多用法详情,请参考故障排查

性能分析#

在本节中,我们将演示如何对使用 Inductor CPU 后端编译的模型进行性能分析。在下例中,我们对 Hugging Face Transformer 模型 MobileBertForQuestionAnswering 分别在 eager 模式和 Inductor 图模式下进行基准测试。基准测试后会打印执行时间和 Inductor 的加速比。我们使用 Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz 并在第一个插槽上运行基准测试以演示本节中的优化。为了在 Intel(R) CPU 上获得最佳基准测试效果,我们设置了以下环境变量:

export KMP_BLOCKTIME=1
export KMP_SETTINGS=1
export KMP_AFFINITY=granularity=fine,compact,1,0
export LD_PRELOAD=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}/lib/libiomp5.so:${CONDA_PREFIX:-"$(dirname $(which conda))/../"}/lib/libjemalloc.so
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1"
numactl -C 0-31 -m 0 python bench.py
# bench.py
from transformers import MobileBertForQuestionAnswering
# Initialize an eager model
model = MobileBertForQuestionAnswering.from_pretrained("csarron/mobilebert-uncased-squad-v2")
seq_length = 128
bs = 128
vocab_size = model.config.vocab_size
input = torch.randint(0, vocab_size, (bs, seq_length), dtype=torch.int64)
input_dict = {"input_ids": input}

# Initialize the inductor model
compiled_model = torch.compile(model)
with torch.no_grad():
    compiled_model(**input_dict)

NUM_ITERS=50
import timeit
with torch.no_grad():
    # warmup
    for _ in range(10):
        model(**input_dict)
    eager_t = timeit.timeit("model(**input_dict)", number=NUM_ITERS, globals=globals())

with torch.no_grad():
    # warmup
    for _ in range(10):
        compiled_model(**input_dict)
    inductor_t = timeit.timeit("compiled_model(**input_dict)", number=NUM_ITERS, globals=globals())
# print(f"eager use: {eager_t * 1000 / NUM_ITERS} ms/iter")
# print(f"inductor use: {inductor_t * 1000 / NUM_ITERS} ms/iter")
# print(f"speed up ratio: {eager_t / inductor_t}")
Loading weights:   0%|          | 0/1113 [00:00<?, ?it/s]
Loading weights:  32%|███▏      | 354/1113 [00:00<00:00, 3495.27it/s]
Loading weights:  63%|██████▎   | 704/1113 [00:00<00:00, 3351.83it/s]
Loading weights:  93%|█████████▎| 1040/1113 [00:00<00:00, 3337.67it/s]
Loading weights: 100%|██████████| 1113/1113 [00:00<00:00, 3374.99it/s]

输出

eager use: 802.1023553796113 ms/iter
inductor use: 339.95180135127157 ms/iter
speed up ratio: 2.359459053287382

在我们的测试中,发现 Inductor CPU 后端将模型速度提升了约 2.355 倍。

接下来,让我们深入操作层面了解加速的来源。PyTorch Profiler 是一个很好的工具。Inductor CPU 后端支持通过 enable_kernel_profile 配置选项将融合内核的时间记录到 profiler 中。

from torch._inductor import config
config.cpp.enable_kernel_profile = True

按照 PyTorch Profiler 的步骤,我们可以得到分析表和 trace 文件。

# bench.py
from torch.profiler import profile, schedule, ProfilerActivity
RESULT_DIR = "./prof_trace"
my_schedule = schedule(
    skip_first=10,
    wait=5,
    warmup=5,
    active=1,
    repeat=5)

def trace_handler(p):
    output = p.key_averages().table(sort_by="self_cpu_time_total", row_limit=20)
    # print(output)
    p.export_chrome_trace(f"{RESULT_DIR}/{p.step_num}.json")

for _ in range(10):
    model(**input_dict)  # compiled_model(**input_dict) to get inductor model profiling

total = 0
with profile(
    activities=[ProfilerActivity.CPU],
    schedule=my_schedule,
    on_trace_ready=trace_handler
) as p:
    for _ in range(50):
        model(**input_dict)  # compiled_model(**input_dict) to get inductor model profiling
        p.step()
/usr/local/lib/python3.10/dist-packages/torch/profiler/profiler.py:272: UserWarning: Warning: Profiler clears events at the end of each cycle.Only events from the current cycle will be reported.To keep events across cycles, set acc_events=True.
  _warn_once(

我们得到了 eager 模式模型的性能分析表(省略部分列):

-------------------------  ------------  ------------  ------------
                     Name   CPU total %     CPU total    # of Calls
-------------------------  ------------  ------------  ------------
              aten::addmm        45.73%     370.814ms           362
                aten::add        19.89%     161.276ms           363
              aten::copy_        14.97%     121.416ms           488
                aten::mul         9.02%      73.154ms           194
          aten::clamp_min         8.81%      71.444ms            96
                aten::bmm         5.46%      44.258ms            48
            ProfilerStep*       100.00%     810.920ms             1
                aten::div         2.89%      23.447ms            24
           aten::_softmax         1.00%       8.087ms            24
             aten::linear        46.48%     376.888ms           362
              aten::clone         2.77%      22.430ms            98
                  aten::t         0.31%       2.502ms           362
               aten::view         0.14%       1.161ms           850
          aten::transpose         0.17%       1.377ms           386
       aten::index_select         0.12%     952.000us             3
             aten::expand         0.12%     986.000us           458
             aten::matmul         8.31%      67.420ms            48
                aten::cat         0.09%     703.000us             1
         aten::as_strided         0.08%     656.000us           963
               aten::relu         8.86%      71.864ms            96
-------------------------  ------------  ------------  ------------
Self CPU time total: 810.920ms

类似地,我们还获得了使用 Inductor 编译后的模型分析表(省略部分列):

-----------------------------------------------  ------------  ------------  ------------
                                           Name   CPU total %     CPU total    # of Calls
-----------------------------------------------  ------------  ------------  ------------
                               mkl::_mkl_linear        68.79%     231.573ms           362
                                      aten::bmm         8.02%      26.992ms            48
                                  ProfilerStep*       100.00%     336.642ms             1
  graph_0_cpp_fused_constant_pad_nd_embedding_0         0.27%     915.000us             1
                                    aten::empty         0.27%     911.000us           362
 graph_0_cpp_fused__mkl_linear_add_mul_relu_151         0.27%     901.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_226         0.27%     899.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_361         0.27%     898.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_121         0.27%     895.000us             1
  graph_0_cpp_fused__mkl_linear_add_mul_relu_31         0.27%     893.000us             1
  graph_0_cpp_fused__mkl_linear_add_mul_relu_76         0.26%     892.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_256         0.26%     892.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_346         0.26%     892.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_241         0.26%     891.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_316         0.26%     891.000us             1
  graph_0_cpp_fused__mkl_linear_add_mul_relu_91         0.26%     890.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_106         0.26%     890.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_211         0.26%     890.000us             1
  graph_0_cpp_fused__mkl_linear_add_mul_relu_61         0.26%     889.000us             1
 graph_0_cpp_fused__mkl_linear_add_mul_relu_286         0.26%     889.000us             1
-----------------------------------------------  ------------  ------------  ------------
Self CPU time total: 336.642ms

从 eager 模型的分析表中,我们可以看到最耗时的算子是 [aten::addmm, aten::add, aten::copy_, aten::mul, aten::clamp_min, aten::bmm]。与 Inductor 模型分析表相比,我们注意到多了一个 mkl::_mkl_linear 条目以及多个形如 graph_0_cpp_fused_* 的融合内核。这些正是 Inductor 模型所做的主要优化。让我们分别讨论它们。

(1) 关于 mkl::_mkl_linear:你可能注意到该内核的调用次数为 362,与 eager 模型分析表中的 aten::linear 完全相同。该算子的 CPU 总时间在 eager 中为 376.888ms,而在 mkl::_mkl_linear 中为 231.573ms。这表明“线性”部分有约 1.63 倍的加速。加速主要来自于将权重张量打包为块内存格式,并在 Inductor CPU 后端调用 cblas_sgemm_compute 以在 GEMM 计算期间获得更好的缓存行为。

(2) 关于其他内存密集型算子:在我们的测试中,eager/inductor 模型的端到端延迟分别为 802/339ms。因此我们可以大致推断,其他内存密集型算子的加速比约为 3.94 倍。让我们阅读生成的代码,以理解 Inductor 是如何实现这种令人印象深刻的优化的。你可以在 output_code.py 中搜索 cpp_fused__mkl_linear_add_mul_relu_151 来找到生成的代码。

cpp_fused__mkl_linear_add_mul_relu_151 = async_compile.cpp('''
#include <ATen/record_function.h>
#include "/tmp/torchinductor_root/lr/clrlgu27q4ggd472umdzwsu6qcpqxcuusjxqvx2hwitjbujiiz7z.h"
extern "C" void kernel(float* in_out_ptr0,
                       const float* in_ptr0,
                       const float* in_ptr1,
                       const float* in_ptr2,
                       const float* in_ptr3)
{
    RECORD_FUNCTION("graph_0_cpp_fused__mkl_linear_add_mul_relu_151", c10::ArrayRef<c10::IValue>({}));
    #pragma omp parallel num_threads(32)
    {
        {
            #pragma omp for
            for(long i0=static_cast<long>(0L); i0<static_cast<long>(16384L); i0+=static_cast<long>(1L))
            {
                for(long i1=static_cast<long>(0L); i1<static_cast<long>(512L); i1+=static_cast<long>(8L))
                {
                    auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(i1 + (512L*i0)));
                    auto tmp1 = at::vec::Vectorized<float>::loadu(in_ptr1 + static_cast<long>(i1));
                    auto tmp3 = at::vec::Vectorized<float>::loadu(in_out_ptr0 + static_cast<long>(i1 + (512L*i0)));
                    auto tmp5 = at::vec::Vectorized<float>::loadu(in_ptr2 + static_cast<long>(i1));
                    auto tmp7 = at::vec::Vectorized<float>::loadu(in_ptr3 + static_cast<long>(i1));
                    auto tmp2 = tmp0 + tmp1;
                    auto tmp4 = tmp2 + tmp3;
                    auto tmp6 = tmp4 * tmp5;
                    auto tmp8 = tmp6 + tmp7;
                    tmp8.store(in_out_ptr0 + static_cast<long>(i1 + (512L*i0)));
                }
            }
        }
    }
}''')

从上面的生成代码中,我们可以看到该内核对 [add, add, mul, add] 执行了典型的循环融合(Loop Fusion)。这是一个导致性能瓶颈的内存绑定操作。为了对这种优化有更直观的感受,我们可以推断输入的尺寸和步长,并进一步对这个 [add, add, mul, add] 模式进行基准测试。

# bench.py
def func(arg_0, arg_1, arg_2, arg_3, arg_4):
    add_0 = arg_0 + arg_1
    add_1 = add_0 + arg_2
    mul_1 = add_1 * arg_3
    add_2 = mul_1 + arg_4
    arg_2 = add_2
    return arg_2

arg_0 = torch.rand(16384, 512)
arg_1 = torch.rand(1, 512)
arg_2 = torch.zeros(16384, 512)
arg_3 = torch.rand(1, 512)
arg_4 = torch.rand(1, 512)

input = (arg_0, arg_1, arg_2, arg_3, arg_4)
inductor_func = torch.compile(func)
with torch.no_grad():
    inductor_func(*input)

import timeit
NUM_ITERS=100
with torch.no_grad():
    # warmup
    for _ in range(10):
        func(*input)
    eager_t = timeit.timeit("func(*input)", number=NUM_ITERS, globals=globals())

with torch.no_grad():
    # warmup
    for _ in range(10):
        inductor_func(*input)
    inductor_t = timeit.timeit("inductor_func(*input)", number=NUM_ITERS, globals=globals())
# print(f"eager use: {eager_t * 1000 / NUM_ITERS} ms/iter")
# print(f"inductor use: {inductor_t * 1000 / NUM_ITERS} ms/iter")
# print(f"speed up ratio: {eager_t / inductor_t}")

输出

eager use: 5.780875144992024 ms/iter
inductor use: 0.9588955780491233 ms/iter
speed up ratio: 6.0286805751604735

这只是一个例子。分析表显示该模型中所有逐元素算子都在 Inductor 中自动融合。你可以在 output_code.py 中阅读更多的内核实现。

结论#

本文档提供了关于 Inductor CPU 后端的深入教程。

通过动机示例,我们逐步介绍了调试和性能分析的过程。核心思想是缩小问题范围。

我们借助调试日志和 Minifier 工具,一步步演示了如何深入研究问题并找到失败的根源。首先确定失败发生的组件,然后尝试生成可以重现失败的最小代码片段。

当 Inductor 的性能优于 eager 模式时,我们提供了一种稳健的性能分析分析方法。我们展示了如何使用 PyTorch Profiler 找到耗时的热点,并从算子级或内核级分析原因来解释这一现象。

脚本总运行时间:(10 分 34.821 秒)