Support for multi-device system descriptor #220

jnie-TT · 2024-07-23T21:09:27Z

#162
Example of system descriptor for n300:
{'version': {'major': 0, 'minor': 0, 'patch': 39}, 'ttmlir_git_hash': '484963fb271e0817d5aa5fd6e6c4c3626787c2f2', 'product_identifier': 'unknown', 'system_desc': {'chip_descs': [{'arch': 'Wormhole_b0', 'grid_size': {'y': 8, 'x': 8}, 'l1_size': 1499136, 'num_dram_channels': 12, 'dram_channel_size': 1073741824, 'noc_l1_address_align_bytes': 16, 'pcie_address_align_bytes': 32, 'noc_dram_address_align_bytes': 32}, {'arch': 'Wormhole_b0', 'grid_size': {'y': 8, 'x': 8}, 'l1_size': 1499136, 'num_dram_channels': 12, 'dram_channel_size': 1073741824, 'noc_l1_address_align_bytes': 16, 'pcie_address_align_bytes': 32, 'noc_dram_address_align_bytes': 32}], 'chip_desc_indices': [0, 1], 'chip_capabilities': [3, 0], 'chip_coords': [{'rack': 0, 'shelf': 0, 'y': 0, 'x': 0}, {'rack': 0, 'shelf': 0, 'y': 0, 'x': 1}], 'chip_channels': [{'endpoint0': 1, 'endpoint1': -1}, {'endpoint0': 0, 'endpoint1': -1}]}}

This requires metal API updates, specifically to support querying the chip location. I added it locally for testing, but currently this API is missing so the build will fail. @nsmithtt What's the process for adding such APIs to metal?

nsmithtt · 2024-07-24T13:38:33Z

This requires metal API updates, specifically to support querying the chip location. I added it locally for testing, but currently this API is missing so the build will fail. @nsmithtt What's the process for adding such APIs to metal?

I think we can just create a PR and it shouldn't be an issue. Can you send me your branch name and I'll take a look?

nsmithtt · 2024-07-24T13:40:55Z

runtime/lib/ttnn/runtime.cpp

+    chipCapabilities.push_back(chipCapability);
+    // Extract chip location
+    int x, y, rack, shelf;
+    std::tie(x, y, rack, shelf) = device.get_chip_location();


Is get_chip_location the only new API you added to metal? If so let's file an issue on metal (to have an issue number in the commit) something like "Add device chip location API".

Cool sounds good

@nsmithtt Metal PR here: https://github.com/tenstorrent/tt-metal/pull/10674/files

issue here: tenstorrent/tt-metal#10671

nsmithtt · 2024-07-24T13:58:40Z

runtime/lib/ttnn/runtime.cpp

+    std::tie(x, y, rack, shelf) = device.get_chip_location();
+    chipCoords.emplace_back(rack, shelf, y, x);
+    // Extract chip connected channels
+    std::pair<int, int> connectedChips = getConnectedChips(device);


Can't there be more than 1 connect chips? Seems like a remote chip in the middle of a galaxy could have up to 4 connections.

Hmmmm that's a good point, I'll update the chipChannels schema to be an array of ints then instead of having only 2 endpoints hardcoded

The intention for chipChannels was to be just a flat list of connections, not tied to the other arrays, but maybe this is more cumbersome to query. If you think of it like a graph you have a list of nodes chipIds and a list of edges chipChannels where each edge is just a pair of node ids (but could be extended with new attributes in the future, like number of streams).

@nsmithtt Ahhh I see, I misinterpreted what chipChannels was intended to do. I pushed an update that now constructs a flat list of connections - also ported to a metal device mesh API that handles opening/closing devices safely.

Updated system descriptor looks like:

{'version': {'major': 0, 'minor': 0, 'patch': 40}, 'ttmlir_git_hash': '1fdd4a0970b17925db82c74c163834ccb24c737f', 'product_identifier': 'unknown', 'system_desc': {'chip_descs': [{'arch': 'Wormhole_b0', 'grid_size': {'y': 8, 'x': 8}, 'l1_size': 1499136, 'num_dram_channels': 12, 'dram_channel_size': 1073741824, 'noc_l1_address_align_bytes': 16, 'pcie_address_align_bytes': 32, 'noc_dram_address_align_bytes': 32}, {'arch': 'Wormhole_b0', 'grid_size': {'y': 8, 'x': 8}, 'l1_size': 1499136, 'num_dram_channels': 12, 'dram_channel_size': 1073741824, 'noc_l1_address_align_bytes': 16, 'pcie_address_align_bytes': 32, 'noc_dram_address_align_bytes': 32}], 'chip_desc_indices': [0, 1], 'chip_capabilities': [3, 0], 'chip_coords': [{'rack': 0, 'shelf': 0, 'y': 0, 'x': 0}, {'rack': 0, 'shelf': 0, 'y': 0, 'x': 1}], 'chip_channels': [{'endpoint0': 0, 'endpoint1': 1}]}}

