diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index f7cb7c048fa66..65e9680d017dc 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -437,6 +437,15 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, uint8_t, ReduceMin); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, GatherND); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 12, Einsum); +// REVIEW(codemzs): ConstEigenVectorArrayMap.cast, BuildKernelCreateInfo, BuildKernelCreateInfo, + // REVIEW(codemzs): ConstEigenVectorArrayMap.cast, + //BuildKernelCreateInfo, + //BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/cpu/nn/dropout_op.cc b/onnxruntime/core/providers/cpu/nn/dropout_op.cc new file mode 100644 index 0000000000000..224638afa4313 --- /dev/null +++ b/onnxruntime/core/providers/cpu/nn/dropout_op.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cpu/nn/dropout_op.h" + +namespace onnxruntime { + +// Dropout +#define REGISTER_KERNEL_TYPED(OpName, VER, T1, T2, Trainable) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + OpName, \ + kOnnxDomain, \ + VER, \ + T1##_##T2, \ + kCpuExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ + Dropout); + +// REVIEW(mzs): ConstEigenVectorArrayMap.cast +#include +#include "core/util/math_cpuonly.h" + +namespace onnxruntime { + +template +class Dropout final: public OpKernel { + public: + Dropout(const OpKernelInfo& info) : OpKernel{info} { + int64_t seed = 0; + if (info.GetAttr("seed", &seed).IsOK()) { + generator_ = onnxruntime::make_unique(seed); + } + } + + Status Compute(OpKernelContext* context) const override; + + private: + mutable std::unique_ptr generator_; +}; + +namespace { +constexpr float k_default_ratio{0.5f}; + +template +float GetRatioOrDefault(const Tensor* ratio_tensor) { + if (ratio_tensor) { + ORT_ENFORCE(ratio_tensor->Shape().Size() == 1, "ratio input should have a single value."); +#ifdef _WIN32 +#pragma warning(disable : 4244) +#endif + const float ratio_value = *ratio_tensor->Data(); + ORT_ENFORCE(0.0f <= ratio_value && ratio_value < 1.0f, "ratio must be in the range [0, 1)"); + return ratio_value; + } + return k_default_ratio; +} +} // namespace + +template +Status Dropout::Compute(OpKernelContext* context) const { + const Tensor* X = context->Input(0); + auto X_span = X->DataAsSpan(); + const Tensor* ratio = context->Input(1); // optional + const float ratio_value = GetRatioOrDefault(ratio); + const auto& X_shape = X->Shape(); + Tensor* Y = context->Output(0, X_shape); + auto Y_span = Y->MutableDataAsSpan(); + Tensor* mask = context->Output(1, X_shape); // optional + std::unique_ptr temp_mask_buffer{}; // temporary buffer to use if mask input is not provided + auto mask_span = [&X_shape, mask, &temp_mask_buffer]() { + if (mask) return mask->MutableDataAsSpan(); + temp_mask_buffer = onnxruntime::make_unique(X_shape.Size()); + return gsl::make_span(temp_mask_buffer.get(), X_shape.Size()); + }(); + + ORT_ENFORCE(!mask || mask->Shape() == X_shape, "X and mask should have the same shape"); + + const Tensor* training_mode = context->Input(2); + if ((0 == ratio_value /*Backward compat with TrainableDropout*/) || + !trainable_dropout && (training_mode == nullptr || *(training_mode->Data()) == false)) { + // drop none + if (X_span.data() != Y_span.data()) { + std::copy(X_span.begin(), X_span.end(), Y_span.begin()); + } + + if (mask != nullptr) { + std::fill(mask_span.begin(), mask_span.end(), true); + } + + } else { + // drop some + ConstEigenVectorArrayMap X_arr(X_span.data(), X_span.size()); + EigenVectorArrayMap Y_arr(Y_span.data(), Y_span.size()); + EigenVectorArrayMap mask_arr(mask_span.data(), mask_span.size()); + + // generate mask + { + RandomGenerator& generator = generator_ != nullptr ? *generator_.get() : RandomGenerator::Default(); + std::default_random_engine rng(generator.NextSeed()); + std::uniform_real_distribution dist{0.0f, 1.0f}; + mask_arr = Eigen::ArrayX::NullaryExpr( + mask_arr.size(), + [ratio_value, &dist, &rng]() { return dist(rng) >= ratio_value; }); + } + + Y_arr = mask_arr.cast() * X_arr / (1.0f - ratio_value); + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index f710344f5a2d6..0ae686d920284 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -773,6 +773,16 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, int64_t, GatherND); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_MLFloat16, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_float, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_double, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_MLFloat16, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_float, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_double, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_MLFloat16, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_float, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_double, Dropout); + static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { BuildKernelCreateInfo, @@ -1289,6 +1299,16 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/cuda/nn/dropout.cc b/onnxruntime/core/providers/cuda/nn/dropout.cc new file mode 100644 index 0000000000000..199b987c81660 --- /dev/null +++ b/onnxruntime/core/providers/cuda/nn/dropout.cc @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/nn/dropout.h" + +namespace onnxruntime { +namespace cuda { + +#define REGISTER_KERNEL_TYPED(T1, T2) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + Dropout, \ + kOnnxDomain, \ + 12, \ + T1##_##T2, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(1) \ + .InputMemoryType(2), \ + Dropout); + +REGISTER_KERNEL_TYPED(MLFloat16, MLFloat16) +REGISTER_KERNEL_TYPED(MLFloat16, float) +REGISTER_KERNEL_TYPED(MLFloat16, double) +REGISTER_KERNEL_TYPED(float, MLFloat16) +REGISTER_KERNEL_TYPED(float, float) +REGISTER_KERNEL_TYPED(float, double) +REGISTER_KERNEL_TYPED(double, MLFloat16) +REGISTER_KERNEL_TYPED(double, float) +REGISTER_KERNEL_TYPED(double, double) + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/dropout.h b/onnxruntime/core/providers/cuda/nn/dropout.h new file mode 100644 index 0000000000000..a9fc19573848a --- /dev/null +++ b/onnxruntime/core/providers/cuda/nn/dropout.h @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/nn/dropout_impl.h" +#include "core/providers/cuda/nn/dropout.h" +#include "core/providers/common.h" +#include "core/framework/random_seed.h" + +namespace onnxruntime { +namespace cuda { + +template +class Dropout final : public CudaKernel { + public: + Dropout(const OpKernelInfo& info) : CudaKernel(info), default_ratio_(0.5) { + int64_t seed = 0; + if (info.GetAttr("seed", &seed).IsOK()) { + generator_ = onnxruntime::make_unique(static_cast(seed)); + } + } + + Status ComputeInternal(OpKernelContext* context) const override; + + private: + mutable std::unique_ptr generator_; + const float default_ratio_; +}; + +template +Status Dropout::ComputeInternal(OpKernelContext* context) const { + typedef typename ToCudaType::MappedType CudaT; + + //Get X_data + const Tensor* X = context->Input(0); + if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available."); + const TensorShape& shape = X->Shape(); + auto X_data = reinterpret_cast(X->template Data()); + const int64_t N = shape.Size(); + + //Get Y_data + auto Y = context->Output(0, shape); + auto Y_data = reinterpret_cast(Y->template MutableData()); + + //Get mask_data + auto mask = context->Output(1, shape); + ORT_ENFORCE(!mask || mask->Shape().Size() == N); + + //Get the ratio_data + float ratio_data; + auto ratio = context->Input(1); + + static_assert(std::is_same::value || std::is_same::value || std::is_same::value, + "T2 must be float16 or float or double"); + + if (ratio) { + ratio_data = static_cast(*(ratio->template Data())); + } else { + ratio_data = default_ratio_; + } + ORT_ENFORCE(ratio_data >= 0.0f && ratio_data < 1.0f); + + const Tensor* training_mode = context->Input(2); + //Check for inference mode. + if ((0 == ratio_data /*Backward compat with TrainableDropout*/) || + (!trainable_dropout && (training_mode == nullptr || *(training_mode->Data()) == false))) { + if (Y_data != X_data) { + CUDA_CALL_THROW(cudaMemcpyAsync(Y_data, X_data, N * sizeof(T1), cudaMemcpyDeviceToDevice)); + } + + // If mask is requested, return all 1s. + if (mask != nullptr) { + ORT_ENFORCE(cudaMemset(mask->MutableData(), true, N * sizeof(bool)) == cudaSuccess); + } + + return Status::OK(); + } + + IAllocatorUniquePtr temp_mask_buffer{}; // buffer to use if mask is not provided + bool* const mask_data = [this, N, mask, &temp_mask_buffer]() { + if (mask) return mask->MutableData(); + temp_mask_buffer = GetScratchBuffer(N); + return temp_mask_buffer.get(); + }(); + + PhiloxGenerator& generator = generator_ != nullptr ? *generator_.get() : PhiloxGenerator::Default(); + DropoutKernelImpl(GetDeviceProp(), N, ratio_data, generator, X_data, Y_data, mask_data); + + return Status::OK(); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/dropout_impl.cu b/onnxruntime/core/providers/cuda/nn/dropout_impl.cu new file mode 100644 index 0000000000000..d0ebc21851be5 --- /dev/null +++ b/onnxruntime/core/providers/cuda/nn/dropout_impl.cu @@ -0,0 +1,103 @@ +/** +* Copyright (c) 2016-present, Facebook, Inc. +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +/* Modifications Copyright (c) Microsoft. */ + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/nn/dropout_impl.h" +#include +#include + +namespace onnxruntime { +namespace cuda { + +constexpr int UNROLL = 4; + +template +__global__ void DropoutKernel( + const int64_t N, + const float ratio, + const std::pair seeds, + const T* X_data, + T* Y_data, + bool* mask_data) { + const float p = 1.0f - ratio; + const T scale = T(1.0f / p); + + CUDA_LONG idx = blockDim.x * blockIdx.x + threadIdx.x; + CUDA_LONG step_size = gridDim.x * blockDim.x * UNROLL; + CUDA_LONG rounded_size = ((N - 1) / step_size + 1) * step_size; + + curandStatePhilox4_32_10_t state; + curand_init(seeds.first, idx, seeds.second, &state); + + // 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; id < rounded_size; id += step_size) { + float4 rand = curand_uniform4(&state); + + for (CUDA_LONG i = 0; i < UNROLL; i++) { + CUDA_LONG li = id + gridDim.x * blockDim.x * i; + if (li < N) { + mask_data[li] = (&rand.x)[i] < p; + Y_data[li] = X_data[li] * T(mask_data[li]) * scale; + } + } + + __syncthreads(); + } +} + +template +void DropoutKernelImpl( + const cudaDeviceProp& prop, + const int64_t N, + const float ratio, + PhiloxGenerator& generator, + const T* X_data, + T* Y_data, + bool* mask_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(CeilDiv(N, block_size))); + + // 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(((N - 1) / (block_size * grid_size * UNROLL) + 1) * UNROLL); + auto seeds = generator.NextPhiloxSeeds(counter_offset); + + DropoutKernel<<>>(N, ratio, seeds, X_data, Y_data, mask_data); +} + +#define SPECIALIZED_DROPOUT_IMPL(T) \ + template void DropoutKernelImpl( \ + const cudaDeviceProp& prop, \ + const int64_t N, \ + const float ratio, \ + PhiloxGenerator& generator, \ + const T* X_data, \ + T* Y_data, \ + bool* mask_data); + +SPECIALIZED_DROPOUT_IMPL(float) +SPECIALIZED_DROPOUT_IMPL(double) +SPECIALIZED_DROPOUT_IMPL(half) + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/dropout_impl.h b/onnxruntime/core/providers/cuda/nn/dropout_impl.h new file mode 100644 index 0000000000000..5c52af13184a7 --- /dev/null +++ b/onnxruntime/core/providers/cuda/nn/dropout_impl.h @@ -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 { + +template +void DropoutKernelImpl( + const cudaDeviceProp& prop, + const int64_t N, + const float ratio, + PhiloxGenerator& generator, + const T* X_data, + T* Y_data, + bool* mask_data); + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc b/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc index 1348f370e0b45..4d6eef7b52dcd 100644 --- a/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc +++ b/orttraining/orttraining/test/training_ops/cpu/nn/dropout_op_test.cc @@ -129,10 +129,6 @@ TEST(DropoutTest, EmptyRatio) { RunDropoutTest("Dropout", true, {1000}); } -TEST(DropoutTest, Float16Ratio) { - RunDropoutTest("Dropout", true, {1000}, 0.0f, true, true); -} - TEST(TrainableDropoutTest, Basic) { RunDropoutTest("TrainableDropout", false, {10, 10, 10}, 0.75); } @@ -149,10 +145,6 @@ TEST(TrainableDropoutTest, EmptyRatio) { RunDropoutTest("TrainableDropout", true, {1000}, -1); } -TEST(TrainableDropoutTest, Float16Ratio) { - RunDropoutTest("TrainableDropout", true, {1000}, 0.0f, true, true); -} - namespace { void RunDropoutGradTest(const char* op, float ratio, const std::vector& input_dims, bool default_ratio = true) { const auto input_shape = TensorShape(input_dims); diff --git a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc index 69fef4b6c81d7..f24b2b28764df 100644 --- a/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc @@ -56,17 +56,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_MLFloat16, TrainableDropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_float, TrainableDropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double_double, TrainableDropoutGrad); -// REVIEW(mzs): ConstEigenVectorArrayMap.cast, // REVIEW(mzs): ConstEigenVectorArrayMap.cast, - //BuildKernelCreateInfo, - //BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - // REVIEW(mzs): ConstEigenVectorArrayMap.cast, //BuildKernelCreateInfo, //BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cpu/nn/dropout_op.cc b/orttraining/orttraining/training_ops/cpu/nn/dropout_op.cc index dd031977ac8f2..0ad839a8192f4 100644 --- a/orttraining/orttraining/training_ops/cpu/nn/dropout_op.cc +++ b/orttraining/orttraining/training_ops/cpu/nn/dropout_op.cc @@ -2,13 +2,13 @@ // Licensed under the MIT License. #include "orttraining/training_ops/cpu/nn/dropout_op.h" +#include "core/providers/cpu/nn/dropout_op.h" #include #include #include "core/util/math_cpuonly.h" namespace onnxruntime { namespace contrib { - namespace { constexpr float k_default_ratio{0.5f}; @@ -39,7 +39,7 @@ float GetRatioOrDefault(const Tensor* ratio_tensor) { .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ - Dropout); + onnxruntime::Dropout); // Temporary for backward compatibility, will eventually get rid of TrainableDropout when PyTorch exporter will move to // opset-12. @@ -50,72 +50,6 @@ REGISTER_KERNEL_TYPED(TrainableDropout, 9, double, MLFloat16, true) REGISTER_KERNEL_TYPED(TrainableDropout, 9, double, float, true) REGISTER_KERNEL_TYPED(TrainableDropout, 9, double, double, true) -// REVIEW(mzs): ConstEigenVectorArrayMap.cast -Status Dropout::Compute(OpKernelContext* context) const { - const Tensor* X = context->Input(0); - auto X_span = X->DataAsSpan(); - const Tensor* ratio = context->Input(1); // optional - const float ratio_value = GetRatioOrDefault(ratio); - const auto& X_shape = X->Shape(); - Tensor* Y = context->Output(0, X_shape); - auto Y_span = Y->MutableDataAsSpan(); - Tensor* mask = context->Output(1, X_shape); // optional - std::unique_ptr temp_mask_buffer{}; // temporary buffer to use if mask input is not provided - auto mask_span = [&X_shape, mask, &temp_mask_buffer]() { - if (mask) return mask->MutableDataAsSpan(); - temp_mask_buffer = onnxruntime::make_unique(X_shape.Size()); - return gsl::make_span(temp_mask_buffer.get(), X_shape.Size()); - }(); - - ORT_ENFORCE(!mask || mask->Shape() == X_shape, "X and mask should have the same shape"); - - const Tensor* training_mode = context->Input(2); - if ((0 == ratio_value /*Backward compat with TrainableDropout*/) || - !trainable_dropout && (training_mode == nullptr || *(training_mode->Data()) == false)) { - // drop none - if (X_span.data() != Y_span.data()) { - std::copy(X_span.begin(), X_span.end(), Y_span.begin()); - } - - if (mask != nullptr) { - std::fill(mask_span.begin(), mask_span.end(), true); - } - - } else { - // drop some - ConstEigenVectorArrayMap X_arr(X_span.data(), X_span.size()); - EigenVectorArrayMap Y_arr(Y_span.data(), Y_span.size()); - EigenVectorArrayMap mask_arr(mask_span.data(), mask_span.size()); - - // generate mask - { - RandomGenerator& generator = generator_ != nullptr ? *generator_.get() : RandomGenerator::Default(); - std::default_random_engine rng(generator.NextSeed()); - std::uniform_real_distribution dist{0.0f, 1.0f}; - mask_arr = Eigen::ArrayX::NullaryExpr( - mask_arr.size(), - [ratio_value, &dist, &rng]() { return dist(rng) >= ratio_value; }); - } - - Y_arr = mask_arr.cast() * X_arr / (1.0f - ratio_value); - } - - return Status::OK(); -} - #define REGISTER_GRADIENT_KERNEL_TYPED(OpName, T1, T2) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ OpName, \ @@ -181,6 +115,5 @@ Status DropoutGrad::Compute(OpKernelContext* context) const { return Status::OK(); } - } // namespace contrib } // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/nn/dropout_op.h b/orttraining/orttraining/training_ops/cpu/nn/dropout_op.h index 54d0d6f351498..69244c4d11f92 100644 --- a/orttraining/orttraining/training_ops/cpu/nn/dropout_op.h +++ b/orttraining/orttraining/training_ops/cpu/nn/dropout_op.h @@ -9,22 +9,6 @@ namespace onnxruntime { namespace contrib { -template -class Dropout final: public OpKernel { - public: - Dropout(const OpKernelInfo& info) : OpKernel{info} { - int64_t seed = 0; - if (info.GetAttr("seed", &seed).IsOK()) { - generator_ = onnxruntime::make_unique(seed); - } - } - - Status Compute(OpKernelContext* context) const override; - - private: - mutable std::unique_ptr generator_; -}; - template class DropoutGrad final : public OpKernel { public: diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 537864f1a1d29..7ce8c95aa0517 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -70,15 +70,6 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_MLFloat16, TrainableDropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_float, TrainableDropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double_double, TrainableDropoutGrad); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_MLFloat16, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_float, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, MLFloat16_double, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_MLFloat16, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_float, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, float_double, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_MLFloat16, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_float, Dropout); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 12, double_double, Dropout); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_MLFloat16, DropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_float, DropoutGrad); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_double, DropoutGrad); @@ -194,15 +185,6 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout.cc b/orttraining/orttraining/training_ops/cuda/nn/dropout.cc index 8f4ecb77e0983..54cdddeb7f051 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout.cc +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout.cc @@ -3,37 +3,12 @@ #include "core/framework/random_seed.h" #include "orttraining/training_ops/cuda/nn/dropout.h" - +#include "core/providers/cuda/nn/dropout.h" #include "core/providers/common.h" namespace onnxruntime { namespace cuda { -#define REGISTER_KERNEL_TYPED(T1, T2) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - Dropout, \ - kOnnxDomain, \ - 12, \ - T1##_##T2, \ - kCudaExecutionProvider, \ - KernelDefBuilder() \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(1) \ - .InputMemoryType(2), \ - Dropout); - -REGISTER_KERNEL_TYPED(MLFloat16, MLFloat16) -REGISTER_KERNEL_TYPED(MLFloat16, float) -REGISTER_KERNEL_TYPED(MLFloat16, double) -REGISTER_KERNEL_TYPED(float, MLFloat16) -REGISTER_KERNEL_TYPED(float, float) -REGISTER_KERNEL_TYPED(float, double) -REGISTER_KERNEL_TYPED(double, MLFloat16) -REGISTER_KERNEL_TYPED(double, float) -REGISTER_KERNEL_TYPED(double, double) - #define REGISTER_TRAINABLE_KERNEL_TYPED(T1, T2) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ TrainableDropout, \ @@ -59,68 +34,6 @@ REGISTER_TRAINABLE_KERNEL_TYPED(double, MLFloat16) REGISTER_TRAINABLE_KERNEL_TYPED(double, float) REGISTER_TRAINABLE_KERNEL_TYPED(double, double) -template -Status Dropout::ComputeInternal(OpKernelContext* context) const { - typedef typename ToCudaType::MappedType CudaT; - - //Get X_data - const Tensor* X = context->Input(0); - if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "X Input is not available."); - const TensorShape& shape = X->Shape(); - auto X_data = reinterpret_cast(X->template Data()); - const int64_t N = shape.Size(); - - //Get Y_data - auto Y = context->Output(0, shape); - auto Y_data = reinterpret_cast(Y->template MutableData()); - - //Get mask_data - auto mask = context->Output(1, shape); - ORT_ENFORCE(!mask || mask->Shape().Size() == N); - - //Get the ratio_data - float ratio_data; - auto ratio = context->Input(1); - - static_assert(std::is_same::value || std::is_same::value || std::is_same::value, - "T2 must be float16 or float or double"); - - if (ratio) { - ratio_data = static_cast(*(ratio->template Data())); - } else { - ratio_data = default_ratio_; - } - ORT_ENFORCE(ratio_data >= 0.0f && ratio_data < 1.0f); - - const Tensor* training_mode = context->Input(2); - //Check for inference mode. - if ((0 == ratio_data /*Backward compat with TrainableDropout*/) || - (!trainable_dropout && (training_mode == nullptr || *(training_mode->Data()) == false))) { - if (Y_data != X_data) { - CUDA_CALL_THROW(cudaMemcpyAsync(Y_data, X_data, N * sizeof(T1), cudaMemcpyDeviceToDevice)); - } - - // If mask is requested, return all 1s. - if (mask != nullptr) { - ORT_ENFORCE(cudaMemset(mask->MutableData(), true, N * sizeof(bool)) == cudaSuccess); - } - - return Status::OK(); - } - - IAllocatorUniquePtr temp_mask_buffer{}; // buffer to use if mask is not provided - bool* const mask_data = [this, N, mask, &temp_mask_buffer]() { - if (mask) return mask->MutableData(); - temp_mask_buffer = GetScratchBuffer(N); - return temp_mask_buffer.get(); - }(); - - PhiloxGenerator& generator = generator_ != nullptr ? *generator_.get() : PhiloxGenerator::Default(); - DropoutKernelImpl(GetDeviceProp(), N, ratio_data, generator, X_data, Y_data, mask_data); - - return Status::OK(); -} - #define REGISTER_GRADIENT_KERNEL_TYPED(OpName, T1, T2) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ OpName, \ diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout.h b/orttraining/orttraining/training_ops/cuda/nn/dropout.h index dca18128d07d2..12084bfc7fd73 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout.h +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout.h @@ -9,23 +9,6 @@ namespace onnxruntime { namespace cuda { -template -class Dropout final : public CudaKernel { - public: - Dropout(const OpKernelInfo& info) : CudaKernel(info), default_ratio_(0.5) { - int64_t seed = 0; - if (info.GetAttr("seed", &seed).IsOK()) { - generator_ = onnxruntime::make_unique(static_cast(seed)); - } - } - - Status ComputeInternal(OpKernelContext* context) const override; - - private: - mutable std::unique_ptr generator_; - const float default_ratio_; -}; - template class DropoutGrad final : public CudaKernel { public: diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu b/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu index 3bbb08ba0877d..df296f0fb1cab 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.cu @@ -24,81 +24,6 @@ namespace onnxruntime { namespace cuda { -constexpr int UNROLL = 4; - -template -__global__ void DropoutKernel( - const int64_t N, - const float ratio, - const std::pair seeds, - const T* X_data, - T* Y_data, - bool* mask_data) { - const float p = 1.0f - ratio; - const T scale = T(1.0f / p); - - CUDA_LONG idx = blockDim.x * blockIdx.x + threadIdx.x; - CUDA_LONG step_size = gridDim.x * blockDim.x * UNROLL; - CUDA_LONG rounded_size = ((N - 1) / step_size + 1) * step_size; - - curandStatePhilox4_32_10_t state; - curand_init(seeds.first, idx, seeds.second, &state); - - // 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; id < rounded_size; id += step_size) { - float4 rand = curand_uniform4(&state); - - for (CUDA_LONG i = 0; i < UNROLL; i++) { - CUDA_LONG li = id + gridDim.x * blockDim.x * i; - if (li < N) { - mask_data[li] = (&rand.x)[i] < p; - Y_data[li] = X_data[li] * T(mask_data[li]) * scale; - } - } - - __syncthreads(); - } -} - -template -void DropoutKernelImpl( - const cudaDeviceProp& prop, - const int64_t N, - const float ratio, - PhiloxGenerator& generator, - const T* X_data, - T* Y_data, - bool* mask_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(CeilDiv(N, block_size))); - - // 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(((N - 1) / (block_size * grid_size * UNROLL) + 1) * UNROLL); - auto seeds = generator.NextPhiloxSeeds(counter_offset); - - DropoutKernel<<>>(N, ratio, seeds, X_data, Y_data, mask_data); -} - -#define SPECIALIZED_DROPOUT_IMPL(T) \ - template void DropoutKernelImpl( \ - const cudaDeviceProp& prop, \ - const int64_t N, \ - const float ratio, \ - PhiloxGenerator& generator, \ - const T* X_data, \ - T* Y_data, \ - bool* mask_data); - -SPECIALIZED_DROPOUT_IMPL(float) -SPECIALIZED_DROPOUT_IMPL(double) -SPECIALIZED_DROPOUT_IMPL(half) - template __global__ void DropoutGradientKernel( const int64_t N, diff --git a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h b/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h index ad128b15ca52a..b75ee462a6b0a 100644 --- a/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h +++ b/orttraining/orttraining/training_ops/cuda/nn/dropout_impl.h @@ -8,16 +8,6 @@ namespace onnxruntime { namespace cuda { -template -void DropoutKernelImpl( - const cudaDeviceProp& prop, - const int64_t N, - const float ratio, - PhiloxGenerator& generator, - const T* X_data, - T* Y_data, - bool* mask_data); - template void DropoutGradientKernelImpl( const int64_t N,