saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Optimize OneHot CUDA Kernel (#4012)

Browse source 
* Optimize for OneHot with zero off value.

* Add test cases for indices out of range.

Co-authored-by: Vincent Wang <weicwang@microsoft.com>
Co-authored-by: Vincent Wang <weicwang@OrtDevTest2v100.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Vincent Wang 2020-05-26 18:12:11 +08:00 committed by GitHub
parent 0a6d9dd301
commit 9d0534c0eb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 371 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
onnxruntime/core/providers/cuda/tensor/onehot.cc
View file

@ -59,12 +59,22 @@ Status OneHotOp<in_type, out_type, depth_type>::ComputeInternal(OpKernelContext*
if (output->Shape().Size() == 0)
return Status::OK();
const fast_divmod fdm_depth_suffix(gsl::narrow_cast<int>(depth_val * suffix_dim_size));
const fast_divmod fdm_suffix(gsl::narrow_cast<int>(suffix_dim_size));
const auto* indices_data = indices->Data<in_type>();
auto* output_data = reinterpret_cast<CudaT_Out*>(output->MutableData<out_type>());
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<int>(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<in_type, out_type, depth_type>::ComputeInternal(OpKernelContext*
}
} // namespace cuda
} // namespace onnxruntime
} // namespace onnxruntime

55
onnxruntime/core/providers/cuda/tensor/onehot.cu
View file

@ -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<typename in_type, typename out_type>
__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 <typename in_type, typename out_type>
void OneHotImpl(
const in_type* indices_data,
@ -60,6 +77,25 @@ void OneHotImpl(
N);
}
template <typename in_type, typename out_type>
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<float>(count) / GridDim::maxThreadsPerBlock));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
_OneHotWithZeroOffValueImpl<in_type, out_type><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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
} // namespace onnxruntime

9
onnxruntime/core/providers/cuda/tensor/onehot.h
View file

@ -20,6 +20,15 @@ void OneHotImpl(
out_type* output,
size_t count);
template <typename in_type, typename out_type>
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 <typename in_type, typename out_type, typename depth_type>
class OneHotOp final : public CudaKernel {
public:

296
onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc
View file

@ -204,6 +204,201 @@ TEST(OneHotOpTest, FloatInt64) {
test.Run();
}
TEST(OneHotOpTest, DefaultAxis_NonZeroOffValue) {
OpTester test("OneHot", 11);
test.AddInput<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<float>("indices", {2, 3}, {1., 9., 8., 2., 4., 6.});
test.AddInput<float>("depth", {1}, {10.});
test.AddInput<float>("values", {2}, {2., 3.});
test.AddOutput<float>("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<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<float>("depth", {1}, {10.});
test.AddInput<int32_t>("values", {2}, {2, 3});
test.AddOutput<int32_t>("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<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<float>("values", {2}, {2, 3});
test.AddOutput<float>("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<int32_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<float>("depth", {1}, {10.0f});
test.AddInput<float>("values", {2}, {2.0f, 3.0f});
test.AddOutput<float>("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<int32_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int32_t>("depth", {1}, {10});
test.AddInput<float>("values", {2}, {2.0f, 3.0f});
test.AddOutput<float>("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<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<float>("indices", {2, 3}, {1.f, 9.f, 8.f, 2.f, 4.f, 6.f});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<float>("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<int64_t>("indices", {2, 3}, {1, -1, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<int64_t>("indices", {2, 3}, {1, -1, 8, 13, 4, -12});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {0, 1});
test.AddOutput<int64_t>("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<int64_t>("indices", {2, 3}, {1, -1, 8, 13, 4, -12});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<int64_t>("values", {2}, {2, 3});
test.AddOutput<int64_t>("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<float> values{2.0f, 3.0f};
std::vector<MLFloat16> fp16_values(values.size());
ConvertFloatToMLFloat16(values.data(), fp16_values.data(), static_cast<int>(values.size()));
std::vector<float> 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<MLFloat16> fp16_output(output.size());
ConvertFloatToMLFloat16(output.data(), fp16_output.data(), static_cast<int>(output.size()));
test.AddInput<int64_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int64_t>("depth", {1}, {10});
test.AddInput<MLFloat16>("values", {2}, fp16_values);
test.AddOutput<MLFloat16>("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<float> values{2.0f, 3.0f};
std::vector<MLFloat16> fp16_values(values.size());
ConvertFloatToMLFloat16(values.data(), fp16_values.data(), static_cast<int>(values.size()));
std::vector<float> 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<MLFloat16> fp16_output(output.size());
ConvertFloatToMLFloat16(output.data(), fp16_output.data(), static_cast<int>(output.size()));
test.AddInput<int32_t>("indices", {2, 3}, {1, 9, 8, 2, 4, 6});
test.AddInput<int32_t>("depth", {1}, {10});
test.AddInput<MLFloat16>("values", {2}, fp16_values);
test.AddOutput<MLFloat16>("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