mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
838fdd6f99
Summary: Closes https://github.com/pytorch/pytorch/pull/8991 Add Cube and Cbrt Ops Reviewed By: houseroad Differential Revision: D8678848 fbshipit-source-id: 051dd475e45ad9f1d11a8b32ae3acd1f7459b930
56 lines
1.2 KiB
Plaintext
56 lines
1.2 KiB
Plaintext
#include "caffe2/operators/cube_op.h"
|
|
|
|
#include <algorithm>
|
|
#include <functional>
|
|
|
|
#include "caffe2/core/context_gpu.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
namespace {
|
|
|
|
template <typename T>
|
|
__global__ void
|
|
CubeGradientCUDAKernel(const int N, const T* dY, const T* X, T* dX) {
|
|
CUDA_1D_KERNEL_LOOP(i, N) {
|
|
#if __CUDA_ARCH__ >= 350
|
|
dX[i] = __ldg(dY + i) * __ldg(X + i) * __ldg(X + i) * T(3);
|
|
#else
|
|
dX[i] = dY[i] * X[i] * X[i] * T(3);
|
|
#endif
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
template <>
|
|
template <typename T>
|
|
bool CubeGradientFunctor<CUDAContext>::Forward(
|
|
const std::vector<int>& dY_dims,
|
|
const std::vector<int>& /* X_dims */,
|
|
const T* dY,
|
|
const T* X,
|
|
T* dX,
|
|
CUDAContext* context) const {
|
|
const int size = std::accumulate(
|
|
dY_dims.cbegin(), dY_dims.cend(), 1, std::multiplies<int>());
|
|
CubeGradientCUDAKernel<T>
|
|
<<<CAFFE_GET_BLOCKS(size),
|
|
CAFFE_CUDA_NUM_THREADS,
|
|
0,
|
|
context->cuda_stream()>>>(size, dY, X, dX);
|
|
return true;
|
|
}
|
|
|
|
REGISTER_CUDA_OPERATOR(
|
|
Cube,
|
|
UnaryElementwiseOp<NumericTypes, CUDAContext, CubeFunctor<CUDAContext>>);
|
|
REGISTER_CUDA_OPERATOR(
|
|
CubeGradient,
|
|
BinaryElementwiseOp<
|
|
NumericTypes,
|
|
CUDAContext,
|
|
CubeGradientFunctor<CUDAContext>>);
|
|
|
|
} // namespace caffe2
|