Add Light Metal capture/replay initial changes to tt-metal for some workloads

.github/workflows/cpp-post-commit.yaml

@@ -63,6 +63,7 @@ jobs:
           {name: llk, cmd: "./build/test/tt_metal/unit_tests_llk"},
           {name: stl, cmd: "./build/test/tt_metal/unit_tests_stl"},
           {name: distributed, cmd: "./build/test/tt_metal/distributed/distributed_unit_tests --gtest_filter=MeshDeviceSuite.*"},
+          {name: lightmetal, cmd: "./build/test/tt_metal/unit_tests_lightmetal"},
 
           {name: dispatch multicmd queue, cmd: "TT_METAL_GTEST_NUM_HW_CQS=2 ./build/test/tt_metal/unit_tests_dispatch_${{ inputs.arch }} --gtest_filter=MultiCommandQueue*Fixture.*"},
 

CMakeLists.txt

@@ -107,6 +107,7 @@ message(STATUS "Build TT METAL Tests: ${TT_METAL_BUILD_TESTS}")
 message(STATUS "Build TTNN Tests: ${TTNN_BUILD_TESTS}")
 message(STATUS "Build with Unity builds: ${TT_UNITY_BUILDS}")
 message(STATUS "Build with Shared TTNN Sublibraries: ${ENABLE_TTNN_SHARED_SUBLIBS}")
+message(STATUS "Build with LightMetal Trace Enabled: ${TT_ENABLE_LIGHT_METAL_TRACE}")
 
 ############################################################################################################################
 
@@ -228,6 +229,10 @@ add_link_options(
     "$<$<BOOL:${ENABLE_UBSAN}>:-fsanitize=undefined>"
 )
 
+if(TT_ENABLE_LIGHT_METAL_TRACE)
+    add_compile_definitions(TT_ENABLE_LIGHT_METAL_TRACE)
+endif()
+
 if(ENABLE_CODE_TIMERS)
     add_compile_definitions(TT_ENABLE_CODE_TIMERS)
 endif()

build_metal.sh

@@ -29,6 +29,7 @@ show_help() {
     echo "  --clean                          Remove build workspaces."
     echo "  --build-static-libs              Build tt_metal (not ttnn) as a static lib (BUILD_SHARED_LIBS=OFF)"
     echo "  --disable-unity-builds           Disable Unity builds"
+    echo "  --disable-light-metal-trace      Disable Light Metal tracing to binary."
     echo "  --cxx-compiler-path              Set path to C++ compiler."
     echo "  --c-compiler-path                Set path to C++ compiler."
     echo "  --ttnn-shared-sub-libs           Use shared libraries for ttnn."
@@ -58,6 +59,7 @@ build_programming_examples="OFF"
 build_tt_train="OFF"
 build_static_libs="OFF"
 unity_builds="ON"
+light_metal_trace="ON"
 build_all="OFF"
 cxx_compiler_path=""
 c_compiler_path=""
@@ -88,6 +90,7 @@ build-programming-examples
 build-tt-train
 build-static-libs
 disable-unity-builds
+disable-light-metal-trace
 release
 development
 debug
@@ -155,6 +158,8 @@ while true; do
             ttnn_shared_sub_libs="ON";;
         --disable-unity-builds)
 	    unity_builds="OFF";;
+        --disable-light-metal-trace)
+            light_metal_trace="OFF";;
         --cxx-compiler-path)
             cxx_compiler_path="$2";shift;;
         --c-compiler-path)
@@ -218,6 +223,7 @@ echo "INFO: Install Prefix: $cmake_install_prefix"
 echo "INFO: Build tests: $build_tests"
 echo "INFO: Enable Unity builds: $unity_builds"
 echo "INFO: TTNN Shared sub libs : $ttnn_shared_sub_libs"
+echo "INFO: Enable Light Metal Trace: $light_metal_trace"
 
 # Prepare cmake arguments
 cmake_args+=("-B" "$build_dir")
@@ -308,6 +314,12 @@ else
     cmake_args+=("-DTT_UNITY_BUILDS=OFF")
 fi
 
+if [ "$light_metal_trace" = "ON" ]; then
+    cmake_args+=("-DTT_ENABLE_LIGHT_METAL_TRACE=ON")
+else
+    cmake_args+=("-DTT_ENABLE_LIGHT_METAL_TRACE=OFF")
+fi
+
 if [ "$build_all" = "ON" ]; then
     cmake_args+=("-DTT_METAL_BUILD_TESTS=ON")
     cmake_args+=("-DTTNN_BUILD_TESTS=ON")

