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

@@ -487,29 +487,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.
@@ -518,7 +506,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

@@ -96,7 +96,7 @@ Tensor AutoFormat::format_input_tensor(
                     padded_shape.to_array_4D(),
                     tt::tt_metal::Array4D({0, 0, 0, 0}),
                     pad_value,
-                    false,
+                    false, /* multicore */
                     mem_config);
             }
         } else if (convert_layout && pad_input) {
@@ -117,7 +117,7 @@ Tensor AutoFormat::format_input_tensor(
                     padded_shape.to_array_4D(),
                     tt::tt_metal::Array4D({0, 0, 0, 0}),
                     pad_value,
-                    false,
+                    false, /* multicore */
                     mem_config);
             }
         }
@@ -182,24 +182,23 @@ 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))) {
+                (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}),
@@ -255,4 +254,49 @@ Tensor AutoFormat::format_output_tensor(
     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::legal_tile_shape(const ttnn::Shape& shape) {
+    return (shape[2] % tt::constants::TILE_HEIGHT == 0 && shape[3] % tt::constants::TILE_WIDTH == 0);
+}
+
+bool AutoFormat::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;
+    }
+}
+
+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

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

@@ -27,99 +27,102 @@ class AutoFormat {
     AutoFormat() {}
 
 public:
-    static void SetDefaultDevice(tt::tt_metal::IDevice* dev) { device = dev; }
-    static tt::tt_metal::IDevice* GetDefaultDevice() { return device; }
-
-    static ttnn::Shape pad_to_tile_shape(
-        const ttnn::Shape& unpadded_shape,
-        bool pad_c = false,
-        bool pad_n = false,
-        bool pad_h = true,
-        bool pad_w = true) {
-        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 = pad_w ? tt::round_up(unpadded_shape[rank - 1], TILE_WIDTH) : unpadded_shape[rank - 1];
-            padded_shape_vec[rank - 1] = w;
-        }
-        if (rank >= 2) {
-            auto h = pad_h ? tt::round_up(unpadded_shape[rank - 2], TILE_HEIGHT) : unpadded_shape[rank - 2];
-            padded_shape_vec[rank - 2] = h;
-        }
-        if (rank >= 3) {
-            auto c = pad_c ? tt::round_up(unpadded_shape[rank - 3], TILE_WIDTH) : unpadded_shape[rank - 3];
-            padded_shape_vec[rank - 3] = c;
-        }
-        if (rank >= 4) {
-            auto n = pad_n ? tt::round_up(unpadded_shape[rank - 4], TILE_HEIGHT) : unpadded_shape[rank - 4];
-            padded_shape_vec[rank - 4] = n;
-        }
-        return Shape(padded_shape_vec);
-    }
-
-    static ttnn::Shape pad_to_rm_shape(const ttnn::Shape& unpadded_shape) {
-        ttnn::Shape padded_shape = unpadded_shape;
-        padded_shape[3] = tt::round_up(unpadded_shape[3], 2);
-        return padded_shape;
-    }
-
-    static ttnn::Shape pad_to_legal_shape(const ttnn::Shape& unpadded_shape, tt::tt_metal::Layout layout) {
-        ttnn::Shape padded_shape = unpadded_shape;
-        switch (layout) {
-            case tt::tt_metal::Layout::ROW_MAJOR: padded_shape = pad_to_rm_shape(unpadded_shape); break;
-            case tt::tt_metal::Layout::TILE: padded_shape = pad_to_tile_shape(unpadded_shape);
-            default: break;
-        }
-        return padded_shape;
-    }
-
-    // TODO: These legal checks should probably be somewhere else like tensor class, since it is common logic not just
-    // for autoformat
-    static bool legal_tile_shape(const ttnn::Shape& shape) {
-        return (shape[2] % tt::constants::TILE_HEIGHT == 0 && shape[3] % tt::constants::TILE_WIDTH == 0);
-    }
-
-    static bool legal_rm_shape(const ttnn::Shape& shape) { return (shape[3] % 2 == 0); }
-
-    static bool legal_device_shape(const ttnn::Shape& shape, tt::tt_metal::Layout layout) {
-        switch (layout) {
-            case tt::tt_metal::Layout::ROW_MAJOR: return legal_rm_shape(shape);
-            case tt::tt_metal::Layout::TILE: return legal_tile_shape(shape);
-            default: return true;
-        }
-    }
-
+    /**
+     * Sets the default device to be used for auto-formatting operations
+     * @param dev Pointer to the device to be used
+     */
+    static void SetDefaultDevice(tt::tt_metal::IDevice* dev);
+
+    /**
+     * Gets the default device used for auto-formatting operations
+     * @return Pointer to the default device
+     */
+    static tt::tt_metal::IDevice* GetDefaultDevice();
+
+    /**
+     * Pads a shape to align with tile dimensions
+     * @param unpadded_shape Original shape to be padded
+     * @param pad_c Whether to pad the channel dimension
+     * @param pad_n Whether to pad the batch dimension
+     * @param pad_h Whether to pad the height dimension
+     * @param pad_w Whether to pad the width dimension
+     * @return Padded shape aligned to tile dimensions
+     */
+    static ttnn::Shape pad_to_tile_shape(const ttnn::Shape& unpadded_shape);
+
+    /**
+     * 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
+     */
+    static bool legal_tile_shape(const ttnn::Shape& shape);
+
+    /**
+     * 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
+     */
+    static bool legal_device_shape(const ttnn::Shape& shape, tt::tt_metal::Layout layout);
+
+    /**
+     * Checks if a tensor matches the required format specifications
+     * @param a Input tensor to check
+     * @param shape Required shape
+     * @param target_layout Required layout
+     * @return True if tensor matches all format requirements
+     */
     static bool check_input_tensor_format(
-        const Tensor& a, const ttnn::Shape& shape, tt::tt_metal::Layout target_layout = tt::tt_metal::Layout::TILE) {
-        if (a.get_layout() == target_layout && a.get_padded_shape() == shape &&
-            a.storage_type() == tt::tt_metal::StorageType::DEVICE) {
-            return true;
-        }
-        return false;
-    }
+        const Tensor& a, const ttnn::Shape& shape, tt::tt_metal::Layout target_layout = tt::tt_metal::Layout::TILE);
 
     // This code is a workaround for cases where we need to remove autoformat but other dependent ops
     // are not quite ready. So here we basically just put the tensor back on device.
     // Used in backward_ops.cpp
     // See: Remove auto format within permute_op.cpp #9404
