已合并
增加mc2单卡测试mock server的能力 #3368
YuZhengzhong创建于 3月26日
增加mc2单卡测试mock server的能力 #3368
已合并
共 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 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 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 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 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 | + | ||
| @@ -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 | + | ||
| 24 | + | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 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 | +} | ||