Skip to content

Commit a9f0051

Browse files
andflo-Armandflo-Arm
committed
fix!: Add tensor size check to kernels
The size check implies a tensor size restriction to 2^31-1 bytes. Kernel configurations larger than that will no longer validate. Resolves: COMPMID-8697 Signed-off-by: Andreas Flöjt <andreas.floejt@arm.com> Change-Id: I54f73ade5cb4a0d34d831505d83d1d7ef526b5db
1 parent e7ed1af commit a9f0051

File tree

147 files changed

+1740
-475
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

147 files changed

+1740
-475
lines changed

‎src/core/CL/kernels/CLArgMinMaxLayerKernel.cpp‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2019-2021, 2023 Arm Limited.
2+
* Copyright (c) 2019-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -30,9 +30,9 @@
3030
#include "arm_compute/core/TensorInfo.h"
3131
#include "arm_compute/core/Utils.h"
3232
#include "arm_compute/core/utils/helpers/AdjustVecSize.h"
33-
#include "arm_compute/core/Validate.h"
3433

3534
#include "src/core/CL/CLValidate.h"
35+
#include "src/core/CPP/Validate.h"
3636
#include "src/core/helpers/AutoConfiguration.h"
3737
#include "src/core/helpers/WindowHelpers.h"
3838
#include "support/StringSupport.h"
@@ -44,10 +44,12 @@ namespace
4444
Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, unsigned int axis, ReductionOperation op)
4545
{
4646
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
47+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(input);
4748
ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(input);
48-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 1, DataType::QASYMM8, DataType::QASYMM8_SIGNED,
49-
DataType::S32, DataType::F16, DataType::F32);
50-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::S32, DataType::S64);
49+
const size_t one_channel = 1u;
50+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(
51+
input, one_channel, DataType::QASYMM8, DataType::QASYMM8_SIGNED, DataType::S32, DataType::F16, DataType::F32);
52+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, one_channel, DataType::S32, DataType::S64);
5153

5254
ARM_COMPUTE_RETURN_ERROR_ON_MSG(op != ReductionOperation::ARG_IDX_MAX && op != ReductionOperation::ARG_IDX_MIN,
5355
"Only ARG_IDX_MAX and ARG_IDX_MIN are supported");
@@ -57,8 +59,13 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, u
5759

5860
if (output->total_size() != 0)
5961
{
60-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::U32, DataType::S32, DataType::S64,
61-
DataType::U64);
62+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
63+
}
64+
else
65+
{
66+
// Assume largest possible data type since we don't know.
67+
const auto output_info = TensorInfo(input->tensor_shape(), one_channel, DataType::S64);
68+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output_info);
6269
}
6370

6471
return Status{};

‎src/core/CL/kernels/CLBatchNormalizationLayerKernel.cpp‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2017-2021, 2023 Arm Limited.
2+
* Copyright (c) 2017-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -34,6 +34,7 @@
3434
#include "arm_compute/core/utils/StringUtils.h"
3535

3636
#include "src/core/CL/CLValidate.h"
37+
#include "src/core/CPP/Validate.h"
3738
#include "src/core/helpers/AutoConfiguration.h"
3839
#include "src/core/helpers/WindowHelpers.h"
3940
#include "support/StringSupport.h"
@@ -52,19 +53,24 @@ Status validate_arguments(const ITensorInfo *input,
5253
ActivationLayerInfo act_info)
5354
{
5455
ARM_COMPUTE_UNUSED(epsilon);
56+
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, mean, var);
57+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(input, mean, var);
5558
ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(input);
56-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, 1, DataType::F16, DataType::F32);
59+
const size_t one_channel = 1u;
60+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, one_channel, DataType::F16, DataType::F32);
5761
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(mean, var);
5862
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, mean, var);
5963
ARM_COMPUTE_RETURN_ERROR_ON(input->dimension(get_data_layout_dimension_index(
6064
input->data_layout(), DataLayoutDimension::CHANNEL)) != mean->dimension(0));
6165
if (beta != nullptr)
6266
{
67+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(beta);
6368
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(mean, beta);
6469
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, beta);
6570
}
6671
if (gamma != nullptr)
6772
{
73+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(gamma);
6874
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(mean, gamma);
6975
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, gamma);
7076
}
@@ -82,10 +88,16 @@ Status validate_arguments(const ITensorInfo *input,
8288

8389
if (output != nullptr && output->total_size() != 0)
8490
{
91+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
8592
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(input, output);
8693
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_LAYOUT(input, output);
8794
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
8895
}
96+
else
97+
{
98+
// No configured output. Since `output` is expected to match `input`,
99+
// there's nothing extra to check in this case.
100+
}
89101

90102
return Status{};
91103
}

