[SHMEM] support npu shmem-pytorch-AtomGit

Atest/distributed/shmem/test_shmem.py+153-0

@@ -0,0 +1,153 @@
	1	+import os
	2	+from unittest import skip
	3	+import torch
	4	+import torch.distributed as dist
	5	+import torch.distributed._symmetric_memory as symm_mem
	6	+from torch.testing._internal.common_distributed import MultiProcContinousTest
	7	+from torch.testing._internal.common_utils import instantiate_parametrized_tests
	8	+import torch_npu
	9	+from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
	10	+
	11	+
	12	+# So that tests are written in device-agnostic way
	13	+device_type = "npu"
	14	+device_module = torch.get_device_module(device_type)
	15	+
	16	+
	17	+@instantiate_parametrized_tests
	18	+@skip("request shmem")
	19	+class NPUSHMEMSymmetricMemoryTest(MultiProcContinousTest):
	20	+ @classmethod
	21	+ def backend_str(cls) -> str:
	22	+ # Testing with HCCL backend
	23	+ return "hccl"
	24	+
	25	+ @classmethod
	26	+ def setUpClass(cls):
	27	+ """
	28	+ Class-scope test fixture. Run once for entire test class, before any test starts.
	29	+ Set up the device.
	30	+ """
	31	+ super().setUpClass()
	32	+ dev_id = cls.rank % torch.npu.device_count()
	33	+ cls.device = torch.device(f"npu:{dev_id}")
	34	+
	35	+ def _init_device(self) -> None:
	36	+ device_module.set_device(self.device)
	37	+ torch.empty(1, device=self.device)
	38	+
	39	+ @property
	40	+ def device(self) -> torch.device:
	41	+ return torch.device(device_type, self.rank)
	42	+
	43	+ @skipIfUnsupportMultiNPU(2)
	44	+ def test_alloc(self) -> None:
	45	+ self._init_device()
	46	+
	47	+ group_name = dist.group.WORLD.group_name
	48	+ symm_mem.enable_symm_mem_for_group(group_name)
	49	+
	50	+ dtype = torch.float
	51	+ numel = 1024
	52	+
	53	+ def foo():
	54	+ inp = symm_mem.empty(numel, dtype=dtype, device=self.device)
	55	+ symm_mem.rendezvous(inp, group=group_name)
	56	+
	57	+ foo()
	58	+
	59	+ out = symm_mem.empty(numel, dtype=dtype, device=self.device)
	60	+ symm_mem.rendezvous(out, group=group_name)
	61	+
	62	+ @skipIfUnsupportMultiNPU(2)
	63	+ def test_alloc_free(self) -> None:
	64	+ self._init_device()
	65	+
	66	+ group_name = dist.group.WORLD.group_name
	67	+ symm_mem.enable_symm_mem_for_group(group_name)
	68	+
	69	+ dtype = torch.float
	70	+ numel = 1024
	71	+
	72	+ out = symm_mem.empty(numel, dtype=dtype, device=self.device)
	73	+ symm_mem.rendezvous(out, group=group_name)
	74	+ del out
	75	+
	76	+ @skipIfUnsupportMultiNPU(2)
	77	+ def test_shmem_copy(self) -> None:
	78	+ self._init_device()
	79	+
	80	+ group_name = dist.group.WORLD.group_name
	81	+ symm_mem.enable_symm_mem_for_group(group_name)
	82	+
	83	+ dtype = torch.float
	84	+ shape = (512, 512)
	85	+
	86	+ tensor = torch.randn(shape, dtype=dtype, device=self.device)
	87	+
	88	+ shmem_tensor = symm_mem.empty(shape, dtype=dtype, device=self.device)
	89	+ shmem_tensor.copy_(tensor)
	90	+ self.assertEqual(shmem_tensor, tensor)
	91	+
	92	+ @skipIfUnsupportMultiNPU(2)
	93	+ def test_shmem_matmul(self) -> None:
	94	+ self._init_device()
	95	+
	96	+ group_name = dist.group.WORLD.group_name
	97	+ symm_mem.enable_symm_mem_for_group(group_name)
	98	+
	99	+ dtype = torch.float
	100	+ shape = (512, 512)
	101	+
	102	+ tensor = torch.randn(shape, dtype=dtype, device=self.device)
	103	+ tensor1 = torch.randn(shape, dtype=dtype, device=self.device)
	104	+
	105	+ matmul = torch.matmul(tensor, tensor1)
	106	+
	107	+ shmem_tensor = symm_mem.empty(shape, dtype=dtype, device=self.device)
	108	+ shmem_tensor.copy_(tensor)
	109	+
	110	+ shmem_matmul = torch.matmul(shmem_tensor, tensor1)
	111	+ self.assertEqual(shmem_matmul, matmul)
	112	+
	113	+ @skipIfUnsupportMultiNPU(2)
	114	+ def test_shmem_matmul1(self) -> None:
	115	+ self._init_device()
	116	+
	117	+ group_name = dist.group.WORLD.group_name
	118	+ symm_mem.enable_symm_mem_for_group(group_name)
	119	+
	120	+ dtype = torch.float
	121	+ shape = (512, 512)
	122	+
	123	+ tensor = torch.randn(shape, dtype=dtype, device=self.device)
	124	+ tensor1 = torch.randn(shape, dtype=dtype, device=self.device)
	125	+
	126	+ matmul = torch.matmul(tensor, tensor1)
	127	+
	128	+ shmem_tensor = symm_mem.empty(shape, dtype=dtype, device=self.device)
	129	+ shmem_tensor.copy_(tensor)
	130	+ shmem_tensor1 = symm_mem.empty(shape, dtype=dtype, device=self.device)
	131	+ shmem_tensor1.copy_(tensor1)
	132	+
	133	+ shmem_matmul = torch.matmul(shmem_tensor, shmem_tensor1)
	134	+ self.assertEqual(shmem_matmul, matmul)
	135	+
	136	+
	137	+if __name__ == "__main__":
	138	+ os.environ['MASTER_ADDR'] = '127.0.0.1'
	139	+ os.environ['MASTER_PORT'] = '29500'
	140	+ rank = int(os.getenv("RANK", -1))
	141	+ world_size = int(os.getenv("WORLD_SIZE", 2))
	142	+
	143	+ if rank != -1:
	144	+ # Launched with torchrun or other multi-proc launchers. Directly run the test.
	145	+ NPUSHMEMSymmetricMemoryTest.run_rank(rank, world_size)
	146	+ else:
	147	+ # Launched as a single process. Spawn subprocess to run the tests.
	148	+ # Also need a rendezvous file for `init_process_group` purpose.
	149	+ torch.multiprocessing.spawn(
	150	+ NPUSHMEMSymmetricMemoryTest.run_rank,
	151	+ nprocs=world_size,
	152	+ args=(world_size,),
	153	+ )

Mtest/distributed/test_fault_mode.py+2-0

	@@ -1,5 +1,6 @@
