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

[kmabeeTT]

Ticket

I am breaking up original PR I put up recently into smaller bite-sized chunks as @omilyutin-tt suggested. Here is round 5.

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

Problem description

This is initial/bootstrapping changes for "Light Metal" capture/replay feature that uses Flatbuffers as serialization/deserialization library. See my previous bigger PR (#16573) if you want to see how this will be used by followup merge, or context.

What's changed

• This is round 5, builds upon previous 4 merges for LightMetal that brought flatbuffer cmake/infra, begin/end APIs, LoadTrace() API, flatbuffer/schema serialization/deserialization

• This adds light-metal Capture support for and instruments with LIGHT_METAL_TRACE_FUNCTION_CALL() and LIGHT_METAL_TRACE_FUNCTION_ENTRY() to many popular (not exuahstive) APIs used by unit tests. The former TRACE_FUNCTION_ENTRY() is more recent, used to protect against host APIs recursively calling other host APIs (only trace top most level). Two macros not always called back-to-back.

• Support Capture/Replay of the following ~14 host APIs

  EnqueueTrace(), ReplayTrace(), ReleaseTrace() CreateBuffer(), EnqueueWriteBuffer(), EnqueueReadBuffer(), DeallocateBuffer CreateKernel(), CreateCircularBuffer() SetRuntimeArgs(uint32) SetRuntimeArgs(Kernel,RuntimeArgs) CreateProgram(), EnqueueProgram() Finish()

• During capture, complex objects like Programs, Kernels, Buffers, CBHandle are assigned unique global_id, and referred to by their global_id in capture when used by functions

• When "Metal Trace" is enabled, don't capture EnqueueProgram(), instead inject ReplayTrace(), would be used alongside LoadTrace()

• Can be optionally disabled at compile time using build_metal.sh --disable-light-metal-trace which will set C++ define TT_ENABLE_LIGHT_METAL_TRACE=0 (trace functions become NOP)

• New Verif APIs LightMetalCompareToCapture() / LightMetalCompareToGolden(). Put them in lightmetal_capture_utils.hpp instead of host_api.hpp since they are purely used at capture time, and not worthy enough to be inside host_api.h since just for verif

• Test fixture hardcoded to run capture-only and skip replay until replay code is merged next

Checklist

• Build and docs and new unit tests passing for wh, gs, bh on original PR, will check again
• Post commit CI passes
• Blackhole Post commit (if applicable)
• Model regression CI testing passes (if applicable)
• Device performance regression CI testing passes (if applicable)
• (For models and ops writers) Full new models tests passes
• New/Existing tests provide coverage for changes

[kmabeeTT]

Whoops, had a couple dependencies on lightmetal_replay.hpp (included in next PR) that I forgot to purge, going to force push fixed commit.

[kmabeeTT]

cmake/C++ define temporary for next little bit, just here as quick way to disable feature from compile time. Plan to remove this in a few weeks once it's settled.

[kmabeeTT]

General note: applied all previous feedback from parent PR to this test, and renamed it.

[omilyutin-tt]

This pattern performs a double look up in a map, consider getting rid of is_in_map functions (unless we need them to be public?) and doing something like:

void LightMetalCaptureContext::remove_from_map(const T* obj) {
  if (auto it = map_.find(obj); it != map_.end()) {
    map_->erase(it);
  } else {
    log_warning(...);
  }
}
...
void LightMetalCaptureContext::add_to_map(const T* obj) {
  auto [it, inserted] = map_.emplace(obj, next_global_id_ + 1);
  if (inserted) {
    ++next_global_id_;
  } else {
    log_warning(...);
  }
}

As I was writing it, I realized we can have a helper template for these methods.

[kmabeeTT]

Yes is_in_map() does have a public usage today. I'll have to think about this, may defer some of this item to later.

[omilyutin-tt]

Can we split this into 1. creation of a device and 2. starting light metal capture? Perhaps move the device creation into SetUp?

[kmabeeTT]

Like some other gtests I based mine off, prefer to keep CreateDevice() logic seperate from Setup() so I can tweak settings in test (SetUp() not called by tests, happens automatically through gtest framework afaik). I will try to move LightMetalBeginCapture() call to SetUp() to be consistent with how LightMetalEndCapture() is automatically called inTearDown().

[kmabeeTT]

Scratch that, I can't move it - I forgot that the placement right now (LightMetalBeginCapture() after CreateDevice() gets called) is intentional. Will add a comment, but basically CreateDevices() calls CreateKernel(program) on a program object that wasn't created from CreateProgram() and so it exposes this failure:

             MetalTrace | TRACE    | /localdev/kmabee/metal1/tt_metal/tt_metal.cpp:1043 - LIGHT_METAL_TRACE_FUNCTION_CALL: CaptureCreateKernel via CreateKernel istracing: true depth: 1
libc++abi: terminating due to uncaught exception of type std::runtime_error: Program not found in global_id map.

A task for later to improve things here, and follow up on fixes.

[omilyutin-tt]

Can remove now?

[kmabeeTT]

Yes, but prefer to keep, it's something I am playing with in test, and need to revisit in next round of tests.

[omilyutin-tt]

Can these two be specified by a single env variable that specifies trace_bin_path? If it is empty - don't trace anything, if not - trace and output into the value of the variable?

Alternatively, this can be boolean, and we can always write to a random location in /tmp. It's just the current method requires re-compilation to specify a custom path, and would be great to get away with just one state variable.

[kmabeeTT]

I iterated on this a bit last year before ending up in current state. It might change in the future as tests evolve, but for now prefer to keep it as is because:

1. plan to have a test soon that generates multiple binary files with specific names, and runs them all in sequence
2. trace_bin_path being empty is what triggers meaningful name of output file based on testname, and it is nice currently for debug/dev to have consistent names of the .bin file based testname , plus a binary-to-json format flow (through flatc tool) tends to write to an output file name based on the input file name
3. Never found a need to run a test or test(s) and force specific path for a run via env-var.

[omilyutin-tt]

Now used? Was this for Enqueue* APIs? Consider annotating like:

EnqueueWriteBuffer(command_queue, *buffer, input_data.data(), /*blocking=*/ true);

[kmabeeTT]

Yeah right, it was for the Enqueue*() APIs, I need to revisit it in test later, will add a TODO, and update annotations like you mentioned.

[omilyutin-tt]

We are making a copy here, right? This is interpreted as calling a constructor from a range?

Is it possible that we have something else at top level, instead of void*?

[kmabeeTT]

Yes and yes, flatbuffer API needs it in vector form. Don't think we should change it from void* golden_data because this function is called after/alongside and matches by design the interface of EnqueueReadBuffer(cq, buffer, dst) where dst is void * to give same user experience:

void LightMetalCompareToCapture(
    CommandQueue& cq, 
    const std::variant<std::reference_wrapper<Buffer>, 
    std::shared_ptr<Buffer>>& buffer, 
    void* dst) {
    
<snip>

    EnqueueReadBuffer(cq, buffer, dst, true);  // Blocking read to get golden value.
    LIGHT_METAL_TRACE_FUNCTION_CALL(CaptureLightMetalCompare, cq, buffer, dst, false);
}

[omilyutin-tt]

Yeah, these void* we use are really annoying and unsafe, I would really love to modernize this. Thats for later though:)

[kmabeeTT]

Hi @omilyutin-tt - re-review please - I rebased on newer passing commit, force pushed, and included a commit addressing most of the change requests, but left some for later that will require more time/work/thought, to hopefully get this in. I have a conflict on device.cpp after Joseph's changes in main that I didn't include in rebase yet, will do some build testing on this current rebase, then fix conflict and force push again later.

[omilyutin-tt]

Yeah, these void* we use are really annoying and unsafe, I would really love to modernize this. Thats for later though:)

[kmabeeTT]

all post-commit run again overnight, just force pushed a fixup to convert a few throw to TT_THROW, just waiting for @tenstorrent/metalium-developers-infra review for cmake change. It was me adding temporary C++ define (enabled by default) to wrap feature and build_metal.sh -> cmake -> c++ define to disable it, as discussed with Wilder/Andrew offline few weeks ago. Plan to remove it in a few weeks if no issues.

[afuller-TT]

Discussed offline about moving the conditional to an option() in CMake so ingress via paths other than the shell script also default to on, but devs can opt-out if needed.

Approving in anticipation of that change.

[kmabeeTT]

Thanks guys, made last feedback re cmake, rebased on near TOTT can compile and run my tests locally, should be good, merge in a few.