| #include "caffe2/utils/math/elementwise.h" |
| |
| #include <type_traits> |
| |
| #include <thrust/execution_policy.h> |
| #include <thrust/fill.h> |
| #include <thrust/functional.h> |
| #include <thrust/transform.h> |
| |
| #include "caffe2/core/context_gpu.h" |
| #include "caffe2/utils/conversions.h" |
| #include "caffe2/utils/math/half_utils.h" |
| #include "caffe2/utils/math/utils.h" |
| |
| namespace caffe2 { |
| namespace math { |
| |
| namespace { |
| |
| template <typename T> |
| __global__ void SinCosCUDAKernel(const int N, const T* X, T* S, T* C) { |
| const int i = blockIdx.x * CAFFE_CUDA_NUM_THREADS + threadIdx.x; |
| if (i < N) { |
| #if __CUDA_ARCH__ >= 350 || defined(USE_ROCM) |
| c10::cuda::compat::sincos(__ldg(X + i), S + i, C + i); |
| #else |
| c10::cuda::compat::sincos(X[i], S + i, C + i); |
| #endif |
| } |
| } |
| |
| #if defined(USE_ROCM) |
| |
| template <typename TAlpha, typename TData> |
| __global__ void AxpyCUDAKernel( |
| const std::int64_t N, |
| const TAlpha alpha, |
| const TData* X, |
| TData* Y) { |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + |
| static_cast<int64_t>(threadIdx.x); |
| if (index < N) { |
| Y[index] += static_cast<TData>(alpha) * __ldg(X + index); |
| } |
| } |
| |
| template <typename TAlpha, typename TData> |
| __global__ void AxpyCUDAKernel( |
| const std::int64_t N, |
| const TAlpha* alpha, |
| const TData* X, |
| TData* Y) { |
| __shared__ TData a; |
| if (threadIdx.x == 0) { |
| a = static_cast<TData>(__ldg(alpha)); |
| } |
| __syncthreads(); |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + |
| static_cast<int64_t>(threadIdx.x); |
| if (index < N) { |
| Y[index] += a * __ldg(X + index); |
| } |
| } |
| |
| #define DELEGATE_HALF_AXPY_CUDA_KERNEL(TAlpha, FMAFunc) \ |
| template <> \ |
| __global__ void AxpyCUDAKernel<TAlpha, at::Half>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const at::Half* X, \ |
| at::Half* Y) { \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = convert::To<TAlpha, at::Half>(FMAFunc( \ |
| alpha, \ |
| convert::To<at::Half, TAlpha>(X[index]), \ |
| convert::To<at::Half, TAlpha>(Y[index]))); \ |
| } \ |
| } \ |
| template <> \ |
| __global__ void AxpyCUDAKernel<TAlpha, at::Half>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const at::Half* X, \ |
| at::Half* Y) { \ |
| __shared__ TAlpha a; \ |
| if (threadIdx.x == 0) { \ |
| a = __ldg(alpha); \ |
| } \ |
| __syncthreads(); \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = convert::To<TAlpha, at::Half>(FMAFunc( \ |
| a, \ |
| convert::To<at::Half, TAlpha>(X[index]), \ |
| convert::To<at::Half, TAlpha>(Y[index]))); \ |
| } \ |
| } |
| DELEGATE_HALF_AXPY_CUDA_KERNEL(float, fmaf) |
| #undef DELEGATE_HALF_AXPY_CUDA_KERNEL |
| |
| #endif // USE_ROCM |
| |
| template <typename TAlpha, typename TData> |
| __global__ void AxpbyCUDAKernel( |
| const std::int64_t N, |
| const TAlpha alpha, |
| const TData* X, |
| const TAlpha beta, |
| TData* Y); |
| |
| template <typename TAlpha, typename TData> |
| __global__ void AxpbyCUDAKernel( |
| const std::int64_t N, |
| const TAlpha* alpha, |
| const TData* X, |
| const TAlpha* beta, |
| TData* Y); |
| |
| #define DELEGATE_AXPBY_CUDA_KERNEL(TAlpha, TData, FMAFunc) \ |
| template <> \ |
| __global__ void AxpbyCUDAKernel<TAlpha, TData>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const TData* X, \ |
| const TAlpha beta, \ |
| TData* Y) { \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = FMAFunc( \ |
| static_cast<TData>(alpha), \ |
| X[index], \ |
| static_cast<TData>(beta) * Y[index]); \ |
| } \ |
| } \ |
| template <> \ |
| __global__ void AxpbyCUDAKernel<TAlpha, TData>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const TData* X, \ |
| const TAlpha* beta, \ |
| TData* Y) { \ |
| __shared__ TData a; \ |
| __shared__ TData b; \ |
| if (threadIdx.x == 0) { \ |
| a = static_cast<TData>(*alpha); \ |
| b = static_cast<TData>(*beta); \ |
| } \ |
| __syncthreads(); \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = FMAFunc(a, X[index], b * Y[index]); \ |
| } \ |
| } |
| DELEGATE_AXPBY_CUDA_KERNEL(float, float, fmaf) |
| DELEGATE_AXPBY_CUDA_KERNEL(float, double, fma) |
| #undef DELEGATE_AXPBY_CUDA_KERNEL |
| |
| #define DELEGATE_HALF_AXPBY_CUDA_KERNEL(TAlpha, FMAFunc) \ |
| template <> \ |
| __global__ void AxpbyCUDAKernel<TAlpha, at::Half>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const at::Half* X, \ |
| const TAlpha beta, \ |
| at::Half* Y) { \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = convert::To<TAlpha, at::Half>(FMAFunc( \ |
| alpha, \ |
| convert::To<at::Half, TAlpha>(X[index]), \ |
| beta * convert::To<at::Half, TAlpha>(Y[index]))); \ |
| } \ |
| } \ |
| template <> \ |
| __global__ void AxpbyCUDAKernel<TAlpha, at::Half>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const at::Half* X, \ |
| const TAlpha* beta, \ |
| at::Half* Y) { \ |
| __shared__ TAlpha a; \ |
| __shared__ TAlpha b; \ |
| if (threadIdx.x == 0) { \ |
| a = *alpha; \ |
| b = *beta; \ |
| } \ |
| __syncthreads(); \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = convert::To<TAlpha, at::Half>(FMAFunc( \ |
| a, \ |
| convert::To<at::Half, TAlpha>(X[index]), \ |
| b * convert::To<at::Half, TAlpha>(Y[index]))); \ |
| } \ |
| } |
| DELEGATE_HALF_AXPBY_CUDA_KERNEL(float, fmaf) |
| #undef DELEGATE_HALF_AXPBY_CUDA_KERNEL |
| |
| template <typename TAlpha, typename TData> |
| __global__ void ScaleCUDAKernel( |
| const std::int64_t N, |
| const TAlpha alpha, |
| const TData* X, |
| TData* Y); |
| |
| template <typename TAlpha, typename TData> |
| __global__ void ScaleCUDAKernel( |
| const std::int64_t N, |
| const TAlpha* alpha, |
| const TData* X, |
| TData* Y); |
| |
| #define CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL(TAlpha, TData) \ |
| template <> \ |
| __global__ void ScaleCUDAKernel<TAlpha, TData>( \ |
| const std::int64_t N, const TAlpha alpha, const TData* X, TData* Y) { \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = static_cast<TData>(alpha) * X[index]; \ |
| } \ |
| } \ |
| template <> \ |
| __global__ void ScaleCUDAKernel<TAlpha, TData>( \ |
| const std::int64_t N, const TAlpha* alpha, const TData* X, TData* Y) { \ |
| __shared__ TData a; \ |
| if (threadIdx.x == 0) { \ |
| a = static_cast<TData>(*alpha); \ |
| } \ |
| __syncthreads(); \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = a * X[index]; \ |
| } \ |
| } |
| CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL(float, float) |
| CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL(double, double) |
| CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL(float, double) |
| CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL(std::int32_t, std::int32_t) |
| CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL(std::int64_t, std::int64_t) |
| #undef CAFFE2_SPECIALIZED_SCALE_CUDA_KERNEL |
| |
| #define CAFFE2_SPECIALIZED_HALF_SCALE_CUDA_KERNEL(TAlpha) \ |
| template <> \ |
| __global__ void ScaleCUDAKernel<TAlpha, at::Half>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const at::Half* X, \ |
| at::Half* Y) { \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = convert::To<TAlpha, at::Half>( \ |
| alpha * convert::To<at::Half, TAlpha>(X[index])); \ |
| } \ |
| } \ |
| template <> \ |
| __global__ void ScaleCUDAKernel<TAlpha, at::Half>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const at::Half* X, \ |
| at::Half* Y) { \ |
| __shared__ TAlpha a; \ |
| if (threadIdx.x == 0) { \ |
| a = *alpha; \ |
| } \ |
| __syncthreads(); \ |
| const int64_t index = static_cast<int64_t>(blockIdx.x) * \ |
| static_cast<int64_t>(CAFFE_CUDA_NUM_THREADS) + \ |
| static_cast<int64_t>(threadIdx.x); \ |
| if (index < N) { \ |
| Y[index] = convert::To<TAlpha, at::Half>( \ |
| a * convert::To<at::Half, TAlpha>(X[index])); \ |
| } \ |
| } |
| CAFFE2_SPECIALIZED_HALF_SCALE_CUDA_KERNEL(float) |
| #undef CAFFE2_SPECIALIZED_HALF_SCALE_CUDA_KERNEL |
| |
| } // namespace |
| |
| #define CAFFE2_SPECIALIZED_CUDA_SET(T) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Set<T, CUDAContext>( \ |
| const std::int64_t N, const T alpha, T* Y, CUDAContext* context) { \ |
| if (N == 0) { \ |
| return; \ |
| } \ |
| if (alpha == T(0)) { \ |
| C10_CUDA_CHECK(cudaMemsetAsync(Y, 0, sizeof(T) * N, context->cuda_stream())); \ |
| } else { \ |
| thrust::fill( \ |
| thrust::cuda::par.on(context->cuda_stream()), Y, Y + N, alpha); \ |
| } \ |
| } |
| CAFFE2_SPECIALIZED_CUDA_SET(bool) |
| CAFFE2_SPECIALIZED_CUDA_SET(char) |
| CAFFE2_SPECIALIZED_CUDA_SET(std::int8_t) |
| CAFFE2_SPECIALIZED_CUDA_SET(std::int16_t) |
| CAFFE2_SPECIALIZED_CUDA_SET(std::int32_t) |
| CAFFE2_SPECIALIZED_CUDA_SET(std::int64_t) |
| CAFFE2_SPECIALIZED_CUDA_SET(std::uint8_t) |
| CAFFE2_SPECIALIZED_CUDA_SET(std::uint16_t) |
| CAFFE2_SPECIALIZED_CUDA_SET(float) |
| CAFFE2_SPECIALIZED_CUDA_SET(double) |
| CAFFE2_SPECIALIZED_CUDA_SET(at::Half) |
| CAFFE2_SPECIALIZED_CUDA_SET(at::BFloat16) |
| #undef CAFFE2_SPECIALIZED_CUDA_SET |
| |
| #define DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(T, Func, DeviceFunc) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Func<T, CUDAContext>( \ |
| const int N, const T* X, T* Y, CUDAContext* context) { \ |
| if (N > 0) { \ |
| thrust::transform( \ |
| thrust::cuda::par.on(context->cuda_stream()), \ |
| X, \ |
| X + N, \ |
| Y, \ |
| [] __device__(const T x) { return DeviceFunc(x); }); \ |
| } \ |
| } |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Exp, expf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Log, logf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Log1p, log1pf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Sin, sinf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Asin, asinf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Cos, cosf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Acos, acosf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Tan, tanf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Atan, atanf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Sinh, sinhf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Cosh, coshf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Tanh, tanhf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Abs, fabsf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Inv, utils::Inv<float>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(double, Inv, utils::Inv<double>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Sqr, utils::Square<float>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Sqrt, sqrtf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Rsqrt, rsqrtf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION( |
| std::int32_t, |
| Cube, |
| utils::Cube<std::int32_t>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION( |
| std::int64_t, |
| Cube, |
| utils::Cube<std::int64_t>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Cube, utils::Cube<float>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(double, Cube, utils::Cube<double>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Cbrt, cbrtf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Erf, erff) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(double, Erf, erf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, CdfNorm, normcdff) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(double, CdfNorm, normcdf) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(bool, Not, utils::Not<bool>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION( |
| std::int32_t, |
| Neg, |
| utils::Negate<std::int32_t>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION( |
| std::int64_t, |
| Neg, |
| utils::Negate<std::int64_t>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Neg, utils::Negate<float>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(double, Neg, utils::Negate<double>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION( |
| std::int32_t, |
| Sign, |
| utils::Sign<std::int32_t>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION( |
| std::int64_t, |
| Sign, |
| utils::Sign<std::int64_t>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(float, Sign, utils::Sign<float>) |
| DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION(double, Sign, utils::Sign<double>) |
| #undef DELEGATE_SIMPLE_CUDA_UNARY_FUNCTION |
| |
| #define DELEGATE_CUDA_POWX(T, DeviceFunc) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Powx<T, CUDAContext>( \ |
| const int N, const T* A, const T b, T* Y, CUDAContext* context) { \ |
| thrust::transform( \ |
| thrust::cuda::par.on(context->cuda_stream()), \ |
| A, \ |
| A + N, \ |
| Y, \ |
| [b] __device__(const T x) { return DeviceFunc(x, b); }); \ |
| } |
| DELEGATE_CUDA_POWX(float, powf) |
| #undef DELEGATE_CUDA_POWX |
| |
| #define CAFFE2_SPECIALIZED_CUDA_SINCOS(T) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void SinCos<T, CUDAContext>( \ |
| const int N, const T* X, T* S, T* C, CUDAContext* context) { \ |
| if (N > 0) { \ |
| const int K = DivUp(N, CAFFE_CUDA_NUM_THREADS); \ |
| SinCosCUDAKernel<T> \ |
| <<<K, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, X, S, C); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } |
| CAFFE2_SPECIALIZED_CUDA_SINCOS(float) |
| CAFFE2_SPECIALIZED_CUDA_SINCOS(double) |
| #undef CAFFE2_SPECIALIZED_CUDA_SINCOS |
| |
| #define DELEGATE_CUDA_SCALE(T, CuBLASFunc) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Scale<T, T, CUDAContext>( \ |
| const std::int64_t N, \ |
| const T alpha, \ |
| const T* X, \ |
| T* Y, \ |
| CUDAContext* context) { \ |
| if (N == 0) { \ |
| return; \ |
| } \ |
| if (Y == X) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_HOST)); \ |
| CUBLAS_ENFORCE(CuBLASFunc(context->cublas_handle(), N, &alpha, Y, 1)); \ |
| } else { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| ScaleCUDAKernel<T, T> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Scale<T, T, CUDAContext>( \ |
| const std::int64_t N, \ |
| const T* alpha, \ |
| const T* X, \ |
| T* Y, \ |
| CUDAContext* context) { \ |
| if (N == 0) { \ |
| return; \ |
| } \ |
| if (Y == X) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_DEVICE)); \ |
| CUBLAS_ENFORCE(CuBLASFunc(context->cublas_handle(), N, alpha, Y, 1)); \ |
| } else { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| ScaleCUDAKernel<T, T> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } |
| DELEGATE_CUDA_SCALE(float, cublasSscal) |
| DELEGATE_CUDA_SCALE(double, cublasDscal) |
| #undef DELEGATE_CUDA_SCALE |
| |
| #if !defined(USE_ROCM) |
| |
| #define DELEGATE_CUDA_SCALE_EX( \ |
| TAlpha, TData, kAlphaType, kDataType, kExecutionType) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Scale<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| if (N == 0) { \ |
| return; \ |
| } \ |
| if (Y == X) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_HOST)); \ |
| CUBLAS_ENFORCE(cublasScalEx( \ |
| context->cublas_handle(), \ |
| N, \ |
| &alpha, \ |
| kAlphaType, \ |
| Y, \ |
| kDataType, \ |
| 1, \ |
| kExecutionType)); \ |
| } else { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| ScaleCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Scale<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| if (N == 0) { \ |
| return; \ |
| } \ |
| if (Y == X) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_DEVICE)); \ |
| CUBLAS_ENFORCE(cublasScalEx( \ |
| context->cublas_handle(), \ |
| N, \ |
| alpha, \ |
| kAlphaType, \ |
| Y, \ |
| kDataType, \ |
| 1, \ |
| kExecutionType)); \ |
| } else { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| ScaleCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } |
| DELEGATE_CUDA_SCALE_EX(float, double, CUDA_R_32F, CUDA_R_64F, CUDA_R_64F) |
| DELEGATE_CUDA_SCALE_EX(float, at::Half, CUDA_R_32F, CUDA_R_16F, CUDA_R_32F) |
| #undef DELEGATE_CUDA_SCALE_EX |
| |
| #endif // USE_ROCM |
| |
| #define CAFFE2_SPECIALIZED_CUDA_SCALE(TAlpha, TData) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Scale<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| if (N > 0) { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| ScaleCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Scale<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| if (N > 0) { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| ScaleCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, *alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| } |
| CAFFE2_SPECIALIZED_CUDA_SCALE(std::int32_t, std::int32_t) |
| CAFFE2_SPECIALIZED_CUDA_SCALE(std::int64_t, std::int64_t) |
| |
| #if defined(USE_ROCM) |
| CAFFE2_SPECIALIZED_CUDA_SCALE(float, double) |
| CAFFE2_SPECIALIZED_CUDA_SCALE(float, at::Half) |
| #endif // USE_ROCM |
| #undef CAFFE2_SPECIALIZED_CUDA_SCALE |
| |
| #define DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(T, Func, DeviceFunc) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Func<T, CUDAContext>( \ |
| const int N, const T* A, const T* B, T* C, CUDAContext* context) { \ |
| if (N > 0) { \ |
| thrust::transform( \ |
| thrust::cuda::par.on(context->cuda_stream()), \ |
| A, \ |
| A + N, \ |
| B, \ |
| C, \ |
| DeviceFunc); \ |
| } \ |
| } |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| Add, |
| thrust::plus<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| Add, |
| thrust::plus<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(float, Add, thrust::plus<float>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(double, Add, thrust::plus<double>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(at::Half, Add, utils::HalfAddFunctor()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| Sub, |
| thrust::minus<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| Sub, |
| thrust::minus<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(float, Sub, thrust::minus<float>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(double, Sub, thrust::minus<double>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(at::Half, Sub, utils::HalfSubFunctor()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| Mul, |
| thrust::multiplies<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| Mul, |
| thrust::multiplies<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(float, Mul, thrust::multiplies<float>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(double, Mul, thrust::multiplies<double>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(at::Half, Mul, utils::HalfMulFunctor()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| Div, |
| thrust::divides<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| Div, |
| thrust::divides<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(float, Div, thrust::divides<float>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(double, Div, thrust::divides<double>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(at::Half, Div, utils::HalfDivFunctor()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(float, Min, thrust::minimum<float>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(double, Min, thrust::minimum<double>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(float, Max, thrust::maximum<float>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(double, Max, thrust::maximum<double>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(bool, And, thrust::logical_and<bool>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(bool, Or, thrust::logical_or<bool>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(bool, Xor, thrust::bit_xor<bool>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(bool, BitwiseAnd, thrust::bit_and<bool>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| BitwiseAnd, |
| thrust::bit_and<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| BitwiseAnd, |
| thrust::bit_and<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(bool, BitwiseOr, thrust::bit_or<bool>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| BitwiseOr, |
| thrust::bit_or<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| BitwiseOr, |
| thrust::bit_or<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION(bool, BitwiseXor, thrust::bit_xor<bool>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int32_t, |
| BitwiseXor, |
| thrust::bit_xor<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION( |
| std::int64_t, |
| BitwiseXor, |
| thrust::bit_xor<std::int64_t>()) |
| #undef DELEGATE_SIMPLE_CUDA_BINARY_FUNCTION |
| |
| #define DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(T, Func, DeviceComp) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Func<T, CUDAContext>( \ |
| const int N, const T* A, const T* B, bool* C, CUDAContext* context) { \ |
| if (N > 0) { \ |
| thrust::transform( \ |
| thrust::cuda::par.on(context->cuda_stream()), \ |
| A, \ |
| A + N, \ |
| B, \ |
| C, \ |
| DeviceComp); \ |
| } \ |
| } |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(bool, EQ, thrust::equal_to<bool>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int32_t, |
| EQ, |
| thrust::equal_to<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int64_t, |
| EQ, |
| thrust::equal_to<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(float, EQ, thrust::equal_to<float>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(double, EQ, thrust::equal_to<double>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(bool, NE, thrust::not_equal_to<bool>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int32_t, |
| NE, |
| thrust::not_equal_to<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int64_t, |
| NE, |
| thrust::not_equal_to<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(float, NE, thrust::not_equal_to<float>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| double, |
| NE, |
| thrust::not_equal_to<double>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(bool, LT, thrust::less<bool>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int32_t, |
| LT, |
| thrust::less<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int64_t, |
| LT, |
| thrust::less<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(float, LT, thrust::less<float>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(double, LT, thrust::less<double>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(bool, LE, thrust::less_equal<bool>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int32_t, |
| LE, |
| thrust::less_equal<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int64_t, |
| LE, |
| thrust::less_equal<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(float, LE, thrust::less_equal<float>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(double, LE, thrust::less_equal<double>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(bool, GT, thrust::greater<bool>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int32_t, |
| GT, |
| thrust::greater<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int64_t, |
| GT, |
| thrust::greater<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(float, GT, thrust::greater<float>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(double, GT, thrust::greater<double>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(bool, GE, thrust::greater_equal<bool>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int32_t, |
| GE, |
| thrust::greater_equal<std::int32_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| std::int64_t, |
| GE, |
| thrust::greater_equal<std::int64_t>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION(float, GE, thrust::greater_equal<float>()) |
| DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION( |
| double, |
| GE, |
| thrust::greater_equal<double>()) |
| #undef DELEGATE_SIMPLE_CUDA_COMPARE_FUNCTION |
| |
| #define DELEGATE_CUDA_AXPY(T, CuBLASFunc) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpy<T, T, CUDAContext>( \ |
| const std::int64_t N, \ |
| const T alpha, \ |
| const T* X, \ |
| T* Y, \ |
| CUDAContext* context) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_HOST)); \ |
| CUBLAS_ENFORCE( \ |
| CuBLASFunc(context->cublas_handle(), N, &alpha, X, 1, Y, 1)); \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpy<T, T, CUDAContext>( \ |
| const std::int64_t N, \ |
| const T* alpha, \ |
| const T* X, \ |
| T* Y, \ |
| CUDAContext* context) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_DEVICE)); \ |
| CUBLAS_ENFORCE( \ |
| cublasSaxpy(context->cublas_handle(), N, alpha, X, 1, Y, 1)); \ |
| } |
| DELEGATE_CUDA_AXPY(float, cublasSaxpy) |
| #undef DELEGATE_CUDA_AXPY |
| |
| #if !defined(USE_ROCM) |
| |
| #define DELEGATE_CUDA_AXPY_EX( \ |
| TAlpha, TData, kAlphaType, kDataType, kExecutionType) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpy<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_HOST)); \ |
| CUBLAS_ENFORCE(cublasAxpyEx( \ |
| context->cublas_handle(), \ |
| N, \ |
| &alpha, \ |
| kAlphaType, \ |
| X, \ |
| kDataType, \ |
| 1, \ |
| Y, \ |
| kDataType, \ |
| 1, \ |
| kExecutionType)); \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpy<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| CUBLAS_ENFORCE(cublasSetPointerMode( \ |
| context->cublas_handle(), CUBLAS_POINTER_MODE_DEVICE)); \ |
| CUBLAS_ENFORCE(cublasAxpyEx( \ |
| context->cublas_handle(), \ |
| N, \ |
| alpha, \ |
| kAlphaType, \ |
| X, \ |
| kDataType, \ |
| 1, \ |
| Y, \ |
| kDataType, \ |
| 1, \ |
| kExecutionType)); \ |
| } |
| DELEGATE_CUDA_AXPY_EX(float, double, CUDA_R_32F, CUDA_R_64F, CUDA_R_64F) |
| DELEGATE_CUDA_AXPY_EX(float, at::Half, CUDA_R_32F, CUDA_R_16F, CUDA_R_32F) |
| #undef DELEGATE_CUDA_AXPY_EX |
| |
| #else // USE_ROCM |
| |
| #define CAFFE2_SPECIALIZED_CUDA_AXPY(TAlpha, TData) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpy<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| AxpyCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpy<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const TData* X, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| AxpyCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } |
| CAFFE2_SPECIALIZED_CUDA_AXPY(float, double) |
| CAFFE2_SPECIALIZED_CUDA_AXPY(float, at::Half) |
| #undef CAFFE2_SPECIALIZED_CUDA_AXPY |
| |
| #endif // USE_ROCM |
| |
| #define CAFFE2_SPECIALIZED_CUDA_AXPBY(TAlpha, TData) \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpby<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha alpha, \ |
| const TData* X, \ |
| const TAlpha beta, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| AxpbyCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, beta, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } \ |
| template <> \ |
| CAFFE2_CUDA_EXPORT void Axpby<TAlpha, TData, CUDAContext>( \ |
| const std::int64_t N, \ |
| const TAlpha* alpha, \ |
| const TData* X, \ |
| const TAlpha* beta, \ |
| TData* Y, \ |
| CUDAContext* context) { \ |
| const std::int64_t M = DivUp<std::int64_t>(N, CAFFE_CUDA_NUM_THREADS); \ |
| AxpbyCUDAKernel<TAlpha, TData> \ |
| <<<M, CAFFE_CUDA_NUM_THREADS, 0, context->cuda_stream()>>>( \ |
| N, alpha, X, beta, Y); \ |
| C10_CUDA_KERNEL_LAUNCH_CHECK(); \ |
| } |
| CAFFE2_SPECIALIZED_CUDA_AXPBY(float, float) |
| CAFFE2_SPECIALIZED_CUDA_AXPBY(float, double) |
| CAFFE2_SPECIALIZED_CUDA_AXPBY(float, at::Half) |
| #undef CAFFE2_SPECIALIZED_CUDA_AXPBY |
| |
| } // namespace math |
| } // namespace caffe2 |
