Op Logging in MAX

Op logging is a diagnostic feature that allows developers to trace the execution of operations (ops) in the MAX runtime. When enabled, it outputs detailed information about each operation's launch and completion to stderr, helping with debugging and performance analysis.

Enabling Op Logging

Op logging is disabled by default and can be enabled in several ways:

1. Via Python InferenceSession API

When using the MAX Python API, set the mojo logging level via set_mojo_log_level on the InferenceSession object:

from max.engine import InferenceSession, LogLevel

# Enable op logging for this session
session = InferenceSession()
session.set_mojo_log_level(LogLevel.TRACE)

2. Via Bazel Build Flag

Enable op logging by running your command with the --config=mojo-trace flag:

./bazelw run --config=mojo-trace //your:target
./bazelw test --config=mojo-trace //your:test

3. Via Mojo Compilation Flag

When compiling Mojo code directly, pass the compilation define:

mojo -D LOGGING_LEVEL=trace your_file.mojo

How Op Logging Works

Op logging leverages Mojo's tracing infrastructure to emit log messages for any Trace statements that use TraceLevel.OP. When an operation starts and completes, it produces structured log output with unique IDs and optional target device details.

Log Output Format

Example output:

[OP] LAUNCH elementwise [id=40028] target=cpu:0
[OP] COMPLETE elementwise [id=40028] target=cpu:0
[OP] LAUNCH rms_norm [id=40029] target=gpu:0
[OP] COMPLETE rms_norm [id=40029] target=gpu:0

Implementation Details

Mojo Tracing Infrastructure

Op logging is implemented using the Trace struct from runtime.tracing with TraceLevel.OP:

from std.runtime.tracing import Trace, TraceLevel

# Create a trace for an operation
with Trace[TraceLevel.OP](op_name, detail, task_id=task_id):
    # Operation implementation to log

Unique ID Generation

Each operation gets a unique, incrementing ID generated by the runtime. This ID is used to correlate LAUNCH and COMPLETE events for the same operation.

Target Information

When available, traces can include target device information (including the device ID) that gets included in the log output:

# Trace with target information
with Trace[TraceLevel.OP, target=StaticString("gpu")](op_name, task_id:Int(ctx.id())):
    # Operation implementation to log

If you see an op log without a target or a device ID, it is likely that this information is not provided in the Trace call. The device target name is passed in the target parameter and the device ID is passed as the task_id argument. Adding these to the Trace call should allow them to show up in the log.