Draft
Conversation
…ag/wait_flag Made-with: Cursor
…e_runner Made-with: Cursor
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! 此拉取请求通过引入两个新示例显著扩展了示例套件：一个多卡集体 gather ( Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Comment on lines
+1
to
+166
| # 多卡 Gather Case 实现计划（comm_gather�? | ||
|
|
||
| �?simpler-PTO 中新增一�?*仅跑多卡 TGATHER** �?case，当前默认配置为 **2 卡 uint8_t、每卡 512 元素**，不接计算任务，参�?pto-comm-isa �?`tests/npu/a2a3/comm/st/testcase/tgather` 的流程与结构�? | ||
|
|
||
| --- | ||
|
|
||
| ## 一、目标与范围 | ||
|
|
||
| - **Case �?*：`comm_gather`（建议路径：`examples/host_build_graph/comm_gather/`�? | ||
| - **功能**�? 2 张卡�?collective gather（root 从各 rank 收数据），无 GEMM 等计算任�? | ||
| - **平台**：a2a3 真机（依�?HCCL + ACL，不�?a2a3sim�? | ||
| - **参考实�?*：pto-comm-isa �?tgather（`ForkAndRunWithHcclRootInfo` + `TestContext` + `RunGatherKernel` + `TGatherKernelImpl`�? | ||
|
|
||
| --- | ||
|
|
||
| ## 二、代码结构分�? | ||
|
|
||
| ### 2.1 当前 simpler-PTO 单卡流程 | ||
|
|
||
| - **入口**：`run_example.py` �?`code_runner.run()` | ||
| - **参数**：`-k kernels_dir`、`-g golden.py`、`-d device`（默�?0）、`-p platform` | ||
| - **流程**�? | ||
| 1. 加载 `kernel_config.py`、`golden.py` | ||
| 2. 构建 runtime（host + aicpu + aicore�? | ||
| 3. 编译 orchestration �?SO，编译各 kernel �?binary | ||
| 4. 单进程：`Runtime()` �?`initialize(orch_so, func_args, kernel_binaries)` �?`launch_runtime()` �?`finalize()` �?Python 侧与 golden 比对 | ||
| - **特点**：单进程、单 device、无 HCCL、无 rank 概念 | ||
|
|
||
| ### 2.2 pto-comm-isa 多卡 tgather 流程 | ||
|
|
||
| - **入口**：GTest �?main 调用 `RunGather<T, count>(n_ranks, n_devices, first_rank_id, first_device_id)`（例�?4 卡：`RunGather<float,256>(4,4,0,0)`�? | ||
| - **多进�?*：`ForkAndRunWithHcclRootInfo(n_ranks, first_rank_id, first_device_id, lambda)` | ||
| - fork �?n_ranks 个子进程 | ||
| - rank0 子进程：`HcclGetRootInfo` 写入 mmap �?`SharedBootstrap` | ||
| - 其余 rank 子进程：轮询等待 `state==1` 后读 `rootInfo` | ||
| - 每个子进程执�?lambda：`RunGatherKernel(rankId, n_ranks, n_devices, first_device_id, rootInfo, root)` | ||
| - **�?rank �?*（`RunGatherKernel`）： | ||
| 1. **建联**：`TestContext.Init(rank_id, n_ranks, n_devices, first_device_id, rootInfo)` | ||
| - `aclInit` �?`aclrtSetDevice(deviceId)` �?`aclrtCreateStream` | ||
| - `HcclCommInitRootInfo` �?`HcclGetCommName` �?`HcomGetCommHandleByGroup` | ||
| - `HcclAllocComResourceByTiling` �?得到 `deviceCtx`（含 `windowsIn` 等），再 `aclrtMemcpy` �?`hostCtx` | ||
| 2. **内存**：`WindowAlloc(localWinBase, winOffset, src_size)` / `dst_size`（在 HCCL 注册�?window 内分配） | ||
| 3. **数据**：host 上初始化 `src_host`，拷贝到 device �?`src_ptr`；root 再初始化 `dst_host` 并拷�?`dst_ptr` | ||
| 4. **同步**：`HcclHostBarrier(comm, stream)` | ||
| 5. **�?kernel**：`TGatherKernelImpl<<<1, nullptr, ctx.stream>>>(dst_ptr, src_ptr, ctx.deviceCtx, n_ranks, root)` | ||
| 6. **同步**：`aclrtSynchronizeStream` �?`HcclHostBarrier` | ||
| 7. **校验**：root �?`dst` 拷回 host，在 C++ 内逐元素与期望值比�? | ||
| 8. **收尾**：`TestContext.Finalize()` | ||
|
|
||
| **结论**�? | ||
| - **建联时机**：在每个 fork 出的子进程内、在�?kernel **之前**，由 `TestContext.Init` 完成（HCCL 通信�?+ RDMA window 等）�? | ||
| - **谁先谁后**：rank0 先执�?`HcclGetRootInfo` 并写�?mmap；其�?rank �?`state==1` 后再�?`rootInfo` �?`HcclCommInitRootInfo`，因此建联是“rank0 提供 rootInfo，各 rank 各自 Init”的顺序�? | ||
|
|
||
| ### 2.3 与现�?simpler-PTO 的差�? | ||
|
|
||
| | 维度 | 现有 host_build_graph case | comm_gather（本 case�? | | ||
| |--------------|-----------------------------|--------------------------------| | ||
| | 进程/设备 | 单进程、单 device | 多进程（fork）、每进程一 device | | ||
| | 运行�? | Runtime + orchestration | �?Runtime；独�?host 可执�? | | ||
| | 建联 | �? | �?rank �?Init 时建�?HCCL | | ||
| | Kernel 调用 | 通过 launch_runtime 统一�? | �?rank 自己 aclrtLaunchKernel | | ||
| | 校验 | Python golden 比对 | C++ �?root 自检（或再写文件�? | | ||
|
|
||
| --- | ||
|
|
||
| ## 三、需要修�?新增的内�? | ||
|
|
||
| ### 3.1 启动命令与参�? | ||
|
|
||
| - **目标用法**（示例）�? | ||
| ```bash | ||
| python run_example.py -k examples/host_build_graph/comm_gather/kernels \ | ||
| -g examples/host_build_graph/comm_gather/golden.py \ | ||
| --n-ranks 2 --first-device 0 | ||
| ``` | ||
| - **修改�?*�? | ||
| - �?**run_example.py** 中增加参数：`--n-ranks`（默�?4）、`--first-device`（默�?0）�? | ||
| - 将这两个参数传入 `create_code_runner(..., n_ranks=..., first_device_id=...)`（或等价接口）�? | ||
|
|
||
| ### 3.2 何时走“多�?gather”分�? | ||
|
|
||
| - �?**code_runner** 中根�?`kernel_config.py` �?`RUNTIME_CONFIG["runtime"] == "comm_gather"` 判断�? | ||
| - 若为 `comm_gather`�? | ||
| - **�?*构建默认 runtime�?*�?*编译 orchestration�?*�?*�?`Runtime()` / `initialize()` / `launch_runtime()`�? | ||
| - 改为：构�?**comm_gather 专用 host 可执�?* + 编译 **AIV gather kernel**，然后执行该 host 可执行（内部 fork 4 个进程），根据进程退出码或约定输出判断成�?失败�? | ||
|
|
||
| ### 3.3 通信建联的调度时机（再强调） | ||
|
|
||
| - 建联**不在** Python 主进程里做，�?*不在**“先起好 4 个进程再一�?Init”的集中步骤�? | ||
| - 建联�?**每个 fork 出的子进程内**、在调用 `RunGatherKernel(rankId, ...)` 时，�?**RunGatherKernel 开头的 `TestContext.Init(...)`** 完成�? | ||
| - 顺序：rank0 生成并写�?`HcclRootInfo` �?其他 rank 读到 `rootInfo` �?�?rank 各自 `Init`（含 `HcclCommInitRootInfo`、`HcclAllocComResourceByTiling`）→ 再各自做 `WindowAlloc`、拷贝、barrier、发 kernel�? | ||
|
|
||
| ### 3.4 目录与文件规�? | ||
|
|
||
| ``` | ||
| examples/host_build_graph/comm_gather/ | ||
| ├── README.md # 说明：四�?gather、无计算、如何跑、依�? | ||
| ├── golden.py # 多卡时以 C++ 校验为主，此处可占位或仅做形�?接口兼容 | ||
| ├── PLAN_multicard_gather.md # 本计划文档（可后续删除或保留作记录） | ||
| └── kernels/ | ||
| ├── kernel_config.py # runtime = "comm_gather", n_ranks = 4 �? | ||
| ├── aiv/ | ||
| �? └── kernel_gather_comm.cpp # 设备侧：TGatherKernelImpl(dst, src, HcclDeviceContext*, nranks, root) | ||
| └── host/ # �?pto-comm-isa 对齐�?host 侧逻辑 | ||
| ├── comm_common.hpp # TestContext, ForkAndRunWithHcclRootInfo, WindowAlloc, HcclRemotePtr 等（可从 pto-isa �?tests/npu/a2a3/comm/st/testcase/common.hpp 复制或引用） | ||
| ├── run_gather_kernel.cpp # RunGatherKernel<T,count>：Init、WindowAlloc、拷贝、Barrier、发 kernel、Sync、Barrier、root 校验、Finalize | ||
| └── main_comm_gather.cpp # main：解�?n_ranks/first_device（或�?env/argv 读），调�?ForkAndRunWithHcclRootInfo(n_ranks, 0, first_device, RunGatherKernel<...>) | ||
| ``` | ||
|
|
||
| - **kernel_config.py**�? | ||
| - `RUNTIME_CONFIG = {"runtime": "comm_gather", "n_ranks": 4, "first_device_id": 0}` | ||
| - 只登�?AIV gather kernel（如 `kernel_gather_comm.cpp`），**不需�?* orchestration（可�?`ORCHESTRATION` 或填占位，code_runner �?comm_gather 分支下不�?orchestration）�? | ||
|
|
||
| ### 3.5 构建 | ||
|
|
||
| - **AIV kernel**：沿用现�?`KernelCompiler.compile_incore(..., core_type="aiv")`，产�?kernel binary（与现有 case 一致，�?host 可执行加载）�? | ||
| - **Host 可执�?*（comm_gather 专用）： | ||
| - 需能编�?C++（如 g++）并链接 **ACL**�?*HCCL**（及 CANN 提供�?`hccl_context.h` 等）�? | ||
| - 源文件：`main_comm_gather.cpp`、`run_gather_kernel.cpp`（以�?`comm_common.hpp` 头文件）�? | ||
| - �?simpler-PTO 当前没有 HCCL 的构建配置，需要在�?case 或公共脚本中增加：HCCL 头路径、库路径、链�?`hccl` 等（�?pto-comm-isa 的编译方式对齐）�? | ||
| - 可执行需要能**加载 kernel binary**（例如从文件路径，由 Python 在运行前写入临时文件并传�?argv），并在�?rank 内用 ACL �?launch 接口（如 `aclrtLaunchKernel` 或当�?CANN 推荐方式）发�?kernel�? | ||
|
|
||
| ### 3.6 运行流程（comm_gather 分支�? | ||
|
|
||
| 1. **code_runner** 检测到 `runtime == "comm_gather"`�? | ||
| 2. 使用 `n_ranks`、`first_device_id`（来�?`kernel_config` 或命令行覆盖）�? | ||
| 3. 编译 AIV kernel，将得到�?kernel binary 写到临时文件（如 `comm_gather_kernel.bin`）�? | ||
| 4. 编译并链�?host 可执行（若尚未编译或源有更新）�? | ||
| 5. 执行该可执行，传入参数（例如 kernel binary 路径、n_ranks、first_device_id）；**可执行内�?*执行 `ForkAndRunWithHcclRootInfo(n_ranks, 0, first_device_id, ...)`，在子进程中完成建联�?gather�? | ||
| 6. 根据可执行退出码判断成功/失败；可选：若可执行�?root �?result 写到文件，再�?Python 里用 `golden.py` �?`compute_golden` 做一次比对（本计划建议先采用 C++ 内校验，退出码即可）�? | ||
|
|
||
| ### 3.7 Golden 校验 | ||
|
|
||
| - **建议**：与 pto-comm-isa 一致，**仅在 C++ �?*�?root �?`RunGatherKernel` 内校�?`dst_host`；可执行返回 0 即表示通过�? | ||
| - **golden.py**：保�?`generate_inputs` / `compute_golden` 接口以满�?code_runner 的通用加载逻辑，但多卡时可不依赖其输出做比对（例如 `generate_inputs` 返回空或占位，`compute_golden` 空实现）�? | ||
|
|
||
| ### 3.8 依赖 | ||
|
|
||
| - **pto-isa（_deps�?*：已�?`pto/comm`（TGATHER、ParallelGroup）、`pto_comm_inst.hpp`；需能包�?`tests/npu/a2a3/comm/st/testcase/common.hpp`，或�?`comm_common.hpp` 复制到本 case �?`kernels/host/` 并适配 include�? | ||
| - **CANN / HCCL**：`hccl.h`、`hccl_types.h`、`hccl_context.h`、`hcom.h`（若存在）；目标环境需能编译并链接�? | ||
|
|
||
| --- | ||
|
|
||
| ## 四、实施步骤小结（Plan�? | ||
|
|
||
| | 步骤 | 内容 | | ||
| |------|------| | ||
| | 1 | 新增目录 `examples/host_build_graph/comm_gather/` �?`kernels/`、`kernels/aiv/`、`kernels/host/`�?| | ||
| | 2 | �?pto-comm-isa 抽取并适配 **comm_common.hpp**（TestContext、ForkAndRunWithHcclRootInfo、WindowAlloc、HcclRemotePtr）到 `kernels/host/comm_common.hpp`（或引用 pto-isa �?common.hpp 并保�?include 路径）�?| | ||
| | 3 | �?pto-comm-isa �?tgather 抽取 **RunGatherKernel** �?**TGatherKernelImpl**：设备侧放入 **kernel_gather_comm.cpp**，host 侧放�?**run_gather_kernel.cpp**�?*main_comm_gather.cpp** 实现 main，调�?`ForkAndRunWithHcclRootInfo(4, 0, first_device_id, ...)`�?| | ||
| | 4 | 编写 **kernel_config.py**（`runtime: comm_gather`，n_ranks=4，仅注册 AIV kernel）�?| | ||
| | 5 | 编写 **golden.py**（满足接口即可，校验�?C++ 为主）�?| | ||
| | 6 | �?**run_example.py** 中增�?`--n-ranks`、`--first-device`，并传入 code_runner�?| | ||
| | 7 | �?**code_runner.run()** 中增加分支：�?`runtime == "comm_gather"` 时，不构建默�?runtime/orchestration，改为构�?comm_gather �?host 可执�?+ AIV kernel，执行可执行并根据退出码判断成功�?| | ||
| | 8 | 实现 **comm_gather 的构建逻辑**：编�?host 源文件并链接 ACL/HCCL；编�?AIV kernel；约�?kernel binary 路径或通过参数传入可执行�?| | ||
| | 9 | 编写 **README.md**：说明四�?gather、如何运行、环境与依赖�?| | ||
|
|
||
| --- | ||
|
|
||
| ## 五、风险与注意�? | ||
|
|
||
| - **环境**：需�?a2a3 真机且安�?CANN/HCCL，保证头与库可用；CI 若仅单卡可考虑跳过�?mock�? | ||
| - **HCCL �?*：`hccl_context.h` 等可能随 CANN 版本变化，若 pto-isa 与当�?CANN 不一致，需�?`comm_common.hpp` 或本 case �?include 中做兼容�? | ||
| - **Windows**：当前方案依�?fork/mmap，仅适用�?Linux；若需�?Windows 跑多卡，需另设计（如多进程 + 共享内存�?MPI 式启动）�? | ||
|
|
||
| 以上为多�?gather case 的完整分析与实现计划，确认后即可按步骤修改代码�? |
There was a problem hiding this comment.
这个计划文档（PLAN_multicard_gather.md）中的一些设计与最终的实现不一致。例如：
- 计划中提到使用
KernelCompiler.compile_incore编译设备侧 kernel（kernel_gather_comm.cpp），但实际实现是将设备代码和主机代码合并在run_gather_kernel.cpp中，并通过Makefile进行混合编译。 kernel_config.py中的KERNELS列表为空，也印证了这一点。
为了避免后续维护者产生困惑，建议在本次 PR 中将这个计划文档删除，因为 README.md 已经准确地描述了当前的实现方式。
| # __global__ AICORE device code. g++ cannot compile it. Use CANN mixed compiler | ||
| # if available (e.g. set CXX to CANN wrapper). Otherwise build on device with CANN. | ||
| $(TARGET): $(SRCS) comm_common.hpp run_gather_kernel.cpp | ||
| $(CXX) $(INC) -o $@ main_comm_gather.cpp $(LDFLAGS) |
There was a problem hiding this comment.
这里的 CXX ?= g++ 默认值可能会导致编译失败，因为 g++ 无法处理包含 __global__ AICORE 设备代码的混合源文件。如果用户直接运行 make，错误信息可能不够清晰。
为了提升用户体验，建议在执行编译前增加一个检查。如果 CXX 是 g++，就给出一个明确的错误提示，引导用户使用 CANN 混合编译器。
$(if $(filter g++,$(CXX)), \
$(error CXX is set to g++, which cannot compile device code. Please use a CANN mixed compiler, e.g., 'make CXX=g++_ccec'), \
$(CXX) $(INC) -o $@ main_comm_gather.cpp $(LDFLAGS))
| StrideDyn dstStride(nranks * count, nranks * count, nranks * count, count, 1); | ||
| Global dstG(dst, dstShape, dstStride); | ||
|
|
||
| Global tensors[16]; |
There was a problem hiding this comment.
这里的 Global tensors[16]; 硬编码了 ParallelGroup 最大支持16个rank。虽然目前 comm_gather 的用例是4卡，没有问题，但如果将来有人复用此代码并使用超过16个rank，TGATHER 将会只收集前16个rank的数据，可能导致非预期的行为或校验失败。
为了提高代码的可维护性和可复用性，建议在此处添加注释，明确指出这个16个rank的限制。例如：
// Note: ParallelGroup is limited to 16 ranks here.
Global tensors[16];| /** | ||
| * GEMM + Gather orchestration (a2a3). | ||
| * | ||
| * Two independent tasks (no dependency): |
There was a problem hiding this comment.
文件顶部的注释指出 "Two independent tasks (no dependency)"，但这与代码实现不符。
在第107行，代码明确地添加了任务依赖 runtime->add_successor(t0, t1);，使得 Gather 任务 (t1) 必须在 GEMM 任务 (t0) 完成后才能开始。
为了保持代码和文档的一致性，请更新此处的注释，使其准确描述任务间的依赖关系。
Suggested change
| * Two independent tasks (no dependency): | |
| * Task 0: C = A @ B (64x64 GEMM, AIC) |
Made-with: Cursor
Made-with: Cursor
- 新增 PTOCompiler、create_compiler()，compile() 返回 artifacts - CodeRunner 支持 compiled_artifacts/prebuilt_dir，跳过 build 分支 - run_example.py: 先 compile，n_devices>1 时写临时目录、spawn 子进程 - 新增 --run-only、--prebuilt-dir 参数 - 更新 multi_bgemm README 说明编译一次多进程运行 Made-with: Cursor
- 移除 --run-only、--prebuilt-dir 参数 - 新增 run_on_device() worker，多进程并行执行 CodeRunner.run() - 不再 spawn Python 子进程，直接多进程运行 runner Made-with: Cursor
- hccl_bindings.py: HCCL ctypes bindings (get_root_info, init_comm, barrier) - comm_include: hccl_context.h, hccl_helpers.h for kernel use - cpt_and_comm case: GEMM -> WindowMemCopyIn -> TGATHER -> WindowMemCopyOut - multi_card_code_runner: run_on_device_comm, comm_context injection - multi_card_run_example: requires_comm path with Barrier + shared rootInfo Made-with: Cursor
- hccl_helper: C++ lib linked like pto-comm-isa (ascendcl, hcomm, runtime) - Python no longer loads libacl.so/libhccl.so; uses libhccl_helper.so only - Add hccl_helper_get_root_info, hccl_helper_init_comm, hccl_helper_barrier - cpt_and_comm README: build hccl_helper before run Made-with: Cursor
Made-with: Cursor
Made-with: Cursor
Made-with: Cursor
… tensor Made-with: Cursor
…_comm Made-with: Cursor
Made-with: Cursor
Made-with: Cursor
Made-with: Cursor
…he flush and pipe_barrier Made-with: Cursor
…instead of raising Made-with: Cursor
…er in between Made-with: Cursor
… device-side TNOTIFY/TWAIT barrier Made-with: Cursor
…le launch) Made-with: Cursor
0311
20260312
Kernel allgather
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.