Align with cuda HostCachingAllocator 2.8-pytorch-AtomGit

已合并

Align with cuda HostCachingAllocator 2.8 #26736

zihao65创建于 2025年11月19日从gcw_5tF58QLT/pytorch： host2.8-2.9 合入到 Ascend/pytorch： v2.8.0

Align with cuda HostCachingAllocator 2.8 #26736

已合并

zihao65创建于 2025年11月19日

从gcw_5tF58QLT/pytorch host2.8-2.9合入到Ascend/pytorch v2.8.0

共 5 个文件变更+465-285

Atest/allocator/test_hostallocator.py+323-0

@@ -0,0 +1,323 @@
	1	+import faulthandler
	2	+import gc
	3	+import importlib
	4	+import os
	5	+import signal
	6	+import sys
	7	+import threading
	8	+import unittest
	9	+from itertools import product
	10	+
	11	+import torch
	12	+from torch.testing._internal.common_utils import IS_WINDOWS
	13	+from torch.utils.data import DataLoader, Dataset, TensorDataset
	14	+
	15	+import torch_npu
	16	+from torch_npu.testing.testcase import TestCase, run_tests
	17	+
	18	+TEST_NPU = torch.npu.is_available()
	19	+
	20	+
	21	+def _collect():
	22	+ gc.collect()
	23	+ torch.npu.empty_cache()
	24	+
	25	+
	26	+class TestHostCachingAllocatorBasic(TestCase):
	27	+ def test_pin_memory_on_non_blocking_copy(self):
	28	+ t_acc = torch.randn(100).to(torch.accelerator.current_accelerator())
	29	+ t_host = t_acc.to("cpu", non_blocking=True)
	30	+ torch.accelerator.synchronize()
	31	+ self.assertTrue(t_host.is_pinned())
	32	+ self.assertEqual(t_acc.cpu(), t_host)
	33	+
	34	+ def test_pin_memory_reuse(self):
	35	+ t = torch.FloatTensor([1]).pin_memory()
	36	+ ptr = t.data_ptr()
	37	+ del t
	38	+ t_new = torch.FloatTensor([1]).pin_memory()
	39	+ self.assertEqual(t_new.data_ptr(), ptr)
	40	+
	41	+ def test_to_non_blocking(self):
	42	+ stream = torch_npu.npu.current_stream()
	43	+
	44	+ def _test_to_non_blocking(a, non_blocking, dst):
	45	+ torch_npu.npu.synchronize()
	46	+ b = a.to(device=dst, non_blocking=non_blocking)
	47	+ stream.synchronize()
	48	+ self.assertEqual(a, b)
	49	+ self.assertTrue(b.is_pinned() == (non_blocking and dst == "cpu"))
	50	+
	51	+ for dst, try_non_blocking in [
	52	+ ("npu", True),
	53	+ ("npu", False),
	54	+ ("cpu", True), # pinned is true only when non_blocking=True and dst="cpu"
	55	+ ("cpu", False),
	56	+ ]:
	57	+ # Creates source on the opposite device from destination.
	58	+ src = torch.randn(1000, 1000, 2, 100,
	59	+ device="npu" if dst == "cpu" else "cpu",
	60	+ pin_memory=True if dst == "npu" else False)
	61	+ _test_to_non_blocking(src, try_non_blocking, dst)
	62	+
	63	+ def test_pin_memory_basic(self):
	64	+ a = torch.Tensor([1])
	65	+ b = a.pin_memory()
	66	+ c = a.pin_memory()
	67	+ d = b.pin_memory()
	68	+ self.assertTrue(a.data_ptr() != b.data_ptr())
	69	+ self.assertTrue(b.data_ptr() != c.data_ptr())
	70	+ self.assertTrue(b.data_ptr() == d.data_ptr())
	71	+
	72	+ def test_malloc_copykernel(self):
	73	+ a = torch.Tensor([1])
	74	+ b = torch.Tensor([1])
	75	+ c = a.to('npu:0', non_blocking=True)
	76	+ d = b.to('npu:0', non_blocking=True)
	77	+ # we do not synchronize here
	78	+ # the above to will call synchronize internally
	79	+ self.assertEqual(c.item(), d.item())
	80	+
	81	+ def test_fragmentation_resilience_varied_sizes(self):
	82	+ # Allocate/release varied sizes and ensure subsequent allocation succeeds and is pinned
	83	+ sizes = [10_000, 200_000, 50_000, 1_000_000, 75_000]
	84	+ tensors = [torch.empty(s, dtype=torch.float32).pin_memory() for s in sizes]
	85	+ for t in tensors:
	86	+ self.assertTrue(t.is_pinned())
	87	+ # Free in different order to stress cache
	88	+ for idx in [2, 0, 4, 1, 3]:
	89	+ tensors[idx] = None
	90	+ _collect()
	91	+
	92	+ # Allocate a size that fits into previously freed blocks
	93	+ t_new = torch.empty(150_000, dtype=torch.float32).pin_memory()
	94	+ self.assertTrue(t_new.is_pinned())
	95	+ # If you have allocator stats, assert on reuse; otherwise, just sanity check
	96	+ del t_new
	97	+ _collect()
	98	+
	99	+ def test_concurrent_pinned_allocations_threaded(self):
	100	+ errs = []
	101	+
	102	+ def worker(n):
	103	+ try:
	104	+ for _ in range(n):
	105	+ t = torch.empty(256, 256, dtype=torch.float32).pin_memory()
	106	+ self.assertTrue(t.is_pinned())
	107	+ # Move to GPU non-blocking to simulate pipeline
	108	+ _ = t.to("npu", non_blocking=True)
	109	+ except Exception as e:
	110	+ errs.append(e)
	111	+
	112	+ threads = [threading.Thread(target=worker, args=(50,)) for _ in range(4)]
	113	+ for th in threads:
	114	+ th.start()
	115	+ for th in threads:
	116	+ th.join()
	117	+
	118	+ torch.npu.synchronize()
	119	+ self.assertTrue(not errs)
	120	+
	121	+ def test_pin_memory_on_views_and_clones(self):
	122	+ base = torch.randn(1024, 1024)
	123	+ view = base[:512, :].pin_memory()
	124	+ clone = base.clone().pin_memory()
	125	+ self.assertTrue(view.is_pinned())
	126	+ self.assertTrue(clone.is_pinned())
	127	+ # Ensure their contents are consistent after to('cuda')
	128	+ yv = view.to("npu", non_blocking=True)
	129	+ yc = clone.to("npu", non_blocking=True)
	130	+ torch.npu.synchronize()
	131	+ self.assertTrue(yv.device and yc.device)
	132	+
	133	+ def test_pin_memory_on_dtypes_and_non_contiguous(self):
	134	+ x = torch.randn(128, 128, dtype=torch.float64).t() # non-contiguous
	135	+ xp = x.pin_memory()
	136	+ self.assertTrue(xp.is_pinned())
	137	+ # Transfer works even if non-contiguous (PyTorch will handle copy)
	138	+ y = xp.to("npu", non_blocking=True)
	139	+ torch.npu.synchronize()
	140	+ self.assertTrue(y.is_npu)
	141	+
	142	+
	143	+def set_faulthander_if_available(_=None):
	144	+ faulthandler.enable(sys.__stderr__)
	145	+ if not IS_WINDOWS:
	146	+ faulthandler.register(signal.SIGUSR1, file=sys.__stderr__, chain=False)
	147	+
	148	+set_faulthander_if_available()
	149	+
	150	+
	151	+class CountingDataset(Dataset):
	152	+ def __init__(self, n):
	153	+ super().__init__()
	154	+ self.n = n
	155	+
	156	+ def __getitem__(self, i):
	157	+ return i
	158	+
	159	+ def __len__(self):
	160	+ return self.n
	161	+
	162	+
	163	+class DictDataset(Dataset):
	164	+ def __len__(self):
	165	+ return 4
	166	+
	167	+ def __getitem__(self, ndx):
	168	+ return {
	169	+ 'a_tensor': torch.empty(4, 2).fill_(ndx),
	170	+ 'another_dict': {
	171	+ 'a_number': torch.tensor(ndx),
	172	+ },
	173	+ }
	174	+
	175	+
	176	+class StringDataset(Dataset):
	177	+ def __init__(self):
	178	+ self.s = '12345'
	179	+
	180	+ def __len__(self):
	181	+ return len(self.s)
	182	+
	183	+ def __getitem__(self, ndx):
	184	+ return (self.s[ndx], ndx)
	185	+
	186	+
	187	+class SimpleCustomBatch:
	188	+ def __init__(self, data):
	189	+ transposed_data = list(zip(*data))
	190	+ self.inp = torch.stack(transposed_data[0], 0)
	191	+ self.tgt = torch.stack(transposed_data[1], 0)
	192	+
	193	+ def pin_memory(self):
	194	+ self.inp = self.inp.pin_memory()
	195	+ self.tgt = self.tgt.pin_memory()
	196	+ return self
	197	+
	198	+ def is_pinned(self):
	199	+ return self.inp.is_pinned() and self.tgt.is_pinned()
	200	+
	201	+module_name = os.path.splitext(os.path.basename(__file__))[0]
	202	+self_module = importlib.import_module(module_name)
	203	+
	204	+
	205	+def collate_wrapper(batch):
	206	+ return self_module.SimpleCustomBatch(batch)
	207	+
	208	+
	209	+def collate_into_packed_sequence(batch):
	210	+ data = torch.stack([sample[0] for sample in batch], 1)
	211	+ t, b = data.size()
	212	+ lengths = torch.randint(1, t, size=(b,), dtype=torch.int64)
	213	+ return torch.nn.utils.rnn.pack_padded_sequence(data, lengths, enforce_sorted=False)
	214	+
	215	+
	216	+def collate_into_packed_sequence_batch_first(batch):
	217	+ data = torch.stack([sample[0] for sample in batch], 0)
	218	+ b, t = data.size()
	219	+ lengths = torch.randint(1, t, size=(b,), dtype=torch.int64)
	220	+ return torch.nn.utils.rnn.pack_padded_sequence(data, lengths, batch_first=True, enforce_sorted=False)
	221	+
	222	+
	223	+class TestDataLoader(TestCase):
	224	+ def setUp(self):
	225	+ super().setUp()
	226	+ self.data = torch.randn(100, 2, 3, 5)
	227	+ self.labels = torch.randperm(50).repeat(2)
	228	+ self.dataset = TensorDataset(self.data, self.labels)
	229	+
	230	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	231	+ def test_sequential_pin_memory(self):
	232	+ loader = DataLoader(self.dataset, batch_size=2, pin_memory=True, pin_memory_device='npu')
	233	+ for input_, target in loader:
	234	+ self.assertTrue(input_.is_pinned())
	235	+ self.assertTrue(target.is_pinned())
	236	+
	237	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	238	+ def test_shuffle_pin_memory(self):
	239	+ loader = DataLoader(self.dataset, batch_size=2, shuffle=True, num_workers=4,
	240	+ pin_memory=True, pin_memory_device='npu')
	241	+ for input_, target in loader:
	242	+ self.assertTrue(input_.is_pinned())
	243	+ self.assertTrue(target.is_pinned())
	244	+
	245	+
	246	+class TestStringDataLoader(TestCase):
	247	+ def setUp(self):
	248	+ super().setUp()
	249	+ self.dataset = StringDataset()
	250	+
	251	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	252	+ def test_shuffle_pin_memory(self):
	253	+ loader = DataLoader(self.dataset, batch_size=2, shuffle=True, num_workers=4, pin_memory=True)
	254	+ for (s, n) in loader:
	255	+ self.assertIsInstance(s[0], str)
	256	+ self.assertTrue(n.is_pinned())
	257	+
	258	+
	259	+class TestDictDataLoader(TestCase):
	260	+ def setUp(self):
	261	+ super().setUp()
	262	+ self.dataset = DictDataset()
	263	+
	264	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	265	+ def test_pin_memory(self):
	266	+ loader = DataLoader(self.dataset, batch_size=2, pin_memory=True)
	267	+ for sample in loader:
	268	+ self.assertTrue(sample['a_tensor'].is_pinned())
	269	+ self.assertTrue(sample['another_dict']['a_number'].is_pinned())
	270	+
	271	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	272	+ def test_pin_memory_device(self):
	273	+ loader = DataLoader(self.dataset, batch_size=2, pin_memory=True, pin_memory_device='npu')
	274	+ for sample in loader:
	275	+ self.assertTrue(sample['a_tensor'].is_pinned(device='npu'))
	276	+ self.assertTrue(sample['another_dict']['a_number'].is_pinned(device='npu'))
	277	+
	278	+
	279	+class TestCustomPinFn(TestCase):
	280	+ def setUp(self):
	281	+ super().setUp()
	282	+ inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
	283	+ tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
	284	+ self.dataset = TensorDataset(inps, tgts)
	285	+
	286	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	287	+ def test_custom_batch_pin(self):
	288	+ test_cases = [
	289	+ (collate_wrapper, self_module.SimpleCustomBatch),
	290	+ (collate_into_packed_sequence, torch.nn.utils.rnn.PackedSequence),
	291	+ (collate_into_packed_sequence_batch_first, torch.nn.utils.rnn.PackedSequence),
	292	+ ]
	293	+ for collate_fn, elem_cls in test_cases:
	294	+ loader = DataLoader(self.dataset, batch_size=2, collate_fn=collate_fn,
	295	+ pin_memory=True, pin_memory_device='npu')
	296	+ for sample in loader:
	297	+ self.assertIsInstance(sample, elem_cls)
	298	+ # 对于 PackedSequence，is_pinned 在 DataLoader 中会递归 pin 其 data
	299	+ if hasattr(sample, 'is_pinned'):
	300	+ self.assertTrue(sample.is_pinned())
	301	+ else:
	302	+ # PackedSequence: 检查其 data
	303	+ self.assertTrue(sample.data.is_pinned())
	304	+
	305	+ @unittest.skipIf(not TEST_NPU, "NPU unavailable")
	306	+ def test_custom_batch_pin_worker(self):
	307	+ test_cases = [
	308	+ (collate_wrapper, self_module.SimpleCustomBatch),
	309	+ (collate_into_packed_sequence, torch.nn.utils.rnn.PackedSequence),
	310	+ (collate_into_packed_sequence_batch_first, torch.nn.utils.rnn.PackedSequence),
	311	+ ]
	312	+ for collate_fn, elem_cls in test_cases:
	313	+ loader = DataLoader(self.dataset, batch_size=2, collate_fn=collate_fn,
	314	+ pin_memory=True, num_workers=1, pin_memory_device='npu')
	315	+ for sample in loader:
	316	+ self.assertIsInstance(sample, elem_cls)
	317	+ if hasattr(sample, 'is_pinned'):
	318	+ self.assertTrue(sample.is_pinned())
	319	+ else:
	320	+ self.assertTrue(sample.data.is_pinned())
	321	+
	322	+if __name__ == "__main__":
	323	+ run_tests()