jnie-TT · 2024-08-06T17:04:39Z

@tapspatel I had to comment out the strict system descriptor check or else the assertion always fires. I'll open an issue for the compiler to properly generate a system descriptor.

@nsmithtt does this match what you were expecting? I ignored chip coords and updated chip channels to include the two endpoint devices and their connected ethernet coords.

Multi-device system descriptor looks like this for an x2 card:

{
  "chip_descs": [
    {
      "arch": "Wormhole_b0",
      "grid_size": {
        "y": 8,
        "x": 8
      },
      "l1_size": 1499136,
      "num_dram_channels": 12,
      "dram_channel_size": 1073741824,
      "noc_l1_address_align_bytes": 16,
      "pcie_address_align_bytes": 32,
      "noc_dram_address_align_bytes": 32
    },
    {
      "arch": "Wormhole_b0",
      "grid_size": {
        "y": 8,
        "x": 8
      },
      "l1_size": 1499136,
      "num_dram_channels": 12,
      "dram_channel_size": 1073741824,
      "noc_l1_address_align_bytes": 16,
      "pcie_address_align_bytes": 32,
      "noc_dram_address_align_bytes": 32
    }
  ],
  "chip_desc_indices": [0, 1],
  "chip_capabilities": [3, 0],
  "chip_coords": [],
  "chip_channels": [
    {
      "device_id0": 0,
      "ethernet_core_coord0": {
        "y": 8,
        "x": 0
      },
      "device_id1": 1,
      "ethernet_core_coord1": {
        "y": 0,
        "x": 0
      }
    }
  ]
}

tapspatel · 2024-08-06T17:30:24Z

@jnie-TT yeah, I added that assert since we were technically generating for n300 but running on n150. I have an issue open that I'm working on right now to address that: #306 (although this is more from test POV). Be default, compiler will generate for n300.

nsmithtt · 2024-08-08T03:31:12Z

