自定义 C++ 和 CUDA 运算符

TORCH_LIBRARY(extension_cpp,m){
// Note that "float" in the schema corresponds to the C++ double type
// and the Python float type.
m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
}

TORCH_LIBRARY_IMPL(extension_cpp,CPU,m){
m.impl("mymuladd",&mymuladd_cpu);
}

__global__voidmuladd_kernel(intnumel,constfloat*a,constfloat*b,floatc,float*result){
intidx=blockIdx.x*blockDim.x+threadIdx.x;
if(idx<numel)result[idx]=a[idx]*b[idx]+c;
}

at::Tensormymuladd_cuda(constat::Tensor&a,constat::Tensor&b,doublec){
TORCH_CHECK(a.sizes()==b.sizes());
TORCH_CHECK(a.dtype()==at::kFloat);
TORCH_CHECK(b.dtype()==at::kFloat);
TORCH_INTERNAL_ASSERT(a.device().type()==at::DeviceType::CUDA);
TORCH_INTERNAL_ASSERT(b.device().type()==at::DeviceType::CUDA);
at::Tensora_contig=a.contiguous();
at::Tensorb_contig=b.contiguous();
at::Tensorresult=torch::empty(a_contig.sizes(),a_contig.options());
constfloat*a_ptr=a_contig.data_ptr<float>();
constfloat*b_ptr=b_contig.data_ptr<float>();
float*result_ptr=result.data_ptr<float>();

intnumel=a_contig.numel();
muladd_kernel<<<(numel+255)/256,256>>>(numel,a_ptr,b_ptr,c,result_ptr);
returnresult;
}

TORCH_LIBRARY_IMPL(extension_cpp,CUDA,m){
m.impl("mymuladd",&mymuladd_cuda);
}

# Important: the C++ custom operator definitions should be loaded first
# before calling ``torch.library`` APIs that add registrations for the
# C++ custom operator(s). The following import loads our
# C++ custom operator definitions.
# Note that if you are striving for Python agnosticism, you should use
# the ``load_library(...)`` API call instead. See the next section for
# more details.
from.import _C

@torch.library.register_fake("extension_cpp::mymuladd")
def_(a, b, c):
    torch._check(a.shape == b.shape)
    torch._check(a.dtype == torch.float)
    torch._check(b.dtype == torch.float)
    torch._check(a.device == b.device)
    return torch.empty_like(a)

// in, say, not_agnostic/csrc/extension_BAD.cpp
#include<pybind11/pybind11.h>

PYBIND11_MODULE("_C",m){}

# in, say, extension/__init__.py
from.import _C

#include<Python.h>

extern"C"{
/* Creates a dummy empty _C module that can be imported from Python.
    The import from Python will load the .so consisting of this file
    in this extension, so that the TORCH_LIBRARY static initializers
    below are run. */
PyObject*PyInit__C(void)
{
staticstructPyModuleDefmodule_def={
PyModuleDef_HEAD_INIT,
"_C",/* name of module */
NULL,/* module documentation, may be NULL */
*1,/* size of per-interpreter state of the module,
                    or -1 if the module keeps state in global variables. */
NULL,/* methods */
};
returnPyModule_Create(&module_def);
}
}

# in, say, extension/__init__.py
from.import _C

importtorch
frompathlibimport Path

so_files = list(Path(__file__).parent.glob("_C*.so"))
assert (
    len(so_files) == 1
), f"Expected one _C*.so file, found {len(so_files)}"
torch.ops.load_library(so_files[0])

from.import ops

def_backward(ctx, grad):
    a, b = ctx.saved_tensors
    grad_a, grad_b = None, None
    if ctx.needs_input_grad[0]:
        grad_a = grad * b
    if ctx.needs_input_grad[1]:
        grad_b = grad * a
    return grad_a, grad_b, None

def_setup_context(ctx, inputs, output):
    a, b, c = inputs
    saved_a, saved_b = None, None
    if ctx.needs_input_grad[0]:
        saved_b = b
    if ctx.needs_input_grad[1]:
        saved_a = a
    ctx.save_for_backward(saved_a, saved_b)

# This code adds training support for the operator. You must provide us
# the backward formula for the operator and a `setup_context` function
# to save values to be used in the backward.
torch.library.register_autograd(
    "extension_cpp::mymuladd", _backward, setup_context=_setup_context)

// New! a mymul_cpu kernel
at::Tensormymul_cpu(constat::Tensor&a,constat::Tensor&b){
TORCH_CHECK(a.sizes()==b.sizes());
TORCH_CHECK(a.dtype()==at::kFloat);
TORCH_CHECK(b.dtype()==at::kFloat);
TORCH_CHECK(a.device().type()==at::DeviceType::CPU);
TORCH_CHECK(b.device().type()==at::DeviceType::CPU);
at::Tensora_contig=a.contiguous();
at::Tensorb_contig=b.contiguous();
at::Tensorresult=torch::empty(a_contig.sizes(),a_contig.options());
constfloat*a_ptr=a_contig.data_ptr<float>();
constfloat*b_ptr=b_contig.data_ptr<float>();
float*result_ptr=result.data_ptr<float>();
for(int64_ti=0;i<result.numel();i++){
result_ptr[i]=a_ptr[i]*b_ptr[i];
}
returnresult;
}

TORCH_LIBRARY(extension_cpp,m){
m.def("mymuladd(Tensor a, Tensor b, float c) -> Tensor");
// New! defining the mymul operator
m.def("mymul(Tensor a, Tensor b) -> Tensor");
}


TORCH_LIBRARY_IMPL(extension_cpp,CPU,m){
m.impl("mymuladd",&mymuladd_cpu);
// New! registering the cpu kernel for the mymul operator
m.impl("mymul",&mymul_cpu);
}

def_backward(ctx, grad):
    a, b = ctx.saved_tensors
    grad_a, grad_b = None, None
    if ctx.needs_input_grad[0]:
        grad_a = torch.ops.extension_cpp.mymul.default(grad, b)
    if ctx.needs_input_grad[1]:
        grad_b = torch.ops.extension_cpp.mymul.default(grad, a)
    return grad_a, grad_b, None


def_setup_context(ctx, inputs, output):
    a, b, c = inputs
    saved_a, saved_b = None, None
    if ctx.needs_input_grad[0]:
        saved_b = b
    if ctx.needs_input_grad[1]:
        saved_a = a
    ctx.save_for_backward(saved_a, saved_b)


# This code adds training support for the operator. You must provide us
# the backward formula for the operator and a `setup_context` function
# to save values to be used in the backward.
torch.library.register_autograd(
    "extension_cpp::mymuladd", _backward, setup_context=_setup_context)