[DEV][PROFILER] Add dot overriders to skip tl.dot (#221)

Author: mark14wu

triton_viz/clients/profiler/profiler.py

@@ -1,6 +1,6 @@
 from ...core.client import Client
 from ...core.callbacks import OpCallbacks, ForLoopCallbacks
-from ...core.data import Op, Load, Store, AddPtr
+from ...core.data import Op, Load, Store, AddPtr, Dot
 from .data import LoadStoreBytes
 from triton.runtime.interpreter import _get_np_dtype, TensorHandle
 import numpy as np
@@ -203,6 +203,11 @@ def store_overrider(ptr, value, mask, cache_modifier, eviction_policy):
             # Skip actual store
             pass
 
+        def dot_overrider(a, b, d, input_precision, max_num_imprecise_acc):
+            # Skip actual dot operation, return zeros with same shape as d
+            # This replaces np.matmul(a_data, b_data, dtype=d.data.dtype) + d.data
+            return TensorHandle(np.zeros_like(d.data), d.dtype.scalar)
+
         def pre_addptr_callback(ptr, offset):
             dtype_tt = ptr.get_element_ty()
             element_bitwidth = dtype_tt.primitive_bitwidth
@@ -232,6 +237,11 @@ def pre_addptr_callback(ptr, offset):
                 return OpCallbacks(
                     before_callback=pre_store_callback, op_overrider=store_overrider
                 )
+        elif op_type is Dot:
+            if self.disable_load_store_skipping:
+                return OpCallbacks()
+            else:
+                return OpCallbacks(op_overrider=dot_overrider)
         elif op_type is AddPtr:
             return OpCallbacks(before_callback=pre_addptr_callback)