+    /**
+     * Moves a tensor to device memory and pads if necessary
+     * @param input Input tensor
+     * @param device Target device
+     * @param target_layout Desired layout
+     * @param target_mem_config Optional memory configuration
+     * @return Formatted tensor on device
+     */
     static Tensor move_tensor_to_device_and_pad(
         const Tensor& input,
         tt::tt_metal::IDevice* device,
         tt::tt_metal::Layout target_layout,
         std::optional<tt::tt_metal::MemoryConfig> target_mem_config);
 
+    /**
+     * Moves a tensor to device memory
+     * @param input Input tensor
+     * @param device Target device
+     * @param mem_config Memory configuration
+     * @return Tensor on device
+     */
     static Tensor move_tensor_to_device(
         const Tensor& input,
         tt::tt_metal::IDevice* device,
         const tt::tt_metal::MemoryConfig& mem_config = tt::tt_metal::operation::DEFAULT_OUTPUT_MEMORY_CONFIG);
 
+    /**
+     * Updates tensor memory configuration
+     * @param input Input tensor
+     * @param mem_config Target memory configuration
+     * @return Tensor with updated memory configuration
+     */
     static Tensor move_tensor_to_mem_config(const Tensor& input, const tt::tt_metal::MemoryConfig& mem_config);
 
+    /**
+     * Formats an input tensor to meet device and layout requirements
+     * @param input Input tensor
+     * @param device Target device
+     * @param padded_shape Required padded shape
+     * @param pad_value Value to use for padding
+     * @param target_layout Desired layout
+     * @param target_mem_config Optional memory configuration
+     * @return Formatted tensor
+     */
     static Tensor format_input_tensor(
         const Tensor& input,
         tt::tt_metal::IDevice* device,
@@ -128,6 +131,15 @@ class AutoFormat {
         tt::tt_metal::Layout target_layout,
         std::optional<tt::tt_metal::MemoryConfig> target_mem_config = std::nullopt);
 
+    /**
+     * Formats an output tensor to meet shape and layout requirements
+     * @param output Output tensor
+     * @param shape Target shape
+     * @param device Target device
+     * @param target_layout Desired layout
+     * @param target_mem_config Optional memory configuration
+     * @return Formatted output tensor
+     */
     static Tensor format_output_tensor(
         const Tensor& output,
         const ttnn::Shape& shape,

ttnn/cpp/ttnn/operations/reduction/prod/prod.cpp

@@ -22,7 +22,7 @@ inline Tensor change_layout_to_tile(const Tensor& temp, const MemoryConfig& outp
     using ttnn::operations::experimental::auto_format::AutoFormat;
     auto formatted_input_tensor = temp;
     if (formatted_input_tensor.get_layout() == Layout::ROW_MAJOR) {
-        auto a_pad_shape = AutoFormat::pad_to_tile_shape(temp.get_padded_shape(), false, false, true, true);
+        auto a_pad_shape = AutoFormat::pad_to_tile_shape(temp.get_padded_shape());
         if (!AutoFormat::check_input_tensor_format(temp, a_pad_shape)) {
             formatted_input_tensor =
                 AutoFormat::format_input_tensor(temp, temp.device(), a_pad_shape, 1.0, Layout::TILE);
@@ -35,7 +35,7 @@ inline Tensor prod_all(const Tensor& input_a, const MemoryConfig& output_mem_con
     using ttnn::operations::experimental::auto_format::AutoFormat;
     auto formatted_input_tensor = input_a;
     if (formatted_input_tensor.get_layout() == Layout::ROW_MAJOR) {
-        auto a_pad_shape = AutoFormat::pad_to_tile_shape(input_a.get_padded_shape(), false, false, true, true);
+        auto a_pad_shape = AutoFormat::pad_to_tile_shape(input_a.get_padded_shape());
         auto out_shape = input_a.get_padded_shape();
         out_shape = ttnn::Shape({out_shape[0], out_shape[1], out_shape[2], out_shape[3]});
         if (!AutoFormat::check_input_tensor_format(input_a, a_pad_shape)) {
@@ -51,7 +51,7 @@ inline Tensor prod_nc(const Tensor& temp, int64_t dim, const MemoryConfig& outpu
     // layout conversion
     auto formatted_input_tensor = temp;
     if (formatted_input_tensor.get_layout() == Layout::ROW_MAJOR) {
-        auto a_pad_shape = AutoFormat::pad_to_tile_shape(temp.get_padded_shape(), false, false, true, true);
+        auto a_pad_shape = AutoFormat::pad_to_tile_shape(temp.get_padded_shape());
         auto out_shape = temp.get_padded_shape();
         out_shape = ttnn::Shape({out_shape[0], out_shape[1], out_shape[2], out_shape[3]});
         if (!AutoFormat::check_input_tensor_format(temp, a_pad_shape)) {