Expose conv2d weights/biases as ops #14566

LPanosTT · 2024-11-01T15:46:37Z

Tickets

Problem description

Modelling convolutions in MLIR is difficult when you need to execute conv2d once in order to get the proper inputs for the next iteration.

What's changed

Moving weight and bias preparation to dedicated ops. The logic that prepared the ops on host already exists and is functional. The issue is that it is done within conv2d. Most of the code changes is copy/pasted.

Checklist

Post commit CI passes
Blackhole Post commit (if applicable)
Model regression CI testing passes (if applicable)
Device performance regression CI testing passes (if applicable)
New/Existing tests provide coverage for changes

github-actions

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

github-actions · 2024-11-01T15:51:01Z

ttnn/cpp/ttnn/operations/conv/conv2d/prepare_conv2d_weights.cpp

+#include <sys/types.h>
+#include <cstdint>
+
+using namespace tt;


⚠️ google-build-using-namespace ⚠️
do not use namespace using-directives; use using-declarations instead

github-actions · 2024-11-01T15:51:01Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d_utils.cpp

+#include "ttnn/cpp/ttnn/operations/data_movement/reshape_view/reshape.hpp"
+#include "ttnn/tensor/tensor.hpp"
+
+using namespace tt;


⚠️ google-build-using-namespace ⚠️
do not use namespace using-directives; use using-declarations instead

bbradelTT · 2024-11-05T14:25:01Z

ttnn/ttnn/operations/conv2d.py

