#17477: Introduce ND coordinate system for TT-distributed #17745

omilyutin-tt · 2025-02-07T22:49:26Z

Ticket

Problem description

Existing mesh infra assumes 2D. This assumption won't hold in the future.

What's changed

Introduce a new SimpleMeshShape that will gradually replace the existing MeshShape, after which it will be renamed to MeshShape.

Introduce MeshCoordinate, MeshCoordinateRange, and MeshContainer - primitives designed to work with the new ND coordinate system.

MeshContainer allows efficient flat representation of various metadata that matches the mesh shape. Iterators are available to make it easy to use. MeshCoordinate along with strides that are precomputed on SimpleMeshShape allows for an easy point access. The integration with MeshBuffer demonstrates the use case.

Next steps:

Replace the existing MeshShape, MeshOffset, and the related aliases with the new SimpleMeshShape, and MeshCoordinate.
No plans to generalize with CoreCoord, for now. Cores are fundamentally in 2D, so a more specialized system can be used for efficiency. Also it is not desired to make CoreCoord to interop with MeshCoordinate - the 2 sets of coordinates mean entirely different concepts.
More functionality might be added, as we continue working on TT-distributed.

Checklist

All post commit
New/Existing tests provide coverage for changes

omilyutin-tt · 2025-02-07T23:00:57Z

ttnn/cpp/ttnn/operations/data_movement/pad/pad.cpp

-bool eq_spans(const ArrayType& a, const ArrayType& b) {
-    return std::equal(a.begin(), a.end(), b.begin(), b.end());
-}
+bool eq_spans(const auto a, const auto b) { return std::equal(a.begin(), a.end(), b.begin(), b.end()); }


I had to do this because a is now tt::stl::Span while b is std::span. Annoying, but this keeps Metal at cpp17.

tt-asaigal

Looks great overall. One thing to keep in mind is that our existing 2D coordinate system in Metal (exposed through CoreCoord, CoreRange and CoreRangeSet) provides a bunch of utility functions, allowing users to compute set/range intersections, adjacency, etc.
It would be very useful for us to expose similar APIs for this new ND coordinate system as well. Especially as we start introducing more heterogeneity in our workloads.

tt-asaigal · 2025-02-08T22:58:40Z

tt_metal/distributed/mesh_buffer.cpp

-        mesh_device_->num_cols());
-    return buffers_[device_coord.row][device_coord.col];
+    return get_device_buffer(MeshCoordinate(device_coord.row, device_coord.col));
+}


We should add a comment saying that this overload is kept around to be compatible with existing infra and will be removed once everything is migrated to use MeshCoordinate.

Discussed offline that the overload itself is useful to keep as a shorthand (in case user just wants to use 2D cluster). But we will drop the "num rows" "num cols" terminology, and the infra will be generic.

This is for MeshCoordinate overload - the std::shared_ptr<Buffer> MeshBuffer::get_device_buffer(const Coordinate& device_coord) will go away entirely.

tt-asaigal · 2025-02-08T23:00:00Z

tt_metal/distributed/mesh_device.cpp

@@ -218,7 +218,11 @@ std::vector<IDevice*> MeshDevice::get_devices() const { return view_->get_device

 // TODO: Remove this function once we have a proper view interface
 IDevice* MeshDevice::get_device(size_t row_idx, size_t col_idx) const {
-    return this->get_device_index(row_idx * num_cols() + col_idx);


Same here, I don't see why we need this overload to expose physical devices using [row, col] in the long term. This doesn't make sense for an ND mesh anyway

Yep all of this will go away.

yeah, we should remove this in another PR otherwise it'll be confusing with (y, x) vs. (x,y). Actually - is this a bug? Should it be MeshCoordinate{col_idx, row_idx}?

I got rid of rows/cols terminology, this logic should be consistent with itself -- we end up computing:

linear_index += coord[dim] * shape.get_stride(dim);

Over all dims. Since I initialize the shape with [num_rows, num_cols] and coord with [row_idx, col_idx], I am effectively getting the same thing.

Unless I am misreading something :)

tt-asaigal · 2025-02-08T23:09:14Z

tt_metal/common/mesh_coord.cpp

+}
+
+MeshCoordinate::MeshCoordinate(uint32_t coord) : value_({coord}) {}
+MeshCoordinate::MeshCoordinate(uint32_t row, uint32_t col) : value_({row, col}) {}


The use of [row, col] should be removed eventually. I expect that users will be extremely confused by two different notations being exposed by the same data structure.
Fundamentally, removing this requires MeshDevice to start using a Cartesian scheme.

Updated to use x y as discussed offline.

tt-asaigal · 2025-02-08T23:09:56Z

tt_metal/common/mesh_coord.cpp

+}
+
+MeshCoordinateRange::MeshCoordinateRange(const SimpleMeshShape& shape) :
+    MeshCoordinateRange(zero_coordinate(shape.dims()), shape_back(shape)) {}


jvegaTT

Approving on pad file changes

cfjchu

Looks great -- left a few comments to resolve before you merge

cfjchu · 2025-02-14T22:39:09Z

tt_metal/distributed/mesh_device.cpp

