logsumexpv2/act.cpp at main · zacmar/logsumexpv2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
#include <torch/extension.h>
#include <vector>
#define GLOG_USE_GLOG_EXPORT

// CUDA forward declarations
std::vector<torch::Tensor> cuda_logsumexp_pot_act_forward(
    const torch::Tensor x, const torch::Tensor weight, const torch::Tensor mu,
    const torch::Tensor sigma);

torch::Tensor cuda_categorical_sampler(
    const torch::Tensor x, const torch::Tensor weight, const torch::Tensor mu,
    const torch::Tensor sigma);

torch::Tensor cuda_gig_sampler(
    const torch::Tensor a, const torch::Tensor b, const torch::Tensor p);

std::vector<torch::Tensor> cuda_logsumexp_pot_act_backward(
    const torch::Tensor x, const torch::Tensor weight,
    const torch::Tensor grad_out, const torch::Tensor grad_out_prime,
    const torch::Tensor mu, const torch::Tensor sigma,
    const torch::Tensor max_exp);

// C++ interface
#define CHECK_CUDA(x)                                                          \
  TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x)                                                    \
  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x)

std::vector<torch::Tensor>
logsumexp_pot_act_forward(const torch::Tensor x, const torch::Tensor weight,
                          const torch::Tensor mu, const torch::Tensor sigma) {
  CHECK_INPUT(x);
  CHECK_INPUT(weight);
  CHECK_INPUT(sigma);
  CHECK_INPUT(mu);

  return cuda_logsumexp_pot_act_forward(x, weight, mu, sigma);
}

torch::Tensor
categorical_sampler(const torch::Tensor x, const torch::Tensor weight,
                    const torch::Tensor mu, const torch::Tensor sigma) {
  CHECK_INPUT(x);
  CHECK_INPUT(weight);
  CHECK_INPUT(sigma);
  CHECK_INPUT(mu);

  return cuda_categorical_sampler(x, weight, mu, sigma);
}

torch::Tensor
gig_sampler(const torch::Tensor a, const torch::Tensor b, const torch::Tensor p) {
  CHECK_INPUT(a);
  CHECK_INPUT(b);
  CHECK_INPUT(p);

  return cuda_gig_sampler(a, b, p);
}

std::vector<torch::Tensor>
logsumexp_pot_act_backward(const torch::Tensor x, const torch::Tensor weight,
                           const torch::Tensor f_grad,
                           const torch::Tensor f_grad_prime,
                           const torch::Tensor mu, const torch::Tensor sigma,
                           const torch::Tensor max_exp) {
  CHECK_INPUT(x);
  CHECK_INPUT(weight);
  CHECK_INPUT(f_grad);
  CHECK_INPUT(f_grad_prime);
  CHECK_INPUT(sigma);
  CHECK_INPUT(mu);

  return cuda_logsumexp_pot_act_backward(x, weight, f_grad, f_grad_prime, mu,
                                         sigma, max_exp);
}

// python interface
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("lse_forward", &logsumexp_pot_act_forward,
        "logsumexp potential activation forward");
  m.def("categorical_sampler", &categorical_sampler,
        "categorical sampler");
  m.def("gig_sampler", &gig_sampler,
        "GIG sampler");
  m.def("lse_backward", &logsumexp_pot_act_backward,
        "logsumexp potential activation backward");
}