From f1f57e1e54593b279dcb856cc9fd8d3cfcd4e44b Mon Sep 17 00:00:00 2001
From: Richard Li <openrichard@meta.com>
Date: Mon, 5 Jun 2023 15:10:31 +0000
Subject: [PATCH] trigger tracing for MTIA events (#102288)

Summary: trigger tracing for MTIA events on python side when ProfilerActivity.MTIA is specified

Test Plan:
Test diff: D45437426

```
hg graft D45437426
```
- in one terminal

```
cd ~/fbsource/fbcode
buck2 run -j 8 \
    //infra_asic_fpga/firmware/tools/mad/service:mad_service
```
- in another terminal

Pytorch profiler
```
buck run mode/dev-nosan -j 8 //caffe2/torch/fb/acc_runtime/afg/tests:test_afg  -- -m kernel_add
```

Differential Revision: D46122853

Pull Request resolved: https://github.com/pytorch/pytorch/pull/102288
Approved by: https://github.com/aaronenyeshi
---
 torch/_C/_profiler.pyi     | 1 +
 torch/autograd/profiler.py | 4 ++++
 torch/profiler/profiler.py | 1 +
 3 files changed, 6 insertions(+)

diff --git a/torch/_C/_profiler.pyi b/torch/_C/_profiler.pyi
index b98db3d41d7..a1c0d720e7a 100644
--- a/torch/_C/_profiler.pyi
+++ b/torch/_C/_profiler.pyi
@@ -37,6 +37,7 @@ class ActiveProfilerType(Enum):
 class ProfilerActivity(Enum):
     CPU = ...
     CUDA = ...
+    MTIA = ...
 
 class _EventType(Enum):
     TorchOp = ...
diff --git a/torch/autograd/profiler.py b/torch/autograd/profiler.py
index 91ef2be92ed..7082b200eab 100644
--- a/torch/autograd/profiler.py
+++ b/torch/autograd/profiler.py
@@ -179,6 +179,7 @@ class profile:
             with_modules=False,
             use_kineto=False,
             use_cpu=True,
+            use_mtia=False,
             experimental_config=None):
         self.enabled: bool = enabled
         if not self.enabled:
@@ -194,6 +195,7 @@ class profile:
         self.with_stack = with_stack
         self.with_modules = with_modules
         self.use_cpu = use_cpu
+        self.use_mtia = use_mtia
         if experimental_config is None:
             experimental_config = _ExperimentalConfig()
         self.experimental_config = experimental_config
@@ -210,6 +212,8 @@ class profile:
         self.kineto_activities = set()
         if self.use_cpu:
             self.kineto_activities.add(ProfilerActivity.CPU)
+        if self.use_mtia:
+            self.kineto_activities.add(ProfilerActivity.MTIA)
 
         self.profiler_kind = ProfilerState.KINETO
         if self.use_cuda:
diff --git a/torch/profiler/profiler.py b/torch/profiler/profiler.py
index d0061b87644..70595cc0f7c 100644
--- a/torch/profiler/profiler.py
+++ b/torch/profiler/profiler.py
@@ -105,6 +105,7 @@ class _KinetoProfile:
         self.profiler = prof.profile(
             use_cuda=(ProfilerActivity.CUDA in self.activities),
             use_cpu=(ProfilerActivity.CPU in self.activities),
+            use_mtia=(ProfilerActivity.MTIA in self.activities),
             record_shapes=self.record_shapes,
             with_flops=self.with_flops,
             profile_memory=self.profile_memory,