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

已合并

Support allocator trace tracker in NPU caching allocator #38687

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

Support allocator trace tracker in NPU caching allocator #38687

已合并

zzhongmin创建于 6月16日

从zzhongmin/pytorch v2.10.0_tratra合入到Ascend/pytorch v2.10.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
	@@ -227,6 +228,8 @@ struct BlockPool {
227	is_small(small),	228	is_small(small),
228	owner_PrivatePool(private_pool)	229	owner_PrivatePool(private_pool)
229	{}	230	{}
		231	+
		232	+ MempoolId_t owner_MempoolId() const;
230	};	233	};
231		234
232	struct ExpandableSegment;	235	struct ExpandableSegment;
	@@ -824,10 +827,12 @@ private:
824		827
825	// NPU graphs helper	828	// NPU graphs helper
826	struct PrivatePool {	829	struct PrivatePool {
827	- PrivatePool() : large_blocks(false, this), small_blocks(true, this) {}	830	+ explicit PrivatePool(MempoolId_t id)
		831	+ : id(std::move(id)), large_blocks(false, this), small_blocks(true, this) {}
828	PrivatePool(const PrivatePool &) = delete;	832	PrivatePool(const PrivatePool &) = delete;
829	PrivatePool(PrivatePool &&) = delete;	833	PrivatePool(PrivatePool &&) = delete;
830	PrivatePool &operator = (const PrivatePool &) = delete;	834	PrivatePool &operator = (const PrivatePool &) = delete;
		835	+ MempoolId_t id{ 0, 0 };
831	// Number of live graphs using this pool	836	// Number of live graphs using this pool
832	int use_count{ 1 };	837	int use_count{ 1 };
833	// Number of unfreed npuMallocs made for this pool. When use_count and	838	// Number of unfreed npuMallocs made for this pool. When use_count and
	@@ -843,6 +848,14 @@ struct PrivatePool {
843	BlockPool small_blocks;	848	BlockPool small_blocks;
844	};	849	};
845		850
		851	+MempoolId_t BlockPool::owner_MempoolId() const
		852	+{
		853	+ if (owner_PrivatePool) {
		854	+ return owner_PrivatePool->id;
		855	+ }
		856	+ return {0, 0};
		857	+}
		858	+
846	BlockState::BlockState(Block* block)	859	BlockState::BlockState(Block* block)
847	: device(block->device),	860	: device(block->device),
848	stream(block->stream),	861	stream(block->stream),
	@@ -1027,6 +1040,7 @@ private:
1027		1040
1028	// XXX - maybe we should generalize and have multiple events	1041	// XXX - maybe we should generalize and have multiple events
1029	std::vector<OutOfMemoryObserver> oom_observers_;	1042	std::vector<OutOfMemoryObserver> oom_observers_;
		1043	+ std::vector<AllocatorTraceTracker> trace_trackers_;
1030	std::shared_ptr<c10d_npu::HCCLComm> hcclComm_;	1044	std::shared_ptr<c10d_npu::HCCLComm> hcclComm_;
1031		1045
1032	// Private pools for NPU graphs	1046	// Private pools for NPU graphs
	@@ -1101,6 +1115,12 @@ public:
1101	oom_observers_.emplace_back(observer);	1115	oom_observers_.emplace_back(observer);
1102	}	1116	}
1103		1117
		1118	+ void attachAllocatorTraceTracker(AllocatorTraceTracker tracker)
		1119	+ {
		1120	+ std::unique_lock<std::recursive_mutex> lock(mutex);
		1121	+ trace_trackers_.emplace_back(std::move(tracker));
		1122	+ }
		1123	+
1104	bool checkUceInMemPool()	1124	bool checkUceInMemPool()
1105	{	1125	{
1106	auto memUceInfo_ = c10_npu::get_mem_uce_info();	1126	auto memUceInfo_ = c10_npu::get_mem_uce_info();
	@@ -1279,7 +1299,7 @@ public:
1279	stats.num_ooms += 1;	1299	stats.num_ooms += 1;
1280		1300
1281	record_trace(TraceEntry::OOM, device_free, params.size(), params.stream(), params.device(),	1301	record_trace(TraceEntry::OOM, device_free, params.size(), params.stream(), params.device(),
1282	- std::move(context));	1302	+ params.pool->owner_MempoolId(), std::move(context));
1283	auto observers_local = oom_observers_;	1303	auto observers_local = oom_observers_;
1284		1304
1285	// Make sure we do not have the device lock before calling our	1305	// Make sure we do not have the device lock before calling our
	@@ -1421,7 +1441,7 @@ public:
1421		1441
1422	block->context_when_allocated = std::move(context);	1442	block->context_when_allocated = std::move(context);
1423	record_trace(TraceEntry::ALLOC, int64_t(block->ptr), orig_size, block->stream, block->device,	1443	record_trace(TraceEntry::ALLOC, int64_t(block->ptr), orig_size, block->stream, block->device,
1424	- block->context_when_allocated);	1444	+ block->pool->owner_MempoolId(), block->context_when_allocated);
1425		1445
1426	active_blocks.insert(block);	1446	active_blocks.insert(block);
1427		1447
	@@ -1484,7 +1504,8 @@ public:
1484	});	1504	});
1485		1505
1486	record_trace(TraceEntry::FREE_REQUESTED, int64_t(block->ptr), block->requested_size, block->stream,	1506	record_trace(TraceEntry::FREE_REQUESTED, int64_t(block->ptr), block->requested_size, block->stream,
1487	- block->device, context ? context : block->context_when_allocated);	1507	+ block->device, block->pool->owner_MempoolId(),
		1508	+ context ? context : block->context_when_allocated);
