[sync] PR-37136: [feat]profiler add ProfilerStep range-pytorch-AtomGit

Mtest/profiler/test_npu_profiler.py+12-2

	@@ -221,10 +221,15 @@ class TestNpuProfiler(TestCase):
221	worker_name = self.worker_name	221	worker_name = self.worker_name
222	with torch_npu.profiler.profile(	222	with torch_npu.profiler.profile(
223	activities=[torch_npu.profiler.ProfilerActivity.NPU],	223	activities=[torch_npu.profiler.ProfilerActivity.NPU],
		224	+ schedule=torch_npu.profiler.schedule(
		225	+ wait=0, warmup=0, active=1, repeat=1, skip_first=0
		226	+ ),
224	on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(	227	on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
225	self.results_path, worker_name=worker_name	228	self.results_path, worker_name=worker_name
226	),	229	),
227	- experimental_config=torch_npu.profiler._ExperimentalConfig(l2_cache=True),	230	+ experimental_config=torch_npu.profiler._ExperimentalConfig(
		231	+ mstx=True, l2_cache=True
		232	+ ),
228	) as prof:	233	) as prof:
229	for step in range(self.small_steps):	234	for step in range(self.small_steps):
230	self.model_train.train_one_step()	235	self.model_train.train_one_step()
	@@ -246,7 +251,12 @@ class TestNpuProfiler(TestCase):
246	True,	251	True,
247	self._has_view_result(self.results_path, worker_name, self.SOC_PMU),	252	self._has_view_result(self.results_path, worker_name, self.SOC_PMU),
248	)	253	)
249	- # self.assertEqual(False, self._check_trace_view_keywords(worker_name, ["async_npu"]))	254	+ self.assertEqual(
		255	+ True,
		256	+ self._check_trace_view_keywords(
		257	+ self.results_path, worker_name, ["ProfilerStep#0"]
		258	+ ),
		259	+ )
250		260
251	def test_record_shapes(self):	261	def test_record_shapes(self):
252	worker_name = self.worker_name	262	worker_name = self.worker_name

Athird_party/mspti/mspti_activity.h+34-0

@@ -0,0 +1,34 @@
	1	+#ifndef INC_EXTERNAL_MSPTI_ACTIVITY_H_
	2	+#define INC_EXTERNAL_MSPTI_ACTIVITY_H_
	3	+
	4	+#include "acl/acl_base.h"
	5	+
	6	+#ifdef __cplusplus
	7	+extern "C" {
	8	+#endif
	9	+
	10	+typedef enum msptiActivityKind {
	11	+ MSPTI_ACTIVITY_KIND_INVALID = 0,
	12	+ MSPTI_ACTIVITY_KIND_MARKER = 1,
	13	+ MSPTI_ACTIVITY_KIND_KERNEL = 2,
	14	+ MSPTI_ACTIVITY_KIND_API = 3,
	15	+ MSPTI_ACTIVITY_KIND_HCCL = 4,
	16	+ MSPTI_ACTIVITY_KIND_MEMORY = 5,
	17	+ MSPTI_ACTIVITY_KIND_MEMSET = 6,
	18	+ MSPTI_ACTIVITY_KIND_MEMCPY = 7,
	19	+ MSPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION = 8,
	20	+ MSPTI_ACTIVITY_KIND_COMMUNICATION = 9,
	21	+ MSPTI_ACTIVITY_KIND_ACL_API = 10,
	22	+ MSPTI_ACTIVITY_KIND_NODE_API = 11,
	23	+ MSPTI_ACTIVITY_KIND_RUNTIME_API = 12,
	24	+ MSPTI_ACTIVITY_KIND_COUNT,
	25	+ MSPTI_ACTIVITY_KIND_FORCE_INT = 0x7fffffff
	26	+} msptiActivityKind;
	27	+
	28	+ACL_FUNC_VISIBILITY bool msptiActivityIsEnabled(msptiActivityKind kind);
	29	+
	30	+#ifdef __cplusplus
	31	+}
	32	+#endif
	33	+
	34	+#endif

Atorch_npu/csrc/framework/interface/MsptiInterface.cpp+62-0

