mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-08 17:17:15 +00:00
Re-work hierarchy, fix virtual method overload/hiding (#10160)
Re-work hierarchy, fix virtual method overload/hiding Use std::optional with a clear comment on the member thread-safety.
This commit is contained in:
parent
d5742f3a43
commit
28ce2a5a78
2 changed files with 102 additions and 70 deletions
|
|
@ -2,6 +2,7 @@
|
|||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/providers/cuda/generator/random.h"
|
||||
#include "core/providers/cuda/generator/random_impl.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace cuda {
|
||||
|
|
@ -28,50 +29,80 @@ ONNX_OPERATOR_KERNEL_EX(RandomUniformLike, kOnnxDomain, 1, kCudaExecutionProvide
|
|||
.TypeConstraint("T2", DataTypeImpl::AllIEEEFloatTensorTypes()),
|
||||
RandomUniformLike);
|
||||
|
||||
Status RandomNormalBase::Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const {
|
||||
Tensor& Y = *p_ctx->Output(0, shape);
|
||||
#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
|
||||
|
||||
Status RandomNormalBase::ComputeNormal(const CudaKernel& cuda_kernel, OpKernelContext& ctx, const TensorShape& shape, int dtype) const {
|
||||
Tensor& Y = *ctx.Output(0, shape);
|
||||
const int64_t N = shape.Size();
|
||||
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
|
||||
PhiloxGenerator& generator = GetPhiloxGenerator();
|
||||
utils::MLTypeCallDispatcher<float, MLFloat16, double> t_disp(dtype);
|
||||
t_disp.Invoke<RandomNormalComputeImpl>(GetDeviceProp(), Stream(), N, scale_, mean_, generator, Y);
|
||||
t_disp.Invoke<RandomNormalComputeImpl>(cuda_kernel.GetDeviceProp(), cuda_kernel.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>() &&
|
||||
|
||||
if (!p_X) {
|
||||
return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
|
||||
}
|
||||
|
||||
int dtype = GetDType();
|
||||
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());
|
||||
|
||||
if (dtype == TensorProto_DataType_UNDEFINED)
|
||||
dtype = p_X->GetElementType();
|
||||
|
||||
return ComputeNormal(*this, *p_ctx, p_X->Shape(), dtype);
|
||||
}
|
||||
|
||||
Status RandomUniformBase::Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const {
|
||||
Tensor& Y = *p_ctx->Output(0, shape);
|
||||
Status RandomUniformBase::ComputeUniform(const CudaKernel& cuda_kernel, OpKernelContext& ctx, const TensorShape& shape, int dtype) const {
|
||||
Tensor& Y = *ctx.Output(0, shape);
|
||||
const int64_t N = shape.Size();
|
||||
PhiloxGenerator& generator = generator_ ? *generator_ : PhiloxGenerator::Default();
|
||||
PhiloxGenerator& generator = GetPhiloxGenerator();
|
||||
utils::MLTypeCallDispatcher<float, MLFloat16, double> t_disp(dtype);
|
||||
t_disp.Invoke<RandomUniformComputeImpl>(GetDeviceProp(), Stream(), N, range_, from_, generator, Y);
|
||||
t_disp.Invoke<RandomUniformComputeImpl>(cuda_kernel.GetDeviceProp(), cuda_kernel.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>() &&
|
||||
|
||||
if (!p_X) {
|
||||
return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available.");
|
||||
}
|
||||
|
||||
int dtype = GetDType();
|
||||
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());
|
||||
|
||||
if (dtype == TensorProto_DataType_UNDEFINED)
|
||||
dtype = p_X->GetElementType();
|
||||
|
||||
return ComputeUniform(*this, *p_ctx, p_X->Shape(), dtype);
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
|
|
|
|||
|
|
@ -3,126 +3,127 @@
|
|||
|
||||
#pragma once
|
||||
|
||||
#include "core/framework/random_generator.h"
|
||||
#include "core/providers/cuda/cuda_kernel.h"
|
||||
|
||||
#include "core/providers/cuda/generator/random_impl.h"
|
||||
#include <optional>
|
||||
|
||||
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 {
|
||||
class RandomBase {
|
||||
protected:
|
||||
RandomBase(const OpKernelInfo& info) : CudaKernel(info) {
|
||||
explicit RandomBase(const OpKernelInfo& info) {
|
||||
float seed = 0.f;
|
||||
if (info.GetAttr<float>("seed", &seed).IsOK()) {
|
||||
generator_ = std::make_unique<PhiloxGenerator>(static_cast<uint64_t>(seed));
|
||||
generator_.emplace(static_cast<uint64_t>(seed));
|
||||
}
|
||||
|
||||
int64_t dtype;
|
||||
if (info.GetAttr<int64_t>("dtype", &dtype).IsOK()) {
|
||||
ORT_ENFORCE(ONNX_NAMESPACE::TensorProto::DataType_IsValid(gsl::narrow<int>(dtype)) &&
|
||||
dtype != ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED,
|
||||
"Invalid dtype of ", dtype);
|
||||
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_;
|
||||
|
||||
void SetDTypeIfUndefined(ONNX_NAMESPACE::TensorProto::DataType dtype) noexcept {
|
||||
if (dtype_ == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED) {
|
||||
dtype_ = dtype;
|
||||
}
|
||||
}
|
||||
|
||||
ONNX_NAMESPACE::TensorProto::DataType GetDType() const noexcept { return dtype_; }
|
||||
|
||||
PhiloxGenerator& GetPhiloxGenerator() const {
|
||||
return (generator_.has_value()) ? *generator_ : PhiloxGenerator::Default();
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
ONNX_NAMESPACE::TensorProto::DataType dtype_ =
|
||||
ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED; // optional and may be inferred
|
||||
|
||||
// This member is thread-safe, ensuring proper synchronization
|
||||
mutable std::optional<PhiloxGenerator> generator_;
|
||||
};
|
||||
|
||||
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());
|
||||
ORT_THROW_IF_ERROR(info.GetAttr<float>("scale", &scale_));
|
||||
ORT_THROW_IF_ERROR(info.GetAttr<float>("mean", &mean_));
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const;
|
||||
Status ComputeNormal(const CudaKernel& cuda_kernel, OpKernelContext& ctx, const TensorShape& shape, int dtype) const;
|
||||
|
||||
protected:
|
||||
private:
|
||||
float scale_;
|
||||
float mean_;
|
||||
};
|
||||
|
||||
class RandomNormal final : public RandomNormalBase {
|
||||
class RandomNormal final : public CudaKernel, protected RandomNormalBase {
|
||||
public:
|
||||
explicit RandomNormal(const OpKernelInfo& info) : RandomNormalBase(info) {
|
||||
if (dtype_ == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED) {
|
||||
dtype_ = ONNX_NAMESPACE::TensorProto_DataType_FLOAT;
|
||||
}
|
||||
explicit RandomNormal(const OpKernelInfo& info) : CudaKernel(info), RandomNormalBase(info) {
|
||||
SetDTypeIfUndefined(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
|
||||
std::vector<int64_t> shape;
|
||||
ORT_ENFORCE(info.GetAttrs<int64_t>("shape", shape).IsOK());
|
||||
ORT_THROW_IF_ERROR(info.GetAttrs<int64_t>("shape", shape));
|
||||
shape_ = TensorShape(shape);
|
||||
}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override {
|
||||
return ComputeNormal(*this, *p_ctx, shape_, GetDType());
|
||||
}
|
||||
|
||||
private:
|
||||
TensorShape shape_;
|
||||
};
|
||||
|
||||
class RandomNormalLike final : public RandomNormalBase {
|
||||
class RandomNormalLike final : public CudaKernel, protected RandomNormalBase {
|
||||
public:
|
||||
explicit RandomNormalLike(const OpKernelInfo& info) : RandomNormalBase(info) {}
|
||||
explicit RandomNormalLike(const OpKernelInfo& info) : CudaKernel(info), RandomNormalBase(info) {}
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
};
|
||||
|
||||
class RandomUniformBase : public RandomBase {
|
||||
protected:
|
||||
RandomUniformBase(const OpKernelInfo& info) : RandomBase(info) {
|
||||
explicit 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());
|
||||
ORT_THROW_IF_ERROR(info.GetAttr<float>("low", &low));
|
||||
ORT_THROW_IF_ERROR(info.GetAttr<float>("high", &high));
|
||||
from_ = low;
|
||||
range_ = high - low;
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* p_ctx, const TensorShape& shape, int dtype) const;
|
||||
Status ComputeUniform(const CudaKernel& cuda_kernel, OpKernelContext& ctx, const TensorShape& shape, int dtype) const;
|
||||
|
||||
protected:
|
||||
private:
|
||||
float range_;
|
||||
float from_;
|
||||
};
|
||||
|
||||
class RandomUniform final : public RandomUniformBase {
|
||||
class RandomUniform final : public CudaKernel, protected RandomUniformBase {
|
||||
public:
|
||||
explicit RandomUniform(const OpKernelInfo& info) : RandomUniformBase(info) {
|
||||
if (dtype_ == ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED) {
|
||||
dtype_ = ONNX_NAMESPACE::TensorProto_DataType_FLOAT;
|
||||
}
|
||||
explicit RandomUniform(const OpKernelInfo& info) : CudaKernel(info), RandomUniformBase(info) {
|
||||
SetDTypeIfUndefined(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
|
||||
std::vector<int64_t> shape;
|
||||
ORT_ENFORCE(info.GetAttrs<int64_t>("shape", shape).IsOK());
|
||||
ORT_THROW_IF_ERROR(info.GetAttrs<int64_t>("shape", shape));
|
||||
shape_ = TensorShape(shape);
|
||||
}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override {
|
||||
return ComputeUniform(*this, *p_ctx, shape_, GetDType());
|
||||
}
|
||||
|
||||
private:
|
||||
TensorShape shape_;
|
||||
};
|
||||
|
||||
class RandomUniformLike final : public RandomUniformBase {
|
||||
class RandomUniformLike final : public CudaKernel, protected RandomUniformBase {
|
||||
public:
|
||||
explicit RandomUniformLike(const OpKernelInfo& info) : RandomUniformBase(info) {}
|
||||
explicit RandomUniformLike(const OpKernelInfo& info) : CudaKernel(info), RandomUniformBase(info) {}
|
||||
Status ComputeInternal(OpKernelContext* p_ctx) const override;
|
||||
};
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue