Memory Debugging Guide: TT-Metal Bank Manager

[bmalesevicTT]

🐞 Background

While running a benchmark model, I ran into an L1 memory leakage between two loop iterations. By enabling Bank Manager logging in TT-Metal, dumping the logs, and comparing runs, I was able to locate the issue:
➡️ the output of the entire model was allocated in L1 but never deallocated.

Since this logging approach might help in similar situations, here’s a rough guide with snippets and notes.

🔧 Setup

To see TT-Metal Bank Manager logs from MLIR / Forge-FE, you may need to set the environment variable:

export TT_METAL_LOGGER_LEVEL=info
# check it with:
echo $TT_METAL_LOGGER_LEVEL

📋 Overview

The Bank Manager in TT-Metal provides the foundation for tracking memory allocation, deallocation, and usage patterns across different memory types (L1, L1_SMALL, DRAM). This guide shows you how to add comprehensive logging capabilities for effective memory debugging.

🎯 Available Memory Information

1. Basic Allocation Details
   When implementing allocation logging, you can capture:
   - Buffer size and type (L1, L1_SMALL, DRAM)
   - Number of banks and bytes per bank
   - Allocated address
   - Sharding status (sharded vs interleaved)

2. Memory Statistics via Allocator
   The underlying allocator provides access to:
   - Total allocated bytes (stats.total_allocated_bytes)
   - Total free bytes (stats.total_free_bytes)
   - Largest free block (stats.largest_free_block_bytes)
   - Memory block table (allocator_->get_memory_block_table())
   - Block-level dumps (allocator_->dump_blocks())

3. Bank-Level Information
   Through BankManager methods:
   - Bank count (num_banks())
   - Bank size (bank_size())
   - Bank offsets (bank_offset(bank_id))
   - Lowest occupied address per bank (lowest_occupied_address())

🔧 Implementing Memory Debugging

Step 1: Add Statistics Helper Method
Add a helper method to BankManager to format comprehensive statistics:

// Add to BankManager class (private section)
std::string get_statistics_string() {
    if (!allocator_) {
        return "allocator_not_initialized";
    }
    
    auto stats = allocator_->get_statistics();
    auto memory_table = allocator_->get_memory_block_table();
    
    std::stringstream ss;
    ss << "allocated=" << stats.total_allocated_bytes << "B, "
       << "free=" << stats.total_free_bytes << "B, "
       << "blocks=" << memory_table.size() << ", "
       << "fragmentation=" << std::fixed << std::setprecision(2)
       << (stats.total_free_bytes > 0 ? 
           (float)stats.largest_free_block_bytes / stats.total_free_bytes * 100 : 0)
       << "%";
    
    // Optionally include detailed block dump
    allocator_->dump_blocks(ss);
    
    return ss.str();
}

⚠️ Note: Logged addresses differ from the memory tables (offsets: +32B for DRAM, +99,776 for L1). Logs show data location, tables show metadata + data.

Step 2: Enhanced Allocation Logging
Add comprehensive logging to the allocate_buffer method:

// In BankManager::allocate_buffer(), after successful allocation
log_info(
    tt::LogAlways,
    "Allocated {} B {} buffer across {} banks, where each bank needs to store {} B\n"
    "Address: {}\n"
    "Memory Stats: {}",
    size,
    enchantum::to_string(buffer_type_),
    num_banks,
    size_per_bank,
    address.has_value() ? address.value() : 0,
    get_statistics_string());

Step 3: Deallocation Logging
Add logging to track memory cleanup:

// In BankManager::deallocate_buffer(), after deallocation
log_info(
    tt::LogAlways,
    "Deallocated address: {}\n"
    "Memory Stats: {}",
    address,
    get_statistics_string());

Step 4: Bulk Cleanup Logging
Add logging for bulk operations:

// In BankManager::deallocate_all()
log_info(
    tt::LogAlways,
    "Deallocating all {} buffers\n"
    "Pre-cleanup stats: {}",
    allocated_buffers_.size(),
    get_statistics_string());

// ... perform deallocation ...

log_info(
    tt::LogAlways,
    "Completed bulk deallocation\n"
    "Post-cleanup stats: {}",
    get_statistics_string());