@@ -0,0 +1,62 @@
	1	+#include "torch_npu/csrc/framework/interface/MsptiInterface.h"
	2	+
	3	+#include "torch_npu/csrc/core/npu/NPUException.h"
	4	+#include "torch_npu/csrc/core/npu/register/FunctionLoader.h"
	5	+#include "torch_npu/csrc/core/npu/npu_log.h"
	6	+#include "torch_npu/csrc/toolkit/profiler/common/utils.h"
	7	+
	8	+namespace at_npu {
	9	+namespace native {
	10	+
	11	+#undef TORCH_NPU_LOAD_FUNC
	12	+#define TORCH_NPU_LOAD_FUNC(funcName) \
	13	+ TORCH_NPU_REGISTER_FUNCTION(libmspti, funcName)
	14	+
	15	+#undef TORCH_NPU_GET_FUNC
	16	+#define TORCH_NPU_GET_FUNC(funcName) \
	17	+ TORCH_NPU_GET_FUNCTION(libmspti, funcName)
	18	+
	19	+TORCH_NPU_REGISTER_LIBRARY(libmspti)
	20	+TORCH_NPU_LOAD_FUNC(msptiActivityIsEnabled)
	21	+
	22	+static bool IsSupportMsptiFuncImpl()
	23	+{
	24	+ static auto checkSupport = []() -> bool {
	25	+ char* path = std::getenv("ASCEND_HOME_PATH");
	26	+ if (path != nullptr) {
	27	+ std::string soPath = std::string(path) + "/lib64/libmspti.so";
	28	+ soPath = torch_npu::toolkit::profiler::Utils::RealPath(soPath);
	29	+ return !soPath.empty();
	30	+ }
	31	+ return false;
	32	+ };
	33	+ return checkSupport();
	34	+}
	35	+
	36	+bool IsSupportMsptiFunc()
	37	+{
	38	+ static bool isSupport = IsSupportMsptiFuncImpl();
	39	+ return isSupport;
	40	+}
	41	+
	42	+bool MsptiActivityIsEnabled(msptiActivityKind kind)
	43	+{
	44	+ using MsptiActivityIsEnabledFunc = bool (*)(msptiActivityKind);
	45	+ static MsptiActivityIsEnabledFunc func = nullptr;
	46	+ static bool noFuncFlag = false;
	47	+ if (noFuncFlag) {
	48	+ return false;
	49	+ }
	50	+ if (func == nullptr) {
	51	+ func = (MsptiActivityIsEnabledFunc)TORCH_NPU_GET_FUNC(msptiActivityIsEnabled);
	52	+ if (func == nullptr) {
	53	+ ASCEND_LOGW("Failed to get func msptiActivityIsEnabled");
	54	+ noFuncFlag = true;
	55	+ return false;
	56	+ }
	57	+ }
	58	+ return func(kind);
	59	+}
	60	+
	61	+} // namespace native
	62	+} // namespace at_npu

Atorch_npu/csrc/framework/interface/MsptiInterface.h+16-0

@@ -0,0 +1,16 @@
	1	+#ifndef __TORCH_NPU_MSPTIINTERFACE__
	2	+#define __TORCH_NPU_MSPTIINTERFACE__
	3	+
	4	+#include <third_party/mspti/mspti_activity.h>
	5	+
	6	+namespace at_npu {
	7	+namespace native {
	8	+
	9	+bool IsSupportMsptiFunc();
	10	+
	11	+bool MsptiActivityIsEnabled(msptiActivityKind kind);
	12	+
	13	+} // namespace native
	14	+} // namespace at_npu
	15	+
	16	+#endif

Mtorch_npu/csrc/profiler/mstx_mgr.cpp+21-17

	@@ -2,6 +2,7 @@
2	#include "torch_npu/csrc/core/npu/NPUStream.h"	2	#include "torch_npu/csrc/core/npu/NPUStream.h"
3	#include "torch_npu/csrc/core/npu/NPUFunctions.h"	3	#include "torch_npu/csrc/core/npu/NPUFunctions.h"
4	#include "torch_npu/csrc/framework/interface/MstxInterface.h"	4	#include "torch_npu/csrc/framework/interface/MstxInterface.h"
		5	+#include "torch_npu/csrc/framework/interface/MsptiInterface.h"
5	#include "torch_npu/csrc/core/npu/npu_log.h"	6	#include "torch_npu/csrc/core/npu/npu_log.h"
6	#include "torch_npu/csrc/framework/OpCommand.h"	7	#include "torch_npu/csrc/framework/OpCommand.h"
7	#include "torch_npu/csrc/profiler/profiler_mgr.h"	8	#include "torch_npu/csrc/profiler/profiler_mgr.h"
	@@ -297,28 +298,31 @@ bool MstxMgr::isProfTxEnable()
