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

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+19-1

	@@ -1025,6 +1025,7 @@ private:
1025		1025
1026	// XXX - maybe we should generalize and have multiple events	1026	// XXX - maybe we should generalize and have multiple events
1027	std::vector<OutOfMemoryObserver> oom_observers_;	1027	std::vector<OutOfMemoryObserver> oom_observers_;
		1028	+ std::vector<AllocatorTraceTracker> trace_trackers_;
1028	std::shared_ptr<c10d_npu::HCCLComm> hcclComm_;	1029	std::shared_ptr<c10d_npu::HCCLComm> hcclComm_;
1029		1030
1030	// Private pools for NPU graphs	1031	// Private pools for NPU graphs
	@@ -1098,6 +1099,12 @@ public:
1098	oom_observers_.emplace_back(observer);	1099	oom_observers_.emplace_back(observer);
1099	}	1100	}
1100		1101
		1102	+ void attachAllocatorTraceTracker(AllocatorTraceTracker tracker)
		1103	+ {
		1104	+ std::unique_lock<std::recursive_mutex> lock(mutex);
		1105	+ trace_trackers_.emplace_back(std::move(tracker));
		1106	+ }
		1107	+
1101	bool checkUceInMemPool()	1108	bool checkUceInMemPool()
1102	{	1109	{
1103	auto memUceInfo_ = c10_npu::get_mem_uce_info();	1110	auto memUceInfo_ = c10_npu::get_mem_uce_info();
	@@ -3114,13 +3121,17 @@ private:
3114	void record_trace(TraceEntry::Action action, int64_t addr, size_t size, aclrtStream stream, int device,	3121	void record_trace(TraceEntry::Action action, int64_t addr, size_t size, aclrtStream stream, int device,
3115	std::shared_ptr<c10::GatheredContext> context)	3122	std::shared_ptr<c10::GatheredContext> context)
3116	{	3123	{
3117	- if (!record_history) {	3124	+ if (!record_history && trace_trackers_.empty()) {
3118	return;	3125	return;
3119	}	3126	}
3120		3127
3121	auto te = TraceEntry(action, device, addr, size, stream,	3128	auto te = TraceEntry(action, device, addr, size, stream,
3122	record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);	3129	record_context_ >= RecordContext::ALLOC ? std::move(context) : nullptr);
3123		3130
		3131	+ for (const auto& cb : trace_trackers_) {
		3132	+ cb(te);
		3133	+ }
		3134	+
3124	if (record_history) {	3135	if (record_history) {
3125	if (alloc_trace->size() < alloc_trace_max_entries_) {	3136	if (alloc_trace->size() < alloc_trace_max_entries_) {
3126	alloc_trace->emplace_back(te);	3137	alloc_trace->emplace_back(te);
	@@ -3278,6 +3289,13 @@ public:
3278	}	3289	}
3279	}	3290	}
3280		3291
		3292	+ void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override
		3293	+ {
		3294	+ for (auto& allocator : device_allocator) {
		3295	+ allocator->attachAllocatorTraceTracker(tracker);
		3296	+ }
		3297	+ }
		3298	+
3281	bool checkUceInMemPool(int device) override	3299	bool checkUceInMemPool(int device) override
3282	{	3300	{
3283	return device_allocator[device]->checkUceInMemPool();	3301	return device_allocator[device]->checkUceInMemPool();

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.h+7-0

	@@ -219,6 +219,7 @@ enum struct RecordContext {
219	using OutOfMemoryObserver =	219	using OutOfMemoryObserver =
220	std::function<void(int64_t device, int64_t allocated, int64_t device_total,	220	std::function<void(int64_t device, int64_t allocated, int64_t device_total,
221	int64_t device_free)>;	221	int64_t device_free)>;
		222	+using AllocatorTraceTracker = std::function<void(const TraceEntry&)>;
222		223
223	struct ShareableHandle {	224	struct ShareableHandle {
224	ptrdiff_t offset;	225	ptrdiff_t offset;
	@@ -284,6 +285,7 @@ public:
284	size_t alloc_trace_max_entries,	285	size_t alloc_trace_max_entries,
285	RecordContext when) = 0;	286	RecordContext when) = 0;
286	virtual void attachOutOfMemoryObserver(OutOfMemoryObserver observer) = 0;	287	virtual void attachOutOfMemoryObserver(OutOfMemoryObserver observer) = 0;
		288	+ virtual void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) = 0;
287	virtual bool checkUceInMemPool(int device) = 0;	289	virtual bool checkUceInMemPool(int device) = 0;
288	virtual bool checkBlockIsSafe(const c10::DataPtr& ptr) = 0;	290	virtual bool checkBlockIsSafe(const c10::DataPtr& ptr) = 0;
289	virtual void markAllBlockUnsafe(int device) = 0;	291	virtual void markAllBlockUnsafe(int device) = 0;
	@@ -502,6 +504,11 @@ inline void attachOutOfMemoryObserver(OutOfMemoryObserver observer)
502	return get()->attachOutOfMemoryObserver(observer);	504	return get()->attachOutOfMemoryObserver(observer);
503	}	505	}
504		506
		507	+inline void attachAllocatorTraceTracker(AllocatorTraceTracker tracker)
		508	+{
		509	+ return get()->attachAllocatorTraceTracker(std::move(tracker));
		510	+}
		511	+
505	inline bool checkUceInMemPool(int device)	512	inline bool checkUceInMemPool(int device)
506	{	513	{
507	return get()->checkUceInMemPool(device);	514	return get()->checkUceInMemPool(device);

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;