#include <torch/extension.h>
#include <vector>

// 前向傳播
torch::Tensor Test_forward_cpu(const torch::Tensor& inputA,
              const torch::Tensor& inputB);
// 反向傳播
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput);

#include "test.h"

// 前向傳播，兩個(gè) Tensor 相加。這里只關(guān)注 C++ 擴(kuò)展的流程，具體實(shí)現(xiàn)不深入探討。
torch::Tensor Test_forward_cpu(const torch::Tensor& x,
              const torch::Tensor& y) {
  AT_ASSERTM(x.sizes() == y.sizes(), "x must be the same size as y");
  torch::Tensor z = torch::zeros(x.sizes());
  z = 2 * x + y;
  return z;
}

// 反向傳播
// 在這個(gè)例子中，z對(duì)x的導(dǎo)數(shù)是2，z對(duì)y的導(dǎo)數(shù)是1。
// 至于這個(gè)backward函數(shù)的接口（參數(shù)，返回值）為何要這樣設(shè)計(jì)，后面會(huì)講。
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput) {
  torch::Tensor gradOutputX = 2 * gradOutput * torch::ones(gradOutput.sizes());
  torch::Tensor gradOutputY = gradOutput * torch::ones(gradOutput.sizes());
  return {gradOutputX, gradOutputY};
}

// pybind11 綁定
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 m.def("forward", &Test_forward_cpu, "TEST forward");
 m.def("backward", &Test_backward_cpu, "TEST backward");
}

└── csrc
  ├── cpu
  │  ├── test.cpp
  │  └── test.h
  └── setup.py

from setuptools import setup
import os
import glob
from torch.utils.cpp_extension import BuildExtension, CppExtension

# 頭文件目錄
include_dirs = os.path.dirname(os.path.abspath(__file__))
# 源代碼目錄
source_cpu = glob.glob(os.path.join(include_dirs, 'cpu', '*.cpp'))

setup(
  name='test_cpp', # 模塊名稱，需要在python中調(diào)用
  version="0.1",
  ext_modules=[
    CppExtension('test_cpp', sources=source_cpu, include_dirs=[include_dirs]),
  ],
  cmdclass={
    'build_ext': BuildExtension
  }
)

python setup.py install

from torch.autograd import Function

import test_cpp

class TestFunction(Function):

  @staticmethod
  def forward(ctx, x, y):
    return test_cpp.forward(x, y)

  @staticmethod
  def backward(ctx, gradOutput):
    gradX, gradY = test_cpp.backward(gradOutput)
    return gradX, gradY

class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin)):
 
  ...

  @staticmethod
  def forward(ctx, *args, **kwargs):
    r"""Performs the operation.

    This function is to be overridden by all subclasses.

    It must accept a context ctx as the first argument, followed by any
    number of arguments (tensors or other types).

    The context can be used to store tensors that can be then retrieved
    during the backward pass.
    """
    raise NotImplementedError

  @staticmethod
  def backward(ctx, *grad_outputs):
    r"""Defines a formula for differentiating the operation.

    This function is to be overridden by all subclasses.

    It must accept a context :attr:`ctx` as the first argument, followed by
    as many outputs did :func:`forward` return, and it should return as many
    tensors, as there were inputs to :func:`forward`. Each argument is the
    gradient w.r.t the given output, and each returned value should be the
    gradient w.r.t. the corresponding input.

    The context can be used to retrieve tensors saved during the forward
    pass. It also has an attribute :attr:`ctx.needs_input_grad` as a tuple
    of booleans representing whether each input needs gradient. E.g.,
    :func:`backward` will have ``ctx.needs_input_grad[0] = True`` if the
    first input to :func:`forward` needs gradient computated w.r.t. the
    output.
    """
    raise NotImplementedError

import torch

class Test(torch.nn.Module):

  def __init__(self):
    super(Test, self).__init__()

  def forward(self, inputA, inputB):
    return TestFunction.apply(inputA, inputB)

├── csrc
│  ├── cpu
│  │  ├── test.cpp
│  │  └── test.h
│  └── setup.py
└── test.py

import torch
from torch.autograd import Variable

from test import Test

x = Variable(torch.Tensor([1,2,3]), requires_grad=True)
y = Variable(torch.Tensor([4,5,6]), requires_grad=True)
test = Test()
z = test(x, y)
z.sum().backward()
print('x: ', x)
print('y: ', y)
print('z: ', z)
print('x.grad: ', x.grad)
print('y.grad: ', y.grad)