Inductor C++ Wrapper 教程

创建日期：2023 年 10 月 2 日 | 最后更新：2024 年 1 月 16 日 | 最后验证：2024 年 11 月 5 日

作者：Chunyuan Wu, Bin Bao, Jiong Gong

先决条件：

• PyTorch 中的 torch.compile 和 TorchInductor 概念

引言

Python 作为 PyTorch 的主要接口，易于使用且对开发和调试高效。Inductor 的默认 Wrapper 生成 Python 代码来调用生成的内核和外部内核。然而，在需要高性能的部署场景中，Python 作为一种解释型语言，相比编译型语言运行速度较慢。

我们通过利用 PyTorch C++ API 实现了一个 Inductor C++ Wrapper，以生成纯 C++ 代码，结合生成的内核和外部内核。这允许在纯 C++ 中执行每个捕获的 Dynamo 图，从而减少图内的 Python 开销。

启用 API

此功能仍处于原型阶段。要激活此功能，请在代码中添加以下内容：

import torch._inductor.config as config
config.cpp_wrapper = True

这将通过减少 Inductor Wrapper 的 Python 开销来加速你的模型。

示例代码

我们将使用下面的前端代码作为示例

import torch

def fn(x):
    return torch.tensor(list(range(2, 40, 2)), device=x.device) + x

x = torch.randn(1)
opt_fn = torch.compile()(fn)
y = opt_fn(x)

对于 CPU

使用默认 Python Wrapper 的 Inductor 生成代码的主要部分将如下所示

def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (1, ), (1, ))
    buf0 = empty_strided((19, ), (1, ), device='cpu', dtype=torch.float32)
    cpp_fused_add_lift_fresh_0(c_void_p(constant0.data_ptr()), c_void_p(arg0_1.data_ptr()), c_void_p(buf0.data_ptr()))
    del arg0_1
    return (buf0, )

启用 C++ Wrapper 后，call 函数生成的代码将变成 C++ 扩展 module 的一个 C++ 函数 inductor_entry_cpp

std::vector<at::Tensor> inductor_entry_cpp(const std::vector<at::Tensor>& args) {
    at::Tensor arg0_1 = args[0];
    at::Tensor constant0 = args[1];
    auto buf0 = at::empty_strided({19L, }, {1L, }, at::device(at::kCPU).dtype(at::kFloat));
    cpp_fused_add_lift_fresh_0((long*)(constant0.data_ptr()), (float*)(arg0_1.data_ptr()), (float*)(buf0.data_ptr()));
    arg0_1.reset();
    return {buf0};
}

module = CppWrapperCodeCache.load(cpp_wrapper_src, 'inductor_entry_cpp', 'c2buojsvlqbywxe3itb43hldieh4jqulk72iswa2awalwev7hjn2', False)

def _wrap_func(f):
    def g(args):
        args_tensor = [arg if isinstance(arg, torch.Tensor) else torch.tensor(arg) for arg in args]
        constants_tensor = [constant0]
        args_tensor.extend(constants_tensor)

        return f(args_tensor)
    return g
call = _wrap_func(module.inductor_entry_cpp)

对于 GPU

基于相同的示例代码，针对 GPU 生成的代码将如下所示

def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (1, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0) # no-op to ensure context
        buf0 = empty_strided((19, ), (1, ), device='cuda', dtype=torch.float32)
        # Source Nodes: [add, tensor], Original ATen: [aten.add, aten.lift_fresh]
        stream0 = get_cuda_stream(0)
        triton_poi_fused_add_lift_fresh_0.run(constant0, arg0_1, buf0, 19, grid=grid(19), stream=stream0)
        run_intermediate_hooks('add', buf0)
        del arg0_1
        return (buf0, )

启用 C++ Wrapper 后，将生成以下等效的 C++ 代码

std::vector<at::Tensor> inductor_entry_cpp(const std::vector<at::Tensor>& args) {
    at::Tensor arg0_1 = args[0];
    at::Tensor constant0 = args[1];

    at::cuda::CUDAGuard device_guard(0);
    auto buf0 = at::empty_strided({19L, }, {1L, }, at::TensorOptions(c10::Device(at::kCUDA, 0)).dtype(at::kFloat));
    // Source Nodes: [add, tensor], Original ATen: [aten.add, aten.lift_fresh]
    if (triton_poi_fused_add_lift_fresh_0 == nullptr) {
        triton_poi_fused_add_lift_fresh_0 = loadKernel("/tmp/torchinductor_user/mm/cmm6xjgijjffxjku4akv55eyzibirvw6bti6uqmfnruujm5cvvmw.cubin", "triton_poi_fused_add_lift_fresh_0_0d1d2d3");
    }
    CUdeviceptr var_0 = reinterpret_cast<CUdeviceptr>(constant0.data_ptr());
    CUdeviceptr var_1 = reinterpret_cast<CUdeviceptr>(arg0_1.data_ptr());
    CUdeviceptr var_2 = reinterpret_cast<CUdeviceptr>(buf0.data_ptr());
    auto var_3 = 19;
    void* kernel_args_var_0[] = {&var_0, &var_1, &var_2, &var_3};
    cudaStream_t stream0 = at::cuda::getCurrentCUDAStream(0);
    launchKernel(triton_poi_fused_add_lift_fresh_0, 1, 1, 1, 1, 0, kernel_args_var_0, stream0);
    arg0_1.reset();
    return {buf0};
}

module = CppWrapperCodeCache.load(cpp_wrapper_src, 'inductor_entry_cpp', 'czbpeilh4qqmbyejdgsbpdfuk2ss5jigl2qjb7xs4gearrjvuwem', True)

def _wrap_func(f):
    def g(args):
        args_tensor = [arg if isinstance(arg, torch.Tensor) else torch.tensor(arg) for arg in args]
        constants_tensor = [constant0]
        args_tensor.extend(constants_tensor)

        return f(args_tensor)
    return g
call = _wrap_func(module.inductor_entry_cpp)

结论

在本教程中，我们介绍了 TorchInductor 中的新 C++ Wrapper，只需更改两行代码即可加速你的模型。我们解释了这项新功能的动机，并介绍了易于使用的 API 来激活此实验性功能。此外，我们展示了使用默认 Python Wrapper 和新 C++ Wrapper 在 CPU 和 GPU 上生成的 Inductor 代码，以直观地展示这两种 Wrapper 之间的差异。

此功能仍处于原型阶段。如果你有任何功能请求或遇到任何问题，请在GitHub issues上提交错误报告。