From 9d0534c0ebce73265cbdcb7863c474c26daf98c9 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 26 May 2020 18:12:11 +0800 Subject: [PATCH] Optimize OneHot CUDA Kernel (#4012) * Optimize for OneHot with zero off value. * Add test cases for indices out of range. Co-authored-by: Vincent Wang Co-authored-by: Vincent Wang --- .../core/providers/cuda/tensor/onehot.cc | 16 +- .../core/providers/cuda/tensor/onehot.cu | 55 +++- .../core/providers/cuda/tensor/onehot.h | 9 + .../providers/cpu/tensor/onehot_op_test.cc | 296 ++++++++++++++++++ 4 files changed, 371 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/onehot.cc b/onnxruntime/core/providers/cuda/tensor/onehot.cc index 5e8fde280d..6f302ee273 100644 --- a/onnxruntime/core/providers/cuda/tensor/onehot.cc +++ b/onnxruntime/core/providers/cuda/tensor/onehot.cc @@ -59,12 +59,22 @@ Status OneHotOp::ComputeInternal(OpKernelContext* if (output->Shape().Size() == 0) return Status::OK(); - const fast_divmod fdm_depth_suffix(gsl::narrow_cast(depth_val * suffix_dim_size)); const fast_divmod fdm_suffix(gsl::narrow_cast(suffix_dim_size)); - const auto* indices_data = indices->Data(); auto* output_data = reinterpret_cast(output->MutableData()); + if (values_data[0] == CudaT_Out(0.f)) { + CUDA_RETURN_IF_ERROR(cudaMemset(output->MutableDataRaw(), 0, output->SizeInBytes())); + OneHotWithZeroOffValueImpl(indices_data, + fdm_suffix, + depth_val, + values_data[1], + output_data, + indices->Shape().Size()); + return Status::OK(); + } + + const fast_divmod fdm_depth_suffix(gsl::narrow_cast(depth_val * suffix_dim_size)); OneHotImpl(indices_data, fdm_depth_suffix, fdm_suffix, depth_val, values_data[1], values_data[0], @@ -75,4 +85,4 @@ Status OneHotOp::ComputeInternal(OpKernelContext* } } // namespace cuda -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/onehot.cu b/onnxruntime/core/providers/cuda/tensor/onehot.cu index fadb0d7547..88cf5576dc 100644 --- a/onnxruntime/core/providers/cuda/tensor/onehot.cu +++ b/onnxruntime/core/providers/cuda/tensor/onehot.cu @@ -29,7 +29,7 @@ __global__ void _OneHotImpl( CUDA_LONG indices_index = prefix_index * fdm_suffix.d_ + suffix_index; // handle index outside the range [-depth, depth-1] case - bool is_valid_range = indices_data[indices_index] >= -depth_val || indices_data[indices_index] < depth_val; + bool is_valid_range = indices_data[indices_index] >= -depth_val && indices_data[indices_index] < depth_val; // handle negative index in_type adjusted_indice = (indices_data[indices_index] + depth_val) % depth_val; @@ -37,6 +37,23 @@ __global__ void _OneHotImpl( output_data[id] = (is_valid_range && adjusted_indice == in_type(depth_index)) ? on_value : off_value; } +template +__global__ void _OneHotWithZeroOffValueImpl( + const in_type* indices_data, + const fast_divmod fdm_suffix, + const int64_t depth_val, + const out_type on_value, + out_type* output_data, + CUDA_LONG N) { + CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); + if (indices_data[id] >= -depth_val && indices_data[id] < depth_val) { + in_type adjusted_index = indices_data[id] >= 0 ? indices_data[id] : indices_data[id] + depth_val; + int q, r; + fdm_suffix.divmod(id, q, r); + output_data[(q * depth_val + adjusted_index) * fdm_suffix.d_ + r] = on_value; + } +} + template void OneHotImpl( const in_type* indices_data, @@ -60,6 +77,25 @@ void OneHotImpl( N); } +template +void OneHotWithZeroOffValueImpl( + const in_type* indices_data, + const fast_divmod fdm_suffix, + const int64_t depth_val, + const out_type on_value, + out_type* output_data, + size_t count) { + int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); + CUDA_LONG N = static_cast(count); + _OneHotWithZeroOffValueImpl<<>>( + indices_data, + fdm_suffix, + depth_val, + on_value, + output_data, + N); +} + #define SPECIALIZED_OneHotImpl(in_type, out_type) \ template void OneHotImpl( \ const in_type* indices_data, \ @@ -77,5 +113,20 @@ SPECIALIZED_OneHotImpl(int32_t, float) SPECIALIZED_OneHotImpl(int64_t, half) SPECIALIZED_OneHotImpl(int32_t, half) +#define SPECIALIZED_OneHotWithZeroOffValueImpl(in_type, out_type) \ + template void OneHotWithZeroOffValueImpl( \ + const in_type* indices_data, \ + const fast_divmod fdm_suffix, \ + const int64_t depth_val, \ + const out_type on_value, \ + out_type* output_data, \ + size_t count); + +SPECIALIZED_OneHotWithZeroOffValueImpl(int64_t, int64_t) +SPECIALIZED_OneHotWithZeroOffValueImpl(int64_t, float) +SPECIALIZED_OneHotWithZeroOffValueImpl(int32_t, float) +SPECIALIZED_OneHotWithZeroOffValueImpl(int64_t, half) +SPECIALIZED_OneHotWithZeroOffValueImpl(int32_t, half) + } // namespace cuda -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/onehot.h b/onnxruntime/core/providers/cuda/tensor/onehot.h index dbb4ea5fe2..b58458abec 100644 --- a/onnxruntime/core/providers/cuda/tensor/onehot.h +++ b/onnxruntime/core/providers/cuda/tensor/onehot.h @@ -20,6 +20,15 @@ void OneHotImpl( out_type* output, size_t count); +template +void OneHotWithZeroOffValueImpl( + const in_type* indices, + const fast_divmod fdm_suffix, + const int64_t depth_val, + const out_type on_value, + out_type* output, + size_t count); + template class OneHotOp final : public CudaKernel { public: diff --git a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc index 1fb814578d..140198d78c 100644 --- a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc @@ -204,6 +204,201 @@ TEST(OneHotOpTest, FloatInt64) { test.Run(); } +TEST(OneHotOpTest, DefaultAxis_NonZeroOffValue) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 3, 10}, + {2, 3, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, + 2, 2, 3, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 3, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_float_float_float_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1., 9., 8., 2., 4., 6.}); + test.AddInput("depth", {1}, {10.}); + test.AddInput("values", {2}, {2., 3.}); + test.AddOutput("output", {2, 3, 10}, + {2., 3., 2., 2., 2., 2., 2., 2., 2., 2., + 2., 2., 2., 2., 2., 2., 2., 2., 2., 3., + 2., 2., 2., 2., 2., 2., 2., 2., 3., 2., + 2., 2., 3., 2., 2., 2., 2., 2., 2., 2., + 2., 2., 2., 2., 3., 2., 2., 2., 2., 2., + 2., 2., 2., 2., 2., 2., 3., 2., 2., 2.}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_int64_int32_float_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10.}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 3, 10}, + {2, 3, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, + 2, 2, 3, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 3, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_int64_float_int64_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 3, 10}, + {2, 3, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, + 2, 2, 3, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 3, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_int32_float_float_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10.0f}); + test.AddInput("values", {2}, {2.0f, 3.0f}); + test.AddOutput("output", {2, 3, 10}, + {2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, + 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_int32_float_int32_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2.0f, 3.0f}); + test.AddOutput("output", {2, 3, 10}, + {2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, + 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_0_NonZeroOffValue) { + OpTester test("OneHot", 11); + int64_t axis = 0; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {10, 2, 3}, + {2, 2, 2, 2, 2, 2, + 3, 2, 2, 2, 2, 2, + 2, 2, 2, 3, 2, 2, + 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, + 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, + 2, 2, 3, 2, 2, 2, + 2, 3, 2, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_1_NonZeroOffValue) { + OpTester test("OneHot", 11); + int64_t axis = 1; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 10, 3}, + {2, 2, 2, + 3, 2, 2, + 2, 2, 2, + 2, 2, 2, + 2, 2, 2, + 2, 2, 2, + 2, 2, 2, + 2, 2, 2, + 2, 2, 3, + 2, 3, 2, + 2, 2, 2, + 2, 2, 2, + 3, 2, 2, + 2, 2, 2, + 2, 3, 2, + 2, 2, 2, + 2, 2, 3, + 2, 2, 2, + 2, 2, 2, + 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_2_NonZeroOffValue) { + OpTester test("OneHot", 11); + int64_t axis = 2; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 3, 10}, + {2, 3, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, + 2, 2, 3, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 3, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_Negative_NonDefault_NonZeroOffValue) { + OpTester test("OneHot", 11); + int64_t axis = -3; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {10, 2, 3}, + {2, 2, 2, 2, 2, 2, + 3, 2, 2, 2, 2, 2, + 2, 2, 2, 3, 2, 2, + 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, + 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, + 2, 2, 3, 2, 2, 2, + 2, 3, 2, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, FloatInt64_NonZeroOffValue) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1.f, 9.f, 8.f, 2.f, 4.f, 6.f}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 3, 10}, + {2, 3, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, + 2, 2, 3, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 3, 2, 2, 2}); + test.Run(); +} + TEST(OneHotOpTest, FloatString) { OpTester test("OneHot", 11); test.AddInput("indices", {2, 3}, {1.f, 9.f, 8.f, 2.f, 4.f, 6.f}); @@ -240,6 +435,57 @@ TEST(OneHotOpTest, Axis_Negative_NegIndex_NonDefault) { test.Run(); } +TEST(OneHotOpTest, Axis_Negative_NegIndex_NonDefault_NonZeroOffValue) { + OpTester test("OneHot", 11); + int64_t axis = -3; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, -1, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {10, 2, 3}, + {2, 2, 2, 2, 2, 2, + 3, 2, 2, 2, 2, 2, + 2, 2, 2, 3, 2, 2, + 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, + 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, + 2, 2, 3, 2, 2, 2, + 2, 3, 2, 2, 2, 2}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_IndicesOutOfRange) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, -1, 8, 13, 4, -12}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 10}, + {0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_IndicesOutOfRange_NonZeroOffValue) { + OpTester test("OneHot", 11); + test.AddInput("indices", {2, 3}, {1, -1, 8, 13, 4, -12}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {2, 3}); + test.AddOutput("output", {2, 3, 10}, + {2, 3, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 3, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}); + test.Run(); +} + TEST(OneHotOpTest, DimWithZero) { OpTester test("OneHot", 11); int64_t axis = 0; @@ -303,6 +549,56 @@ TEST(OneHotOpTest, DefaultAxis_int32_MLFloat16_int32 /*indices, output, depth*/) test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider}); } +TEST(OneHotOpTest, DefaultAxis_int64_MLFloat16_int64_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + + std::vector values{2.0f, 3.0f}; + std::vector fp16_values(values.size()); + ConvertFloatToMLFloat16(values.data(), fp16_values.data(), static_cast(values.size())); + + std::vector output{2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, + 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f}; + std::vector fp16_output(output.size()); + ConvertFloatToMLFloat16(output.data(), fp16_output.data(), static_cast(output.size())); + + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, fp16_values); + test.AddOutput("output", {2, 3, 10}, fp16_output); + + // exclude CPU Execution Provider as MLFloat16 is not supported in CPU + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider}); +} + +TEST(OneHotOpTest, DefaultAxis_int32_MLFloat16_int32_NonZeroOffValue /*indices, output, depth*/) { + OpTester test("OneHot", 11); + + std::vector values{2.0f, 3.0f}; + std::vector fp16_values(values.size()); + ConvertFloatToMLFloat16(values.data(), fp16_values.data(), static_cast(values.size())); + + std::vector output{2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, + 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 3.0f, 2.0f, 2.0f, 2.0f}; + std::vector fp16_output(output.size()); + ConvertFloatToMLFloat16(output.data(), fp16_output.data(), static_cast(output.size())); + + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, fp16_values); + test.AddOutput("output", {2, 3, 10}, fp16_output); + + // exclude CPU Execution Provider as MLFloat16 is not supported in CPU + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider}); +} + #endif } // namespace test