From 28ce2a5a785804d6244d78274c6fbbda06f48b57 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Mon, 3 Jan 2022 10:24:49 -0800 Subject: [PATCH] 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. --- .../core/providers/cuda/generator/random.cc | 67 ++++++++--- .../core/providers/cuda/generator/random.h | 105 +++++++++--------- 2 files changed, 102 insertions(+), 70 deletions(-) diff --git a/onnxruntime/core/providers/cuda/generator/random.cc b/onnxruntime/core/providers/cuda/generator/random.cc index 643da1579d..e5560b5acc 100644 --- a/onnxruntime/core/providers/cuda/generator/random.cc +++ b/onnxruntime/core/providers/cuda/generator/random.cc @@ -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 \ + 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::MappedType CudaT; \ + CudaT* Y_data = reinterpret_cast(Y.template MutableData()); \ + name##KernelImpl(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 t_disp(dtype); - t_disp.Invoke(GetDeviceProp(), Stream(), N, scale_, mean_, generator, Y); + t_disp.Invoke(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(0); - if (!p_X) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available."); - if (dtype_ == TensorProto_DataType_UNDEFINED && !p_X->IsDataType() && !p_X->IsDataType() && + + 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() && !p_X->IsDataType() && !p_X->IsDataType()) { 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 t_disp(dtype); - t_disp.Invoke(GetDeviceProp(), Stream(), N, range_, from_, generator, Y); + t_disp.Invoke(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(0); - if (!p_X) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available."); - if (dtype_ == TensorProto_DataType_UNDEFINED && !p_X->IsDataType() && !p_X->IsDataType() && + + 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() && !p_X->IsDataType() && !p_X->IsDataType()) { 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 diff --git a/onnxruntime/core/providers/cuda/generator/random.h b/onnxruntime/core/providers/cuda/generator/random.h index f1b58b1636..3b7fcb58ca 100644 --- a/onnxruntime/core/providers/cuda/generator/random.h +++ b/onnxruntime/core/providers/cuda/generator/random.h @@ -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 namespace onnxruntime { namespace cuda { -#define RANDOM_COMPUTE_IMPL(name) \ - template \ - 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::MappedType CudaT; \ - CudaT* Y_data = reinterpret_cast(Y.template MutableData()); \ - name##KernelImpl(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("seed", &seed).IsOK()) { - generator_ = std::make_unique(static_cast(seed)); + generator_.emplace(static_cast(seed)); } int64_t dtype; if (info.GetAttr("dtype", &dtype).IsOK()) { + ORT_ENFORCE(ONNX_NAMESPACE::TensorProto::DataType_IsValid(gsl::narrow(dtype)) && + dtype != ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED, + "Invalid dtype of ", dtype); dtype_ = static_cast(dtype); - ORT_ENFORCE(ONNX_NAMESPACE::TensorProto::DataType_IsValid(dtype_) && - dtype_ != ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED, - "Invalid dtype of ", dtype_); } } protected: - std::unique_ptr 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 generator_; }; class RandomNormalBase : public RandomBase { protected: RandomNormalBase(const OpKernelInfo& info) : RandomBase(info) { - ORT_ENFORCE(info.GetAttr("scale", &scale_).IsOK()); - ORT_ENFORCE(info.GetAttr("mean", &mean_).IsOK()); + ORT_THROW_IF_ERROR(info.GetAttr("scale", &scale_)); + ORT_THROW_IF_ERROR(info.GetAttr("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 shape; - ORT_ENFORCE(info.GetAttrs("shape", shape).IsOK()); + ORT_THROW_IF_ERROR(info.GetAttrs("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("low", &low).IsOK()); - ORT_ENFORCE(info.GetAttr("high", &high).IsOK()); + ORT_THROW_IF_ERROR(info.GetAttr("low", &low)); + ORT_THROW_IF_ERROR(info.GetAttr("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 shape; - ORT_ENFORCE(info.GetAttrs("shape", shape).IsOK()); + ORT_THROW_IF_ERROR(info.GetAttrs("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; };