Mtorch_npu/csrc/aten/common/CopyKernel.cpp+7-2

	@@ -1,4 +1,5 @@
1	#include <ATen/ATen.h>	1	#include <ATen/ATen.h>
		2	+#include <ATen/core/CachingHostAllocator.h>
2		3
3	#include "op_plugin/OpInterface.h"	4	#include "op_plugin/OpInterface.h"
4	#include "torch_npu/csrc/core/npu/NPUException.h"	5	#include "torch_npu/csrc/core/npu/NPUException.h"
	@@ -98,8 +99,12 @@ void copy_between_host_and_device(
98	auto ret = CalcuOpUtil::LaunchAsyncCopyTaskWithModeSwitch(dst, nbytes, src, nbytes, kind);	99	auto ret = CalcuOpUtil::LaunchAsyncCopyTaskWithModeSwitch(dst, nbytes, src, nbytes, kind);
99	NPU_CHECK_ERROR(ret);	100	NPU_CHECK_ERROR(ret);
100	ASCEND_LOGD("non_blocking copy without StreamSynchronize.");	101	ASCEND_LOGD("non_blocking copy without StreamSynchronize.");
101	- void* ptr = torch_npu::utils::is_npu(dst) ? src.storage().mutable_data() : dst.storage().mutable_data();	102	+ auto& storage = torch_npu::utils::is_npu(dst) ? src.storage() : dst.storage();
102	- NPU_CHECK_ERROR(CachingHostAllocator_recordEvent(ptr, kind, stream), "aclrtSynchronizeStreamWithTimeout");	103	+ if (!at_npu::native::ptr_exist(storage.mutable_data()) && (!c10_npu::acl::AclrtMemcpyAsyncWithConditionExist() \|\| (kind != aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST))) {
		104	+ // Sync when host memory is allocated by malloc
		105	+ NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream), "ACL stream synchronize failed.");
		106	+ }
		107	+ at::getHostAllocator(at::kPrivateUse1)->record_event(storage.mutable_data(), storage.data_ptr().get_context(), stream);