‎src/core/CL/kernels/CLBatchToSpaceLayerKernel.cpp‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2018-2021, 2023 Arm Limited.
2+
* Copyright (c) 2018-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -30,6 +30,7 @@
3030
#include "arm_compute/core/utils/StringUtils.h"
3131

3232
#include "src/core/CL/CLValidate.h"
33+
#include "src/core/CPP/Validate.h"
3334
#include "src/core/helpers/AutoConfiguration.h"
3435
#include "src/core/helpers/WindowHelpers.h"
3536
#include "support/StringSupport.h"
@@ -42,16 +43,23 @@ namespace
4243
Status validate_arguments(const ITensorInfo *input, const ITensorInfo *block_info, const ITensorInfo *output)
4344
{
4445
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, block_info, output);
45-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(block_info, 1, DataType::S32);
46+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(input, block_info);
47+
const size_t one_channel = 1u;
48+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(block_info, one_channel, DataType::S32);
4649
ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
4750
ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
4851

4952
// Validate output if initialized
5053
if (output->total_size() != 0)
5154
{
55+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
5256
ARM_COMPUTE_RETURN_ERROR_ON(output->num_dimensions() > 4);
5357
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
5458
}
59+
else
60+
{
61+
// Ignored; dynamic block is deprecated.
62+
}
5563

5664
return Status{};
5765
}
@@ -62,6 +70,9 @@ Status validate_arguments_static(const ITensorInfo *input,
6270
const CropInfo &crop_info)
6371
{
6472
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
73+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(input);
74+
const size_t one_channel = 1u;
75+
ARM_COMPUTE_RETURN_ERROR_ON(input->num_channels() != one_channel);
6576
ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
6677
ARM_COMPUTE_RETURN_ERROR_ON(block_shape_x <= 0);
6778
ARM_COMPUTE_RETURN_ERROR_ON(block_shape_y <= 0);
@@ -70,16 +81,23 @@ Status validate_arguments_static(const ITensorInfo *input,
7081
const int idx_batch = get_data_layout_dimension_index(data_layout, DataLayoutDimension::BATCHES);
7182
ARM_COMPUTE_RETURN_ERROR_ON(input->tensor_shape()[idx_batch] % (block_shape_x * block_shape_y) != 0);
7283

84+
const TensorShape expected_output_shape = compute_batch_to_space_shape(input->data_layout(), input->tensor_shape(),
85+
block_shape_x, block_shape_y, crop_info);
86+
7387
// Validate output if initialized
7488
if (output->total_size() != 0)
7589
{
76-
const TensorShape expected_output_shape = compute_batch_to_space_shape(
77-
input->data_layout(), input->tensor_shape(), block_shape_x, block_shape_y, crop_info);
90+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
7891
const TensorInfo expected_output = output->clone()->set_tensor_shape(expected_output_shape);
7992
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(output, &expected_output);
8093
ARM_COMPUTE_RETURN_ERROR_ON(output->num_dimensions() > 4);
8194
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
8295
}
96+
else
97+
{
98+
const auto output_info = TensorInfo(expected_output_shape, one_channel, input->data_type());
99+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output_info);
100+
}
83101

84102
return Status{};
85103
}

‎src/core/CL/kernels/CLBitwiseKernel.cpp‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2020-2021, 2023 Arm Limited.
2+
* Copyright (c) 2020-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -31,6 +31,7 @@
3131
#include "arm_compute/core/utils/helpers/AdjustVecSize.h"
3232
#include "arm_compute/core/Validate.h"
3333

