Tt Lang To Tt Metal

Prerequisites

Ensure you can run TT-Lang kernels on hardware. If using remote tools, the tt-connect-remote-device skill provides run-test.sh, copy-file.sh, copy-from-remote.sh, and remote-run.sh. If running directly on a machine with HW access, use python directly.

Task

Export a TT-Lang kernel to standalone TT-Metal C++ code with a Python entry point using ttnn.generic_op. The primary goal is a working, correct kernel that can run independently of ttlang.

Input

$ARGUMENTS

Priorities

P1 - Must Work: Generate correct C++ kernels and a working Python runner P2 - Readability: Rename generated variables (v1 -> input_lhs), collapse redundant casts P3 - Low Priority: Use clearer APIs only if huge clarity win with zero risk

Non-goals: Performance optimizations, architectural changes, risky transformations

Process

Step 1: Compile the TT-Lang Kernel

First, verify the TT-Lang kernel works by running it on hardware:

# Via run-test.sh:
run-test.sh --hw /path/to/kernel.py

# Or directly:
python /path/to/kernel.py

If the kernel fails, STOP and tell the user to fix it before proceeding.

Then run with TTLANG_EMIT_RUNNER=1 to generate the C++ kernels and Python runner:

# Via run-test.sh:
run-test.sh --hw --emit-runner /path/to/kernel.py

# Or directly:
TTLANG_EMIT_RUNNER=1 python /path/to/kernel.py

This generates files in /tmp/$USER/ on the remote (one set per @ttl.kernel function):

• ttlang_kernel_compute_<hash>.cpp - Compute kernel
• ttlang_kernel_dm_read_<hash>.cpp - Data movement reader
• ttlang_kernel_dm_write_<hash>.cpp - Data movement writer
• ttlang_kernel_compute_<hash>_runner.py - Python runner with all CB/tensor setup

Find the exact file paths in the output. If using run-test.sh, grep the log:

remote-run.sh cat /tmp/ttlang_test_output.log | grep -i "written"

If running directly, pipe the output or redirect it to a file to search:

TTLANG_EMIT_RUNNER=1 python /path/to/kernel.py 2>&1 | grep -i "written"