@@ -104,8 +176,16 @@ def conv2d(
    memory_config: ttnn.MemoryConfig = None,  # memory config overrides by user
    conv_op_cache={},  # basic conv object caching in python needed for intermediate refactoring. Not needed after full op refactoring in C++.
    debug=False,  # ignored
+    return_output_size=False,
+    return_prepared_device_weights=False,


Why complicate the API like this?

If the values are computed then there's no additional cost to returning them.

bbradelTT · 2024-11-05T14:26:18Z

ttnn/ttnn/operations/conv2d.py

+    )
+
+
+def prepare_conv_bias(


Shouldn't new ops use the new op framework?

If you use that, then bind_registered_operation should remove the need for this extra code.

bbradelTT · 2024-11-05T14:27:04Z

ttnn/ttnn/operations/conv1d.py

@@ -30,6 +30,8 @@ def Conv1d(
    conv_config: Conv1dConfig = None,  # config overrides by user
    conv_op_cache={},  # basic conv object caching in python needed for intermediate refactoring. Not needed after full op refactoring in C++.
    debug=False,
+    return_output_length=False,
+    return_prepared_device_weights=False,


Why complicate the API like this?

If the values are computed then there's no additional cost to returning them.

Same comment below for 2d.

github-actions

⚠️ Clang-Tidy found issue(s) with the introduced code (1/2)

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.cpp

-        conv_config.act_block_h_override = constants::TILE_HEIGHT;
-    }
-}
-
 template <typename T>
 Result conv2d(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.cpp

-
-Result Conv2dOperation::invoke(
-    uint8_t queue_id,
+template std::tuple<ttnn::Tensor, uint32_t, uint32_t, ttnn::Tensor, std::optional<ttnn::Tensor>> conv2d<Device>(


⚠️ clang-diagnostic-error ⚠️
explicit template instantiation cannot have a definition; if this definition is meant to be an explicit specialization, add <> after the template keyword

Suggested change

template std::tuple<ttnn::Tensor, uint32_t, uint32_t, ttnn::Tensor, std::optional<ttnn::Tensor>> conv2d<Device>(

template<> std::tuple<ttnn::Tensor, uint32_t, uint32_t, ttnn::Tensor, std::optional<ttnn::Tensor>> conv2d<Device>(

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.cpp

-
-Result Conv2dOperation::invoke(
-    uint8_t queue_id,
+template std::tuple<ttnn::Tensor, uint32_t, uint32_t, ttnn::Tensor, std::optional<ttnn::Tensor>> conv2d<Device>(


⚠️ clang-diagnostic-error ⚠️
no function template matches function template specialization conv2d

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.hpp

-
-template <typename T>
-std::pair<ttnn::Tensor, std::optional<ttnn::Tensor>> prepare_conv_weights_biases_and_move_to_device(const ttnn::Tensor& weight_tensor, std::optional<const ttnn::Tensor>& bias_tensor, uint32_t input_channels_alignment, DataType weights_bias_dtype, uint32_t weight_block_h_ntiles, uint32_t weight_block_w_ntiles, const sliding_window::ParallelConfig& parallel_config, T * device, uint32_t groups, uint32_t act_block_h_ntiles, uint32_t input_width);
-
 template <typename T>
 Result conv2d(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.hpp

@@ -215,7 +45,6 @@
    std::optional<const Conv2dConfig> conv_config_ = std::nullopt,
    const std::optional<const MemoryConfig> memory_config = std::nullopt);

-
 struct Conv2dOperation{
    static Result invoke(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.hpp

-
-template <typename T>
-std::pair<ttnn::Tensor, std::optional<ttnn::Tensor>> prepare_conv_weights_biases_and_move_to_device(const ttnn::Tensor& weight_tensor, std::optional<const ttnn::Tensor>& bias_tensor, uint32_t input_channels_alignment, DataType weights_bias_dtype, uint32_t weight_block_h_ntiles, uint32_t weight_block_w_ntiles, const sliding_window::ParallelConfig& parallel_config, T * device, uint32_t groups, uint32_t act_block_h_ntiles, uint32_t input_width);
-
 template <typename T>
 Result conv2d(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.hpp

@@ -215,7 +45,6 @@
    std::optional<const Conv2dConfig> conv_config_ = std::nullopt,
    const std::optional<const MemoryConfig> memory_config = std::nullopt);

-
 struct Conv2dOperation{
    static Result invoke(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.hpp

-
-template <typename T>
-std::pair<ttnn::Tensor, std::optional<ttnn::Tensor>> prepare_conv_weights_biases_and_move_to_device(const ttnn::Tensor& weight_tensor, std::optional<const ttnn::Tensor>& bias_tensor, uint32_t input_channels_alignment, DataType weights_bias_dtype, uint32_t weight_block_h_ntiles, uint32_t weight_block_w_ntiles, const sliding_window::ParallelConfig& parallel_config, T * device, uint32_t groups, uint32_t act_block_h_ntiles, uint32_t input_width);
-
 template <typename T>
 Result conv2d(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:29Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.hpp

@@ -215,7 +45,6 @@
    std::optional<const Conv2dConfig> conv_config_ = std::nullopt,
    const std::optional<const MemoryConfig> memory_config = std::nullopt);

-
 struct Conv2dOperation{
    static Result invoke(


⚠️ clang-diagnostic-error ⚠️
unknown type name Result

github-actions · 2024-11-07T22:03:30Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d_utils.cpp

+        if (num_cores_nhw < compute_grid_size.x && out_nhw_ntiles > compute_grid_size.x) {
+            num_cores_nhw = find_closest_largest_divisor_with_num_padding(out_nhw_ntiles, compute_grid_size.x);
+        }
+        grid = num_cores_to_corerange_set(num_cores_nhw, compute_grid_size, true);


⚠️ clang-diagnostic-error ⚠️
use of undeclared identifier num_cores_to_corerange_set; did you mean num_cores_to_corerangeset?

Suggested change

grid = num_cores_to_corerange_set(num_cores_nhw, compute_grid_size, true);

grid = num_cores_to_corerangeset(num_cores_nhw, compute_grid_size, true);

github-actions

⚠️ Clang-Tidy found issue(s) with the introduced code (2/2)

github-actions · 2024-11-07T22:03:41Z

ttnn/cpp/ttnn/operations/conv/conv2d/conv2d_utils.cpp

+    } else if (shard_layout == TensorMemoryLayout::WIDTH_SHARDED) {
+        num_cores_nhw = 1;
+        uint32_t num_cores_c = find_closest_common_largest_divisor(out_c_ntiles, std::ceil((float)input_channels / effective_tile_width), max_num_cores);
+        grid = num_cores_to_corerange_set(num_cores_c, compute_grid_size, true);


⚠️ clang-diagnostic-error ⚠️
use of undeclared identifier num_cores_to_corerange_set; did you mean num_cores_to_corerangeset?

Suggested change

grid = num_cores_to_corerange_set(num_cores_c, compute_grid_size, true);

grid = num_cores_to_corerangeset(num_cores_c, compute_grid_size, true);

Added new weight and bias preparation ops, and new conv op that can only take pre-prepared weights. Working conv test with pre-prepared weights added return weight/output dims kwargs Only auto-shard if shard_layout not specified Pass input memory config to prepare functions Organize utility functions into their own files

github-actions bot reviewed Nov 1, 2024

View reviewed changes

bbradelTT reviewed Nov 5, 2024

View reviewed changes

shwetankTT mentioned this pull request Nov 7, 2024

Delete convd_host_weights and update all tests using conv2d #14179

Open

LPanosTT force-pushed the lpanos/conv_weight_prepare branch from c368626 to 638c23e Compare November 7, 2024 21:57

github-actions bot reviewed Nov 7, 2024

View reviewed changes

LPanosTT force-pushed the lpanos/conv_weight_prepare branch from 638c23e to f22c36d Compare November 7, 2024 22:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Expose conv2d weights/biases as ops #14566

Expose conv2d weights/biases as ops #14566

LPanosTT commented Nov 1, 2024

github-actions bot left a comment

github-actions bot Nov 1, 2024

github-actions bot Nov 1, 2024

bbradelTT Nov 5, 2024

bbradelTT Nov 5, 2024

bbradelTT Nov 5, 2024

github-actions bot left a comment

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot Nov 7, 2024

github-actions bot left a comment

github-actions bot Nov 7, 2024

	template std::tuple<ttnn::Tensor, uint32_t, uint32_t, ttnn::Tensor, std::optional<ttnn::Tensor>> conv2d<Device>(
	template<> std::tuple<ttnn::Tensor, uint32_t, uint32_t, ttnn::Tensor, std::optional<ttnn::Tensor>> conv2d<Device>(

	grid = num_cores_to_corerange_set(num_cores_nhw, compute_grid_size, true);
	grid = num_cores_to_corerangeset(num_cores_nhw, compute_grid_size, true);

Expose conv2d weights/biases as ops #14566

Are you sure you want to change the base?

Expose conv2d weights/biases as ops #14566

Conversation

LPanosTT commented Nov 1, 2024

Tickets

Problem description

What's changed

Checklist

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Nov 1, 2024

Choose a reason for hiding this comment

github-actions bot Nov 1, 2024

Choose a reason for hiding this comment

bbradelTT Nov 5, 2024

Choose a reason for hiding this comment

bbradelTT Nov 5, 2024

Choose a reason for hiding this comment

bbradelTT Nov 5, 2024

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Nov 7, 2024

Choose a reason for hiding this comment