Add CUDA Kernels of RandomNormal[Like], RandomUniform[Like] (#9761)

This commit is contained in:
Vincent Wang 2021-11-19 08:18:34 +08:00 committed by GitHub
parent 6284cbe833
commit f390347c11
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 602 additions and 4 deletions

View file

@ -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)|

View file

@ -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)>,

View 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

View 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

View 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

View 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

View file

@ -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)>,

View file

@ -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