torch_logs.py

"""
(beta) Using TORCH_LOGS python API with torch.compile
==========================================================================================
**Author:** `Michael Lazos <https://github.com/mlazos>`_
"""

import logging

######################################################################
#
# This tutorial introduces the ``TORCH_LOGS`` environment variable, as well as the Python API, and
# demonstrates how to apply it to observe the phases  of ``torch.compile``.
#
# .. note::
#
#   This tutorial requires PyTorch 2.2.0 or later.
#
#


######################################################################
# Setup
# ~~~~~~~~~~~~~~~~~~~~~
# In this example, we'll set up a simple Python function which performs an elementwise
# add and observe the compilation process with ``TORCH_LOGS`` Python API.
#
# .. note::
#
#   There is also an environment variable ``TORCH_LOGS``, which can be used to
#   change logging settings at the command line. The equivalent environment
#   variable setting is shown for each example.

import torch
import sys


def env_setup():
    """Set up environment for running the example. Exit cleanly if CUDA is not available."""
    if not torch.cuda.is_available():
        print("CUDA is not available. Exiting.")
        sys.exit(0)
        
    if torch.cuda.get_device_capability() < (7, 0):
        print("Skipping because torch.compile is not supported on this device.")
        sys.exit(0)


def separator(name):
    """Print separator and reset dynamo between each example"""
    print(f"\n{'='*20} {name} {'='*20}")
    torch._dynamo.reset()


def run_debugging_suite():
    """Run the complete debugging suite with all logging options"""
    env_setup()

    @torch.compile()
    def fn(x, y):
        z = x + y
        return z + 2

    inputs = (
        torch.ones(2, 2, device="cuda"),
        torch.zeros(2, 2, device="cuda")
    )
    
    logging_scenarios = [
        # View dynamo tracing; TORCH_LOGS="+dynamo"
        ("Dynamo Tracing", {"dynamo": logging.DEBUG}), 

        # View traced graph; TORCH_LOGS="graph"
        ("Traced Graph", {"graph": True}), 
        
        # View fusion decisions; TORCH_LOGS="fusion"
        ("Fusion Decisions", {"fusion": True}), 
        
        # View output code generated by inductor; TORCH_LOGS="output_code"
        ("Output Code", {"output_code": True}) 
    ]
    
    for name, log_config in logging_scenarios:
        separator(name)
        torch._logging.set_logs(**log_config)
        try:
            result = fn(*inputs)
            print(f"Function output shape: {result.shape}")
        except Exception as e:
            print(f"Error during {name}: {str(e)}")


run_debugging_suite()

######################################################################
# Use TORCH_TRACE/tlparse to produce produce compilation reports
# ~~~~~~~~~~
#
# In this section, we introduce the usage of TORCH_TRACE and tlparse to produce reports.
# First, we run `TORCH_TRACE="/tmp/tracedir" python script.py` to generate the raw trace logs.
# We have replace `/tmp/tracedir` with a path to a directory you want to store the trace logs
# and reaplce script with the name of your script.
# 
# Next, we are going to pass the trace log to `tlparse` to generate compilation reports. We run
# `pip install tlparse` and then `tlparse /tmp/tracedir`. This will open up your browser with 
# HTML like generated above.
#
# By default, reports generated by `tlparse` are 
# stored in the directory `tl_out`. You can change that by doing `tlparse /tmp/tracedir -o output_dir/`.

######################################################################
# Conclusion
# ~~~~~~~~~~
#
# In this tutorial we introduced the TORCH_LOGS environment variable and python API
# by experimenting with a small number of the available logging options.
# To view descriptions of all available options, run any python script
# which imports torch and set TORCH_LOGS to "help".
#
# Alternatively, you can view the `torch._logging documentation`_ to see
# descriptions of all available logging options.
#
# For more information on torch.compile, see the `torch.compile tutorial`_.
#
# .. _torch._logging documentation: https://pytorch.org/docs/main/logging.html
# .. _torch.compile tutorial: https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html