297		298
298	bool MstxMgr::isMsptiTxEnableImpl()	299	bool MstxMgr::isMsptiTxEnableImpl()
299	{	300	{
300	- bool ret = false;	301	+ static bool isMsptiSoInLdPreload = []() -> bool {
301	- const char* envVal = std::getenv("LD_PRELOAD");	302	+ const char* envVal = std::getenv("LD_PRELOAD");
302	- if (envVal == nullptr) {	303	+ if (envVal == nullptr) {
303	- return ret;	304	+ return false;
304	- }
305	- static const std::string soName = "libmspti.so";
306	- std::stringstream ss(envVal);
307	- std::string path;
308	- while (std::getline(ss, path, ':')) {
309	- path = torch_npu::toolkit::profiler::Utils::RealPath(path);
310	- if ((path.size() > soName.size()) && (path.substr(path.size() - soName.size()) == soName)) {
311	- ret = true;
312	- break;
313	}	305	}
		306	+ static const std::string soName = "libmspti.so";
		307	+ std::stringstream ss(envVal);
		308	+ std::string path;
		309	+ while (std::getline(ss, path, ':')) {
		310	+ path = torch_npu::toolkit::profiler::Utils::RealPath(path);
		311	+ if ((path.size() > soName.size()) && (path.substr(path.size() - soName.size()) == soName)) {
		312	+ return true;
		313	+ }
		314	+ }
		315	+ return false;
		316	+ }();
		317	+ if (isMsptiSoInLdPreload) {
		318	+ return true;
314	}	319	}
315	- return ret;	320	+ return at_npu::native::IsSupportMsptiFunc() && at_npu::native::MsptiActivityIsEnabled(MSPTI_ACTIVITY_KIND_MARKER);
316	}	321	}
317		322
318	bool MstxMgr::isMsptiTxEnable()	323	bool MstxMgr::isMsptiTxEnable()
319	{	324	{
320	- static bool isEnable = isMsptiTxEnableImpl();	325	+ return isMsptiTxEnableImpl();
321	- return isEnable;
322	}	326	}
323		327
324	bool MstxMgr::isMstxEnable()	328	bool MstxMgr::isMstxEnable()
	@@ -334,4 +338,4 @@ bool MstxMgr::isMstxTxDomainEnable(const std::string &domainName)
334	return true;	338	return true;
335	}	339	}
336	}	340	}
337	-}	341	+}

Mtorch_npu/profiler/analysis/prof_view/prof_db_parse/_fwk_api_db_parser.py+6-5

	@@ -78,12 +78,13 @@ class FwkApiDbParser(BaseParser):
78		78
79	# update connection id for mstx mark op	79	# update connection id for mstx mark op
80	for mstx_mark_api in mstx_mark_apis:	80	for mstx_mark_api in mstx_mark_apis:
81	- if mstx_mark_api[TorchOpDataOri.CONNECTION_ID]:	81	+ mstx_mark_api[TorchOpDataOri.CONNECTION_ID] = (
82	- mstx_mark_api[TorchOpDataOri.CONNECTION_ID] = (	82	+ connectionId_manager.get_id_from_connection_ids(
83	- connectionId_manager.get_id_from_connection_ids(	83	+ mstx_mark_api[TorchOpDataOri.CONNECTION_ID]
84	- mstx_mark_api[TorchOpDataOri.CONNECTION_ID]
85	- )
86	)	84	)
		85	+ if mstx_mark_api[TorchOpDataOri.CONNECTION_ID]
		86	+ else None
		87	+ )