include/ttmlir/Target/Utils/MLIRToFlatbuffer.h

 }

 inline ::tt::target::Dim2d toFlatbuffer(FlatbufferObjectCache &cache,
-                                        GridAttr arch) {
+                                        const GridAttr &arch) {


Not blocking the review, but typically we only pass the MLIR types, i.e. Value, *Attr, *Type, by value. MLIR already enforces that they are wrappers around some underlying storage and it is convention to treat them as value types. It doesn't hurt to pass it by const&, but it's not necessary. All MLIR types are also immutable, so they are truly value types in this sense.

Ahhh I see, thanks for the explanation, I'll undo these changes!

nsmithtt · 2024-08-08T03:33:17Z

runtime/lib/ttnn/runtime.cpp

+  std::vector<::tt::target::ChipChannel> allConnections;
+  allConnections.resize(connectionSet.size());
+
+  std::transform(


nsmithtt

This looks great! Couple comments inline

nsmithtt · 2024-08-08T03:43:26Z

runtime/lib/ttnn/runtime.cpp

+
+  for (const ::ttnn::Device *device : devices) {
+    // Construct chip descriptor
+    ::tt::target::Dim2d deviceGrid = toFlatbuffer(device->logical_grid_size());


device->logical_grid_size() -> device.compute_with_storage_grid_size()

tapspatel · 2024-08-08T14:19:59Z

runtime/tools/python/ttrt/common/api.py

@@ -153,7 +153,8 @@ def run(args):

    # execution
    print("executing action for all provided flatbuffers")
-    device = ttrt.runtime.open_device(device_ids)
+    system_desc, device_ids = ttrt.runtime.get_current_system_desc()


this is already called earlier on in the function, you can remove

Great catch! This was leftover after the rebase, removed it

tapspatel · 2024-08-08T15:25:57Z

ttrt code lgtm!

nsmithtt

Awesome refactor! Both runtimes have a shared system desc! 🙌

kmabeeTT · 2024-08-13T04:33:24Z

Hey @jnie-TT , am I nuts or is there some issue here at the merged commit 0b393ec in main?

Looks good in this format:

% ttrt query --save-artifacts --system-desc
<snip>

    "chip_capabilities": [
      3,
      0
    ],

Not so good here though, empty brackets causes error:

% ./build/bin/ttmlir-opt --ttir-to-ttnn-backend-pipeline --ttir-load-system-desc="path=./ttrt-artifacts/system_desc.ttsys" test/ttmlir/Dialect/TTNN/simple_subtract.mlir | ./build/bin/ttmlir-translate --ttnn-to-flatbuffer

<stdin>:6:572: error: expected valid keyword
module attributes {tt.device = #tt.device<#tt.grid<8x8, (d0, d1) -> (0, d0, d1)>, [0]>, tt.system_desc = #tt.system_desc<[{arch = <wormhole_b0>, grid = 7x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32}, {arch = <wormhole_b0>, grid = 7x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32}], [0, 1], [<pcie|host_mmio>, <>], [<0, 0, 0, 0>]>} {
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           ^
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      ```

jnie-TT · 2024-08-13T13:43:38Z

@kmabeeTT Seems like chip_capabilities=0 is causing an issue, I think it's an issue on the mlir side, I'll take a closer look at this. In the meantime you can try running on an n150 machine, the issue should go away.

jnie-TT · 2024-08-13T14:24:53Z

Hey @jnie-TT , am I nuts or is there some issue here at the merged commit 0b393ec in main?

Looks good in this format:

% ttrt query --save-artifacts --system-desc
<snip>

    "chip_capabilities": [
      3,
      0
    ],

Not so good here though, empty brackets causes error:

% ./build/bin/ttmlir-opt --ttir-to-ttnn-backend-pipeline --ttir-load-system-desc="path=./ttrt-artifacts/system_desc.ttsys" test/ttmlir/Dialect/TTNN/simple_subtract.mlir | ./build/bin/ttmlir-translate --ttnn-to-flatbuffer

<stdin>:6:572: error: expected valid keyword
module attributes {tt.device = #tt.device<#tt.grid<8x8, (d0, d1) -> (0, d0, d1)>, [0]>, tt.system_desc = #tt.system_desc<[{arch = <wormhole_b0>, grid = 7x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32}, {arch = <wormhole_b0>, grid = 7x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32}], [0, 1], [<pcie|host_mmio>, <>], [<0, 0, 0, 0>]>} {
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           ^
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      ```

Created a ticket #380

jnie-TT requested a review from nsmithtt July 23, 2024 21:09

jnie-TT requested review from tapspatel, kmabeeTT and AleksKnezevic as code owners July 23, 2024 21:09

nsmithtt reviewed Jul 24, 2024

View reviewed changes

jnie-TT force-pushed the jnie/sys_desc branch from 1fdd4a0 to 74b8070 Compare July 24, 2024 19:26

jnie-TT force-pushed the jnie/sys_desc branch from 74b8070 to e065071 Compare August 6, 2024 16:59

jnie-TT requested review from sdjordjevicTT, rpavlovicTT, mrakitaTT and nobradovictt as code owners August 6, 2024 16:59

jnie-TT force-pushed the jnie/sys_desc branch from e065071 to 5bf6652 Compare August 8, 2024 03:17

nsmithtt reviewed Aug 8, 2024

View reviewed changes

tapspatel requested changes Aug 8, 2024

View reviewed changes

jnie-TT force-pushed the jnie/sys_desc branch from 5bf6652 to 54e77a7 Compare August 8, 2024 15:23

tapspatel approved these changes Aug 8, 2024

View reviewed changes

jnie-TT force-pushed the jnie/sys_desc branch 2 times, most recently from 71cde8c to 1bf8d22 Compare August 8, 2024 19:42

#162: TTNN: Support for multi-device system descriptor.

a38f35e

jnie-TT force-pushed the jnie/sys_desc branch from 1bf8d22 to 259278e Compare August 10, 2024 03:54

nsmithtt approved these changes Aug 10, 2024

View reviewed changes

Commonize system descriptor API, update metal APIs

c3fc7dd

jnie-TT force-pushed the jnie/sys_desc branch from 259278e to c3fc7dd Compare August 10, 2024 16:26

jnie-TT merged commit 0b393ec into main Aug 10, 2024
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support for multi-device system descriptor #220

Support for multi-device system descriptor #220

jnie-TT commented Jul 23, 2024

nsmithtt commented Jul 24, 2024

nsmithtt Jul 24, 2024

jnie-TT Jul 24, 2024

jnie-TT Jul 24, 2024

nsmithtt Jul 24, 2024

jnie-TT Jul 24, 2024

nsmithtt Jul 24, 2024

jnie-TT Jul 24, 2024

jnie-TT commented Aug 6, 2024

tapspatel commented Aug 6, 2024

nsmithtt Aug 8, 2024

jnie-TT Aug 8, 2024

nsmithtt Aug 8, 2024

nsmithtt left a comment

nsmithtt Aug 8, 2024

tapspatel Aug 8, 2024

jnie-TT Aug 8, 2024

tapspatel commented Aug 8, 2024

nsmithtt left a comment

kmabeeTT commented Aug 13, 2024

jnie-TT commented Aug 13, 2024

jnie-TT commented Aug 13, 2024

Support for multi-device system descriptor #220

Support for multi-device system descriptor #220

Conversation

jnie-TT commented Jul 23, 2024

nsmithtt commented Jul 24, 2024

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

jnie-TT commented Aug 6, 2024

tapspatel commented Aug 6, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nsmithtt 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

tapspatel commented Aug 8, 2024

nsmithtt left a comment

Choose a reason for hiding this comment

kmabeeTT commented Aug 13, 2024

jnie-TT commented Aug 13, 2024

jnie-TT commented Aug 13, 2024