[Feature Request] Add Light Metal capture/replay initial changes to tt-metal for some workloads

[kmabeeTT]

Light Metal Feature Parent Ticket: #17037

Request

It's a new feature - put together initial/boostrapping changes to:

• infra changes to add flatbuffers and schema files into tt-metal env, build changes, etc
• add some initial support for capturing / replaying various tt-metal host APIs
• show that we can capture/replay them from binary / standalone runner
• use "metal trace" to replay on-device portion
• add some initial unit tests that use enough capture/replay to be useful
• whatever else comes up that can be included in this round-1 issue/PR

Should have opened this ticket last year, have already started/completed this chunk of work, has gone through a few rounds of reviews, will reply with details to this ticket and roll out PR more globally shortly.

[kmabeeTT]

Here is high level overview of the changes I made for this ticket so far that I will roll out PR for. It basically adds limited (unsupported) support for tracing a bunch of popular APIs, and adds a few tests with correctness checking (reading output bufers, comparing between capture time and replay time, or to expected values) including some tests that use capture a program using "metal trace" and then replay it via LoadTrace() + ReplayTrace() APIs alongside the replay of all the captured host APIs.

Misc

• Rebasing regularly has been a bit tough, lots of refactors past month or so, eager to get this initial change merged. Rebased as recently as Jan 22.
• Attempted to use google-coding style where possible to have function names in PascalCase after brief chat with Wilder offline.

Initial Infra

• New Host APIs LightMetalBeginCapture(), LightMetalEndCapture(), LoadTrace(), see docs/commit for descriptions, but basically start/end surround workload to enable feature and LoadTrace is used during replay to transfer a metal-trace TraceDescriptor from flatbuffer binary to device.
• Introduce flatbuffers as CPM package, same version UMD uses, had to add CMAKE_POSITION_INDEPENDENT_CODE=ON to solve linker errors when bringing these changes to tt-mlir project, chatgpt said it made sense. Had to also demote a warning for possible gcc bug seen in g++12 build.
• Add few initial flatbuffer schema files (.fbs files) to represent flatbuffer binary, and host commands, and few commands for example
• In follow up commit, made compile-time define to guard this feature enabled by default, disabled with flag to build_metal.sh -> cmake -> C++. If things go well, no unexpected fallout, will remove compile time define, mainly here as quick way to have tracing macros become compile time NOPs (they are lightweight anyways).
• Programs traced with "metal trace" are captured in flatbuffer binary in order they are traced, and replayed using existing metal trace feature APIs.
• Standalone runner "lightmetal_runner" tool added as a way to run binaries as they currently exist. Example: ./build/tools/lightmetal_runner /tmp/SingleProgramTraceCapture.bin - e2e cpp unit tests also run directly from binary using same underlying replay executor that standalone runner tool uses.

Host APIs Traced

• Added TRACE_FUNCTION_CALL() macro to trace host API functions, and later TRACE_FUNCTION_ENTRY() to protect against recursively called host API functions (those that call others) so only the top level is captured.

Here is light from command.fbs, remove Command and these are pretty much the matching host APIs from host_api.hpp. With this came several new enums, structs, tables in types_*.fbs files and their corresponding serialization / deserialization functions in the form of ToFlatbuffer() and FromFlatbuffer() functions.

    ReplayTraceCommand,
    EnqueueTraceCommand,
    LoadTraceCommand,
    ReleaseTraceCommand,
    CreateBufferCommand,
    DeallocateBufferCommand,
    EnqueueWriteBufferCommand,
    EnqueueReadBufferCommand,
    FinishCommand,
    CreateProgramCommand,
    EnqueueProgramCommand,
    CreateKernelCommand,
    SetRuntimeArgsUint32Command,
    SetRuntimeArgsCommand,
    CreateCircularBufferCommand,
    LightMetalCompareCommand,

Note that when Metal Trace is enabled, don't capture EnqueueProgram(), instead inject ReplayTrace(), would be used alongside new API LoadTrace()

Testing

• Added new cpp unit test tt_metal/lightmetal/lightmetal_fixture.hpp and lightmetal_sanity.cpp which helped me test several features during development here
• Few very basic tests, but all have light-metal capture/end APIs called in setup/teardown respectively to capture a LightMetalBinary and replay it afterwards using the same replay library that is called from standalone runner.
• Have buffer write/read tests, datamovement program test, compute program test, compute program with metal-trace using single program, two programs. More to come in the future.
• Added APIs mostly for verif called LightMetalCompareToCapture() and LightMetalCompareToGolden() used in tests as way to binarize a buffer value at capture time, or a user-provided expected value, and compare the same buffer to these values at replay time (from light metal binary)
• The 6 tests added run capture + replay end-to-end and pass for grayskull, wormhole, and blackhole

Caveats/Restrictions

• Covers only tt-metal for now, no ttnn support yet (it has issues, bypasses host_api in several cases, need to follow up)
• Left a few TODO items with my username in code for things not well covered, will address in follow up ticket. Mainly related to not properly capturing/replaying same device/cq configuration, and hardcoding things there including CreateDevices() call.

[kmabeeTT]

Splitting this into some smaller PR's to make review easier by request.

• Round 1: cmake/flatbuffers/infra : LightMetal - Add Flatbuffers into cmake infra/build as cpm package (#17039) #17157
• Round 2: new capture begin/end APIs and LightMetalBinary class: LightMetal - New APIs LightMetalBeginCapture() and LightMetalEndCapture() and docs (#17039) #17262
• Round 3: new LoadTrace() API and move TraceDescriptor out of detail namespace: LightMetal - Add LoadTrace() API and move TraceDescriptor out of detail namespace (#17039) #17313
• Round 4: Flatbuffer schema files and to_flatbuffer(), from_flatbuffer() functions : LightMetal - Add flatbuffer schema and conversion to/from functions for various types (#17039) #17353
• Round 5: Host/Device capture library + unit tests : LightMetal - Initial Host API + Device Capture infra/library and unit tests (#17039) #17514
• Round 6: Replay library + e2e unit tests + standalone runner : LightMetal - Initial Replay infra/library for LightMetalBinary and standalone runner (#17039) #17524

[kmabeeTT]

This is now closed - these initial light-metal capture/replay changes are now merged across 6 PR's (see list in comment above) over last 10 days, after incorporating a bunch of feedback, was around 3600 lines/additions as reported by git. Special thanks to Oleg for thorough review. Will add more work into parent ticket next that is being scoped out (some of it has already begun).