1 #include "caffe2/core/context_gpu.h" 2 #include "caffe2/operators/flatten_op.h" 3 #include "caffe2/operators/utility_ops.h" 4 #include "caffe2/utils/math.h" 9 bool WeightedSumOp<CUDAContext>::RunOnDevice() {
10 if (Input(0).IsType<float>()) {
11 return DoRunWithType<float>();
12 }
else if (Input(0).IsType<float16>()) {
13 return DoRunWithType<float16>();
15 CAFFE_THROW(
"Unsupported inputs");
21 bool SumOp<CUDAContext>::RunOnDevice() {
22 if (Input(0).IsType<float>()) {
23 return DoRunWithType<float, float>();
24 }
else if (Input(0).IsType<float16>()) {
25 return DoRunWithType<float16, float16>();
27 CAFFE_THROW(
"Unsupported inputs");
40 bool RunOnDevice()
override {
41 auto& input = Input(0);
42 auto* output = OperatorBase::Output<Tensor<CUDAContext>>(0);
44 output->ResizeLike(input);
45 context.template CopyItems<CUDAContext, CUDAContext>(
49 output->raw_mutable_data(input.meta()));
62 REGISTER_CUDA_OPERATOR(
66 REGISTER_CUDA_OPERATOR(
72 REGISTER_CUDA_OPERATOR(
75 REGISTER_CUDA_OPERATOR(
82 REGISTER_CUDA_OPERATOR(
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
int GetGPUIDForPointer(const void *ptr)
Gets the GPU id that the current pointer is located at.
Alias op makes the output and the input share the same underlying storage.