34+
#include "src/core/CPP/Validate.h"
3435
#include "src/core/helpers/AutoConfiguration.h"
3536
#include "src/core/helpers/WindowHelpers.h"
3637
#include "support/StringSupport.h"
@@ -49,17 +50,23 @@ void CLBitwiseKernel::configure(const CLCompileContext &compile_context,
4950
BitwiseOperation op)
5051
{
5152
ARM_COMPUTE_ERROR_ON_NULLPTR(input1);
52-
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input1, 1, DataType::U8);
53+
ARM_COMPUTE_ERROR_ON_SIZE_UNSUPPORTED(input1->info());
54+
#ifdef ARM_COMPUTE_ASSERTS_ENABLED
55+
const size_t one_channel = 1u;
56+
#endif
57+
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input1, one_channel, DataType::U8);
5358
if (op != BitwiseOperation::NOT)
5459
{
5560
ARM_COMPUTE_ERROR_ON_NULLPTR(input2);
56-
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input2, 1, DataType::U8);
61+
ARM_COMPUTE_ERROR_ON_SIZE_UNSUPPORTED(input2->info());
62+
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input2, one_channel, DataType::U8);
5763
}
5864
ARM_COMPUTE_ERROR_ON_NULLPTR(output);
59-
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::U8);
65+
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, one_channel, DataType::U8);
6066

61-
// Output auto inizialitation if not yet initialized
67+
// Output auto initialization if not yet initialized
6268
auto_init_if_empty(*(output->info()), *(input1->info()));
69+
ARM_COMPUTE_ERROR_ON_SIZE_UNSUPPORTED(output->info());
6370
auto padding_info = get_padding_info({input1, input2, output});
6471

6572
// Configure kernel window

‎src/core/CL/kernels/CLBoundingBoxTransformKernel.cpp‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2018-2021, 2023 Arm Limited.
2+
* Copyright (c) 2018-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -33,6 +33,7 @@
3333
#include "arm_compute/core/utils/StringUtils.h"
3434