1488		1509
1489	if (block->size >= NPUAllocatorConfig::max_split_size()) {	1510	if (block->size >= NPUAllocatorConfig::max_split_size()) {
1490	update_stat(stats.oversize_allocations, -1);	1511	update_stat(stats.oversize_allocations, -1);
	@@ -2066,7 +2087,7 @@ public:
2066	std::sort(result.begin(), result.end(),	2087	std::sort(result.begin(), result.end(),
2067	[](const SegmentInfo &a, const SegmentInfo &b) { return a.address < b.address; });	2088	[](const SegmentInfo &a, const SegmentInfo &b) { return a.address < b.address; });
2068		2089
2069	- record_trace(TraceEntry::SNAPSHOT, 0, total_active, nullptr, 0, nullptr);	2090	+ record_trace(TraceEntry::SNAPSHOT, 0, total_active, nullptr, 0, {0, 0}, nullptr);
2070	return result;	2091	return result;
2071	}	2092	}
2072		2093
	@@ -2126,26 +2147,41 @@ public:
2126	// See Note [Interaction with NPU graph capture]	2147	// See Note [Interaction with NPU graph capture]
2127		2148
2128	// Called by NPUGraph::capture_begin	2149	// Called by NPUGraph::capture_begin
2129	- void beginAllocateToPool(MempoolId_t mempool_id, std::function<bool(aclrtStream)> filter)	2150	+ void create_or_incref_pool(MempoolId_t mempool_id)
2130	{	2151	{
2131	- std::lock_guard<std::recursive_mutex> lock(mutex);
2132	auto it = graph_pools.find(mempool_id);	2152	auto it = graph_pools.find(mempool_id);
2133	if (it == graph_pools.end()) {	2153	if (it == graph_pools.end()) {
2134	- // mempool_id does not reference an existing pool. Make a new pool for	2154	+ // mempool_id does not reference an existing pool.
2135	- // this capture.	2155	+ // Make a new pool for NPUGraph capture or torch.npu.use_mem_pool
2136	- graph_pools.emplace(mempool_id, std::make_unique<PrivatePool>());	2156	+ // usage. use_count is initially 1, which means the pool is
		2157	+ // being used since somebody called createOrIncrefPool.
		2158	+ graph_pools.emplace(mempool_id, std::make_unique<PrivatePool>(mempool_id));
2137	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: new pool, "	2159	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: new pool, "
2138	"mempool_id=(%lu,%lu)", mempool_id.first, mempool_id.second);	2160	"mempool_id=(%lu,%lu)", mempool_id.first, mempool_id.second);
2139	} else {	2161	} else {
2140	- // mempool_id references an existing pool, which the current capture will	2162	+ // mempool_id references an existing pool, which the current NPUGraph
		2163	+ // capture or torch.npu.use_mem_pool will
2141	// share. Check this pool is live (at least one other capture already	2164	// share. Check this pool is live (at least one other capture already
2142	- // references it).	2165	+ // references it). Increment it to establish the usage.
2143	TORCH_INTERNAL_ASSERT(it->second->use_count > 0);	2166	TORCH_INTERNAL_ASSERT(it->second->use_count > 0);
2144	it->second->use_count++;	2167	it->second->use_count++;
2145	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: reuse pool, "	2168	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator beginAllocateToPool: reuse pool, "
2146	"mempool_id=(%lu,%lu), use_count=%d",	2169	"mempool_id=(%lu,%lu), use_count=%d",
2147	mempool_id.first, mempool_id.second, it->second->use_count);	2170	mempool_id.first, mempool_id.second, it->second->use_count);
2148	}	2171	}
		2172	+ }
		2173	+
		2174	+ PrivatePool* get_private_pool(MempoolId_t mempool_id) const
		2175	+ {
		2176	+ auto it = graph_pools.find(mempool_id);
		2177	+ TORCH_INTERNAL_ASSERT(it != graph_pools.end());
		2178	+ return it->second.get();
		2179	+ }
		2180	+
		2181	+ void beginAllocateToPool(MempoolId_t mempool_id, std::function<bool(aclrtStream)> filter)
		2182	+ {
		2183	+ std::lock_guard<std::recursive_mutex> lock(mutex);
		2184	+ create_or_incref_pool(mempool_id);
2149	for (auto it2 = captures_underway.begin(); it2 != captures_underway.end(); ++it2) {	2185	for (auto it2 = captures_underway.begin(); it2 != captures_underway.end(); ++it2) {
2150	TORCH_CHECK(it2->first != mempool_id, "beginAllocateToPool: already recording to mempool_id");	2186	TORCH_CHECK(it2->first != mempool_id, "beginAllocateToPool: already recording to mempool_id");
2151	}	2187	}
	@@ -2186,15 +2222,14 @@ public:
2186	// mempool. When the count reaches 0, we tell free_cached_blocks it may now	2222	// mempool. When the count reaches 0, we tell free_cached_blocks it may now
2187	// npuFree blocks from this graph's pool when it discovers they're unused	2223	// npuFree blocks from this graph's pool when it discovers they're unused
2188	// (unsplit).	2224	// (unsplit).
2189	- auto it = graph_pools.find(mempool_id);	2225	+ auto pp = get_private_pool(mempool_id);
2190	- TORCH_INTERNAL_ASSERT(it != graph_pools.end());	2226	+ auto uc = --(pp->use_count);
2191	- auto uc = --(it->second->use_count);
2192	TORCH_INTERNAL_ASSERT(uc >= 0);	2227	TORCH_INTERNAL_ASSERT(uc >= 0);
2193	if (uc == 0) {	2228	if (uc == 0) {
2194	// Allows free_cached_blocks to begin npuFreeing this pool's memory,	2229	// Allows free_cached_blocks to begin npuFreeing this pool's memory,
2195	// and makes sure this pool wasn't somehow made freeable already.	2230	// and makes sure this pool wasn't somehow made freeable already.
2196	// NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)	2231	// NOLINTNEXTLINE(clang-analyzer-deadcode.DeadStores)
2197	- bool inserted = graph_pools_freeable.insert({ mempool_id, it->second.get() }).second;	2232	+ bool inserted = graph_pools_freeable.insert({ mempool_id, pp }).second;
2198	TORCH_INTERNAL_ASSERT(inserted);	2233	TORCH_INTERNAL_ASSERT(inserted);
2199	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator releasePool: mempool_id=(%lu,%lu), "	2234	TORCH_NPU_MEMORY_LOGD("NPUCachingAllocator releasePool: mempool_id=(%lu,%lu), "
2200	"use_count reached 0, pool marked freeable",	2235	"use_count reached 0, pool marked freeable",
	@@ -2343,7 +2378,7 @@ private:
2343	for_each_selected_stat_type(stat_types,	2378	for_each_selected_stat_type(stat_types,
2344	[&](size_t stat_type) { update_stat(stats.reserved_bytes[stat_type], mapped_range.size); });	2379	[&](size_t stat_type) { update_stat(stats.reserved_bytes[stat_type], mapped_range.size); });
2345	record_trace(TraceEntry::SEGMENT_MAP, int64_t(mapped_range.ptr), mapped_range.size, to_map->stream,	2380	record_trace(TraceEntry::SEGMENT_MAP, int64_t(mapped_range.ptr), mapped_range.size, to_map->stream,
2346	- to_map->device, ctx);	2381	+ to_map->device, to_map->pool->owner_MempoolId(), ctx);
2347	if (!to_map->prev && !to_map->context_when_segment_allocated) {	2382	if (!to_map->prev && !to_map->context_when_segment_allocated) {
2348	to_map->context_when_segment_allocated = ctx;	2383	to_map->context_when_segment_allocated = ctx;
2349	}	2384	}
	@@ -2390,7 +2425,8 @@ private:
2390	AT_ASSERT(!block->allocated && block->event_count == 0, PTA_ERROR(ErrCode::VALUE));	2425	AT_ASSERT(!block->allocated && block->event_count == 0, PTA_ERROR(ErrCode::VALUE));
2391		2426
2392	record_trace(TraceEntry::FREE_COMPLETED, int64_t(block->ptr), block->requested_size, block->stream,	2427	record_trace(TraceEntry::FREE_COMPLETED, int64_t(block->ptr), block->requested_size, block->stream,
2393	- block->device, context ? context : block->context_when_allocated);	2428	+ block->device, block->pool->owner_MempoolId(),
		2429	+ context ? context : block->context_when_allocated);
2394		2430
2395	block->context_when_allocated = nullptr;	2431	block->context_when_allocated = nullptr;
2396	block->hccl_work_ptr = nullptr;	2432	block->hccl_work_ptr = nullptr;
	@@ -2746,7 +2782,8 @@ private:
2746	// p.block came from new, not npuMalloc. It should not be nullptr here.	2782	// p.block came from new, not npuMalloc. It should not be nullptr here.
2747	TORCH_INTERNAL_ASSERT(p.block != nullptr && p.block->ptr != nullptr);	2783	TORCH_INTERNAL_ASSERT(p.block != nullptr && p.block->ptr != nullptr);
2748		2784
2749	- record_trace(TraceEntry::SEGMENT_ALLOC, int64_t(p.block->ptr), p.block->size, p.stream(), p.device(), ctx);	2785	+ record_trace(TraceEntry::SEGMENT_ALLOC, int64_t(p.block->ptr), p.block->size, p.stream(), p.device(),
		2786	+ p.block->pool->owner_MempoolId(), ctx);
2750	p.block->context_when_segment_allocated = ctx;	2787	p.block->context_when_segment_allocated = ctx;
2751	return true;	2788	return true;
2752	}	2789	}
	@@ -2853,6 +2890,7 @@ private:
2853	block->size, block->ptr, block->device);	2890	block->size, block->ptr, block->device);
2854		2891
2855	record_trace(TraceEntry::SEGMENT_FREE, int64_t(block->ptr), block->size, block->stream, block->device,	2892	record_trace(TraceEntry::SEGMENT_FREE, int64_t(block->ptr), block->size, block->stream, block->device,
		2893	+ block->pool->owner_MempoolId(),
2856	context ? context : block->context_when_segment_allocated);	2894	context ? context : block->context_when_segment_allocated);
2857		2895
2858	auto it = ipc_handle_map.find(block->ptr);	2896	auto it = ipc_handle_map.find(block->ptr);
	@@ -2937,6 +2975,7 @@ private:
