round_size pad 32 byte-pytorch-AtomGit

已合并

round_size pad 32 byte #28328

round_size pad 32 byte #28328

已合并

zhaoyu65创建于 2025年12月19日

共 3 个文件变更+3-30

Mtest/npu/test_allocator_envs.py+2-4

	@@ -15,8 +15,7 @@ class TestAllocator(TestCase):
15	a = torch.rand(1024 * 1024 * 40, dtype=torch.float32).npu()	15	a = torch.rand(1024 * 1024 * 40, dtype=torch.float32).npu()
16	torch.npu.synchronize()	16	torch.npu.synchronize()
17	# 实际申请1G内存	17	# 实际申请1G内存
18	- nbytes = 1024 * 1024 * 40 * 4	18	+ self.assertEqual(torch_npu.npu.memory_allocated(), prev + ((4 * 40 * 1024 * 1024 + 32) // 512 + 1) * 512)
19	- self.assertEqual(torch_npu.npu.memory_allocated(), prev + ((nbytes + 511) // 512) * 512)
20		19
21	@SupportedDevices(['Ascend910B'])	20	@SupportedDevices(['Ascend910B'])
22	def test_huge_memory_alloc_512B(self):	21	def test_huge_memory_alloc_512B(self):
	@@ -24,8 +23,7 @@ class TestAllocator(TestCase):
24	a = torch.rand(8 * 8 * 16, dtype=torch.float32).npu() # 512B	23	a = torch.rand(8 * 8 * 16, dtype=torch.float32).npu() # 512B
25	torch.npu.synchronize()	24	torch.npu.synchronize()
26	# 实际申请1M内存	25	# 实际申请1M内存
27	- nbytes = 8 * 8 * 16 * 4	26	+ self.assertEqual(torch_npu.npu.memory_allocated(), prev + ((8 * 8 * 16 * 4 + 32) // 512 + 1) * 512)
28	- self.assertEqual(torch_npu.npu.memory_allocated(), prev + ((nbytes + 511) // 512) * 512)
29		27
30	if __name__ == '__main__':	28	if __name__ == '__main__':
31	run_tests()	29	run_tests()

Mtest/npu/test_roundup_power2_divisions.py+0-10

	@@ -73,16 +73,6 @@ class Test_roundup_power2_divisions(TestCase):
73	self.assertEqual(allocated_bytes, active_bytes)	73	self.assertEqual(allocated_bytes, active_bytes)
74	self.assertEqual(allocated_bytes, 1024)	74	self.assertEqual(allocated_bytes, 1024)
75		75
76	- @SupportedDevices(['Ascend910B', 'Ascend910C'])
77	- def test_without_add_32(self):
78	- torch.npu.memory.empty_cache()
79	-
80	- x = torch.from_numpy(np.random.randn(512).astype(np.uint8)).to("npu")
81	- allocated_bytes = torch.npu.memory_stats()["allocated_bytes.all.current"]
82	- active_bytes = torch.npu.memory_stats()["active_bytes.all.current"]
83	- self.assertEqual(allocated_bytes, active_bytes)
84	- self.assertEqual(allocated_bytes, 512)
85	-
86		76
87	if __name__ == '__main__':	77	if __name__ == '__main__':
88	run_tests()	78	run_tests()

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+1-16

	@@ -1988,19 +1988,6 @@ public:
1988	return result;	1988	return result;
1989	}	1989	}
1990		1990
1991	- static bool should_apply_legacy_round_padding()
1992	- {
1993	- static bool should_pad = []() -> bool {
1994	- const auto &soc_version = c10_npu::GetSocVersion();
1995	- const bool disable_padding =
1996	- (soc_version >= c10_npu::SocVersion::Ascend910B1 &&
1997	- soc_version <= c10_npu::SocVersion::Ascend910B4_1) \|\|
1998	- (soc_version >= c10_npu::SocVersion::Ascend910_9391);
1999	- return !disable_padding;
2000	- }();
2001	- return should_pad;
2002	- }
2003	-
2004	// This function takes the size and number of divisions argument and rounds	1991	// This function takes the size and number of divisions argument and rounds
2005	// up the size argument for the nearest power-of-2 division.	1992	// up the size argument for the nearest power-of-2 division.
2006	// For example, if we need to round-up 1200 and number of divisions is 4,	1993	// For example, if we need to round-up 1200 and number of divisions is 4,
	@@ -2031,9 +2018,7 @@ public:
2031	static size_t round_size(size_t size)	2018	static size_t round_size(size_t size)
2032	{	2019	{
2033	constexpr size_t kPadSize = 32;	2020	constexpr size_t kPadSize = 32;
2034	- if (should_apply_legacy_round_padding()) {	2021	+ size += kPadSize;
2035	- size += kPadSize;
2036	- }
2037		2022
2038	if (size < kMinBlockSize) {	2023	if (size < kMinBlockSize) {
2039	return kMinBlockSize;	2024	return kMinBlockSize;