From 0d9755fb2a94770f43bde67a7e758d2b306a29c2 Mon Sep 17 00:00:00 2001 From: georgebisbas Date: Tue, 21 Jul 2026 15:51:32 +0200 Subject: [PATCH] feat(l3): add Max/Min/Prod to allreduce/reduce_scatter kernels Wire a reduce_op scalar (args[5]) into the hand-written onephase/twophase/ring allreduce and reduce_scatter kernels, dispatching to TADD/TMAX/TMIN/TMUL. Mirrors pypto's ReduceOp via a local CollectiveReduceOp enum header (simpler_setup/incore/collectives_reduce_op.hpp), avoiding a cross-repo include dependency. bidirectional_ring and ibing reject non-Sum at entry (TPUT only; no AtomicMax/Min in the ISA). Orchestration entries and scene-test helpers thread the scalar through; new P=2 Max/Min/Prod sim cases for onephase/ring/reduce_scatter with golden dispatch. --- .../incore/collectives_reduce_op.hpp | 21 +++ tests/st/worker/collectives/_helpers.py | 85 +++++++-- .../allreduce_bidirectional_ring_kernel.cpp | 7 + .../kernels/aiv/allreduce_ibing_kernel.cpp | 7 + .../kernels/aiv/allreduce_onephase_kernel.cpp | 18 +- .../kernels/aiv/allreduce_ring_kernel.cpp | 18 +- .../kernels/aiv/allreduce_twophase_kernel.cpp | 18 +- .../allreduce_bidirectional_ring_orch.cpp | 1 + .../orchestration/allreduce_ibing_orch.cpp | 1 + .../orchestration/allreduce_onephase_orch.cpp | 1 + .../orchestration/allreduce_ring_orch.cpp | 1 + .../orchestration/allreduce_twophase_orch.cpp | 1 + .../collectives/allreduce/test_allreduce.py | 168 ++++++++++++++---- .../kernels/aiv/reduce_scatter_kernel.cpp | 18 +- .../orchestration/reduce_scatter_orch.cpp | 1 + .../reduce_scatter/test_reduce_scatter.py | 64 ++++++- 16 files changed, 375 insertions(+), 55 deletions(-) create mode 100644 simpler_setup/incore/collectives_reduce_op.hpp diff --git a/simpler_setup/incore/collectives_reduce_op.hpp b/simpler_setup/incore/collectives_reduce_op.hpp new file mode 100644 index 0000000000..53e4a55b81 --- /dev/null +++ b/simpler_setup/incore/collectives_reduce_op.hpp @@ -0,0 +1,21 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +#pragma once + +/// Reduction operator for simpler hand-written collective kernels. +/// Mirrors pypto's ReduceOp (include/pypto/ir/comm.h) without a cross-repo +/// include dependency. +enum class CollectiveReduceOp : int { + kSum = 0, + kMax = 1, + kMin = 2, + kProd = 3, +}; diff --git a/tests/st/worker/collectives/_helpers.py b/tests/st/worker/collectives/_helpers.py index 3336487478..f63acbd2fc 100644 --- a/tests/st/worker/collectives/_helpers.py +++ b/tests/st/worker/collectives/_helpers.py @@ -8,14 +8,16 @@ # ----------------------------------------------------------------------------------------------------------- """Shared helpers for collective scene tests. -Provides comm-window scratch-parameter computation and orch-function -building so each collective test (allreduce, allgather, reduce_scatter, +Provides comm-window scratch-parameter computation, orch-function +building, reduction operator enumeration, and golden output helpers +so each collective test (allreduce, allgather, reduce_scatter, broadcast, all_to_all) can reuse the same domain-allocation pattern. """ from __future__ import annotations import ctypes +from enum import IntEnum import torch from simpler.task_interface import CommBufferSpec, DataType, TaskArgs, TensorArgType @@ -25,6 +27,18 @@ _F32 = DataType.FLOAT32 + +class CollectiveReduceOp(IntEnum): + """Mirror of C++ CollectiveReduceOp (simpler_setup/incore/collectives_reduce_op.hpp).""" + + SUM = 0 + MAX = 1 + MIN = 2 + PROD = 3 + + +_REDUCE_OP_NAMES = {0: "Sum", 1: "Max", 2: "Min", 3: "Prod"} + # --------------------------------------------------------------------------- # Allreduce constants (must match kernel COUNT) # --------------------------------------------------------------------------- @@ -90,8 +104,9 @@ def _allreduce_scratch_params(mode: str, nranks: int) -> tuple[int, int, int]: def allreduce_orch_fn(orch, callables, task_args, config): """L3 orch: allocate a domain and submit all allreduce ranks as one group. - Reads nranks and mode_id from task_args scalars. Selects the - ChipCallable by mode name (e.g. ``allreduce_onephase``). + Reads nranks, mode_id, and optional reduce_op from task_args + scalars. Selects the ChipCallable by mode name (e.g. + ``allreduce_onephase``). """ nranks = int(task_args.nranks.value) if not (2 <= nranks <= ALLREDUCE_MAX_RANKS): @@ -100,6 +115,9 @@ def allreduce_orch_fn(orch, callables, task_args, config): if not (0 <= mode_id < len(_ALLREDUCE_MODE_NAMES)): raise ValueError(f"invalid allreduce mode_id: {mode_id}") mode = _ALLREDUCE_MODE_NAMES[mode_id] + reduce_op_val = 0 + if hasattr(task_args, "reduce_op"): + reduce_op_val = int(task_args.reduce_op.value) # ibing is only supported for P=2 if mode == "ibing" and nranks != 2: @@ -132,6 +150,7 @@ def allreduce_orch_fn(orch, callables, task_args, config): chip_args.add_tensor(domain.buffers["scratch"].tensor((float_elems,), _F32), TensorArgType.INOUT) chip_args.add_scalar(domain.domain_size) chip_args.add_scalar(domain.device_ctx) + chip_args.add_scalar(reduce_op_val) args_list.append(chip_args) orch.submit_next_level_group(chip, args_list, config, workers=list(range(nranks))) @@ -141,9 +160,27 @@ def allreduce_orch_fn(orch, callables, task_args, config): # --------------------------------------------------------------------------- -def allreduce_expected_output(nranks: int) -> list[float]: - """output[i] = nranks*i + 100*nranks*(nranks-1)//2.""" - return [float(nranks * i + 100 * nranks * (nranks - 1) // 2) for i in range(ALLREDUCE_COUNT)] +def allreduce_expected_output(nranks: int, reduce_op: CollectiveReduceOp = CollectiveReduceOp.SUM) -> list[float]: + """output = reduce_op over all rank inputs. + + Input[rank][i] = i + rank*100. Return [golden[i] for i in range(C)]. + """ + if reduce_op == CollectiveReduceOp.SUM: + return [float(nranks * i + 100 * nranks * (nranks - 1) // 2) for i in range(ALLREDUCE_COUNT)] + if reduce_op == CollectiveReduceOp.MAX: + return [float((nranks - 1) * 100 + i) for i in range(ALLREDUCE_COUNT)] + if reduce_op == CollectiveReduceOp.MIN: + return [float(i) for i in range(ALLREDUCE_COUNT)] + if reduce_op == CollectiveReduceOp.PROD: + return [allreduce_prod_expected(nranks, i) for i in range(ALLREDUCE_COUNT)] + raise ValueError(f"unsupported reduce_op: {reduce_op}") + + +def allreduce_prod_expected(nranks: int, i: int) -> float: + p = 1.0 + for r in range(nranks): + p *= float(r * 100 + i) + return p # --------------------------------------------------------------------------- @@ -162,6 +199,7 @@ def generic_collective_orch_fn( scratch_nbytes: int, window_size: int, extra_scalars: list | None = None, + post_scalars: list | None = None, ): """Generic L3 orch for single-mode collectives (allgather, reduce_scatter, broadcast, all_to_all). @@ -197,6 +235,8 @@ def generic_collective_orch_fn( for s in extras: chip_args.add_scalar(s) chip_args.add_scalar(domain.device_ctx) + for s in post_scalars or []: + chip_args.add_scalar(s) args_list.append(chip_args) orch.submit_next_level_group(chip, args_list, config, workers=list(range(nranks))) @@ -211,11 +251,27 @@ def allgather_expected_output(nranks: int) -> list[float]: return [float(r * 100 + i) for r in range(nranks) for i in range(COUNT_PER_RANK)] -def reduce_scatter_expected_output(nranks: int, dest: int) -> list[float]: - """output[j] = sum_r (dest*C+j + r*100) = nranks*(dest*C+j) + 100*nranks*(nranks-1)/2.""" - return [ - float(nranks * (dest * COUNT_PER_RANK + j) + 100 * nranks * (nranks - 1) // 2) for j in range(COUNT_PER_RANK) - ] +def reduce_scatter_expected_output( + nranks: int, dest: int, reduce_op: CollectiveReduceOp = CollectiveReduceOp.SUM +) -> list[float]: + """golden[dest][j] = reduce_op over r of input[r][dest*C+j].""" + base = dest * COUNT_PER_RANK + if reduce_op == CollectiveReduceOp.SUM: + return [float(nranks * (base + j) + 100 * nranks * (nranks - 1) // 2) for j in range(COUNT_PER_RANK)] + if reduce_op == CollectiveReduceOp.MAX: + return [float((nranks - 1) * 100 + base + j) for j in range(COUNT_PER_RANK)] + if reduce_op == CollectiveReduceOp.MIN: + return [float(base + j) for j in range(COUNT_PER_RANK)] + if reduce_op == CollectiveReduceOp.PROD: + return [reduce_scatter_prod_expected(nranks, base, j) for j in range(COUNT_PER_RANK)] + raise ValueError(f"unsupported reduce_op: {reduce_op}") + + +def reduce_scatter_prod_expected(nranks: int, base: int, j: int) -> float: + p = 1.0 + for r in range(nranks): + p *= float(r * 100 + base + j) + return p def broadcast_expected_output(root: int) -> list[float]: @@ -233,8 +289,8 @@ def all_to_all_expected_output(nranks: int, rank: int) -> list[float]: # --------------------------------------------------------------------------- -def make_allreduce_args(nranks: int, mode_id: int) -> TaskArgsBuilder: - """Build per-rank input/output tensors + nranks/mode_id scalars. +def make_allreduce_args(nranks: int, mode_id: int, reduce_op: int = 0) -> TaskArgsBuilder: + """Build per-rank input/output tensors + nranks/mode_id/reduce_op scalars. input[rank][i] = i + rank*100. Output initially zeros. """ @@ -251,4 +307,5 @@ def make_allreduce_args(nranks: int, mode_id: int) -> TaskArgsBuilder: builder_specs.append(STensor(f"out_{rank}", out)) builder_specs.append(SScalar("nranks", ctypes.c_int64(nranks))) builder_specs.append(SScalar("mode_id", ctypes.c_int64(mode_id))) + builder_specs.append(SScalar("reduce_op", ctypes.c_int64(reduce_op))) return TaskArgsBuilder(*builder_specs) diff --git a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_bidirectional_ring_kernel.cpp b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_bidirectional_ring_kernel.cpp index 9073f7d94c..dcc768bbe9 100644 --- a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_bidirectional_ring_kernel.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_bidirectional_ring_kernel.cpp @@ -36,6 +36,7 @@ #include "pto/comm/pto_comm_inst.hpp" #include "platform_comm/comm_context.h" #include "tensor.h" +#include "collectives_reduce_op.hpp" #ifndef __gm__ #define __gm__ @@ -93,6 +94,12 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); int nranks = static_cast(args[3]); __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + CollectiveReduceOp reduce_op = static_cast(args[5]); + // TPUT only supports Sum reduction. + if (reduce_op != CollectiveReduceOp::kSum) { + pipe_barrier(PIPE_ALL); + return; + } __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; diff --git a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ibing_kernel.cpp b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ibing_kernel.cpp index 38b3beeafa..b3d89a46c0 100644 --- a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ibing_kernel.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ibing_kernel.cpp @@ -49,6 +49,7 @@ #include "pto/comm/pto_comm_inst.hpp" #include "platform_comm/comm_context.h" #include "tensor.h" +#include "collectives_reduce_op.hpp" #if defined(__CPU_SIM) || defined(__COSTMODEL) #include "pto/comm/async_common/async_types.hpp" @@ -129,6 +130,12 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); int nranks = static_cast(args[3]); __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + CollectiveReduceOp reduce_op = static_cast(args[5]); + // TPUT only supports Sum reduction. + if (reduce_op != CollectiveReduceOp::kSum) { + pipe_barrier(PIPE_ALL); + return; + } __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; diff --git a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_onephase_kernel.cpp b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_onephase_kernel.cpp index c629d053d0..4730ec6001 100644 --- a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_onephase_kernel.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_onephase_kernel.cpp @@ -28,6 +28,7 @@ * tensor(2) = scratch (HCCL window slot, cross-rank addressable) * scalar(0) = nranks * scalar(1) = CommContext device pointer + * scalar(2) = reduce_op (CollectiveReduceOp: 0=Sum, 1=Max, 2=Min, 3=Prod) */ #include @@ -36,6 +37,7 @@ #include "pto/comm/pto_comm_inst.hpp" #include "platform_comm/comm_context.h" #include "tensor.h" +#include "collectives_reduce_op.hpp" #ifndef __gm__ #define __gm__ @@ -61,6 +63,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); int nranks = static_cast(args[3]); __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + CollectiveReduceOp reduce_op = static_cast(args[5]); __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; @@ -142,7 +145,20 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in TLOAD(recvTile, remoteG); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TADD(accTile, accTile, recvTile); + switch (reduce_op) { + case CollectiveReduceOp::kSum: + TADD(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kMax: + TMAX(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kMin: + TMIN(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kProd: + TMUL(accTile, accTile, recvTile); + break; + } set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); } diff --git a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ring_kernel.cpp b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ring_kernel.cpp index 3d4c030f68..6a4fb5b273 100644 --- a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ring_kernel.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ring_kernel.cpp @@ -34,6 +34,7 @@ * tensor(2) = scratch (HCCL window slot, cross-rank addressable) * scalar(0) = nranks * scalar(1) = CommContext device pointer + * scalar(2) = reduce_op (CollectiveReduceOp: 0=Sum, 1=Max, 2=Min, 3=Prod) */ #include @@ -42,6 +43,7 @@ #include "pto/comm/pto_comm_inst.hpp" #include "platform_comm/comm_context.h" #include "tensor.h" +#include "collectives_reduce_op.hpp" #ifndef __gm__ #define __gm__ @@ -86,6 +88,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); int nranks = static_cast(args[3]); __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + CollectiveReduceOp reduce_op = static_cast(args[5]); __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; @@ -163,7 +166,20 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in TLOAD(chunkTile, accG); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TADD(chunkTile, chunkTile, recvTile); + switch (reduce_op) { + case CollectiveReduceOp::kSum: + TADD(chunkTile, chunkTile, recvTile); + break; + case CollectiveReduceOp::kMax: + TMAX(chunkTile, chunkTile, recvTile); + break; + case CollectiveReduceOp::kMin: + TMIN(chunkTile, chunkTile, recvTile); + break; + case CollectiveReduceOp::kProd: + TMUL(chunkTile, chunkTile, recvTile); + break; + } set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); TSTORE(accG, chunkTile); diff --git a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_twophase_kernel.cpp b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_twophase_kernel.cpp index f350b5526b..6b89c84490 100644 --- a/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_twophase_kernel.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_twophase_kernel.cpp @@ -27,6 +27,7 @@ * tensor(2) = scratch (HCCL window slot, cross-rank addressable) * scalar(0) = nranks * scalar(1) = CommContext device pointer + * scalar(2) = reduce_op (CollectiveReduceOp: 0=Sum, 1=Max, 2=Min, 3=Prod) */ #include @@ -35,6 +36,7 @@ #include "pto/comm/pto_comm_inst.hpp" #include "platform_comm/comm_context.h" #include "tensor.h" +#include "collectives_reduce_op.hpp" #ifndef __gm__ #define __gm__ @@ -75,6 +77,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); int nranks = static_cast(args[3]); __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + CollectiveReduceOp reduce_op = static_cast(args[5]); __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; @@ -150,7 +153,20 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in TLOAD(recvTile, remoteG); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TADD(accTile, accTile, recvTile); + switch (reduce_op) { + case CollectiveReduceOp::kSum: + TADD(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kMax: + TMAX(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kMin: + TMIN(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kProd: + TMUL(accTile, accTile, recvTile); + break; + } set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); } diff --git a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_bidirectional_ring_orch.cpp b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_bidirectional_ring_orch.cpp index 4a5389c568..d71a54d1f8 100644 --- a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_bidirectional_ring_orch.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_bidirectional_ring_orch.cpp @@ -50,6 +50,7 @@ __attribute__((visibility("default"))) void allreduce_bidirectional_ring_orchest params.add_inout(scratch); params.add_scalar(orch_args.scalar(0)); // nranks params.add_scalar(orch_args.scalar(1)); // CommContext + params.add_scalar(orch_args.scalar(2)); // reduce_op rt_submit_aiv_task(0, params); } diff --git a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ibing_orch.cpp b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ibing_orch.cpp index 9326f379ef..f8c2c8a16b 100644 --- a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ibing_orch.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ibing_orch.cpp @@ -49,6 +49,7 @@ __attribute__((visibility("default"))) void allreduce_ibing_orchestration(const params.add_inout(scratch); params.add_scalar(orch_args.scalar(0)); // nranks params.add_scalar(orch_args.scalar(1)); // CommContext + params.add_scalar(orch_args.scalar(2)); // reduce_op rt_submit_aiv_task(0, params); } diff --git a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_onephase_orch.cpp b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_onephase_orch.cpp index 544502e480..f5b96eea5b 100644 --- a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_onephase_orch.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_onephase_orch.cpp @@ -50,6 +50,7 @@ __attribute__((visibility("default"))) void allreduce_orchestration(const ChipTa params.add_inout(scratch); params.add_scalar(orch_args.scalar(0)); // nranks params.add_scalar(orch_args.scalar(1)); // CommContext + params.add_scalar(orch_args.scalar(2)); // reduce_op rt_submit_aiv_task(0, params); } diff --git a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ring_orch.cpp b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ring_orch.cpp index f6e804ebeb..9998f0c3a9 100644 --- a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ring_orch.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_ring_orch.cpp @@ -49,6 +49,7 @@ __attribute__((visibility("default"))) void allreduce_ring_orchestration(const C params.add_inout(scratch); params.add_scalar(orch_args.scalar(0)); // nranks params.add_scalar(orch_args.scalar(1)); // CommContext + params.add_scalar(orch_args.scalar(2)); // reduce_op rt_submit_aiv_task(0, params); } diff --git a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_twophase_orch.cpp b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_twophase_orch.cpp index 29381ead01..9813046dee 100644 --- a/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_twophase_orch.cpp +++ b/tests/st/worker/collectives/allreduce/kernels/orchestration/allreduce_twophase_orch.cpp @@ -49,6 +49,7 @@ __attribute__((visibility("default"))) void allreduce_twophase_orchestration(con params.add_inout(scratch); params.add_scalar(orch_args.scalar(0)); // nranks params.add_scalar(orch_args.scalar(1)); // CommContext + params.add_scalar(orch_args.scalar(2)); // reduce_op rt_submit_aiv_task(0, params); } diff --git a/tests/st/worker/collectives/allreduce/test_allreduce.py b/tests/st/worker/collectives/allreduce/test_allreduce.py index 093874c37f..dfed974c14 100644 --- a/tests/st/worker/collectives/allreduce/test_allreduce.py +++ b/tests/st/worker/collectives/allreduce/test_allreduce.py @@ -20,7 +20,12 @@ from simpler_setup import SceneTestCase, scene_test -from .._helpers import allreduce_expected_output, allreduce_orch_fn, make_allreduce_args +from .._helpers import ( + CollectiveReduceOp, + allreduce_expected_output, + allreduce_orch_fn, + make_allreduce_args, +) def _orch_entry(source, func_name, config_name=""): @@ -95,15 +100,18 @@ class TestAllreduceOnephaseP2(SceneTestCase): "name": "onephase", "platforms": ["a2a3sim", "a2a3", "a5sim", "a5"], "config": {"device_count": 2}, - "params": {"nranks": 2, "mode_id": 0}, + "params": {"nranks": 2, "mode_id": 0, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -119,15 +127,18 @@ class TestAllreduceTwophaseP2(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim", "a5"], "manual": ["a2a3sim", "a5sim"], "config": {"device_count": 2}, - "params": {"nranks": 2, "mode_id": 1}, + "params": {"nranks": 2, "mode_id": 1, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -143,15 +154,18 @@ class TestAllreduceRingP2(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim", "a5"], "manual": ["a2a3sim", "a5sim"], "config": {"device_count": 2}, - "params": {"nranks": 2, "mode_id": 2}, + "params": {"nranks": 2, "mode_id": 2, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -167,15 +181,18 @@ class TestAllreduceBidirectionalRingP2(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim", "a5"], "manual": ["a2a3sim", "a5sim"], "config": {"device_count": 2}, - "params": {"nranks": 2, "mode_id": 3}, + "params": {"nranks": 2, "mode_id": 3, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -191,15 +208,18 @@ class TestAllreduceIbingP2(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim", "a5"], "manual": ["a2a3sim", "a5sim"], "config": {"device_count": 2}, - "params": {"nranks": 2, "mode_id": 4}, + "params": {"nranks": 2, "mode_id": 4, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -215,15 +235,18 @@ class TestAllreduceOnephaseP4(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": True, "config": {"device_count": 4}, - "params": {"nranks": 4, "mode_id": 0}, + "params": {"nranks": 4, "mode_id": 0, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -239,15 +262,18 @@ class TestAllreduceTwophaseP4(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": True, "config": {"device_count": 4}, - "params": {"nranks": 4, "mode_id": 1}, + "params": {"nranks": 4, "mode_id": 1, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -263,15 +289,18 @@ class TestAllreduceRingP4(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": True, "config": {"device_count": 4}, - "params": {"nranks": 4, "mode_id": 2}, + "params": {"nranks": 4, "mode_id": 2, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -287,15 +316,18 @@ class TestAllreduceBidirectionalRingP4(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": True, "config": {"device_count": 4}, - "params": {"nranks": 4, "mode_id": 3}, + "params": {"nranks": 4, "mode_id": 3, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): - expected = torch.tensor(allreduce_expected_output(params["nranks"]), dtype=torch.float32) + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) for rank in range(params["nranks"]): getattr(args, f"out_{rank}").copy_(expected) @@ -311,12 +343,12 @@ class TestAllreduceIbingNranksError(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": True, "config": {"device_count": 4}, - "params": {"nranks": 4, "mode_id": 4}, + "params": {"nranks": 4, "mode_id": 4, "reduce_op": 0}, } ] def generate_args(self, params): - return make_allreduce_args(params["nranks"], params["mode_id"]) + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) def compute_golden(self, args, params): pass @@ -327,5 +359,81 @@ def test_run(self, st_platform, st_worker, request): super().test_run(st_platform, st_worker, request) +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestAllreduceOnephaseP2MaxMinProd(SceneTestCase): + """Allreduce onephase — 2-rank, non-Sum reduce ops.""" + + CALLABLE = {"orchestration": allreduce_orch_fn, "callables": _ALLREDUCE_MODES} + CASES = [ + { + "name": "onephase_max", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "mode_id": 0, "reduce_op": 1}, + }, + { + "name": "onephase_min", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "mode_id": 0, "reduce_op": 2}, + }, + { + "name": "onephase_prod", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "mode_id": 0, "reduce_op": 3}, + }, + ] + + def generate_args(self, params): + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) + + def compute_golden(self, args, params): + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) + for rank in range(params["nranks"]): + getattr(args, f"out_{rank}").copy_(expected) + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestAllreduceRingP2MaxMinProd(SceneTestCase): + """Allreduce ring — 2-rank, non-Sum reduce ops.""" + + CALLABLE = {"orchestration": allreduce_orch_fn, "callables": _ALLREDUCE_MODES} + CASES = [ + { + "name": "ring_max", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "mode_id": 2, "reduce_op": 1}, + }, + { + "name": "ring_min", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "mode_id": 2, "reduce_op": 2}, + }, + { + "name": "ring_prod", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "mode_id": 2, "reduce_op": 3}, + }, + ] + + def generate_args(self, params): + return make_allreduce_args(params["nranks"], params["mode_id"], params["reduce_op"]) + + def compute_golden(self, args, params): + expected = torch.tensor( + allreduce_expected_output(params["nranks"], CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) + for rank in range(params["nranks"]): + getattr(args, f"out_{rank}").copy_(expected) + + if __name__ == "__main__": SceneTestCase.run_module(__name__) diff --git a/tests/st/worker/collectives/reduce_scatter/kernels/aiv/reduce_scatter_kernel.cpp b/tests/st/worker/collectives/reduce_scatter/kernels/aiv/reduce_scatter_kernel.cpp index 7e49e67f44..a1bfdfc2c4 100644 --- a/tests/st/worker/collectives/reduce_scatter/kernels/aiv/reduce_scatter_kernel.cpp +++ b/tests/st/worker/collectives/reduce_scatter/kernels/aiv/reduce_scatter_kernel.cpp @@ -26,6 +26,7 @@ * tensor(2) = scratch HCCL window slot (INOUT) * scalar(0) = nranks * scalar(1) = CommContext device pointer + * scalar(2) = reduce_op (CollectiveReduceOp: 0=Sum, 1=Max, 2=Min, 3=Prod) */ #include @@ -34,6 +35,7 @@ #include "pto/comm/pto_comm_inst.hpp" #include "platform_comm/comm_context.h" #include "tensor.h" +#include "collectives_reduce_op.hpp" #ifndef __gm__ #define __gm__ @@ -59,6 +61,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in __gm__ Tensor *scratch_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); int nranks = static_cast(args[3]); __gm__ CommContext *commCtx = reinterpret_cast<__gm__ CommContext *>(args[4]); + CollectiveReduceOp reduce_op = static_cast(args[5]); __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; __gm__ float *output = reinterpret_cast<__gm__ float *>(output_tensor->buffer.addr) + output_tensor->start_offset; @@ -144,7 +147,20 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in TLOAD(recvTile, remoteG); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TADD(accTile, accTile, recvTile); + switch (reduce_op) { + case CollectiveReduceOp::kSum: + TADD(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kMax: + TMAX(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kMin: + TMIN(accTile, accTile, recvTile); + break; + case CollectiveReduceOp::kProd: + TMUL(accTile, accTile, recvTile); + break; + } set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); } diff --git a/tests/st/worker/collectives/reduce_scatter/kernels/orchestration/reduce_scatter_orch.cpp b/tests/st/worker/collectives/reduce_scatter/kernels/orchestration/reduce_scatter_orch.cpp index 339aaaa89d..fcf4cef8cc 100644 --- a/tests/st/worker/collectives/reduce_scatter/kernels/orchestration/reduce_scatter_orch.cpp +++ b/tests/st/worker/collectives/reduce_scatter/kernels/orchestration/reduce_scatter_orch.cpp @@ -43,6 +43,7 @@ __attribute__((visibility("default"))) void reduce_scatter_orchestration(const C params.add_inout(scratch); params.add_scalar(orch_args.scalar(0)); // nranks params.add_scalar(orch_args.scalar(1)); // CommContext + params.add_scalar(orch_args.scalar(2)); // reduce_op rt_submit_aiv_task(0, params); } diff --git a/tests/st/worker/collectives/reduce_scatter/test_reduce_scatter.py b/tests/st/worker/collectives/reduce_scatter/test_reduce_scatter.py index a433afc19f..075392d697 100644 --- a/tests/st/worker/collectives/reduce_scatter/test_reduce_scatter.py +++ b/tests/st/worker/collectives/reduce_scatter/test_reduce_scatter.py @@ -27,6 +27,7 @@ COUNT_PER_RANK, DTYPE_NBYTES, SIGNAL_TAIL_NBYTES, + CollectiveReduceOp, generic_collective_orch_fn, reduce_scatter_expected_output, ) @@ -36,6 +37,9 @@ def reduce_scatter_orch_fn(orch, callables, task_args, config): nranks = int(task_args.nranks.value) scratch_nbytes = nranks * COUNT_PER_RANK * DTYPE_NBYTES + SIGNAL_TAIL_NBYTES window_size = max(scratch_nbytes, 4 * 1024) + reduce_op_val = 0 + if hasattr(task_args, "reduce_op"): + reduce_op_val = int(task_args.reduce_op.value) generic_collective_orch_fn( orch, callables, @@ -45,6 +49,7 @@ def reduce_scatter_orch_fn(orch, callables, task_args, config): float_elems=nranks * COUNT_PER_RANK, scratch_nbytes=scratch_nbytes, window_size=window_size, + post_scalars=[reduce_op_val], ) @@ -72,7 +77,7 @@ def reduce_scatter_orch_fn(orch, callables, task_args, config): } -def _make_args(nranks): +def _make_args(nranks, reduce_op=0): specs = [] for r in range(nranks): inp = torch.tensor([i + r * 100 for i in range(nranks * COUNT_PER_RANK)], dtype=torch.float32).share_memory_() @@ -80,6 +85,7 @@ def _make_args(nranks): specs.append(STensor(f"in_{r}", inp)) specs.append(STensor(f"out_{r}", out)) specs.append(SScalar("nranks", ctypes.c_int64(nranks))) + specs.append(SScalar("reduce_op", ctypes.c_int64(reduce_op))) return TaskArgsBuilder(*specs) @@ -92,16 +98,19 @@ class TestReduceScatterP2(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": ["a2a3sim", "a5sim"], "config": {"device_count": 2}, - "params": {"nranks": 2}, + "params": {"nranks": 2, "reduce_op": 0}, } ] def generate_args(self, params): - return _make_args(params["nranks"]) + return _make_args(params["nranks"], params["reduce_op"]) def compute_golden(self, args, params): for r in range(params["nranks"]): - expected = torch.tensor(reduce_scatter_expected_output(params["nranks"], r), dtype=torch.float32) + expected = torch.tensor( + reduce_scatter_expected_output(params["nranks"], r, CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) getattr(args, f"out_{r}").copy_(expected) @@ -114,16 +123,57 @@ class TestReduceScatterP4(SceneTestCase): "platforms": ["a2a3sim", "a2a3", "a5sim"], "manual": True, "config": {"device_count": 4}, - "params": {"nranks": 4}, + "params": {"nranks": 4, "reduce_op": 0}, } ] def generate_args(self, params): - return _make_args(params["nranks"]) + return _make_args(params["nranks"], params["reduce_op"]) def compute_golden(self, args, params): for r in range(params["nranks"]): - expected = torch.tensor(reduce_scatter_expected_output(params["nranks"], r), dtype=torch.float32) + expected = torch.tensor( + reduce_scatter_expected_output(params["nranks"], r, CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) + getattr(args, f"out_{r}").copy_(expected) + + +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestReduceScatterP2MaxMinProd(SceneTestCase): + """ReduceScatter — 2-rank, non-Sum reduce ops.""" + + CALLABLE = _CALLABLE + CASES = [ + { + "name": "p2_max", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "reduce_op": 1}, + }, + { + "name": "p2_min", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "reduce_op": 2}, + }, + { + "name": "p2_prod", + "platforms": ["a2a3sim"], + "config": {"device_count": 2}, + "params": {"nranks": 2, "reduce_op": 3}, + }, + ] + + def generate_args(self, params): + return _make_args(params["nranks"], params["reduce_op"]) + + def compute_golden(self, args, params): + for r in range(params["nranks"]): + expected = torch.tensor( + reduce_scatter_expected_output(params["nranks"], r, CollectiveReduceOp(params["reduce_op"])), + dtype=torch.float32, + ) getattr(args, f"out_{r}").copy_(expected)