@@ -218,7 +218,11 @@ std::vector<IDevice*> MeshDevice::get_devices() const { return view_->get_device

 // TODO: Remove this function once we have a proper view interface
 IDevice* MeshDevice::get_device(size_t row_idx, size_t col_idx) const {
-    return this->get_device_index(row_idx * num_cols() + col_idx);


yeah, we should remove this in another PR otherwise it'll be confusing with (y, x) vs. (x,y). Actually - is this a bug? Should it be MeshCoordinate{col_idx, row_idx}?

cfjchu · 2025-02-14T22:43:48Z

tt_metal/api/tt-metalium/mesh_coord.hpp

+    size_t get_stride(size_t dim) const;
+
+    // Returns the total number of elements in the mesh.
+    size_t mesh_size() const;


maybe just size()?

That was my thinking initially. There is a size() member from ShapeBase that refers to the dimensionality of the shape, so I wanted to be explicit here.

cfjchu · 2025-02-14T23:01:55Z

tt_metal/api/tt-metalium/mesh_coord.hpp

+// Throws if `coord` is out of bounds of `shape`.
+size_t to_linear_index(const SimpleMeshShape& shape, const MeshCoordinate& coord);
+
+// Represents a range of MeshCoordinates. Requires that mesh coordinates have the same dimensionality.


would be good just to add a comment saying it's inclusive start/end

I have it further down, where I mention specific "start" and "end" (without these a bit hard to describe what exactly does range "include").

Also added end_coord() documentation.

tt_metal/api/tt-metalium/mesh_coord.hpp

### Ticket #17477 ### Problem description Existing mesh infra assumes 2D. This assumption won't hold in the future. ### What's changed Introduce a new `SimpleMeshShape` that will gradually replace the existing `MeshShape`, after which it will be renamed to `MeshShape`. Introduce `MeshCoordinate`, `MeshCoordinateRange`, and `MeshContainer` - primitives designed to work with the new ND coordinate system. `MeshContainer` allows efficient flat representation of various metadata that matches the mesh shape. Iterators are available to make it easy to use. `MeshCoordinate` along with strides that are precomputed on `SimpleMeshShape` allows for an easy point access. The integration with `MeshBuffer` demonstrates the use case. Next steps: * Replace the existing `MeshShape`, `MeshOffset`, and the related aliases with the new `SimpleMeshShape`, and `MeshCoordinate`. * No plans to generalize with `CoreCoord`, for now. Cores are fundamentally in 2D, so a more specialized system can be used for efficiency. Also it is not desired to make `CoreCoord` to interop with `MeshCoordinate` - the 2 sets of coordinates mean entirely different concepts. * More functionality might be added, as we continue working on TT-distributed. ### Checklist - [X] [All post commit](https://github.com/tenstorrent/tt-metal/actions/runs/13347753550) - [X] New/Existing tests provide coverage for changes

omilyutin-tt requested review from ntarafdar, sjameelTT, jaykru-tt, yugi957, jvegaTT, llongTT, nardoTT, ayerofieiev-tt, dmakoviichuk-tt, cfjchu, TT-BrianLiu, tt-asaigal, aliuTT, abhullar-tt, pgkeller, tt-aho, tt-dma and ubcheema as code owners February 7, 2025 22:49

omilyutin-tt commented Feb 7, 2025

View reviewed changes

tt-asaigal reviewed Feb 8, 2025

View reviewed changes

jvegaTT approved these changes Feb 9, 2025

View reviewed changes

omilyutin-tt added 7 commits February 14, 2025 17:40

Remove std::span dependency from shapes

d53909a

ND mesh coordinate system

a9b0079

Revert ttnn changes

1eb08d4

Remove unnecessary alias

800e1f3

Better tests, 1D case

86dc0cc

Value proxy class for better container iteration

24f249e

Use x, y instead of rows, cols.

5e5b71a

omilyutin-tt force-pushed the omilyutin/nd branch from be2fd92 to 5e5b71a Compare February 14, 2025 18:00

cfjchu approved these changes Feb 14, 2025

View reviewed changes

cfjchu reviewed Feb 14, 2025

View reviewed changes

tt_metal/api/tt-metalium/mesh_coord.hpp Show resolved Hide resolved

omilyutin-tt added 3 commits February 15, 2025 05:56

Add const iterator

5c8b650

Better to_linear_index test

16be729

Fix shape base view

5d46d98

omilyutin-tt merged commit 52c53d5 into main Feb 16, 2025
236 of 247 checks passed

omilyutin-tt deleted the omilyutin/nd branch February 16, 2025 01:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

#17477: Introduce ND coordinate system for TT-distributed #17745

#17477: Introduce ND coordinate system for TT-distributed #17745

omilyutin-tt commented Feb 7, 2025 •

edited

Loading

omilyutin-tt Feb 7, 2025 •

edited

Loading

tt-asaigal left a comment

tt-asaigal Feb 8, 2025

omilyutin-tt Feb 14, 2025

omilyutin-tt Feb 14, 2025

tt-asaigal Feb 8, 2025

omilyutin-tt Feb 14, 2025

cfjchu Feb 14, 2025

omilyutin-tt Feb 15, 2025

tt-asaigal Feb 8, 2025

omilyutin-tt Feb 14, 2025

tt-asaigal Feb 8, 2025

jvegaTT left a comment

cfjchu left a comment

cfjchu Feb 14, 2025

cfjchu Feb 14, 2025

omilyutin-tt Feb 15, 2025

cfjchu Feb 14, 2025

omilyutin-tt Feb 15, 2025

#17477: Introduce ND coordinate system for TT-distributed #17745

#17477: Introduce ND coordinate system for TT-distributed #17745

Conversation

omilyutin-tt commented Feb 7, 2025 • edited Loading

Ticket

Problem description

What's changed

Checklist

omilyutin-tt Feb 7, 2025 • edited Loading

Choose a reason for hiding this comment

tt-asaigal left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jvegaTT left a comment

Choose a reason for hiding this comment

cfjchu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

omilyutin-tt commented Feb 7, 2025 •

edited

Loading

omilyutin-tt Feb 7, 2025 •

edited

Loading