#380,#388: Handle NONE type chip capability, dump chip channel when …

…translating flatbuffer to mlir (#402)
tenstorrent · Aug 15, 2024 · 29f03dd · 29f03dd
1 parent e2a4b7a
commit 29f03dd
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 11 deletions.
diff --git a/include/ttmlir/Dialect/TT/IR/TTOpsEnums.td b/include/ttmlir/Dialect/TT/IR/TTOpsEnums.td
@@ -135,7 +135,7 @@ def TT_ChipCapability : I32BitEnumAttr<"ChipCapability", "TT Chip Capabilities",
                             TT_ChipCapabilityPCIE,
                             TT_ChipCapabilityHostMMIO,
                            ]> {
-  let genSpecializedAttr = 0;
+  let genSpecializedAttr = 1;
   let cppNamespace = "::mlir::tt";
 }
 

diff --git a/include/ttmlir/Dialect/TT/IR/TTOpsTypes.td b/include/ttmlir/Dialect/TT/IR/TTOpsTypes.td
@@ -47,10 +47,6 @@ def TT_GridAttr : TT_Attr<"Grid", "grid"> {
   }];
 }
 
-def TT_ChipCapabilityAttr : EnumAttr<TT_Dialect, TT_ChipCapability, "Chip_capability"> {
-  let assemblyFormat = "`<` $value `>`";
-}
-
 def TT_ArchAttr : EnumAttr<TT_Dialect, TT_Arch, "arch"> {
   let assemblyFormat = "`<` $value `>`";
 }
@@ -106,7 +102,7 @@ def TT_ChipChannelAttr : TT_Attr<"ChipChannel", "chip_channel"> {
                         ArrayRefParameter<"int64_t">:$ethernetCoreCoord0,
                         "unsigned":$deviceId1,
                         ArrayRefParameter<"int64_t">:$ethernetCoreCoord1);
-  let assemblyFormat = "`<` $deviceId0 `,` $ethernetCoreCoord0 `,` $deviceId1 `,` $ethernetCoreCoord1 `>`";
+  let assemblyFormat = "`<` `[` $deviceId0 `,` $ethernetCoreCoord0 `]` `,` `[` $deviceId1 `,` $ethernetCoreCoord1 `]` `>`";
 }
 
 def TT_SystemDescAttr : TT_Attr<"SystemDesc", "system_desc"> {

diff --git a/lib/Dialect/TT/IR/TTOpsTypes.cpp b/lib/Dialect/TT/IR/TTOpsTypes.cpp
@@ -78,6 +78,7 @@ mlir::tt::SystemDescAttr::getFromPath(MLIRContext *context, std::string &path) {
       binary_system_desc->chip_desc_indices();
   auto const *chip_capabilities = binary_system_desc->chip_capabilities();
   auto const *binary_chip_coords = binary_system_desc->chip_coords();
+  auto const *chip_channel_connections = binary_system_desc->chip_channels();
 
   // Acquire chip descs
   std::vector<tt::ChipDescAttr> chip_desc_list;
@@ -151,10 +152,26 @@ mlir::tt::SystemDescAttr::getFromPath(MLIRContext *context, std::string &path) {
     chip_coordinate_list.push_back(chip_coordinate_attr);
   }
 
+  std::vector<tt::ChipChannelAttr> chip_channel_list;
+  for (auto element : *chip_channel_connections) {
+    std::vector<int64_t> ethernet_core_coord0_vec = {
+        element->ethernet_core_coord0().y(),
+        element->ethernet_core_coord0().x()};
+
+    std::vector<int64_t> ethernet_core_coord1_vec = {
+        element->ethernet_core_coord1().y(),
+        element->ethernet_core_coord1().x()};
+
+    auto chip_channel_attr = tt::ChipChannelAttr::get(
+        context, element->device_id0(), ethernet_core_coord0_vec,
+        element->device_id1(), ethernet_core_coord1_vec);
+    chip_channel_list.push_back(chip_channel_attr);
+  }
+
   // Generate system desc attribute
-  auto system_desc_attr =
-      tt::SystemDescAttr::get(context, chip_desc_list, chip_indices_list,
-                              chip_capabilities_list, chip_coordinate_list, {});
+  auto system_desc_attr = tt::SystemDescAttr::get(
+      context, chip_desc_list, chip_indices_list, chip_capabilities_list,
+      chip_coordinate_list, chip_channel_list);
 
   return system_desc_attr;
 }

diff --git a/lib/Dialect/TTNN/Transforms/Passes.cpp b/lib/Dialect/TTNN/Transforms/Passes.cpp
@@ -34,8 +34,6 @@ class TTNNOpenDevice : public impl::TTNNOpenDeviceBase<TTNNOpenDevice> {
     OpBuilder builder(module);
     auto systemDesc = llvm::cast<tt::SystemDescAttr>(
         module->getAttr(tt::SystemDescAttr::name));
-    auto chipDescIndices = systemDesc.getChipDescIndices();
-    assert(chipDescIndices.size() == 1 && "Multiple chips not supported yet");
 
     module->walk([&](func::FuncOp func) {
       // For now just push the open and close device ops to the beginning and