Generate schema metrics and input variation metrics

tests/models/bert/test_bert.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering

tests/models/bloom/test_bloom.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForCausalLM

tests/models/falcon/test_falcon.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForCausalLM

tests/models/gpt2/test_gpt2.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForSequenceClassification

tests/models/llama/test_llama.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 
 # Load model directly
 from transformers import AutoTokenizer, AutoModelForCausalLM

tests/models/mnist/test_mnist.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 from tests.utils import check_with_pcc
 from torchvision import transforms, datasets
 from torch.utils.data import DataLoader

tests/models/resnet/test_resnet.py

@@ -4,7 +4,7 @@
 import unittest
 import ttnn
 import collections
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 from tests.utils import check_with_pcc
 
 

tests/models/yolos/test_yolos.py

@@ -2,7 +2,7 @@
 import torch_ttnn
 import unittest
 import ttnn
-from torch_ttnn.utils import RunTimeMetrics
+from torch_ttnn.metrics import RunTimeMetrics
 from PIL import Image
 import requests
 

tools/collect_metrics.py

@@ -4,6 +4,7 @@
 import csv
 from pathlib import Path
 import pandas as pd
+import numpy as np
 
 from tests.utils import comp_pcc
 
@@ -62,6 +63,9 @@ def load_pickle(path: str):
     # Holds the concatenation of all the metrics for each model
     all_metrics = []
 
+    # Holds the concatenation of input variation metrics for all models
+    all_input_var_metrics = {}
+
     # Assumed directory structure example. Some files will not exist if test failed.
     """
     pytorch2.0_ttnn
@@ -156,9 +160,9 @@ def load_pickle(path: str):
                 f"{compiled_outputs}\n"
             )
         accuracy_metric = {
-            "accuracy": round(accuracy, 2)
-            if not isinstance(accuracy, str)
-            else accuracy
+            "accuracy": (
+                round(accuracy, 2) if not isinstance(accuracy, str) else accuracy
+            )
         }
 
         # Add links that point to the directory of the model in the model name
@@ -203,6 +207,63 @@ def load_pickle(path: str):
 
         all_metrics.append(cat_metrics_remapped)
 
+        # Process input variation metrics. Currently, this is not per model, but per op.
+        input_var_metrics_path = model_path / "aten_ops_input_variations.pickle"
+        input_var_metrics = (
+            load_pickle(input_var_metrics_path)
+            if os.path.isfile(input_var_metrics_path)
+            else {}
+        )
+
+        for key, val in input_var_metrics.items():
+            if key not in all_input_var_metrics:
+                all_input_var_metrics[key] = val
+            else:
+                # Only append if shape and value combination have not been collected
+                for shape, value in zip(val["input_shapes"], val["input_values"]):
+                    if (
+                        shape not in all_input_var_metrics[key]["input_shapes"]
+                        and value not in all_input_var_metrics[key]["input_values"]
+                    ):
+                        all_input_var_metrics[key]["input_shapes"].append(shape)
+                        all_input_var_metrics[key]["input_values"].append(value)
+
+    # Write input variation metrics to csv
+    if all_input_var_metrics:
+        # Holds the rows to generate csv
+        input_var_list_for_csv = {}
+        # turn input_shapes and input_values into individual columns
+        for val in list(all_input_var_metrics.values()):
+            # holds the variations of input string
+            input_var_list = []
+            for shapes, values in zip(val["input_shapes"], val["input_values"]):
+                # holds each individual input to be joined to a string
+                input_string_list = []
+                for i, (shape, value) in enumerate(zip(shapes, values)):
+                    # This instance is a kwarg
+                    if isinstance(value, tuple):
+                        arg_name = value[0]
+                        arg_type = val["schema"]["kwargs"][arg_name]
+                        arg_val = f" = {value[1]}"
+                    else:
+                        arg_type = val["schema"]["args"][i][0]
+                        arg_name = val["schema"]["args"][i][1]
+                        arg_val = f" = {value}" if value else ""
+
+                    arg_shape = f"<{shape}>" if shape else ""
+
+                    input_string_list.append(
+                        f"{arg_type}{arg_shape} {arg_name}{arg_val}"
+                    )
+                input_var_list.append(", ".join(input_string_list))
+            input_var_list_for_csv[val["opname"]] = input_var_list
+
+        df = pd.DataFrame(
+            {key: pd.Series(value) for key, value in input_var_list_for_csv.items()}
+        )
+        df.to_csv("input_variations.csv", encoding="utf-8", index=False)
+        print(f"Data written to input_variations.csv")
+
     # Write metrics to csv
     if all_metrics:
         with open(f"metrics.csv", "w", newline="") as f:

torch_ttnn/backend.py

@@ -7,6 +7,7 @@
 import pickle
 from pathlib import Path
 import os
+import torch_ttnn.metrics as metrics
 
 torch._dynamo.config.suppress_errors = False
 torch._dynamo.config.verbose = True
@@ -74,24 +75,19 @@ def aten_backend(
 
     option: TorchTtnnOption = options["torch_ttnn_option"]
 
-    # Helper function to count the number of aten ops in the graph currently
-    # Returns a tuple of (total ops, total unique ops)
-    def count_aten_ops():
-        aten_ops = [
-            str(node.target)
-            for node in list(gm.graph.nodes)
-            if node.op in ["call_function", "call_method"]
-            and isinstance(node.target, torch._ops.OpOverload)
-            and "aten" in str(node.target)
-        ]
-        return (len(aten_ops), len(set(aten_ops)))
-
     # Save the number of aten ops before compilation
     if option.metrics_path:
         (
             option._metrics["torch_ops_before"],
             option._metrics["torch_ops_unique_before"],
-        ) = count_aten_ops()
+        ) = metrics.count_aten_ops(gm.graph.nodes)
+
+        input_variations = metrics.collect_input_variations_from_nodes(gm.graph.nodes)
+        # Save the input variation data
+        p = Path(f"metrics/{option.metrics_path}")
+        pickle_out_path = p / "aten_ops_input_variations.pickle"
+        with open(pickle_out_path, "wb") as f:
+            pickle.dump(input_variations, f)
 
     # Register ttnn objects as graph globals
     register_ttnn_objects(option)
@@ -140,17 +136,11 @@ def count_aten_ops():
         (
             option._metrics["torch_ops_remain"],
             option._metrics["torch_ops_unique_remain"],
-        ) = count_aten_ops()
-        # Save the number of to/from_device ops in current graph
-        to_from_device_ops = [
-            node.target.__name__
-            for node in list(gm.graph.nodes)
-            if node.op in ["call_function", "call_method"]
-            and (
-                "ttnn.to" in node.target.__name__ or "ttnn.from" in node.target.__name__
-            )
-        ]
-        option._metrics["to_from_device_ops"] = len(to_from_device_ops)
+        ) = metrics.count_aten_ops(gm.graph.nodes)
+        option._metrics["to_from_device_ops"] = metrics.count_to_from_device_ops(
+            gm.graph.nodes
+        )
+
         # Save the data as pickle files
         p = Path(f"metrics/{option.metrics_path}")
         pickle_out_path = p / "compiled-op_metrics.pickle"