unstable/inductor_cpp_wrapper_tutorial

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

Inductor C++ Wrapper 教程

作者: Chunyuan Wu, Bin Bao, Jiong Gong

先决条件:

• PyTorch 中的 torch.compile 和 TorchInductor 概念

简介

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

我们通过利用 PyTorch C++ API 实现了一个 Inductor C++ 包装器，该包装器生成纯 C++ 代码，将生成的内核和外部内核结合起来。这使得每个捕获的 Dynamo 图都可以用纯 C++ 执行，从而减少了图内的 Python 开销。

启用 API

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

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

这将通过减少 Inductor 包装器的 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 包装器的 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++ 包装器，为 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++ 包装器，将生成以下等效的 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++ 包装器，只需更改两行代码即可加速您的模型。我们解释了此新功能的动机，并介绍了易于使用的 API 来激活此实验性功能。此外，我们还演示了在 CPU 和 GPU 上使用默认 Python 包装器和新的 C++ 包装器的 Inductor 生成代码，直观地展示了这两种包装器之间的差异。

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