已合并
增加mc2单卡测试mock server的能力 #3368
YuZhengzhong创建于 3月26日
增加mc2单卡测试mock server的能力 #3368
已合并
YuZhengzhong创建于 3月26日
4 个文件变更+1660-0
@@ -0,0 +1,224 @@
1+# Mock HCCL Server Framework
2+ 
3+## 概览
4+ 
5+本框架提供了一个**通用的 HCCL Server 端模拟实现**,用于在单卡环境下替代真实的 CCU Server / HCCL 通信后端。它与具体的 MC2 算子无关 — 任何通过 workspace 消息协议与 HCCL Server 交互的算子(如 AllGatherMatmul、MatmulReduceScatter 等)都可以使用本框架进行单卡测试。
6+ 
7+### 背景:hcclClient / Server 架构
8+ 
9+在真实的多卡环境中,MC2 类算子的通信流程如下:
10+ 
11+```
12+Device (kernel) CCU (HCCL Server)
13+ │ │
14+ │ 1. 填 sendMsgs[slot] │
15+ │ 2. 写 commitTurnCnt[slot] │
16+ │ ──────── workspace ──────────→ │
17+ │ │ 3. 读 commitTurnCnt, 发现有效
18+ │ │ 4. 读 sendMsgs, 执行集合通信
19+ │ │ 5. 写 finishedTurnCnt[slot]
20+ │ ←──────── workspace ───────── │
21+ │ 6. 轮询 finishedTurnCnt, 继续 │
22+ │ │
23+```
24+ 
25+kernel 通过 workspace 中的消息协议向 Server 提交通信请求,Server 执行实际的跨 rank 数据搬移,完成后通过 workspace 回写完成标志。这个 **workspace 消息协议是算子无关的** — 不同的 MC2 算子(AllGather、ReduceScatter 等)使用相同的协议格式,只是 commType 字段不同。
26+ 
27+**本框架就是在 host 侧用一个轮询线程模拟这个 Server 的行为。**
28+ 
29+```
30+目录结构:
31+tests/mc2_mock_test_frame/ # 本目录 — 通用 mock 框架
32+├── mock_framework.h # C++ 核心 (MockContextBuilder + MultiRankMockContext + MockHcclServer)
33+├── mock_framework_test.cpp # C++ 单元测试 (协议交互验证)
34+├── mock_framework.cpp # pybind11 torch extension (Python binding)
35+└── README.md # 本文档
36+ 
37+tests/torch_extension_tests/mc2/all_gather_matmul_v3/ # V3 算子测试
38+├── conftest.py # mock 测试公共工具
39+├── test_all_gather_matmul_v3_mock.py # 单 rank mock 测试
40+├── test_all_gather_matmul_v3_mock_multirank.py # 多 rank 多流 mock 测试
41+└── test_all_gather_matmul_v3.py # 真实多卡测试 (torchrun)
42+```
43+ 
44+---
45+ 
46+## 1. Workspace 消息协议
47+ 
48+MockHcclServer 模拟的核心是 workspace 消息协议。这是 HCCL 定义的 kernel ↔ Server 通信接口,与具体算子无关。
49+ 
50+### 1.1 Workspace 布局
51+ 
52+workspace 需要 512B 对齐,包含 4 个区域:
53+ 
54+```
55+Workspace (512B 对齐):
56+┌──────────────────────────────────────┐
57+│ sendMsgs[64] (each 112B) │ +0x0000 kernel → server 的通信请求
58+├──────────────────────────────────────┤
59+│ recvMsgs[64] (each 112B) │ +0x1C00 server → kernel 的响应 (预留)
60+├──────────────────────────────────────┤
61+│ commitTurnCnt[64] (each 64B) │ +0x5800 kernel 提交标志
62+├──────────────────────────────────────┤
63+│ finishedTurnCnt[64](each 64B) │ +0x6800 server 完成标志
64+└──────────────────────────────────────┘
65+```
66+ 
67+### 1.2 消息结构
68+ 
69+**HcclMsg (112 bytes)** — kernel 填写的通信请求:
70+ 
71+| 偏移 | 字段 | 说明 |
72+|------|------|------|
73+| +0x00 | commType (u32) | 通信类型: AllGather=6, ReduceScatter=7 |
74+| +0x04 | opType (u32) | reduce 操作类型 |
75+| +0x08 | sendBuffer (u64) | device 源地址 |
76+| +0x10 | recvBuffer (u64) | device 目标地址 |
77+| +0x18 | dataCnt (u64) | 元素数 |
78+| +0x20 | strideCount (u64) | recvBuf 中 rank 间步长 |
79+| +0x28 | msgValid (u32) | HCCL_MSG_VALID_MASK (0x5CDF123A) |
80+| +0x2C | hcclDataType (u32) | 数据类型: FP32=0, FP16=1, BF16=5, ... |
81+| +0x30 | rest[64] | 剩余字段 |
82+ 
83+**TurnCnt (64 bytes)** — 提交/完成计数器:
84+ 
85+| 偏移 | 字段 | 说明 |
86+|------|------|------|
87+| +0x00 | valid (u64) | COMMIT_VALID_MASK (987654321) 表示有效 |
88+| +0x08 | cnt (u64) | 提交/完成计数 |
89+| +0x10 | reserved[6] | 填充至 cache line 对齐 |
90+ 
91+### 1.3 协议流程
92+ 
93+**常规通信(AllGather/ReduceScatter):**
94+ 
95+```
96+kernel:
97+ 1. 填写 sendMsgs[slot] (commType, sendBuf, recvBuf, dataCnt, ...)
98+ 2. 写 commitTurnCnt[slot].valid = COMMIT_VALID_MASK
99+ 
100+server (MockHcclServer):
101+ 3. 轮询 commitTurnCnt, 发现 valid → 读 sendMsgs[slot]
102+ 4. 根据 commType 执行操作 (单卡: D2D memcpy sendBuf → recvBuf)
103+ 5. 写 finishedTurnCnt[slot].cnt = commitCnt
104+ 6. 清除 commitTurnCnt[slot].valid
105+ 
106+kernel:
107+ 7. 轮询 finishedTurnCnt[slot].cnt >= 期望值 → 继续
108+```
109+ 
110+**Finalize(通信结束握手):**
111+ 
112+```
113+kernel Finalize():
114+ 1. 提交 finalize commit (最后一条消息)
115+ 
116+server:
117+ 2. 检测到 finalize commit
118+ 3. 写 finishedTurnCnt[slot].cnt = FINALIZE_FINISH_CNT (1234567899999999999)
119+ 
120+kernel:
121+ 4. 检测到 FINALIZE_FINISH_CNT → 通信流程结束
122+```
123+ 
124+---
125+ 
126+## 2. 框架组件
127+ 
128+### 2.1 MockHcclServer — 通用 HCCL Server 模拟
129+ 
130+**核心职责:** 在 host 端启动一个轮询线程,监听 workspace 中的通信请求,并使用传入的各 rank 输入 tensor 模拟真实的集合通信语义。
131+ 
132+构造时需传入每个 rank 的输入 tensor(device 指针),server 在处理通信请求时读取这些 tensor,模拟从远端 rank 获取数据。
133+ 
134+**支持的通信类型:**
135+ 
136+| commType | 操作 | Mock 行为 |
137+|----------|------|-----------|
138+| 6 (AllGather) | 各 rank chunk 拼接 | rankInputs[r] → recvBuf[r * stride],本卡用 sendBuf |
139+| 7 (ReduceScatter) | reduce + scatter | D2H 读各 rank → host 端 element-wise reduce → 取 chunk[localRank] H2D 写回 |
140+| 2 (AllReduce) | reduce 全量 | D2H 读各 rank → host 端 element-wise reduce → 全量 H2D 写回 |
141+| 12 (AlltoAll) | 块重排 | 各 rank 的 block[localRank] → recvBuf[r] |
142+ 
143+**Reduce 操作:** 支持 SUM(0) / PROD(1) / MAX(2) / MIN(3),host 侧通过 FP16/BF16 ↔ float 转换后计算。
144+ 
145+### 2.2 MockContextBuilder — 单 rank HCCL Context 构造
146+ 
147+在 device 侧构造 kernel 需要的 `HcclA2CombineOpParam` 结构(即 commContext)。所有 rank 的 `windowsIn[i]` 指向同一块 device 内存,使得 910B AIV 模式的 flag 同步机制自洽。
148+ 
149+### 2.3 MultiRankMockContext — 多 rank Context 构造
150+ 
151+分配 N 个独立 window(184MB 每个)+ N 个 context,context 中 windowsIn[] 互相交叉引用。配合多 stream 并发 kernel 实现真实多 rank 通信。
152+ 
153+### 2.4 Python Binding (mock_framework.cpp)
154+ 
155+通过 pybind11 暴露给 Python:
156+ 
157+```python
158+import mock_hccl_ext
159+ 
160+# 单 rank 测试
161+ctx = mock_hccl_ext.MockContext(rank_num=2, rank_id=0, device_id=0)
162+ctx.build()
163+server = mock_hccl_ext.MockServer(workspace_ptr=ctx.workspace_ptr(),
164+ rank_inputs=[t0, t1], local_rank_id=0)
165+server.start()
166+# ... 调用算子 ...
167+server.wait_for_finalize(slot=0, timeout_ms=10000)
168+server.stop()
169+ 
170+# 多 rank 测试
171+mctx = mock_hccl_ext.MultiRankContext(rank_num=2, device_id=0)
172+mctx.build()
173+ctx_r0 = mctx.context_tensor(0)
174+ctx_r1 = mctx.context_tensor(1)
175+# Launch on separate streams...
176+```
177+ 
178+---
179+ 
180+## 3. 单元测试
181+ 
182+`mock_framework_test.cpp` 验证 mock server 的协议正确性(与具体算子无关):
183+ 
184+| Test | 验证点 |
185+|------|--------|
186+| TestContextBuilder | context H2D/D2H,字段 readback 正确 |
187+| TestServerProtocol | 手写 commitTurnCnt → server 响应 finishedTurnCnt |
188+| TestServerDataCopy | AllGather msg → server 执行 D2D copy → recvBuf 内容正确 |
189+| TestFinalizeProtocol | regular msg + finalize commit → FINALIZE_FINISH_CNT 响应 |
190+ 
191+编译运行:
192+```bash
193+ASCEND_HOME=~/Ascend/ascend-toolkit/latest
194+g++ -std=c++17 -O2 -I${ASCEND_HOME}/include \
195+ mock_framework_test.cpp \
196+ -L${ASCEND_HOME}/lib64 -lascendcl -lpthread \
197+ -Wl,-rpath,${ASCEND_HOME}/lib64 \
198+ -o /tmp/mock_framework_test
199+/tmp/mock_framework_test
200+```
201+ 
202+---
203+ 
204+## 4. 如何为新算子添加测试
205+ 
206+本框架与算子无关,为新的 MC2 算子添加测试只需:
207+ 
208+1. **构造 mock context**`MockContextBuilder.Build(rankNum, rankId)`
209+2. **启动 mock server**`MockHcclServer(workspace).Start()`
210+3. **调用算子** — 将 mock context 作为 commContext 输入传给算子
211+4. **等待完成**`WaitForFinalize()` 等待 kernel 的 Finalize 握手
212+5. **验证结果** — 检查算子的计算输出是否正确
213+ 
214+---
215+ 
216+## 5. 已知限制
217+ 
218+| 限制 | 说明 |
219+|------|------|
220+| 单卡 only | 多卡需要真实 HCCL 或多 device mock |
221+| D2D memcpy 代替集合通信 | 单卡无跨 rank 数据,AllGather/ReduceScatter 退化为 memcpy |
222+| Mock context 无 RDMA | IbVerbsData / aiRMAInfo 为 0,不支持跨节点场景 |
223+| Server 轮询延迟 | host D2H/H2D 轮询有 ~50us 级延迟,不反映真实通信时序 |
224+| Flag 自解不验证时序 | 单卡 mock 的 flag 写后立即可读,无法验证多卡竞争条件 |
@@ -0,0 +1,262 @@
1+/**
2+ * Mock HCCL Server — PyTorch C++ Extension (pybind11)
3+ *
4+ * Exposes MockContextBuilder and MockHcclServer to Python.
5+ * The MockServer accepts per-rank input tensors to simulate
6+ * real collective communication semantics.
7+ *
8+ * Usage:
9+ * import mock_hccl_ext
10+ * ctx = mock_hccl_ext.MockContext(rank_num=2, rank_id=0, device_id=0)
11+ * ctx.build()
12+ *
13+ * # rank_inputs: list of npu tensors, one per rank (including local rank)
14+ * server = mock_hccl_ext.MockServer(
15+ * workspace_ptr=ctx.workspace_ptr(),
16+ * rank_inputs=[tensor_rank0, tensor_rank1],
17+ * local_rank_id=0,
18+ * device_id=0)
19+ * server.start()
20+ * ...
21+ * server.stop()
22+ */
23+ 
24+#include <torch/extension.h>
25+#include <pybind11/pybind11.h>
26+#include <pybind11/stl.h>
27+ 
28+#include "mock_framework.h"
29+ 
30+namespace py = pybind11;
31+using namespace mock_hccl;
32+ 
33+// ============================================================
34+// MockContext — Python-facing context builder
35+// ============================================================
36+class MockContext {
37+public:
38+ MockContext(uint32_t rankNum, uint32_t rankId, int deviceId)
39+ : rankNum_(rankNum), rankId_(rankId), deviceId_(deviceId) {}
40+ 
41+ ~MockContext() { destroy(); }
42+ 
43+ bool build() {
44+ devWindowMem_ = DevMalloc(WINDOW_TOTAL_SIZE);
45+ if (!devWindowMem_) return false;
46+ 
47+ size_t wsAllocSize = MIN_WORKSPACE_SIZE + 512;
48+ devWorkspaceRaw_ = DevMalloc(wsAllocSize);
49+ if (!devWorkspaceRaw_) return false;
50+ devWorkspaceAligned_ = AlignUp512(devWorkspaceRaw_);
51+ 
52+ MockHcclContext hostCtx;
53+ memset(&hostCtx, 0, sizeof(hostCtx));
54+ hostCtx.rankId = rankId_;
55+ hostCtx.rankNum = rankNum_;
56+ hostCtx.winSize = WINDOW_TOTAL_SIZE;
57+ hostCtx.workSpace = reinterpret_cast<uint64_t>(devWorkspaceAligned_);
58+ hostCtx.workSpaceSize = MIN_WORKSPACE_SIZE;
59+ for (uint32_t i = 0; i < rankNum_; i++) {
60+ hostCtx.windowsIn[i] = reinterpret_cast<uint64_t>(devWindowMem_);
61+ hostCtx.windowsOut[i] = reinterpret_cast<uint64_t>(devWindowMem_);
62+ }
63+ 
64+ devContext_ = DevMalloc(sizeof(MockHcclContext));
65+ if (!devContext_) return false;
66+ aclrtMemcpy(devContext_, sizeof(hostCtx), &hostCtx, sizeof(hostCtx),
67+ ACL_MEMCPY_HOST_TO_DEVICE);
68+ return true;
69+ }
70+ 
71+ void destroy() {
72+ DevFree(devContext_); devContext_ = nullptr;
73+ DevFree(devWorkspaceRaw_); devWorkspaceRaw_ = nullptr;
74+ DevFree(devWindowMem_); devWindowMem_ = nullptr;
75+ devWorkspaceAligned_ = nullptr;
76+ }
77+ 
78+ void clear_flags() {
79+ if (!devWindowMem_) return;
80+ void* flagArea = OffsetPtr(devWindowMem_, FLAG_OFFSET_BYTES);
81+ aclrtMemset(flagArea, FLAG_AREA_SIZE, 0, FLAG_AREA_SIZE);
82+ }
83+ 
84+ at::Tensor as_tensor() {
85+ if (!devContext_) throw std::runtime_error("Context not built");
86+ auto opts = at::TensorOptions().dtype(at::kChar).device(at::kPrivateUse1, deviceId_);
87+ size_t nbytes = sizeof(MockHcclContext);
88+ return at::from_blob(devContext_, {static_cast<int64_t>(nbytes)}, opts);
89+ }
90+ 
91+ int64_t context_ptr() const { return reinterpret_cast<int64_t>(devContext_); }
92+ int64_t workspace_ptr() const { return reinterpret_cast<int64_t>(devWorkspaceAligned_); }
93+ int64_t window_ptr() const { return reinterpret_cast<int64_t>(devWindowMem_); }
94+ uint32_t rank_num() const { return rankNum_; }
95+ 
96+private:
97+ uint32_t rankNum_, rankId_;
98+ int deviceId_;
99+ void* devContext_{nullptr};
100+ void* devWorkspaceRaw_{nullptr};
101+ void* devWorkspaceAligned_{nullptr};
102+ void* devWindowMem_{nullptr};
103+};
104+ 
105+// ============================================================
106+// PyMockServer — Python-facing HCCL server wrapper
107+//
108+// Accepts a list of torch.Tensor (one per rank) and delegates
109+// to MockHcclServer for real collective simulation.
110+// ============================================================
111+class PyMockServer {
112+public:
113+ /**
114+ * @param workspacePtr int64 device address of 512B-aligned workspace
115+ * @param rankInputs list of npu tensors, one per rank
116+ * @param localRankId rank ID of the kernel under test
117+ * @param deviceId NPU device ID
118+ */
119+ PyMockServer(int64_t workspacePtr,
120+ std::vector<at::Tensor> rankInputs,
121+ uint32_t localRankId,
122+ int deviceId)
123+ : deviceId_(deviceId)
124+ {
125+ // Hold references to tensors so they stay alive
126+ tensorRefs_ = std::move(rankInputs);
127+ 
128+ // Extract device pointers → RankData
129+ std::vector<RankData> rankData;
130+ rankData.reserve(tensorRefs_.size());
131+ for (auto& t : tensorRefs_) {
132+ rankData.push_back({t.data_ptr(), (size_t)(t.nbytes())});
133+ }
134+ 
135+ server_ = std::make_unique<MockHcclServer>(
136+ reinterpret_cast<void*>(workspacePtr),
137+ std::move(rankData),
138+ localRankId,
139+ deviceId);
140+ }
141+ 
142+ ~PyMockServer() { stop(); }
143+ 
144+ void start() { server_->Start(); }
145+ void stop() { server_->Stop(); }
146+ 
147+ uint32_t msg_count() const { return server_->GetMsgCount(); }
148+ bool is_finalized() const { return server_->IsFinalized(); }
149+ 
150+ bool wait_for_finalize(uint32_t slot, uint32_t timeoutMs) {
151+ return server_->WaitForFinalize(slot, timeoutMs);
152+ }
153+ 
154+private:
155+ int deviceId_;
156+ std::vector<at::Tensor> tensorRefs_; // prevent GC
157+ std::unique_ptr<MockHcclServer> server_;
158+};
159+ 
160+// ============================================================
161+// PyMultiRankContext — Python-facing multi-rank context builder
162+//
163+// Allocates N independent windows + N contexts with cross-referenced
164+// windowsIn[]. For use with multi-stream concurrent kernel launches.
165+// ============================================================
166+class PyMultiRankContext {
167+public:
168+ PyMultiRankContext(uint32_t rankNum, int deviceId)
169+ : rankNum_(rankNum), deviceId_(deviceId) {}
170+ 
171+ ~PyMultiRankContext() { destroy(); }
172+ 
173+ bool build() { return ctx_.Build(rankNum_); }
174+ void destroy() { ctx_.Destroy(); }
175+ void clear_flags() { ctx_.ClearAllFlags(); }
176+ 
177+ at::Tensor context_tensor(uint32_t rank) {
178+ void* addr = ctx_.GetContextAddr(rank);
179+ if (!addr) throw std::runtime_error("Context not built for rank " + std::to_string(rank));
180+ auto opts = at::TensorOptions().dtype(at::kChar).device(at::kPrivateUse1, deviceId_);
181+ return at::from_blob(addr, {static_cast<int64_t>(sizeof(MockHcclContext))}, opts);
182+ }
183+ 
184+ // D2H read flag values from rank's window.
185+ // Returns list of int32 flag values at FLAG_OFFSET + [0..num_flags-1]
186+ std::vector<int32_t> read_flags(uint32_t rank, uint32_t num_flags = 8) {
187+ std::vector<int32_t> result(num_flags, -1);
188+ void* win = ctx_.GetWindowMem(rank);
189+ if (!win) return result;
190+ // flags are at (int32_t*)win + FLAG_OFFSET, where FLAG_OFFSET = 180*1024*1024/4
191+ constexpr int64_t FLAG_OFFSET = 180LL * 1024 * 1024 / sizeof(int32_t);
192+ void* flagAddr = reinterpret_cast<void*>(
193+ reinterpret_cast<int32_t*>(win) + FLAG_OFFSET);
194+ size_t bytes = num_flags * sizeof(int32_t);
195+ aclrtMemcpy(result.data(), bytes, flagAddr, bytes, ACL_MEMCPY_DEVICE_TO_HOST);
196+ return result;
197+ }
198+ 
199+ // D2H read first N bytes from rank's window data area (for debugging)
200+ std::vector<uint8_t> read_window_data(uint32_t rank, size_t offset, size_t nbytes) {
201+ std::vector<uint8_t> result(nbytes, 0);
202+ void* win = ctx_.GetWindowMem(rank);
203+ if (!win) return result;
204+ void* src = OffsetPtr(win, offset);
205+ aclrtMemcpy(result.data(), nbytes, src, nbytes, ACL_MEMCPY_DEVICE_TO_HOST);
206+ return result;
207+ }
208+ 
209+ uint32_t rank_num() const { return rankNum_; }
210+ 
211+private:
212+ uint32_t rankNum_;
213+ int deviceId_;
214+ MultiRankMockContext ctx_;
215+};
216+ 
217+// ============================================================
218+// pybind11 module
219+// ============================================================
220+PYBIND11_MODULE(mock_hccl_ext, m) {
221+ m.doc() = "Mock HCCL server extension for single-rank MC2 testing";
222+ 
223+ py::class_<MockContext>(m, "MockContext")
224+ .def(py::init<uint32_t, uint32_t, int>(),
225+ py::arg("rank_num") = 1, py::arg("rank_id") = 0, py::arg("device_id") = 0)
226+ .def("build", &MockContext::build)
227+ .def("destroy", &MockContext::destroy)
228+ .def("as_tensor", &MockContext::as_tensor)
229+ .def("clear_flags", &MockContext::clear_flags)
230+ .def("context_ptr", &MockContext::context_ptr)
231+ .def("workspace_ptr", &MockContext::workspace_ptr)
232+ .def("window_ptr", &MockContext::window_ptr)
233+ .def("rank_num", &MockContext::rank_num);
234+ 
235+ py::class_<PyMockServer>(m, "MockServer")
236+ .def(py::init<int64_t, std::vector<at::Tensor>, uint32_t, int>(),
237+ py::arg("workspace_ptr"),
238+ py::arg("rank_inputs"),
239+ py::arg("local_rank_id") = 0,
240+ py::arg("device_id") = 0)
241+ .def("start", &PyMockServer::start)
242+ .def("stop", &PyMockServer::stop)
243+ .def("msg_count", &PyMockServer::msg_count)
244+ .def("is_finalized", &PyMockServer::is_finalized)
245+ .def("wait_for_finalize", &PyMockServer::wait_for_finalize,
246+ py::arg("slot") = 0, py::arg("timeout_ms") = 5000);
247+ 
248+ py::class_<PyMultiRankContext>(m, "MultiRankContext")
249+ .def(py::init<uint32_t, int>(),
250+ py::arg("rank_num"), py::arg("device_id") = 0)
251+ .def("build", &PyMultiRankContext::build)
252+ .def("destroy", &PyMultiRankContext::destroy)
253+ .def("clear_flags", &PyMultiRankContext::clear_flags)
254+ .def("context_tensor", &PyMultiRankContext::context_tensor,
255+ py::arg("rank"))
256+ .def("read_flags", &PyMultiRankContext::read_flags,
257+ py::arg("rank"), py::arg("num_flags") = 8)
258+ .def("read_window_data", &PyMultiRankContext::read_window_data,
259+ py::arg("rank"), py::arg("offset") = 0, py::arg("nbytes") = 64)
260+ .def("rank_num", &PyMultiRankContext::rank_num);
261+ 
262+}
@@ -0,0 +1,826 @@
1+/**
2+ * Mock HCCL Server Framework — Generic hcclClient Server-Side Simulator
3+ *
4+ * Simulates the CCU server (HCCL Server) on the host CPU side for single-card
5+ * testing. This framework is operator-agnostic — any MC2 operator that
6+ * communicates via the workspace message protocol can use it.
7+ *
8+ * Components:
9+ * 1. MockContextBuilder — constructs device-side HcclCombinOpParam (commContext)
10+ * 2. MockHcclServer — host polling thread that responds to kernel's
11+ * communication requests with real collective semantics
12+ *
13+ * Communication simulation:
14+ * The server accepts per-rank input tensors at construction. When the kernel
15+ * issues a collective operation, the server reads these tensors as if they
16+ * came from remote ranks, and performs the real collective semantics:
17+ * - AllGather: gather each rank's chunk into recvBuf
18+ * - ReduceScatter: element-wise reduce all ranks, scatter chunks
19+ * - AllReduce: element-wise reduce all ranks, full result to recvBuf
20+ * - AlltoAll: block redistribution
21+ */
22+ 
23+#ifndef MOCK_FRAMEWORK_H_
24+#define MOCK_FRAMEWORK_H_
25+ 
26+#include <cstdint>
27+#include <cstdio>
28+#include <cstring>
29+#include <cmath>
30+#include <algorithm>
31+#include <atomic>
32+#include <thread>
33+#include <chrono>
34+#include <functional>
35+#include <vector>
36+#include <unistd.h>
37+#include "acl/acl.h"
38+ 
39+namespace mock_hccl {
40+ 
41+// ============================================================
42+// HCCL Protocol Constants
43+// ============================================================
44+constexpr uint32_t COMMIT_VALID_MASK = 987654321U; // 0x3ADE68B1
45+constexpr uint64_t FINALIZE_FINISH_CNT = 1234567899999999999ULL;
46+constexpr uint32_t HCCL_MSG_VALID_MASK = 0x5CDF123AU;
47+constexpr uint32_t HCCL_MSG_CNT = 64;
48+ 
49+// HCCL command types
50+constexpr uint32_t HCCL_CMD_ALLREDUCE = 2;
51+constexpr uint32_t HCCL_CMD_ALLGATHER = 6;
52+constexpr uint32_t HCCL_CMD_REDUCE_SCATTER = 7;
53+constexpr uint32_t HCCL_CMD_ALLTOALL = 12; // verify against actual HCCL
54+constexpr uint32_t HCCL_CMD_FINALIZE = 100;
55+ 
56+// HCCL reduce operation types
57+constexpr uint32_t HCCL_REDUCE_SUM = 0;
58+constexpr uint32_t HCCL_REDUCE_PROD = 1;
59+constexpr uint32_t HCCL_REDUCE_MAX = 2;
60+constexpr uint32_t HCCL_REDUCE_MIN = 3;
61+ 
62+// ============================================================
63+// Workspace Layout Offsets (from 512B-aligned base)
64+// ============================================================
65+constexpr size_t SEND_MSGS_OFFSET = 0x0000; // sendMsgs[64], each 112 bytes
66+constexpr size_t RECV_MSGS_OFFSET = 0x1C00; // recvMsgs[64]
67+constexpr size_t COMMIT_TURNCNT_OFFSET = 0x5800; // commitTurnCnt[64], each 64 bytes
68+constexpr size_t FINISH_TURNCNT_OFFSET = 0x6800; // finishedTurnCnt[64], each 64 bytes
69+ 
70+constexpr size_t MSG_STRIDE = 112; // bytes per HcclMsg
71+constexpr size_t TURNCNT_STRIDE = 64; // bytes per TurnCnt (cache-line aligned)
72+ 
73+// Minimum workspace size (single queue mode)
74+constexpr size_t MIN_WORKSPACE_SIZE = 2 * 1024 * 1024; // 2MB
75+ 
76+// ============================================================
77+// Workspace Structures (host-side mirrors)
78+// ============================================================
79+ 
80+struct HcclMsg {
81+ uint32_t commType; // +0x00 HCCL_CMD_ALLGATHER=6, etc.
82+ uint32_t opType; // +0x04 reduce operation type (SUM=0, MAX=2, ...)
83+ uint64_t sendBuffer; // +0x08 device source address
84+ uint64_t recvBuffer; // +0x10 device destination address
85+ uint64_t dataCnt; // +0x18 element count
86+ uint64_t strideCount; // +0x20 stride between ranks in recvBuf (elements)
87+ uint32_t msgValid; // +0x28 HCCL_MSG_VALID_MASK when valid
88+ uint32_t hcclDataType; // +0x2C data type enum
89+ uint8_t rest[64]; // +0x30 remaining fields to fill 112 bytes total
90+};
91+static_assert(sizeof(HcclMsg) == MSG_STRIDE, "HcclMsg size mismatch");
92+ 
93+struct TurnCnt {
94+ uint64_t valid; // +0x00 COMMIT_VALID_MASK or 0
95+ uint64_t cnt; // +0x08 commit/finish count
96+ uint64_t reserved[6]; // +0x10 padding to 64 bytes
97+};
98+static_assert(sizeof(TurnCnt) == TURNCNT_STRIDE, "TurnCnt size mismatch");
99+ 
100+// ============================================================
101+// Context Structure (910B: HcclA2CombineOpParam)
102+// ============================================================
103+constexpr uint32_t MAX_RANK = 32;
104+constexpr int64_t FLAG_OFFSET_BYTES = 180LL * 1024 * 1024;
105+constexpr int64_t FLAG_AREA_SIZE = 4LL * 1024 * 1024;
106+constexpr int64_t WINDOW_TOTAL_SIZE = FLAG_OFFSET_BYTES + FLAG_AREA_SIZE;
107+ 
108+struct MockHcclContext {
109+ uint64_t workSpace; // +0
110+ uint64_t workSpaceSize; // +8
111+ uint32_t rankId; // +16
112+ uint32_t rankNum; // +20
113+ uint64_t winSize; // +24
114+ uint64_t windowsIn[MAX_RANK]; // +32
115+ uint64_t windowsOut[MAX_RANK]; // +288
116+ uint8_t res[8328]; // +544
117+ uint8_t multiFlag; // +8872
118+ uint64_t data; // +8873
119+ uint64_t dataSize; // +8881
120+ uint64_t sizeOfAiRMAInfo; // +8889
121+ uint64_t aiRMAInfo; // +8897
122+};
123+ 
124+// ============================================================
125+// HCCL Data Type → byte size mapping
126+// ============================================================
127+inline size_t HcclDataTypeSize(uint32_t hcclType) {
128+ switch (hcclType) {
129+ case 0: return 4; // FP32
130+ case 1: return 2; // FP16
131+ case 2: return 1; // INT8
132+ case 3: return 4; // INT32
133+ case 5: return 2; // BF16
134+ case 12: return 1; // FP8_E4M3
135+ case 13: return 1; // FP8_E5M2
136+ default: return 2;
137+ }
138+}
139+ 
140+// ============================================================
141+// FP16 / BF16 ↔ float conversion (host-side computation)
142+// ============================================================
143+inline float Fp16ToFloat(uint16_t h) {
144+ uint32_t sign = (h >> 15) & 1;
145+ uint32_t exp = (h >> 10) & 0x1f;
146+ uint32_t mant = h & 0x3ff;
147+ uint32_t f;
148+ if (exp == 0) {
149+ if (mant == 0) {
150+ f = sign << 31;
151+ } else {
152+ // denorm → norm
153+ exp = 1;
154+ while (!(mant & 0x400)) { mant <<= 1; exp--; }
155+ mant &= 0x3ff;
156+ f = (sign << 31) | ((exp + 112) << 23) | (mant << 13);
157+ }
158+ } else if (exp == 31) {
159+ f = (sign << 31) | 0x7f800000 | (mant << 13); // inf/nan
160+ } else {
161+ f = (sign << 31) | ((exp + 112) << 23) | (mant << 13);
162+ }
163+ float result;
164+ memcpy(&result, &f, 4);
165+ return result;
166+}
167+ 
168+inline uint16_t FloatToFp16(float val) {
169+ uint32_t u;
170+ memcpy(&u, &val, 4);
171+ uint32_t sign = (u >> 31) & 1;
172+ int32_t exp = ((u >> 23) & 0xff) - 127 + 15;
173+ uint32_t mant = (u >> 13) & 0x3ff;
174+ if (exp <= 0) return (uint16_t)(sign << 15);
175+ if (exp >= 31) return (uint16_t)((sign << 15) | 0x7c00);
176+ return (uint16_t)((sign << 15) | (exp << 10) | mant);
177+}
178+ 
179+inline float Bf16ToFloat(uint16_t bf) {
180+ uint32_t f = (uint32_t)bf << 16;
181+ float result;
182+ memcpy(&result, &f, 4);
183+ return result;
184+}
185+ 
186+inline uint16_t FloatToBf16(float val) {
187+ uint32_t u;
188+ memcpy(&u, &val, 4);
189+ return (uint16_t)(u >> 16);
190+}
191+ 
192+// ============================================================
193+// Host-side element read/write by hcclDataType
194+// ============================================================
195+inline float ReadElement(const void* buf, size_t idx, uint32_t hcclType) {
196+ switch (hcclType) {
197+ case 0: return ((const float*)buf)[idx]; // FP32
198+ case 1: return Fp16ToFloat(((const uint16_t*)buf)[idx]); // FP16
199+ case 3: return (float)((const int32_t*)buf)[idx]; // INT32
200+ case 5: return Bf16ToFloat(((const uint16_t*)buf)[idx]); // BF16
201+ default: return 0.0f;
202+ }
203+}
204+ 
205+inline void WriteElement(void* buf, size_t idx, float val, uint32_t hcclType) {
206+ switch (hcclType) {
207+ case 0: ((float*)buf)[idx] = val; break; // FP32
208+ case 1: ((uint16_t*)buf)[idx] = FloatToFp16(val); break; // FP16
209+ case 3: ((int32_t*)buf)[idx] = (int32_t)val; break; // INT32
210+ case 5: ((uint16_t*)buf)[idx] = FloatToBf16(val); break; // BF16
211+ default: break;
212+ }
213+}
214+ 
215+// ============================================================
216+// Host-side reduce: dst = reduce(srcs[0..N-1])
217+// All buffers are host memory, count = number of elements.
218+// ============================================================
219+inline void HostReduce(void* dst,
220+ const std::vector<const void*>& srcs,
221+ size_t count,
222+ uint32_t hcclDataType,
223+ uint32_t reduceOp) {
224+ for (size_t i = 0; i < count; i++) {
225+ float acc = ReadElement(srcs[0], i, hcclDataType);
226+ for (size_t r = 1; r < srcs.size(); r++) {
227+ float val = ReadElement(srcs[r], i, hcclDataType);
228+ switch (reduceOp) {
229+ case HCCL_REDUCE_SUM: acc += val; break;
230+ case HCCL_REDUCE_PROD: acc *= val; break;
231+ case HCCL_REDUCE_MAX: acc = std::max(acc, val); break;
232+ case HCCL_REDUCE_MIN: acc = std::min(acc, val); break;
233+ default: acc += val; break; // default to sum
234+ }
235+ }
236+ WriteElement(dst, i, acc, hcclDataType);
237+ }
238+}
239+ 
240+// ============================================================
241+// Device Memory Helper
242+// ============================================================
243+inline void* DevMalloc(size_t size) {
244+ void* ptr = nullptr;
245+ if (aclrtMalloc(&ptr, size, ACL_MEM_MALLOC_HUGE_FIRST) != 0) return nullptr;
246+ aclrtMemset(ptr, size, 0, size);
247+ return ptr;
248+}
249+ 
250+inline void DevFree(void* ptr) {
251+ if (ptr) aclrtFree(ptr);
252+}
253+ 
254+inline void* AlignUp512(void* ptr) {
255+ uint64_t addr = reinterpret_cast<uint64_t>(ptr);
256+ if (addr & 0x1ff) {
257+ addr = (addr & (~(uint64_t)0x1ff)) + 0x200;
258+ }
259+ return reinterpret_cast<void*>(addr);
260+}
261+ 
262+inline void* OffsetPtr(void* base, size_t offset) {
263+ return reinterpret_cast<void*>(reinterpret_cast<uint8_t*>(base) + offset);
264+}
265+ 
266+// ============================================================
267+// RankData — per-rank input tensor descriptor
268+// ============================================================
269+struct RankData {
270+ void* devicePtr; // device memory pointer
271+ size_t byteSize; // size in bytes
272+};
273+ 
274+// ============================================================
275+// MockContextBuilder
276+// ============================================================
277+class MockContextBuilder {
278+public:
279+ ~MockContextBuilder() { Destroy(); }
280+ 
281+ bool Build(uint32_t rankNum = 1, uint32_t rankId = 0) {
282+ rankNum_ = rankNum;
283+ rankId_ = rankId;
284+ 
285+ // Allocate window memory
286+ devWindowMem_ = DevMalloc(WINDOW_TOTAL_SIZE);
287+ if (!devWindowMem_) {
288+ printf("[MockContext] Window memory alloc failed (%lldMB)\n",
289+ (long long)(WINDOW_TOTAL_SIZE / (1024*1024)));
290+ return false;
291+ }
292+ 
293+ // Allocate workspace (extra 512B for alignment)
294+ size_t wsAllocSize = MIN_WORKSPACE_SIZE + 512;
295+ devWorkspaceRaw_ = DevMalloc(wsAllocSize);
296+ if (!devWorkspaceRaw_) {
297+ printf("[MockContext] Workspace alloc failed\n");
298+ return false;
299+ }
300+ devWorkspaceAligned_ = AlignUp512(devWorkspaceRaw_);
301+ 
302+ // Construct context on host
303+ MockHcclContext hostCtx;
304+ memset(&hostCtx, 0, sizeof(hostCtx));
305+ hostCtx.rankId = rankId_;
306+ hostCtx.rankNum = rankNum_;
307+ hostCtx.winSize = WINDOW_TOTAL_SIZE;
308+ hostCtx.workSpace = reinterpret_cast<uint64_t>(devWorkspaceAligned_);
309+ hostCtx.workSpaceSize = MIN_WORKSPACE_SIZE;
310+ 
311+ for (uint32_t i = 0; i < rankNum_; i++) {
312+ hostCtx.windowsIn[i] = reinterpret_cast<uint64_t>(devWindowMem_);
313+ hostCtx.windowsOut[i] = reinterpret_cast<uint64_t>(devWindowMem_);
314+ }
315+ 
316+ // Copy to device
317+ devContext_ = DevMalloc(sizeof(MockHcclContext));
318+ if (!devContext_) return false;
319+ if (aclrtMemcpy(devContext_, sizeof(hostCtx), &hostCtx, sizeof(hostCtx),
320+ ACL_MEMCPY_HOST_TO_DEVICE) != 0) {
321+ printf("[MockContext] H2D copy failed\n");
322+ return false;
323+ }
324+ 
325+ printf("[MockContext] Built: rankId=%u, rankNum=%u\n", rankId_, rankNum_);
326+ printf(" context @ %p\n", devContext_);
327+ printf(" workspace @ %p (aligned: %p)\n", devWorkspaceRaw_, devWorkspaceAligned_);
328+ printf(" window @ %p (%lldMB)\n", devWindowMem_,
329+ (long long)(WINDOW_TOTAL_SIZE / (1024*1024)));
330+ return true;
331+ }
332+ 
333+ void Destroy() {
334+ DevFree(devContext_); devContext_ = nullptr;
335+ DevFree(devWorkspaceRaw_); devWorkspaceRaw_ = nullptr;
336+ DevFree(devWindowMem_); devWindowMem_ = nullptr;
337+ devWorkspaceAligned_ = nullptr;
338+ }
339+ 
340+ void* GetContextAddr() const { return devContext_; }
341+ void* GetWorkspaceAligned() const { return devWorkspaceAligned_; }
342+ void* GetWindowMem() const { return devWindowMem_; }
343+ uint32_t GetRankNum() const { return rankNum_; }
344+ 
345+private:
346+ void* devContext_{nullptr};
347+ void* devWorkspaceRaw_{nullptr};
348+ void* devWorkspaceAligned_{nullptr};
349+ void* devWindowMem_{nullptr};
350+ uint32_t rankNum_{1};
351+ uint32_t rankId_{0};
352+};
353+ 
354+// ============================================================
355+// MultiRankMockContext — allocates N windows + N contexts
356+//
357+// Each rank gets its own window (184MB). All contexts share the
358+// same windowsIn[] array pointing to all ranks' windows.
359+// Combined with multi-stream + core binding, this enables real
360+// multi-rank AllGather on a single card.
361+//
362+// Usage:
363+// MultiRankMockContext mctx;
364+// mctx.Build(4); // 4 ranks
365+// void* ctx_r0 = mctx.GetContextAddr(0); // rank 0's context
366+// void* ctx_r1 = mctx.GetContextAddr(1); // rank 1's context
367+// // Launch V3 on stream[0] with ctx_r0, stream[1] with ctx_r1, ...
368+// ============================================================
369+class MultiRankMockContext {
370+public:
371+ ~MultiRankMockContext() { Destroy(); }
372+ 
373+ bool Build(uint32_t rankNum) {
374+ rankNum_ = rankNum;
375+ 
376+ // Allocate per-rank windows
377+ devWindows_.resize(rankNum, nullptr);
378+ for (uint32_t r = 0; r < rankNum; r++) {
379+ devWindows_[r] = DevMalloc(WINDOW_TOTAL_SIZE);
380+ if (!devWindows_[r]) {
381+ printf("[MultiRankMock] Window alloc failed for rank %u\n", r);
382+ return false;
383+ }
384+ }
385+ 
386+ // Allocate per-rank workspaces
387+ devWorkspaceRaw_.resize(rankNum, nullptr);
388+ devWorkspaceAligned_.resize(rankNum, nullptr);
389+ for (uint32_t r = 0; r < rankNum; r++) {
390+ devWorkspaceRaw_[r] = DevMalloc(MIN_WORKSPACE_SIZE + 512);
391+ if (!devWorkspaceRaw_[r]) {
392+ printf("[MultiRankMock] Workspace alloc failed for rank %u\n", r);
393+ return false;
394+ }
395+ devWorkspaceAligned_[r] = AlignUp512(devWorkspaceRaw_[r]);
396+ }
397+ 
398+ // Build per-rank contexts with cross-referenced windows
399+ devContexts_.resize(rankNum, nullptr);
400+ for (uint32_t r = 0; r < rankNum; r++) {
401+ MockHcclContext hostCtx;
402+ memset(&hostCtx, 0, sizeof(hostCtx));
403+ hostCtx.rankId = r;
404+ hostCtx.rankNum = rankNum;
405+ hostCtx.winSize = WINDOW_TOTAL_SIZE;
406+ hostCtx.workSpace = reinterpret_cast<uint64_t>(devWorkspaceAligned_[r]);
407+ hostCtx.workSpaceSize = MIN_WORKSPACE_SIZE;
408+ 
409+ // All ranks see all windows
410+ for (uint32_t j = 0; j < rankNum; j++) {
411+ hostCtx.windowsIn[j] = reinterpret_cast<uint64_t>(devWindows_[j]);
412+ hostCtx.windowsOut[j] = reinterpret_cast<uint64_t>(devWindows_[j]);
413+ }
414+ 
415+ devContexts_[r] = DevMalloc(sizeof(MockHcclContext));
416+ if (!devContexts_[r]) return false;
417+ if (aclrtMemcpy(devContexts_[r], sizeof(hostCtx), &hostCtx, sizeof(hostCtx),
418+ ACL_MEMCPY_HOST_TO_DEVICE) != 0) {
419+ printf("[MultiRankMock] H2D copy failed for rank %u\n", r);
420+ return false;
421+ }
422+ }
423+ 
424+ printf("[MultiRankMock] Built %u ranks, each window %lldMB\n",
425+ rankNum, (long long)(WINDOW_TOTAL_SIZE / (1024*1024)));
426+ for (uint32_t r = 0; r < rankNum; r++) {
427+ printf(" rank %u: context=%p window=%p workspace=%p\n",
428+ r, devContexts_[r], devWindows_[r], devWorkspaceAligned_[r]);
429+ }
430+ return true;
431+ }
432+ 
433+ void ClearAllFlags() {
434+ for (uint32_t r = 0; r < rankNum_; r++) {
435+ if (devWindows_[r]) {
436+ void* flagArea = OffsetPtr(devWindows_[r], FLAG_OFFSET_BYTES);
437+ aclrtMemset(flagArea, FLAG_AREA_SIZE, 0, FLAG_AREA_SIZE);
438+ }
439+ }
440+ }
441+ 
442+ void Destroy() {
443+ for (auto p : devContexts_) DevFree(p);
444+ for (auto p : devWorkspaceRaw_) DevFree(p);
445+ for (auto p : devWindows_) DevFree(p);
446+ devContexts_.clear();
447+ devWorkspaceRaw_.clear();
448+ devWorkspaceAligned_.clear();
449+ devWindows_.clear();
450+ }
451+ 
452+ void* GetContextAddr(uint32_t rank) const { return devContexts_[rank]; }
453+ void* GetWindowMem(uint32_t rank) const { return devWindows_[rank]; }
454+ void* GetWorkspaceAligned(uint32_t rank) const { return devWorkspaceAligned_[rank]; }
455+ uint32_t GetRankNum() const { return rankNum_; }
456+ 
457+private:
458+ uint32_t rankNum_{0};
459+ std::vector<void*> devWindows_;
460+ std::vector<void*> devWorkspaceRaw_;
461+ std::vector<void*> devWorkspaceAligned_;
462+ std::vector<void*> devContexts_;
463+};
464+ 
465+// ============================================================
466+// MockHcclServer — host thread simulating CCU server
467+//
468+// Accepts per-rank input tensors (device pointers) and uses them
469+// to simulate real collective communication semantics.
470+// ============================================================
471+class MockHcclServer {
472+public:
473+ /**
474+ * @param workspaceAligned 512B-aligned workspace device pointer
475+ * @param rankInputs per-rank input data (rankInputs[i] = rank i's tensor)
476+ * @param localRankId the rank ID of the kernel under test
477+ * @param deviceId NPU device ID
478+ */
479+ MockHcclServer(void* workspaceAligned,
480+ std::vector<RankData> rankInputs,
481+ uint32_t localRankId = 0,
482+ int deviceId = 0)
483+ : wsBase_(workspaceAligned)
484+ , rankInputs_(std::move(rankInputs))
485+ , localRankId_(localRankId)
486+ , rankNum_(static_cast<uint32_t>(rankInputs_.size()))
487+ , deviceId_(deviceId) {}
488+ 
489+ ~MockHcclServer() { Stop(); }
490+ 
491+ void Start() {
492+ running_ = true;
493+ serverThread_ = std::thread(&MockHcclServer::ServerLoop, this);
494+ printf("[MockServer] Started: rankNum=%u, localRankId=%u, workspace @ %p\n",
495+ rankNum_, localRankId_, wsBase_);
496+ }
497+ 
498+ void Stop() {
499+ running_ = false;
500+ if (serverThread_.joinable()) {
501+ serverThread_.join();
502+ }
503+ printf("[MockServer] Stopped. Processed %u messages.\n", msgCount_.load());
504+ }
505+ 
506+ bool IsFinalized() const { return finalized_.load(); }
507+ uint32_t GetMsgCount() const { return msgCount_.load(); }
508+ 
509+ // Block until the kernel signals Finalize, then respond.
510+ bool WaitForFinalize(uint32_t slot = 0, uint32_t timeoutMs = 5000) {
511+ auto start = std::chrono::steady_clock::now();
512+ while (true) {
513+ auto elapsed = std::chrono::steady_clock::now() - start;
514+ if (std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count() >=
515+ timeoutMs) {
516+ printf("[MockServer] WaitForFinalize: timeout after %ums\n", timeoutMs);
517+ return false;
518+ }
519+ 
520+ TurnCnt commitCnt;
521+ void* commitAddr = OffsetPtr(wsBase_,
522+ COMMIT_TURNCNT_OFFSET + slot * TURNCNT_STRIDE);
523+ if (aclrtMemcpy(&commitCnt, sizeof(TurnCnt), commitAddr, sizeof(TurnCnt),
524+ ACL_MEMCPY_DEVICE_TO_HOST) != 0) {
525+ usleep(100);
526+ continue;
527+ }
528+ 
529+ if (commitCnt.valid != COMMIT_VALID_MASK ||
530+ commitCnt.cnt <= lastProcessedCnt_[slot]) {
531+ usleep(100);
532+ continue;
533+ }
534+ 
535+ lastProcessedCnt_[slot] = commitCnt.cnt;
536+ 
537+ // Respond with FINALIZE_FINISH_CNT
538+ TurnCnt finishCnt;
539+ memset(&finishCnt, 0, sizeof(finishCnt));
540+ finishCnt.cnt = FINALIZE_FINISH_CNT;
541+ void* finishAddr = OffsetPtr(wsBase_,
542+ FINISH_TURNCNT_OFFSET + slot * TURNCNT_STRIDE);
543+ aclrtMemcpy(finishAddr, sizeof(TurnCnt), &finishCnt, sizeof(TurnCnt),
544+ ACL_MEMCPY_HOST_TO_DEVICE);
545+ 
546+ // Clear commit
547+ TurnCnt clearCnt;
548+ memset(&clearCnt, 0, sizeof(clearCnt));
549+ clearCnt.cnt = commitCnt.cnt;
550+ aclrtMemcpy(commitAddr, sizeof(TurnCnt), &clearCnt, sizeof(TurnCnt),
551+ ACL_MEMCPY_HOST_TO_DEVICE);
552+ 
553+ printf("[MockServer] Finalize response sent on slot %u "
554+ "(cnt=%lu → FINALIZE_FINISH_CNT)\n", slot, commitCnt.cnt);
555+ finalized_ = true;
556+ return true;
557+ }
558+ }
559+ 
560+private:
561+ void ServerLoop() {
562+ aclrtSetDevice(deviceId_);
563+ while (running_) {
564+ for (uint32_t i = 0; i < HCCL_MSG_CNT && running_; i++) {
565+ PollSlot(i);
566+ }
567+ usleep(50); // 50us polling interval
568+ }
569+ }
570+ 
571+ void PollSlot(uint32_t slot) {
572+ // Read commitTurnCnt[slot] from device
573+ TurnCnt commitCnt;
574+ void* commitAddr = OffsetPtr(wsBase_, COMMIT_TURNCNT_OFFSET + slot * TURNCNT_STRIDE);
575+ if (aclrtMemcpy(&commitCnt, sizeof(TurnCnt), commitAddr, sizeof(TurnCnt),
576+ ACL_MEMCPY_DEVICE_TO_HOST) != 0) {
577+ return;
578+ }
579+ 
580+ if (commitCnt.valid != COMMIT_VALID_MASK) return;
581+ if (commitCnt.cnt <= lastProcessedCnt_[slot]) return;
582+ 
583+ // Read the message
584+ HcclMsg msg;
585+ void* msgAddr = OffsetPtr(wsBase_, SEND_MSGS_OFFSET + slot * MSG_STRIDE);
586+ if (aclrtMemcpy(&msg, sizeof(HcclMsg), msgAddr, sizeof(HcclMsg),
587+ ACL_MEMCPY_DEVICE_TO_HOST) != 0) {
588+ return;
589+ }
590+ 
591+ printf("[MockServer] slot=%u cnt=%lu commType=%u dataCnt=%lu opType=%u\n",
592+ slot, commitCnt.cnt, msg.commType, msg.dataCnt, msg.opType);
593+ 
594+ // Dispatch by commType
595+ HandleMessage(msg);
596+ 
597+ // Record processed count
598+ lastProcessedCnt_[slot] = commitCnt.cnt;
599+ 
600+ // Write finishedTurnCnt[slot]
601+ TurnCnt finishCnt;
602+ memset(&finishCnt, 0, sizeof(finishCnt));
603+ finishCnt.cnt = commitCnt.cnt;
604+ void* finishAddr = OffsetPtr(wsBase_, FINISH_TURNCNT_OFFSET + slot * TURNCNT_STRIDE);
605+ aclrtMemcpy(finishAddr, sizeof(TurnCnt), &finishCnt, sizeof(TurnCnt),
606+ ACL_MEMCPY_HOST_TO_DEVICE);
607+ 
608+ // Clear commit valid flag
609+ TurnCnt clearCnt;
610+ memset(&clearCnt, 0, sizeof(clearCnt));
611+ clearCnt.cnt = commitCnt.cnt;
612+ aclrtMemcpy(commitAddr, sizeof(TurnCnt), &clearCnt, sizeof(TurnCnt),
613+ ACL_MEMCPY_HOST_TO_DEVICE);
614+ 
615+ msgCount_++;
616+ }
617+ 
618+ // --------------------------------------------------------
619+ // Message dispatch
620+ // --------------------------------------------------------
621+ void HandleMessage(const HcclMsg& msg) {
622+ switch (msg.commType) {
623+ case HCCL_CMD_ALLGATHER:
624+ HandleAllGather(msg);
625+ break;
626+ case HCCL_CMD_REDUCE_SCATTER:
627+ HandleReduceScatter(msg);
628+ break;
629+ case HCCL_CMD_ALLREDUCE:
630+ HandleAllReduce(msg);
631+ break;
632+ case HCCL_CMD_ALLTOALL:
633+ HandleAlltoAll(msg);
634+ break;
635+ default:
636+ printf("[MockServer] Unknown commType=%u, no-op\n", msg.commType);
637+ break;
638+ }
639+ }
640+ 
641+ // --------------------------------------------------------
642+ // AllGather: each rank contributes dataCnt elements
643+ // recvBuf = [rank0_chunk | rank1_chunk | ... | rankN-1_chunk]
644+ // ^stride ^stride
645+ // --------------------------------------------------------
646+ void HandleAllGather(const HcclMsg& msg) {
647+ size_t elemSize = HcclDataTypeSize(msg.hcclDataType);
648+ size_t chunkBytes = msg.dataCnt * elemSize;
649+ size_t strideBytes = msg.strideCount * elemSize;
650+ if (strideBytes == 0) strideBytes = chunkBytes; // default: tightly packed
651+ 
652+ void* recvBuf = reinterpret_cast<void*>(msg.recvBuffer);
653+ void* sendBuf = reinterpret_cast<void*>(msg.sendBuffer);
654+ 
655+ for (uint32_t r = 0; r < rankNum_; r++) {
656+ void* dst = OffsetPtr(recvBuf, r * strideBytes);
657+ if (r == localRankId_) {
658+ // Local rank: copy from kernel's sendBuf
659+ if (sendBuf && sendBuf != dst) {
660+ aclrtMemcpy(dst, chunkBytes, sendBuf, chunkBytes,
661+ ACL_MEMCPY_DEVICE_TO_DEVICE);
662+ }
663+ } else {
664+ // Remote rank: copy from rankInputs_[r]
665+ if (r < rankInputs_.size() && rankInputs_[r].devicePtr) {
666+ size_t copyBytes = std::min(chunkBytes, rankInputs_[r].byteSize);
667+ aclrtMemcpy(dst, copyBytes, rankInputs_[r].devicePtr, copyBytes,
668+ ACL_MEMCPY_DEVICE_TO_DEVICE);
669+ }
670+ }
671+ }
672+ 
673+ printf("[MockServer] AllGather: %u ranks × %zu bytes → recvBuf %p\n",
674+ rankNum_, chunkBytes, recvBuf);
675+ }
676+ 
677+ // --------------------------------------------------------
678+ // ReduceScatter: each rank has dataCnt elements (total input),
679+ // element-wise reduce across all ranks, then rank i gets
680+ // chunk i of the result. Output size = dataCnt / rankNum.
681+ // --------------------------------------------------------
682+ void HandleReduceScatter(const HcclMsg& msg) {
683+ size_t elemSize = HcclDataTypeSize(msg.hcclDataType);
684+ size_t totalElems = msg.dataCnt;
685+ size_t totalBytes = totalElems * elemSize;
686+ size_t chunkElems = totalElems / rankNum_;
687+ size_t chunkBytes = chunkElems * elemSize;
688+ 
689+ // D2H read all ranks' data
690+ std::vector<std::vector<uint8_t>> hostBufs(rankNum_);
691+ for (uint32_t r = 0; r < rankNum_; r++) {
692+ hostBufs[r].resize(totalBytes, 0);
693+ if (r == localRankId_) {
694+ // Local rank: read from kernel's sendBuf
695+ aclrtMemcpy(hostBufs[r].data(), totalBytes,
696+ reinterpret_cast<void*>(msg.sendBuffer), totalBytes,
697+ ACL_MEMCPY_DEVICE_TO_HOST);
698+ } else if (r < rankInputs_.size() && rankInputs_[r].devicePtr) {
699+ size_t readBytes = std::min(totalBytes, rankInputs_[r].byteSize);
700+ aclrtMemcpy(hostBufs[r].data(), readBytes,
701+ rankInputs_[r].devicePtr, readBytes,
702+ ACL_MEMCPY_DEVICE_TO_HOST);
703+ }
704+ }
705+ 
706+ // Host-side element-wise reduce
707+ std::vector<uint8_t> reducedBuf(totalBytes);
708+ std::vector<const void*> srcs(rankNum_);
709+ for (uint32_t r = 0; r < rankNum_; r++) srcs[r] = hostBufs[r].data();
710+ HostReduce(reducedBuf.data(), srcs, totalElems, msg.hcclDataType, msg.opType);
711+ 
712+ // Extract local rank's chunk and H2D to recvBuf
713+ aclrtMemcpy(reinterpret_cast<void*>(msg.recvBuffer), chunkBytes,
714+ reducedBuf.data() + localRankId_ * chunkBytes, chunkBytes,
715+ ACL_MEMCPY_HOST_TO_DEVICE);
716+ 
717+ printf("[MockServer] ReduceScatter: %u ranks × %zu elems → reduce → chunk[%u] %zu elems\n",
718+ rankNum_, totalElems, localRankId_, chunkElems);
719+ }
720+ 
721+ // --------------------------------------------------------
722+ // AllReduce: each rank has dataCnt elements,
723+ // element-wise reduce across all ranks, full result to recvBuf.
724+ // --------------------------------------------------------
725+ void HandleAllReduce(const HcclMsg& msg) {
726+ size_t elemSize = HcclDataTypeSize(msg.hcclDataType);
727+ size_t totalElems = msg.dataCnt;
728+ size_t totalBytes = totalElems * elemSize;
729+ 
730+ // D2H read all ranks' data
731+ std::vector<std::vector<uint8_t>> hostBufs(rankNum_);
732+ for (uint32_t r = 0; r < rankNum_; r++) {
733+ hostBufs[r].resize(totalBytes, 0);
734+ if (r == localRankId_) {
735+ aclrtMemcpy(hostBufs[r].data(), totalBytes,
736+ reinterpret_cast<void*>(msg.sendBuffer), totalBytes,
737+ ACL_MEMCPY_DEVICE_TO_HOST);
738+ } else if (r < rankInputs_.size() && rankInputs_[r].devicePtr) {
739+ size_t readBytes = std::min(totalBytes, rankInputs_[r].byteSize);
740+ aclrtMemcpy(hostBufs[r].data(), readBytes,
741+ rankInputs_[r].devicePtr, readBytes,
742+ ACL_MEMCPY_DEVICE_TO_HOST);
743+ }
744+ }
745+ 
746+ // Host-side element-wise reduce
747+ std::vector<uint8_t> reducedBuf(totalBytes);
748+ std::vector<const void*> srcs(rankNum_);
749+ for (uint32_t r = 0; r < rankNum_; r++) srcs[r] = hostBufs[r].data();
750+ HostReduce(reducedBuf.data(), srcs, totalElems, msg.hcclDataType, msg.opType);
751+ 
752+ // H2D full result to recvBuf
753+ aclrtMemcpy(reinterpret_cast<void*>(msg.recvBuffer), totalBytes,
754+ reducedBuf.data(), totalBytes,
755+ ACL_MEMCPY_HOST_TO_DEVICE);
756+ 
757+ printf("[MockServer] AllReduce: %u ranks × %zu elems → reduce → full result\n",
758+ rankNum_, totalElems);
759+ }
760+ 
761+ // --------------------------------------------------------
762+ // AlltoAll: rank i's block j → rank j's block i
763+ // Each rank has dataCnt elements total (rankNum blocks of
764+ // dataCnt/rankNum elements each).
765+ // For local rank: collect block[localRankId] from each rank.
766+ // --------------------------------------------------------
767+ void HandleAlltoAll(const HcclMsg& msg) {
768+ size_t elemSize = HcclDataTypeSize(msg.hcclDataType);
769+ size_t totalElems = msg.dataCnt;
770+ size_t totalBytes = totalElems * elemSize;
771+ size_t blockElems = totalElems / rankNum_;
772+ size_t blockBytes = blockElems * elemSize;
773+ 
774+ // D2H read all ranks' data
775+ std::vector<std::vector<uint8_t>> hostBufs(rankNum_);
776+ for (uint32_t r = 0; r < rankNum_; r++) {
777+ hostBufs[r].resize(totalBytes, 0);
778+ if (r == localRankId_) {
779+ aclrtMemcpy(hostBufs[r].data(), totalBytes,
780+ reinterpret_cast<void*>(msg.sendBuffer), totalBytes,
781+ ACL_MEMCPY_DEVICE_TO_HOST);
782+ } else if (r < rankInputs_.size() && rankInputs_[r].devicePtr) {
783+ size_t readBytes = std::min(totalBytes, rankInputs_[r].byteSize);
784+ aclrtMemcpy(hostBufs[r].data(), readBytes,
785+ rankInputs_[r].devicePtr, readBytes,
786+ ACL_MEMCPY_DEVICE_TO_HOST);
787+ }
788+ }
789+ 
790+ // Reassemble: recvBuf[j] = rank_j's block[localRankId]
791+ // i.e., collect what each rank would send to localRankId
792+ std::vector<uint8_t> resultBuf(totalBytes);
793+ for (uint32_t r = 0; r < rankNum_; r++) {
794+ // rank r sends its block[localRankId_] to us
795+ const uint8_t* srcBlock = hostBufs[r].data() + localRankId_ * blockBytes;
796+ uint8_t* dstBlock = resultBuf.data() + r * blockBytes;
797+ memcpy(dstBlock, srcBlock, blockBytes);
798+ }
799+ 
800+ // H2D to recvBuf
801+ aclrtMemcpy(reinterpret_cast<void*>(msg.recvBuffer), totalBytes,
802+ resultBuf.data(), totalBytes,
803+ ACL_MEMCPY_HOST_TO_DEVICE);
804+ 
805+ printf("[MockServer] AlltoAll: %u ranks × %zu blocks → reassemble for rank %u\n",
806+ rankNum_, blockElems, localRankId_);
807+ }
808+ 
809+ // --------------------------------------------------------
810+ // Member data
811+ // --------------------------------------------------------
812+ void* wsBase_{nullptr};
813+ std::vector<RankData> rankInputs_;
814+ uint32_t localRankId_{0};
815+ uint32_t rankNum_{1};
816+ int deviceId_{0};
817+ std::atomic<bool> running_{false};
818+ std::thread serverThread_;
819+ std::atomic<uint32_t> msgCount_{0};
820+ std::atomic<bool> finalized_{false};
821+ uint64_t lastProcessedCnt_[HCCL_MSG_CNT] = {};
822+};
823+ 
824+} // namespace mock_hccl
825+ 
826+#endif // MOCK_FRAMEWORK_H_
@@ -0,0 +1,348 @@
1+/**
2+ * Mock V3 Kernel Integration Test
3+ *
4+ * Tests AllGatherMatmulV3 on a single card by:
5+ * 1. Constructing mock comm_context (MockContextBuilder)
6+ * 2. Starting mock HCCL server (MockHcclServer) on host thread
7+ * 3. Calling aclnnAllGatherMatmulV3 with the mock context
8+ * 4. Verifying output shapes and basic data flow
9+ *
10+ * Build:
11+ * g++ -std=c++17 -o mock_v3_kernel_test mock_v3_kernel_test.cpp \
12+ * -I$ASCEND_HOME/include -I../op_api \
13+ * -L$ASCEND_HOME/lib64 -lascendcl \
14+ * -L$CUST_PKG/lib -lcust_opapi \
15+ * -Wl,--allow-shlib-undefined -lpthread
16+ *
17+ * Run:
18+ * LD_LIBRARY_PATH=$ASCEND_HOME/lib64:$CUST_PKG/lib \
19+ * LD_PRELOAD=$ASCEND_HOME/lib64/libhcomm.so \
20+ * ./mock_v3_kernel_test
21+ */
22+ 
23+#include <cstdio>
24+#include <cstdlib>
25+#include <cstring>
26+#include <vector>
27+#include <chrono>
28+#include "acl/acl.h"
29+#include "mock_framework.h"
30+ 
31+using namespace mock_hccl;
32+ 
33+// ============================================================
34+// Test Configuration
35+// ============================================================
36+struct TestConfig {
37+ int64_t M = 256; // per-rank M dimension
38+ int64_t K = 512; // K dimension
39+ int64_t N = 256; // N dimension
40+ uint32_t rankSize = 2; // simulated rank count (context has rankNum=1 but tiling uses this)
41+ bool isTransA = false;
42+ bool isTransB = false;
43+ int device = 0;
44+};
45+ 
46+// ============================================================
47+// Tensor utilities
48+// ============================================================
49+struct DevTensor {
50+ void* data{nullptr};
51+ int64_t dims[2]{0, 0};
52+ size_t elemSize{2}; // FP16
53+ 
54+ size_t ByteSize() const { return dims[0] * dims[1] * elemSize; }
55+ 
56+ bool Alloc(int64_t d0, int64_t d1, size_t eSize = 2) {
57+ dims[0] = d0; dims[1] = d1; elemSize = eSize;
58+ data = DevMalloc(ByteSize());
59+ return data != nullptr;
60+ }
61+ 
62+ void Free() { DevFree(data); data = nullptr; }
63+};
64+ 
65+// Fill device tensor with pattern (host-side fill then copy)
66+static bool FillPattern(DevTensor& t, uint8_t pattern) {
67+ size_t sz = t.ByteSize();
68+ std::vector<uint8_t> host(sz, pattern);
69+ return aclrtMemcpy(t.data, sz, host.data(), sz, ACL_MEMCPY_HOST_TO_DEVICE) == 0;
70+}
71+ 
72+// ============================================================
73+// Test: MockContextBuilder verification
74+// ============================================================
75+static bool TestContextBuilder() {
76+ printf("\n=== Test 1: MockContextBuilder ===\n");
77+ 
78+ MockContextBuilder ctx;
79+ if (!ctx.Build(1, 0)) {
80+ printf("[FAIL] Context build\n");
81+ return false;
82+ }
83+ 
84+ // Readback verification
85+ MockHcclContext readBack;
86+ if (aclrtMemcpy(&readBack, sizeof(readBack), ctx.GetContextAddr(), sizeof(readBack),
87+ ACL_MEMCPY_DEVICE_TO_HOST) != 0) {
88+ printf("[FAIL] Context readback\n");
89+ return false;
90+ }
91+ 
92+ bool ok = (readBack.rankId == 0) && (readBack.rankNum == 1) &&
93+ (readBack.workSpace == reinterpret_cast<uint64_t>(ctx.GetWorkspaceAligned()));
94+ 
95+ printf(" rankId=%u rankNum=%u workSpace=%p\n",
96+ readBack.rankId, readBack.rankNum, (void*)readBack.workSpace);
97+ printf("[%s] Context builder\n", ok ? "PASS" : "FAIL");
98+ return ok;
99+}
100+ 
101+// ============================================================
102+// Test: MockHcclServer basic protocol
103+// ============================================================
104+static bool TestServerProtocol() {
105+ printf("\n=== Test 2: MockHcclServer Protocol ===\n");
106+ 
107+ MockContextBuilder ctx;
108+ if (!ctx.Build(1, 0)) return false;
109+ 
110+ void* wsAligned = ctx.GetWorkspaceAligned();
111+ 
112+ // Manually write a commit message to slot 0
113+ // Simulate what the kernel's AllGather() would do:
114+ 
115+ // 1. Write a HcclMsg at sendMsgs[0]
116+ HcclMsg msg;
117+ memset(&msg, 0, sizeof(msg));
118+ msg.commType = HCCL_CMD_ALLGATHER;
119+ // For this test, just check protocol — no actual data copy
120+ msg.sendBuffer = 0;
121+ msg.recvBuffer = 0;
122+ msg.dataCnt = 100;
123+ msg.hcclDataType = 1; // FP16
124+ 
125+ void* msgAddr = reinterpret_cast<uint8_t*>(wsAligned) + SEND_MSGS_OFFSET;
126+ if (aclrtMemcpy(msgAddr, sizeof(msg), &msg, sizeof(msg),
127+ ACL_MEMCPY_HOST_TO_DEVICE) != 0) {
128+ printf("[FAIL] Write sendMsg\n");
129+ return false;
130+ }
131+ 
132+ // 2. Write commitTurnCnt[0]
133+ TurnCnt commit;
134+ memset(&commit, 0, sizeof(commit));
135+ commit.valid = COMMIT_VALID_MASK;
136+ commit.cnt = 1;
137+ 
138+ void* commitAddr = reinterpret_cast<uint8_t*>(wsAligned) + COMMIT_TURNCNT_OFFSET;
139+ if (aclrtMemcpy(commitAddr, sizeof(commit), &commit, sizeof(commit),
140+ ACL_MEMCPY_HOST_TO_DEVICE) != 0) {
141+ printf("[FAIL] Write commitTurnCnt\n");
142+ return false;
143+ }
144+ 
145+ // 3. Verify commit data is on device (before starting server)
146+ {
147+ TurnCnt readCommit;
148+ aclrtMemcpy(&readCommit, sizeof(readCommit), commitAddr, sizeof(readCommit),
149+ ACL_MEMCPY_DEVICE_TO_HOST);
150+ printf(" [debug] commitTurnCnt[0] on device: valid=%lu cnt=%lu\n",
151+ readCommit.valid, readCommit.cnt);
152+ printf(" [debug] expected valid=%u\n", COMMIT_VALID_MASK);
153+ }
154+ 
155+ // 3b. Start server and wait for it to process
156+ MockHcclServer server(wsAligned);
157+ server.Start();
158+ 
159+ // Wait for processing
160+ usleep(500000); // 500ms
161+ 
162+ server.Stop();
163+ 
164+ // 4. Check finishedTurnCnt[0]
165+ TurnCnt finish;
166+ void* finishAddr = reinterpret_cast<uint8_t*>(wsAligned) + FINISH_TURNCNT_OFFSET;
167+ if (aclrtMemcpy(&finish, sizeof(finish), finishAddr, sizeof(finish),
168+ ACL_MEMCPY_DEVICE_TO_HOST) != 0) {
169+ printf("[FAIL] Read finishedTurnCnt\n");
170+ return false;
171+ }
172+ 
173+ bool ok = (finish.cnt >= 1);
174+ printf(" finishedTurnCnt[0].cnt = %lu (expect >= 1)\n", finish.cnt);
175+ printf("[%s] Server protocol\n", ok ? "PASS" : "FAIL");
176+ return ok;
177+}
178+ 
179+// ============================================================
180+// Test: MockHcclServer data copy (AllGather simulation)
181+// ============================================================
182+static bool TestServerDataCopy() {
183+ printf("\n=== Test 3: MockHcclServer Data Copy ===\n");
184+ 
185+ MockContextBuilder ctx;
186+ if (!ctx.Build(1, 0)) return false;
187+ 
188+ void* wsAligned = ctx.GetWorkspaceAligned();
189+ 
190+ // Allocate source and destination buffers
191+ const size_t elemCount = 1024;
192+ const size_t elemSize = 2; // FP16
193+ const size_t byteCount = elemCount * elemSize;
194+ 
195+ void* sendBuf = DevMalloc(byteCount);
196+ void* recvBuf = DevMalloc(byteCount);
197+ if (!sendBuf || !recvBuf) {
198+ printf("[FAIL] Buffer alloc\n");
199+ DevFree(sendBuf); DevFree(recvBuf);
200+ return false;
201+ }
202+ 
203+ // Fill sendBuf with pattern
204+ std::vector<uint8_t> pattern(byteCount, 0xAB);
205+ aclrtMemcpy(sendBuf, byteCount, pattern.data(), byteCount, ACL_MEMCPY_HOST_TO_DEVICE);
206+ 
207+ // Write AllGather message
208+ HcclMsg msg;
209+ memset(&msg, 0, sizeof(msg));
210+ msg.commType = HCCL_CMD_ALLGATHER;
211+ msg.sendBuffer = reinterpret_cast<uint64_t>(sendBuf);
212+ msg.recvBuffer = reinterpret_cast<uint64_t>(recvBuf);
213+ msg.dataCnt = elemCount;
214+ msg.hcclDataType = 1; // FP16
215+ 
216+ void* msgAddr = reinterpret_cast<uint8_t*>(wsAligned) + SEND_MSGS_OFFSET;
217+ aclrtMemcpy(msgAddr, sizeof(msg), &msg, sizeof(msg), ACL_MEMCPY_HOST_TO_DEVICE);
218+ 
219+ TurnCnt commit;
220+ memset(&commit, 0, sizeof(commit));
221+ commit.valid = COMMIT_VALID_MASK;
222+ commit.cnt = 1;
223+ void* commitAddr = reinterpret_cast<uint8_t*>(wsAligned) + COMMIT_TURNCNT_OFFSET;
224+ aclrtMemcpy(commitAddr, sizeof(commit), &commit, sizeof(commit), ACL_MEMCPY_HOST_TO_DEVICE);
225+ 
226+ // Start server
227+ MockHcclServer server(wsAligned);
228+ server.Start();
229+ usleep(200000);
230+ server.Stop();
231+ 
232+ // Verify recvBuf has the data
233+ std::vector<uint8_t> result(byteCount, 0);
234+ aclrtMemcpy(result.data(), byteCount, recvBuf, byteCount, ACL_MEMCPY_DEVICE_TO_HOST);
235+ 
236+ bool ok = (memcmp(result.data(), pattern.data(), byteCount) == 0);
237+ printf(" recvBuf[0..3] = 0x%02x 0x%02x 0x%02x 0x%02x (expect 0xAB)\n",
238+ result[0], result[1], result[2], result[3]);
239+ printf("[%s] Server data copy\n", ok ? "PASS" : "FAIL");
240+ 
241+ DevFree(sendBuf);
242+ DevFree(recvBuf);
243+ return ok;
244+}
245+ 
246+// ============================================================
247+// Test: Finalize protocol
248+// ============================================================
249+static bool TestFinalizeProtocol() {
250+ printf("\n=== Test 4: Finalize Protocol ===\n");
251+ 
252+ MockContextBuilder ctx;
253+ if (!ctx.Build(1, 0)) return false;
254+ 
255+ void* wsAligned = ctx.GetWorkspaceAligned();
256+ 
257+ // Start server
258+ MockHcclServer server(wsAligned);
259+ server.Start();
260+ 
261+ // Simulate: kernel sends a regular AllGather commit on slot 0
262+ {
263+ HcclMsg msg;
264+ memset(&msg, 0, sizeof(msg));
265+ msg.commType = HCCL_CMD_ALLGATHER;
266+ msg.dataCnt = 64;
267+ msg.hcclDataType = 1;
268+ 
269+ void* msgAddr = reinterpret_cast<uint8_t*>(wsAligned) + SEND_MSGS_OFFSET;
270+ aclrtMemcpy(msgAddr, sizeof(msg), &msg, sizeof(msg), ACL_MEMCPY_HOST_TO_DEVICE);
271+ 
272+ TurnCnt commit;
273+ memset(&commit, 0, sizeof(commit));
274+ commit.valid = COMMIT_VALID_MASK;
275+ commit.cnt = 1;
276+ void* commitAddr = reinterpret_cast<uint8_t*>(wsAligned) + COMMIT_TURNCNT_OFFSET;
277+ aclrtMemcpy(commitAddr, sizeof(commit), &commit, sizeof(commit), ACL_MEMCPY_HOST_TO_DEVICE);
278+ 
279+ usleep(200000); // let server process the regular message
280+ }
281+ 
282+ printf(" Regular message processed: %u\n", server.GetMsgCount());
283+ 
284+ // Now simulate: kernel sends Finalize commit on slot 0
285+ // Write a new commit with cnt=2 (next sequence number)
286+ {
287+ TurnCnt commit;
288+ memset(&commit, 0, sizeof(commit));
289+ commit.valid = COMMIT_VALID_MASK;
290+ commit.cnt = 2;
291+ void* commitAddr = reinterpret_cast<uint8_t*>(wsAligned) + COMMIT_TURNCNT_OFFSET;
292+ aclrtMemcpy(commitAddr, sizeof(commit), &commit, sizeof(commit), ACL_MEMCPY_HOST_TO_DEVICE);
293+ }
294+ 
295+ // Use WaitForFinalize from main thread — it will detect the new commit
296+ // and respond with FINALIZE_FINISH_CNT
297+ bool finalizeOk = server.WaitForFinalize(0, 2000);
298+ 
299+ server.Stop();
300+ 
301+ // Verify finishedTurnCnt[0].cnt == FINALIZE_FINISH_CNT
302+ TurnCnt finish;
303+ void* finishAddr = reinterpret_cast<uint8_t*>(wsAligned) + FINISH_TURNCNT_OFFSET;
304+ aclrtMemcpy(&finish, sizeof(finish), finishAddr, sizeof(finish), ACL_MEMCPY_DEVICE_TO_HOST);
305+ 
306+ bool ok = finalizeOk && (finish.cnt == FINALIZE_FINISH_CNT);
307+ printf(" finishedTurnCnt[0].cnt = %lu\n", finish.cnt);
308+ printf(" expected FINALIZE_FINISH_CNT = %lu\n", FINALIZE_FINISH_CNT);
309+ printf(" server.IsFinalized() = %s\n", server.IsFinalized() ? "true" : "false");
310+ printf("[%s] Finalize protocol\n", ok ? "PASS" : "FAIL");
311+ return ok;
312+}
313+ 
314+// ============================================================
315+// Main
316+// ============================================================
317+int main(int argc, char* argv[])
318+{
319+ printf("========================================\n");
320+ printf(" Mock HCCL Framework Test\n");
321+ printf("========================================\n");
322+ 
323+ int device = 0;
324+ if (argc > 1) device = atoi(argv[1]);
325+ 
326+ if (aclInit(nullptr) != 0) { printf("[FAIL] aclInit\n"); return 1; }
327+ if (aclrtSetDevice(device) != 0) { printf("[FAIL] aclrtSetDevice(%d)\n", device); aclFinalize(); return 1; }
328+ 
329+ aclrtStream stream = nullptr;
330+ aclrtCreateStream(&stream);
331+ printf("Device=%d, stream=%p\n", device, stream);
332+ 
333+ int passed = 0, failed = 0;
334+ 
335+ if (TestContextBuilder()) passed++; else failed++;
336+ if (TestServerProtocol()) passed++; else failed++;
337+ if (TestServerDataCopy()) passed++; else failed++;
338+ if (TestFinalizeProtocol()) passed++; else failed++;
339+ 
340+ printf("\n========================================\n");
341+ printf(" Results: %d passed, %d failed\n", passed, failed);
342+ printf("========================================\n");
343+ 
344+ if (stream) aclrtDestroyStream(stream);
345+ aclrtResetDevice(device);
346+ aclFinalize();
347+ return failed > 0 ? 1 : 0;
348+}