#11838: Remove num_cores_to_core_range_set

tenstorrent · Sep 7, 2024 · 014940b · 014940b
1 parent 3802793
commit 014940b
Show file tree

Hide file tree

Showing 8 changed files with 8 additions and 38 deletions.
diff --git a/tests/sweep_framework/sweeps/matmul/full/matmul_default_height_sharded.py b/tests/sweep_framework/sweeps/matmul/full/matmul_default_height_sharded.py
@@ -88,7 +88,7 @@ def run(
     # TODO: row_wise=False and ROW_MAJOR shard orientation gives bad PCC
     # TODO: COL_MAJOR shard orientation doesn't work for get_matmul_program_config
     input_a_memory_config.shard_spec = ttnn.ShardSpec(
-        ttnn.num_cores_to_core_range_set(num_cores_height, core_grid, row_wise=True),
+        ttnn.num_cores_to_corerange_set(num_cores_height, core_grid, row_wise=True),
         (per_core_height, k_size),
         ttnn.ShardOrientation.ROW_MAJOR,
         False,

diff --git a/tests/sweep_framework/sweeps/matmul/full/matmul_default_width_sharded.py b/tests/sweep_framework/sweeps/matmul/full/matmul_default_width_sharded.py
@@ -81,7 +81,7 @@ def run(
     # TODO: row_wise=False and ROW_MAJOR shard orientation gives bad PCC
     # TODO: COL_MAJOR shard orientation doesn't work for get_matmul_program_config
     input_a_memory_config.shard_spec = ttnn.ShardSpec(
-        ttnn.num_cores_to_core_range_set(num_cores_width, core_grid, row_wise=True),
+        ttnn.num_cores_to_corerange_set(num_cores_width, core_grid, row_wise=True),
         (total_height, per_core_width),
         ttnn.ShardOrientation.ROW_MAJOR,
         False,

diff --git a/tests/sweep_framework/sweeps/matmul/short/matmul_user_program_config_mcast_1d.py b/tests/sweep_framework/sweeps/matmul/short/matmul_user_program_config_mcast_1d.py
@@ -77,7 +77,7 @@ class TensorMemoryConfigs(enum.Enum):
                 memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
                 buffer_type=ttnn.BufferType.L1,
                 shard_spec=ttnn.ShardSpec(
-                    ttnn.num_cores_to_core_range_set(28, core_grid, row_wise=True),
+                    ttnn.num_cores_to_corerange_set(28, core_grid, row_wise=True),
                     (64, IN0_INNER_DIM_PER_CORE),
                     ttnn.ShardOrientation.ROW_MAJOR,
                     False,
@@ -105,7 +105,7 @@ class TensorMemoryConfigs(enum.Enum):
                 memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
                 buffer_type=ttnn.BufferType.L1,
                 shard_spec=ttnn.ShardSpec(
-                    ttnn.num_cores_to_core_range_set(35, core_grid, row_wise=True),
+                    ttnn.num_cores_to_corerange_set(35, core_grid, row_wise=True),
                     (64, IN0_INNER_DIM_PER_CORE),
                     ttnn.ShardOrientation.ROW_MAJOR,
                     False,
@@ -134,7 +134,7 @@ class TensorMemoryConfigs(enum.Enum):
                 memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
                 buffer_type=ttnn.BufferType.L1,
                 shard_spec=ttnn.ShardSpec(
-                    ttnn.num_cores_to_core_range_set(28, core_grid, row_wise=True),
+                    ttnn.num_cores_to_corerange_set(28, core_grid, row_wise=True),
                     (64, IN0_INNER_DIM_PER_CORE),
                     ttnn.ShardOrientation.ROW_MAJOR,
                     False,
@@ -163,7 +163,7 @@ class TensorMemoryConfigs(enum.Enum):
                 memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
                 buffer_type=ttnn.BufferType.L1,
                 shard_spec=ttnn.ShardSpec(
-                    ttnn.num_cores_to_core_range_set(30, core_grid, row_wise=True),
+                    ttnn.num_cores_to_corerange_set(30, core_grid, row_wise=True),
                     (64, IN0_INNER_DIM_PER_CORE),
                     ttnn.ShardOrientation.ROW_MAJOR,
                     False,

diff --git a/tt_metal/common/work_split.hpp b/tt_metal/common/work_split.hpp
@@ -156,12 +156,6 @@ inline std::set<CoreRange> num_cores_to_corerange_set(
     return num_cores_to_corerange_set({0, 0}, target_num_cores, grid_size, row_wise);
 }
 
-// TODO: Switch num_cores_to_corerange_set to always return CoreRangeSet
-inline CoreRangeSet num_cores_to_core_range_set(
-    const uint32_t target_num_cores, const CoreCoord grid_size, const bool row_wise = false) {
-    return CoreRangeSet(num_cores_to_corerange_set({0, 0}, target_num_cores, grid_size, row_wise));
-}
-
 // This function takes in the core grid size, as well as the number of units of work to divide between the cores
 // This function returns the number of cores, the CoreRangeSet of all cores, and then the CoreRangeSet that does
 // the greater amount of work, and the CoreRangeSet that does less work if work cannot be evenly divided

diff --git a/ttnn/cpp/pybind11/operations/core.hpp b/ttnn/cpp/pybind11/operations/core.hpp
@@ -311,14 +311,6 @@ void py_module(py::module& module) {
             py::arg("memory_config") = std::nullopt,
             py::arg("device") = nullptr});
 
-    module.def(
-        "num_cores_to_core_range_set",
-        &tt::tt_metal::num_cores_to_core_range_set,
-        py::arg("target_num_cores"),
-        py::arg("grid_size"),
-        py::arg("row_wise")= false,
-        R"doc(Returns a CoreRangeSet from number of cores)doc");
-
     module.def(
         "num_cores_to_corerange_set",
         py::overload_cast<const uint32_t, const CoreCoord, const bool>(&tt::tt_metal::num_cores_to_corerange_set),

diff --git a/ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.cpp b/ttnn/cpp/ttnn/operations/conv/conv2d/conv2d.cpp
@@ -89,14 +89,14 @@ ParallelConfig determine_parallel_config(
 
     auto calculate_grid = [&](uint32_t num_cores_nhw) {
         if (shard_layout == TensorMemoryLayout::HEIGHT_SHARDED) {
-            CoreRangeSet grid = num_cores_to_core_range_set(num_cores_nhw, device_grid_size_coord, true);
+            CoreRangeSet grid = num_cores_to_corerange_set(num_cores_nhw, device_grid_size_coord, true);
             return grid;
 
         } else if(shard_layout == TensorMemoryLayout::WIDTH_SHARDED) {
             uint32_t num_cores_channels = find_closest_common_largest_divisor(
                 conv_out_2d_matrix_width_ntiles, std::ceil((double)input_channels / (double)tt::constants::TILE_WIDTH), max_num_cores);
              log_debug(LogOp, "Num cores for Width Sharding : {}", num_cores_channels);
-            CoreRangeSet grid = num_cores_to_core_range_set(num_cores_channels, device_grid_size_coord, true);
+            CoreRangeSet grid = num_cores_to_corerange_set(num_cores_channels, device_grid_size_coord, true);
             return grid;
 
         } else if(shard_layout == TensorMemoryLayout::BLOCK_SHARDED) {

diff --git a/ttnn/ttnn/__init__.py b/ttnn/ttnn/__init__.py
@@ -205,7 +205,6 @@ def manage_config(name, value):
     dump_memory_config,
     load_memory_config,
     dump_stack_trace_on_segfault,
-    num_cores_to_core_range_set,
     num_cores_to_corerange_set,
 )
 

diff --git a/ttnn/ttnn/core.py b/ttnn/ttnn/core.py
@@ -44,21 +44,6 @@ def is_sharded(tensor) -> bool:
 get_memory_config = ttnn._ttnn.core.get_memory_config
 
 
-def num_cores_to_core_range_set(
-    target_num_cores: int,
-    grid_size: ttnn.CoreCoord,
-    row_wise: bool = False,
-):
-    """
-    Returns a CoreRangeSet from number of cores
-    """
-    return ttnn._ttnn.operations.core.num_cores_to_core_range_set(
-        target_num_cores,
-        grid_size,
-        row_wise,
-    )
-
-
 def num_cores_to_corerange_set(
     target_num_cores: int,
     grid_size: ttnn.CoreCoord,