[mpact][profiler] Add utils for profiling Python programs and torch ops (#59)
diff --git a/benchmark/README.md b/benchmark/README.md
index 4749547..e870ae9 100644
--- a/benchmark/README.md
+++ b/benchmark/README.md
@@ -24,3 +24,14 @@
```shell
python path/to/the/kernels_benchmark.py --benchmark-filter=add
```
+
+### Profiler
+
+Utils for profiling python scripts and pytorch models could be found in
+`benchmark/python/utils/profiler.py`.
+
+To run the profiling example, use the following command:
+
+```shell
+python benchmark/python/utils/profiler.py
+```
diff --git a/benchmark/python/utils/profiler.py b/benchmark/python/utils/profiler.py
new file mode 100644
index 0000000..8bca04b
--- /dev/null
+++ b/benchmark/python/utils/profiler.py
@@ -0,0 +1,66 @@
+import torch
+import cProfile
+from pstats import Stats
+
+
+def profile_torch(
+ func, args, row_limit=10, save_output=False, func_name=None, file_name="trace"
+):
+ """Use PyTorch's profiler to profile torch ops.
+
+ To see the graph: upload trace.json to chrome://tracing
+
+ More details about PyTorch profiler:
+ https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
+ """
+ func_name = func_name if func_name else func.__name__
+ with torch.profiler.profile() as prof:
+ with torch.profiler.record_function(func_name):
+ func(*args)
+ print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=row_limit))
+ if save_output:
+ prof.export_chrome_trace(f"{file_name}.json")
+
+
+def profile_python(func, args, row_limit=10, save_output=False, file_name="stats"):
+ """Use cProfile to profile python function calls.
+
+ To see the graph, run the following commands:
+ 1. python -m pip install snakeviz
+ 2. snakeviz stats.prof
+ """
+ pr = cProfile.Profile()
+ pr.enable()
+ func(*args)
+ pr.disable()
+ stats = Stats(pr)
+ stats.sort_stats("tottime").print_stats(row_limit)
+ if save_output:
+ pr.dump_stats(f"{file_name}.prof")
+
+
+if __name__ == "__main__":
+ # Example usage of the profiler.
+ from mpact.models.kernels import MMNet
+ from mpact_benchmark.utils.tensor_generator import generate_tensor
+ from mpact.mpactbackend import mpact_jit
+
+ # Generate input tensors.
+ dense_tensor1 = generate_tensor(seed=0, shape=(32, 32), sparsity=0.8)
+ dense_tensor2 = generate_tensor(seed=1, shape=(32, 32), sparsity=0.8)
+ sparse_tensor1 = dense_tensor1.to_sparse_csr()
+ sparse_tensor2 = dense_tensor2.to_sparse_csr()
+
+ # Profile with PyTorch profiler for torch operators.
+ # MPACT sparse.
+ profile_torch(mpact_jit, (MMNet(), sparse_tensor1, sparse_tensor2))
+ # Torch sparse.
+ profile_torch(
+ MMNet(), (sparse_tensor1, sparse_tensor2), func_name="sparsexsparse matmul"
+ )
+
+ # Profile with cProfile for Python function calls.
+ # MPACT sparse.
+ profile_python(mpact_jit, (MMNet(), sparse_tensor1, sparse_tensor2))
+ # Torch sparse.
+ profile_python(MMNet(), (sparse_tensor1, sparse_tensor2))
