forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
FillKernel.cu
30 lines (25 loc) · 922 Bytes
/
FillKernel.cu
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
#define TORCH_ASSERT_NO_OPERATORS
#include <ATen/Dispatch.h>
#include <ATen/Dispatch_v2.h>
#include <ATen/native/cuda/Loops.cuh>
#include <ATen/native/DispatchStub.h>
#include <ATen/native/TensorIterator.h>
#include <ATen/native/Fill.h>
#include <c10/core/Scalar.h>
namespace at::native {
template<typename scalar_t>
struct FillFunctor {
FillFunctor(scalar_t v): value(v) {}
__device__ __forceinline__ scalar_t operator() () const {
return value;
}
private:
scalar_t value;
};
void fill_kernel_cuda(TensorIterator& iter, const Scalar& value) {
AT_DISPATCH_V2(iter.dtype(), "fill_cuda", AT_WRAP([&]() {
gpu_kernel(iter, FillFunctor<scalar_t>(value.to<scalar_t>()));
}), AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), kComplexHalf, kBool, kHalf, kBFloat16, kFloat8_e4m3fn, kFloat8_e5m2, AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));
}
REGISTER_DISPATCH(fill_stub, &fill_kernel_cuda);
} // namespace at::native