Support allocator trace tracker in NPU caching allocator-pytorch-AtomGit

已合并

Support allocator trace tracker in NPU caching allocator #38685

zzhongmin创建于 6月16日从zzhongmin/pytorch： v2.12.0_tratra 合入到 Ascend/pytorch： v2.12.0

Support allocator trace tracker in NPU caching allocator #38685

已合并

zzhongmin创建于 6月16日

从zzhongmin/pytorch v2.12.0_tratra合入到Ascend/pytorch v2.12.0

共 7 个文件变更+301-29

Atest/cpp_extensions/allocator_trace_tracker_extension.cpp+62-0

@@ -0,0 +1,62 @@
	1	+#include <atomic>
	2	+#include <mutex>
	3	+
	4	+#include <torch/extension.h>
	5	+
	6	+#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
	7	+
	8	+namespace py = pybind11;
	9	+
	10	+namespace {
	11	+
	12	+using TraceEntry = c10_npu::NPUCachingAllocator::TraceEntry;
	13	+
	14	+std::once_flag tracker_registration_once;
	15	+std::atomic<int64_t> segment_alloc_count{0};
	16	+std::atomic<int64_t> segment_free_count{0};
	17	+
	18	+void update_trace_tracker_state(const TraceEntry& te)
	19	+{
	20	+ if (te.action_ == TraceEntry::SNAPSHOT) {
	21	+ return;
	22	+ }
	23	+
	24	+ if (te.action_ == TraceEntry::SEGMENT_ALLOC) {
	25	+ segment_alloc_count.fetch_add(1, std::memory_order_relaxed);
	26	+ } else if (te.action_ == TraceEntry::SEGMENT_FREE) {
	27	+ segment_free_count.fetch_add(1, std::memory_order_relaxed);
	28	+ }
	29	+}
	30	+
	31	+void attach_trace_tracker()
	32	+{
	33	+ std::call_once(tracker_registration_once, []() {
	34	+ c10_npu::NPUCachingAllocator::attachAllocatorTraceTracker(
	35	+ &update_trace_tracker_state);
	36	+ });
	37	+}
	38	+
	39	+void reset_trace_tracker_state()
	40	+{
	41	+ segment_alloc_count.store(0, std::memory_order_relaxed);
	42	+ segment_free_count.store(0, std::memory_order_relaxed);
	43	+}
	44	+
	45	+py::dict get_trace_tracker_state()
	46	+{
	47	+ py::dict result;
	48	+ result["segment_alloc_count"] =
	49	+ segment_alloc_count.load(std::memory_order_relaxed);
	50	+ result["segment_free_count"] =
	51	+ segment_free_count.load(std::memory_order_relaxed);
	52	+ return result;
	53	+}
	54	+
	55	+} // namespace
	56	+
	57	+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
	58	+{
	59	+ m.def("attach_trace_tracker", &attach_trace_tracker);
	60	+ m.def("reset_trace_tracker_state", &reset_trace_tracker_state);
	61	+ m.def("get_trace_tracker_state", &get_trace_tracker_state);
	62	+}

Atest/npu/test_allocator_trace_tracker.py+127-0

