Add function to stringify tensor shapes

First step toward fixing #7902

ghstack-source-id: d869c4e
ghstack-comment-id: 2613189307
Pull Request resolved: #7943

Author: swolchok

runtime/core/exec_aten/util/targets.bzl

@@ -50,7 +50,7 @@ def define_common_targets():
 
         runtime.cxx_library(
             name = "tensor_util" + aten_suffix,
-            srcs = ["tensor_util_aten.cpp"] if aten_mode else ["tensor_util_portable.cpp"],
+            srcs = ["tensor_util.cpp"] + (["tensor_util_aten.cpp"] if aten_mode else ["tensor_util_portable.cpp"]),
             exported_headers = [
                 "tensor_util.h",
             ],

runtime/core/exec_aten/util/tensor_util.cpp

@@ -0,0 +1,44 @@
+#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::aten::ArrayRef<executorch::aten::SizesType> shape) {
+  std::array<char, kTensorShapeStringSizeLimit> out;
+  char* p = out.data();
+  *p++ = '(';
+  for (const auto elem : shape) {
+    if (elem < 0 || elem > kMaximumPrintableTensorShapeElement) {
+      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

runtime/core/exec_aten/util/tensor_util.h

@@ -1120,6 +1120,28 @@ bool extract_scalar_tensor(executorch::aten::Tensor tensor, BOOL_T* out_val) {
   return true;
 }
 
+/**
+ * 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
+        */
+
+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();
+
+std::array<char, kTensorShapeStringSizeLimit> tensor_shape_to_c_string(
+    executorch::aten::ArrayRef<executorch::aten::SizesType> shape);
+
 /// These APIs should not be used outside of Executor.cpp.
 namespace internal {
 /**

runtime/core/exec_aten/util/test/tensor_util_test.cpp

@@ -12,12 +12,17 @@
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/test/utils/DeathTest.h>
 #include <cmath>
+#include <cstring>
 #include <limits>
 
 using namespace ::testing;
 using executorch::aten::ScalarType;
 using executorch::aten::Tensor;
 using executorch::runtime::extract_scalar_tensor;
+using executorch::runtime::kMaximumPrintableTensorShapeElement;
+using executorch::runtime::kTensorDimensionLimit;
+using executorch::runtime::kTensorShapeStringSizeLimit;
+using executorch::runtime::tensor_shape_to_c_string;
 using executorch::runtime::testing::TensorFactory;
 
 class TensorUtilTest : public ::testing::Test {
@@ -605,3 +610,36 @@ TEST_F(TensorUtilTest, SameShapesDifferentDimOrder) {
   EXPECT_FALSE(tensors_have_same_dim_order(a, c, b));
   EXPECT_FALSE(tensors_have_same_dim_order(c, b, a));
 }
+
+TEST_F(TensorUtilTest, TensorShapeToCStringBasic) {
+  std::array<executorch::aten::SizesType, 3> sizes = {123, 456, 789};
+  auto str = tensor_shape_to_c_string(sizes);
+  EXPECT_STREQ(str.data(), "(123, 456, 789)");
+
+  std::array<executorch::aten::SizesType, 1> one_size = {1234567890};
+  str = tensor_shape_to_c_string(one_size);
+  EXPECT_STREQ(str.data(), "(1234567890)");
+}
+
+TEST_F(TensorUtilTest, TensorShapeToCStringMaximumLength) {
+  std::array<
+      executorch::aten::SizesType,
+      executorch::runtime::kTensorDimensionLimit>
+      sizes;
+  std::fill(
+      sizes.begin(),
+      sizes.end(),
+      executorch::runtime::kMaximumPrintableTensorShapeElement);
+
+  auto str = executorch::runtime::tensor_shape_to_c_string(sizes);
+
+  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());
+}