Skip to content

Commit 7d0d25f

Browse files
andflo-Armgunes-arm
authored andcommitted
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 09bda62 commit 7d0d25f

File tree

148 files changed

+1738
-481
lines changed

Some content is hidden

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

148 files changed

+1738
-481
lines changed

arm_compute/core/ITensorInfo.h

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2016-2025 Arm Limited.
2+
* Copyright (c) 2016-2026 Arm Limited.
33
*
44
* SPDX-License-Identifier: MIT
55
*
@@ -56,6 +56,10 @@ class ITensorInfo : public misc::ICloneable<ITensorInfo>
5656
using Id = int32_t;
5757
/** An invalid tensor id within a domain */
5858
static constexpr Id invalid_tensor_id = 0;
59+
/** Named constant for commonly used value */
60+
static constexpr size_t one_channel = 1u;
61+
/** Named constant for commonly used value */
62+
static constexpr size_t two_channels = 2u;
5963
/** Get the value representing dynamic dimension state
6064
*
6165
* @return Value representing dynamic dimension state

src/core/CL/kernels/CLArgMinMaxLayerKernel.cpp

Lines changed: 15 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,13 @@ 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+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, ITensorInfo::one_channel, DataType::QASYMM8,
50+
DataType::QASYMM8_SIGNED, DataType::S32, DataType::F16,
51+
DataType::F32);
52+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, ITensorInfo::one_channel, DataType::S32,
53+
DataType::S64);
5154

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

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

6472
return Status{};

src/core/CL/kernels/CLBatchNormalizationLayerKernel.cpp

Lines changed: 13 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,23 @@ 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+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input, ITensorInfo::one_channel, DataType::F16, DataType::F32);
5760
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(mean, var);
5861
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, mean, var);
5962
ARM_COMPUTE_RETURN_ERROR_ON(input->dimension(get_data_layout_dimension_index(
6063
input->data_layout(), DataLayoutDimension::CHANNEL)) != mean->dimension(0));
6164
if (beta != nullptr)
6265
{
66+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(beta);
6367
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(mean, beta);
6468
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, beta);
6569
}
6670
if (gamma != nullptr)
6771
{
72+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(gamma);
6873
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(mean, gamma);
6974
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, gamma);
7075
}
@@ -82,10 +87,16 @@ Status validate_arguments(const ITensorInfo *input,
8287

8388
if (output != nullptr && output->total_size() != 0)
8489
{
90+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
8591
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(input, output);
8692
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_LAYOUT(input, output);
8793
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
8894
}
95+
else
96+
{
97+
// No configured output. Since `output` is expected to match `input`,
98+
// there's nothing extra to check in this case.
99+
}
89100

90101
return Status{};
91102
}

src/core/CL/kernels/CLBatchToSpaceLayerKernel.cpp

Lines changed: 20 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,22 @@ 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+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(block_info, ITensorInfo::one_channel, DataType::S32);
4648
ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
4749
ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
4850

4951
// Validate output if initialized
5052
if (output->total_size() != 0)
5153
{
54+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
5255
ARM_COMPUTE_RETURN_ERROR_ON(output->num_dimensions() > 4);
5356
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
5457
}
58+
else
59+
{
60+
// Ignored; dynamic block is deprecated.
61+
}
5562

5663
return Status{};
5764
}
@@ -62,6 +69,8 @@ Status validate_arguments_static(const ITensorInfo *input,
6269
const CropInfo &crop_info)
6370
{
6471
ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
72+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(input);
73+
ARM_COMPUTE_RETURN_ERROR_ON(input->num_channels() != ITensorInfo::one_channel);
6574
ARM_COMPUTE_RETURN_ERROR_ON(input->num_dimensions() > 4);
6675
ARM_COMPUTE_RETURN_ERROR_ON(block_shape_x <= 0);
6776
ARM_COMPUTE_RETURN_ERROR_ON(block_shape_y <= 0);
@@ -70,16 +79,23 @@ Status validate_arguments_static(const ITensorInfo *input,
7079
const int idx_batch = get_data_layout_dimension_index(data_layout, DataLayoutDimension::BATCHES);
7180
ARM_COMPUTE_RETURN_ERROR_ON(input->tensor_shape()[idx_batch] % (block_shape_x * block_shape_y) != 0);
7281

82+
const TensorShape expected_output_shape = compute_batch_to_space_shape(input->data_layout(), input->tensor_shape(),
83+
block_shape_x, block_shape_y, crop_info);
84+
7385
// Validate output if initialized
7486
if (output->total_size() != 0)
7587
{
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);
88+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(output);
7889
const TensorInfo expected_output = output->clone()->set_tensor_shape(expected_output_shape);
7990
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(output, &expected_output);
8091
ARM_COMPUTE_RETURN_ERROR_ON(output->num_dimensions() > 4);
8192
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(input, output);
8293
}
94+
else
95+
{
96+
const auto output_info = TensorInfo(expected_output_shape, ITensorInfo::one_channel, input->data_type());
97+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output_info);
98+
}
8399

84100
return Status{};
85101
}

src/core/CL/kernels/CLBitwiseKernel.cpp

Lines changed: 9 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,20 @@ 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+
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input1, ITensorInfo::one_channel, DataType::U8);
5355
if (op != BitwiseOperation::NOT)
5456
{
5557
ARM_COMPUTE_ERROR_ON_NULLPTR(input2);
56-
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input2, 1, DataType::U8);
58+
ARM_COMPUTE_ERROR_ON_SIZE_UNSUPPORTED(input2->info());
59+
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(input2, ITensorInfo::one_channel, DataType::U8);
5760
}
5861
ARM_COMPUTE_ERROR_ON_NULLPTR(output);
59-
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, 1, DataType::U8);
62+
ARM_COMPUTE_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(output, ITensorInfo::one_channel, DataType::U8);
6063

61-
// Output auto inizialitation if not yet initialized
64+
// Output auto initialization if not yet initialized
6265
auto_init_if_empty(*(output->info()), *(input1->info()));
66+
ARM_COMPUTE_ERROR_ON_SIZE_UNSUPPORTED(output->info());
6367
auto padding_info = get_padding_info({input1, input2, output});
6468

6569
// Configure kernel window

src/core/CL/kernels/CLBoundingBoxTransformKernel.cpp

Lines changed: 14 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,12 @@ 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+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(boxes, ITensorInfo::one_channel, DataType::QASYMM16,
54+
DataType::F32, DataType::F16);
55+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(deltas, ITensorInfo::one_channel, DataType::QASYMM8,
56+
DataType::F32, DataType::F16);
5357
ARM_COMPUTE_RETURN_ERROR_ON(deltas->tensor_shape()[1] != boxes->tensor_shape()[1]);
5458
ARM_COMPUTE_RETURN_ERROR_ON(deltas->tensor_shape()[0] % 4 != 0);
5559
ARM_COMPUTE_RETURN_ERROR_ON(boxes->tensor_shape()[0] != 4);
@@ -62,7 +66,7 @@ Status validate_arguments(const ITensorInfo *boxes,
6266
const UniformQuantizationInfo boxes_qinfo = boxes->quantization_info().uniform();
6367
ARM_COMPUTE_RETURN_ERROR_ON(boxes_qinfo.scale != 0.125f);
6468
ARM_COMPUTE_RETURN_ERROR_ON(boxes_qinfo.offset != 0);
65-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_NOT_IN(deltas, DataType::QASYMM8);
69+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(deltas, ITensorInfo::one_channel, DataType::QASYMM8);
6670
}
6771
else
6872
{
@@ -71,6 +75,7 @@ Status validate_arguments(const ITensorInfo *boxes,
7175

7276
if (pred_boxes->total_size() > 0)
7377
{
78+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(pred_boxes);
7479
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DIMENSIONS(pred_boxes->tensor_shape(), deltas->tensor_shape());
7580
ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(pred_boxes, boxes);
7681
ARM_COMPUTE_RETURN_ERROR_ON(pred_boxes->num_dimensions() > 2);
@@ -81,6 +86,11 @@ Status validate_arguments(const ITensorInfo *boxes,
8186
ARM_COMPUTE_RETURN_ERROR_ON(pred_boxes_qinfo.offset != 0);
8287
}
8388
}
89+
else
90+
{
91+
const auto pred_boxes_info = TensorInfo(deltas->tensor_shape(), ITensorInfo::one_channel, boxes->data_type());
92+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&pred_boxes_info);
93+
}
8494
ARM_COMPUTE_RETURN_ERROR_ON(info.scale() <= 0);
8595

8696
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: 11 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);
@@ -62,10 +65,16 @@ Status validate_arguments(const ITensorInfo &input1,
6265
// Validate in case of configured output
6366
if (output.total_size() > 0)
6467
{
65-
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(&output, 1, DataType::U8);
68+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output);
69+
ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(&output, ITensorInfo::one_channel, DataType::U8);
6670
ARM_COMPUTE_RETURN_ERROR_ON_MSG(detail::have_different_dimensions(out_shape, output.tensor_shape(), 0),
6771
"Wrong shape for output");
6872
}
73+
else
74+
{
75+
const auto output_info = TensorInfo(out_shape, ITensorInfo::one_channel, DataType::U8);
76+
ARM_COMPUTE_RETURN_ERROR_ON_SIZE_UNSUPPORTED(&output_info);
77+
}
6978

7079
return Status{};
7180
}

0 commit comments

Comments
 (0)