From 5a880cf7a4c1c5015d9b4c61a8ead800f3d0210d Mon Sep 17 00:00:00 2001
From: ClairePankk <1136126752@qq.com>
Date: Tue, 12 Sep 2023 14:34:47 +0800
Subject: [PATCH] extend usage of torch_npu.profler

---
 torch_npu/utils/profiler.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/torch_npu/utils/profiler.py b/torch_npu/utils/profiler.py
index 9569c65c1b..f82331166b 100644
--- a/torch_npu/utils/profiler.py
+++ b/torch_npu/utils/profiler.py
@@ -1,8 +1,10 @@
 import sys
 import os
 import time
+import json
 import torch
 import torch_npu
+from torch_npu.profiler.analysis.prof_common_func.constant import Constant
 
 def Singleton(cls):
     _instances = {}
@@ -52,6 +54,8 @@ class Profile(object):
             torch_args_set = set(["enabled", "use_cuda", "record_shapes", "with_flops", "profile_memory",
                                   "with_stack", "use_kineto", "use_cpu", "use_npu_simple"])
             cann_ge_args_set = set(["config"])
+            profiler_args_set = set(["avtivities", "record_shapes", "with_flops", "profile_memory", "with_stack" 
+                                     "with_flops", "with_modules", "experimental_config", "use_cuda"])
             if self.profile_type == "TORCH":
                 if not set(kwargs.keys()).issubset(torch_args_set):
                     raise ValueError("Args '%s' invaild, expect args '%s' ." % (kwargs.keys(), torch_args_set))
@@ -64,7 +68,26 @@ class Profile(object):
                 if not set(kwargs.keys()).issubset(cann_ge_args_set):
                     raise ValueError("Args '%s' invaild, expect args '%s' ." % (kwargs.keys(), cann_ge_args_set))
                 self.prof = torch.npu.profile(self.save_path, self.use_e2e_profiler, **kwargs)
-
+            elif self.profile_type == 'ASCEND_PROFILER' :
+                profile_option = eval(str(os.getnev('PROFILE_OPTION','{"":""}')))
+                if not set(kwargs.keys()).issubset(profiler_args_set):
+                    raise  ValueError("Args '%s' invaild, expect args '%s' ." % (kwargs.keys(), profiler_args_set))
+                experimental_config = torch_npu.profiler._ExperimentalConfig(
+                    aic_metrics = profile_option.get('aic_metrics') or Constant.AicMetricsNone,
+                    profiler_level = profile_option.get('trace_level') or Constant.LEVEL0,
+                    l2_cache = profile_option.get('l2_cache') or False,
+                    record_op_args = profile_option.get('record_op_args') or False
+                )
+                self.prof = torch_npu.profiler.profile(
+                    activities = profile_option.get('activities'),
+                    record_shapes = profile_option.get('record_shapes') or False,
+                    profile_memory = profile_option.get('profile_memory') or False,
+                    with_stack = profile_option.get('with_stack') or False,
+                    with_flops = profile_option.get('with_flops') or False,
+                    with_modules = profile_option.get('with_modules') or False,
+                    experimental_config= experimental_config,
+                    on_trace_ready = torch_npu.profiler.tensorboard_trace_handler("./npu_profiling/result_dir",**kwargs)
+                )
             try:
                 os.makedirs(self.save_path, exist_ok=True)
             except Exception:
@@ -90,6 +113,8 @@ class Profile(object):
         self.count = 0
 
         if self.enable and self.step_count == self.start_step:
+            if self.profile_type == 'ASCEND_PROFILER':
+                self.prof.stop()
             self.prof.__exit__(None, None, None)
             if self.profile_type == "TORCH":
                 filename = "torch_prof_" + time.strftime("%Y%m%d%H%M%S") + ".json"
-- 
Gitee