3535
#include "src/core/CL/CLValidate.h"
36+
#include "src/core/CPP/Validate.h"
3637
#include "src/core/helpers/AutoConfiguration.h"
3738
#include "src/core/helpers/WindowHelpers.h"
3839
#include "support/StringSupport.h"
@@ -47,9 +48,13 @@ Status validate_arguments(const ITensorInfo *boxes,
4748
const BoundingBoxTransformInfo &info)
4849
{
4950
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(boxes, pred_boxes, deltas);
51+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(boxes, deltas);
5052
ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(boxes);
51-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_NOT_IN(boxes, DataType::QASYMM16, DataType::F32, DataType::F16);
52-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_NOT_IN(deltas, DataType::QASYMM8, DataType::F32, DataType::F16);
53+
const size_t one_channel = 1u;
54+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(boxes, one_channel, DataType::QASYMM16, DataType::F32,
55+
DataType::F16);
56+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(deltas, one_channel, DataType::QASYMM8, DataType::F32,
57+
DataType::F16);
5358
ARM_COMPUTE_RETURN_ERROR_ON(deltas->tensor_shape()[1] != boxes->tensor_shape()[1]);
5459
ARM_COMPUTE_RETURN_ERROR_ON(deltas->tensor_shape()[0] % 4 != 0);
5560
ARM_COMPUTE_RETURN_ERROR_ON(boxes->tensor_shape()[0] != 4);
@@ -62,7 +67,7 @@ Status validate_arguments(const ITensorInfo *boxes,
6267
const UniformQuantizationInfo boxes_qinfo = boxes->quantization_info().uniform();
6368
ARM_COMPUTE_RETURN_ERROR_ON(boxes_qinfo.scale != 0.125f);
6469
ARM_COMPUTE_RETURN_ERROR_ON(boxes_qinfo.offset != 0);
65-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_NOT_IN(deltas, DataType::QASYMM8);
70+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(deltas, one_channel, DataType::QASYMM8);
6671
}
6772
else
6873
{
@@ -71,6 +76,7 @@ Status validate_arguments(const ITensorInfo *boxes,
7176

7277
if (pred_boxes->total_size() > 0)
7378
{
79+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(pred_boxes);
7480
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DIMENSIONS(pred_boxes->tensor_shape(), deltas->tensor_shape());
7581
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(pred_boxes, boxes);
7682
ARM_COMPUTE_RETURN_ERROR_ON(pred_boxes->num_dimensions() > 2);
@@ -81,6 +87,11 @@ Status validate_arguments(const ITensorInfo *boxes,
8187
ARM_COMPUTE_RETURN_ERROR_ON(pred_boxes_qinfo.offset != 0);
8288
}
8389
}
90+
else
91+
{
92+
const auto pred_boxes_info = TensorInfo(deltas->tensor_shape(), one_channel, boxes->data_type());
93+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&pred_boxes_info);
94+
}
8495
ARM_COMPUTE_RETURN_ERROR_ON(info.scale() <= 0);
8596

8697
return Status{};

‎src/core/CL/kernels/CLChannelShuffleLayerKernel.cpp‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2018-2021, 2023 Arm Limited.
2+
* Copyright (c) 2018-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -33,6 +33,7 @@
3333
#include "arm_compute/core/utils/StringUtils.h"
3434

3535
#include "src/core/CL/CLValidate.h"
36+
#include "src/core/CPP/Validate.h"
3637
#include "src/core/helpers/AutoConfiguration.h"
3738
#include "src/core/helpers/WindowHelpers.h"
3839
#include "support/StringSupport.h"
@@ -43,6 +44,8 @@ namespace
4344
{
4445
Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, unsigned int num_groups)
4546
{
47+
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
48+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(input);
4649
ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(input);
4750
ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
4851
ARM_COMPUTE_RETURN_ERROR_ON_MSG(num_groups < 2, "Channel shuffling with less than 2 groups would be inefficient");
@@ -61,10 +64,16 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, u
6164
// Checks performed when output is configured
6265
if (output->total_size() != 0)
6366
{
67+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
6468
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(input, output);
6569
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_QUANTIZATION_INFO(input, output);
6670
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
6771
}
72+
else
73+
{
74+
// No configured output. Since `output` is expected to match `input`,
75+
// there's nothing extra to check in this case.
76+
}
6877

6978
return Status{};
7079
}

‎src/core/CL/kernels/CLComparisonKernel.cpp‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2018-2021, 2023 Arm Limited.
2+
* Copyright (c) 2018-2021, 2023, 2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -25,10 +25,12 @@
2525

2626
#include "arm_compute/core/CL/CLHelpers.h"
2727
#include "arm_compute/core/CL/ICLTensor.h"
28+
#include "arm_compute/core/TensorInfo.h"
2829
#include "arm_compute/core/utils/helpers/AdjustVecSize.h"
2930
#include "arm_compute/core/utils/StringUtils.h"
3031

3132
#include "src/core/CL/CLValidate.h"
33+
#include "src/core/CPP/Validate.h"
3234
#include "src/core/helpers/AutoConfiguration.h"
3335
#include "src/core/helpers/WindowHelpers.h"
3436
#include "support/StringSupport.h"
@@ -51,6 +53,7 @@ Status validate_arguments(const ITensorInfo &input1,
5153
const ITensorInfo &output,
5254
ComparisonOperation operation)
5355
{
56+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&input1, &input2);
5457
ARM_COMPUTE_RETURN_ERROR_ON_F16_UNSUPPORTED(&input1);
5558
ARM_COMPUTE_RETURN_ERROR_ON(input1.data_type() == DataType::UNKNOWN);
5659
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(&input1, &input2);
@@ -59,13 +62,21 @@ Status validate_arguments(const ITensorInfo &input1,
5962
const TensorShape out_shape = TensorShape::broadcast_shape(input1.tensor_shape(), input2.tensor_shape());
6063
ARM_COMPUTE_RETURN_ERROR_ON_MSG(out_shape.total_size() == 0, "Inputs are not broadcast compatible");
6164

65+
const size_t one_channel = 1u;
66+
6267
// Validate in case of configured output
6368
if (output.total_size() > 0)
6469
{
65-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(&output, 1, DataType::U8);
70+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output);
71+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(&output, one_channel, DataType::U8);
6672
ARM_COMPUTE_RETURN_ERROR_ON_MSG(detail::have_different_dimensions(out_shape, output.tensor_shape(), 0),
6773
"Wrong shape for output");
6874
}
75+
else
76+
{
77+
const auto output_info = TensorInfo(out_shape, one_channel, DataType::U8);
78+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output_info);
79+
}
6980

7081
return Status{};
7182
}

0 commit comments

Comments
 (0)