Add function to stringify tensor shapes (#7943)

First step toward fixing #7902

Author: swolchok

kernels/portable/cpu/util/broadcast_util.cpp

@@ -7,9 +7,9 @@
  */
 
 #include <executorch/kernels/portable/cpu/util/repeat_util.h>
+#include <executorch/kernels/portable/cpu/util/tensor_util.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
-#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
 #include <string.h>
 
 namespace torch {

kernels/portable/cpu/util/targets.bzl

@@ -73,6 +73,7 @@ def define_common_targets():
         compiler_flags = ["-Wno-missing-prototypes"],
         deps = [
             ":repeat_util",
+            ":tensor_util",
             "//executorch/runtime/kernel:kernel_includes",
             "//executorch/runtime/core/exec_aten/util:tensor_util",
         ],
@@ -267,6 +268,16 @@ def define_common_targets():
         visibility = ["//executorch/kernels/portable/cpu/..."],
     )
 
+    runtime.cxx_library(
+        name = "tensor_util",
+        srcs = ["tensor_util.cpp"],
+        exported_headers = ["tensor_util.h"],
+        deps = [
+            "//executorch/runtime/kernel:kernel_includes",
+        ],
+        visibility = ["//executorch/kernels/portable/cpu/..."],
+    )
+
     runtime.cxx_library(
         name = "upsample_util",
         srcs = ["upsample_util.cpp"],

kernels/portable/cpu/util/tensor_util.cpp

@@ -0,0 +1,54 @@
+#include "tensor_util.h"
+
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+
+#include <executorch/runtime/platform/assert.h>
+
+namespace executorch::runtime {
+/**
+ * Shared implementation for tensor_util.h, may only contain code that
+ * works whether or not ATen mode is active.
+ */
+std::array<char, kTensorShapeStringSizeLimit> tensor_shape_to_c_string(
+    executorch::runtime::Span<const executorch::aten::SizesType> shape) {
+  std::array<char, kTensorShapeStringSizeLimit> out;
+  char* p = out.data();
+  if ET_UNLIKELY (shape.size() > kTensorDimensionLimit) {
+    static constexpr char kLimitExceededError[] =
+        "(ERR: tensor ndim exceeds limit)";
+    static_assert(sizeof(kLimitExceededError) <= kTensorShapeStringSizeLimit);
+    std::strcpy(p, kLimitExceededError);
+    return out;
+  }
+  *p++ = '(';
+  for (const auto elem : shape) {
+    if (elem < 0 || elem > internal::kMaximumPrintableTensorShapeElement) {
+      static_assert(
+          internal::kMaximumPrintableTensorShapeElement > 99999,
+          "must have room for error string!");
+      strcpy(p, "ERR, ");
+      p += strlen("ERR, ");
+    } else {
+      // snprintf returns characters *except* the NUL terminator, which is what
+      // we want.
+      p += snprintf(
+          p,
+          kTensorShapeStringSizeLimit - (p - out.data()),
+          "%" PRIu32 ", ",
+          static_cast<uint32_t>(elem));
+    }
+  }
+  *(p - 2) = ')';
+  *(p - 1) = '\0';
+  return out;
+}
+
+} // namespace executorch::runtime

kernels/portable/cpu/util/tensor_util.h

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <array>
+#include <cstddef>
+#include <limits>
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
+#include <executorch/runtime/core/span.h>
+
+namespace executorch::runtime {
+
+/**
+ * Maximum size of a string returned by tensor_shape_to_c_string, for
+ * stack allocation.
+ */
+constexpr size_t kTensorShapeStringSizeLimit = 1 + /* opening parenthesis */
+    10 * kTensorDimensionLimit + /* maximum digits we will print; update
+                                  * kMaximumPrintableTensorShapeElement
+                                  * if changing */
+    2 * kTensorDimensionLimit + /* comma and space after each item,
+                                 * overwritten with closing paren and
+                                 * NUL terminator for last element */
+    1; /* padding for temporary NUL terminator for simplicity of implementation
+        */
+
+namespace internal {
+constexpr size_t kMaximumPrintableTensorShapeElement =
+    std::is_same_v<executorch::aten::SizesType, int32_t>
+    ? std::numeric_limits<int32_t>::max()
+    : std::numeric_limits<uint32_t>::max();
+} // namespace internal
+
+/**
+ * Convert a shape to a NUL-terminated C string with limited size. If
+ * elements of the shape are larger than
+ * kMaximumPrintableTensorShapeElement, those elements will be
+ * rendered as ERR instead.
+ */
+std::array<char, kTensorShapeStringSizeLimit> tensor_shape_to_c_string(
+    executorch::runtime::Span<const executorch::aten::SizesType> shape);
+
+} // namespace executorch::runtime

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

@@ -19,7 +19,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..)
 
 include(${EXECUTORCH_ROOT}/build/Test.cmake)
 
-set(_test_srcs broadcast_test.cpp reduce_test.cpp)
+set(_test_srcs broadcast_test.cpp reduce_test.cpp tensor_util_test.cpp)
 
 et_cxx_test(
   kernels_portable_cpu_util_test SOURCES ${_test_srcs} EXTRA_LIBS

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

@@ -21,3 +21,11 @@ def define_common_targets():
             "//executorch/kernels/portable/cpu/util:reduce_util",
         ],
     )
