Skip to content

Commit 4a1bf29

Browse files
swolchokswolchok
authored
Add function to stringify tensor shapes (#7943)
First step toward fixing #7902
1 parent 48c1395 commit 4a1bf29

File tree

9 files changed

+217
-3
lines changed

9 files changed

+217
-3
lines changed

‎kernels/portable/cpu/util/broadcast_util.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@
77
*/
88

99
#include <executorch/kernels/portable/cpu/util/repeat_util.h>
10+
#include <executorch/kernels/portable/cpu/util/tensor_util.h>
1011
#include <executorch/runtime/core/exec_aten/exec_aten.h>
1112
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
12-
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
1313
#include <string.h>
1414

1515
namespace torch {

‎kernels/portable/cpu/util/targets.bzl‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ def define_common_targets():
7373
compiler_flags = ["-Wno-missing-prototypes"],
7474
deps = [
7575
":repeat_util",
76+
":tensor_util",
7677
"//executorch/runtime/kernel:kernel_includes",
7778
"//executorch/runtime/core/exec_aten/util:tensor_util",
7879
],
@@ -267,6 +268,16 @@ def define_common_targets():
267268
visibility = ["//executorch/kernels/portable/cpu/..."],
268269
)
269270

271+
runtime.cxx_library(
272+
name = "tensor_util",
273+
srcs = ["tensor_util.cpp"],
274+
exported_headers = ["tensor_util.h"],
275+
deps = [
276+
"//executorch/runtime/kernel:kernel_includes",
277+
],
278+
visibility = ["//executorch/kernels/portable/cpu/..."],
279+
)
280+
270281
runtime.cxx_library(
271282
name = "upsample_util",
272283
srcs = ["upsample_util.cpp"],

‎kernels/portable/cpu/util/tensor_util.cpp‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#include "tensor_util.h"
2+
3+
/*
4+
* Copyright (c) Meta Platforms, Inc. and affiliates.
5+
* All rights reserved.
6+
*
7+
* This source code is licensed under the BSD-style license found in the
8+
* LICENSE file in the root directory of this source tree.
9+
*/
10+
11+
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
12+
13+
#include <executorch/runtime/platform/assert.h>
14+
15+
namespace executorch::runtime {
16+
/**
17+
* Shared implementation for tensor_util.h, may only contain code that
18+
* works whether or not ATen mode is active.
19+
*/
20+
std::array<char, kTensorShapeStringSizeLimit> tensor_shape_to_c_string(
21+
executorch::runtime::Span<const executorch::aten::SizesType> shape) {
22+
std::array<char, kTensorShapeStringSizeLimit> out;
23+
char* p = out.data();
24+
if ET_UNLIKELY (shape.size() > kTensorDimensionLimit) {
25+
static constexpr char kLimitExceededError[] =
26+
"(ERR: tensor ndim exceeds limit)";
27+
static_assert(sizeof(kLimitExceededError) <= kTensorShapeStringSizeLimit);
28+
std::strcpy(p, kLimitExceededError);
29+
return out;
30+
}
31+
*p++ = '(';
32+
for (const auto elem : shape) {
33+
if (elem < 0 || elem > internal::kMaximumPrintableTensorShapeElement) {
34+
static_assert(
35+
internal::kMaximumPrintableTensorShapeElement > 99999,
36+
"must have room for error string!");
37+
strcpy(p, "ERR, ");
38+
p += strlen("ERR, ");
39+
} else {
40+
// snprintf returns characters *except* the NUL terminator, which is what
41+
// we want.
42+
p += snprintf(
43+
p,
44+
kTensorShapeStringSizeLimit - (p - out.data()),
45+
"%" PRIu32 ", ",
46+
static_cast<uint32_t>(elem));
47+
}
48+
}
49+
*(p - 2) = ')';
50+
*(p - 1) = '\0';
51+
return out;
52+
}
53+
54+
} // namespace executorch::runtime

‎kernels/portable/cpu/util/tensor_util.h‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <array>
12+
#include <cstddef>
13+
#include <limits>
14+
15+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
16+
#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
17+
#include <executorch/runtime/core/span.h>
18+
19+
namespace executorch::runtime {
20+
21+
/**
22+
* Maximum size of a string returned by tensor_shape_to_c_string, for
23+
* stack allocation.
24+
*/
25+
constexpr size_t kTensorShapeStringSizeLimit = 1 + /* opening parenthesis */
26+
10 * kTensorDimensionLimit + /* maximum digits we will print; update
27+
* kMaximumPrintableTensorShapeElement
28+
* if changing */
29+
2 * kTensorDimensionLimit + /* comma and space after each item,
30+
* overwritten with closing paren and
31+
* NUL terminator for last element */
32+
1; /* padding for temporary NUL terminator for simplicity of implementation
33+
*/
34+
35+
namespace internal {
36+
constexpr size_t kMaximumPrintableTensorShapeElement =
37+
std::is_same_v<executorch::aten::SizesType, int32_t>
38+
? std::numeric_limits<int32_t>::max()
39+
: std::numeric_limits<uint32_t>::max();
40+
} // namespace internal
41+
42+
/**
43+
* Convert a shape to a NUL-terminated C string with limited size. If
44+
* elements of the shape are larger than
45+
* kMaximumPrintableTensorShapeElement, those elements will be
46+
* rendered as ERR instead.
47+
*/
48+
std::array<char, kTensorShapeStringSizeLimit> tensor_shape_to_c_string(
49+
executorch::runtime::Span<const executorch::aten::SizesType> shape);
50+
51+
} // namespace executorch::runtime

‎kernels/portable/cpu/util/test/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..)
1919

2020
include(${EXECUTORCH_ROOT}/build/Test.cmake)
2121

22-
set(_test_srcs broadcast_test.cpp reduce_test.cpp)
22+
set(_test_srcs broadcast_test.cpp reduce_test.cpp tensor_util_test.cpp)
2323

2424
et_cxx_test(
2525
kernels_portable_cpu_util_test SOURCES ${_test_srcs} EXTRA_LIBS

‎kernels/portable/cpu/util/test/targets.bzl‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,11 @@ def define_common_targets():
2121
"//executorch/kernels/portable/cpu/util:reduce_util",
2222
],
2323
)
24+
25+
runtime.cxx_test(
26+
name = "tensor_util_test",
27+
srcs = ["tensor_util_test.cpp"],
28+
deps = [
29+
"//executorch/kernels/portable/cpu/util:tensor_util",
30+
],
31+
)