@@ -0,0 +1,127 @@
	1	+import gc
	2	+import os
	3	+import shutil
	4	+import subprocess
	5	+import unittest
	6	+
	7	+import torch
	8	+import torch.utils.cpp_extension
	9	+
	10	+import torch_npu
	11	+from torch_npu.testing.testcase import TestCase, run_tests
	12	+
	13	+
	14	+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
	15	+PYTORCH_INSTALL_PATH = os.path.dirname(os.path.realpath(torch.__file__))
	16	+PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))
	17	+
	18	+
	19	+def create_build_path(build_directory):
	20	+ if os.path.exists(build_directory):
	21	+ shutil.rmtree(build_directory, ignore_errors=True)
	22	+ os.makedirs(build_directory, exist_ok=True)
	23	+
	24	+
	25	+def build_stub(base_dir):
	26	+ build_stub_cmd = ["sh", os.path.join(base_dir, "third_party/acl/libs/build_stub.sh")]
	27	+ if subprocess.call(build_stub_cmd) != 0:
	28	+ raise RuntimeError(f"Failed to build stub: {build_stub_cmd}")
	29	+
	30	+
	31	+@unittest.skipIf(not torch_npu.npu.is_available(), "npu not available, skipping tests")
	32	+class TestAllocatorTraceTracker(TestCase):
	33	+ module = None
	34	+ build_directory = os.path.join(REPO_ROOT, "test", "build", "allocator_trace_tracker")
	35	+
	36	+ @classmethod
	37	+ def setUpClass(cls):
	38	+ super().setUpClass()
	39	+ build_stub(REPO_ROOT)
	40	+ create_build_path(cls.build_directory)
	41	+
	42	+ cann_lib_path = os.path.join(REPO_ROOT, "third_party", "acl", "libs")
	43	+ torch_npu_lib_path = os.path.join(PYTORCH_NPU_INSTALL_PATH, "lib")
	44	+ extra_include_paths = [
	45	+ os.path.join(PYTORCH_NPU_INSTALL_PATH, "include"),
	46	+ os.path.join(PYTORCH_NPU_INSTALL_PATH, "include", "third_party", "acl", "inc"),
	47	+ ]
	48	+ extra_ldflags = [
	49	+ f"-L{cann_lib_path}",
	50	+ "-lascendcl",
	51	+ f"-L{torch_npu_lib_path}",
	52	+ "-ltorch_npu",
	53	+ f"-Wl,-rpath,{torch_npu_lib_path}",
	54	+ "-lc10",
	55	+ f"-L{PYTORCH_INSTALL_PATH}",
	56	+ ]
	57	+
	58	+ cls.module = torch.utils.cpp_extension.load(
	59	+ name="allocator_trace_tracker_extension",
	60	+ sources=[
	61	+ os.path.join(REPO_ROOT, "test", "cpp_extensions", "allocator_trace_tracker_extension.cpp"),
	62	+ ],
	63	+ extra_include_paths=extra_include_paths,
	64	+ extra_cflags=["-g"],
	65	+ extra_ldflags=extra_ldflags,
	66	+ build_directory=cls.build_directory,
	67	+ verbose=False,
	68	+ )
	69	+
	70	+ torch.empty(1, device="npu")
	71	+ cls.module.attach_trace_tracker()
	72	+ gc.collect()
	73	+ torch_npu.npu.empty_cache()
	74	+
	75	+ def tearDown(self):
	76	+ self.module.reset_trace_tracker_state()
	77	+ torch_npu.npu.memory._record_memory_history(None)
	78	+ gc.collect()
	79	+ torch_npu.npu.empty_cache()
	80	+ super().tearDown()
	81	+
	82	+ @staticmethod
	83	+ def _allocate_large_buffer():
	84	+ return torch.empty(64 * 1024 * 1024, dtype=torch.uint8, device="npu")
	85	+
	86	+ def test_trace_tracker_callbacks_without_history(self):
	87	+ torch_npu.npu.memory._record_memory_history(None)
	88	+ self.assertFalse(torch_npu._C._npu_isHistoryEnabled())
	89	+
	90	+ torch_npu.npu.empty_cache()
	91	+ self.module.reset_trace_tracker_state()
	92	+
	93	+ buffer = self._allocate_large_buffer()
	94	+ state_after_alloc = self.module.get_trace_tracker_state()
	95	+
	96	+ self.assertGreaterEqual(state_after_alloc["segment_alloc_count"], 1)
	97	+
	98	+ del buffer
	99	+ gc.collect()
	100	+ torch_npu.npu.empty_cache()
	101	+
	102	+ state_after_free = self.module.get_trace_tracker_state()
	103	+ self.assertGreaterEqual(state_after_free["segment_free_count"], 1)
	104	+
	105	+ def test_trace_tracker_callbacks_with_history_enabled(self):
	106	+ torch_npu.npu.memory._record_memory_history(
	107	+ "all",
	108	+ context="alloc",
	109	+ stacks="python",
	110	+ max_entries=128,
	111	+ )
	112	+ self.assertTrue(torch_npu._C._npu_isHistoryEnabled())
	113	+
	114	+ self.module.reset_trace_tracker_state()
	115	+
	116	+ buffer = self._allocate_large_buffer()
	117	+ del buffer
	118	+ gc.collect()
	119	+ torch_npu.npu.empty_cache()
	120	+
	121	+ state = self.module.get_trace_tracker_state()
	122	+ self.assertGreaterEqual(state["segment_alloc_count"], 1)
	123	+ self.assertGreaterEqual(state["segment_free_count"], 1)
	124	+
	125	+
	126	+if __name__ == "__main__":
	127	+ run_tests()

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+79-23

	@@ -6,6 +6,7 @@
6	#include <mutex>	6	#include <mutex>
7	#include <regex>	7	#include <regex>
8	#include <set>	8	#include <set>
		9	+#include <unordered_set>