2937	}	2975	}
2938		2976
2939	record_trace(TraceEntry::SEGMENT_UNMAP, int64_t(unmapped.ptr), unmapped.size, block->stream, block->device,	2977	record_trace(TraceEntry::SEGMENT_UNMAP, int64_t(unmapped.ptr), unmapped.size, block->stream, block->device,
		2978	+ block->pool->owner_MempoolId(),
2940	context ? context : block->context_when_segment_allocated);	2979	context ? context : block->context_when_segment_allocated);
2941	}	2980	}
2942		2981
	@@ -3129,15 +3168,25 @@ private:
3129	}	3168	}
3130		3169
3131	void record_trace(TraceEntry::Action action, int64_t addr, size_t size, aclrtStream stream, int device,	3170	void record_trace(TraceEntry::Action action, int64_t addr, size_t size, aclrtStream stream, int device,
3132	- std::shared_ptr<c10::GatheredContext> context)	3171	+ MempoolId_t mempool_id = {0, 0}, std::shared_ptr<c10::GatheredContext> context = nullptr)
3133	{	3172	{
3134	- if (!record_history) {	3173	+ if (!record_history && trace_trackers_.empty()) {
3135	return;	3174	return;
3136	}	3175	}
3137	-	3176	+ auto te = TraceEntry(
3138	- auto te = TraceEntry(action, device, addr, size, stream,	3177	+ action,
		3178	+ device,
		3179	+ addr,
		3180	+ size,
		3181	+ stream,
		3182	+ mempool_id,
3139	record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);	3183	record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);
3140		3184
		3185	+ // Callbacks should not include any Pytorch call
		3186	+ for (const auto& cb : trace_trackers_) {
		3187	+ cb(te);
		3188	+ }
		3189	+
3141	if (record_history) {	3190	if (record_history) {
3142	if (alloc_trace->size() < alloc_trace_max_entries_) {	3191	if (alloc_trace->size() < alloc_trace_max_entries_) {
3143	alloc_trace->emplace_back(te);	3192	alloc_trace->emplace_back(te);
	@@ -3295,6 +3344,13 @@ public:
3295	}	3344	}
3296	}	3345	}
3297		3346
		3347	+ void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override
		3348	+ {
		3349	+ for (auto& allocator : device_allocator) {
		3350	+ allocator->attachAllocatorTraceTracker(tracker);
		3351	+ }
		3352	+ }
		3353	+
3298	bool checkUceInMemPool(int device) override	3354	bool checkUceInMemPool(int device) override
3299	{	3355	{
3300	return device_allocator[device]->checkUceInMemPool();	3356	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;