interface.cpp 3.32 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
#include <torch/extension.h>
#include <ATen/Generator.h>
#include <ATen/cuda/CUDAGeneratorImpl.h>
#include <vector>

std::vector<c10::optional<torch::Tensor>> fwd_cuda(
    bool is_training,
8
9
10
    torch::Tensor &input,
    const c10::optional<torch::Tensor> &attn_mask,
    const c10::optional<torch::Tensor> &bias,
Guolin Ke's avatar
Guolin Ke committed
11
    float dropout_prob,
12
    c10::optional<at::Generator> gen_);
Guolin Ke's avatar
Guolin Ke committed
13
14
15
16
17

torch::Tensor bwd_cuda(
    torch::Tensor &output_grads,
    const torch::Tensor &softmax_results,
    const c10::optional<torch::Tensor> &dropout_mask,
18
    float dropout_prob);
Guolin Ke's avatar
Guolin Ke committed
19
20
21
22
23

// C++ interface

#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
24
25
26
#define CHECK_INPUT(x) \
    CHECK_CUDA(x);     \
    CHECK_CONTIGUOUS(x)
Guolin Ke's avatar
Guolin Ke committed
27
28
29

std::vector<c10::optional<torch::Tensor>> fwd(
    bool is_training,
Guolin Ke's avatar
Guolin Ke committed
30
    torch::Tensor &input,
31
32
    const c10::optional<torch::Tensor> &attn_mask,
    const c10::optional<torch::Tensor> &bias,
Guolin Ke's avatar
Guolin Ke committed
33
    float dropout_prob,
34
35
    c10::optional<at::Generator> gen_)
{
Guolin Ke's avatar
Guolin Ke committed
36
    CHECK_INPUT(input);
37
38
39
40
41
42
43
44
45
46
47
48
    if (attn_mask)
    {
        CHECK_INPUT(attn_mask.value());
        AT_ASSERTM(attn_mask->dim() == 3, "expected 3D tensor");
    }
    if (bias)
    {
        CHECK_INPUT(bias.value());
        AT_ASSERTM(bias->dim() == 3, "expected 3D tensor");
        AT_ASSERTM(input.size(0) % bias->size(0) == 0, "wrong first dim of bias.");
        AT_ASSERTM(bias->size(1) == input.size(1) && bias->size(2) == input.size(2), "the last two dims of bias and input should be the same.");
    }
Guolin Ke's avatar
Guolin Ke committed
49
50
    AT_ASSERTM(input.dim() == 3, "expected 3D tensor");
    AT_ASSERTM(input.scalar_type() == at::ScalarType::Half ||
51
52
53
54
                   input.scalar_type() == at::ScalarType::BFloat16 ||
                   input.scalar_type() == at::ScalarType::Float,
               "Only HALF/BFloat16/Float is supported");
    return fwd_cuda(is_training, input, attn_mask, bias, dropout_prob, gen_);
Guolin Ke's avatar
Guolin Ke committed
55
56
57
}

torch::Tensor bwd(
58
    torch::Tensor &output_grads,
Guolin Ke's avatar
Guolin Ke committed
59
60
    const torch::Tensor &softmax_results,
    const c10::optional<torch::Tensor> &dropout_mask,
61
62
    float dropout_prob)
{
Guolin Ke's avatar
Guolin Ke committed
63
64
    CHECK_INPUT(output_grads);
    CHECK_INPUT(softmax_results);
65
66
    if (dropout_mask)
    {
Guolin Ke's avatar
Guolin Ke committed
67
68
        CHECK_INPUT(dropout_mask.value());
    }
69
    AT_ASSERTM(output_grads.dim() == 3, "expected 3D tensor");
Guolin Ke's avatar
Guolin Ke committed
70
71
72
73
    AT_ASSERTM(softmax_results.dim() == 3, "expected 3D tensor");
    AT_ASSERTM(!dropout_mask || dropout_mask->dim() == 1, "expected 1D tensor");

    AT_ASSERTM(output_grads.scalar_type() == at::ScalarType::Half ||
74
75
76
                   output_grads.scalar_type() == at::ScalarType::BFloat16 ||
                   output_grads.scalar_type() == at::ScalarType::Float,
               "Only HALF/BFloat16/Float is supported");
Guolin Ke's avatar
Guolin Ke committed
77
    AT_ASSERTM(softmax_results.scalar_type() == at::ScalarType::Half ||
78
79
80
                   softmax_results.scalar_type() == at::ScalarType::BFloat16 ||
                   softmax_results.scalar_type() == at::ScalarType::Float,
               "Only HALF/BFloat16/Float is supported");
Guolin Ke's avatar
Guolin Ke committed
81
82
83
84
    AT_ASSERTM(output_grads.scalar_type() == softmax_results.scalar_type(), "the types mismatch");
    return bwd_cuda(output_grads, softmax_results, dropout_mask, dropout_prob);
}

85
86
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
{
Guolin Ke's avatar
Guolin Ke committed
87
88
89
    m.def("forward", &fwd, "softmax dropout -- Forward.");
    m.def("backward", &bwd, "softmax dropout -- Backward.");
}