cmake/flatbuffers.cmake

@@ -0,0 +1,17 @@
+# Function to generate FlatBuffers C++ headers from schema files
+function(GENERATE_FBS_HEADER FBS_FILE)
+    get_filename_component(FBS_FILE_NAME ${FBS_FILE} NAME)
+    set(FBS_GENERATED_HEADER_DIR "${CMAKE_CURRENT_BINARY_DIR}/flatbuffers")
+    set(FBS_GENERATED_HEADER_FILE "${FBS_GENERATED_HEADER_DIR}/${FBS_FILE_NAME}_generated.h")
+    add_custom_command(
+        OUTPUT
+            ${FBS_GENERATED_HEADER_FILE}
+        COMMAND
+            flatc --cpp --scoped-enums -I ${CMAKE_CURRENT_SOURCE_DIR} -o "${FBS_GENERATED_HEADER_DIR}" ${FBS_FILE}
+        DEPENDS
+            flatc
+            ${FBS_FILE}
+        COMMENT "Building C++ header for ${FBS_FILE}"
+    )
+    set(FBS_GENERATED_HEADER_FILE ${FBS_GENERATED_HEADER_FILE} PARENT_SCOPE)
+endfunction()

dependencies/CMakeLists.txt

@@ -140,3 +140,25 @@ endif()
 
 CPMAddPackage(NAME taskflow GITHUB_REPOSITORY taskflow/taskflow GIT_TAG v3.7.0)
 add_library(Taskflow::Taskflow ALIAS Taskflow)
+
+############################################################################################################################
+# flatbuffers : https://github.com/google/flatbuffers
+############################################################################################################################
+
+CPMAddPackage(
+    NAME flatbuffers
+    GITHUB_REPOSITORY google/flatbuffers
+    GIT_TAG v24.3.25
+    OPTIONS
+        "FLATBUFFERS_BUILD_FLATC ON"
+        "FLATBUFFERS_BUILD_TESTS OFF"
+        "FLATBUFFERS_SKIP_MONSTER_EXTRA ON"
+        "FLATBUFFERS_STRICT_MODE ON"
+        "CMAKE_POSITION_INDEPENDENT_CODE ON"
+)
+
+if(flatbuffers_ADDED)
+    # Few files including idl_gen_dart.cpp:175:18, Possibly related: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=105329
+    target_compile_options(flatc PRIVATE -Wno-restrict)
+    target_compile_options(flatbuffers PRIVATE -Wno-restrict)
+endif()

docs/source/tt-metalium/tt_metal/apis/host_apis/command_queue/LightMetalBeginCapture.rst

@@ -0,0 +1,4 @@
+LightMetalBeginCapture
+======================
+
+.. doxygenfunction:: tt::tt_metal::v0::LightMetalBeginCapture

docs/source/tt-metalium/tt_metal/apis/host_apis/command_queue/LightMetalEndCapture.rst

@@ -0,0 +1,4 @@
+LightMetalEndCapture
+====================
+
+.. doxygenfunction:: tt::tt_metal::v0::LightMetalEndCapture

docs/source/tt-metalium/tt_metal/apis/host_apis/command_queue/LoadTrace.rst

@@ -0,0 +1,4 @@
+LoadTrace
+=========
+
+.. doxygenfunction:: tt::tt_metal::v0::LoadTrace

docs/source/tt-metalium/tt_metal/apis/host_apis/command_queue/command_queue.rst

@@ -16,5 +16,8 @@ CommandQueue
   ReplayTrace
   ReleaseTrace
   EnqueueTrace
+  LoadTrace
+  LightMetalBeginCapture
+  LightMetalEndCapture
   Finish
   Synchronize

tests/tt_metal/tt_metal/CMakeLists.txt