+
+    runtime.cxx_test(
+        name = "tensor_util_test",
+        srcs = ["tensor_util_test.cpp"],
+        deps = [
+            "//executorch/kernels/portable/cpu/util:tensor_util",
+        ],
+    )

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

@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#include <gtest/gtest.h>
+
+#include <executorch/kernels/portable/cpu/util/tensor_util.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/platform/runtime.h>
+#include <array>
+
+using executorch::runtime::kTensorDimensionLimit;
+using executorch::runtime::Span;
+using executorch::runtime::tensor_shape_to_c_string;
+using executorch::runtime::internal::kMaximumPrintableTensorShapeElement;
+
+TEST(TensorUtilTest, TensorShapeToCStringBasic) {
+  std::array<executorch::aten::SizesType, 3> sizes = {123, 456, 789};
+  auto str = tensor_shape_to_c_string(
+      Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
+  EXPECT_STREQ(str.data(), "(123, 456, 789)");
+
+  std::array<executorch::aten::SizesType, 1> one_size = {1234567890};
+  str = tensor_shape_to_c_string(Span<const executorch::aten::SizesType>(
+      one_size.data(), one_size.size()));
+  EXPECT_STREQ(str.data(), "(1234567890)");
+}
+
+TEST(TensorUtilTest, TensorShapeToCStringNegativeItems) {
+  std::array<executorch::aten::SizesType, 4> sizes = {-1, -3, -2, 4};
+  auto str = tensor_shape_to_c_string(
+      Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
+  EXPECT_STREQ(str.data(), "(ERR, ERR, ERR, 4)");
+
+  std::array<executorch::aten::SizesType, 1> one_size = {-1234567890};
+  str = tensor_shape_to_c_string(Span<const executorch::aten::SizesType>(
+      one_size.data(), one_size.size()));
+  if constexpr (std::numeric_limits<executorch::aten::SizesType>::is_signed) {
+    EXPECT_STREQ(str.data(), "(ERR)");
+  } else {
+    EXPECT_EQ(str.data(), "(" + std::to_string(one_size[0]) + ")");
+  }
+}
+TEST(TensorUtilTest, TensorShapeToCStringMaximumElement) {
+  std::array<executorch::aten::SizesType, 3> sizes = {
+      123, std::numeric_limits<executorch::aten::SizesType>::max(), 789};
+  auto str = tensor_shape_to_c_string(
+      Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
+  std::ostringstream expected;
+  expected << '(';
+  for (const auto elem : sizes) {
+    expected << elem << ", ";
+  }
+  auto expected_str = expected.str();
+  expected_str.pop_back();
+  expected_str.back() = ')';
+  EXPECT_EQ(str.data(), expected_str);
+}
+
+TEST(TensorUtilTest, TensorShapeToCStringMaximumLength) {
+  std::array<executorch::aten::SizesType, kTensorDimensionLimit> sizes;
+  std::fill(sizes.begin(), sizes.end(), kMaximumPrintableTensorShapeElement);
+
+  auto str = tensor_shape_to_c_string(
+      Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
+
+  std::ostringstream expected;
+  expected << '(' << kMaximumPrintableTensorShapeElement;
+  for (int ii = 0; ii < kTensorDimensionLimit - 1; ++ii) {
+    expected << ", " << kMaximumPrintableTensorShapeElement;
+  }
+  expected << ')';
+  auto expected_str = expected.str();
+
+  EXPECT_EQ(expected_str, str.data());
+}
+
+TEST(TensorUtilTest, TensorShapeToCStringExceedsDimensionLimit) {
+  std::array<executorch::aten::SizesType, kTensorDimensionLimit + 1> sizes;
+  std::fill(sizes.begin(), sizes.end(), kMaximumPrintableTensorShapeElement);
+
+  auto str = tensor_shape_to_c_string(
+      Span<const executorch::aten::SizesType>(sizes.data(), sizes.size()));
+
+  EXPECT_STREQ(str.data(), "(ERR: tensor ndim exceeds limit)");
+}

runtime/core/exec_aten/util/tensor_util.h

@@ -20,6 +20,7 @@
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/span.h>
 #include <executorch/runtime/platform/assert.h>
 #include <executorch/runtime/platform/compiler.h>
 

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

@@ -19,7 +19,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..)
 
 include(${EXECUTORCH_ROOT}/build/Test.cmake)
 
-set(_test_srcs tensor_util_test.cpp scalar_type_util_test.cpp
+set(_test_srcs scalar_type_util_test.cpp
                operator_impl_example_test.cpp dim_order_util_test.cpp
 )