LightMetal - Initial Host API + Device Capture infra/library and unit tests (#17039)

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

@@ -64,6 +64,7 @@ jobs:
           {name: stl, cmd: "./build/test/tt_metal/unit_tests_stl"},
           {name: distributed, cmd: "./build/test/tt_metal/distributed/distributed_unit_tests_${{ inputs.arch }} --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.*"},
 
           {name: ttnn cpp unit tests, cmd: ./build/test/ttnn/unit_tests_ttnn},

CMakeLists.txt

@@ -111,6 +111,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}")
 
 ############################################################################################################################
 
@@ -232,6 +233,13 @@ add_link_options(
     "$<$<BOOL:${ENABLE_UBSAN}>:-fsanitize=undefined>"
 )
 
+# Planned to be temporary, remove later.
+if(TT_ENABLE_LIGHT_METAL_TRACE)
+    add_compile_definitions(TT_ENABLE_LIGHT_METAL_TRACE=1)
+else()
+    add_compile_definitions(TT_ENABLE_LIGHT_METAL_TRACE=0)
+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/project_options.cmake

@@ -19,6 +19,7 @@ option(ENABLE_CCACHE "Build with compiler cache" FALSE)
 option(TT_UNITY_BUILDS "Build with Unity builds" ON)
 option(BUILD_TT_TRAIN "Enables build of tt-train" OFF)
 option(ENABLE_TTNN_SHARED_SUBLIBS "Use shared libraries for ttnn to speed up incremental builds" OFF)
+option(TT_ENABLE_LIGHT_METAL_TRACE "Enable Light Metal Trace" ON)
 
 ###########################################################################################
 

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.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,62 @@
+// 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 "command_queue_fixture.hpp"
+#include <lightmetal_binary.hpp>
+
+class SingleDeviceLightMetalFixture : public CommandQueueFixture {
+protected:
+    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 CreateDeviceAndBeginCapture(
+        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);
+        // TODO (kmabee) - revisit placement. CreateDevice() path calls CreateKernel() on programs not
+        // created with CreateProgram() traced API which leads to "program not in global_id map"
+        LightMetalBeginCapture();
+    }
+
+    // End light metal tracing, write to optional filename and optionally run from binary blob
+    void TearDown() override {
+        LightMetalBinary binary = LightMetalEndCapture();
+
+        if (binary.is_empty()) {
+            FAIL() << "Light Metal Binary is empty for test, unexpected.";
+        }
+        if (write_bin_to_disk_ && !this->trace_bin_path_.empty() && !binary.is_empty()) {
+            log_info(tt::LogTest, "Writing light metal binary {} bytes to {}", binary.size(), this->trace_bin_path_);
+            binary.save_to_file(this->trace_bin_path_);
+        }
+
+        if (!this->IsSlowDispatch()) {
+            tt::tt_metal::CloseDevice(this->device_);
+        }
+    }
+};