Function main

advanced_source/cuda_graph_annotations_tutorial.py:541–595

End-to-end CUDA graph annotation and profiling demo.

()

# annotations, profile it, and post-process the trace.

def main():
    """End-to-end CUDA graph annotation and profiling demo."""
    if not torch.cuda.is_available():
        raise SystemExit("CUDA required for this tutorial")

    # Check if annotation support is available
    # PyTorch will log a warning if cuda-bindings version is too old
    supported = not _is_tools_id_unavailable()
    print(f"Annotation support available: {supported}")
    if not supported:
        print("NOTE: Annotation API not available.")
        print("This could be due to:")
        print("  - Driver/CUDA-compat < 13.1")
        print("  - Outdated cuda-bindings (check PyTorch warnings above)")
        print("Annotations will not be recorded, but the demo will still run.")
        print("Kernels will be reassigned to the default lane, not semantic lanes.\n")

    output_dir = Path("traces")

    # Build the model
    print("\n1. Building transformer block model...")
    model_fn = build_transformer_block()

    # Capture graph with annotations
    print("\n2. Capturing CUDA graph with annotations...")
    graph, output = capture_graph_with_annotations(model_fn)

    # Save annotations
    print("\n3. Saving annotation metadata...")
    annotations_path = save_annotations(output_dir)

    # Profile the graph
    print("\n4. Profiling graph replays...")
    raw_trace_path = profile_graph(graph, output_dir)

    # Post-process the trace
    print("\n5. Post-processing: merging annotations into trace...")
    annotated_path, raw_trace, annotated_trace = post_process_trace(
        raw_trace_path, annotations_path, output_dir
    )

    # Compare before and after
    print("\n6. Comparing traces...")
    compare_traces(raw_trace, annotated_trace)

    # Summary
    print("\n" + "="*60)
    print("SUMMARY")
    print("="*60)
    print(f"Raw trace:       {raw_trace_path}")
    print(f"Annotated trace: {annotated_path}")
    print(f"Annotations:     {annotations_path}")
    print("\nOpen the annotated trace in https://ui.perfetto.dev/ to visualize")
    print("the semantic kernel lanes.")
    print("="*60)

# Example output:
# if __name__ == "__main__":