9	#include <vector>	10	#include <vector>
10	#include <fstream>	11	#include <fstream>
11		12
	@@ -226,6 +227,8 @@ struct BlockPool {
226	is_small(small),	227	is_small(small),
227	owner_PrivatePool(private_pool)	228	owner_PrivatePool(private_pool)
228	{}	229	{}
		230	+
		231	+ MempoolId_t owner_MempoolId() const;
229	};	232	};
230		233
231	struct ExpandableSegment;	234	struct ExpandableSegment;
	@@ -823,10 +826,12 @@ private:
823		826
824	// NPU graphs helper	827	// NPU graphs helper
825	struct PrivatePool {	828	struct PrivatePool {
826	- PrivatePool() : large_blocks(false, this), small_blocks(true, this) {}	829	+ explicit PrivatePool(MempoolId_t id)
		830	+ : id(std::move(id)), large_blocks(false, this), small_blocks(true, this) {}
827	PrivatePool(const PrivatePool &) = delete;	831	PrivatePool(const PrivatePool &) = delete;
828	PrivatePool(PrivatePool &&) = delete;	832	PrivatePool(PrivatePool &&) = delete;
829	PrivatePool &operator = (const PrivatePool &) = delete;	833	PrivatePool &operator = (const PrivatePool &) = delete;
		834	+ MempoolId_t id{ 0, 0 };
830	// Number of live graphs using this pool	835	// Number of live graphs using this pool
831	int use_count{ 1 };	836	int use_count{ 1 };
832	// Number of unfreed npuMallocs made for this pool. When use_count and	837	// Number of unfreed npuMallocs made for this pool. When use_count and
	@@ -842,6 +847,14 @@ struct PrivatePool {
842	BlockPool small_blocks;	847	BlockPool small_blocks;
843	};	848	};
844		849
		850	+MempoolId_t BlockPool::owner_MempoolId() const
		851	+{
		852	+ if (owner_PrivatePool) {
		853	+ return owner_PrivatePool->id;
		854	+ }
		855	+ return {0, 0};
		856	+}
		857	+
845	BlockState::BlockState(Block* block)	858	BlockState::BlockState(Block* block)
846	: device(block->device),	859	: device(block->device),
847	stream(block->stream),	860	stream(block->stream),
	@@ -1026,6 +1039,7 @@ private:
1026		1039
1027	// XXX - maybe we should generalize and have multiple events	1040	// XXX - maybe we should generalize and have multiple events
1028	std::vector<OutOfMemoryObserver> oom_observers_;	1041	std::vector<OutOfMemoryObserver> oom_observers_;
		1042	+ std::vector<AllocatorTraceTracker> trace_trackers_;
1029	std::shared_ptr<c10d_npu::HCCLComm> hcclComm_;	1043	std::shared_ptr<c10d_npu::HCCLComm> hcclComm_;
1030		1044
1031	// Private pools for NPU graphs	1045	// Private pools for NPU graphs
	@@ -1100,6 +1114,12 @@ public:
1100	oom_observers_.emplace_back(observer);	1114	oom_observers_.emplace_back(observer);
1101	}	1115	}
1102		1116
		1117	+ void attachAllocatorTraceTracker(AllocatorTraceTracker tracker)
		1118	+ {
		1119	+ std::unique_lock<std::recursive_mutex> lock(mutex);
		1120	+ trace_trackers_.emplace_back(std::move(tracker));
		1121	+ }
		1122	+
1103	bool checkUceInMemPool()	1123	bool checkUceInMemPool()
1104	{	1124	{
1105	auto memUceInfo_ = c10_npu::get_mem_uce_info();	1125	auto memUceInfo_ = c10_npu::get_mem_uce_info();
	@@ -1278,7 +1298,7 @@ public:
1278	stats.num_ooms += 1;	1298	stats.num_ooms += 1;
1279		1299
1280	record_trace(TraceEntry::OOM, device_free, params.size(), params.stream(), params.device(),	1300	record_trace(TraceEntry::OOM, device_free, params.size(), params.stream(), params.device(),
1281	- std::move(context));	1301	+ params.pool->owner_MempoolId(), std::move(context));
1282	auto observers_local = oom_observers_;	1302	auto observers_local = oom_observers_;
1283		1303
1284	// Make sure we do not have the device lock before calling our	1304	// Make sure we do not have the device lock before calling our
	@@ -1420,7 +1440,7 @@ public:
1420		1440
1421	block->context_when_allocated = std::move(context);	1441	block->context_when_allocated = std::move(context);
1422	record_trace(TraceEntry::ALLOC, int64_t(block->ptr), orig_size, block->stream, block->device,	1442	record_trace(TraceEntry::ALLOC, int64_t(block->ptr), orig_size, block->stream, block->device,
1423	- block->context_when_allocated);	1443	+ block->pool->owner_MempoolId(), block->context_when_allocated);
1424		1444
1425	active_blocks.insert(block);	1445	active_blocks.insert(block);
1426		1446
	@@ -1483,7 +1503,8 @@ public:
1483	});	1503	});
1484		1504
1485	record_trace(TraceEntry::FREE_REQUESTED, int64_t(block->ptr), block->requested_size, block->stream,	1505	record_trace(TraceEntry::FREE_REQUESTED, int64_t(block->ptr), block->requested_size, block->stream,
1486	- block->device, context ? context : block->context_when_allocated);	1506	+ block->device, block->pool->owner_MempoolId(),
		1507	+ context ? context : block->context_when_allocated);
