Simplify autoformat facilities in preparation for removal

models/experimental/vgg/vgg_utils.py

@@ -16,7 +16,7 @@ def format_tensor(x, target_layout, device, output_mem_config, pad_value=0.0):
         return x
 
     if x.get_layout() == ttnn.ROW_MAJOR_LAYOUT and target_layout == ttnn.TILE_LAYOUT:
-        x_padded_shape = ttnn.pad_to_tile_shape(x.padded_shape, False, False, True, True)
+        x_padded_shape = ttnn.pad_to_tile_shape(x.padded_shape)
         if x.padded_shape != x_padded_shape:
             return ttnn.format_input_tensor(x, device, x_padded_shape, pad_value, target_layout, output_mem_config)
         else:

tests/tt_eager/ops/test_average_pool.cpp

@@ -17,7 +17,7 @@ using tt::tt_metal::Tensor;
 Tensor run_avg_pool_2d_resnet(ttnn::Shape& tensor_shape, IDevice* device) {
     using ttnn::operations::experimental::auto_format::AutoFormat;
     auto input_tensor = ttnn::random::random(tensor_shape, DataType::BFLOAT16);
-    auto padded_input_shape = AutoFormat::pad_to_tile_shape(tensor_shape, false, false);
+    auto padded_input_shape = AutoFormat::pad_to_tile_shape(tensor_shape);
     Tensor padded_input_tensor = input_tensor;
     if (!AutoFormat::check_input_tensor_format(input_tensor, padded_input_shape)) {
         padded_input_tensor =

ttnn/cpp/pybind11/device.cpp

@@ -493,29 +493,17 @@ void device_module(py::module& m_device) {
 
     m_device.def(
         "pad_to_tile_shape",
-        [](const std::array<uint32_t, 4>& unpadded_shape,
-           bool pad_c = false,
-           bool pad_n = false,
-           bool pad_h = true,
-           bool pad_w = true) -> std::vector<uint32_t> {
-            auto result = ttnn::operations::experimental::auto_format::AutoFormat::pad_to_tile_shape(
-                ttnn::Shape(unpadded_shape), pad_c, pad_n, pad_h, pad_w);
+        [](const std::array<uint32_t, 4>& unpadded_shape) -> std::vector<uint32_t> {
+            auto result =
+                ttnn::operations::experimental::auto_format::AutoFormat::pad_to_tile_shape(ttnn::Shape(unpadded_shape));
             return std::vector<uint32_t>(result.cbegin(), result.cend());
         },
         py::arg("unpadded_shape"),
-        py::arg("pad_c") = false,
-        py::arg("pad_n") = false,
-        py::arg("pad_h") = true,
-        py::arg("pad_w") = true,
         R"doc(
         Pads the given shape to tile shape based on specified padding options.
 
         Args:
             unpadded_shape (List of [int]): The original shape of the tensor to pad.
-            pad_c (bool, optional): Pad the channel dimension. Defaults to `False`.
-            pad_n (bool, optional): Pad the batch dimension. Defaults to `False`.
-            pad_h (bool, optional): Pad the height dimension. Defaults to `True`.
-            pad_w (bool, optional): Pad the width dimension. Defaults to `True`.
 
         Returns:
             List of [int]: The padded shape.
@@ -524,7 +512,7 @@ void device_module(py::module& m_device) {
             This functionality is planned for deprecation in the future.
 
         Example:
-            >>> padded_shape = ttnn.pad_to_tile_shape(unpadded_shape=[1, 2, 2, 2], pad_c=False, pad_n=False, pad_h=True, pad_w=True)
+            >>> padded_shape = ttnn.pad_to_tile_shape(unpadded_shape=[1, 2, 2, 2])
 
         )doc");
 

ttnn/cpp/ttnn/operations/data_movement/bcast/bcast.cpp

@@ -74,8 +74,7 @@ Tensor BcastOperation::invoke(
                 {input_tensor_a, input_tensor_b},
                 {},
                 {output_tensor},
-                0,     /* pad_value*/
-                false, /*pad_c*/
+                0, /* pad_value*/
                 queue_id);
         },
         {input_tensor_a, input_tensor_b},

ttnn/cpp/ttnn/operations/eltwise/unary_backward/unary_backward.cpp

@@ -1759,8 +1759,8 @@ std::vector<Tensor> ExecuteUnaryBackwardRepeat::invoke(
 Tensor change_layout_to_tile(const Tensor& temp, const MemoryConfig& output_mem_config) {
     auto formatted_input_tensor = temp;
     if (formatted_input_tensor.get_layout() == Layout::ROW_MAJOR) {
-        auto a_pad_shape = ttnn::operations::experimental::auto_format::AutoFormat::pad_to_tile_shape(
-            temp.get_padded_shape(), false, false, true, true);
+        auto a_pad_shape =
+            ttnn::operations::experimental::auto_format::AutoFormat::pad_to_tile_shape(temp.get_padded_shape());
         if (!ttnn::operations::experimental::auto_format::AutoFormat::check_input_tensor_format(temp, a_pad_shape)) {
             formatted_input_tensor = ttnn::operations::experimental::auto_format::AutoFormat::format_input_tensor(
                 temp, temp.device(), a_pad_shape, 1.0, Layout::TILE);

ttnn/cpp/ttnn/operations/experimental/auto_format/auto_format.cpp

@@ -18,6 +18,35 @@
 
 namespace ttnn::operations::experimental::auto_format {
 
+namespace {
+namespace CMAKE_UNIQUE_NAMESPACE {
+
+/**
+ * Checks if a shape is legal for tile layout
+ * @param shape Shape to check
+ * @return True if shape dimensions are properly aligned for tile layout
+ */
+bool legal_tile_shape(const ttnn::Shape& shape) {
+    return (shape[2] % tt::constants::TILE_HEIGHT == 0 && shape[3] % tt::constants::TILE_WIDTH == 0);
+}
+
+/**
+ * Checks if a shape is legal for a specific device layout
+ * @param shape Shape to check
+ * @param layout Target layout
+ * @return True if shape is legal for the specified layout
+ */
+bool legal_device_shape(const ttnn::Shape& shape, tt::tt_metal::Layout layout) {
+    switch (layout) {
+        case tt::tt_metal::Layout::ROW_MAJOR: return true;
+        case tt::tt_metal::Layout::TILE: return legal_tile_shape(shape);
+        default: return true;
+    }
+}
+
+}  // namespace CMAKE_UNIQUE_NAMESPACE
+}  // anonymous namespace
+
 Tensor AutoFormat::move_tensor_to_device(const Tensor& input, IDevice* device, const MemoryConfig& mem_config) {
     if (input.storage_type() != StorageType::DEVICE) {
         return input.to_device(device, mem_config);
@@ -94,7 +123,7 @@ Tensor AutoFormat::format_input_tensor(
                     padded_shape.to_array_4D(),
                     tt::tt_metal::Array4D({0, 0, 0, 0}),
                     pad_value,
-                    false,
+                    /* multicore */ false,
                     mem_config);
             }
         } else if (convert_layout && pad_input) {
@@ -115,7 +144,7 @@ Tensor AutoFormat::format_input_tensor(
                     padded_shape.to_array_4D(),
                     tt::tt_metal::Array4D({0, 0, 0, 0}),
                     pad_value,
-                    false,
+                    /* multicore */ false,
                     mem_config);
             }
         }
@@ -168,7 +197,7 @@ Tensor AutoFormat::format_output_tensor(
         if (!unpad_output && convert_layout) {
             // If target layout is tile but shape does not support tile, we don't do any conversions
             if (target_layout == Layout::TILE && formatted_output.get_layout() == Layout::ROW_MAJOR) {
-                if (AutoFormat::legal_tile_shape(formatted_output.get_padded_shape())) {
+                if (CMAKE_UNIQUE_NAMESPACE::legal_tile_shape(formatted_output.get_padded_shape())) {
                     formatted_output = ttnn::tilize(formatted_output, mem_config);
                 }
                 return formatted_output;
@@ -179,33 +208,32 @@ Tensor AutoFormat::format_output_tensor(
 
         } else if (unpad_output && !convert_layout) {
             // Output can be unpadded and layout supports the shape
-            if ((formatted_output.get_layout() == Layout::TILE && AutoFormat::legal_tile_shape(shape)) ||
-                (formatted_output.get_layout() == Layout::ROW_MAJOR && AutoFormat::legal_rm_shape(shape))) {
+            if ((formatted_output.get_layout() == Layout::TILE && CMAKE_UNIQUE_NAMESPACE::legal_tile_shape(shape)) ||
+                (formatted_output.get_layout() == Layout::ROW_MAJOR)) {
                 auto begins = std::array<uint32_t, 4>({0, 0, 0, 0});
                 auto ends = std::array<uint32_t, 4>({shape[0], shape[1], shape[2], shape[3]});
                 auto step = std::array<uint32_t, 4>({1, 1, 1, 1});
 
                 formatted_output = ttnn::slice(formatted_output, begins, ends, step, mem_config);
                 return formatted_output;
                 // Output is tile but shape cannot be tile. We leave in RM
-            } else if (formatted_output.get_layout() == Layout::TILE && AutoFormat::legal_rm_shape(shape)) {
+            } else if (formatted_output.get_layout() == Layout::TILE) {
                 formatted_output = ttnn::untilize_with_unpadding(
                     formatted_output,
                     ttnn::Shape({shape[0] - 1, shape[1] - 1, shape[2] - 1, shape[3] - 1}),
                     mem_config);
                 return formatted_output;
             }
         } else if (unpad_output && convert_layout) {
-            if (formatted_output.get_layout() == Layout::TILE && target_layout == Layout::ROW_MAJOR &&
-                AutoFormat::legal_rm_shape(shape)) {
+            if (formatted_output.get_layout() == Layout::TILE && target_layout == Layout::ROW_MAJOR) {
                 formatted_output = ttnn::untilize_with_unpadding(
                     formatted_output,
                     ttnn::Shape({shape[0] - 1, shape[1] - 1, shape[2] - 1, shape[3] - 1}),
                     mem_config);
                 return formatted_output;
             } else if (
                 formatted_output.get_layout() == Layout::ROW_MAJOR && target_layout == Layout::TILE &&
-                AutoFormat::legal_tile_shape(shape)) {
+                CMAKE_UNIQUE_NAMESPACE::legal_tile_shape(shape)) {
                 auto begins = std::array<uint32_t, 4>({0, 0, 0, 0});
                 auto ends = std::array<uint32_t, 4>({shape[0], shape[1], shape[2], shape[3]});
                 auto step = std::array<uint32_t, 4>({1, 1, 1, 1});
@@ -233,7 +261,8 @@ Tensor AutoFormat::format_output_tensor(
 
     if (convert_layout) {
         // Default to RM layout if we can't match the formatted_input layout
-        if (target_layout == Layout::TILE && !AutoFormat::legal_tile_shape(formatted_output.get_padded_shape())) {
+        if (target_layout == Layout::TILE &&
+            !CMAKE_UNIQUE_NAMESPACE::legal_tile_shape(formatted_output.get_padded_shape())) {
             if (formatted_output.get_layout() != Layout::ROW_MAJOR) {
                 formatted_output = formatted_output.to_layout(Layout::ROW_MAJOR);
             }
@@ -245,12 +274,46 @@ Tensor AutoFormat::format_output_tensor(
     // Send formatted_output to device if possible
     // Check that shape is supported on device
     if (formatted_output.storage_type() != StorageType::DEVICE) {
-        if (AutoFormat::legal_device_shape(formatted_output.get_padded_shape(), formatted_output.get_layout())) {
+        if (CMAKE_UNIQUE_NAMESPACE::legal_device_shape(
+                formatted_output.get_padded_shape(), formatted_output.get_layout())) {
             formatted_output = AutoFormat::move_tensor_to_device(formatted_output, device, mem_config);
         }
     }
 
     return formatted_output;
 }
 
+void AutoFormat::SetDefaultDevice(tt::tt_metal::IDevice* dev) { device = dev; }
+
+tt::tt_metal::IDevice* AutoFormat::GetDefaultDevice() { return device; }
+
+ttnn::Shape AutoFormat::pad_to_tile_shape(const ttnn::Shape& unpadded_shape) {
+    using namespace tt::constants;
+    auto rank = unpadded_shape.rank();
+    TT_ASSERT(rank >= 1, "rank of shape to pad to tile shape must be at least 1.");
+    SmallVector<uint32_t> padded_shape_vec(rank);
+
+    for (auto i = 0; i < rank; ++i) {
+        padded_shape_vec[i] = unpadded_shape[i];
+    }
+    if (rank >= 1) {
+        auto w = tt::round_up(unpadded_shape[rank - 1], TILE_WIDTH);
+        padded_shape_vec[rank - 1] = w;
+    }
+    if (rank >= 2) {
+        auto h = tt::round_up(unpadded_shape[rank - 2], TILE_HEIGHT);
+        padded_shape_vec[rank - 2] = h;
+    }
+    return Shape(padded_shape_vec);
+}
+
+bool AutoFormat::check_input_tensor_format(
+    const Tensor& a, const ttnn::Shape& shape, tt::tt_metal::Layout target_layout) {
+    if (a.get_layout() == target_layout && a.get_padded_shape() == shape &&
+        a.storage_type() == tt::tt_metal::StorageType::DEVICE) {
+        return true;
+    }
+    return false;
+}
+
 }  // namespace ttnn::operations::experimental::auto_format