Initial changes for updated triton-viz

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="triton-viz",
-    version="0.1",
+    version="0.2",
     packages=find_packages(),
     description="A visualization tool for Triton",
     author="Deep Learning Profiling Tools Team",
@@ -11,8 +11,7 @@
     install_requires=[
         "setuptools",
         "triton",
-        "gradio",
-        "chalk-diagrams @ git+https://github.com/chalk-diagrams/chalk.git",
+        "flask",
         "pyarrow",
         "pre-commit",
         "pytest",

triton_viz/__init__.py

@@ -1,5 +1,5 @@
 from .trace import trace, dump, sample
-from .draw import collect_grid, draw_record
+from .draw import collect_grid
 from .interface import launch
 
-__all__ = ["trace", "launch", "dump", "sample", "collect_grid", "draw_record"]
+__all__ = ["trace", "launch", "dump", "sample", "collect_grid"]

triton_viz/data.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass, field
-from typing import List, Tuple, Any
+from typing import List, Tuple, Any, Dict
 import traceback
 import numpy.typing as npt
 import numpy as np
 
+import torch
+
 
 @dataclass
 class Op:
@@ -73,6 +75,15 @@ class Dot(Op):
     input_shape: Tuple
     other_shape: Tuple
     output_shape: Tuple
+    input_data: List[List[float]]
+    other_data: List[List[float]]
+    intermediate_results: Dict[Tuple[int, int], float] = field(
+        default_factory=dict
+    )  # Only storing the result now
+
+    def update_intermediate(self, row: int, col: int, result: float):
+        # Store only the result as a float
+        self.intermediate_results[(row, col)] = result
 
 
 @dataclass
@@ -91,6 +102,7 @@ class Tensor:
     stride: Tuple
     shape: Tuple
     element_size: int
+    data: torch.Tensor
 
 
 @dataclass