1487		1508
1488	if (block->size >= NPUAllocatorConfig::max_split_size()) {	1509	if (block->size >= NPUAllocatorConfig::max_split_size()) {
1489	update_stat(stats.oversize_allocations, -1);	1510	update_stat(stats.oversize_allocations, -1);
	@@ -2065,7 +2086,7 @@ public:
2065	std::sort(result.begin(), result.end(),	2086	std::sort(result.begin(), result.end(),
2066	[](const SegmentInfo &a, const SegmentInfo &b) { return a.address < b.address; });	2087	[](const SegmentInfo &a, const SegmentInfo &b) { return a.address < b.address; });
2067		2088
2068	- record_trace(TraceEntry::SNAPSHOT, 0, total_active, nullptr, 0, nullptr);	2089	+ record_trace(TraceEntry::SNAPSHOT, 0, total_active, nullptr, 0, {0, 0}, nullptr);
2069	return result;	2090	return result;
2070	}	2091	}
2071		2092
	@@ -2125,26 +2146,41 @@ public:
2125	// See Note [Interaction with NPU graph capture]	2146	// See Note [Interaction with NPU graph capture]
2126		2147
2127	// Called by NPUGraph::capture_begin	2148	// Called by NPUGraph::capture_begin
2128	- void beginAllocateToPool(MempoolId_t mempool_id, std::function<bool(aclrtStream)> filter)	2149	+ void create_or_incref_pool(MempoolId_t mempool_id)
2129	{	2150	{
2130	- std::lock_guard<std::recursive_mutex> lock(mutex);
2131	auto it = graph_pools.find(mempool_id);	2151	auto it = graph_pools.find(mempool_id);
2132	if (it == graph_pools.end()) {	2152	if (it == graph_pools.end()) {
2133	- // mempool_id does not reference an existing pool. Make a new pool for	2153	+ // mempool_id does not reference an existing pool.
2134	- // this capture.	2154	+ // Make a new pool for NPUGraph capture or torch.npu.use_mem_pool
2135	- graph_pools.emplace(mempool_id, std::make_unique<PrivatePool>());	2155	+ // usage. use_count is initially 1, which means the pool is
		2156	+ // being used since somebody called createOrIncrefPool.
		2157	+ graph_pools.emplace(mempool_id, std::make_unique<PrivatePool>(mempool_id));
2136	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: new pool, "	2158	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: new pool, "
2137	"mempool_id=(%lu,%lu)", mempool_id.first, mempool_id.second);	2159	"mempool_id=(%lu,%lu)", mempool_id.first, mempool_id.second);
2138	} else {	2160	} else {
2139	- // mempool_id references an existing pool, which the current capture will	2161	+ // mempool_id references an existing pool, which the current NPUGraph
		2162	+ // capture or torch.npu.use_mem_pool will
2140	// share. Check this pool is live (at least one other capture already	2163	// share. Check this pool is live (at least one other capture already
2141	- // references it).	2164	+ // references it). Increment it to establish the usage.
2142	TORCH_INTERNAL_ASSERT(it->second->use_count > 0);	2165	TORCH_INTERNAL_ASSERT(it->second->use_count > 0);
2143	it->second->use_count++;	2166	it->second->use_count++;
2144	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: reuse pool, "	2167	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: reuse pool, "
2145	"mempool_id=(%lu,%lu), use_count=%d",	2168	"mempool_id=(%lu,%lu), use_count=%d",
2146	mempool_id.first, mempool_id.second, it->second->use_count);	2169	mempool_id.first, mempool_id.second, it->second->use_count);
2147	}	2170	}
		2171	+ }
		2172	+
		2173	+ PrivatePool* get_private_pool(MempoolId_t mempool_id) const
		2174	+ {
		2175	+ auto it = graph_pools.find(mempool_id);
		2176	+ TORCH_INTERNAL_ASSERT(it != graph_pools.end());
		2177	+ return it->second.get();
		2178	+ }
		2179	+
		2180	+ void beginAllocateToPool(MempoolId_t mempool_id, std::function<bool(aclrtStream)> filter)
		2181	+ {
		2182	+ std::lock_guard<std::recursive_mutex> lock(mutex);
		2183	+ create_or_incref_pool(mempool_id);
2148	for (auto it2 = captures_underway.begin(); it2 != captures_underway.end(); ++it2) {	2184	for (auto it2 = captures_underway.begin(); it2 != captures_underway.end(); ++it2) {
2149	TORCH_CHECK(it2->first != mempool_id, "beginAllocateToPool: already recording to mempool_id");	2185	TORCH_CHECK(it2->first != mempool_id, "beginAllocateToPool: already recording to mempool_id");
2150	}	2186	}
	@@ -2185,15 +2221,14 @@ public:
2185	// mempool. When the count reaches 0, we tell free_cached_blocks it may now	2221	// mempool. When the count reaches 0, we tell free_cached_blocks it may now
2186	// npuFree blocks from this graph's pool when it discovers they're unused	2222	// npuFree blocks from this graph's pool when it discovers they're unused
2187	// (unsplit).	2223	// (unsplit).
2188	- auto it = graph_pools.find(mempool_id);	2224	+ auto pp = get_private_pool(mempool_id);
2189	- TORCH_INTERNAL_ASSERT(it != graph_pools.end());	2225	+ auto uc = --(pp->use_count);
2190	- auto uc = --(it->second->use_count);
2191	TORCH_INTERNAL_ASSERT(uc >= 0);	2226	TORCH_INTERNAL_ASSERT(uc >= 0);
2192	if (uc == 0) {	2227	if (uc == 0) {
2193	// Allows free_cached_blocks to begin npuFreeing this pool's memory,	2228	// Allows free_cached_blocks to begin npuFreeing this pool's memory,
2194	// and makes sure this pool wasn't somehow made freeable already.	2229	// and makes sure this pool wasn't somehow made freeable already.
2195	// NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)	2230	// NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
2196	- bool inserted = graph_pools_freeable.insert({ mempool_id, it->second.get() }).second;	2231	+ bool inserted = graph_pools_freeable.insert({ mempool_id, pp }).second;
2197	TORCH_INTERNAL_ASSERT(inserted);	2232	TORCH_INTERNAL_ASSERT(inserted);
2198	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator releasePool: mempool_id=(%lu,%lu), "	2233	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator releasePool: mempool_id=(%lu,%lu), "
2199	"use_count reached 0, pool marked freeable",	2234	"use_count reached 0, pool marked freeable",
	@@ -2341,7 +2376,7 @@ private:
2341	for_each_selected_stat_type(stat_types,	2376	for_each_selected_stat_type(stat_types,
2342	[&](size_t stat_type) { update_stat(stats.reserved_bytes[stat_type], mapped_range.size); });	2377	[&](size_t stat_type) { update_stat(stats.reserved_bytes[stat_type], mapped_range.size); });
2343	record_trace(TraceEntry::SEGMENT_MAP, int64_t(mapped_range.ptr), mapped_range.size, to_map->stream,	2378	record_trace(TraceEntry::SEGMENT_MAP, int64_t(mapped_range.ptr), mapped_range.size, to_map->stream,
2344	- to_map->device, ctx);	2379	+ to_map->device, to_map->pool->owner_MempoolId(), ctx);
2345	if (!to_map->prev && !to_map->context_when_segment_allocated) {	2380	if (!to_map->prev && !to_map->context_when_segment_allocated) {
2346	to_map->context_when_segment_allocated = ctx;	2381	to_map->context_when_segment_allocated = ctx;
2347	}	2382	}
	@@ -2388,7 +2423,8 @@ private:
2388	AT_ASSERT(!block->allocated && block->event_count == 0, PTA_ERROR(ErrCode::VALUE));	2423	AT_ASSERT(!block->allocated && block->event_count == 0, PTA_ERROR(ErrCode::VALUE));
2389		2424
2390	record_trace(TraceEntry::FREE_COMPLETED, int64_t(block->ptr), block->requested_size, block->stream,	2425	record_trace(TraceEntry::FREE_COMPLETED, int64_t(block->ptr), block->requested_size, block->stream,
2391	- block->device, context ? context : block->context_when_allocated);	2426	+ block->device, block->pool->owner_MempoolId(),
		2427	+ context ? context : block->context_when_allocated);
2392		2428
2393	block->context_when_allocated = nullptr;	2429	block->context_when_allocated = nullptr;
2394	block->hccl_work_ptr = nullptr;	2430	block->hccl_work_ptr = nullptr;
	@@ -2754,7 +2790,8 @@ private:
2754	// p.block came from new, not npuMalloc. It should not be nullptr here.	2790	// p.block came from new, not npuMalloc. It should not be nullptr here.
2755	TORCH_INTERNAL_ASSERT(p.block != nullptr && p.block->ptr != nullptr);	2791	TORCH_INTERNAL_ASSERT(p.block != nullptr && p.block->ptr != nullptr);
2756		2792
2757	- record_trace(TraceEntry::SEGMENT_ALLOC, int64_t(p.block->ptr), p.block->size, p.stream(), p.device(), ctx);	2793	+ record_trace(TraceEntry::SEGMENT_ALLOC, int64_t(p.block->ptr), p.block->size, p.stream(), p.device(),
		2794	+ p.block->pool->owner_MempoolId(), ctx);
2758	p.block->context_when_segment_allocated = ctx;	2795	p.block->context_when_segment_allocated = ctx;
2759	return true;	2796	return true;
2760	}	2797	}
	@@ -2861,6 +2898,7 @@ private:
2861	block->size, block->ptr, block->device);	2898	block->size, block->ptr, block->device);
2862		2899
2863	record_trace(TraceEntry::SEGMENT_FREE, int64_t(block->ptr), block->size, block->stream, block->device,	2900	record_trace(TraceEntry::SEGMENT_FREE, int64_t(block->ptr), block->size, block->stream, block->device,
		2901	+ block->pool->owner_MempoolId(),
2864	context ? context : block->context_when_segment_allocated);	2902	context ? context : block->context_when_segment_allocated);
2865		2903
2866	auto it = ipc_handle_map.find(block->ptr);	2904	auto it = ipc_handle_map.find(block->ptr);
	@@ -2945,6 +2983,7 @@ private:
