Skip to content

Commit

Permalink
Implement DepthToSpace uint8_t and Enable DropQDQNodesRules
Browse files Browse the repository at this point in the history 
- Implemented the DepthToSpace uint8_t kernel.
- Enabled DropQDQNodesRules for DepthToSpace.
- Added unit tests for the DepthToSpace uint8_t kernel.
  • Loading branch information
@yihonglyu
yihonglyu committed Dec 5, 2024
1 parent 6ed77cc commit f020f6c
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 9 deletions.
5 changes: 3 additions & 2 deletions ...ntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,10 +91,11 @@ void DropQDQNodesRules(SelectorActionRegistry& qdq_selector_action_registry) {
std::move(drop_action_no_int16_and_positive_scale));

std::unique_ptr<NodeSelector> selector = std::make_unique<QDQ::DropQDQNodesSelector>(true, false, true, providers);
// DepthToSpace and SpaceToDepth not included because there are no integer implementations.
// SpaceToDepth not included because there are no integer implementations.
// https://github.com/microsoft/onnxruntime/issues/21287
qdq_selector_action_registry.RegisterSelectorAndAction(drop_action_name,
{{"Expand", {}},
{{"DepthToSpace", {}},
{"Expand", {}},
{"Flatten", {}},
{"Gather", {}},
{"GatherElements", {}},
Expand Down
16 changes: 15 additions & 1 deletion onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,8 @@ ONNX_CPU_OPERATOR_KERNEL(
13,
KernelDefBuilder()
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>()}),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<uint8_t>()}),
DepthToSpace);

