advanced/cpp_custom_ops_sycl

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

自定义 SYCL 算子

您将学到什么

• 如何将使用 SYCL 编写的自定义算子集成到 PyTorch 中

先决条件

• 适用于 Linux 的 PyTorch 2.8 或更高版本

• 适用于 Windows 的 PyTorch 2.10 或更高版本

• 对 SYCL 编程的基本了解

注意

SYCL 是英特尔 GPU（设备标签 xpu）的后端编程语言。有关配置详情，请参阅：英特尔 GPU 入门。随 英特尔深度学习基础套件 一同提供的英特尔编译器负责处理 SYCL 编译。在执行本教程中的代码示例之前，请确保已安装并激活了编译器环境。

PyTorch 提供了一个庞大的算子库，可用于处理张量（例如 torch.add、torch.sum 等）。然而，您可能希望将新的自定义算子引入 PyTorch。本教程演示了编写 SYCL 自定义算子的最佳路径。有关 C++ 和 CUDA 算子的教程，请参阅 自定义 C++ 和 CUDA 算子。

按照以下结构创建自定义 SYCL 算子

sycl_example/
├── setup.py
├── sycl_extension
│   ├── __init__.py
│   ├── muladd.sycl
│   └── ops.py
└── test_sycl_extension.py

设置构建系统

如果您需要编译 SYCL 代码（注意扩展名应为 .sycl），请使用 torch.utils.cpp_extension.SyclExtension。设置过程与 C++/CUDA 非常相似，只是需要针对 SYCL 调整编译参数。

使用 sycl_extension 就像编写如下的 setup.py 一样简单

import os
import torch
import glob
import platform
from setuptools import find_packages, setup
from torch.utils.cpp_extension import SyclExtension, BuildExtension

library_name = "sycl_extension"
py_limited_api = True

IS_WINDOWS = (platform.system() == 'Windows')

if IS_WINDOWS:
    cxx_args = [
        "/O2",
        "/std:c++17",
        "/DPy_LIMITED_API=0x03090000",
    ]
    sycl_args = ["/O2", "/std:c++17"]
else:
    cxx_args = [
        "-O3",
        "-fdiagnostics-color=always",
        "-DPy_LIMITED_API=0x03090000"
    ]
    sycl_args = ["-O3"]

extra_compile_args = {
    "cxx": cxx_args,
    "sycl": sycl_args
}

assert(torch.xpu.is_available()), "XPU is not available, please check your environment"

# Source files collection
this_dir = os.path.dirname(os.path.curdir)
extensions_dir = os.path.join(this_dir, library_name)
sources = list(glob.glob(os.path.join(extensions_dir, "*.sycl")))

# Construct extension
ext_modules = [
    SyclExtension(
        f"{library_name}._C",
        sources,
        extra_compile_args=extra_compile_args,
        py_limited_api=py_limited_api,
    )
]

setup(
    name=library_name,
    packages=find_packages(),
    ext_modules=ext_modules,
    install_requires=["torch"],
    description="Simple Example of PyTorch Sycl extensions",
    cmdclass={"build_ext": BuildExtension},
    options={"bdist_wheel": {"py_limited_api": "cp39"}} if py_limited_api else {},
)

定义自定义算子并添加后端实现

首先，让我们编写一个计算 mymuladd 的 SYCL 函数

为了从 PyTorch 的 Python 前端使用它，我们需要使用 TORCH_LIBRARY API 将其注册为 PyTorch 算子。这将自动将算子绑定到 Python。

如果您还有一个 myaddmul 的 SYCL 实现，也可以将其注册在单独的 TORCH_LIBRARY_IMPL 代码块中

#include <c10/xpu/XPUStream.h>
#include <sycl/sycl.hpp>
#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>


#include <Python.h>