‎kernels/portable/cpu/util/test/tensor_util_test.cpp‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
#include <gtest/gtest.h>
9+
10+
#include <executorch/kernels/portable/cpu/util/tensor_util.h>
11+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
12+
#include <executorch/runtime/platform/runtime.h>
13+
#include <array>
14+
15+
using executorch::runtime::kTensorDimensionLimit;
16+
using executorch::runtime::Span;
17+
using executorch::runtime::tensor_shape_to_c_string;
18+
using executorch::runtime::internal::kMaximumPrintableTensorShapeElement;
19+
20+
TEST(TensorUtilTest, TensorShapeToCStringBasic) {
21+
std::array<executorch::aten::SizesType, 3> sizes = {123, 456, 789};
22+
auto str = tensor_shape_to_c_string(
23+
Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
24+
EXPECT_STREQ(str.data(), "(123, 456, 789)");
25+
26+
std::array<executorch::aten::SizesType, 1> one_size = {1234567890};
27+
str = tensor_shape_to_c_string(Span<const executorch::aten::SizesType>(
28+
one_size.data(), one_size.size()));
29+
EXPECT_STREQ(str.data(), "(1234567890)");
30+
}
31+
32+
TEST(TensorUtilTest, TensorShapeToCStringNegativeItems) {
33+
std::array<executorch::aten::SizesType, 4> sizes = {-1, -3, -2, 4};
34+
auto str = tensor_shape_to_c_string(
35+
Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
36+
EXPECT_STREQ(str.data(), "(ERR, ERR, ERR, 4)");
37+
38+
std::array<executorch::aten::SizesType, 1> one_size = {-1234567890};
39+
str = tensor_shape_to_c_string(Span<const executorch::aten::SizesType>(
40+
one_size.data(), one_size.size()));
41+
if constexpr (std::numeric_limits<executorch::aten::SizesType>::is_signed) {
42+
EXPECT_STREQ(str.data(), "(ERR)");
43+
} else {
44+
EXPECT_EQ(str.data(), "(" + std::to_string(one_size[0]) + ")");
45+
}
46+
}
47+
TEST(TensorUtilTest, TensorShapeToCStringMaximumElement) {
48+
std::array<executorch::aten::SizesType, 3> sizes = {
49+
123, std::numeric_limits<executorch::aten::SizesType>::max(), 789};
50+
auto str = tensor_shape_to_c_string(
51+
Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
52+
std::ostringstream expected;
53+
expected << '(';
54+
for (const auto elem : sizes) {
55+
expected << elem << ", ";
56+
}
57+
auto expected_str = expected.str();
58+
expected_str.pop_back();
59+
expected_str.back() = ')';
60+
EXPECT_EQ(str.data(), expected_str);
61+
}
62+
63+
TEST(TensorUtilTest, TensorShapeToCStringMaximumLength) {
64+
std::array<executorch::aten::SizesType, kTensorDimensionLimit> sizes;
65+
std::fill(sizes.begin(), sizes.end(), kMaximumPrintableTensorShapeElement);
66+
67+
auto str = tensor_shape_to_c_string(
68+
Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
69+
70+
std::ostringstream expected;
71+
expected << '(' << kMaximumPrintableTensorShapeElement;
72+
for (int ii = 0; ii < kTensorDimensionLimit - 1; ++ii) {
73+
expected << ", " << kMaximumPrintableTensorShapeElement;
74+
}
75+
expected << ')';
76+
auto expected_str = expected.str();
77+
78+
EXPECT_EQ(expected_str, str.data());
79+
}
80+
81+
TEST(TensorUtilTest, TensorShapeToCStringExceedsDimensionLimit) {
82+
std::array<executorch::aten::SizesType, kTensorDimensionLimit + 1> sizes;
83+
std::fill(sizes.begin(), sizes.end(), kMaximumPrintableTensorShapeElement);
84+
85+
auto str = tensor_shape_to_c_string(
86+
Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
87+
88+
EXPECT_STREQ(str.data(), "(ERR: tensor ndim exceeds limit)");
89+
}

‎runtime/core/exec_aten/util/tensor_util.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include <executorch/runtime/core/exec_aten/exec_aten.h>
2121
#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
2222
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
23+
#include <executorch/runtime/core/span.h>
2324
#include <executorch/runtime/platform/assert.h>
2425
#include <executorch/runtime/platform/compiler.h>
2526

‎runtime/core/exec_aten/util/test/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..)
1919

2020
include(${EXECUTORCH_ROOT}/build/Test.cmake)
2121

22-
set(_test_srcs tensor_util_test.cpp scalar_type_util_test.cpp
22+
set(_test_srcs scalar_type_util_test.cpp
2323
operator_impl_example_test.cpp dim_order_util_test.cpp
2424
)
2525

0 commit comments

Comments
 (0)