Profiling-helper (#2321)

* change profiler

* remove unused imports

* remove unused imports

* change lazybuffer references

* remove unused line

* remove unused import

* remove unused stuff

* add types

* typing

* typing

* typing

* trigger actions

* -1 loc

* fixup

* trigger actions

* revert lazy typing changes

* WIP profiler helper

* replace old start & stop profiler

* fixup

* linting

* Update llama.py

---------

Co-authored-by: George Hotz <[email protected]>

examples/llama.py

@@ -9,7 +9,7 @@ import numpy as np
 np.set_printoptions(linewidth=200)
 from typing import Optional, Tuple, Union
 
-from tinygrad.helpers import Timing, getenv, DEBUG, dtypes, CI
+from tinygrad.helpers import Timing, Profiling, getenv, DEBUG, dtypes, CI
 from tinygrad.ops import Device
 from tinygrad.tensor import Tensor
 from tinygrad.nn import Embedding, Linear
@@ -514,10 +514,6 @@ After you are done speaking, output [EOS]. You are not Chad.
   sys.stdout.write(outputted)
   sys.stdout.flush()
 
-  if args.profile:
-    import cProfile, pstats
-    profiler = cProfile.Profile()
-
   # chatbot loop
   while 1:
     # add tokens from user in chatbot mode
@@ -533,17 +529,17 @@ After you are done speaking, output [EOS]. You are not Chad.
     last_break = len(outputted)
     for i in range(args.count):
       GlobalCounters.reset()
-      if args.profile and i == 2: profiler.enable()
 
-      if args.timing: print("")
+      if args.timing or args.profile: print("")
       st = GlobalCounters.time_sum_s
-      with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
-        with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
-                    f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
-                    (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_count*1e-9*2/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
-          probs = llama.model(Tensor([toks[start_pos:]]), start_pos, args.temperature).realize()
-        # TODO: fix JIT rand so we can put this in the JIT
-        tok = probs.multinomial().item()
+      with Profiling(enabled=args.profile):
+        with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
+          with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
+                      f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
+                      (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_count*1e-9*2/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
+            probs = llama.model(Tensor([toks[start_pos:]]), start_pos, args.temperature).realize()
+          # TODO: fix JIT rand so we can put this in the JIT
+          tok = probs.multinomial().item()
 
       # use the kv cache
       start_pos = len(toks)
@@ -560,8 +556,3 @@ After you are done speaking, output [EOS]. You are not Chad.
       # stop after you have your answer
       if chatbot and outputted.endswith(end_delim): break
     if not chatbot: break
-
-  if args.profile:
-    profiler.disable()
-    stats = pstats.Stats(profiler)
-    stats.dump_stats("out.prof")

test/external/external_test_speed_llama.py

@@ -2,12 +2,11 @@
 import unittest, time
 import numpy as np
 from examples.llama import Transformer, MODEL_PARAMS
-from test.test_net_speed import start_profile, stop_profile
 from tinygrad.tensor import Tensor
 from tinygrad.ops import Device
 from tinygrad.nn.state import get_state_dict
 from tinygrad.ops import Compiled
-from tinygrad.helpers import dtypes, prod
+from tinygrad.helpers import dtypes, prod, Profiling
 from tinygrad.runtime.lib import RawBuffer
 
 class FakeProgram:
@@ -46,9 +45,8 @@ class TestLLaMASpeed(unittest.TestCase):
     run_llama("codegen")
     run_llama("methodcache", False)
 
-    pr = start_profile()
-    run_llama("profile")
-    stop_profile(pr, sort='time', frac=0.1)
+    with Profiling(sort='time', frac=0.1):
+      run_llama("profile")
 
     Device[Device.DEFAULT].runtime = backup_program
     Device[Device.DEFAULT].buffer = backup_buffer

test/test_net_speed.py

@@ -1,27 +1,13 @@
 #!/usr/bin/env python
 import time
-import cProfile
-import pstats
 import unittest
 import torch
-from tinygrad.tensor import Tensor, Device
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import Profiling
 import pytest
 
 pytestmark = [pytest.mark.exclude_cuda, pytest.mark.exclude_gpu, pytest.mark.exclude_clang]
 
-def start_profile():
-  import time
-  pr = cProfile.Profile(timer=lambda: int(time.time()*1e9), timeunit=1e-6)
-  pr.enable()
-  return pr
-
-def stop_profile(pr, sort='cumtime', frac=0.2):
-  pr.disable()
-  ps = pstats.Stats(pr)
-  ps.strip_dirs()
-  ps.sort_stats(sort)
-  ps.print_stats(frac)
-
 class TestConvSpeed(unittest.TestCase):
 
   def test_mnist(self):
@@ -86,12 +72,13 @@ class TestConvSpeed(unittest.TestCase):
       [x.grad.realize() for x in [c1, c2, l1]]
       et2 = time.time()
       if i == 0:
-        pr = start_profile()
+        pr = Profiling(sort='time', frac=0.2)
+        pr.__enter__()
       else:
         fpt += (et1-et0)
         bpt += (et2-et1)
 
-    stop_profile(pr, sort='time')
+    pr.__exit__()
     fpt = (fpt*1000/cnt)
     bpt = (bpt*1000/cnt)
     print("forward pass:  %.3f ms, %.2fx off baseline %.3f ms" % (fpt, fpt/fpt_baseline, fpt_baseline))

test/test_winograd.py

@@ -1,9 +1,8 @@
 import unittest
-from tinygrad.helpers import Timing, CI
+from tinygrad.helpers import Timing, CI, Profiling
 from tinygrad.tensor import Tensor
 from tinygrad.ops import LoadOps
 from tinygrad.codegen.linearizer import Linearizer
-from test.test_net_speed import start_profile, stop_profile
 
 class TestWinograd(unittest.TestCase):
   def setUp(self):
@@ -31,9 +30,8 @@ class TestWinograd(unittest.TestCase):
 
   def test_profile(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
-    if not CI: pr = start_profile()
-    out = Tensor.conv2d(x,w).realize()
-    if not CI: stop_profile(pr, sort='time')
+    with Profiling(enabled=not CI, sort='time'):
+      out = Tensor.conv2d(x,w).realize()
     out.numpy()
 
 if __name__ == '__main__':

tinygrad/helpers.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3
+import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, cProfile, pstats
 import numpy as np
 from typing import Dict, Tuple, Union, List, NamedTuple, Final, Iterator, ClassVar, Optional, Iterable, Any, TypeVar, TYPE_CHECKING
 if TYPE_CHECKING:  # TODO: remove this and import TypeGuard from typing once minimum python supported version is 3.10
@@ -67,10 +67,20 @@ GRAPH, GRAPHPATH = getenv("GRAPH", 0), getenv("GRAPHPATH", "/tmp/net")
 class Timing(contextlib.ContextDecorator):
   def __init__(self, prefix="", on_exit=None, enabled=True): self.prefix, self.on_exit, self.enabled = prefix, on_exit, enabled
   def __enter__(self): self.st = time.perf_counter_ns()
-  def __exit__(self, exc_type, exc_val, exc_tb):
+  def __exit__(self, *exc):
     self.et = time.perf_counter_ns() - self.st
     if self.enabled: print(f"{self.prefix}{self.et*1e-6:.2f} ms"+(self.on_exit(self.et) if self.on_exit else ""))
 
+class Profiling(contextlib.ContextDecorator):
+  def __init__(self, enabled=True, sort='cumtime', frac=0.2): self.enabled, self.sort, self.frac = enabled, sort, frac
+  def __enter__(self):
+    self.pr = cProfile.Profile(timer=lambda: int(time.time()*1e9), timeunit=1e-6)
+    if self.enabled: self.pr.enable()
+  def __exit__(self, *exc):
+    if self.enabled:
+      self.pr.disable()
+      pstats.Stats(self.pr).strip_dirs().sort_stats(self.sort).print_stats(self.frac)
+
 # **** tinygrad now supports dtypes! *****
 
 class DType(NamedTuple):