239 lines
11 KiB
C++
239 lines
11 KiB
C++
/**
|
|
* This file includes code originally from the Pytorch Scatter repository:
|
|
* https://github.com/rusty1s/pytorch_scatter
|
|
* Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
|
|
*/
|
|
|
|
#include <torch/extension.h>
|
|
#include <vector>
|
|
|
|
#include "cuda/scatter_cuda.h"
|
|
|
|
inline std::vector<int64_t> list2vec(const c10::List<int64_t> list) {
|
|
std::vector<int64_t> result;
|
|
result.reserve(list.size());
|
|
for (size_t i = 0; i < list.size(); i++) result.push_back(list[i]);
|
|
return result;
|
|
}
|
|
|
|
torch::Tensor broadcast(torch::Tensor src, torch::Tensor other, int64_t dim) {
|
|
if (src.dim() == 1)
|
|
for (auto i = 0; i < dim; i++) src = src.unsqueeze(0);
|
|
for (auto i = src.dim(); i < other.dim(); i++) src = src.unsqueeze(-1);
|
|
src = src.expand(other.sizes().vec());
|
|
return src;
|
|
}
|
|
|
|
std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_fw(torch::Tensor src, torch::Tensor index, int64_t dim,
|
|
std::optional<torch::Tensor> optional_out,
|
|
std::optional<int64_t> dim_size,
|
|
std::string reduce) {
|
|
return scatter_cuda(src, index, dim, optional_out, dim_size, reduce);
|
|
}
|
|
|
|
using torch::autograd::AutogradContext;
|
|
using torch::autograd::Variable;
|
|
using torch::autograd::variable_list;
|
|
|
|
class ScatterSum : public torch::autograd::Function<ScatterSum> {
|
|
public:
|
|
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
|
std::optional<Variable> optional_out, std::optional<int64_t> 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});
|
|
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
|
return {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);
|
|
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
|
}
|
|
};
|
|
|
|
class ScatterMul : public torch::autograd::Function<ScatterMul> {
|
|
public:
|
|
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
|
std::optional<Variable> optional_out, std::optional<int64_t> 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, "mul");
|
|
auto out = std::get<0>(result);
|
|
ctx->save_for_backward({src, index, out});
|
|
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
|
return {out};
|
|
}
|
|
|
|
static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
|
|
auto grad_out = grad_outs[0];
|
|
auto saved = ctx->get_saved_variables();
|
|
auto src = saved[0];
|
|
auto index = saved[1];
|
|
auto out = saved[2];
|
|
auto dim = ctx->saved_data["dim"].toInt();
|
|
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
|
auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);
|
|
grad_in.masked_fill_(grad_in.isnan(), 0);
|
|
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
|
}
|
|
};
|
|
|
|
class ScatterMean : public torch::autograd::Function<ScatterMean> {
|
|
public:
|
|
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
|
std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
|
|
dim = dim < 0 ? src.dim() + dim : dim;
|
|
ctx->saved_data["dim"] = dim;
|
|
ctx->saved_data["src_shape"] = src.sizes();
|
|
|
|
auto old_index = index;
|
|
|
|
index = broadcast(index, src, dim);
|
|
auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
|
|
auto out = std::get<0>(result);
|
|
|
|
auto ones = torch::ones(old_index.sizes(), src.options());
|
|
result = scatter_fw(ones, old_index, old_index.dim() <= dim ? old_index.dim() - 1 : dim, std::nullopt,
|
|
out.size(dim), "sum");
|
|
auto count = std::get<0>(result);
|
|
count.masked_fill_(count < 1, 1);
|
|
count = broadcast(count, out, dim);
|
|
if (out.is_floating_point())
|
|
out.true_divide_(count);
|
|
else
|
|
out.div_(count, "floor");
|
|
|
|
ctx->save_for_backward({index, count});
|
|
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
|
return {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 count = saved[1];
|
|
auto dim = ctx->saved_data["dim"].toInt();
|
|
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
|
count = torch::gather(count, dim, index, false);
|
|
auto grad_in = torch::gather(grad_out, dim, index, false);
|
|
grad_in.true_divide_(count);
|
|
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
|
}
|
|
};
|
|
|
|
class ScatterMin : public torch::autograd::Function<ScatterMin> {
|
|
public:
|
|
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
|
std::optional<Variable> optional_out, std::optional<int64_t> 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, "min");
|
|
auto out = std::get<0>(result);
|
|
auto arg_out = std::get<1>(result).value();
|
|
ctx->save_for_backward({index, arg_out});
|
|
ctx->mark_non_differentiable({arg_out});
|
|
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
|
return {out, arg_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 arg_out = saved[1];
|
|
auto dim = ctx->saved_data["dim"].toInt();
|
|
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
|
src_shape[dim] += 1;
|
|
auto grad_in = torch::zeros(src_shape, grad_out.options());
|
|
grad_in.scatter_(dim, arg_out, grad_out);
|
|
grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
|
|
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
|
}
|
|
};
|
|
|
|
class ScatterMax : public torch::autograd::Function<ScatterMax> {
|
|
public:
|
|
static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
|
|
std::optional<Variable> optional_out, std::optional<int64_t> 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, "max");
|
|
auto out = std::get<0>(result);
|
|
auto arg_out = std::get<1>(result).value();
|
|
ctx->save_for_backward({index, arg_out});
|
|
ctx->mark_non_differentiable({arg_out});
|
|
if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
|
|
return {out, arg_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 arg_out = saved[1];
|
|
auto dim = ctx->saved_data["dim"].toInt();
|
|
auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
|
|
src_shape[dim] += 1;
|
|
auto grad_in = torch::zeros(src_shape, grad_out.options());
|
|
grad_in.scatter_(dim, arg_out, grad_out);
|
|
grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
|
|
return {grad_in, Variable(), Variable(), Variable(), Variable()};
|
|
}
|
|
};
|
|
|
|
torch::Tensor scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
|
|
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
|
return ScatterSum::apply(src, index, dim, optional_out, dim_size)[0];
|
|
}
|
|
|
|
torch::Tensor scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
|
|
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
|
return ScatterMul::apply(src, index, dim, optional_out, dim_size)[0];
|
|
}
|
|
|
|
torch::Tensor scatter_mean(torch::Tensor src, torch::Tensor index, int64_t dim,
|
|
std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
|
|
return ScatterMean::apply(src, index, dim, optional_out, dim_size)[0];
|
|
}
|
|
|
|
std::tuple<torch::Tensor, torch::Tensor> scatter_min(torch::Tensor src, torch::Tensor index, int64_t dim,
|
|
std::optional<torch::Tensor> optional_out,
|
|
std::optional<int64_t> dim_size) {
|
|
auto result = ScatterMin::apply(src, index, dim, optional_out, dim_size);
|
|
return std::make_tuple(result[0], result[1]);
|
|
}
|
|
|
|
std::tuple<torch::Tensor, torch::Tensor> scatter_max(torch::Tensor src, torch::Tensor index, int64_t dim,
|
|
std::optional<torch::Tensor> optional_out,
|
|
std::optional<int64_t> dim_size) {
|
|
auto result = ScatterMax::apply(src, index, dim, optional_out, dim_size);
|
|
return std::make_tuple(result[0], result[1]);
|
|
}
|
|
|
|
void pybind_scatter_ext(py::module &m) {
|
|
m.def("scatter_sum", &scatter_sum, "Scatter sum operation");
|
|
m.def("scatter_mul", &scatter_mul, "Scatter multiplication operation");
|
|
m.def("scatter_mean", &scatter_mean, "Scatter mean operation");
|
|
m.def("scatter_min", &scatter_min, "Scatter min operation");
|
|
m.def("scatter_max", &scatter_max, "Scatter max operation");
|
|
}
|