[DEV] First pass at Analyzer

triton_viz/analysis.py

@@ -5,7 +5,7 @@
 
 
 # Function to compute metrics in the analysis shown during visualization.
-def analyzer():
+def analyze_records():
     launch_data = record_builder.launches[0]
     op_type_counts = defaultdict(int)
     grid_size = launch_data.grid
@@ -17,32 +17,30 @@ def analyzer():
         tensor.ptr: tensor.element_size for tensor in launch_data.tensors
     }
     for record in launch_data.records:
-        if isinstance(record, Grid):
-            pass
-        else:
+        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
-    )
+            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])

triton_viz/interface.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import triton_viz
 import tempfile
-from .analysis import analyzer
+from .analysis import analyze_records
 import pandas as pd
 
 
 def launch(share=True):
     cache = {}
-    analysis_data = analyzer()
+    analysis_data = analyze_records()
     program_records, tt, failures = triton_viz.collect_grid()
     m = [0, 0, 0]
     size = [0, 0]