namespace sycl_extension {

// ==========================================================
// 1. Kernel
// ==========================================================
static void muladd_kernel(
    int numel, const float* a, const float* b, float c, float* result,
    const sycl::nd_item<1>& item) {
    int idx = item.get_global_id(0);
    if (idx < numel) {
        result[idx] = a[idx] * b[idx] + c;
    }
}

class MulAddKernelFunctor {
public:
    MulAddKernelFunctor(int _numel, const float* _a, const float* _b, float _c, float* _result)
        : numel(_numel), a(_a), b(_b), c(_c), result(_result) {}
    void operator()(const sycl::nd_item<1>& item) const {
        muladd_kernel(numel, a, b, c, result, item);
    }

private:
    int numel;
    const float* a;
    const float* b;
    float c;
    float* result;
};

// ==========================================================
// 2. Wrapper
// ==========================================================
at::Tensor mymuladd_xpu(const at::Tensor& a, const at::Tensor& b, double c) {
    TORCH_CHECK(a.sizes() == b.sizes(), "a and b must have the same shape");
    TORCH_CHECK(a.dtype() == at::kFloat, "a must be a float tensor");
    TORCH_CHECK(b.dtype() == at::kFloat, "b must be a float tensor");
    TORCH_CHECK(a.device().is_xpu(), "a must be an XPU tensor");
    TORCH_CHECK(b.device().is_xpu(), "b must be an XPU tensor");

    at::Tensor a_contig = a.contiguous();
    at::Tensor b_contig = b.contiguous();
    at::Tensor result = at::empty_like(a_contig);

    const float* a_ptr = a_contig.data_ptr<float>();
    const float* b_ptr = b_contig.data_ptr<float>();
    float* res_ptr = result.data_ptr<float>();
    int numel = a_contig.numel();

    sycl::queue& queue = c10::xpu::getCurrentXPUStream().queue();
    constexpr int threads = 256;
    int blocks = (numel + threads - 1) / threads;

    queue.submit([&](sycl::handler& cgh) {
        cgh.parallel_for<MulAddKernelFunctor>(
            sycl::nd_range<1>(blocks * threads, threads),
            MulAddKernelFunctor(numel, a_ptr, b_ptr, static_cast<float>(c), res_ptr)
        );
    });

    return result;
}

// ==========================================================
// 3. Registration
// ==========================================================
TORCH_LIBRARY(sycl_extension, m) {
m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
}

TORCH_LIBRARY_IMPL(sycl_extension, XPU, m) {
    m.impl("mymuladd", &mymuladd_xpu);
}

} // namespace sycl_extension

// ==========================================================
// 4. Windows Linker
// ==========================================================
extern "C" {
    #ifdef _WIN32
    __declspec(dllexport)
    #endif
    PyObject* PyInit__C(void) {
        static struct PyModuleDef moduledef = {
            PyModuleDef_HEAD_INIT,
            "_C",
            "XPU Extension Shim",
            -1,
            NULL
        };
        return PyModule_Create(&moduledef);
    }
}

创建 Python 接口

在 sycl_extension/ops.py 文件中为我们的算子创建 Python 接口

import torch
from torch import Tensor
__all__ = ["mymuladd"]

def mymuladd(a: Tensor, b: Tensor, c: float) -> Tensor:
    """Performs a * b + c in an efficient fused kernel"""
    return torch.ops.sycl_extension.mymuladd.default(a, b, c)

初始化包

创建 sycl_extension/__init__.py 文件以使该包可导入

import ctypes
import platform
from pathlib import Path

import torch

current_dir = Path(__file__).parent.parent
build_dir = current_dir / "build"

if platform.system() == 'Windows':
    file_pattern = "**/*.pyd"
else:
    file_pattern = "**/*.so"

lib_files = list(build_dir.glob(file_pattern))

if not lib_files:
    current_package_dir = Path(__file__).parent
    lib_files = list(current_package_dir.glob(file_pattern))

assert len(lib_files) > 0, f"Could not find any {file_pattern} file in {build_dir} or {current_dir}"
lib_file = lib_files[0]


with torch._ops.dl_open_guard():
    loaded_lib = ctypes.CDLL(str(lib_file))

from . import ops

__all__ = [
    "loaded_lib",
    "ops",
]

测试 SYCL 扩展算子

使用简单的测试来验证算子是否工作正常。

import torch
from torch.testing._internal.common_utils import TestCase
import unittest
import sycl_extension

def reference_muladd(a, b, c):
    return a * b + c

class TestMyMulAdd(TestCase):
    def sample_inputs(self, device, *, requires_grad=False):
        def make_tensor(*size):
            return torch.randn(size, device=device, requires_grad=requires_grad)

        def make_nondiff_tensor(*size):
            return torch.randn(size, device=device, requires_grad=False)

        return [
            [make_tensor(3), make_tensor(3), 1],
            [make_tensor(20), make_tensor(20), 3.14],
            [make_tensor(20), make_nondiff_tensor(20), -123],
            [make_nondiff_tensor(2, 3), make_tensor(2, 3), -0.3],
        ]

    def _test_correctness(self, device):
        samples = self.sample_inputs(device)
        for args in samples:
            result = sycl_extension.ops.mymuladd(*args)
            expected = reference_muladd(*args)
            torch.testing.assert_close(result, expected)

    @unittest.skipIf(not torch.xpu.is_available(), "requires Intel GPU")
    def test_correctness_xpu(self):
        self._test_correctness("xpu")

if __name__ == "__main__":
    unittest.main()

此测试通过将自定义算子的输出与参考实现进行比较，来检查其正确性。

结论

在本教程中，我们演示了如何为 PyTorch 实现和编译自定义 SYCL 算子。我们特别展示了一个推理操作 muladd。如需添加反向传播支持或启用 torch.compile 兼容性，请参考 自定义 C++ 和 CUDA 算子。