1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
|
#include <torch/extension.h>
// Declare the function from cuda_dlink_extension.cu.
void add_cuda(const float* a, const float* b, float* output, int size);
at::Tensor add(at::Tensor a, at::Tensor b) {
TORCH_CHECK(a.device().is_cuda(), "a is a cuda tensor");
TORCH_CHECK(b.device().is_cuda(), "b is a cuda tensor");
TORCH_CHECK(a.dtype() == at::kFloat, "a is a float tensor");
TORCH_CHECK(b.dtype() == at::kFloat, "b is a float tensor");
TORCH_CHECK(a.sizes() == b.sizes(), "a and b should have same size");
at::Tensor output = at::empty_like(a);
add_cuda(a.data_ptr<float>(), b.data_ptr<float>(), output.data_ptr<float>(), a.numel());
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("add", &add, "a + b");
}
|