2945	}	2983	}
2946		2984
2947	record_trace(TraceEntry::SEGMENT_UNMAP, int64_t(unmapped.ptr), unmapped.size, block->stream, block->device,	2985	record_trace(TraceEntry::SEGMENT_UNMAP, int64_t(unmapped.ptr), unmapped.size, block->stream, block->device,
		2986	+ block->pool->owner_MempoolId(),
2948	context ? context : block->context_when_segment_allocated);	2987	context ? context : block->context_when_segment_allocated);
2949	}	2988	}
2950		2989
	@@ -3137,15 +3176,25 @@ private:
3137	}	3176	}
3138		3177
3139	void record_trace(TraceEntry::Action action, int64_t addr, size_t size, aclrtStream stream, int device,	3178	void record_trace(TraceEntry::Action action, int64_t addr, size_t size, aclrtStream stream, int device,
3140	- std::shared_ptr<c10::GatheredContext> context)	3179	+ MempoolId_t mempool_id = {0, 0}, std::shared_ptr<c10::GatheredContext> context = nullptr)
3141	{	3180	{
3142	- if (!record_history) {	3181	+ if (!record_history && trace_trackers_.empty()) {
3143	return;	3182	return;
3144	}	3183	}
3145	-	3184	+ auto te = TraceEntry(
3146	- auto te = TraceEntry(action, device, addr, size, stream,	3185	+ action,
		3186	+ device,
		3187	+ addr,
		3188	+ size,
		3189	+ stream,
		3190	+ mempool_id,
3147	record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);	3191	record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);
3148		3192
		3193	+ // Callbacks should not include any Pytorch call
		3194	+ for (const auto& cb : trace_trackers_) {
		3195	+ cb(te);
		3196	+ }
		3197	+
3149	if (record_history) {	3198	if (record_history) {
3150	if (alloc_trace->size() < alloc_trace_max_entries_) {	3199	if (alloc_trace->size() < alloc_trace_max_entries_) {
3151	alloc_trace->emplace_back(te);	3200	alloc_trace->emplace_back(te);
	@@ -3303,6 +3352,13 @@ public:
3303	}	3352	}
3304	}	3353	}
3305		3354
		3355	+ void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override
		3356	+ {
		3357	+ for (auto& allocator : device_allocator) {
		3358	+ allocator->attachAllocatorTraceTracker(tracker);
		3359	+ }
		3360	+ }
		3361	+
3306	bool checkUceInMemPool(int device) override	3362	bool checkUceInMemPool(int device) override
3307	{	3363	{
3308	return device_allocator[device]->checkUceInMemPool();	3364	return device_allocator[device]->checkUceInMemPool();

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.h+18-3

	@@ -183,17 +183,19 @@ struct TraceEntry {
183	};	183	};
184	TraceEntry(Action action, int device, int64_t addr, size_t size,	184	TraceEntry(Action action, int device, int64_t addr, size_t size,
185	aclrtStream stream,	185	aclrtStream stream,
		186	+ MempoolId_t mempool = {0, 0},
186	std::shared_ptr<c10::GatheredContext> context = nullptr)	187	std::shared_ptr<c10::GatheredContext> context = nullptr)
187	: action_(action), device_(device), addr_(addr),	188	: action_(action), device_(device), addr_(addr),
188	- context_(std::move(context)), stream_(stream), size_(size)	189	+ context_(std::move(context)), stream_(stream), size_(size),
189	- {	190	+ mempool_(std::move(mempool))
190	- }	191	+ {}
191	Action action_;	192	Action action_;
192	int device_;	193	int device_;
193	int64_t addr_; // for OOM, this is the amount of free bytes reported by cuda	194	int64_t addr_; // for OOM, this is the amount of free bytes reported by cuda
194	std::shared_ptr<c10::GatheredContext> context_;	195	std::shared_ptr<c10::GatheredContext> context_;
195	aclrtStream stream_;	196	aclrtStream stream_;
196	int64_t size_;	197	int64_t size_;
		198	+ MempoolId_t mempool_;
197	};	199	};
198		200
199	struct SnapshotInfo {	201	struct SnapshotInfo {
	@@ -219,6 +221,7 @@ enum struct RecordContext {
219	using OutOfMemoryObserver =	221	using OutOfMemoryObserver =
220	std::function<void(int64_t device, int64_t allocated, int64_t device_total,	222	std::function<void(int64_t device, int64_t allocated, int64_t device_total,
221	int64_t device_free)>;	223	int64_t device_free)>;
		224	+using AllocatorTraceTracker = std::function<void(const TraceEntry&)>;
222		225
223	struct ShareableHandle {	226	struct ShareableHandle {
224	ptrdiff_t offset;	227	ptrdiff_t offset;
	@@ -296,6 +299,13 @@ public:
296	virtual CheckpointDelta setCheckpointPoolState(	299	virtual CheckpointDelta setCheckpointPoolState(
297	c10::DeviceIndex device,	300	c10::DeviceIndex device,
298	std::shared_ptr<AllocatorState> pps) = 0;	301	std::shared_ptr<AllocatorState> pps) = 0;
		302	+ // Attached AllocatorTraceTracker callbacks will be called while the
		303	+ // per-device allocator lock is held. Any additional locks taken from within
		304	+ // the callback must be proven to always have the lock order that never
		305	+ // triggers a deadlock. In particular, Python's GIL may be held when
		306	+ // calling the allocator so it is unsafe to try to acquire the GIL in this
		307	+ // callback.
		308	+ virtual void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) = 0;
299	};	309	};
300		310
301	// Allocator object, statically initialized	311	// Allocator object, statically initialized
	@@ -502,6 +512,11 @@ inline void attachOutOfMemoryObserver(OutOfMemoryObserver observer)
502	return get()->attachOutOfMemoryObserver(observer);	512	return get()->attachOutOfMemoryObserver(observer);
503	}	513	}
504		514
		515	+inline void attachAllocatorTraceTracker(AllocatorTraceTracker tracker)
		516	+{
		517	+ return get()->attachAllocatorTraceTracker(std::move(tracker));
		518	+}
		519	+
505	inline bool checkUceInMemPool(int device)	520	inline bool checkUceInMemPool(int device)
506	{	521	{
507	return get()->checkUceInMemPool(device);	522	return get()->checkUceInMemPool(device);

Mtorch_npu/csrc/core/npu/NPUWorkspaceAllocator.cpp+3-3

	@@ -287,19 +287,19 @@ public:
287	}	287	}
288	for (const auto& block_pair : blocks) {	288	for (const auto& block_pair : blocks) {
289	auto te = TraceEntry(TraceEntry::WORKSPACE_SNAPSHOT, device, int64_t(block_pair.second->data_ptr),	289	auto te = TraceEntry(TraceEntry::WORKSPACE_SNAPSHOT, device, int64_t(block_pair.second->data_ptr),
290	- block_pair.second->size, block_pair.first,	290	+ block_pair.second->size, block_pair.first, MempoolId_t{0, 0},
291	record_context_ >= RecordContext::ALLOC ? block_pair.second->context_when_allocated	291	record_context_ >= RecordContext::ALLOC ? block_pair.second->context_when_allocated
292	: nullptr);	292	: nullptr);
293	alloc_trace.emplace_back(te);	293	alloc_trace.emplace_back(te);
294		294
295	te = TraceEntry(TraceEntry::SEGMENT_ALLOC, device, int64_t(block_pair.second->data_ptr),	295	te = TraceEntry(TraceEntry::SEGMENT_ALLOC, device, int64_t(block_pair.second->data_ptr),
296	- block_pair.second->size, block_pair.first,	296	+ block_pair.second->size, block_pair.first, MempoolId_t{0, 0},
297	record_context_ >= RecordContext::ALLOC ? block_pair.second->context_when_allocated	297	record_context_ >= RecordContext::ALLOC ? block_pair.second->context_when_allocated
298	: nullptr);	298	: nullptr);
299	alloc_trace.emplace_back(te);	299	alloc_trace.emplace_back(te);
300		300
301	te = TraceEntry(TraceEntry::ALLOC, device, int64_t(block_pair.second->data_ptr), block_pair.second->size,	301	te = TraceEntry(TraceEntry::ALLOC, device, int64_t(block_pair.second->data_ptr), block_pair.second->size,
302	- block_pair.first,	302	+ block_pair.first, MempoolId_t{0, 0},
303	record_context_ >= RecordContext::ALLOC ? block_pair.second->context_when_allocated	303	record_context_ >= RecordContext::ALLOC ? block_pair.second->context_when_allocated
304	: nullptr);	304	: nullptr);
305	alloc_trace.emplace_back(te);	305	alloc_trace.emplace_back(te);

Mtorch_npu/csrc/npu/NPUPluggableAllocator.cpp+11-0

	@@ -433,6 +433,17 @@ void NPUPluggableAllocator::attachOutOfMemoryObserver(
433	"If you need it, please file an issue describing your use case.");	433	"If you need it, please file an issue describing your use case.");
434	}	434	}
435		435
		436	+void NPUPluggableAllocator::attachAllocatorTraceTracker(
		437	+ c10_npu::NPUCachingAllocator::AllocatorTraceTracker tracker)
		438	+{
		439	+ (void)tracker;
		440	+ TORCH_CHECK(
		441	+ false,
		442	+ "NPUPluggableAllocator does not support attachAllocatorTraceTracker. "
		443	+ "attachAllocatorTraceTracker is only used inside Pytorch.",
		444	+ PTA_ERROR(ErrCode::NOT_SUPPORT));
		445	+}
		446	+
