[DEV] First pass at Analyzer

triton_viz/analysis.py

@@ -0,0 +1,48 @@
+from collections import defaultdict
+import numpy as np
+from .interpreter import record_builder
+from .data import Load, Store, Grid
+
+
+# Function to compute metrics in the analysis shown during visualization.
+def analyze_records():
+    launch_data = record_builder.launches[0]
+    op_type_counts = defaultdict(int)
+    grid_size = launch_data.grid
+    total_load_bytes_true = 0
+    total_store_bytes_true = 0
+    total_load_bytes_attempted = 0
+    total_store_bytes_attempted = 0
+    tensor_ptr_to_element_size = {
+        tensor.ptr: tensor.element_size for tensor in launch_data.tensors
+    }
+    for record in launch_data.records:
+        if not isinstance(record, Grid):
+            op_type_counts[type(record).__name__] += 1
+            if isinstance(record, Load):
+                element_size = tensor_ptr_to_element_size[record.ptr]
+                mask_true = np.count_nonzero(record.access_masks)
+                mask_false = np.count_nonzero(np.logical_not(record.access_masks))
+                total_load_bytes_true += mask_true * element_size
+                total_load_bytes_attempted += (mask_true + mask_false) * element_size
+            elif isinstance(record, Store):
+                element_size = tensor_ptr_to_element_size[record.ptr]
+                mask_true = np.count_nonzero(record.access_masks)
+                mask_false = np.count_nonzero(np.logical_not(record.access_masks))
+                total_store_bytes_true += mask_true * element_size
+                total_store_bytes_attempted += (mask_true + mask_false) * element_size
+                overall_load_ratio = (
+                    total_load_bytes_true / total_load_bytes_attempted
+                    if total_load_bytes_attempted > 0
+                    else 0
+                )
+                overall_store_ratio = (
+                    total_store_bytes_true / total_store_bytes_attempted
+                    if total_store_bytes_attempted > 0
+                    else 0
+                )
+    data = [["Grid Size", tuple(grid_size)]]
+    data += [[op_type, count] for op_type, count in op_type_counts.items()]
+    data.append(["Masked Load Ratio", overall_load_ratio])
+    data.append(["Masked Store Ratio", overall_store_ratio])
+    return data

triton_viz/interface.py

@@ -1,10 +1,13 @@
 import gradio as gr
 import triton_viz
 import tempfile
+from .analysis import analyze_records
+import pandas as pd
 
 
 def launch(share=True):
     cache = {}
+    analysis_data = analyze_records()
     program_records, tt, failures = triton_viz.collect_grid()
     m = [0, 0, 0]
     size = [0, 0]
@@ -35,8 +38,10 @@ def launch(share=True):
                 s2 = gr.Slider(0, m[1] - 1, value=0, step=1, label="Program Id 1")
                 s3 = gr.Slider(0, m[2] - 1, value=0, step=1, label="Program Id 2")
                 b1 = gr.Button("Precompute")
-                gr.Markdown(f"## Program Ids: {tuple(m)}")
-
+                # gr.Markdown(f"## Program Ids: {tuple(m)}")
+                df = pd.DataFrame(analysis_data, columns=["Metric", "Value"])
+                gr.Markdown("## Analysis")
+                gr.Dataframe(df)
                 if failures:
                     gr.Markdown(
                         show_label=False,