@@ -69,6 +69,7 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/llk)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/perf_microbenchmark)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/stl)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/noc)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/lightmetal)
 
 add_custom_target(
     metal_tests
@@ -92,4 +93,5 @@ add_custom_target(
         unit_tests_llk
         unit_tests_stl
         unit_tests_noc
+        unit_tests_lightmetal
 )

tests/tt_metal/tt_metal/lightmetal/CMakeLists.txt

@@ -0,0 +1,23 @@
+set(UNIT_TESTS_LIGHTMETAL_SRC ${CMAKE_CURRENT_SOURCE_DIR}/test_lightmetal_sanity.cpp)
+
+add_executable(unit_tests_lightmetal ${UNIT_TESTS_LIGHTMETAL_SRC})
+TT_ENABLE_UNITY_BUILD(unit_tests_lightmetal)
+
+target_link_libraries(unit_tests_lightmetal PUBLIC test_metal_common_libs)
+
+target_include_directories(
+    unit_tests_lightmetal
+    PRIVATE
+        "$<TARGET_PROPERTY:Metalium::Metal,INCLUDE_DIRECTORIES>"
+        ${PROJECT_SOURCE_DIR}/tests
+        ${PROJECT_SOURCE_DIR}/tests/tt_metal/tt_metal/common
+        ${CMAKE_CURRENT_SOURCE_DIR}
+        ${CMAKE_CURRENT_SOURCE_DIR}/common
+)
+
+set_target_properties(
+    unit_tests_lightmetal
+    PROPERTIES
+        RUNTIME_OUTPUT_DIRECTORY
+            ${PROJECT_BINARY_DIR}/test/tt_metal
+)

tests/tt_metal/tt_metal/lightmetal/lightmetal_fixture.hpp

@@ -0,0 +1,80 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "dispatch_fixture.hpp"
+#include <tt-metalium/device_impl.hpp>
+#include <tt-metalium/hal.hpp>
+#include <tt-metalium/host_api.hpp>
+#include <tt-metalium/tt_metal.hpp>
+#include <circular_buffer_constants.h>
+#include <tt-metalium/kernel.hpp>
+#include <tt-metalium/tt_backend_api_types.hpp>
+#include "lightmetal_replay.hpp"
+#include "command_queue_fixture.hpp"
+#include "lightmetal_types.hpp"
+
+class SingleDeviceLightMetalFixture : public CommandQueueFixture {
+protected:
+    bool replay_binary_;
+    std::string trace_bin_path_;
+    bool write_bin_to_disk_;
+
+    void SetUp() override {
+        this->validate_dispatch_mode();
+        this->arch_ = tt::get_arch_from_string(tt::test_utils::get_umd_arch_name());
+    }
+
+    void CreateDevice(
+        const size_t trace_region_size, const bool replay_binary = true, const std::string trace_bin_path = "") {
+        // Skip writing to disk by default, unless user sets env var for local testing
+        write_bin_to_disk_ = tt::parse_env("LIGHTMETAL_SAVE_BINARY", false);
+
+        // If user didn't provide a specific trace bin path, set a default here based on test name
+        if (trace_bin_path == "") {
+            const auto test_info = ::testing::UnitTest::GetInstance()->current_test_info();
+            auto trace_filename = test_info ? std::string(test_info->name()) + ".bin" : "lightmetal_trace.bin";
+            this->trace_bin_path_ = "/tmp/" + trace_filename;
+        }
+
+        this->create_device(trace_region_size);
+        this->replay_binary_ = replay_binary && !tt::parse_env("LIGHTMETAL_DISABLE_RUN", false);
+        LightMetalBeginCapture();
+    }
+
+    // End light metal tracing, write to optional filename and optionally run from binary blob
+    void TearDown() override {
+        LightMetalBinary binary_blob = LightMetalEndCapture();
+
+        if (binary_blob.IsEmpty()) {
+            FAIL() << "Light Metal Binary is empty for test, unexpected.";
+        }
+        if (write_bin_to_disk_ && !this->trace_bin_path_.empty() && !binary_blob.IsEmpty()) {
+            log_info(
+                tt::LogTest, "Writing light metal binary {} bytes to {}", binary_blob.Size(), this->trace_bin_path_);
+            binary_blob.SaveToFile(this->trace_bin_path_);
+        }
+
+        if (!this->IsSlowDispatch()) {
+            tt::tt_metal::CloseDevice(this->device_);
+        }
+
+        // We could gaurd this to not attempt to replay empty binary, and still allow test to pass
+        // but, would rather catch the case if the feature gets disabled at compile time.
+        if (replay_binary_) {
+            RunLightMetalBinary(std::move(binary_blob));
+        }
+    }
+
+    // Mimic the light-metal standalone run replay tool by executing the binary.
+    void RunLightMetalBinary(LightMetalBinary&& binary_blob) {
+        tt::tt_metal::LightMetalReplay lm_replay(std::move(binary_blob));
+        if (!lm_replay.ExecuteLightMetalBinary()) {
+            FAIL() << "Light Metal Binary failed to execute or encountered errors.";
+        } else {
+            log_info(tt::LogMetalTrace, "Light Metal Binary executed successfully!");
+        }
+    }
+};