Star402
274
代码
代码
Issues6
Pull Requests51
Actions
项目讨论
Wiki
分析
Star402
274
已合并
TGET/TPUT ASYNC CPU-SIM support Urma mode #991
Azizbek创建于 5月26日
TGET/TPUT ASYNC CPU-SIM support Urma mode #991
已合并
Azizbek创建于 5月26日
从Azizbek/pto-isaurma合入到cann/pto-isamaster
共 13 个文件变更+649-1
Mtests/cpu/comm/st/testcase/CMakeLists.txt+2-0
@@ -39,6 +39,8 @@ find_package(Threads REQUIRED)
39set(ALL_TESTCASES39set(ALL_TESTCASES
40tscatter40tscatter
41tgather41tgather
42+tget_async_urma
Z
ZZhanghaijian5月26日

TGET/TPUT ASYNC already supported

likedislike
43+tput_async_urma
42)44)
43 45 
44foreach(TESTCASE ${ALL_TESTCASES})46foreach(TESTCASE ${ALL_TESTCASES})
Mtests/cpu/comm/st/testcase/common.hpp+4-1
@@ -33,6 +33,9 @@ inline void *WindowAlloc(uint64_t windowBase, size_t &offset, size_t bytes)
33 return ptr;33 return ptr;
34}34}
35 35 
36+inline void CommMpiBarrier()
37+{}
38+ 
36template <typename T, size_t count>39template <typename T, size_t count>
37struct TestContext {40struct TestContext {
38 int32_t deviceId{-1};41 int32_t deviceId{-1};
@@ -90,4 +93,4 @@ inline bool ForkAndRunWithHcclRootInfo(int nRanks, int firstRankId, int firstDev
90 }93 }
91 94 
92 return res;95 return res;
93-}96+}
Atests/cpu/comm/st/testcase/tget_async_urma/CMakeLists.txt+1-0
@@ -0,0 +1 @@
1+pto_cpu_sim_comm_st(tget_async_urma)
Atests/cpu/comm/st/testcase/tget_async_urma/gen_data.py+19-0
@@ -0,0 +1,19 @@
1+#!/usr/bin/python3
2+# coding=utf-8
3+# --------------------------------------------------------------------------------
4+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# --------------------------------------------------------------------------------
12+ 
13+import os
14+ 
15+def main():
16+ os.makedirs("testcases", exist_ok=True)
17+ 
18+if __name__ == "__main__":
19+ main()
Atests/cpu/comm/st/testcase/tget_async_urma/main.cpp+60-0
@@ -0,0 +1,60 @@
1+/**
2+Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+ 
11+#include <cstddef>
12+#include <cstdint>
13+#include <gtest/gtest.h>
14+ 
15+#include "tget_async_urma_kernel.h"
16+#include "pto/common/cpu_stub.hpp"
17+ 
18+// ============================================================================
19+// 1D Vector Tile Tests (URMA true async GET on A5 3510)
20+// ============================================================================
21+TEST(TGetAsyncUrma, Vec_FloatSmall)
22+{
23+ SKIP_IF_RANKS_LT(2);
24+ ASSERT_TRUE((RunGetAsyncUrmaRootGet<float, 256>(2, 2, 0, 0)));
25+}
26+TEST(TGetAsyncUrma, Vec_Int32Large)
27+{
28+ SKIP_IF_RANKS_LT(2);
29+ ASSERT_TRUE((RunGetAsyncUrmaRootGet<int32_t, 4096>(2, 2, 0, 0)));
30+}
31+TEST(TGetAsyncUrma, Vec_Uint8Small)
32+{
33+ SKIP_IF_RANKS_LT(2);
34+ ASSERT_TRUE((RunGetAsyncUrmaRootGet<uint8_t, 512>(2, 2, 0, 0)));
35+}
36+ 
37+// ============================================================================
38+// Large MR boundary tests
39+// ============================================================================
40+TEST(TGetAsyncUrma, Vec_Float_MR_8MB)
41+{
42+ SKIP_IF_RANKS_LT(2);
43+ // 512K floats → commBytesNeeded ≈ 6MB (3 buffers), allocSize = 8MB (>2MB)
44+ ASSERT_TRUE((RunGetAsyncUrmaRootGet<float, 524288>(2, 2, 0, 0)));
45+}
46+TEST(TGetAsyncUrma, Vec_Int32_MR_Over512MB)
47+{
48+ SKIP_IF_RANKS_LT(2);
49+ // 64M int32 → commBytesNeeded ≈ 768MB (3 buffers), allocSize ≈ 770MB (>512MB)
50+ ASSERT_TRUE((RunGetAsyncUrmaRootGet<int32_t, 67108864>(2, 2, 0, 0)));
51+}
52+ 
53+int main(int argc, char **argv)
54+{
55+ CommMpiInit(&argc, &argv);
56+ ::testing::InitGoogleTest(&argc, argv);
57+ int ret = RUN_ALL_TESTS();
58+ CommMpiFinalize();
59+ return ret;
60+}
Atests/cpu/comm/st/testcase/tget_async_urma/tget_async_urma_kernel.cpp+175-0
@@ -0,0 +1,175 @@
1+/**
2+Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+ 
11+#include <cstddef>
12+#include <cstdint>
13+#include <iostream>
14+#include <pto/pto-inst.hpp>
15+ 
16+#define PTO_URMA_SUPPORTED
17+ 
18+#include "pto/pto-inst.hpp"
19+#include "pto/common/pto_tile.hpp"
20+#include "pto/common/cpu_stub.hpp"
21+#include "../common.hpp"
22+#include "../urma_context.hpp"
23+ 
24+// ============================================================================
25+// TGET_ASYNC via URMA — device kernel.
26+// ============================================================================
27+ 
28+template <typename T, size_t count>
29+__global__ AICORE void TGetAsyncUrmaKernelImpl(__gm__ T *localBuf, int nranks, int my_rank, int first_rank_id,
30+ int root_rank, int elem_offset, int elem_count,
31+ __gm__ uint8_t *urmaWorkspace)
32+{
33+ using ShapeDyn = pto::Shape<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC>;
34+ using StrideDyn = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC>;
35+ using Global = pto::GlobalTensor<T, ShapeDyn, StrideDyn, pto::Layout::ND>;
36+ 
37+ if (elem_count <= 0 || elem_offset < 0 || elem_offset + elem_count > static_cast<int>(count)) {
38+ pipe_barrier(PIPE_ALL);
39+ return;
40+ }
41+ 
42+ ShapeDyn shape(1, 1, 1, 1, elem_count);
43+ StrideDyn stride(elem_count, elem_count, elem_count, elem_count, 1);
44+ 
45+ constexpr size_t kDataOffset = 64 * sizeof(int32_t);
46+ 
47+ __gm__ T *sendBuf = reinterpret_cast<__gm__ T *>(reinterpret_cast<__gm__ uint8_t *>(localBuf) + kDataOffset);
48+ __gm__ T *recvBuf = sendBuf + count;
49+ 
50+ pipe_barrier(PIPE_ALL);
51+ 
52+ if (my_rank == root_rank) {
53+#ifdef PTO_URMA_SUPPORTED
54+ const int my_peer = my_rank - first_rank_id;
55+ for (int target_peer = 0; target_peer < nranks; ++target_peer) {
56+ if (target_peer == my_peer) {
57+ continue;
58+ }
59+ uint64_t peerBase = UrmaPeerMrBaseAddr(urmaWorkspace, static_cast<uint32_t>(target_peer));
60+ __gm__ T *remoteSendBuf = reinterpret_cast<__gm__ T *>(peerBase + kDataOffset) + elem_offset;
61+ __gm__ T *localRecvBuf = recvBuf + target_peer * count + elem_offset;
62+ Global remoteSendG(remoteSendBuf, shape, stride);
63+ Global localRecvG(localRecvBuf, shape, stride);
64+ 
65+ pto::comm::AsyncSession session;
66+ BuildAsyncSession<pto::comm::DmaEngine::URMA>(urmaWorkspace, static_cast<uint32_t>(target_peer), session);
67+ auto event = pto::comm::TGET_ASYNC<pto::comm::DmaEngine::URMA>(localRecvG, remoteSendG, session);
68+ event.Wait(session);
69+ }
70+#endif
71+ }
72+ 
73+ pipe_barrier(PIPE_ALL);
74+}
75+ 
76+// ============================================================================
77+// Verify root-get results: each remote rank's data should match pattern i + rank * 10000.
78+// ============================================================================
79+template <typename T, size_t count>
80+bool VerifyRootGetResults(const uint8_t *output_host, int n_ranks, int first_rank_id, int root_rank, int rank_id,
81+ int deviceId, int syncRet)
82+{
83+ const int root_peer = root_rank - first_rank_id;
84+ for (int src_peer = 0; src_peer < n_ranks; ++src_peer) {
85+ if (src_peer == root_peer)
86+ continue;
87+ const int src_logical = first_rank_id + src_peer;
88+ const size_t base = static_cast<size_t>(src_peer) * count;
89+ for (size_t i = 0; i < count; ++i) {
90+ T value = reinterpret_cast<const T *>(output_host)[base + i];
91+ T expected = static_cast<T>(i + src_logical * 10000);
92+ if (value != expected) {
93+ std::cerr << "Rank " << rank_id << " Device " << deviceId << " SyncRet " << syncRet
94+ << " Expected: " << (float)expected << " Actual: " << (float)value << std::endl;
95+ return false;
96+ }
97+ }
98+ }
99+ return true;
100+}
101+ 
102+// ============================================================================
103+// Host-side runner.
104+// ============================================================================
105+template <typename T, size_t count>
106+bool RunGetAsyncUrmaRootGetKernel(int rank_id, int n_ranks, int n_devices, int first_device_id, int first_rank_id,
107+ int root_rank)
108+{
109+ const size_t recv_elems = static_cast<size_t>(n_ranks) * count;
110+ size_t commBytesNeeded = 64 * sizeof(int32_t) + (static_cast<size_t>(n_ranks) + 1) * count * sizeof(T);
111+ 
112+ UrmaTestContext ctx;
113+ if (!ctx.Setup(rank_id, n_ranks, n_devices, first_device_id, root_rank, commBytesNeeded)) {
114+ return false;
115+ }
116+ 
117+ uint8_t *input_host = nullptr, *output_host = nullptr;
118+ aclrtMallocHost(reinterpret_cast<void **>(&input_host), count * sizeof(T));
119+ aclrtMallocHost(reinterpret_cast<void **>(&output_host), recv_elems * sizeof(T));
120+ if (!input_host || !output_host) {
121+ std::cerr << "[ERROR] aclrtMallocHost failed!" << std::endl;
122+ ctx.Cleanup();
123+ return false;
124+ }
125+ for (size_t i = 0; i < count; ++i)
126+ reinterpret_cast<T *>(input_host)[i] = static_cast<T>(i + rank_id * 10000);
127+ for (size_t i = 0; i < recv_elems; ++i)
128+ reinterpret_cast<T *>(output_host)[i] = static_cast<T>(-1);
129+ 
130+ constexpr size_t kDataOffset = 64 * sizeof(int32_t);
131+ T *sendBuf = reinterpret_cast<T *>(reinterpret_cast<uint8_t *>(ctx.devBuf) + kDataOffset);
132+ T *recvBuf = sendBuf + count;
133+ aclrtMemcpy(sendBuf, count * sizeof(T), input_host, count * sizeof(T), ACL_MEMCPY_HOST_TO_DEVICE);
134+ aclrtMemcpy(recvBuf, recv_elems * sizeof(T), output_host, recv_elems * sizeof(T), ACL_MEMCPY_HOST_TO_DEVICE);
135+ 
136+ CommMpiBarrier();
137+ 
138+ TGetAsyncUrmaKernelImpl<T, count>(reinterpret_cast<T *>(ctx.devBuf), n_ranks, rank_id, first_rank_id, root_rank, 0,
139+ static_cast<int>(count),
140+ reinterpret_cast<uint8_t *>(ctx.urmaMgr.GetWorkspaceAddr()));
141+ int syncRet = aclrtSynchronizeStream(ctx.stream);
142+ 
143+ CommMpiBarrier();
144+ 
145+ bool is_ok = true;
146+ if (rank_id == root_rank) {
147+ aclrtMemcpy(output_host, recv_elems * sizeof(T), recvBuf, recv_elems * sizeof(T), ACL_MEMCPY_DEVICE_TO_HOST);
148+ is_ok = VerifyRootGetResults<T, count>(output_host, n_ranks, first_rank_id, root_rank, rank_id, ctx.deviceId,
149+ syncRet);
150+ }
151+ 
152+ aclrtFreeHost(input_host);
153+ aclrtFreeHost(output_host);
154+ ctx.Cleanup();
155+ 
156+ return is_ok;
157+}
158+ 
159+// ============================================================================
160+// MPI-based multi-rank launch.
161+// ============================================================================
162+template <typename T, size_t count>
163+bool RunGetAsyncUrmaRootGet(int n_ranks, int n_devices, int first_rank_id, int first_device_id)
164+{
165+ return RunUrmaTestMpiLaunch(n_ranks, n_devices, first_rank_id, first_device_id,
166+ RunGetAsyncUrmaRootGetKernel<T, count>);
167+}
168+ 
169+// Explicit instantiations
170+template bool RunGetAsyncUrmaRootGet<float, 256>(int, int, int, int);
171+template bool RunGetAsyncUrmaRootGet<int32_t, 4096>(int, int, int, int);
172+template bool RunGetAsyncUrmaRootGet<uint8_t, 512>(int, int, int, int);
173+template bool RunGetAsyncUrmaRootGet<float, 524288>(int, int, int,
174+ int); // MR = 8MB (>2MB)
175+template bool RunGetAsyncUrmaRootGet<int32_t, 67108864>(int, int, int, int); // MR ≈ 770MB (>512MB)
Atests/cpu/comm/st/testcase/tget_async_urma/tget_async_urma_kernel.h+17-0
@@ -0,0 +1,17 @@
1+/**
2+Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+ 
11+#pragma once
12+ 
13+#include <cstddef>
14+#include <cstdint>
15+ 
16+template <typename T, size_t count>
17+bool RunGetAsyncUrmaRootGet(int n_ranks, int n_devices, int first_rank_id, int first_device_id);
Atests/cpu/comm/st/testcase/tput_async_urma/CMakeLists.txt+1-0
@@ -0,0 +1 @@
1+pto_cpu_sim_comm_st(tput_async_urma)
Atests/cpu/comm/st/testcase/tput_async_urma/gen_data.py+19-0
@@ -0,0 +1,19 @@
1+#!/usr/bin/python3
2+# coding=utf-8
3+# --------------------------------------------------------------------------------
4+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# --------------------------------------------------------------------------------
12+ 
13+import os
14+ 
15+def main():
16+ os.makedirs("testcases", exist_ok=True)
17+ 
18+if __name__ == "__main__":
19+ main()
Atests/cpu/comm/st/testcase/tput_async_urma/main.cpp+83-0
@@ -0,0 +1,83 @@
1+/**
2+Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+ 
11+#include <cstddef>
12+#include <cstdint>
13+#include <gtest/gtest.h>
14+ 
15+#include "tput_async_urma_kernel.h"
16+#include "pto/common/cpu_stub.hpp"
17+ 
18+// ============================================================================
19+// Basic correctness (URMA true async PUT on A5 3510)
20+// ============================================================================
21+TEST(TPutAsyncUrma, Vec_FloatSmall)
22+{
23+ SKIP_IF_RANKS_LT(2);
24+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<float, 256>(2, 2, 0, 0)));
25+}
26+TEST(TPutAsyncUrma, Vec_Int32Large)
27+{
28+ SKIP_IF_RANKS_LT(2);
29+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<int32_t, 4096>(2, 2, 0, 0)));
30+}
31+TEST(TPutAsyncUrma, Vec_Uint8Small)
32+{
33+ SKIP_IF_RANKS_LT(2);
34+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<uint8_t, 512>(2, 2, 0, 0)));
35+}
36+ 
37+// ============================================================================
38+// Boundary scenarios
39+// ============================================================================
40+TEST(TPutAsyncUrma, Vec_Uint8_SingleChunk)
41+{
42+ SKIP_IF_RANKS_LT(2);
43+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<uint8_t, 64>(2, 2, 0, 0)));
44+}
45+TEST(TPutAsyncUrma, Vec_Float_ExactChunk)
46+{
47+ SKIP_IF_RANKS_LT(2);
48+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<float, 64>(2, 2, 0, 0)));
49+}
50+ 
51+// ============================================================================
52+// Multi-rank broadcast
53+// ============================================================================
54+TEST(TPutAsyncUrma, Vec_FloatSmall_4Ranks)
55+{
56+ SKIP_IF_RANKS_LT(4);
57+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<float, 256>(4, 4, 0, 0)));
58+}
59+ 
60+// ============================================================================
61+// Large MR boundary tests
62+// ============================================================================
63+TEST(TPutAsyncUrma, Vec_Float_MR_6MB)
64+{
65+ SKIP_IF_RANKS_LT(2);
66+ // 512K floats → commBytesNeeded ≈ 4MB+256B, allocSize = 6MB (>2MB)
67+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<float, 524288>(2, 2, 0, 0)));
68+}
69+TEST(TPutAsyncUrma, Vec_Int32_MR_Over512MB)
70+{
71+ SKIP_IF_RANKS_LT(2);
72+ // 64M int32 → commBytesNeeded ≈ 512MB+256B, allocSize ≈ 514MB (>512MB)
73+ ASSERT_TRUE((RunPutAsyncUrmaRootPut<int32_t, 67108864>(2, 2, 0, 0)));
74+}
75+ 
76+int main(int argc, char **argv)
77+{
78+ CommMpiInit(&argc, &argv);
79+ ::testing::InitGoogleTest(&argc, argv);
80+ int ret = RUN_ALL_TESTS();
81+ CommMpiFinalize();
82+ return ret;
83+}
Atests/cpu/comm/st/testcase/tput_async_urma/tput_async_urma_kernel.cpp+162-0
@@ -0,0 +1,162 @@
1+/**
2+Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+ 
11+#include <cstddef>
12+#include <cstdint>
13+#include <iostream>
14+#include <pto/pto-inst.hpp>
15+ 
16+#define PTO_URMA_SUPPORTED
17+ 
18+#include "pto/pto-inst.hpp"
19+#include "pto/common/pto_tile.hpp"
20+#include "pto/common/cpu_stub.hpp"
21+#include "../common.hpp"
22+#include "../urma_context.hpp"
23+ 
24+// ============================================================================
25+// TPUT_ASYNC via URMA — device kernel.
26+// ============================================================================
27+ 
28+template <typename T, size_t count>
29+__global__ AICORE void TPutAsyncUrmaKernelImpl(__gm__ T *localBuf, int nranks, int my_rank, int first_rank_id,
30+ int root_rank, int elem_offset, int elem_count,
31+ __gm__ uint8_t *urmaWorkspace)
32+{
33+ using ShapeDyn = pto::Shape<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC>;
34+ using StrideDyn = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC>;
35+ using Global = pto::GlobalTensor<T, ShapeDyn, StrideDyn, pto::Layout::ND>;
36+ 
37+ if (elem_count <= 0 || elem_offset < 0 || elem_offset + elem_count > static_cast<int>(count)) {
38+ pipe_barrier(PIPE_ALL);
39+ return;
40+ }
41+ 
42+ ShapeDyn shape(1, 1, 1, 1, elem_count);
43+ StrideDyn stride(elem_count, elem_count, elem_count, elem_count, 1);
44+ 
45+ constexpr size_t kDataOffset = 64 * sizeof(int32_t);
46+ 
47+ __gm__ T *sendBuf = reinterpret_cast<__gm__ T *>(reinterpret_cast<__gm__ uint8_t *>(localBuf) + kDataOffset);
48+ __gm__ T *sendBufCore = sendBuf + elem_offset;
49+ Global sendG(sendBufCore, shape, stride);
50+ 
51+ pipe_barrier(PIPE_ALL);
52+ 
53+ if (my_rank == root_rank) {
54+#ifdef PTO_URMA_SUPPORTED
55+ const int my_peer = my_rank - first_rank_id;
56+ for (int target_peer = 0; target_peer < nranks; ++target_peer) {
57+ if (target_peer == my_peer) {
58+ continue;
59+ }
60+ uint64_t peerBase = UrmaPeerMrBaseAddr(urmaWorkspace, static_cast<uint32_t>(target_peer));
61+ __gm__ T *remoteRecvBuf = reinterpret_cast<__gm__ T *>(peerBase + kDataOffset) + count + elem_offset;
62+ Global remoteRecvG(remoteRecvBuf, shape, stride);
63+ 
64+ pto::comm::AsyncSession session;
65+ BuildAsyncSession<pto::comm::DmaEngine::URMA>(urmaWorkspace, static_cast<uint32_t>(target_peer), session);
66+ auto event = pto::comm::TPUT_ASYNC<pto::comm::DmaEngine::URMA>(remoteRecvG, sendG, session);
67+ event.Wait(session);
68+ }
69+#endif
70+ }
71+ 
72+ pipe_barrier(PIPE_ALL);
73+}
74+ 
75+// ============================================================================
76+// Host-side runner.
77+// ============================================================================
78+template <typename T, size_t count>
79+bool RunPutAsyncUrmaRootPutKernel(int rank_id, int n_ranks, int n_devices, int first_device_id, int first_rank_id,
80+ int root_rank)
81+{
82+ size_t commBytesNeeded = 64 * sizeof(int32_t) + 2 * count * sizeof(T);
83+ 
84+ UrmaTestContext ctx;
85+ if (!ctx.Setup(rank_id, n_ranks, n_devices, first_device_id, root_rank, commBytesNeeded)) {
86+ return false;
87+ }
88+ 
89+ uint8_t *input_host = nullptr;
90+ uint8_t *output_host = nullptr;
91+ aclrtMallocHost(reinterpret_cast<void **>(&input_host), count * sizeof(T));
92+ aclrtMallocHost(reinterpret_cast<void **>(&output_host), count * sizeof(T));
93+ if (!input_host || !output_host) {
94+ std::cerr << "[ERROR] aclrtMallocHost failed!" << std::endl;
95+ ctx.Cleanup();
96+ return false;
97+ }
98+ 
99+ for (size_t i = 0; i < count; ++i) {
100+ reinterpret_cast<T *>(input_host)[i] = static_cast<T>(i + rank_id * 10000);
101+ reinterpret_cast<T *>(output_host)[i] = static_cast<T>(-1);
102+ }
103+ 
104+ constexpr size_t kDataOffset = 64 * sizeof(int32_t);
105+ uint8_t *commBytes = reinterpret_cast<uint8_t *>(ctx.devBuf);
106+ T *sendBuf = reinterpret_cast<T *>(commBytes + kDataOffset);
107+ T *recvBuf = sendBuf + count;
108+ 
109+ aclrtMemcpy(sendBuf, count * sizeof(T), input_host, count * sizeof(T), ACL_MEMCPY_HOST_TO_DEVICE);
110+ aclrtMemcpy(recvBuf, count * sizeof(T), output_host, count * sizeof(T), ACL_MEMCPY_HOST_TO_DEVICE);
111+ 
112+ CommMpiBarrier();
113+ 
114+ TPutAsyncUrmaKernelImpl<T, count>(reinterpret_cast<T *>(ctx.devBuf), n_ranks, rank_id, first_rank_id, root_rank, 0,
115+ static_cast<int>(count),
116+ reinterpret_cast<uint8_t *>(ctx.urmaMgr.GetWorkspaceAddr()));
117+ int syncRet = aclrtSynchronizeStream(ctx.stream);
118+ 
119+ CommMpiBarrier();
120+ 
121+ aclrtMemcpy(output_host, count * sizeof(T), recvBuf, count * sizeof(T), ACL_MEMCPY_DEVICE_TO_HOST);
122+ 
123+ bool is_ok = true;
124+ if (rank_id != root_rank) {
125+ for (size_t i = 0; i < count; ++i) {
126+ T value = reinterpret_cast<T *>(output_host)[i];
127+ T expected = static_cast<T>(i + root_rank * 10000);
128+ if (value != expected) {
129+ std::cerr << "Rank " << rank_id << " Device " << ctx.deviceId << " SyncRet " << syncRet
130+ << " Expected: " << (float)expected << " Actual: " << (float)value << std::endl;
131+ is_ok = false;
132+ break;
133+ }
134+ }
135+ }
136+ 
137+ aclrtFreeHost(input_host);
138+ aclrtFreeHost(output_host);
139+ ctx.Cleanup();
140+ 
141+ return is_ok;
142+}
143+ 
144+// ============================================================================
145+// MPI-based multi-rank launch.
146+// ============================================================================
147+template <typename T, size_t count>
148+bool RunPutAsyncUrmaRootPut(int n_ranks, int n_devices, int first_rank_id, int first_device_id)
149+{
150+ return RunUrmaTestMpiLaunch(n_ranks, n_devices, first_rank_id, first_device_id,
151+ RunPutAsyncUrmaRootPutKernel<T, count>);
152+}
153+ 
154+// Explicit instantiations
155+template bool RunPutAsyncUrmaRootPut<float, 256>(int, int, int, int);
156+template bool RunPutAsyncUrmaRootPut<int32_t, 4096>(int, int, int, int);
157+template bool RunPutAsyncUrmaRootPut<uint8_t, 512>(int, int, int, int);
158+template bool RunPutAsyncUrmaRootPut<uint8_t, 64>(int, int, int, int);
159+template bool RunPutAsyncUrmaRootPut<float, 64>(int, int, int, int);
160+template bool RunPutAsyncUrmaRootPut<float, 524288>(int, int, int,
161+ int); // MR = 6MB (>2MB)
162+template bool RunPutAsyncUrmaRootPut<int32_t, 67108864>(int, int, int, int); // MR ≈ 514MB (>512MB)
Atests/cpu/comm/st/testcase/tput_async_urma/tput_async_urma_kernel.h+17-0
@@ -0,0 +1,17 @@
1+/**
2+Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+ 
11+#pragma once
12+ 
13+#include <cstddef>
14+#include <cstdint>
15+ 
16+template <typename T, size_t count>
17+bool RunPutAsyncUrmaRootPut(int n_ranks, int n_devices, int first_rank_id, int first_device_id);
Atests/cpu/comm/st/testcase/urma_context.hpp+89-0
@@ -0,0 +1,89 @@
1+/**
2+Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.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+#include "pto/comm/comm_types.hpp"
11+#include <cstdint>
12+#include <cstdlib>
13+#include <cstring>
14+ 
15+#include "pto/common/cpu_stub.hpp"
16+ 
17+class UrmaWorkspaceManager {
18+public:
19+ UrmaWorkspaceManager() = default;
20+ ~UrmaWorkspaceManager()
21+ {
22+ Finalize();
23+ }
24+ 
25+ void Finalize()
26+ {}
27+ 
28+ void *GetWorkspaceAddr() const
29+ {
30+ return urmaInfoDevice_;
31+ }
32+ 
33+ void *urmaInfoDevice_{nullptr};
34+};
35+ 
36+struct UrmaTestContext {
37+ int deviceId{-1};
38+ void *devBuf{nullptr};
39+ UrmaWorkspaceManager urmaMgr;
40+ 
41+ bool AllocHugePageBuffer(size_t commBytesNeeded)
42+ {
43+ return true;
44+ }
45+ 
46+ bool Setup(int rank_id, int n_ranks, int n_devices, int first_device_id, int root_rank, size_t commBytesNeeded)
47+ {
48+ if (commBytesNeeded > 0) {
49+ devBuf = malloc(commBytesNeeded);
50+ urmaMgr.urmaInfoDevice_ = devBuf;
51+ return true;
52+ }
53+ return false;
54+ }
55+ 
56+ void Cleanup()
57+ {
58+ urmaMgr.Finalize();
59+ if (devBuf) {
60+ aclrtFree(devBuf);
61+ devBuf = nullptr;
62+ }
63+ }
64+};
65+ 
66+AICORE inline uint64_t UrmaPeerMrBaseAddr(__gm__ uint8_t *urmaWorkspace, uint32_t peerRank)
67+{
68+ return reinterpret_cast<uint64_t>(urmaWorkspace);
69+}
70+ 
71+template <pto::comm::DmaEngine engine>
72+PTO_INTERNAL bool BuildAsyncSession(__gm__ uint8_t *workspace, uint32_t destRankId, pto::comm::AsyncSession &session)
73+{
74+ static_assert(engine == pto::comm::DmaEngine::URMA, "This overload is for URMA only");
75+ return true;
76+}
77+ 
78+using UrmaKernelFn = bool (*)(int, int, int, int, int, int);
79+ 
80+inline bool RunUrmaTestMpiLaunch(int n_ranks, int n_devices, int first_rank_id, int first_device_id,
81+ UrmaKernelFn kernelFn)
82+{
83+ int root = 0;
84+ bool res = false;
85+ for (size_t i = 0; i < n_ranks; i++) {
86+ res = res || kernelFn(i, n_ranks, n_devices, first_device_id, first_rank_id, root);
87+ }
88+ return res;
89+}