CUDAExtension¶

paddle.utils.cpp_extension. CUDAExtension ( sources, *args, **kwargs ) [source]

The interface is used to config source files of customized operators and complies Op Kernel supporting both CPU and GPU devices. Please use CppExtension if you want to compile Op Kernel that supports only CPU device.

It further encapsulates python built-in setuptools.Extension .The arguments and usage are same as the native interface, except for no need to explicitly specify name .

A simple example:

           # setup.py

# Compiling customized operators supporting CPU and GPU devices
from paddle.utils.cpp_extension import CUDAExtension, setup

setup(
    name='custom_op',
    ext_modules=CUDAExtension(
        sources=['relu_op.cc', 'relu_op.cu']
    )
)

          

Note

It is mainly used in setup and the nama of built shared library keeps same as name argument specified in setup interface.

Parameters

sources (list[str]) – Specify the C++/CUDA source files of customized operators.
*args (list[options], optional) – Specify other arguments same as setuptools.Extension .
**kwargs (dict[option], optional) – Specify other arguments same as setuptools.Extension .

Returns

An instance of setuptools.Extension.

Return type

setuptools.Extension