1	import os.path	1	import os.path
2	import subprocess	2	import subprocess
		3	+from unittest import skip
3	import torch	4	import torch
4	import torch_npu	5	import torch_npu
5	from torch_npu.testing.testcase import TestCase, run_tests	6	from torch_npu.testing.testcase import TestCase, run_tests
	@@ -135,6 +136,7 @@ class TestMode(TestCase):
135	)	136	)
136		137
137	@skipIfUnsupportMultiNPU(2)	138	@skipIfUnsupportMultiNPU(2)
		139	+ @skip("Environmental problem, temporarily skip.")
138	def test_hccl_timeout(self):	140	def test_hccl_timeout(self):
139	path = os.path.join(os.path.dirname(__file__), '_fault_mode_cases/error_hccl_timeout.py')	141	path = os.path.join(os.path.dirname(__file__), '_fault_mode_cases/error_hccl_timeout.py')
140	process = subprocess.Popen(["torchrun", "--nproc-per-node=2", f"{path}"], shell=False, stdout=subprocess.PIPE,	142	process = subprocess.Popen(["torchrun", "--nproc-per-node=2", f"{path}"], shell=False, stdout=subprocess.PIPE,

Mtest/distributed/test_flight_recorder.py+1-0

	@@ -939,6 +939,7 @@ class HCCLTraceTestDumpOnHcclTimeout(HCCLTraceTestBase):
939	return None	939	return None
940		940
941	@parametrize("timing_enabled", [False])	941	@parametrize("timing_enabled", [False])
		942	+ @skipIf(True, "Environmental problem, temporarily skip.")
942	def test_hccl_timeout_dumps(self, timing_enabled):	943	def test_hccl_timeout_dumps(self, timing_enabled):
943	if self.rank != self.MAIN_PROCESS_RANK:	944	if self.rank != self.MAIN_PROCESS_RANK:
944	# dump on heartbeatmonitor thread	945	# dump on heartbeatmonitor thread

Athird_party/shmem/include/shmem_host_def.h+111-0

@@ -0,0 +1,111 @@
	1	+/*
	2	+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+ * This file is a part of the CANN Open Software.
	4	+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+#ifndef SHMEM_HOST_DEF_H
	11	+#define SHMEM_HOST_DEF_H
	12	+#include <climits>
	13	+#include "shmem_types.h"
	14	+
	15	+#ifdef __cplusplus
	16	+extern "C" {
	17	+#endif
	18	+
	19	+/**
	20	+ * @defgroup group_macros Macros
	21	+ * @{
	22	+*/
	23	+
	24	+/// \def SHMEM_XXX_VERSION
	25	+/// \brief macros that define current version info
	26	+#define SHMEM_MAX_IP_PORT_LEN 64
	27	+/*@} / // end of group_macros
	28	+
	29	+/**
	30	+ * @defgroup group_enums Enumerations
	31	+ * @{
	32	+*/
	33	+
	34	+/**
	35	+ * @brief Error code for the SHMEM library.
	36	+*/
	37	+enum shmem_error_code_t : int {
	38	+ SHMEM_SUCCESS = 0, ///< Task execution was successful.
	39	+ SHMEM_INVALID_PARAM = -1, ///< There is a problem with the parameters.
	40	+ SHMEM_INVALID_VALUE = -2, ///< There is a problem with the range of the value of the parameter.
	41	+ SHMEM_SMEM_ERROR = -3, ///< There is a problem with SMEM.
	42	+ SHMEM_INNER_ERROR = -4, ///< This is a problem caused by an internal error.
	43	+ SHMEM_NOT_INITED = -5, ///< This is a problem caused by an uninitialization.
	44	+};
	45	+
	46	+/**
	47	+ * @brief The state of the SHMEM library initialization.
	48	+*/
	49	+enum shmem_init_status_t {
	50	+ SHMEM_STATUS_NOT_INITIALIZED = 0, ///< Uninitialized.
	51	+ SHMEM_STATUS_SHM_CREATED, ///< Shared memory heap creation is complete.
	52	+ SHMEM_STATUS_IS_INITIALIZED, ///< Initialization is complete.
	53	+ SHMEM_STATUS_INVALID = INT_MAX, ///< Invalid status code.
	54	+};
	55	+
	56	+/*@} / // end of group_enums
	57	+
	58	+/**
	59	+ * @defgroup group_structs Structs
	60	+ * @{
	61	+*/
	62	+
	63	+constexpr uint16_t SHMEM_UNIQUE_ID_INNER_LEN = 60;
	64	+
	65	+typedef struct {
	66	+ int32_t version;
	67	+ char internal[SHMEM_UNIQUE_ID_INNER_LEN];
	68	+} shmem_uniqueid_t;
	69	+
	70	+/**
	71	+ * @struct shmem_init_optional_attr_t
	72	+ * @brief Optional parameter for the attributes used for initialization.
	73	+ *
	74	+ * - int version: version
	75	+ * - data_op_engine_type_t data_op_engine_type: data_op_engine_type
	76	+ * - uint32_t shm_init_timeout: shm_init_timeout
	77	+ * - uint32_t shm_create_timeout: shm_create_timeout
	78	+ * - uint32_t control_operation_timeout: control_operation_timeout
	79	+*/
	80	+typedef struct {
	81	+ int version;
	82	+ data_op_engine_type_t data_op_engine_type;
	83	+ uint32_t shm_init_timeout;
	84	+ uint32_t shm_create_timeout;
	85	+ uint32_t control_operation_timeout;
	86	+} shmem_init_optional_attr_t;
	87	+
	88	+/**
	89	+ * @struct shmem_init_attr_t
	90	+ * @brief Mandatory parameter for attributes used for initialization.
	91	+ *
	92	+ * - int my_rank: The rank of the current process.
	93	+ * - int n_ranks: The total rank number of all processes.
	94	+ * - char ip_port[SHMEM_MAX_IP_PORT_LEN]: The ip and port of the communication server. The port must not conflict
	95	+ * with other modules and processes.
	96	+ * - uint64_t local_mem_size: The size of shared memory currently occupied by current rank.
	97	+ * - shmem_init_optional_attr_t option_attr: Optional Parameters.
	98	+*/
	99	+typedef struct {
	100	+ int my_rank;
	101	+ int n_ranks;
	102	+ char ip_port[SHMEM_MAX_IP_PORT_LEN];
	103	+ uint64_t local_mem_size;
	104	+ shmem_init_optional_attr_t option_attr;
	105	+} shmem_init_attr_t;
	106	+
	107	+#ifdef __cplusplus
	108	+}
	109	+#endif
	110	+
	111	+#endif

Athird_party/shmem/include/shmem_types.h+32-0

@@ -0,0 +1,32 @@
	1	+/*
	2	+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+ * This file is a part of the CANN Open Software.
	4	+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+#ifndef SHMEM_TYPES_H
	11	+#define SHMEM_TYPES_H
	12	+
	13	+#ifdef __cplusplus
	14	+extern "C" {
	15	+#endif
	16	+
	17	+/**
	18	+ * @brief Data op engine type.
	19	+*/
	20	+enum data_op_engine_type_t {
	21	+ SHMEM_DATA_OP_MTE = 0x01,
	22	+ SHMEM_DATA_OP_SDMA = 0x02,
	23	+ SHMEM_DATA_OP_ROCE = 0x04,
	24	+};
	25	+
	26	+/*@} / // end of group_typedef
	27	+
	28	+#ifdef __cplusplus
	29	+}
	30	+#endif
	31	+
	32	+#endif /* SHMEM_TYPES_H */

Mtorch_npu/_logging/_internal.py+1-0

	@@ -38,3 +38,4 @@ def _add_logging_module():
38	torch._logging._internal.register_log("silent", "torch_npu.silent_check")	38	torch._logging._internal.register_log("silent", "torch_npu.silent_check")
39	torch._logging._internal.register_log("recovery", "torch_npu.recovery")	39	torch._logging._internal.register_log("recovery", "torch_npu.recovery")
40	torch._logging._internal.register_log("op_plugin", "torch_npu.op_plugin")	40	torch._logging._internal.register_log("op_plugin", "torch_npu.op_plugin")
		41	+ torch._logging._internal.register_log("shmem", "torch_npu.symmetric_memory")

Mtorch_npu/csrc/core/npu/NPUHooksInterface.cpp+5-0

	@@ -15,6 +15,11 @@ TORCH_DECLARE_REGISTRY(PrivateUse1HooksRegistry, NPUHooksInterface, NPUHooksArgs
15		15
16	C10_DEFINE_REGISTRY(PrivateUse1HooksRegistry, NPUHooksInterface, NPUHooksArgs)	16	C10_DEFINE_REGISTRY(PrivateUse1HooksRegistry, NPUHooksInterface, NPUHooksArgs)
17		17
		18	+at::Device NPUHooksInterface::getDeviceFromPtr(void* data) const
		19	+{
		20	+ return {at::DeviceType::PrivateUse1, c10_npu::current_device()};
		21	+}
		22	+
18	void NPUHooksInterface::init() const	23	void NPUHooksInterface::init() const
19	{	24	{
20	#ifndef BUILD_LIBTORCH	25	#ifndef BUILD_LIBTORCH

Mtorch_npu/csrc/core/npu/NPUHooksInterface.h+1-0

	@@ -12,6 +12,7 @@ struct TORCH_API NPUHooksInterface : public at::PrivateUse1HooksInterface {
12	static auto device_gen = at_npu::detail::getDefaultNPUGenerator(device_index);	12	static auto device_gen = at_npu::detail::getDefaultNPUGenerator(device_index);
13	return device_gen;	13	return device_gen;
14	}	14	}
		15	+ at::Device getDeviceFromPtr(void* data) const override;
15	void init() const override;	16	void init() const override;
16	at::Generator getNewGenerator(c10::DeviceIndex device_index = -1) const override;	17	at::Generator getNewGenerator(c10::DeviceIndex device_index = -1) const override;
17	bool hasPrimaryContext(c10::DeviceIndex device_index) const override;	18	bool hasPrimaryContext(c10::DeviceIndex device_index) const override;

Mtorch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.cpp+7-0

	@@ -18,6 +18,7 @@
18	#include "torch_npu/csrc/core/npu/register/OptionRegister.h"	18	#include "torch_npu/csrc/core/npu/register/OptionRegister.h"
19	#include "torch_npu/csrc/core/npu/register/OptionsManager.h"	19	#include "torch_npu/csrc/core/npu/register/OptionsManager.h"
20	#include "torch_npu/csrc/core/npu/NpuVariables.h"	20	#include "torch_npu/csrc/core/npu/NpuVariables.h"
		21	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h"
21	#include "third_party/acl/inc/acl/acl_op_compiler.h"	22	#include "third_party/acl/inc/acl/acl_op_compiler.h"
22	#include "third_party/acl/inc/acl/acl_rt.h"	23	#include "third_party/acl/inc/acl/acl_rt.h"
23	#include "torch_npu/csrc/framework/interface/AclOpCompileInterface.h"	24	#include "torch_npu/csrc/framework/interface/AclOpCompileInterface.h"
	@@ -281,6 +282,12 @@ NpuSysCtrl::SysStatus NpuSysCtrl::Finalize()
281	c10_npu::NPUEventManager::GetInstance().ClearEvent();	282	c10_npu::NPUEventManager::GetInstance().ClearEvent();
282	NPU_CHECK_WARN(c10_npu::DestroyUsedStreams());	283	NPU_CHECK_WARN(c10_npu::DestroyUsedStreams());
283	NPU_CHECK_WARN(c10_npu::ResetUsedDevices());	284	NPU_CHECK_WARN(c10_npu::ResetUsedDevices());
		285	+#ifndef BUILD_LIBTORCH
		286	+ if (c10d::symmetric_memory::Shmem_finalize_exist()) {
		287	+ auto ret = c10d::symmetric_memory::Shmem_finalize();
		288	+ ASCEND_LOGI("shmem_finalize emd, ret is %d", ret);
		289	+ }
		290	+#endif
284	// Maintain a basic point of view, who applies for the resource, the resource is released by whom.	291	// Maintain a basic point of view, who applies for the resource, the resource is released by whom.
285	// If aclInit is not a PTA call, then aclFinalize should not be a PTA call either.	292	// If aclInit is not a PTA call, then aclFinalize should not be a PTA call either.
286	if (repeat_init_acl_flag_) {	293	if (repeat_init_acl_flag_) {

Mtorch_npu/csrc/distributed/CMakeLists.txt+1-1

	@@ -1,4 +1,4 @@
1	-FILE(GLOB _DIST_SRCS .cpp rpc/.cpp)	1	+FILE(GLOB _DIST_SRCS .cpp rpc/.cpp symm_mem/*.cpp)
2		2
3	LIST(APPEND DIST_SRCS ${_DIST_SRCS})	3	LIST(APPEND DIST_SRCS ${_DIST_SRCS})
4		4

Atorch_npu/csrc/distributed/symm_mem/NPUSHMEMExtension.cpp+79-0

@@ -0,0 +1,79 @@
	1	+#include <torch/csrc/distributed/c10d/SymmetricMemory.hpp>
	2	+#include "torch_npu/csrc/core/npu/NPUException.h"
	3	+#include "torch_npu/csrc/core/npu/NPUFunctions.h"
	4	+#include "torch_npu/csrc/logging/LogContext.h"
	5	+#include "torch_npu/csrc/distributed/symm_mem/NPUSymmetricMemoryUtils.hpp"
	6	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h"
	7	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMExtension.h"
	8	+
	9	+namespace c10d::npushmem_extension {
	10	+
	11	+constexpr uint64_t local_mem_size = 1024UL * 1024UL * 1024UL;
	12	+static std::shared_ptr<npu_logging::Logger> logger = npu_logging::logging().getLogger("torch_npu.symmetric_memory");
	13	+using c10d::symmetric_memory::NPUStoreExchange;
	14	+static NPUStoreExchange storeExchange = NPUStoreExchange("npushmem_ext");
	15	+
	16	+void initialize_npushmem_with_store(
	17	+ c10::intrusive_ptr<c10d::Store> store,
	18	+ int rank,
	19	+ int world_size)
	20	+{
	21	+ static bool is_initialized = false;
	22	+ if (is_initialized) {
	23	+ return;
	24	+ }
	25	+
	26	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, rank is %d, world_size is %d.", rank, world_size);
	27	+
	28	+ uint32_t status = c10d::symmetric_memory::Shmem_set_conf_store_tls(false, nullptr, 0);
	29	+ TORCH_CHECK(status == 0, "shmem_set_conf_store_tls failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL));
	30	+
	31	+ shmem_uniqueid_t unique_id;
	32	+ if (rank == 0) {
	33	+ status = c10d::symmetric_memory::Shmem_get_uniqueid(&unique_id);
	34	+ TORCH_CHECK(status == 0, "shmem_get_uniqueid failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL));
	35	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, Shmem_get_uniqueid rank is %d, version %d, internal is %s.",
	36	+ rank, unique_id.version, unique_id.internal);
	37	+ }
	38	+ auto unique_ids = storeExchange.all_gather(store, rank, world_size, unique_id);
	39	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, unique_id rank is %d, version %d, internal is %s.",
	40	+ rank, unique_ids[0].version, unique_ids[0].internal);
	41	+
	42	+ auto env_val = std::getenv("NPU_SHMEM_SYMMETRIC_SIZE");
	43	+ int64_t init_size = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : local_mem_size;
	44	+ shmem_init_attr_t* attributes;
	45	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, start shmem_set_attr rank is %d, world_size is %d, size is %d.",
	46	+ rank, world_size, init_size);
	47	+ status = c10d::symmetric_memory::Shmem_set_attr(rank, world_size, init_size, nullptr, &attributes);
	48	+ TORCH_CHECK(status == 0, "shmem_set_attr failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL));
	49	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, end shmem_set_attr rank is %d, world_size is %d, size is %d.",
	50	+ rank, world_size, init_size);
	51	+
	52	+ status = c10d::symmetric_memory::Shmem_set_attr_uniqueid_args(rank, world_size, &unique_ids[0], attributes);
	53	+ TORCH_CHECK(status == 0, "Shmem_set_attr_uniqueid_args failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL));
	54	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store success, rank is %d, world_size is %d.", rank, world_size);
	55	+ is_initialized = true;
	56	+}
	57	+
	58	+void nvshmem_put(at::Tensor& tensor, int64_t peer)
	59	+{
	60	+ // to be done: support non-contiguous tensors
	61	+ TORCH_CHECK(tensor.is_contiguous(),
	62	+ "put op currently supports contiguous tensors only", DIST_ERROR(ErrCode::PARAM));
	63	+ // to be done: rendezvous should remember the group name
	64	+ auto hdl = c10d::symmetric_memory::rendezvous(tensor, "0");
	65	+ auto rank = hdl->get_rank();
	66	+ void* buffer_ptr = hdl->get_buffer_ptrs()[rank];
	67	+ auto buffer_size = tensor.numel() * tensor.element_size();
	68	+
	69	+ at::DeviceGuard device_guard(tensor.device());
	70	+ // to be done for putmem
	71	+ throw std::runtime_error("NPUSHMEMSymmetricMemory does not support nvshmem_put" + DIST_ERROR(ErrCode::NOT_SUPPORT));
	72	+}
	73	+
	74	+} // namespace c10d::npushmem_extension
	75	+
	76	+
	77	+TORCH_LIBRARY_IMPL(symm_mem, PrivateUse1, m) {
	78	+ m.impl("nvshmem_put", c10d::npushmem_extension::nvshmem_put);
	79	+}

Atorch_npu/csrc/distributed/symm_mem/NPUSHMEMExtension.h+16-0

@@ -0,0 +1,16 @@
	1	+#pragma once
	2	+
	3	+#include <ATen/ATen.h>
	4	+
	5	+#include <torch/csrc/distributed/c10d/Store.hpp>
	6	+
	7	+namespace c10d::npushmem_extension {
	8	+
	9	+void initialize_npushmem_with_store(
	10	+ c10::intrusive_ptr<c10d::Store> store,
	11	+ int rank,
	12	+ int world_size);
	13	+
	14	+TORCH_API void nvshmem_put(at::Tensor& tensor, int64_t peer);
	15	+
	16	+} // namespace c10d::npushmem_extension

Atorch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.cpp+141-0

@@ -0,0 +1,141 @@
	1	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h"
	2	+#include "torch_npu/csrc/core/npu/register/FunctionLoader.h"
	3	+#include "torch_npu/csrc/core/npu/NPUException.h"
	4	+
	5	+namespace c10d {
	6	+namespace symmetric_memory {
	7	+
	8	+#undef LOAD_FUNCTION
	9	+#define LOAD_FUNCTION(funcName) \
	10	+ REGISTER_FUNCTION(libshmem, funcName)
	11	+#undef GET_FUNC
	12	+#define GET_FUNC(funcName) \
	13	+ GET_FUNCTION(libshmem, funcName)
	14	+
	15	+REGISTER_LIBRARY(libshmem)
	16	+LOAD_FUNCTION(shmem_set_conf_store_tls)
	17	+LOAD_FUNCTION(shmem_set_attr)
	18	+LOAD_FUNCTION(shmem_init_attr)
	19	+LOAD_FUNCTION(shmem_get_uniqueid)
	20	+LOAD_FUNCTION(shmem_set_attr_uniqueid_args)
	21	+LOAD_FUNCTION(shmem_malloc)
	22	+LOAD_FUNCTION(shmem_free)
	23	+LOAD_FUNCTION(shmem_ptr)
	24	+LOAD_FUNCTION(shmem_finalize)
	25	+
	26	+int32_t Shmem_set_conf_store_tls(bool enable, const char *tls_info, const uint32_t tls_info_len)
	27	+{
	28	+ typedef int32_t (ShmemApiFunc)(bool, const char , const uint32_t);
	29	+ static ShmemApiFunc shmem_set_conf_store_tls_func = nullptr;
	30	+ if (shmem_set_conf_store_tls_func == nullptr) {
	31	+ shmem_set_conf_store_tls_func = (ShmemApiFunc)GET_FUNC(shmem_set_conf_store_tls);
	32	+ }
	33	+ TORCH_CHECK(shmem_set_conf_store_tls_func, "Failed to find function ", "shmem_set_conf_store_tls", PTA_ERROR(ErrCode::NOT_FOUND));
	34	+ return shmem_set_conf_store_tls_func(enable, tls_info, tls_info_len);
	35	+}
	36	+
	37	+int32_t Shmem_set_attr(int32_t my_rank, int32_t n_ranks, uint64_t local_mem_size, const char *ip_port,
	38	+ shmem_init_attr_t **attributes)
	39	+{
	40	+ typedef int32_t (ShmemApiFunc)(int32_t, int32_t, uint64_t, const char , shmem_init_attr_t **);
	41	+ static ShmemApiFunc shmem_set_attr_func = nullptr;
	42	+ if (shmem_set_attr_func == nullptr) {
	43	+ shmem_set_attr_func = (ShmemApiFunc)GET_FUNC(shmem_set_attr);
	44	+ }
	45	+ TORCH_CHECK(shmem_set_attr_func, "Failed to find function ", "shmem_set_attr", PTA_ERROR(ErrCode::NOT_FOUND));
	46	+ return shmem_set_attr_func(my_rank, n_ranks, local_mem_size, ip_port, attributes);
	47	+}
	48	+
	49	+int32_t Shmem_init_attr(shmem_init_attr_t *attributes)
	50	+{
	51	+ typedef int32_t (ShmemApiFunc)(shmem_init_attr_t );
	52	+ static ShmemApiFunc shmem_init_attr_func = nullptr;
	53	+ if (shmem_init_attr_func == nullptr) {
	54	+ shmem_init_attr_func = (ShmemApiFunc)GET_FUNC(shmem_init_attr);
	55	+ }
	56	+ TORCH_CHECK(shmem_init_attr_func, "Failed to find function ", "shmem_init_attr", PTA_ERROR(ErrCode::NOT_FOUND));
	57	+ return shmem_init_attr_func(attributes);
	58	+}
	59	+
	60	+int32_t Shmem_get_uniqueid(shmem_uniqueid_t *uid)
	61	+{
	62	+ typedef int32_t (ShmemApiFunc)(shmem_uniqueid_t );
	63	+ static ShmemApiFunc shmem_get_uniqueid_func = nullptr;
	64	+ if (shmem_get_uniqueid_func == nullptr) {
	65	+ shmem_get_uniqueid_func = (ShmemApiFunc)GET_FUNC(shmem_get_uniqueid);
	66	+ }
	67	+ TORCH_CHECK(shmem_get_uniqueid_func, "Failed to find function ", "shmem_get_uniqueid", PTA_ERROR(ErrCode::NOT_FOUND));
	68	+ return shmem_get_uniqueid_func(uid);
	69	+}
	70	+
	71	+int Shmem_set_attr_uniqueid_args(int rank_id, int nranks, const shmem_uniqueid_t uid, shmem_init_attr_t attr)
	72	+{
	73	+ typedef int32_t (ShmemApiFunc)(int, int, const shmem_uniqueid_t , shmem_init_attr_t *);
	74	+ static ShmemApiFunc shmem_set_attr_uniqueid_args_func = nullptr;
	75	+ if (shmem_set_attr_uniqueid_args_func == nullptr) {
	76	+ shmem_set_attr_uniqueid_args_func = (ShmemApiFunc)GET_FUNC(shmem_set_attr_uniqueid_args);
	77	+ }
	78	+ TORCH_CHECK(shmem_set_attr_uniqueid_args_func, "Failed to find function ", "shmem_set_attr_uniqueid_args", PTA_ERROR(ErrCode::NOT_FOUND));
	79	+ return shmem_set_attr_uniqueid_args_func(rank_id, nranks, uid, attr);
	80	+}
	81	+
	82	+void *Shmem_malloc(size_t size)
	83	+{
	84	+ typedef void* (*ShmemApiFunc)(size_t);
	85	+ static ShmemApiFunc shmem_malloc_func = nullptr;
	86	+ if (shmem_malloc_func == nullptr) {
	87	+ shmem_malloc_func = (ShmemApiFunc)GET_FUNC(shmem_malloc);
	88	+ }
	89	+ TORCH_CHECK(shmem_malloc_func, "Failed to find function ", "shmem_malloc", PTA_ERROR(ErrCode::NOT_FOUND));
	90	+ return shmem_malloc_func(size);
	91	+}
	92	+
	93	+void Shmem_free(void *ptr)
	94	+{
	95	+ typedef void (ShmemApiFunc)(void );
	96	+ static ShmemApiFunc shmem_free_func = nullptr;
	97	+ if (shmem_free_func == nullptr) {
	98	+ shmem_free_func = (ShmemApiFunc)GET_FUNC(shmem_free);
	99	+ }
	100	+ TORCH_CHECK(shmem_free_func, "Failed to find function ", "shmem_free", PTA_ERROR(ErrCode::NOT_FOUND));
	101	+ return shmem_free_func(ptr);
	102	+}
	103	+
	104	+void Shmem_ptr(void ptr, int pe)
	105	+{
	106	+ typedef void* (ShmemApiFunc)(void , int);
	107	+ static ShmemApiFunc shmem_ptr_func = nullptr;
	108	+ if (shmem_ptr_func == nullptr) {
	109	+ shmem_ptr_func = (ShmemApiFunc)GET_FUNC(shmem_ptr);
	110	+ }
	111	+ TORCH_CHECK(shmem_ptr_func, "Failed to find function ", "shmem_ptr", PTA_ERROR(ErrCode::NOT_FOUND));
	112	+ return shmem_ptr_func(ptr, pe);
	113	+}
	114	+
	115	+bool Shmem_finalize_exist()
	116	+{
	117	+ const static bool shmemApiFuncExist = []() -> bool {
	118	+ try {
	119	+ auto func = GET_FUNC(shmem_finalize)
	120	+ return func != nullptr;
	121	+ } catch (...) {
	122	+ // libshmem.so not exist
	123	+ return false;
	124	+ }
	125	+ }();
	126	+ return shmemApiFuncExist;
	127	+}
	128	+
	129	+int Shmem_finalize(void)
	130	+{
	131	+ typedef int (*ShmemApiFunc)(void);
	132	+ static ShmemApiFunc shmem_finalize_func = nullptr;
	133	+ if (shmem_finalize_func == nullptr) {
	134	+ shmem_finalize_func = (ShmemApiFunc)GET_FUNC(shmem_finalize);
	135	+ }
	136	+ TORCH_CHECK(shmem_finalize_func, "Failed to find function ", "shmem_finalize", PTA_ERROR(ErrCode::NOT_FOUND));
	137	+ return shmem_finalize_func();
	138	+}
	139	+
	140	+} // namespace symmetric_memory
	141	+} // namespace c10d

Atorch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h+32-0

@@ -0,0 +1,32 @@
	1	+#pragma once
	2	+
	3	+#include <cstddef>
	4	+#include <cstdint>
	5	+#include "third_party/shmem/include/shmem_host_def.h"
	6	+
	7	+namespace c10d {
	8	+namespace symmetric_memory {
	9	+
	10	+int32_t Shmem_set_conf_store_tls(bool enable, const char *tls_info, const uint32_t tls_info_len);
	11	+
	12	+int32_t Shmem_set_attr(int32_t my_rank, int32_t n_ranks, uint64_t local_mem_size, const char *ip_port,
	13	+ shmem_init_attr_t **attributes);
	14	+
	15	+int32_t Shmem_init_attr(shmem_init_attr_t *attributes);
	16	+
	17	+int Shmem_get_uniqueid(shmem_uniqueid_t *uid);
	18	+
	19	+int Shmem_set_attr_uniqueid_args(int rank_id, int nranks, const shmem_uniqueid_t uid, shmem_init_attr_t attr);
	20	+
	21	+void *Shmem_malloc(size_t size);
	22	+
	23	+void Shmem_free(void *ptr);
	24	+
	25	+void Shmem_ptr(void ptr, int pe);
	26	+
	27	+bool Shmem_finalize_exist();
	28	+
	29	+int Shmem_finalize(void);
	30	+
	31	+} // namespace symmetric_memory
	32	+} // namespace c10d

Atorch_npu/csrc/distributed/symm_mem/NPUSHMEMSymmetricMemory.cpp+262-0

@@ -0,0 +1,262 @@
	1	+#include "torch_npu/csrc/core/npu/NPUException.h"
	2	+#include "torch_npu/csrc/core/npu/NPUFunctions.h"
	3	+#include "torch_npu/csrc/logging/LogContext.h"
	4	+#include "torch_npu/csrc/distributed/symm_mem/NPUSymmetricMemoryUtils.hpp"
	5	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMExtension.h"
	6	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMSymmetricMemory.hpp"
	7	+
	8	+namespace c10d {
	9	+namespace symmetric_memory {
	10	+
	11	+/* Start of NPUSHMEMSymmetricMemory implementation */
	12	+
	13	+constexpr size_t npu_signal_pad_size = 2048;
	14	+static std::shared_ptr<npu_logging::Logger> logger = npu_logging::logging().getLogger("torch_npu.symmetric_memory");
	15	+static NPUStoreExchange storeExchange = NPUStoreExchange("NPUSHMEMSymmetricMemory");
	16	+
	17	+NPUSHMEMAllocation::~NPUSHMEMAllocation()
	18	+{
	19	+ // Avoid calling NPU functions after driver shutting down
	20	+ if (is_finalizing()) {
	21	+ return;
	22	+ }
	23	+ auto device = c10::Device(at::DeviceType::PrivateUse1, device_idx);
	24	+ at::DeviceGuard device_guard(device);
	25	+ logger->debug("~NPUSHMEMAllocation, start Shmem_free, ptr is %p.", ptr);
	26	+ Shmem_free(ptr); // shmem_free has no return value
	27	+ logger->debug("~NPUSHMEMAllocation, end Shmem_free, ptr is %p.", ptr);
	28	+}
	29	+
	30	+NPUSHMEMSymmetricMemory::NPUSHMEMSymmetricMemory(
	31	+ std::shared_ptr<NPUSHMEMAllocation> allocation,
	32	+ const std::string& group_name)
	33	+ : allocation_(allocation),
	34	+ buffer_size_(allocation->buffer_size),
	35	+ device_idx_(allocation->device_idx),
	36	+ group_name_(group_name)
	37	+{
	38	+ // For logging only
	39	+ static int exchanged_n_times = 0;
	40	+ auto device = c10::Device(at::DeviceType::PrivateUse1, device_idx_);
	41	+ at::DeviceGuard device_guard(device);
	42	+
	43	+ auto global_rank = get_group_info("0").rank;
	44	+ const GroupInfo& group_info = get_group_info(group_name_);
	45	+ auto store = group_info.store;
	46	+ rank_ = group_info.rank;
	47	+ world_size_ = group_info.world_size;
	48	+ // Exchange rank to global rank mapping for this group.
	49	+ // If it is already available, skip the exchange.
	50	+ if (rank_to_global_rank_.empty()) {
	51	+ rank_to_global_rank_ =
	52	+ storeExchange.all_gather(store, rank_, world_size_, global_rank);
	53	+ exchanged_n_times++;
	54	+ if (rank_ == 0) {
	55	+ std::stringstream ss;
	56	+ for (size_t i = 0; i < rank_to_global_rank_.size(); ++i) {
	57	+ ss << rank_to_global_rank_[i];
	58	+ if (i != rank_to_global_rank_.size() - 1) {
	59	+ ss << ", ";
	60	+ }
	61	+ }
	62	+ logger->debug("[rank %d] rank_to_global_rank: %s, group_name: %s, exchanged_n_times: %d.",
	63	+ rank_, (ss.str()).c_str(), group_name_.c_str(), exchanged_n_times);
	64	+ }
	65	+ }
	66	+ TORCH_INTERNAL_ASSERT(!rank_to_global_rank_.empty());
	67	+ for (int r = 0; r < world_size_; ++r) {
	68	+ auto buffer = Shmem_ptr(allocation->ptr, rank_to_global_rank_[r]);
	69	+ buffers_.push_back(buffer);
	70	+ logger->debug("[rank %d] NPUSHMEMSymmetricMemory shmem_ptr, r is %d, rank_to_global_rank is %d, ptr is %p, shmem_ptr is %p.",
	71	+ rank_, r, rank_to_global_rank_[r], allocation->ptr, buffer);
	72	+ }
	73	+
	74	+ // to be done
	75	+ // signal_pads_ buffers_dev_ signal_pads_dev_ rank_to_global_rank_dev_
	76	+ logger->debug("NPUSHMEMSymmetricMemory created, buffer_size is %d, device_idx is %d, group_name is %s.",
	77	+ allocation->buffer_size, allocation->device_idx, group_name_.c_str());
	78	+}
	79	+
	80	+NPUSHMEMSymmetricMemory::~NPUSHMEMSymmetricMemory()
	81	+{
	82	+ // to be done
	83	+ logger->debug("NPUSHMEMSymmetricMemory destroy, group_name is %s", group_name_.c_str());
	84	+}
	85	+
	86	+std::vector<void*> NPUSHMEMSymmetricMemory::get_buffer_ptrs()
	87	+{
	88	+ return buffers_;
	89	+}
	90	+
	91	+std::vector<void*> NPUSHMEMSymmetricMemory::get_signal_pad_ptrs()
	92	+{
	93	+ return signal_pads_;
	94	+}
	95	+
	96	+void** NPUSHMEMSymmetricMemory::get_buffer_ptrs_dev()
	97	+{
	98	+ return buffers_dev_;
	99	+}
	100	+
	101	+void** NPUSHMEMSymmetricMemory::get_signal_pad_ptrs_dev()
	102	+{
	103	+ return signal_pads_dev_;
	104	+}
	105	+
	106	+size_t NPUSHMEMSymmetricMemory::get_buffer_size()
	107	+{
	108	+ return buffer_size_;
	109	+}
	110	+
	111	+size_t NPUSHMEMSymmetricMemory::get_signal_pad_size()
	112	+{
	113	+ return npu_signal_pad_size;
	114	+}
	115	+
	116	+bool NPUSHMEMSymmetricMemory::has_multicast_support()
	117	+{
	118	+ return false;
	119	+}
	120	+
	121	+void* NPUSHMEMSymmetricMemory::get_multicast_ptr()
	122	+{
	123	+ return nullptr;
	124	+}
	125	+
	126	+at::Tensor NPUSHMEMSymmetricMemory::get_buffer(
	127	+ int rank,
	128	+ c10::IntArrayRef sizes,
	129	+ c10::ScalarType dtype,
	130	+ int64_t storage_offset)
	131	+{
	132	+ // to be done
	133	+ throw std::runtime_error("NPUSHMEMSymmetricMemory does not support get_buffer" + DIST_ERROR(ErrCode::NOT_SUPPORT));
	134	+}
	135	+
	136	+at::Tensor NPUSHMEMSymmetricMemory::get_signal_pad(
	137	+ int rank,
	138	+ c10::IntArrayRef sizes,
	139	+ std::optional<c10::ScalarType> dtype,
	140	+ int64_t storage_offset)
	141	+{
	142	+ // to be done
	143	+ throw std::runtime_error("NPUSHMEMSymmetricMemory does not support get_signal_pad" + DIST_ERROR(ErrCode::NOT_SUPPORT));
	144	+}
	145	+
	146	+void NPUSHMEMSymmetricMemory::barrier(int channel, size_t timeout_ms)
	147	+{
	148	+ // to be done
	149	+}
	150	+
	151	+void NPUSHMEMSymmetricMemory::put_signal(
	152	+ int dst_rank,
	153	+ int channel,
	154	+ size_t timeout_ms)
	155	+{
	156	+ // to be done
	157	+}
	158	+
	159	+void NPUSHMEMSymmetricMemory::wait_signal(
	160	+ int src_rank,
	161	+ int channel,
	162	+ size_t timeout_ms)
	163	+{
	164	+ // to be done
	165	+}
	166	+
	167	+int NPUSHMEMSymmetricMemory::get_rank()
	168	+{
	169	+ return rank_;
	170	+}
	171	+
	172	+int NPUSHMEMSymmetricMemory::get_world_size()
	173	+{
	174	+ return world_size_;
	175	+}
	176	+
	177	+void* NPUSHMEMSymmetricMemoryAllocator::alloc(
	178	+ size_t size,
	179	+ int device_idx,
	180	+ const std::optional<std::string>& group_name)
	181	+{
	182	+ TORCH_CHECK(
	183	+ group_name == std::nullopt,
	184	+ "NPUSHMEMSymmetricMemoryAllocator::alloc "
	185	+ "must not be called with a group_name", DIST_ERROR(ErrCode::PARAM));
	186	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator alloc start, size is %d, device is %d, group_name is %s",
	187	+ size, device_idx, group_name == std::nullopt ? "" : (*group_name).c_str());
	188	+
	189	+ c10_npu::LazySetDevice(device_idx);
	190	+ auto group_info = get_group_info("0");
	191	+ auto store = group_info.store;
	192	+ int rank = group_info.rank;
	193	+ int world_size = group_info.world_size;
	194	+ npushmem_extension::initialize_npushmem_with_store(store, rank, world_size);
	195	+
	196	+ auto ptr = Shmem_malloc(size);
	197	+ auto allocation =
	198	+ std::make_shared<NPUSHMEMAllocation>(ptr, size, device_idx);
	199	+ // to be done: thread safety
	200	+ allocations_.try_emplace(ptr, std::move(allocation));
	201	+ logger->debug("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);
	203	+ return ptr;
	204	+}
	205	+
	206	+void NPUSHMEMSymmetricMemoryAllocator::free(void* ptr)
	207	+{
	208	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator free start, ptr is %p", ptr);
	209	+ allocations_.erase(ptr);
	210	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator free end, ptr is %p", ptr);
	211	+}
	212	+
	213	+size_t NPUSHMEMSymmetricMemoryAllocator::get_alloc_size(void* ptr)
	214	+{
	215	+ auto it = allocations_.find(ptr);
	216	+ if (it == allocations_.end()) {
	217	+ TORCH_CHECK(false, ptr, " is not allocated with NPUSHMEMSymmetricMemoryAllocator", DIST_ERROR(ErrCode::PARAM));
	218	+ }
	219	+ return it->second->buffer_size;
	220	+}
	221	+
	222	+c10::intrusive_ptr<SymmetricMemory> NPUSHMEMSymmetricMemoryAllocator::rendezvous(
	223	+ void* ptr,
	224	+ const std::optional<std::string>& group_name)
	225	+{
	226	+ logger->debug("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));
	228	+ {
	229	+ auto it = symm_mems_.find(std::make_tuple(ptr, *group_name));
	230	+ if (it != symm_mems_.end()) {
	231	+ return it->second;
	232	+ }
	233	+ }
	234	+ auto it = allocations_.find(ptr);
	235	+ TORCH_CHECK(it != allocations_.end(), "rendezvous, ptr is invalid.", DIST_ERROR(ErrCode::PARAM));
	236	+ auto symm_mem =
	237	+ c10::make_intrusive<NPUSHMEMSymmetricMemory>(it->second, *group_name);
	238	+
	239	+ symm_mems_[std::make_tuple(ptr, *group_name)] = symm_mem;
	240	+ logger->debug("NPUSHMEMSymmetricMemoryAllocator rendezvous end, ptr is %p, group_name is %s", ptr, (*group_name).c_str());
	241	+ return symm_mem;
	242	+}
	243	+
	244	+bool NPUSHMEMSymmetricMemoryAllocator::has_multicast_support(int device_idx)
	245	+{
	246	+ // to be done
	247	+ throw std::runtime_error("NPUSHMEMSymmetricMemoryAllocator does not support has_multicast_support" + DIST_ERROR(ErrCode::NOT_SUPPORT));
	248	+}
	249	+
	250	+struct RegisterNPUSHMEMSymmetricMemoryAllocator {
	251	+ RegisterNPUSHMEMSymmetricMemoryAllocator()
	252	+ {
	253	+ register_allocator(
	254	+ c10::DeviceType::PrivateUse1,
	255	+ c10::make_intrusive<NPUSHMEMSymmetricMemoryAllocator>());
	256	+ }
	257	+};
	258	+
	259	+static RegisterNPUSHMEMSymmetricMemoryAllocator register_allocator_;
	260	+
	261	+} // namespace symmetric_memory
	262	+} // namespace c10d

Atorch_npu/csrc/distributed/symm_mem/NPUSHMEMSymmetricMemory.hpp+99-0

@@ -0,0 +1,99 @@
	1	+#pragma once
	2	+
	3	+#include <unordered_map>
	4	+#include <map>
	5	+#include <memory>
	6	+#include <torch/csrc/distributed/c10d/SymmetricMemory.hpp>
	7	+#include "torch_npu/csrc/distributed/symm_mem/NPUSHMEMInterface.h"
	8	+
	9	+namespace c10d {
	10	+namespace symmetric_memory {
	11	+
	12	+struct NPUSHMEMAllocation {
	13	+ void* ptr;
	14	+ size_t buffer_size;
	15	+ int device_idx;
	16	+
	17	+ NPUSHMEMAllocation(void* ptr, size_t buffer_size, int device_idx)
	18	+ : ptr(ptr), buffer_size(buffer_size), device_idx(device_idx) {}
	19	+
	20	+ ~NPUSHMEMAllocation();
	21	+};
	22	+
	23	+class NPUSHMEMSymmetricMemory : public SymmetricMemory {
	24	+public:
	25	+ NPUSHMEMSymmetricMemory(
	26	+ std::shared_ptr<NPUSHMEMAllocation> allocation,
	27	+ const std::string& group_name);
	28	+
	29	+ ~NPUSHMEMSymmetricMemory() override;
	30	+
	31	+ std::vector<void*> get_buffer_ptrs() override;
	32	+ std::vector<void*> get_signal_pad_ptrs() override;
	33	+ void** get_buffer_ptrs_dev() override;
	34	+ void** get_signal_pad_ptrs_dev() override;
	35	+ size_t get_buffer_size() override;
	36	+ size_t get_signal_pad_size() override;
	37	+
	38	+ bool has_multicast_support() override;
	39	+ void* get_multicast_ptr() override;
	40	+
	41	+ at::Tensor get_buffer(
	42	+ int rank,
	43	+ c10::IntArrayRef sizes,
	44	+ c10::ScalarType dtype,
	45	+ int64_t storage_offset) override;
	46	+
	47	+ at::Tensor get_signal_pad(
	48	+ int rank,
	49	+ c10::IntArrayRef sizes,
	50	+ std::optional<c10::ScalarType> dtype,
	51	+ int64_t storage_offset) override;
	52	+
	53	+ void barrier(int channel, size_t timeout_ms) override;
	54	+ void put_signal(int dst_rank, int channel, size_t timeout_ms) override;
	55	+ void wait_signal(int src_rank, int channel, size_t timeout_ms) override;
	56	+
	57	+ int get_rank() override;
	58	+ int get_world_size() override;
	59	+
	60	+private:
	61	+ std::shared_ptr<NPUSHMEMAllocation> allocation_;
	62	+ size_t buffer_size_;
	63	+ std::vector<void*> buffers_;
	64	+ std::vector<void*> signal_pads_;
	65	+ int device_idx_;
	66	+ int rank_;
	67	+ int world_size_;
	68	+ void** buffers_dev_;
	69	+ void** signal_pads_dev_;
	70	+ std::string group_name_;
	71	+
	72	+ std::vector<int> rank_to_global_rank_;
	73	+};
	74	+
	75	+class NPUSHMEMSymmetricMemoryAllocator : public SymmetricMemoryAllocator {
	76	+public:
	77	+ void* alloc(
	78	+ size_t size,
	79	+ int device_idx,
	80	+ const std::optional<std::string>& group_name) override;
	81	+
	82	+ void free(void* ptr) override;
	83	+
	84	+ size_t get_alloc_size(void* ptr) override;
	85	+
	86	+ c10::intrusive_ptr<SymmetricMemory> rendezvous(
	87	+ void* ptr,
	88	+ const std::optional<std::string>& group_name) override;
	89	+
	90	+ bool has_multicast_support(int device_idx) override;
	91	+
	92	+private:
	93	+ std::unordered_map<void*, std::shared_ptr<NPUSHMEMAllocation>> allocations_;
	94	+ std::map<std::tuple<void*, std::string>, c10::intrusive_ptr<SymmetricMemory>>
	95	+ symm_mems_;
	96	+};
	97	+
	98	+} // namespace symmetric_memory
	99	+} // namespace c10d

Atorch_npu/csrc/distributed/symm_mem/NPUSymmetricMemoryUtils.hpp+74-0

@@ -0,0 +1,74 @@
	1	+#pragma once
	2	+
	3	+#include <torch/csrc/distributed/c10d/Store.hpp>
	4	+
	5	+namespace c10d {
	6	+namespace symmetric_memory {
	7	+
	8	+// A set of store-based exchange methods with a preset prefix typically type of
	9	+// the SymmetricMemory. Most used as static instances at respective
	10	+// SymmetricMemory implementation files.
	11	+class NPUStoreExchange {
	12	+public:
	13	+ explicit NPUStoreExchange(const std::string& store_prefix)
	14	+ : store_prefix_(store_prefix) {}
	15	+
	16	+ // Put template function in header file so that compiler can easily access it.
	17	+ template <typename T>
	18	+ std::vector<T> all_gather(
	19	+ const c10::intrusive_ptr<c10d::Store>& store,
	20	+ int rank,
	21	+ int world_size,
	22	+ T val)
	23	+ {
	24	+ static_assert(std::is_trivially_copyable_v<T>);
	25	+
	26	+ std::vector<std::string> peer_keys;
	27	+ peer_keys.reserve(world_size);
	28	+ for (int r = 0; r < world_size; ++r) {
	29	+ std::ostringstream oss;
	30	+ oss << store_prefix_ << "/" << seq_id_ << "/" << r;
	31	+ peer_keys.push_back(oss.str());
	32	+ }
	33	+ ++seq_id_;
	34	+
	35	+ {
	36	+ std::vector<uint8_t> payload(
	37	+ reinterpret_cast<uint8_t*>(&val),
	38	+ reinterpret_cast<uint8_t*>(&val) + sizeof(T));
	39	+ store->set(peer_keys[rank], payload);
	40	+ }
	41	+
	42	+ std::vector<T> peer_vals;
	43	+ peer_vals.reserve(world_size);
	44	+ for (int r = 0; r < world_size; ++r) {
	45	+ if (r == rank) {
	46	+ peer_vals.push_back(val);
	47	+ continue;
	48	+ }
	49	+ store->wait({peer_keys[r]});
	50	+ auto payload = store->get(peer_keys[r]);
	51	+ TORCH_CHECK(payload.size() == sizeof(T));
	52	+ T peer_val{};
	53	+ std::memcpy(&peer_val, payload.data(), sizeof(T));
	54	+ peer_vals.push_back(peer_val);
	55	+ }
	56	+ return peer_vals;
	57	+ }
	58	+
	59	+ void barrier(
	60	+ const c10::intrusive_ptr<c10d::Store>& store,
	61	+ int rank,
	62	+ int world_size)
	63	+ {
	64	+ // to be done: implement an efficient one?
	65	+ all_gather(store, rank, world_size, 0);
	66	+ }
	67	+
	68	+private:
	69	+ const std::string store_prefix_;
	70	+ size_t seq_id_ = 0;
	71	+};
	72	+
	73	+} // namespace symmetric_memory
	74	+} // namespace c10d