PyTorch扩展自定义PyThon/C++(CUDA)算子的若干方法总结

// PyG的C++扩展就选择的是直接继承PyTorch的C++端的torch::autograd类进行扩展// 下面是PyG的一个ScatterSum算子的扩展示例// 不用纠结这个算子的具体内容，对扩展的算子的结构有一个大致了解即可class ScatterSum : public torch::autograd::Function {public:  // AutogradContext *ctx指针可以操作  static variable_list forward(AutogradContext *ctx, Variable src,                               Variable index, int64_t dim,                               torch::optional optional_out,                               torch::optional dim_size) {    dim = dim < 0 ? src.dim() + dim : dim;    ctx->saved_data["dim"] = dim;    ctx->saved_data["src_shape"] = src.sizes();    index = broadcast(index, src, dim);    auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");    auto out = std::get<0>(result);    ctx->save_for_backward({index});    // 如果在扩展的C++代码中使用非Aten内建操作修改了tensor的值，需要对其进行脏标记    if (optional_out.has_value())      ctx->mark_dirty({optional_out.value()});      return {out};  } // grad_outs是out参数反传回来的梯度值  static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {    auto grad_out = grad_outs[0];    auto saved = ctx->get_saved_variables();    auto index = saved[0];    auto dim = ctx->saved_data["dim"].toInt();    auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());    auto grad_in = torch::gather(grad_out, dim, index, false);    // 不需要求导的参数需要空Variable占位    return {grad_in, Variable(), Variable(), Variable(), Variable()};  }};