103	} else {	108	} else {
104	aclError error = c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream);	109	aclError error = c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream);
105	auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch(	110	auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch(

Mtorch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp+7-2

	@@ -13,6 +13,7 @@
13	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.	13	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14	// See the License for the specific language governing permissions and	14	// See the License for the specific language governing permissions and
15	// limitations under the License.	15	// limitations under the License.
		16	+#include <ATen/core/CachingHostAllocator.h>
16		17
17	#include "torch_npu/csrc/core/npu/NPUGuard.h"	18	#include "torch_npu/csrc/core/npu/NPUGuard.h"
18	#include "torch_npu/csrc/core/npu/NPUPeerToPeerAccess.h"	19	#include "torch_npu/csrc/core/npu/NPUPeerToPeerAccess.h"
	@@ -44,8 +45,12 @@ void copy_between_host_and_device_opapi(at::Tensor& dst, const at::Tensor& src,
44	auto ret = CalcuOpUtil::LaunchAsyncCopyTaskWithModeSwitch(dst, nbytes, src, nbytes, kind);	45	auto ret = CalcuOpUtil::LaunchAsyncCopyTaskWithModeSwitch(dst, nbytes, src, nbytes, kind);
45	NPU_CHECK_ERROR(ret);	46	NPU_CHECK_ERROR(ret);
46	ASCEND_LOGD("non_blocking copy without StreamSynchronize.");	47	ASCEND_LOGD("non_blocking copy without StreamSynchronize.");
47	- void* ptr = torch_npu::utils::is_npu(dst) ? src.storage().mutable_data() : dst.storage().mutable_data();	48	+ auto& storage = torch_npu::utils::is_npu(dst) ? src.storage() : dst.storage();
48	- NPU_CHECK_ERROR(CachingHostAllocator_recordEvent(ptr, kind, stream), "aclrtSynchronizeStreamWithTimeout");	49	+ if (!at_npu::native::ptr_exist(storage.mutable_data()) && (!c10_npu::acl::AclrtMemcpyAsyncWithConditionExist() \|\| (kind != aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST))) {
		50	+ // Sync when host memory is allocated by malloc
		51	+ NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream), "ACL stream synchronize failed.");
		52	+ }
		53	+ at::getHostAllocator(at::kPrivateUse1)->record_event(storage.mutable_data(), storage.data_ptr().get_context(), stream);
49	} else {	54	} else {
50	aclError error = aclrtSynchronizeStream(stream);	55	aclError error = aclrtSynchronizeStream(stream);
51	auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch(	56	auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch(

Mtorch_npu/csrc/core/npu/CachingHostAllocator.cpp+93-273

	@@ -1,3 +1,7 @@
		1	+#include <ATen/DeviceGuard.h>
		2	+#include <c10/core/thread_pool.h>
		3	+#include <future>
		4	+
1	#include <c10/core/DeviceGuard.h>	5	#include <c10/core/DeviceGuard.h>
2	#include <c10/util/llvmMathExtras.h>	6	#include <c10/util/llvmMathExtras.h>
3	#include "torch_npu/csrc/core/npu/npu_log.h"	7	#include "torch_npu/csrc/core/npu/npu_log.h"
	@@ -24,38 +28,39 @@
24	#include "torch_npu/csrc/core/npu/CachingHostAllocator.h"	28	#include "torch_npu/csrc/core/npu/CachingHostAllocator.h"
25	#include "torch_npu/csrc/core/npu/NPUEvent.h"	29	#include "torch_npu/csrc/core/npu/NPUEvent.h"
26		30
27	-namespace at_npu {	31	+#include <ATen/core/CachingHostAllocator.h>
28	-namespace native {	32	+#include <c10/util/flat_hash_map.h>
29		33
30	-namespace {	34	+#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
31	-struct BlockSize {
32	- size_t size; // allocation size
33	- void *ptr; // host memory pointer
34		35
35	- explicit BlockSize(size_t size, void *ptr = nullptr) : size(size), ptr(ptr) {}	36	+using Block = at::HostBlock<c10_npu::NPUStream>;
36	-};
37		37
38	-struct Block : public BlockSize {	38	+namespace at_npu::native {
39	- bool allocated; // true if the block is currently allocated	39	+
40	- int event_count; // number of outstanding npu events	40	+void innerInitNPU()
41	- std::unordered_set<c10_npu::NPUStream> streams;	41	+{
42	- Block(size_t size, void *ptr, bool allocated)	42	+ // check whether we need std::call_once here
43	- : BlockSize(size, ptr), allocated(allocated), event_count(0), streams() {}	43	+ C10_LOG_API_USAGE_ONCE("aten.init.npu");
44	-};	44	+ c10_npu::NpuSysCtrl::SysStatus status =
		45	+ c10_npu::NpuSysCtrl::GetInstance().Initialize();
		46	+ if (status != c10_npu::NpuSysCtrl::SysStatus::INIT_SUCC) {
		47	+ ASCEND_LOGE("Npu init fail.");
		48	+ }
		49	+}
45		50
46	class EventPool {	51	class EventPool {
47	public:	52	public:
48	using Event = std::unique_ptr<	53	using Event = std::unique_ptr<
49	c10_npu::NPUEvent,	54	c10_npu::NPUEvent,
50	- std::function<void(c10_npu::NPUEvent *)>>;	55	+ std::function<void(c10_npu::NPUEvent*)>>;
51	EventPool() : pools_(c10_npu::device_count()) {}	56	EventPool() : pools_(c10_npu::device_count()) {}
52		57
53	Event get(at::DeviceIndex device)	58	Event get(at::DeviceIndex device)
54	{	59	{
55	TORCH_INTERNAL_ASSERT(0 <= device, PTA_ERROR(ErrCode::PARAM));	60	TORCH_INTERNAL_ASSERT(0 <= device, PTA_ERROR(ErrCode::PARAM));
56	TORCH_INTERNAL_ASSERT(device < static_cast<at::DeviceIndex>(pools_.size()), PTA_ERROR(ErrCode::PARAM));	61	TORCH_INTERNAL_ASSERT(device < static_cast<at::DeviceIndex>(pools_.size()), PTA_ERROR(ErrCode::PARAM));
57	- auto &pool = pools_[device];	62	+ auto& pool = pools_[device];
58	- auto destructor = [&pool](c10_npu::NPUEvent *event) {	63	+ auto destructor = [&pool](c10_npu::NPUEvent* event) {
59	std::lock_guard<std::mutex> g(pool.mutex_);	64	std::lock_guard<std::mutex> g(pool.mutex_);
60	pool.event_pool_.push_back(std::unique_ptr<c10_npu::NPUEvent>(event));	65	pool.event_pool_.push_back(std::unique_ptr<c10_npu::NPUEvent>(event));
61	};	66	};
	@@ -64,7 +69,7 @@ public:
64	{	69	{
65	std::lock_guard<std::mutex> g(pool.mutex_);	70	std::lock_guard<std::mutex> g(pool.mutex_);
66	if (!pool.event_pool_.empty()) {	71	if (!pool.event_pool_.empty()) {
67	- auto *event = pool.event_pool_.back().release();	72	+ auto* event = pool.event_pool_.back().release();
68	pool.event_pool_.pop_back();	73	pool.event_pool_.pop_back();
69	return Event(event, destructor);	74	return Event(event, destructor);
70	}	75	}
	@@ -78,7 +83,7 @@ public:
78		83
79	void empty_cache()	84	void empty_cache()
80	{	85	{
81	- for (auto &pool : pools_) {	86	+ for (auto& pool : pools_) {
82	std::lock_guard<std::mutex> g(pool.mutex_);	87	std::lock_guard<std::mutex> g(pool.mutex_);
83	pool.event_pool_.clear();	88	pool.event_pool_.clear();
84	}	89	}
	@@ -92,307 +97,122 @@ private:
92	std::vector<PerDevicePool> pools_;	97	std::vector<PerDevicePool> pools_;
93	};	98	};
94		99
95	-static bool BlockComparator(const BlockSize &a, const BlockSize &b)	100	+struct NPUCachingHostAllocatorImpl : public at::CachingHostAllocatorImpl<c10_npu::NPUStream, EventPool::Event> {
96	-{	101	+public:
97	- // sort by size, break ties with pointer	102	+ bool ptr_check(void* ptr)
98	- if (a.size != b.size) {
99	- return a.size < b.size;
100	- }
101	- return reinterpret_cast<uintptr_t>(a.ptr) < reinterpret_cast<uintptr_t>(b.ptr);
102	-}
103	-
104	-struct HostAllocator {
105	- using Comparison = bool (*)(const BlockSize &, const BlockSize &);
106	-
107	- HostAllocator() : available(BlockComparator) {}
108	-
109	- aclError malloc(void **ptr, size_t size)
110	{	103	{
111	- std::lock_guard<std::mutex> lock(mutex);	104	+ std::lock_guard<std::mutex> g(npu_ptrs_mutex_);
		105	+ return npu_ptrs_.find(ptr) != npu_ptrs_.end();
		106	+ }
112		107
113	- // process outstanding npu events which may have occurred	108	+private:
114	- aclError err = processEvents();	109	+ void allocate_host_memory(size_t size, void** ptr) override
115	- if (err != ACL_ERROR_NONE) {	110	+ {
116	- return err;	111	+ // alloc needs set device first when using dataloader with pin_memory=True
117	- }
118	-
119	- // search for the smallest block which can hold this allocation
120	- BlockSize search_key(size);
121	- auto it = available.lower_bound(search_key);
122	- if (it != available.end()) {
123	- Block &block = blocks.at(it->ptr);
124	- AT_ASSERT(!block.allocated && block.event_count == 0, PTA_ERROR(ErrCode::PARAM));
125	- block.allocated = true;
126	- *ptr = block.ptr;
127	- available.erase(it);
128	- return ACL_ERROR_NONE;
129	- }
130	-
131	- *ptr = nullptr;
132	- // for pin_memory in dataloader, it should be set device first when new a thread
133	if (c10_npu::GetLocalDevice() < 0) {	112	if (c10_npu::GetLocalDevice() < 0) {
134	c10_npu::SetCurrentDevice();	113	c10_npu::SetCurrentDevice();
135	}	114	}
136		115
137	- // Round up the allocation to the nearest power of two to improve reuse.	116	+ auto start = std::chrono::steady_clock::now();
138	- size_t roundSize = c10::llvm::PowerOf2Ceil(size);	117	+ aclError err = aclrtMallocHost(ptr, size);
139	- // allocate a new block if no cached allocation is found
140	- err = aclrtMallocHost(ptr, roundSize);
141	if (err != ACL_ERROR_NONE) {	118	if (err != ACL_ERROR_NONE) {
142	CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(err);	119	CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(err);
143	- return err;
144	}	120	}
		121	+ if (*ptr != nullptr) { // we add the segment pointer here when initialization, but it does not matter
		122	+ std::lock_guard<std::mutex> g(npu_ptrs_mutex_);
		123	+ npu_ptrs_.insert(*ptr);
		124	+ }
		125	+ auto end = std::chrono::steady_clock::now();
		126	+ auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
145		127
146	- blocks.insert({ptr, Block(roundSize, ptr, true)});	128	+ {
147	- return ACL_ERROR_NONE;	129	+ std::lock_guard<std::mutex> g(stats_.timing_mutex_);
		130	+ stats_.host_alloc_time.increase(duration.count());
		131	+ }
148	}	132	}
149		133
150	- aclError free(void *ptr)	134	+ void free_block(Block* block) override
151	{	135	{
152	- std::lock_guard<std::mutex> lock(mutex);	136	+ auto start = std::chrono::steady_clock::now();
153	- if (!ptr) {	137	+ void* ptr = block->ptr_;
154	- return ACL_ERROR_NONE;	138	+ aclError err = aclrtFreeHost(block->ptr_);
155	- }
156	-
157	- auto it = blocks.find(ptr);
158	- AT_ASSERT(it != blocks.end(), PTA_ERROR(ErrCode::VALUE));
159	-
160	- Block &block = it->second;
161	- AT_ASSERT(block.allocated, PTA_ERROR(ErrCode::VALUE));
162	-
163	- // free (on valid memory) shouldn't fail, so mark unallocated before
164	- // we process the streams.
165	- block.allocated = false;
166	-
167	- // insert npu events for each stream on which this block was used. This
168	- aclError err = insertEvents(block);
169	if (err != ACL_ERROR_NONE) {	139	if (err != ACL_ERROR_NONE) {
170	CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(err);	140	CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(err);
171	- return err;
172	}	141	}
173	-	142	+ if (ptr != nullptr) {
174	- if (block.event_count == 0) {	143	+ std::lock_guard<std::mutex> g(npu_ptrs_mutex_);
175	- // the block can be re-used if there are no outstanding npu events	144	+ npu_ptrs_.erase(block->ptr_);
176	- available.insert(block);
177	}	145	}
178	- return ACL_ERROR_NONE;	146	+ auto end = std::chrono::steady_clock::now();
179	- }	147	+ auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
180	-	148	+ {
181	- aclError recordEvent(void *ptr, aclrtMemcpyKind kind, c10_npu::NPUStream stream)	149	+ std::lock_guard<std::mutex> g(stats_.timing_mutex_);
182	- {	150	+ stats_.host_free_time.increase(duration.count());
183	- std::lock_guard<std::mutex> lock(mutex);
184	-
185	- auto it = blocks.find(ptr);
186	- if (it == blocks.end()) {
187	- if (c10_npu::acl::AclrtMemcpyAsyncWithConditionExist() && kind == aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST) {
188	- return ACL_ERROR_NONE;
189	- }
190	- // Sync when host memory is allocated by malloc
191	- aclError error = c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream);
192	- if (error != ACL_ERROR_NONE) {
193	- CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(error);
194	- C10_NPU_SHOW_ERR_MSG();
195	- AT_ERROR("ACL stream synchronize failed.");
196	- return error;
197	- }
198	- return ACL_ERROR_NONE;
199	- }
200	-
201	- Block &block = it->second;
202	- AT_ASSERT(block.allocated, PTA_ERROR(ErrCode::VALUE));
203	-
204	- block.streams.insert(stream);
205	- return ACL_ERROR_NONE;
206	- }
207	-
208	- bool isPinndPtr(void *ptr)
209	- {
210	- std::lock_guard<std::mutex> lock(mutex);
211	- return blocks.find(ptr) != blocks.end();
212	- }
213	-
214	- aclError processEvents()
215	- {
216	- // Process outstanding npuEvents. Events that are completed are removed
217	- // from the queue, and the 'event_count' for the corresponding allocation
218	- // is decremented. Stops at the first event which has not been completed.
219	- // Since events on different devices or streams may occur out of order,
220	- // the processing of some events may be delayed.
221	- while (!npu_events.empty()) {
222	- auto &e = npu_events.front();
223	- EventPool::Event event = std::move(e.first);
224	- if (!event->query()) {
225	- e.first = std::move(event);
226	- break;
227	- }
228	-
229	- Block &block = blocks.at(e.second);
230	- block.event_count--;
231	- if (block.event_count == 0 && !block.allocated) {
232	- available.insert(block);
233	- }
234	- npu_events.pop_front();
235	- }
236	- return ACL_ERROR_NONE;
237	- }
238	-
239	- void emptyCache()
240	- {
241	- std::lock_guard<std::mutex> lock(mutex);
242	-
243	- // process outstanding npu events which may have occurred
244	- processEvents();
245	-
246	- // Release cached events from the event pool.
247	- event_pool_.empty_cache();
248	-
249	- // clear list of available blocks
250	- available.clear();
251	-
252	- // free and erase non-allocated blocks
253	- for (auto it = blocks.begin(); it != blocks.end();) {
254	- Block &block = it->second;
255	- if (aclrtFreeHost(block.ptr) != ACL_ERROR_NONE) {
256	- ASCEND_LOGE("free host pin failed!");
257	- }
258	- if (!block.allocated) {
259	- it = blocks.erase(it);
260	- } else {
261	- block.streams.clear();
262	- ++it;
263	- }
264	}	151	}
265	}	152	}
266		153
267	- aclError insertEvents(Block &block)	154	+ void record_stream(std::optional<std::vector<EventPool::Event>>& events, c10_npu::NPUStream stream) override
268	{	155	{
269	- aclError err = ACL_ERROR_NONE;	156	+ auto event = create_event_internal(stream.device_index());
		157	+ event->record(stream);
		158	+ events->push_back(std::move(event));
		159	+ }
270		160
271	- int prev_device = 0;	161	+ bool query_event(EventPool::Event& event) override
272	- err = c10_npu::GetDevice(&prev_device);	162	+ {
273	- if (err != ACL_ERROR_NONE) {	163	+ return event->query();
274	- return err;	164	+ }
275	- }
276		165
277	- std::unordered_set<c10_npu::NPUStream> streams(std::move(block.streams));	166	+ EventPool::Event create_event_internal(at::DeviceIndex idx)
278	- for (auto it = streams.begin(); it != streams.end(); ++it) {	167	+ {
279	- err = c10_npu::SetDevice(it->device_index());	168	+ static auto* event_pool = new EventPool();
280	- if (err != ACL_ERROR_NONE) {	169	+ return event_pool->get(idx);
281	- C10_NPU_SHOW_ERR_MSG();
282	- break;
283	- }
284	-
285	- EventPool::Event event = event_pool_.get(it->device_index());
286	- event->record(*it);
287	- ASCEND_LOGI("Event: record HostAllocator is successfully executed, event=%p", event.get());
288	-
289	- block.event_count++;
290	- npu_events.emplace_back(std::move(event), block.ptr);
291	- }
292	-
293	- c10_npu::SetDevice(prev_device);
294	-
295	- return err;
296	}	170	}
297		171
298	private:	172	private:
299	- EventPool event_pool_;	173	+ std::mutex npu_ptrs_mutex_;
300	-	174	+ ska::flat_hash_set<void*> npu_ptrs_;
301	- // lock around all operations
302	- std::mutex mutex;
303	-
304	- // blocks by pointer
305	- std::unordered_map<void *, Block> blocks;
306	-
307	- // pointers that are ready to be allocated (event_count=0)
308	- std::set<BlockSize, Comparison> available;
309	-
310	- // outstanding ACL events
311	- std::deque<std::pair<EventPool::Event, void *>> npu_events;
312	};	175	};
313	-} // namespace
314		176
315	-static HostAllocator& getHostAllocator()	177	+// Note : we do not use the macro DECLARE_HOST_ALLOCATOR here, because we need to access caching_host_allocator with ptr_exist function
316	-{	178	+void raw_local_deleter(void* ptr);
317	- // Construct allocator inside a function to prevent initialization when import
318	- static HostAllocator allocator;
319	- return allocator;
320	-}
321		179
322	-aclError CachingHostAllocator_recordEvent(	180	+struct NPUCachingHostAllocator final : public at::CachingHostAllocatorInterface<NPUCachingHostAllocatorImpl, raw_local_deleter> {};
323	- void *ptr,
324	- aclrtMemcpyKind kind,
325	- c10_npu::NPUStream stream)
326	-{
327	- return getHostAllocator().recordEvent(ptr, kind, stream);
328	-}
329		181
330	-bool CachingHostAllocator_isPinned(void *ptr)	182	+static NPUCachingHostAllocator caching_host_allocator;
331	-{
332	- return getHostAllocator().isPinndPtr(ptr);
333	-}
334		183
335	-void CachingHostAllocator_emptyCache()	184	+void raw_local_deleter(void* ptr)
336	-{
337	- getHostAllocator().emptyCache();
338	-}
339	-
340	-static void CachingHostDeleter(void *ptr)
341	{	185	{
342	#ifndef BUILD_LIBTORCH	186	#ifndef BUILD_LIBTORCH
343	// check the current thread have hold GIL Lock.	187	// check the current thread have hold GIL Lock.
344	if (PyGILState_Check()) {	188	if (PyGILState_Check()) {
345	// the current thread should not hold GIL.	189	// the current thread should not hold GIL.
346	Py_BEGIN_ALLOW_THREADS	190	Py_BEGIN_ALLOW_THREADS
347	- getHostAllocator().free(ptr);	191	+ caching_host_allocator.free(ptr);
348	Py_END_ALLOW_THREADS	192	Py_END_ALLOW_THREADS
349	} else {	193	} else {
350	- getHostAllocator().free(ptr);	194	+ caching_host_allocator.free(ptr);
351	}	195	}
352	#else	196	#else
353	- getHostAllocator().free(ptr);	197	+ caching_host_allocator.free(ptr);
354	#endif	198	#endif
355	}	199	}
		200	+// END of DECLARE_HOST_ALLOCATOR
356		201
357	-struct CachingHostAllocator final : public at::Allocator {	202	+REGISTER_HOST_ALLOCATOR(
358	- at::DataPtr allocate(size_t size) override	203	+ at::kPrivateUse1,
359	- {	204	+ &caching_host_allocator
360	- AT_ASSERT(size >= 0, PTA_ERROR(ErrCode::VALUE));	205	+)
361	- void *ptr = nullptr;
362	- if (size > 0) {
363	- if (getHostAllocator().malloc(&ptr, size) != ACL_ERROR_NONE) {
364	- ASCEND_LOGE("allocate host pinned memory fail");
365	- }
366	- }
367	- return {ptr, ptr, &CachingHostDeleter, at::DeviceType::CPU};
368	- }
369	- at::DeleterFnPtr raw_deleter() const override
370	- {
371	- return &CachingHostDeleter;
372	- }
373	- // Note [COW/lazy_clone is not supported yet]
374	- void copy_data(void* dest, const void* src, std::size_t count) const
375	- {
376	- TORCH_CHECK_NOT_IMPLEMENTED(false, "Not implemented for THNPUCachingHostAllocator", PTA_ERROR(ErrCode::NOT_SUPPORT));
377	- }
378	-};
379		206
380	-static CachingHostAllocator caching_host_allocator;	207	+bool ptr_exist(void* ptr)
381	-at::Allocator *getCachingHostAllocator()
382	{	208	{
383	- return &caching_host_allocator;	209	+ return caching_host_allocator.impl_->ptr_check(ptr);
384	}	210	}
385		211
386	-c10::Allocator *getPinnedMemoryAllocator()	212	+c10::Allocator* getPinnedMemoryAllocator()
387	{	213	{
388	- C10_LOG_API_USAGE_ONCE("aten.init.npu");	214	+ innerInitNPU();
389	- c10_npu::NpuSysCtrl::SysStatus status =	215	+ return at::getHostAllocator(at::kPrivateUse1);
390	- c10_npu::NpuSysCtrl::GetInstance().Initialize();
391	- if (status != c10_npu::NpuSysCtrl::SysStatus::INIT_SUCC) {
392	- ASCEND_LOGE("Npu init fail.");
393	- }
394	- return getCachingHostAllocator();
395	}	216	}
396		217
397	-} // namespace native	218	+} // namespace at_npu::native
398	-} // namespace at_npu

Mtorch_npu/csrc/core/npu/CachingHostAllocator.h+35-8

	@@ -1,3 +1,5 @@
		1	+#pragma once
		2	+
1	#include <c10/core/Allocator.h>	3	#include <c10/core/Allocator.h>
2	#include <c10/util/SmallVector.h>	4	#include <c10/util/SmallVector.h>
3		5
	@@ -5,19 +7,44 @@
5	#include "torch_npu/csrc/core/npu/NPUStream.h"	7	#include "torch_npu/csrc/core/npu/NPUStream.h"
6	#include "torch_npu/csrc/core/npu/NPUException.h"	8	#include "torch_npu/csrc/core/npu/NPUException.h"
7	#include <third_party/acl/inc/acl/acl.h>	9	#include <third_party/acl/inc/acl/acl.h>
		10	+#include <third_party/acl/inc/acl/acl_rt.h>
8		11
9	-namespace at_npu {	12	+#include <c10/core/DeviceGuard.h>
10	-namespace native {	13	+#include <ATen/DeviceGuard.h>
		14	+#include <ATen/Utils.h>
11		15
12	-TORCH_NPU_API c10::Allocator* getCachingHostAllocator();	16	+#include <ATen/core/CachingHostAllocator.h>
		17	+#include <c10/util/Deprecated.h>
13		18
14	-TORCH_NPU_API aclError CachingHostAllocator_recordEvent(void* ptr, aclrtMemcpyKind kind, c10_npu::NPUStream stream);	19	+#include "torch_npu/csrc/core/npu/NPUFunctions.h"
		20	+#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
		21	+
		22	+namespace at_npu::native {
		23	+
		24	+bool ptr_exist(void* ptr);
		25	+
		26	+inline TORCH_NPU_API c10::Allocator* getCachingHostAllocator() {
		27	+ return at::getHostAllocator(at::kPrivateUse1);
		28	+}
		29	+
		30	+inline TORCH_NPU_API bool CachingHostAllocator_recordEvent(void* ptr, void* ctx, c10_npu::NPUStream stream) {
		31	+ return at::getHostAllocator(at::kPrivateUse1)->record_event(ptr, ctx, stream.unwrap());
		32	+}
15		33
16	-TORCH_NPU_API bool CachingHostAllocator_isPinned(void* ptr);
17	// Releases cached pinned memory allocations via npuHostFree	34	// Releases cached pinned memory allocations via npuHostFree
18	-TORCH_NPU_API void CachingHostAllocator_emptyCache();	35	+inline TORCH_NPU_API void CachingHostAllocator_emptyCache() {
		36	+ return at::getHostAllocator(at::kPrivateUse1)->empty_cache();
		37	+}
		38	+
		39	+inline TORCH_NPU_API bool CachingHostAllocator_isPinned(void* ptr) {
		40	+ return at_npu::native::ptr_exist(ptr);
		41	+}
		42	+
		43	+inline at::DataPtr HostAlloc(size_t size)
		44	+{
		45	+ return at::getHostAllocator(at::kPrivateUse1)->allocate(size);
		46	+}
19		47
20	c10::Allocator* getPinnedMemoryAllocator();	48	c10::Allocator* getPinnedMemoryAllocator();
21		49
22	-} // namespace native	50	+} // namespace at_npu::native
23	-} // namespace at_npu