436	bool NPUPluggableAllocator::checkUceInMemPool(int device)	447	bool NPUPluggableAllocator::checkUceInMemPool(int device)
437	{	448	{
438	TORCH_NPU_WARN(	449	TORCH_NPU_WARN(

Mtorch_npu/csrc/npu/NPUPluggableAllocator.h+1-0

	@@ -101,6 +101,7 @@ struct NPUPluggableAllocator
101	size_t alloc_trace_max_entries,	101	size_t alloc_trace_max_entries,
102	c10_npu::NPUCachingAllocator::RecordContext when) override;	102	c10_npu::NPUCachingAllocator::RecordContext when) override;
103	void attachOutOfMemoryObserver(c10_npu::NPUCachingAllocator::OutOfMemoryObserver observer) override;	103	void attachOutOfMemoryObserver(c10_npu::NPUCachingAllocator::OutOfMemoryObserver observer) override;
		104	+ void attachAllocatorTraceTracker(c10_npu::NPUCachingAllocator::AllocatorTraceTracker tracker) override;
104	bool checkUceInMemPool(int device) override;	105	bool checkUceInMemPool(int device) override;
105	bool checkBlockIsSafe(const c10::DataPtr& ptr) override;	106	bool checkBlockIsSafe(const c10::DataPtr& ptr) override;
106	void markAllBlockUnsafe(int device) override;	107	void markAllBlockUnsafe(int device) override;