LightMetal - Add flatbuffer schema and conversion to/from functions f…

…or various types (#17039) - Code is compiled, but not used by anything yet, will be used after subsequent merge of light metal capture/replay libraries. - Make to/from_flatbuffer() friend of CircularBufferConfig to access internals - Remove default case statements from case statements for enums so when new enum values are added, compile error is seen to force updates.
tenstorrent · Jan 30, 2025 · 8abaddb · 8abaddb
1 parent 0971366
commit 8abaddb
Show file tree

Hide file tree

Showing 19 changed files with 1,331 additions and 1 deletion.
diff --git a/tt_metal/api/tt-metalium/circular_buffer_types.hpp b/tt_metal/api/tt-metalium/circular_buffer_types.hpp
@@ -10,6 +10,8 @@
 #include <map>
 #include <optional>
 #include <unordered_set>
+#include <flatbuffers/flatbuffers.h>
+#include "flatbuffers/flatbuffer_builder.h"
 
 #include "logger.hpp"
 #include "tt_backend_api_types.hpp"
@@ -18,12 +20,31 @@
 
 #include "circular_buffer_constants.h"
 
+// Forward declarations for external flatbuffer serialization function
+namespace tt::tt_metal::flatbuffer {
+class CircularBufferConfig;
+}
+namespace tt::tt_metal {
+inline namespace v0 {
+class CircularBufferConfig;
+}
+flatbuffers::Offset<tt::tt_metal::flatbuffer::CircularBufferConfig> to_flatbuffer(
+    const tt::tt_metal::CircularBufferConfig& config, flatbuffers::FlatBufferBuilder& builder);
+CircularBufferConfig from_flatbuffer(
+    const tt::tt_metal::flatbuffer::CircularBufferConfig* config_fb, const Buffer* shadow_global_buffer);
+}  // namespace tt::tt_metal
+
 namespace tt::tt_metal {
 inline namespace v0 {
 
 using CBHandle = uintptr_t;
 
 class CircularBufferConfig {
+    friend flatbuffers::Offset<tt::tt_metal::flatbuffer::CircularBufferConfig> tt::tt_metal::to_flatbuffer(
+        const tt::tt_metal::CircularBufferConfig& config, flatbuffers::FlatBufferBuilder& builder);
+    friend CircularBufferConfig tt::tt_metal::from_flatbuffer(
+        const tt::tt_metal::flatbuffer::CircularBufferConfig* config_fb, const Buffer* shadow_global_buffer);
+
 public:
     // Static circular buffer spec
     CircularBufferConfig(uint32_t total_size, const std::map<uint8_t, tt::DataFormat>& data_format_spec);

diff --git a/tt_metal/impl/CMakeLists.txt b/tt_metal/impl/CMakeLists.txt
@@ -43,12 +43,24 @@ set(IMPL_SRC
     ${CMAKE_CURRENT_SOURCE_DIR}/trace/trace.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/trace/trace_buffer.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/event/event.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/base_types_from_flatbuffer.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/base_types_to_flatbuffer.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/buffer_types_from_flatbuffer.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/buffer_types_to_flatbuffer.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/program_types_from_flatbuffer.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/program_types_to_flatbuffer.cpp
 )
 
 # Include helper functions and generate headers from flatbuffer schemas
 include(flatbuffers)
 
-set(FLATBUFFER_SCHEMAS) # Empty to start, coming soon.
+set(FLATBUFFER_SCHEMAS
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/binary.fbs
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/command.fbs
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/base_types.fbs
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/buffer_types.fbs
+    ${CMAKE_CURRENT_SOURCE_DIR}/flatbuffer/program_types.fbs
+)
 
 foreach(FBS_FILE ${FLATBUFFER_SCHEMAS})
     GENERATE_FBS_HEADER(${FBS_FILE})

diff --git a/tt_metal/impl/flatbuffer/base_types.fbs b/tt_metal/impl/flatbuffer/base_types.fbs
@@ -0,0 +1,102 @@
+namespace tt.tt_metal.flatbuffer;
+
+
+enum Arch: uint {
+  Grayskull = 0,
+  Wormhole_b0 = 1,
+  Blackhole = 2,
+}
+
+enum BufferType: ushort {
+  DRAM = 0,
+  L1 = 1,
+  SystemMemory = 2,
+  L1Small = 3,
+  Trace = 4,
+}
+
+enum TensorMemoryLayout: ushort {
+  None = 0,
+  Interleaved = 1,
+  SingleBank = 2,
+  HeightSharded = 3,
+  WidthSharded = 4,
+  BlockSharded = 5,
+}
+
+
+enum DataMovementProcessor : byte {
+  RISCV_0,
+  RISCV_1
+}
+
+enum NOC : byte {
+  NOC_0,
+  NOC_1
+}
+
+enum NOC_MODE : byte {
+  DM_DEDICATED_NOC,
+  DM_DYNAMIC_NOC
+}
+
+enum Eth : ubyte {
+  SENDER = 0,
+  RECEIVER = 1,
+  IDLE = 2
+}
+
+enum MathFidelity : ubyte {
+  LoFi = 0,
+  HiFi2 = 2,
+  HiFi3 = 3,
+  HiFi4 = 4,
+  Invalid = 255
+}
+
+enum DataFormat : uint8 {
+  Float32 = 0,
+  Float16 = 1,
+  Bfp8 = 2,
+  Bfp4 = 3,
+  Bfp2 = 11,
+  Float16_b = 5,
+  Bfp8_b = 6,
+  Bfp4_b = 7,
+  Bfp2_b = 15,
+  Lf8 = 10,
+  Fp8_e4m3 = 26, // 0x1A in decimal
+  Int8 = 14,
+  Tf32 = 4,
+  UInt8 = 30,
+  UInt16 = 9,
+  Int32 = 8,
+  UInt32 = 24,
+  RawUInt8 = 240, // 0xf0 in decimal
+  RawUInt16 = 241, // 0xf1 in decimal
+  RawUInt32 = 242, // 0xf2 in decimal
+  Invalid = 255
+}
+
+enum UnpackToDestMode : byte {
+  Default,
+  UnpackToDestFp32
+}
+
+table DefineEntry {
+  key: string;
+  value: string;
+}
+
+
+table Tile {
+  tile_shape: [uint32]; // Shape of the tile (e.g., height, width)
+  face_shape: [uint32]; // Shape of the face
+  tile_hw: uint32;      // Tile hardware size
+  face_hw: uint32;      // Face hardware size
+  num_faces: uint32;    // Number of faces
+  partial_face: uint32; // Indicates if this is a partial face
+  narrow_tile: uint32;  // Indicates if this is a narrow tile
+  transpose_within_face: bool; // Transpose within each face
+  transpose_of_faces: bool;    // Transpose face order
+  }
diff --git a/tt_metal/impl/flatbuffer/base_types_from_flatbuffer.cpp b/tt_metal/impl/flatbuffer/base_types_from_flatbuffer.cpp
@@ -0,0 +1,128 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "flatbuffer/base_types_from_flatbuffer.hpp"
+
+namespace tt::tt_metal {
+
+BufferType from_flatbuffer(tt::tt_metal::flatbuffer::BufferType type) {
+    switch (type) {
+        case tt::tt_metal::flatbuffer::BufferType::DRAM: return BufferType::DRAM;
+        case tt::tt_metal::flatbuffer::BufferType::L1: return BufferType::L1;
+        case tt::tt_metal::flatbuffer::BufferType::SystemMemory: return BufferType::SYSTEM_MEMORY;
+        case tt::tt_metal::flatbuffer::BufferType::L1Small: return BufferType::L1_SMALL;
+        case tt::tt_metal::flatbuffer::BufferType::Trace: return BufferType::TRACE;
+    }
+}
+
+tt::tt_metal::DataMovementProcessor from_flatbuffer(tt::tt_metal::flatbuffer::DataMovementProcessor in) {
+    switch (in) {
+        case tt::tt_metal::flatbuffer::DataMovementProcessor::RISCV_0:
+            return tt::tt_metal::DataMovementProcessor::RISCV_0;
+        case tt::tt_metal::flatbuffer::DataMovementProcessor::RISCV_1:
+            return tt::tt_metal::DataMovementProcessor::RISCV_1;
+    }
+}
+
+tt::tt_metal::NOC from_flatbuffer(tt::tt_metal::flatbuffer::NOC in) {
+    switch (in) {
+        case tt::tt_metal::flatbuffer::NOC::NOC_0: return tt::tt_metal::NOC::NOC_0;
+        case tt::tt_metal::flatbuffer::NOC::NOC_1: return tt::tt_metal::NOC::NOC_1;
+    }
+}
+
+tt::tt_metal::NOC_MODE from_flatbuffer(tt::tt_metal::flatbuffer::NOC_MODE in) {
+    switch (in) {
+        case tt::tt_metal::flatbuffer::NOC_MODE::DM_DEDICATED_NOC: return tt::tt_metal::NOC_MODE::DM_DEDICATED_NOC;
+        case tt::tt_metal::flatbuffer::NOC_MODE::DM_DYNAMIC_NOC: return tt::tt_metal::NOC_MODE::DM_DYNAMIC_NOC;
+    }
+}
+
+tt::tt_metal::Eth from_flatbuffer(tt::tt_metal::flatbuffer::Eth in) {
+    switch (in) {
+        case tt::tt_metal::flatbuffer::Eth::SENDER: return tt::tt_metal::Eth::SENDER;
+        case tt::tt_metal::flatbuffer::Eth::RECEIVER: return tt::tt_metal::Eth::RECEIVER;
+        case tt::tt_metal::flatbuffer::Eth::IDLE: return tt::tt_metal::Eth::IDLE;
+    }
+}
+
+MathFidelity from_flatbuffer(tt::tt_metal::flatbuffer::MathFidelity input) {
+    switch (input) {
+        case tt::tt_metal::flatbuffer::MathFidelity::LoFi: return MathFidelity::LoFi;
+        case tt::tt_metal::flatbuffer::MathFidelity::HiFi2: return MathFidelity::HiFi2;
+        case tt::tt_metal::flatbuffer::MathFidelity::HiFi3: return MathFidelity::HiFi3;
+        case tt::tt_metal::flatbuffer::MathFidelity::HiFi4: return MathFidelity::HiFi4;
+        case tt::tt_metal::flatbuffer::MathFidelity::Invalid: return MathFidelity::Invalid;
+    }
+}
+
+UnpackToDestMode from_flatbuffer(tt::tt_metal::flatbuffer::UnpackToDestMode input) {
+    switch (input) {
+        case tt::tt_metal::flatbuffer::UnpackToDestMode::UnpackToDestFp32: return UnpackToDestMode::UnpackToDestFp32;
+        case tt::tt_metal::flatbuffer::UnpackToDestMode::Default: return UnpackToDestMode::Default;
+    }
+}
+
+tt::DataFormat from_flatbuffer(tt::tt_metal::flatbuffer::DataFormat input) {
+    switch (input) {
+        case tt::tt_metal::flatbuffer::DataFormat::Float32: return tt::DataFormat::Float32;
+        case tt::tt_metal::flatbuffer::DataFormat::Float16: return tt::DataFormat::Float16;
+        case tt::tt_metal::flatbuffer::DataFormat::Bfp8: return tt::DataFormat::Bfp8;
+        case tt::tt_metal::flatbuffer::DataFormat::Bfp4: return tt::DataFormat::Bfp4;
+        case tt::tt_metal::flatbuffer::DataFormat::Bfp2: return tt::DataFormat::Bfp2;
+        case tt::tt_metal::flatbuffer::DataFormat::Float16_b: return tt::DataFormat::Float16_b;
+        case tt::tt_metal::flatbuffer::DataFormat::Bfp8_b: return tt::DataFormat::Bfp8_b;
+        case tt::tt_metal::flatbuffer::DataFormat::Bfp4_b: return tt::DataFormat::Bfp4_b;
+        case tt::tt_metal::flatbuffer::DataFormat::Bfp2_b: return tt::DataFormat::Bfp2_b;
+        case tt::tt_metal::flatbuffer::DataFormat::Lf8: return tt::DataFormat::Lf8;
+        case tt::tt_metal::flatbuffer::DataFormat::Fp8_e4m3: return tt::DataFormat::Fp8_e4m3;
+        case tt::tt_metal::flatbuffer::DataFormat::Int8: return tt::DataFormat::Int8;
+        case tt::tt_metal::flatbuffer::DataFormat::Tf32: return tt::DataFormat::Tf32;
+        case tt::tt_metal::flatbuffer::DataFormat::UInt8: return tt::DataFormat::UInt8;
+        case tt::tt_metal::flatbuffer::DataFormat::UInt16: return tt::DataFormat::UInt16;
+        case tt::tt_metal::flatbuffer::DataFormat::Int32: return tt::DataFormat::Int32;
+        case tt::tt_metal::flatbuffer::DataFormat::UInt32: return tt::DataFormat::UInt32;
+        case tt::tt_metal::flatbuffer::DataFormat::RawUInt8: return tt::DataFormat::RawUInt8;
+        case tt::tt_metal::flatbuffer::DataFormat::RawUInt16: return tt::DataFormat::RawUInt16;
+        case tt::tt_metal::flatbuffer::DataFormat::RawUInt32: return tt::DataFormat::RawUInt32;
+        case tt::tt_metal::flatbuffer::DataFormat::Invalid: return tt::DataFormat::Invalid;
+    }
+}
+
+Tile from_flatbuffer(const tt::tt_metal::flatbuffer::Tile* tile_fb) {
+    if (!tile_fb) {
+        throw std::runtime_error("Invalid Tile FlatBuffer object");
+    }
+
+    // Convert FlatBuffer vectors to std::array
+    std::array<uint32_t, 2> tile_shape = {tile_fb->tile_shape()->Get(0), tile_fb->tile_shape()->Get(1)};
+    std::array<uint32_t, 2> face_shape = {tile_fb->face_shape()->Get(0), tile_fb->face_shape()->Get(1)};
+
+    // Create and return the Tile object, explicitly initializing the members
+    Tile tile;
+    tile.tile_shape = tile_shape;
+    tile.face_shape = face_shape;
+    tile.tile_hw = tile_fb->tile_hw();
+    tile.face_hw = tile_fb->face_hw();
+    tile.num_faces = tile_fb->num_faces();
+    tile.partial_face = tile_fb->partial_face();
+    tile.narrow_tile = tile_fb->narrow_tile();
+    tile.transpose_within_face = tile_fb->transpose_within_face();
+    tile.transpose_of_faces = tile_fb->transpose_of_faces();
+
+    return tile;
+}
+
+std::array<std::optional<Tile>, NUM_CIRCULAR_BUFFERS> from_flatbuffer(
+    const flatbuffers::Vector<flatbuffers::Offset<tt::tt_metal::flatbuffer::Tile>>* tiles_fb) {
+    std::array<std::optional<Tile>, NUM_CIRCULAR_BUFFERS> tiles = {};
+    if (tiles_fb) {
+        for (size_t i = 0; i < tiles_fb->size() && i < NUM_CIRCULAR_BUFFERS; ++i) {
+            tiles[i] = from_flatbuffer(tiles_fb->Get(i));
+        }
+    }
+    return tiles;
+}
+
+}  // namespace tt::tt_metal
diff --git a/tt_metal/impl/flatbuffer/base_types_from_flatbuffer.hpp b/tt_metal/impl/flatbuffer/base_types_from_flatbuffer.hpp
@@ -0,0 +1,34 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "base_types_generated.h"
+#include <buffer_constants.hpp>
+#include <kernel_types.hpp>
+#include <data_types.hpp>
+#include <tt_backend_api_types.hpp>
+#include <tile.hpp>
+#include <circular_buffer_constants.h>
+
+namespace tt::tt_metal {
+
+BufferType from_flatbuffer(tt::tt_metal::flatbuffer::BufferType type);
+
+tt::tt_metal::DataMovementProcessor from_flatbuffer(tt::tt_metal::flatbuffer::DataMovementProcessor in);
+
+tt::tt_metal::NOC from_flatbuffer(tt::tt_metal::flatbuffer::NOC in);
+tt::tt_metal::NOC_MODE from_flatbuffer(tt::tt_metal::flatbuffer::NOC_MODE in);
+tt::tt_metal::Eth from_flatbuffer(tt::tt_metal::flatbuffer::Eth in);
+
+MathFidelity from_flatbuffer(tt::tt_metal::flatbuffer::MathFidelity input);
+UnpackToDestMode from_flatbuffer(tt::tt_metal::flatbuffer::UnpackToDestMode input);
+tt::DataFormat from_flatbuffer(tt::tt_metal::flatbuffer::DataFormat input);
+
+Tile from_flatbuffer(const tt::tt_metal::flatbuffer::Tile* tile_fb);
+
+std::array<std::optional<Tile>, NUM_CIRCULAR_BUFFERS> from_flatbuffer(
+    const flatbuffers::Vector<flatbuffers::Offset<tt::tt_metal::flatbuffer::Tile>>* tiles_fb);
+
+}  // namespace tt::tt_metal