// intermediate tensor shapes are:
Expand Down Expand Up @@ -196,6 +197,19 @@ Status DepthToSpace::Compute(OpKernelContext* context) const {
onnxruntime::narrow<std::ptrdiff_t>(blocksize_),
onnxruntime::narrow<std::ptrdiff_t>(input_width),
onnxruntime::narrow<std::ptrdiff_t>(blocksize_));
} else if (input.IsDataType<uint8_t>()) {
SpaceDepthOpCpuImpl<uint8_t>(input, output, permutation,
onnxruntime::narrow<std::ptrdiff_t>(batch),
onnxruntime::narrow<std::ptrdiff_t>(dim1),
onnxruntime::narrow<std::ptrdiff_t>(blocksize_),
onnxruntime::narrow<std::ptrdiff_t>(dim3),
onnxruntime::narrow<std::ptrdiff_t>(input_height),
onnxruntime::narrow<std::ptrdiff_t>(input_width),
onnxruntime::narrow<std::ptrdiff_t>(input_depth / blocksize_ / blocksize_),
onnxruntime::narrow<std::ptrdiff_t>(input_height),
onnxruntime::narrow<std::ptrdiff_t>(blocksize_),
onnxruntime::narrow<std::ptrdiff_t>(input_width),
onnxruntime::narrow<std::ptrdiff_t>(blocksize_));
} else {
// user will not see this as the kernel doesn't claim support for types other than float and double
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported input type in DepthToSpace op: ", input.DataType());
Expand Down
42 changes: 36 additions & 6 deletions onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ template <typename T>
class TensorOpTest : public ::testing::Test {
};

using TensorOpTestTypes = ::testing::Types<float, MLFloat16>;
using TensorOpTestTypes = ::testing::Types<float, MLFloat16, uint8_t>;
TYPED_TEST_SUITE(TensorOpTest, TensorOpTestTypes);

TEST(TensorOpTest, SpaceToDepthTest_1) {
Expand Down Expand Up @@ -224,6 +224,7 @@ TEST(TensorOpTest, DepthToSpaceTest_1_double) {
test.AddOutput<double>("output", {N, C / (blocksize * blocksize), H * blocksize, W * blocksize}, result);
test.Run();
}

TEST(TensorOpTest, DepthToSpaceTest_2) {
OpTester test("DepthToSpace", 7); // create an opset 7 model
constexpr int64_t blocksize = 2;
Expand Down Expand Up @@ -308,14 +309,24 @@ TYPED_TEST(TensorOpTest, DepthToSpaceTest_3) {
if constexpr (std::is_same<TypeParam, float>::value) {
test.AddInput<float>("input", {N, C, H, W}, X);
test.AddOutput<float>("output", {2, 3, 6, 4}, result);
} else {
} else if constexpr (std::is_same<TypeParam, MLFloat16>::value) {
std::vector<TypeParam> X_fp16(X.size());
std::vector<TypeParam> result_fp16(result.size());
ConvertFloatToMLFloat16(result.data(), result_fp16.data(), result.size());
ConvertFloatToMLFloat16(X.data(), X_fp16.data(), X.size());
test.AddOutput<TypeParam>("output", {2, 3, 6, 4}, result_fp16);
ConvertFloatToMLFloat16(result.data(), result_fp16.data(), result.size());
test.AddInput<TypeParam>("input", {N, C, H, W}, X_fp16);
test.AddOutput<TypeParam>("output", {2, 3, 6, 4}, result_fp16);
} else if constexpr (std::is_same<TypeParam, uint8_t>::value) {
std::vector<uint8_t> X_u8(X.size());
std::vector<uint8_t> result_u8(result.size());
ConvertFloatToUint8_t(X.data(), X_u8.data(), X.size());
ConvertFloatToUint8_t(result.data(), result_u8.data(), result.size());
test.AddInput<uint8_t>("input", {N, C, H, W}, X_u8);
test.AddOutput<uint8_t>("output", {2, 3, 6, 4}, result_u8);
} else {
ORT_THROW("Type not supported");
}

// TODO: Test is flaky on QNN EP (CPU backend).
// Re-enable when the QnnCPUBackendTests.DISABLED_SpaceToDepth_Flaky test is fixed.
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kQnnExecutionProvider});
Expand Down Expand Up @@ -363,13 +374,22 @@ TYPED_TEST(TensorOpTest, DepthToSpaceTest_4) {
if constexpr (std::is_same<TypeParam, float>::value) {
test.AddInput<float>("input", {N, C, H, W}, X);
test.AddOutput<float>("output", {2, 3, 6, 4}, result);
} else {
} else if constexpr (std::is_same<TypeParam, MLFloat16>::value) {
std::vector<TypeParam> X_fp16(X.size());
std::vector<TypeParam> result_fp16(result.size());
ConvertFloatToMLFloat16(X.data(), X_fp16.data(), X.size());
ConvertFloatToMLFloat16(result.data(), result_fp16.data(), result.size());
test.AddInput<TypeParam>("input", {N, C, H, W}, X_fp16);
test.AddOutput<TypeParam>("output", {2, 3, 6, 4}, result_fp16);
} else if constexpr (std::is_same<TypeParam, uint8_t>::value) {
std::vector<uint8_t> X_u8(X.size());
std::vector<uint8_t> result_u8(result.size());
ConvertFloatToUint8_t(X.data(), X_u8.data(), X.size());
ConvertFloatToUint8_t(result.data(), result_u8.data(), result.size());
test.AddInput<uint8_t>("input", {N, C, H, W}, X_u8);
test.AddOutput<uint8_t>("output", {2, 3, 6, 4}, result_u8);
} else {
ORT_THROW("Type not supported");
}

// TODO: Test is flaky on QNN EP (CPU backend).
Expand Down Expand Up @@ -401,14 +421,24 @@ TYPED_TEST(TensorOpTest, DepthToSpaceTest_5) {
if constexpr (std::is_same<TypeParam, float>::value) {
test.AddInput<float>("input", {N, C, H, W}, X);
test.AddOutput<float>("output", {1, 1, 4, 6}, result);
} else {
} else if constexpr (std::is_same<TypeParam, MLFloat16>::value) {
std::vector<TypeParam> X_fp16(X.size());
std::vector<TypeParam> result_fp16(result.size());
ConvertFloatToMLFloat16(X.data(), X_fp16.data(), X.size());
ConvertFloatToMLFloat16(result.data(), result_fp16.data(), result.size());
test.AddInput<TypeParam>("input", {N, C, H, W}, X_fp16);
test.AddOutput<TypeParam>("output", {1, 1, 4, 6}, result_fp16);
} else if constexpr (std::is_same<TypeParam, uint8_t>::value) {
std::vector<uint8_t> X_u8(X.size());
std::vector<uint8_t> result_u8(result.size());
ConvertFloatToUint8_t(X.data(), X_u8.data(), X.size());
ConvertFloatToUint8_t(result.data(), result_u8.data(), result.size());
test.AddInput<uint8_t>("input", {N, C, H, W}, X_u8);
test.AddOutput<uint8_t>("output", {1, 1, 4, 6}, result_u8);
} else {
ORT_THROW("Type not supported");
}

// TODO: Test is flaky on QNN EP (CPU backend).
// Re-enable when the QnnCPUBackendTests.DISABLED_SpaceToDepth_Flaky2 test is fixed.
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kQnnExecutionProvider});
Expand Down
6 changes: 6 additions & 0 deletions onnxruntime/test/providers/provider_test_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ inline void ConvertFloatToMLFloat16(const float* f_datat, MLFloat16* h_data, siz
output_vector = in_vector.template cast<Eigen::half>();
}

inline void ConvertFloatToUint8_t(const float* f_datat, uint8_t* u8_data, size_t input_size) {
auto in_vector = ConstEigenVectorMap<float>(f_datat, input_size);
auto output_vector = EigenVectorMap<uint8_t>(static_cast<uint8_t*>(static_cast<void*>(u8_data)), input_size);
output_vector = in_vector.template cast<uint8_t>();
}

inline void ConvertMLFloat16ToFloat(const MLFloat16* h_data, float* f_data, size_t input_size) {
auto in_vector =
ConstEigenVectorMap<Eigen::half>(static_cast<const Eigen::half*>(static_cast<const void*>(h_data)), input_size);
Expand Down

0 comments on commit f020f6c

Please sign in to comment.