NOTE: if there are multiple kernels, you will have a copy of all of these files per kernel (that's OK! many tt-lang programs utilize multiple kernels)

Copy the generated files to inspect them (if running remotely, use copy-from-remote.sh; if running locally, the files are already at /tmp/$USER/):

copy-from-remote.sh /tmp/$USER/ttlang_kernel_compute_<hash>.cpp ./
copy-from-remote.sh /tmp/$USER/ttlang_kernel_compute_<hash>_runner.py ./

Step 2: Create the Output Directory Structure

my_kernel/
├── kernels/
│   ├── compute.cpp      # Compute kernel
│   ├── reader.cpp       # Data movement reader
│   └── writer.cpp       # Data movement writer
└── run_kernel.py        # Python entry point

The run kernel.py should be an adapted version of the program generated by TTLANG_EMIT_RUNNER=1. You should fill in the places that need to be filled in using the main function from the ttlang program to guide you.

Step 3: Test with the New Entry Point

Test the runner you created using the compiler-generated C++ kernels (still in /tmp/$USER/):

# Via run-test.sh:
run-test.sh --hw /path/to/my_kernel/run_kernel.py

# Or directly:
python /path/to/my_kernel/run_kernel.py

If it fails, check:

1. Tensor order matches what the kernel expects (see KERNEL_TENSOR_INDICES in generated runner)
2. The main() function setup matches the original ttlang program
3. Grid dimensions are correct

Iterate until the output matches the original ttlang program's behavior.

Step 4: Copy temporary files and test again

At this point you should have only been using the ttlang-generated cpp kernels, now copy them to the working directory in the structure above so that you can make modifications in the next step. Before you start modifying, make sure you have the remote copying flow down. Make sure the paths point to the right place in the kernel.py runner and you are able to modify the cpp kernels locally, copy them over, and test with the new changes.

Concretely:

1. copy the files locally
2. update cpp kernel paths in your python entry point
3. copy back to the remote and test with new paths

Step 5: P2 Beautification (Safe Transformations Only)

Naming note: TT-Lang calls them "dataflow buffers" (DFBs) at the Python level, but the generated TT-Metal C++ code uses the name "circular buffer" (CB). They are the same thing, just different names at different abstraction levels.

IMPORTANT: only apply one transformation at a time. Running the kernel is cheap. You can iterate as much as you need to. Make targeted changes, test, and rollback if they break something.

Apply ONLY these safe, mechanical transformations:

1. Rename generated variables:

// Before (generated)
int32_t v1 = ...;
int32_t v2 = ...;

// After (readable)
int32_t input_lhs = ...;
int32_t input_rhs = ...;

Note: make sure you read the input tt-lang python kernel (the one that you ran through the compiler). USE THIS TO guide variable names, comments, etc.

2. Collapse redundant casts:

// Before (generated)
int32_t v1 = (int32_t)(int32_t)get_compile_time_arg_val(0);

// After
int32_t cb_lhs = get_compile_time_arg_val(0);

3. Name CB indices meaningfully:

// Before
cb_wait_front(get_compile_time_arg_val(0), 1);
cb_wait_front(get_compile_time_arg_val(1), 1);

// After (add constexpr at top, keep the call identical)
constexpr uint32_t cb_lhs = 0;
constexpr uint32_t cb_rhs = 1;
constexpr uint32_t cb_out = 2;

cb_wait_front(cb_lhs, 1);
cb_wait_front(cb_rhs, 1);

4. Add comments:

• Add comments from the input tt-lang program
• Add other comments that describe how the kernel works
• Comments are zero risk and help clarify what's happening. Make your comments succinct. Only comment what's not immediately obvious from the code.

Bad comment (no extra context provided):

// read from CB 0
cb_wait_front(0, 1);

Good comment (succinct, provides context not obvious):

cb_wait_front(/*lhs*/0, 1);

DO NOT change:

• Loop structures or bounds
• API calls or their arguments (except variable names)
• Order of operations
• Synchronization patterns (barriers, waits, etc.)
• Any logic or control flow

Step 6: Copy Files and Test

The exported kernel must run on a machine with HW access. If using remote tools, copy files to the remote first. If running locally, ensure the kernel paths in run_kernel.py are correct.

File Structure

The kernel paths in run_kernel.py must match where files actually are.

# If using remote tools, copy files:
copy-file.sh my_kernel/run_kernel.py
copy-file.sh my_kernel/kernels/compute.cpp kernels/
copy-file.sh my_kernel/kernels/reader.cpp kernels/
copy-file.sh my_kernel/kernels/writer.cpp kernels/

Update Kernel Paths

The kernel_source paths in run_kernel.py must match where files are on the remote:

# Option 1: Relative paths (if run_kernel.py is in same dir as kernels/)
reader_kernel = ttnn.KernelDescriptor(
    kernel_source="kernels/reader.cpp",  # Relative to run_kernel.py location
    ...
)

# Option 2: Absolute paths on remote
reader_kernel = ttnn.KernelDescriptor(
    kernel_source="/tmp/kernels/reader.cpp",
    ...
)

Run the Test

# Via run-test.sh:
run-test.sh --hw /path/to/my_kernel/run_kernel.py

# Or directly:
python /path/to/my_kernel/run_kernel.py

# Check results (via remote tools or directly):
tail -50 /tmp/ttlang_test_output.log  # if using run-test.sh

Debugging Failures

If it fails, check:

1. CB indices in compute kernel match the CBDescriptor buffer_index values
2. Tensor order in ttnn.generic_op([...]) matches kernel expectations
3. Grid dimensions match between Python and original TTLang kernel
4. Compile-time args order matches get_compile_time_arg_val(N) usage
5. Kernel file paths are correct for the remote filesystem

A working kernel with ugly variable names is better than a broken kernel with nice names.

You MUST:

1. Verify the original TT-Lang kernel works before starting
2. Ensure all files are accessible on the machine with HW (copy to remote if needed)
3. Run the exported kernel on hardware (via run-test.sh --hw or python directly)
4. Read the output and confirm it runs correctly
5. Only mark complete after successful execution on hardware
6. Only make targeted small changes to the kernel and test each change, iterate as much as you need to