mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-24 19:43:35 +00:00
Add CUDA Kernels of RandomNormal[Like], RandomUniform[Like] (#9761)
This commit is contained in:
parent
6284cbe833
commit
f390347c11
8 changed files with 602 additions and 4 deletions
|
|
@ -600,6 +600,10 @@ Do not modify directly.*
|
|||
|QuantizeLinear|*in* x:**T1**<br> *in* y_scale:**tensor(float)**<br> *in* y_zero_point:**T2**<br> *out* y:**T2**|10+|**T1** = tensor(float)<br/> **T2** = tensor(int8), tensor(uint8)|
|
||||
|RNN|*in* X:**T**<br> *in* W:**T**<br> *in* R:**T**<br> *in* B:**T**<br> *in* sequence_lens:**T1**<br> *in* initial_h:**T**<br> *out* Y:**T**<br> *out* Y_h:**T**|14+|**T** = tensor(double), tensor(float), tensor(float16)<br/> **T1** = tensor(int32)|
|
||||
|||[7, 13]|**T** = tensor(double), tensor(float), tensor(float16)<br/> **T1** = tensor(int32)|
|
||||
|RandomNormal|*out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|RandomNormalLike|*in* input:**T1**<br> *out* output:**T2**|1+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T2** = tensor(double), tensor(float), tensor(float16)|
|
||||
|RandomUniform|*out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|RandomUniformLike|*in* input:**T1**<br> *out* output:**T2**|1+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T2** = tensor(double), tensor(float), tensor(float16)|
|
||||
|Range|*in* start:**T**<br> *in* limit:**T**<br> *in* delta:**T**<br> *out* output:**T**|11+|**T** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64)|
|
||||
|Reciprocal|*in* X:**T**<br> *out* Y:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|||[6, 12]|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|
|
|
|||
|
|
@ -716,6 +716,10 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDom
|
|||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, Loop);
|
||||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, DepthToSpace);
|
||||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 12, SpaceToDepth);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomNormal);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomNormalLike);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomUniform);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomUniformLike);
|
||||
|
||||
// opset 10
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool);
|
||||
|
|
@ -1561,6 +1565,10 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, Loop)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, DepthToSpace)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 12, SpaceToDepth)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomNormal)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomNormalLike)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomUniform)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, RandomUniformLike)>,
|
||||
|
||||
// opset 10
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool)>,
|
||||
|
|
|
|||
78
onnxruntime/core/providers/cuda/generator/random.cc
Normal file
78
onnxruntime/core/providers/cuda/generator/random.cc
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/providers/cuda/generator/random.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace cuda {
|
||||
|
||||
using namespace ONNX_NAMESPACE;
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(RandomNormal, kOnnxDomain, 1, kCudaExecutionProvider,
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()),
|
||||
RandomNormal);
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(RandomNormalLike, kOnnxDomain, 1, kCudaExecutionProvider,
|
||||
(*KernelDefBuilder::Create())
|
||||
.TypeConstraint("T1", DataTypeImpl::AllTensorTypes())
|
||||
.TypeConstraint("T2", DataTypeImpl::AllIEEEFloatTensorTypes()),
|
||||
RandomNormalLike);
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(RandomUniform, kOnnxDomain, 1, kCudaExecutionProvider,
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes()),
|
||||
RandomUniform);
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(RandomUniformLike, kOnnxDomain, 1, kCudaExecutionProvider,
|
||||
(*KernelDefBuilder::Create())
|
||||
.TypeConstraint("T1", DataTypeImpl::AllTensorTypes())
|
||||
.TypeConstraint("T2", DataTypeImpl::AllIEEEFloatTensorTypes()),
|
||||
RandomUniformLike);
|
||||
|
||||
Status RandomNormalBase::Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const {
|
||||
Tensor& Y = *p_ctx->Output(0, shape);
|
||||
const int64_t N = shape.Size();
|
||||
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
|
||||
utils::MLTypeCallDispatcher<float, MLFloat16, double> t_disp(dtype);
|
||||
t_disp.Invoke<RandomNormalComputeImpl>(GetDeviceProp(), Stream(), N, scale_, mean_, generator, Y);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status RandomNormal::ComputeInternal(OpKernelContext* p_ctx) const { return Compute(p_ctx, shape_, dtype_); }
|
||||
|
||||
Status RandomNormalLike::ComputeInternal(OpKernelContext* p_ctx) const {
|
||||
const Tensor* p_X = p_ctx->Input<Tensor>(0);
|
||||
if (!p_X) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
|
||||
if (dtype_ == TensorProto_DataType_UNDEFINED && !p_X->IsDataType<float>() && !p_X->IsDataType<double>() &&
|
||||
!p_X->IsDataType<MLFloat16>()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
|
||||
"Output data type is required to be one of float types, but got incompatible data type ",
|
||||
p_X->DataType(), " from input tensor.");
|
||||
}
|
||||
return Compute(p_ctx, p_X->Shape(), dtype_ != TensorProto_DataType_UNDEFINED ? dtype_ : p_X->GetElementType());
|
||||
}
|
||||
|
||||
Status RandomUniformBase::Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const {
|
||||
Tensor& Y = *p_ctx->Output(0, shape);
|
||||
const int64_t N = shape.Size();
|
||||
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
|
||||
utils::MLTypeCallDispatcher<float, MLFloat16, double> t_disp(dtype);
|
||||
t_disp.Invoke<RandomUniformComputeImpl>(GetDeviceProp(), Stream(), N, range_, from_, generator, Y);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status RandomUniform::ComputeInternal(OpKernelContext* p_ctx) const { return Compute(p_ctx, shape_, dtype_); }
|
||||
|
||||
Status RandomUniformLike::ComputeInternal(OpKernelContext* p_ctx) const {
|
||||
const Tensor* p_X = p_ctx->Input<Tensor>(0);
|
||||
if (!p_X) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
|
||||
if (dtype_ == TensorProto_DataType_UNDEFINED && !p_X->IsDataType<float>() && !p_X->IsDataType<double>() &&
|
||||
!p_X->IsDataType<MLFloat16>()) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
|
||||
"Output data type is required to be one of float types, but got incompatible data type ",
|
||||
p_X->DataType(), " from input tensor.");
|
||||
}
|
||||
return Compute(p_ctx, p_X->Shape(), dtype_ != TensorProto_DataType_UNDEFINED ? dtype_ : p_X->GetElementType());
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace onnxruntime
|
||||
130
onnxruntime/core/providers/cuda/generator/random.h
Normal file
130
onnxruntime/core/providers/cuda/generator/random.h
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/providers/cuda/cuda_kernel.h"
|
||||
|
||||
#include "core/providers/cuda/generator/random_impl.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace cuda {
|
||||
|
||||
#define RANDOM_COMPUTE_IMPL(name) \
|
||||
template <typename T> \
|
||||
struct name##ComputeImpl { \
|
||||
void operator()(const cudaDeviceProp& prop, cudaStream_t stream, const int64_t N, const float alpha, \
|
||||
const float beta, PhiloxGenerator& generator, Tensor& Y) const { \
|
||||
typedef typename ToCudaType<T>::MappedType CudaT; \
|
||||
CudaT* Y_data = reinterpret_cast<CudaT*>(Y.template MutableData<T>()); \
|
||||
name##KernelImpl<CudaT>(prop, stream, N, alpha, beta, generator, Y_data); \
|
||||
} \
|
||||
};
|
||||
|
||||
RANDOM_COMPUTE_IMPL(RandomNormal)
|
||||
RANDOM_COMPUTE_IMPL(RandomUniform)
|
||||
|
||||
#undef RANDOM_COMPUTE_IMPL
|
||||
|
||||
class RandomBase : public CudaKernel {
|
||||
protected:
|
||||
RandomBase(const OpKernelInfo& info) : CudaKernel(info) {
|
||||
float seed = 0.f;
|
||||
if (info.GetAttr<float>("seed", &seed).IsOK()) {
|
||||
generator_ = std::make_unique<PhiloxGenerator>(static_cast<uint64_t>(seed));
|
||||
}
|
||||
|
||||
int64_t dtype;
|
||||
if (info.GetAttr<int64_t>("dtype", &dtype).IsOK()) {
|
||||
dtype_ = static_cast<ONNX_NAMESPACE::TensorProto::DataType>(dtype);
|
||||
ORT_ENFORCE(ONNX_NAMESPACE::TensorProto::DataType_IsValid(dtype_) &&
|
||||
dtype_ != ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED,
|
||||
"Invalid dtype of ", dtype_);
|
||||
}
|
||||
}
|
||||
|
||||
protected:
|
||||
std::unique_ptr<PhiloxGenerator> generator_;
|
||||
ONNX_NAMESPACE::TensorProto::DataType dtype_ =
|
||||
ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED; // optional and may be inferred
|
||||
};
|
||||
|
||||
class RandomNormalBase : public RandomBase {
|
||||
protected:
|
||||
RandomNormalBase(const OpKernelInfo& info) : RandomBase(info) {
|
||||
ORT_ENFORCE(info.GetAttr<float>("scale", &scale_).IsOK());
|
||||
ORT_ENFORCE(info.GetAttr<float>("mean", &mean_).IsOK());
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const;
|
||||
|
||||
protected:
|
||||
float scale_;
|
||||
float mean_;
|
||||
};
|
||||
|
||||
class RandomNormal final : public RandomNormalBase {
|
||||
public:
|
||||
explicit RandomNormal(const OpKernelInfo& info) : RandomNormalBase(info) {
|
||||
if (dtype_ == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED) {
|
||||
dtype_ = ONNX_NAMESPACE::TensorProto_DataType_FLOAT;
|
||||
}
|
||||
std::vector<int64_t> shape;
|
||||
ORT_ENFORCE(info.GetAttrs<int64_t>("shape", shape).IsOK());
|
||||
shape_ = TensorShape(shape);
|
||||
}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
|
||||
private:
|
||||
TensorShape shape_;
|
||||
};
|
||||
|
||||
class RandomNormalLike final : public RandomNormalBase {
|
||||
public:
|
||||
explicit RandomNormalLike(const OpKernelInfo& info) : RandomNormalBase(info) {}
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
};
|
||||
|
||||
class RandomUniformBase : public RandomBase {
|
||||
protected:
|
||||
RandomUniformBase(const OpKernelInfo& info) : RandomBase(info) {
|
||||
float low, high;
|
||||
ORT_ENFORCE(info.GetAttr<float>("low", &low).IsOK());
|
||||
ORT_ENFORCE(info.GetAttr<float>("high", &high).IsOK());
|
||||
from_ = low;
|
||||
range_ = high - low;
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const;
|
||||
|
||||
protected:
|
||||
float range_;
|
||||
float from_;
|
||||
};
|
||||
|
||||
class RandomUniform final : public RandomUniformBase {
|
||||
public:
|
||||
explicit RandomUniform(const OpKernelInfo& info) : RandomUniformBase(info) {
|
||||
if (dtype_ == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED) {
|
||||
dtype_ = ONNX_NAMESPACE::TensorProto_DataType_FLOAT;
|
||||
}
|
||||
std::vector<int64_t> shape;
|
||||
ORT_ENFORCE(info.GetAttrs<int64_t>("shape", shape).IsOK());
|
||||
shape_ = TensorShape(shape);
|
||||
}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
|
||||
private:
|
||||
TensorShape shape_;
|
||||
};
|
||||
|
||||
class RandomUniformLike final : public RandomUniformBase {
|
||||
public:
|
||||
explicit RandomUniformLike(const OpKernelInfo& info) : RandomUniformBase(info) {}
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
};
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace onnxruntime
|
||||
145
onnxruntime/core/providers/cuda/generator/random_impl.cu
Normal file
145
onnxruntime/core/providers/cuda/generator/random_impl.cu
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/providers/cuda/generator/random_impl.h"
|
||||
|
||||
#include <curand_kernel.h>
|
||||
#include <algorithm>
|
||||
#include "core/providers/cuda/cu_inc/common.cuh"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace cuda {
|
||||
|
||||
constexpr int UNROLL = 4;
|
||||
|
||||
struct DistFunc_RandomNormal {
|
||||
__device__ __inline__ float4 operator()(curandStatePhilox4_32_10_t* state) const { return curand_normal4(state); }
|
||||
};
|
||||
|
||||
struct DistFunc_RandomUniform {
|
||||
__device__ __inline__ float4 operator()(curandStatePhilox4_32_10_t* state) const { return curand_uniform4(state); }
|
||||
};
|
||||
|
||||
struct TransformFunc_RandomNormal {
|
||||
__device__ __inline__ float operator()(const float value, const float scale, const float mean) const {
|
||||
return value * scale + mean;
|
||||
}
|
||||
};
|
||||
|
||||
struct TransformFunc_RandomUniform {
|
||||
__device__ __inline__ float operator()(const float value, const float range, const float from) const {
|
||||
// reverse the bounds of curand4 from (0, 1] to [0, 1).
|
||||
// ref: https://github.com/pytorch/pytorch/blob/e795315c638228d4170f3797356c09a70b2ed4cd/aten/src/ATen/native/cuda/DistributionTemplates.h#L464
|
||||
float reverse_bound_value = value == 1.0f ? 0.0f : value;
|
||||
return reverse_bound_value * range + from;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename DistFuncT, typename TransformFuncT>
|
||||
__global__ void RandomKernel(const int64_t N, const std::pair<uint64_t, uint64_t> seeds, const DistFuncT& dist_func,
|
||||
const TransformFuncT& transform_func, const float alpha, const float beta, T* Y_data) {
|
||||
CUDA_LONG idx = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
CUDA_LONG step_size = gridDim.x * blockDim.x * UNROLL;
|
||||
|
||||
curandStatePhilox4_32_10_t state;
|
||||
curand_init(seeds.first, idx, seeds.second, &state);
|
||||
float4 rand;
|
||||
|
||||
// We ensure every thread generates the same number of random numbers (by rounding
|
||||
// up the size) and at the same timestep (by syncing threads).
|
||||
// From CUDA curand documentation:
|
||||
// The Philox_4x32_10 algorithm is closely tied to the thread and block count.
|
||||
// Each thread computes 4 random numbers in the same time thus the most efficient
|
||||
// use of Philox_4x32_10 is to generate a multiple of 4 times number of threads.
|
||||
for (CUDA_LONG id = idx * UNROLL; id < N; id += step_size) {
|
||||
rand = dist_func(&state);
|
||||
|
||||
// actual computation
|
||||
#pragma unroll
|
||||
for (int i = 0; i < UNROLL; i++) {
|
||||
CUDA_LONG li = id + i;
|
||||
if (li < N) {
|
||||
Y_data[li] = static_cast<T>(transform_func((&rand.x)[i], alpha, beta));
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename DistFuncT, typename TransformFuncT>
|
||||
__global__ void RandomVectorizedKernel(const int64_t N, const std::pair<uint64_t, uint64_t> seeds,
|
||||
const DistFuncT& dist_func, const TransformFuncT& transform_func,
|
||||
const float alpha, const float beta, T* Y_data) {
|
||||
CUDA_LONG idx = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
CUDA_LONG step_size = gridDim.x * blockDim.x * UNROLL;
|
||||
|
||||
curandStatePhilox4_32_10_t state;
|
||||
curand_init(seeds.first, idx, seeds.second, &state);
|
||||
float4 rand;
|
||||
|
||||
// Using vectorized data load/store approach when N % 4 == 0 since this is typical case for input shape size.
|
||||
using LoadT = aligned_vector<T, UNROLL>;
|
||||
for (CUDA_LONG id = idx * UNROLL; id < N; id += step_size) {
|
||||
rand = dist_func(&state);
|
||||
T r[UNROLL];
|
||||
|
||||
// actual computation
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < UNROLL; ii++) {
|
||||
r[ii] = static_cast<T>(transform_func((&rand.x)[ii], alpha, beta));
|
||||
}
|
||||
|
||||
// Vectorized writes for Y_data
|
||||
*(reinterpret_cast<LoadT*>(&Y_data[id])) = *reinterpret_cast<LoadT*>(&r[0]);
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename DistFuncT, typename TransformFuncT>
|
||||
void RandomKernelImpl(const cudaDeviceProp& prop, cudaStream_t stream, const int64_t N, const DistFuncT& dist_func,
|
||||
const TransformFuncT& transform_func, float alpha, float beta, PhiloxGenerator& generator,
|
||||
T* Y_data) {
|
||||
const int block_size = 256;
|
||||
const int blocks_per_sm = prop.maxThreadsPerMultiProcessor / block_size;
|
||||
const int grid_size =
|
||||
std::min(prop.multiProcessorCount * blocks_per_sm, static_cast<int>(CeilDiv(N, block_size * UNROLL)));
|
||||
|
||||
// Compute the number of random numbers generated by each thread, and increment philox generator offset by that
|
||||
// amount.
|
||||
const uint64_t counter_offset = static_cast<uint64_t>(((N - 1) / (block_size * grid_size * UNROLL) + 1) * UNROLL);
|
||||
auto seeds = generator.NextPhiloxSeeds(counter_offset);
|
||||
|
||||
if (N % UNROLL != 0) {
|
||||
RandomKernel<T><<<grid_size, block_size, 0, stream>>>(N, seeds, dist_func, transform_func, alpha, beta, Y_data);
|
||||
} else {
|
||||
RandomVectorizedKernel<T>
|
||||
<<<grid_size, block_size, 0, stream>>>(N, seeds, dist_func, transform_func, alpha, beta, Y_data);
|
||||
}
|
||||
}
|
||||
|
||||
#define RANDOM_KERNEL_IMPL(name) \
|
||||
template <typename T> \
|
||||
void name##KernelImpl(const cudaDeviceProp& prop, cudaStream_t stream, const int64_t N, const float alpha, \
|
||||
const float beta, PhiloxGenerator& generator, T* Y_data) { \
|
||||
RandomKernelImpl(prop, stream, N, DistFunc_##name(), TransformFunc_##name(), alpha, beta, generator, Y_data); \
|
||||
}
|
||||
|
||||
RANDOM_KERNEL_IMPL(RandomNormal)
|
||||
RANDOM_KERNEL_IMPL(RandomUniform)
|
||||
|
||||
#define SPECIALIZED_RANDOM_KERNEL(name, T) \
|
||||
template void name##KernelImpl(const cudaDeviceProp& prop, cudaStream_t stream, const int64_t N, const float alpha, \
|
||||
const float beta, PhiloxGenerator& generator, T* Y_data);
|
||||
|
||||
#define SPECIALIZED_RANDOM_KERNELS(T) \
|
||||
SPECIALIZED_RANDOM_KERNEL(RandomNormal, T) \
|
||||
SPECIALIZED_RANDOM_KERNEL(RandomUniform, T)
|
||||
|
||||
SPECIALIZED_RANDOM_KERNELS(float)
|
||||
SPECIALIZED_RANDOM_KERNELS(double)
|
||||
SPECIALIZED_RANDOM_KERNELS(half)
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace onnxruntime
|
||||
22
onnxruntime/core/providers/cuda/generator/random_impl.h
Normal file
22
onnxruntime/core/providers/cuda/generator/random_impl.h
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "core/framework/random_generator.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace cuda {
|
||||
|
||||
#define RANDOM_KERNEL_DECLARE(name) \
|
||||
template <typename T> \
|
||||
void name##KernelImpl(const cudaDeviceProp& prop, cudaStream_t stream, const int64_t N, const float alpha, \
|
||||
const float beta, PhiloxGenerator& generator, T* Y_data);
|
||||
|
||||
RANDOM_KERNEL_DECLARE(RandomNormal)
|
||||
RANDOM_KERNEL_DECLARE(RandomUniform)
|
||||
|
||||
#undef RANDOM_KERNEL_DECLARE
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -719,6 +719,10 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDom
|
|||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, Loop);
|
||||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, DepthToSpace);
|
||||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 12, SpaceToDepth);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomNormal);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomNormalLike);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomUniform);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomUniformLike);
|
||||
|
||||
// opset 10
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool);
|
||||
|
|
@ -1559,6 +1563,10 @@ static Status RegisterRocmKernels(KernelRegistry& kernel_registry) {
|
|||
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, Loop)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 10, DepthToSpace)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, 12, SpaceToDepth)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomNormal)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomNormalLike)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomUniform)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 1, RandomUniformLike)>,
|
||||
|
||||
// opset 10
|
||||
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kRocmExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool)>,
|
||||
|
|
|
|||
|
|
@ -33,7 +33,10 @@ TEST(Random, RandomNormal2DDouble) {
|
|||
[&generator, &distribution](double& value) { value = distribution(generator); });
|
||||
|
||||
test.AddOutput<double>("Y", dims, expected_output);
|
||||
test.Run();
|
||||
|
||||
// The expected_output is generated using std lib, which is used by CPU kernel only.
|
||||
// So we need to exclude other EPs here. Ditto for other places.
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider, kRocmExecutionProvider});
|
||||
}
|
||||
|
||||
void RunRandomNormalLike3DFloat(bool infer_dtype = false) {
|
||||
|
|
@ -68,7 +71,7 @@ void RunRandomNormalLike3DFloat(bool infer_dtype = false) {
|
|||
|
||||
test.AddOutput<float>("Y", dims, expected_output);
|
||||
|
||||
test.Run();
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider, kRocmExecutionProvider});
|
||||
}
|
||||
|
||||
TEST(Random, RandomNormalLike3DDouble) {
|
||||
|
|
@ -105,7 +108,7 @@ TEST(Random, RandomUniform1DFloat) {
|
|||
test.AddOutput<float>("Y", dims, expected_output);
|
||||
|
||||
// TensorRT does not support manual seed overrides and there will be result mismatch
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider, kRocmExecutionProvider, kTensorrtExecutionProvider});
|
||||
}
|
||||
|
||||
void RunRandomUniformLikeTest(bool infer_dtype = false) {
|
||||
|
|
@ -138,7 +141,7 @@ void RunRandomUniformLikeTest(bool infer_dtype = false) {
|
|||
test.AddOutput<double>("Y", dims, expected_output);
|
||||
|
||||
// TensorRT does not support seed parameter and there will be result mismatch
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider, kRocmExecutionProvider, kTensorrtExecutionProvider});
|
||||
}
|
||||
|
||||
TEST(Random, RandomUniformLike2DDouble) {
|
||||
|
|
@ -324,5 +327,205 @@ TEST(Random, MultinomialInvalidDtype) {
|
|||
|
||||
test.Run(OpTester::ExpectResult::kExpectFailure, "Output type must be int32 or int64");
|
||||
}
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_ROCM)
|
||||
// We cannot call CUDA lib from UT, so just do some simple verification on output tensor.
|
||||
void RunRandomNormalGpuTest(const std::vector<int64_t> dims, const float mean, const float scale, const float seed,
|
||||
TensorProto_DataType dtype, bool is_random_like, bool infer_dtype) {
|
||||
OpTester test(is_random_like ? "RandomNormalLike" : "RandomNormal");
|
||||
test.AddAttribute("mean", mean);
|
||||
test.AddAttribute("scale", scale);
|
||||
test.AddAttribute("seed", seed);
|
||||
if (!is_random_like) {
|
||||
test.AddAttribute<int64_t>("dtype", dtype);
|
||||
} else if (!infer_dtype) {
|
||||
// For RandomNormalLike, if not infer dtype, use float as target.
|
||||
test.AddAttribute<int64_t>("dtype", TensorProto_DataType::TensorProto_DataType_FLOAT);
|
||||
}
|
||||
size_t size = 1;
|
||||
for (size_t i = 0; i < dims.size(); ++i) {
|
||||
size *= static_cast<size_t>(dims[i]);
|
||||
}
|
||||
if (!is_random_like) {
|
||||
test.AddAttribute("shape", dims);
|
||||
} else {
|
||||
if (dtype == TensorProto_DataType::TensorProto_DataType_FLOAT) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
test.AddInput("X", dims, float_data);
|
||||
} else if (dtype == TensorProto_DataType::TensorProto_DataType_DOUBLE) {
|
||||
std::vector<double> double_data(size, 0.);
|
||||
test.AddInput("X", dims, double_data);
|
||||
} else if (dtype == TensorProto_DataType::TensorProto_DataType_FLOAT16) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
std::vector<MLFloat16> fp16_data(size);
|
||||
ConvertFloatToMLFloat16(float_data.data(), fp16_data.data(), static_cast<int>(size));
|
||||
test.AddInput("X", dims, fp16_data);
|
||||
}
|
||||
}
|
||||
|
||||
// We'll do our own output verification.
|
||||
TensorProto_DataType output_dtype =
|
||||
is_random_like && !infer_dtype ? TensorProto_DataType::TensorProto_DataType_FLOAT : dtype;
|
||||
if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
test.AddOutput("Y", dims, float_data);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_DOUBLE) {
|
||||
std::vector<double> double_data(size, 0.);
|
||||
test.AddOutput("Y", dims, double_data);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT16) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
std::vector<MLFloat16> fp16_data(size);
|
||||
ConvertFloatToMLFloat16(float_data.data(), fp16_data.data(), static_cast<int>(size));
|
||||
test.AddOutput("Y", dims, fp16_data);
|
||||
}
|
||||
|
||||
auto output_verifier = [&](const std::vector<OrtValue>& fetches, const std::string& provider_type) {
|
||||
// Only one output, and mean of output values are near attribute mean.
|
||||
ASSERT_EQ(fetches.size(), 1);
|
||||
const auto& output_tensor = FetchTensor(fetches[0]);
|
||||
if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT) {
|
||||
auto output_span = output_tensor.DataAsSpan<float>();
|
||||
float sum = std::accumulate(output_span.begin(), output_span.end(), 0.f);
|
||||
ASSERT_NEAR(sum / static_cast<float>(size), mean, 0.1f);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_DOUBLE) {
|
||||
auto output_span = output_tensor.DataAsSpan<double>();
|
||||
double sum = std::accumulate(output_span.begin(), output_span.end(), 0.);
|
||||
ASSERT_NEAR(sum / static_cast<double>(size), static_cast<double>(mean), 0.1);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT16) {
|
||||
auto output_span = output_tensor.DataAsSpan<MLFloat16>();
|
||||
float sum = 0.f;
|
||||
for (auto value : output_span) {
|
||||
sum += value.ToFloat();
|
||||
}
|
||||
ASSERT_NEAR(sum / static_cast<float>(size), mean, 0.1f);
|
||||
}
|
||||
};
|
||||
|
||||
test.SetCustomOutputVerifier(output_verifier);
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider, kTensorrtExecutionProvider});
|
||||
}
|
||||
|
||||
TEST(Random, RandomNormalGpu) {
|
||||
// We will call RandomVectorizedKernel if total_size % 4 == 0, so test two input sizes here.
|
||||
std::vector<int64_t> dims1{256, 256};
|
||||
RunRandomNormalGpuTest(dims1, 1.f, 10.f, 123.f, TensorProto_DataType::TensorProto_DataType_FLOAT, false, false);
|
||||
RunRandomNormalGpuTest(dims1, -1.f, 8.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, false, false);
|
||||
RunRandomNormalGpuTest(dims1, 0.f, 16.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, false, false);
|
||||
RunRandomNormalGpuTest(dims1, 1.f, 10.f, 123.f, TensorProto_DataType::TensorProto_DataType_FLOAT, true, true);
|
||||
RunRandomNormalGpuTest(dims1, -1.f, 8.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, true, true);
|
||||
RunRandomNormalGpuTest(dims1, 0.f, 16.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, true, true);
|
||||
RunRandomNormalGpuTest(dims1, -1.f, 8.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, true, false);
|
||||
RunRandomNormalGpuTest(dims1, 0.f, 16.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, true, false);
|
||||
std::vector<int64_t> dims2{255, 255};
|
||||
RunRandomNormalGpuTest(dims2, 1.f, 10.f, 123.f, TensorProto_DataType::TensorProto_DataType_FLOAT, false, false);
|
||||
RunRandomNormalGpuTest(dims2, -1.f, 8.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, true, true);
|
||||
RunRandomNormalGpuTest(dims2, 0.f, 16.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, true, false);
|
||||
}
|
||||
|
||||
void RunRandomUniformGpuTest(const std::vector<int64_t> dims, const float low, const float high, const float seed,
|
||||
TensorProto_DataType dtype, bool is_random_like, bool infer_dtype) {
|
||||
OpTester test(is_random_like ? "RandomUniformLike" : "RandomUniform");
|
||||
test.AddAttribute("low", low);
|
||||
test.AddAttribute("high", high);
|
||||
test.AddAttribute("seed", seed);
|
||||
if (!is_random_like) {
|
||||
test.AddAttribute<int64_t>("dtype", dtype);
|
||||
} else if (!infer_dtype) {
|
||||
// For RandomUniformLike, if not infer dtype, use float as target.
|
||||
test.AddAttribute<int64_t>("dtype", TensorProto_DataType::TensorProto_DataType_FLOAT);
|
||||
}
|
||||
size_t size = 1;
|
||||
for (size_t i = 0; i < dims.size(); ++i) {
|
||||
size *= static_cast<size_t>(dims[i]);
|
||||
}
|
||||
if (!is_random_like) {
|
||||
test.AddAttribute("shape", dims);
|
||||
} else {
|
||||
if (dtype == TensorProto_DataType::TensorProto_DataType_FLOAT) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
test.AddInput("X", dims, float_data);
|
||||
} else if (dtype == TensorProto_DataType::TensorProto_DataType_DOUBLE) {
|
||||
std::vector<double> double_data(size, 0.);
|
||||
test.AddInput("X", dims, double_data);
|
||||
} else if (dtype == TensorProto_DataType::TensorProto_DataType_FLOAT16) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
std::vector<MLFloat16> fp16_data(size);
|
||||
ConvertFloatToMLFloat16(float_data.data(), fp16_data.data(), static_cast<int>(size));
|
||||
test.AddInput("X", dims, fp16_data);
|
||||
}
|
||||
}
|
||||
|
||||
// We'll do our own output verification.
|
||||
TensorProto_DataType output_dtype =
|
||||
is_random_like && !infer_dtype ? TensorProto_DataType::TensorProto_DataType_FLOAT : dtype;
|
||||
if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
test.AddOutput("Y", dims, float_data);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_DOUBLE) {
|
||||
std::vector<double> double_data(size, 0.);
|
||||
test.AddOutput("Y", dims, double_data);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT16) {
|
||||
std::vector<float> float_data(size, 0.f);
|
||||
std::vector<MLFloat16> fp16_data(size);
|
||||
ConvertFloatToMLFloat16(float_data.data(), fp16_data.data(), static_cast<int>(size));
|
||||
test.AddOutput("Y", dims, fp16_data);
|
||||
}
|
||||
|
||||
auto output_verifier = [&](const std::vector<OrtValue>& fetches, const std::string& provider_type) {
|
||||
// Only one output. Each value in output tensoer is between low and high.
|
||||
// Mean of output values are near attribute mean of low and high.
|
||||
ASSERT_EQ(fetches.size(), 1);
|
||||
const auto& output_tensor = FetchTensor(fetches[0]);
|
||||
if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT) {
|
||||
auto output_span = output_tensor.DataAsSpan<float>();
|
||||
for (auto value : output_span) {
|
||||
ASSERT_GE(value, low);
|
||||
ASSERT_LE(value, high);
|
||||
}
|
||||
float sum = std::accumulate(output_span.begin(), output_span.end(), 0.f);
|
||||
ASSERT_NEAR(sum / static_cast<float>(size), (high + low) / 2.f, 0.1f);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_DOUBLE) {
|
||||
auto output_span = output_tensor.DataAsSpan<double>();
|
||||
for (auto value : output_span) {
|
||||
ASSERT_GE(value, static_cast<double>(low));
|
||||
ASSERT_LE(value, static_cast<double>(high));
|
||||
}
|
||||
double sum = std::accumulate(output_span.begin(), output_span.end(), 0.);
|
||||
ASSERT_NEAR(sum / static_cast<double>(size), static_cast<double>((high + low) / 2.f), 0.1);
|
||||
} else if (output_dtype == TensorProto_DataType::TensorProto_DataType_FLOAT16) {
|
||||
auto output_span = output_tensor.DataAsSpan<MLFloat16>();
|
||||
float sum = 0.f;
|
||||
for (auto value : output_span) {
|
||||
float f = value.ToFloat();
|
||||
ASSERT_GE(f, low);
|
||||
ASSERT_LE(f, high);
|
||||
sum += f;
|
||||
}
|
||||
ASSERT_NEAR(sum / static_cast<float>(size), (high + low) / 2.f, 0.1f);
|
||||
}
|
||||
};
|
||||
|
||||
test.SetCustomOutputVerifier(output_verifier);
|
||||
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider, kTensorrtExecutionProvider});
|
||||
}
|
||||
|
||||
TEST(Random, RandomUniformGpu) {
|
||||
// We will call RandomVectorizedKernel if total_size % 4 == 0, so test two input sizes here.
|
||||
std::vector<int64_t> dims1{256, 256};
|
||||
RunRandomUniformGpuTest(dims1, 0.f, 10.f, 123.f, TensorProto_DataType::TensorProto_DataType_FLOAT, false, false);
|
||||
RunRandomUniformGpuTest(dims1, -10.f, 0.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, false, false);
|
||||
RunRandomUniformGpuTest(dims1, -5.f, 5.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, false, false);
|
||||
RunRandomUniformGpuTest(dims1, 0.f, 10.f, 123.f, TensorProto_DataType::TensorProto_DataType_FLOAT, true, true);
|
||||
RunRandomUniformGpuTest(dims1, -10.f, 0.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, true, true);
|
||||
RunRandomUniformGpuTest(dims1, -5.f, 5.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, true, true);
|
||||
RunRandomUniformGpuTest(dims1, -10.f, 0.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, true, false);
|
||||
RunRandomUniformGpuTest(dims1, -5.f, 5.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, true, false);
|
||||
std::vector<int64_t> dims2{255, 255};
|
||||
RunRandomUniformGpuTest(dims2, 0.f, 10.f, 123.f, TensorProto_DataType::TensorProto_DataType_FLOAT, false, false);
|
||||
RunRandomUniformGpuTest(dims2, -10.f, 0.f, 231.f, TensorProto_DataType::TensorProto_DataType_DOUBLE, true, true);
|
||||
RunRandomUniformGpuTest(dims2, -5.f, 5.f, 312.f, TensorProto_DataType::TensorProto_DataType_FLOAT16, true, false);
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
Loading…
Reference in a new issue