diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py index b5c3f0bd8a..948ffd5571 100644 --- a/recipes_source/torch_logs.py +++ b/recipes_source/torch_logs.py @@ -31,52 +31,102 @@ # variable setting is shown for each example. import torch - -# exit cleanly if we are on a device that doesn't support torch.compile -if torch.cuda.get_device_capability() < (7, 0): - print("Skipping because torch.compile is not supported on this device.") -else: - @torch.compile() - def fn(x, y): - z = x + y - return z + 2 +import sys - inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda")) +def env_setup(): + """Set up the environment to run 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) -# print separator and reset dynamo -# between each example - def separator(name): - print(f"==================={name}=========================") - torch._dynamo.reset() +def separator(name): + """Print a separator and reset dynamo between each example""" + print(f"\n{'='*20} {name} {'='*20}") + torch._dynamo.reset() - separator("Dynamo Tracing") -# View dynamo tracing -# TORCH_LOGS="+dynamo" - torch._logging.set_logs(dynamo=logging.DEBUG) - fn(*inputs) +def run_debugging_suite(): + """Run the complete debugging suite with all logging options""" + env_setup() - separator("Traced Graph") -# View traced graph -# TORCH_LOGS="graph" - torch._logging.set_logs(graph=True) - fn(*inputs) - - separator("Fusion Decisions") -# View fusion decisions -# TORCH_LOGS="fusion" - torch._logging.set_logs(fusion=True) - fn(*inputs) + @torch.compile() + def fn(x, y): + z = x + y + return z + 2 - separator("Output Code") -# View output code generated by inductor -# TORCH_LOGS="output_code" - torch._logging.set_logs(output_code=True) - fn(*inputs) + 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() - separator("") +###################################################################### +# Using ``TORCH_TRACE/tlparse`` to produce produce compilation reports (for PyTorch 2) +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# In this section, we introduce ``TORCH_TRACE`` and ``tlparse`` to produce reports. +# +# +# 1. Generate the raw trace logs by running the following command: +# +# .. code-block:: bash +# +# TORCH_TRACE="/tmp/tracedir" python script.py` +# +# Ensure you replace ``/tmp/tracedir`` with the path to the directory where you want +# to store the trace logs and replace the script with the name of your script. +# +# 2. Install ``tlparse`` by running: +# +# .. code-block:: bash +# +# pip install tlparse +# +# 3. Pass the trace log to ``tlparse`` to generate compilation reports: +# +# .. code-block: bash +# +# tlparse /tmp/tracedir +# +# This will open your browser with the HTML-like code generated above. +# +# By default, reports generated by ``tlparse`` are stored in the ``tl_out`` directory. +# You can change that by running: +# +# .. code-block:: bash +# +# tlparse /tmp/tracedir -o output_dir/ ###################################################################### # Conclusion