New profiler API (#48280)

Summary:
Pull Request resolved: #48280

Adding new API for the kineto profiler that supports enable predicate
function

Test Plan: unit test

Reviewed By: ngimel

Differential Revision: D25142220

Pulled By: ilia-cher

fbshipit-source-id: c57fa42855895075328733d7379eaf3dc1743d14

docs/source/index.rst

@@ -66,6 +66,7 @@ Features described in this documentation are classified by release status:
    torch.jit <jit>
    torch.linalg <linalg>
    torch.overrides
+   profiler
    nn.init
    onnx
    optim

docs/source/profiler.rst

@@ -0,0 +1,17 @@
+.. currentmodule:: torch.profiler
+
+torch.profiler
+==============
+
+Overview
+--------
+.. automodule:: torch.profiler
+
+
+API Reference
+-------------
+
+.. autoclass:: torch.profiler.profile
+  :members:
+
+.. autofunction:: torch.profiler.schedule

test/test_profiler.py

@@ -240,5 +240,47 @@ def test_flops(self):
         print(profiler_output)
         self.assertIn("FLOPS", profiler_output)
 
+    @unittest.skipIf(not kineto_available(), "Kineto is required")
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
+    def test_kineto_profiler_api(self):
+        called_num = [0]
+
+        with profile(use_cuda=True, use_kineto=True):
+            self.payload()
+
+        def trace_handler(p):
+            print(p.key_averages().table(
+                sort_by="self_cuda_time_total", row_limit=-1))
+            # p.export_chrome_trace("/tmp/test_trace_" + str(called_num[0]) + ".json")
+            called_num[0] += 1
+
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA],
+            schedule=torch.profiler.schedule(
+                wait=1,
+                warmup=1,
+                active=2),
+            on_trace_ready=trace_handler
+        ) as p:
+            for idx in range(8):
+                self.payload()
+                p.next_step()
+
+        self.assertEqual(called_num[0], 2)
+
+        # case without enable_pred
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA]
+        ) as p:
+            self.payload()
+            self.payload()
+        print(p.key_averages().table(
+            sort_by="self_cuda_time_total", row_limit=-1))
+
+
 if __name__ == '__main__':
     run_tests()

torch/__init__.py

@@ -596,6 +596,7 @@ def _assert(condition, message):
 import torch.utils.data
 import torch.__config__
 import torch.__future__
+import torch.profiler
 
 _C._init_names(list(torch._storage_classes))
 

torch/autograd/profiler.py

@@ -464,6 +464,15 @@ def __enter__(self):
             torch.autograd._enable_profiler_legacy(self.config())
         return self
 
+    def _prepare_kineto_trace(self):
+        assert self.kineto_activities
+        self.entered = True
+        torch.autograd._prepare_profiler(self.config(), self.kineto_activities)
+
+    def _start_kineto_trace(self):
+        assert self.kineto_activities
+        torch.autograd._enable_profiler(self.config(), self.kineto_activities)
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         if not self.enabled:
             return

torch/profiler/__init__.py

@@ -0,0 +1,12 @@
+# type: ignore
+r'''
+PyTorch Profiler is a tool that allows the collecton of the performance metrics during the training and inference.
+Profiler's context manager API can be used to better understand what model operators are the most expensive,
+examine their input shapes and stack traces, study device kernel activity and visualize the execution trace.
+
+.. note::
+    An earlier version of the API in ``torch.autograd`` module is considered legacy and will be deprecated.
+
+'''
+
+from .profiler import profile, schedule, ProfilerAction, ProfilerActivity