87	self._fwk_apis.extend(mstx_mark_apis)	88	self._fwk_apis.extend(mstx_mark_apis)
88		89
89	def get_mstx_mark_op_connection_ids_with_cann_api(	90	def get_mstx_mark_op_connection_ids_with_cann_api(

Mtorch_npu/profiler/profiler.py+22-6

	@@ -7,6 +7,7 @@ from typing import Any, Optional, Union
7		7
8	import torch.autograd.profiler as prof	8	import torch.autograd.profiler as prof
9	import torch_npu.npu	9	import torch_npu.npu
		10	+from torch_npu.npu import current_stream, mstx
10	from torch_npu._C._profiler import (	11	from torch_npu._C._profiler import (
11	_disable_profiler_in_child_thread,	12	_disable_profiler_in_child_thread,
12	_enable_profiler_in_child_thread,	13	_enable_profiler_in_child_thread,
	@@ -259,6 +260,8 @@ class profile(_KinetoProfile):
259	self.current_action = self.schedule(self.step_num)	260	self.current_action = self.schedule(self.step_num)
260	self._step_num_offset = 0	261	self._step_num_offset = 0
261	self.step_rec_fn: Optional[prof.record_function] = None	262	self.step_rec_fn: Optional[prof.record_function] = None
		263	+ self._step_mstx_range_id = 0
		264	+ self._is_dynamic_prof = False
262	if use_cuda is not None:	265	if use_cuda is not None:
263	print_warn_msg("This is npu environment, use_cuda is invalid")	266	print_warn_msg("This is npu environment, use_cuda is invalid")
264	self.stopped = False	267	self.stopped = False
	@@ -283,6 +286,17 @@ class profile(_KinetoProfile):
283	@no_exception_func()	286	@no_exception_func()
284	def _set_step_num_offset_for_dynamic_prof(self, step: int):	287	def _set_step_num_offset_for_dynamic_prof(self, step: int):
285	self._step_num_offset = step	288	self._step_num_offset = step
		289	+ self._is_dynamic_prof = True
		290	+
		291	+ def _start_step_mstx_range(self, message: str):
		292	+ if not self._is_dynamic_prof:
		293	+ self._step_mstx_range_id = mstx.range_start(message, current_stream())
		294	+
		295	+ def _end_step_mstx_range(self):
		296	+ if self._is_dynamic_prof or not self._step_mstx_range_id:
		297	+ return
		298	+ mstx.range_end(self._step_mstx_range_id)
		299	+ self._step_mstx_range_id = 0
286		300
287	@no_exception_func()	301	@no_exception_func()
288	def start(self):	302	def start(self):
	@@ -291,15 +305,16 @@ class profile(_KinetoProfile):
291	ProfPathCreator().init(export_only_mode=True)	305	ProfPathCreator().init(export_only_mode=True)
292	self.action_controller.transit_action(ProfilerAction.NONE, self.current_action)	306	self.action_controller.transit_action(ProfilerAction.NONE, self.current_action)
293	if self.record_steps:	307	if self.record_steps:
294	- self.step_rec_fn = prof.record_function(	308	+ step_name = "ProfilerStep#" + str(self.step_num + self._step_num_offset)
295	- "ProfilerStep#" + str(self.step_num + self._step_num_offset)	309	+ self.step_rec_fn = prof.record_function(step_name)
296	- )
297	self.step_rec_fn.__enter__()	310	self.step_rec_fn.__enter__()
		311	+ self._start_step_mstx_range(step_name)
298		312
299	@no_exception_func()	313	@no_exception_func()
300	def stop(self):	314	def stop(self):
301	if self.record_steps and self.step_rec_fn:	315	if self.record_steps and self.step_rec_fn:
302	self.step_rec_fn.__exit__(None, None, None)	316	self.step_rec_fn.__exit__(None, None, None)
		317	+ self._end_step_mstx_range()
303	self.action_controller.transit_action(self.current_action, None)	318	self.action_controller.transit_action(self.current_action, None)
304	self.stopped = True	319	self.stopped = True
305		320
	@@ -310,15 +325,16 @@ class profile(_KinetoProfile):
310	return	325	return
311	if self.record_steps and self.step_rec_fn:	326	if self.record_steps and self.step_rec_fn:
312	self.step_rec_fn.__exit__(None, None, None)	327	self.step_rec_fn.__exit__(None, None, None)
		328	+ self._end_step_mstx_range()
313	prev_action = self.current_action	329	prev_action = self.current_action
314	self.step_num += 1	330	self.step_num += 1
315	self.current_action = self.schedule(self.step_num)	331	self.current_action = self.schedule(self.step_num)
316	self.action_controller.transit_action(prev_action, self.current_action)	332	self.action_controller.transit_action(prev_action, self.current_action)
317	if self.record_steps:	333	if self.record_steps:
318	- self.step_rec_fn = prof.record_function(	334	+ step_name = "ProfilerStep#" + str(self.step_num + self._step_num_offset)
319	- "ProfilerStep#" + str(self.step_num + self._step_num_offset)	335	+ self.step_rec_fn = prof.record_function(step_name)
320	- )
321	self.step_rec_fn.__enter__()	336	self.step_rec_fn.__enter__()
		337	+ self._start_step_mstx_range(step_name)
322		338
323	@no_exception_func()	339	@no_exception_func()
324	def set_custom_trace_id_callback(self, callback: Callable[[], str]) -> None:	340	def set_custom_trace_id_callback(self, callback: Callable[[], str]) -> None: