[feat]support shmemput/get-pytorch-AtomGit

Mtest/distributed/shmem/test_shmem.py+43-0

	@@ -135,6 +135,49 @@ class NPUSHMEMSymmetricMemoryTest(MultiProcContinuousTest):
135	shmem_matmul = torch.matmul(shmem_tensor, shmem_tensor1)	135	shmem_matmul = torch.matmul(shmem_tensor, shmem_tensor1)
136	self.assertEqual(shmem_matmul, matmul)	136	self.assertEqual(shmem_matmul, matmul)
137		137
		138	+ @skipIfUnsupportMultiNPU(2)
		139	+ def test_shmem_put(self) -> None:
		140	+ self._init_device()
		141	+
		142	+ group_name = dist.group.WORLD.group_name
		143	+ symm_mem.enable_symm_mem_for_group(group_name)
		144	+
		145	+ dtype = torch.float
		146	+ numel = 1024
		147	+
		148	+ tensor = symm_mem.empty(numel, dtype=dtype, device=self.device).fill_(self.rank)
		149	+ symm_mem.rendezvous(tensor, group=group_name)
		150	+
		151	+ if self.rank == 0:
		152	+ torch.ops.symm_mem.nvshmem_put(tensor, 1)
		153	+ dist.barrier(device_ids=[self.rank])
		154	+ elif self.rank == 1:
		155	+ dist.barrier(device_ids=[self.rank])
		156	+ torch.testing.assert_close(
		157	+ tensor, torch.zeros(numel, dtype=dtype, device=self.device)
		158	+ )
		159	+
		160	+ @skipIfUnsupportMultiNPU(2)
		161	+ def test_shmem_get(self) -> None:
		162	+ self._init_device()
		163	+
		164	+ group_name = dist.group.WORLD.group_name
		165	+ symm_mem.enable_symm_mem_for_group(group_name)
		166	+
		167	+ dtype = torch.float
		168	+ numel = 1024
		169	+
		170	+ tensor = symm_mem.empty(numel, dtype=dtype, device=self.device).fill_(self.rank)
		171	+ symm_mem.rendezvous(tensor, group=group_name)
		172	+
		173	+ if self.rank == 0:
		174	+ torch.ops.symm_mem.nvshmem_get(tensor, 1)
		175	+ dist.barrier(device_ids=[self.rank])
		176	+ torch.testing.assert_close(
		177	+ tensor, torch.ones(numel, dtype=dtype, device=self.device)
		178	+ )
		179	+ elif self.rank == 1:
		180	+ dist.barrier(device_ids=[self.rank])
138		181
139	if __name__ == "__main__":	182	if __name__ == "__main__":
140	run_tests()	183	run_tests()

Mthird_party/shmem/include/shmem_host_def.h+9-22

	@@ -74,7 +74,9 @@ constexpr int DEFAULT_TIMEOUT = 120;
74		74
75	typedef struct {	75	typedef struct {
76	int32_t version;	76	int32_t version;
77	- char internal[SHMEM_UNIQUE_ID_INNER_LEN];	77	+ int my_pe;
		78	+ int n_pes;
		79	+ char internal[ACLSHMEM_UNIQUE_ID_INNER_LEN];
78	} shmem_uniqueid_t;	80	} shmem_uniqueid_t;
79		81
80	constexpr int32_t SHMEM_UNIQUEID_VERSION = (1 << 16) + sizeof(shmem_uniqueid_t);	82	constexpr int32_t SHMEM_UNIQUEID_VERSION = (1 << 16) + sizeof(shmem_uniqueid_t);
	@@ -144,26 +146,9 @@ typedef struct {
144	uint32_t shm_init_timeout;	146	uint32_t shm_init_timeout;
145	uint32_t shm_create_timeout;	147	uint32_t shm_create_timeout;
146	uint32_t control_operation_timeout;	148	uint32_t control_operation_timeout;
		149	+ int32_t sockFd;
147	} shmem_init_optional_attr_t;	150	} shmem_init_optional_attr_t;
148		151
149	-/**
150	- * @struct shmem_init_attr_t
151	- * @brief Mandatory parameter for attributes used for initialization.
152	- *
153	- * - int my_rank: The rank of the current process.
154	- * - int n_ranks: The total rank number of all processes.
155	- * - const char* ip_port: The ip and port of the communication server. The port must not conflict with other modules and processes.
156	- * - uint64_t local_mem_size: The size of shared memory currently occupied by current rank.
157	- * - shmem_init_optional_attr_t option_attr: Optional Parameters.
158	-*/
159	-typedef struct {
160	- int my_rank;
161	- int n_ranks;
162	- const char* ip_port;
163	- uint64_t local_mem_size;
164	- shmem_init_optional_attr_t option_attr;
165	-} shmem_init_attr_t;
166	-
167	/**	152	/**
168	* @struct aclshmem_init_optional_attr_t	153	* @struct aclshmem_init_optional_attr_t
169	* @brief Optional parameter for the attributes used for initialization.	154	* @brief Optional parameter for the attributes used for initialization.
	@@ -194,14 +179,16 @@ typedef struct {
194	* - uint64_t local_mem_size: The size of shared memory currently occupied by current pe.	179	* - uint64_t local_mem_size: The size of shared memory currently occupied by current pe.
195	* - aclshmem_init_optional_attr_t option_attr: Optional Parameters.	180	* - aclshmem_init_optional_attr_t option_attr: Optional Parameters.
196	*/	181	*/
197	-typedef struct {	182	+typedef struct aclshmemx_init_attr_t {
198	int my_pe;	183	int my_pe;
199	int n_pes;	184	int n_pes;
200	- char ip_port[ACLSHMEM_MAX_IP_PORT_LEN];	185	+ char ip_port[ACLSHMEM_MAX_IP_PORT_LEN] = {};
201	uint64_t local_mem_size;	186	uint64_t local_mem_size;
202	aclshmem_init_optional_attr_t option_attr = {(1 << 16) + sizeof(aclshmem_init_optional_attr_t), ACLSHMEM_DATA_OP_MTE, DEFAULT_TIMEOUT, DEFAULT_TIMEOUT, DEFAULT_TIMEOUT};	187	aclshmem_init_optional_attr_t option_attr = {(1 << 16) + sizeof(aclshmem_init_optional_attr_t), ACLSHMEM_DATA_OP_MTE, DEFAULT_TIMEOUT, DEFAULT_TIMEOUT, DEFAULT_TIMEOUT};
203	- void *comm_args;	188	+ void *comm_args = nullptr;
		189	+ uint64_t instance_id = 0;
204	} aclshmemx_init_attr_t;	190	} aclshmemx_init_attr_t;
		191	+#define shmem_init_attr_t aclshmemx_init_attr_t
205		192
206	/**	193	/**
207	* @brief Callback function of private key password decryptor, see shmem_register_decrypt_handler	194	* @brief Callback function of private key password decryptor, see shmem_register_decrypt_handler

Mtorch_npu/csrc/distributed/symm_mem/NPUSHMEMExtension.cpp+21-4

	@@ -1,6 +1,7 @@
1	#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>	1	#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
2	#include "torch_npu/csrc/core/npu/NPUException.h"	2	#include "torch_npu/csrc/core/npu/NPUException.h"
3	#include "torch_npu/csrc/core/npu/NPUFunctions.h"	3	#include "torch_npu/csrc/core/npu/NPUFunctions.h"
		4	+#include "torch_npu/csrc/core/npu/NPUStream.h"
4	#include "torch_npu/csrc/core/npu/register/OptionsManager.h"	5	#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
5	#include "torch_npu/csrc/logging/LogContext.h"	6	#include "torch_npu/csrc/logging/LogContext.h"
6	#include "torch_npu/csrc/distributed/symm_mem/NPUSymmetricMemoryUtils.hpp"	7	#include "torch_npu/csrc/distributed/symm_mem/NPUSymmetricMemoryUtils.hpp"
	@@ -89,15 +90,31 @@ void nvshmem_put(at::Tensor& tensor, int64_t peer)
89	auto rank = hdl->get_rank();	90	auto rank = hdl->get_rank();
90	void* buffer_ptr = hdl->get_buffer_ptrs()[rank];	91	void* buffer_ptr = hdl->get_buffer_ptrs()[rank];
91	auto buffer_size = tensor.numel() * tensor.element_size();	92	auto buffer_size = tensor.numel() * tensor.element_size();
92	-	93	+ TORCH_CHECK(peer < hdl->get_world_size(), "peer must be smaller than world size", DIST_ERROR(ErrCode::PARAM));
93	at::DeviceGuard device_guard(tensor.device());	94	at::DeviceGuard device_guard(tensor.device());
94	- // to be done for putmem	95	+ auto stream = c10_npu::getCurrentNPUStream();
95	- throw std::runtime_error("NPUSHMEMSymmetricMemory does not support nvshmem_put" + DIST_ERROR(ErrCode::NOT_SUPPORT));	96	+ c10d::symmetric_memory::Shmem_putmem_on_stream(buffer_ptr, tensor.data_ptr(), buffer_size, peer, stream);
		97	+}
		98	+
		99	+void nvshmem_get(at::Tensor& tensor, int64_t peer)
		100	+{
		101	+ // to be done: support non-contiguous tensors
		102	+ TORCH_CHECK(tensor.is_contiguous(),
		103	+ "put op currently supports contiguous tensors only", DIST_ERROR(ErrCode::PARAM));
		104	+ // to be done: rendezvous should remember the group name
		105	+ auto hdl = c10d::symmetric_memory::rendezvous(tensor, "0");
		106	+ auto rank = hdl->get_rank();
		107	+ void* buffer_ptr = hdl->get_buffer_ptrs()[rank];
		108	+ auto buffer_size = tensor.numel() * tensor.element_size();
		109	+ TORCH_CHECK(peer < hdl->get_world_size(), "peer must be smaller than world size", DIST_ERROR(ErrCode::PARAM));
		110	+ at::DeviceGuard device_guard(tensor.device());
		111	+ auto stream = c10_npu::getCurrentNPUStream();
		112	+ c10d::symmetric_memory::Shmem_getmem_on_stream(tensor.mutable_data_ptr(), buffer_ptr, buffer_size, peer, stream);
96	}	113	}
97		114
98	} // namespace c10d::npushmem_extension	115	} // namespace c10d::npushmem_extension
99		116
100	-
101	TORCH_LIBRARY_IMPL(symm_mem, PrivateUse1, m) {	117	TORCH_LIBRARY_IMPL(symm_mem, PrivateUse1, m) {
102	m.impl("nvshmem_put", c10d::npushmem_extension::nvshmem_put);	118	m.impl("nvshmem_put", c10d::npushmem_extension::nvshmem_put);
		119	+ m.impl("nvshmem_get", c10d::npushmem_extension::nvshmem_get);
103	}	120	}

Mtorch_npu/csrc/distributed/symm_mem/NPUSHMEMExtension.h+2-0

	@@ -13,4 +13,6 @@ void initialize_npushmem_with_store(
13		13
14	TORCH_API void nvshmem_put(at::Tensor& tensor, int64_t peer);	14	TORCH_API void nvshmem_put(at::Tensor& tensor, int64_t peer);
15		15
		16	+TORCH_API void nvshmem_get(at::Tensor& tensor, int64_t peer);
		17	+
16	} // namespace c10d::npushmem_extension	18	} // namespace c10d::npushmem_extension

Mtorch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.cpp+24-0

	@@ -22,6 +22,8 @@ TORCH_NPU_LOAD_FUNC(aclshmem_malloc)
22	TORCH_NPU_LOAD_FUNC(aclshmem_free)	22	TORCH_NPU_LOAD_FUNC(aclshmem_free)
23	TORCH_NPU_LOAD_FUNC(aclshmem_ptr)	23	TORCH_NPU_LOAD_FUNC(aclshmem_ptr)
24	TORCH_NPU_LOAD_FUNC(aclshmem_finalize)	24	TORCH_NPU_LOAD_FUNC(aclshmem_finalize)
		25	+TORCH_NPU_LOAD_FUNC(aclshmemx_putmem_on_stream)
		26	+TORCH_NPU_LOAD_FUNC(aclshmemx_getmem_on_stream)
25	TORCH_NPU_LOAD_FUNC(shmem_set_conf_store_tls)	27	TORCH_NPU_LOAD_FUNC(shmem_set_conf_store_tls)
26	TORCH_NPU_LOAD_FUNC(shmem_set_attr)	28	TORCH_NPU_LOAD_FUNC(shmem_set_attr)
27	TORCH_NPU_LOAD_FUNC(shmem_get_uniqueid)	29	TORCH_NPU_LOAD_FUNC(shmem_get_uniqueid)
	@@ -163,6 +165,28 @@ void Aclshmem_free(void *ptr)
163	return shmem_free_func(ptr);	165	return shmem_free_func(ptr);
164	}	166	}
165		167
		168	+void Shmem_putmem_on_stream(void dst, void src, size_t elem_size, int32_t pe, aclrtStream stream)
		169	+{
		170	+ typedef void (ShmemApiFunc)(void , void *, size_t, int32_t, aclrtStream);
		171	+ static ShmemApiFunc func = nullptr;
		172	+ if (func == nullptr) {
		173	+ func = (ShmemApiFunc)TORCH_NPU_GET_FUNC(aclshmemx_putmem_on_stream);
		174	+ }
		175	+ TORCH_CHECK(func, "Failed to find function ", "aclshmemx_putmem_on_stream", PTA_ERROR(ErrCode::NOT_FOUND));
		176	+ func(dst, src, elem_size, pe, stream);
		177	+}
		178	+
		179	+void Shmem_getmem_on_stream(void dst, void src, size_t elem_size, int32_t pe, aclrtStream stream)
		180	+{
		181	+ typedef void (ShmemApiFunc)(void , void *, size_t, int32_t, aclrtStream);
		182	+ static ShmemApiFunc func = nullptr;
		183	+ if (func == nullptr) {
		184	+ func = (ShmemApiFunc)TORCH_NPU_GET_FUNC(aclshmemx_getmem_on_stream);
		185	+ }
		186	+ TORCH_CHECK(func, "Failed to find function ", "aclshmemx_getmem_on_stream", PTA_ERROR(ErrCode::NOT_FOUND));
		187	+ func(dst, src, elem_size, pe, stream);
		188	+}
		189	+
166	void Aclshmem_ptr(void ptr, int pe)	190	void Aclshmem_ptr(void ptr, int pe)
167	{	191	{
168	typedef void* (ShmemApiFunc)(void , int);	192	typedef void* (ShmemApiFunc)(void , int);

Mtorch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h+5-0

	@@ -3,6 +3,7 @@
3	#include <cstddef>	3	#include <cstddef>
4	#include <cstdint>	4	#include <cstdint>
5	#include "third_party/shmem/include/shmem_host_def.h"	5	#include "third_party/shmem/include/shmem_host_def.h"
		6	+#include "third_party/acl/inc/acl/acl_base.h"
6		7
7	namespace c10d {	8	namespace c10d {
8	namespace symmetric_memory {	9	namespace symmetric_memory {
	@@ -30,6 +31,10 @@ void Aclshmem_free(void *ptr);
30		31
31	void Aclshmem_ptr(void ptr, int pe);	32	void Aclshmem_ptr(void ptr, int pe);
32		33
		34	+void Shmem_putmem_on_stream(void dst, void src, size_t elem_size, int32_t pe, aclrtStream stream);
		35	+
		36	+void Shmem_getmem_on_stream(void dst, void src, size_t elem_size, int32_t pe, aclrtStream stream);
		37	+
33	bool Aclshmem_finalize_exist();	38	bool Aclshmem_finalize_exist();
34		39
35	int Aclshmem_finalize(void);	40	int Aclshmem_finalize(void);

Mtorch_npu/csrc/distributed/symm_mem/NPUSHMEMSymmetricMemory.cpp+10-3

	@@ -196,8 +196,10 @@ void* NPUSHMEMSymmetricMemoryAllocator::alloc(
196	TORCH_CHECK(ptr != nullptr, "shmem_malloc return nullptr with size ", size, DIST_ERROR(ErrCode::MEMORY));	196	TORCH_CHECK(ptr != nullptr, "shmem_malloc return nullptr with size ", size, DIST_ERROR(ErrCode::MEMORY));
197	auto allocation =	197	auto allocation =
198	std::make_shared<NPUSHMEMAllocation>(ptr, size, device_idx);	198	std::make_shared<NPUSHMEMAllocation>(ptr, size, device_idx);
199	- // to be done: thread safety	199	+ {
200	- allocations_.try_emplace(ptr, std::move(allocation));	200	+ std::lock_guard<std::mutex> lock(mutex_);
		201	+ allocations_.try_emplace(ptr, std::move(allocation));
		202	+ }
201	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator alloc end, size is %d, device is %d, group_name is %s, ptr is %p",	203	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator alloc end, size is %d, device is %d, group_name is %s, ptr is %p",
202	size, device_idx, group_name == std::nullopt ? "" : (*group_name).c_str(), ptr);	204	size, device_idx, group_name == std::nullopt ? "" : (*group_name).c_str(), ptr);
203	return ptr;	205	return ptr;
	@@ -206,12 +208,16 @@ void* NPUSHMEMSymmetricMemoryAllocator::alloc(
206	void NPUSHMEMSymmetricMemoryAllocator::free(void* ptr)	208	void NPUSHMEMSymmetricMemoryAllocator::free(void* ptr)
207	{	209	{
208	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator free start, ptr is %p", ptr);	210	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator free start, ptr is %p", ptr);
209	- allocations_.erase(ptr);	211	+ {
		212	+ std::lock_guard<std::mutex> lock(mutex_);
		213	+ allocations_.erase(ptr);
		214	+ }
210	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator free end, ptr is %p", ptr);	215	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator free end, ptr is %p", ptr);
211	}	216	}
212		217
213	size_t NPUSHMEMSymmetricMemoryAllocator::get_alloc_size(void* ptr)	218	size_t NPUSHMEMSymmetricMemoryAllocator::get_alloc_size(void* ptr)
214	{	219	{
		220	+ std::lock_guard<std::mutex> lock(mutex_);
215	auto it = allocations_.find(ptr);	221	auto it = allocations_.find(ptr);
216	if (it == allocations_.end()) {	222	if (it == allocations_.end()) {
217	TORCH_CHECK(false, ptr, " is not allocated with NPUSHMEMSymmetricMemoryAllocator", DIST_ERROR(ErrCode::PARAM));	223	TORCH_CHECK(false, ptr, " is not allocated with NPUSHMEMSymmetricMemoryAllocator", DIST_ERROR(ErrCode::PARAM));
	@@ -225,6 +231,7 @@ c10::intrusive_ptr<SymmetricMemory> NPUSHMEMSymmetricMemoryAllocator::rendezvous
225	{	231	{
226	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator rendezvous start, ptr is %p, group_name is %s", ptr, (*group_name).c_str());	232	TORCH_NPU_SYMMEM_LOGD("NPUSHMEMSymmetricMemoryAllocator rendezvous start, ptr is %p, group_name is %s", ptr, (*group_name).c_str());
227	TORCH_CHECK(group_name.has_value(), "rendezvous, group_name is invalid.", DIST_ERROR(ErrCode::PARAM));	233	TORCH_CHECK(group_name.has_value(), "rendezvous, group_name is invalid.", DIST_ERROR(ErrCode::PARAM));
		234	+ std::lock_guard<std::mutex> lock(mutex_);
228	{	235	{
229	auto it = symm_mems_.find(std::make_tuple(ptr, *group_name));	236	auto it = symm_mems_.find(std::make_tuple(ptr, *group_name));
230	if (it != symm_mems_.end()) {	237	if (it != symm_mems_.end()) {

Mtorch_npu/csrc/distributed/symm_mem/NPUSHMEMSymmetricMemory.hpp+2-0

	@@ -3,6 +3,7 @@
3	#include <unordered_map>	3	#include <unordered_map>
4	#include <map>	4	#include <map>
5	#include <memory>	5	#include <memory>
		6	+#include <mutex>
6	#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>	7	#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
7	#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h"	8	#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h"
8	#include "torch_npu/csrc/logging/LogContext.h"	9	#include "torch_npu/csrc/logging/LogContext.h"
	@@ -96,6 +97,7 @@ public:
96	std::string name() override;	97	std::string name() override;
97		98
98	private:	99	private:
		100	+ std::mutex mutex_;
99	std::unordered_map<void*, std::shared_ptr<NPUSHMEMAllocation>> allocations_;	101	std::unordered_map<void*, std::shared_ptr<NPUSHMEMAllocation>> allocations_;
100	std::map<std::tuple<void*, std::string>, c10::intrusive_ptr<SymmetricMemory>>	102	std::map<std::tuple<void*, std::string>, c10::intrusive_ptr<SymmetricMemory>>
101	symm_mems_;	103	symm_mems_;