Handle Shapes with 0 dimensions in Tensor compute_strides (#12286)

tenstorrent · Sep 6, 2024 · cb68490 · cb68490
1 parent 73b5d39
commit cb68490
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 17 deletions.
diff --git a/ttnn/cpp/ttnn/tensor/tensor.cpp b/ttnn/cpp/ttnn/tensor/tensor.cpp
@@ -494,19 +494,7 @@ StorageType Tensor::storage_type() const {
         this->get_storage());
 }
 
-namespace detail {
-const Shape compute_strides(const Shape& shape) {
-    auto num_elements = compute_volume(shape);
-    std::vector<std::uint32_t> strides;
-    for (std::int32_t index = 0; index < shape.rank(); index++) {
-        num_elements /= shape[index];
-        strides.push_back(num_elements);
-    }
-    return strides;
-}
-}  // namespace detail
-
-const Shape Tensor::strides() const { return detail::compute_strides(this->get_legacy_shape()); }
+const Shape Tensor::strides() const { return Shape(tt::tt_metal::compute_strides(this->get_legacy_shape())); }
 
 uint32_t Tensor::volume() const { return tt::tt_metal::compute_volume(this->get_legacy_shape()); }
 

diff --git a/ttnn/cpp/ttnn/tensor/tensor_utils.hpp b/ttnn/cpp/ttnn/tensor/tensor_utils.hpp
@@ -46,17 +46,26 @@ static std::size_t compute_volume(const T& shape) {
     return volume;
 }
 
-static std::vector<std::size_t> compute_strides(Shape shape) {
+static std::vector<uint32_t> compute_strides(const Shape& shape) {
+    if (shape.rank() == 0)
+        return {};
+
     auto num_elements = compute_volume(shape);
-    std::vector<std::size_t> strides;
+    std::vector<uint32_t> strides;
     for (std::int32_t index = 0; index < shape.rank(); index++) {
+        if (shape[index] == 0) {
+            // Insert 0 to indicate no memory access for this dimension
+            strides.push_back(0);
+            continue;
+        }
+
         num_elements /= shape[index];
         strides.push_back(num_elements);
     }
     return strides;
 }
 
-static int compute_flat_indices(vector<int> indices, vector<std::size_t> strides) {
+static int compute_flat_indices(const vector<int>& indices, const vector<std::uint32_t> strides) {
     int flat_index = 0;
     for (auto i = 0; i < indices.size(); i++) {
         flat_index += indices[i] * strides[i];

diff --git a/ttnn/cpp/ttnn/tensor/types.hpp b/ttnn/cpp/ttnn/tensor/types.hpp
@@ -245,7 +245,6 @@ class Shape {
         }
         return ret_array;
     }
-
 };
 
 inline std::ostream &operator<<(std::ostream &os, const Shape &shape) {