diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 5a0c17e711..8d67cbc8b0 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -10201,7 +10201,7 @@ This section documents PTO communication primitives. PTOAS currently exposes: ##### `pto.comm.tput` - Synchronous Remote Write -**Summary:** Lowers to `pto::comm::TPUT(...)` and copies data from local GM to remote GM through a VEC staging tile. +**Summary:** Copies data synchronously from local GM to remote GM through a VEC staging tile. **Arguments:** @@ -10214,9 +10214,29 @@ This section documents PTO communication primitives. PTOAS currently exposes: **Constraints & Verification:** -- `dst` / `src` must be GM-shaped values with positive static shapes. -- `dst` and `src` must have the same element type and static shape. -- `ping` / `pong` must be local VEC tile-like values whose element type matches `src`. +- `dst` / `src` must be GM-shaped values. Static dimensions must be positive. +- Dynamic dimensions are supported only when both operands are + `pto.partition_tensor_view` values. +- `dst` and `src` must have the same element type and identical static/dynamic + shape signatures. Corresponding dynamic extents must be equal at runtime. +- Runtime extents must be nonnegative, and each partition range must remain + within its backing tensor view. A zero extent denotes an empty transfer. +- `ping` / `pong` must be local VEC tile-like values whose element type matches + `src`. Their physical `rows` / `cols` must be positive static values. +- Staging `v_row` / `v_col` values must be positive and may be static or + dynamic. If a transfer enters the chunked path, a static `v_row` / `v_col` + must exactly divide the corresponding logical row / column extent. Use + dynamic valid dimensions when a partial final chunk is possible. +- When `pong` is present, it must have the same type as `ping`. + +**Semantics:** + +For every logical index in the common `src` / `dst` shape, the operation reads +the local `src` element and writes the corresponding remote `dst` element. +`atomic_none` performs a normal write; an atomic mode such as `atomic_add` +combines the source value with the destination according to that mode. The +staging bundle may divide the logical range into chunks but does not change the +logical transfer extent. **Examples:** @@ -10228,11 +10248,36 @@ pto.comm.tput(%dst, %src, buf(%ping) : !pto.partition_tensor_view<128xf32>, !pto pto.comm.tput(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.tile_buf, !pto.tile_buf) {atomicType = #pto} ``` +For a variable-size transfer, construct both partitions with the same runtime +extent and use dynamic staging valid dimensions when the extent may require a +partial final chunk: + +```mlir +%dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0], sizes = [%rows, %c4096] + : !pto.tensor_view -> !pto.partition_tensor_view +%src_part = pto.partition_view %src_view, + offsets = [%c0, %c0], sizes = [%rows, %c4096] + : !pto.tensor_view -> !pto.partition_tensor_view +%stage = pto.alloc_tile addr = %c0_i64 + valid_row = %c1 valid_col = %c4096 + : !pto.tile_buf +pto.comm.tput(%dst_part, %src_part, buf(%stage) + : !pto.partition_tensor_view, + !pto.partition_tensor_view, + !pto.tile_buf) + {atomicType = #pto} +``` + --- ##### `pto.comm.tget` - Synchronous Remote Read -**Summary:** Lowers to `pto::comm::TGET(...)` and copies data from remote GM to local GM through a VEC staging tile. +**Summary:** Copies data synchronously from remote GM to local GM through a VEC staging tile. **Arguments:** @@ -10245,8 +10290,17 @@ pto.comm.tput(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view<128xf32 **Constraints & Verification:** -- Same GM/global-like and staging constraints as `pto.comm.tput`. -- `dst` and `src` must have the same element type and static shape. +- The GM/global-like, dynamic partition, runtime extent, zero-length, and + staging constraints are the same as for `pto.comm.tput`. +- `dst` and `src` must have the same element type and identical static/dynamic + shape signatures. + +**Semantics:** + +For every logical index in the common `src` / `dst` shape, the operation reads +the remote `src` element and writes the corresponding local `dst` element. The +staging bundle may divide the logical range into chunks but does not change the +logical transfer extent. **Examples:** diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index d4c83a241a..5ec61a7d5d 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3934,8 +3934,14 @@ static bool isCommGlobalLikeType(Type ty) { return isa(ty); } -static LogicalResult verifyCommGlobalLike(Operation *op, Value value, - StringRef name) { +enum class CommGlobalShapePolicy { + StaticOnly, + AllowDynamicPartitionView, +}; + +static LogicalResult verifyCommGlobalLike( + Operation *op, Value value, StringRef name, + CommGlobalShapePolicy policy = CommGlobalShapePolicy::StaticOnly) { Type ty = value.getType(); if (!isCommGlobalLikeType(ty)) return op->emitOpError() @@ -3944,10 +3950,23 @@ static LogicalResult verifyCommGlobalLike(Operation *op, Value value, SmallVector shape = getShapeVec(ty); if (shape.empty()) return op->emitOpError() << "expects " << name << " to have rank >= 1"; + + bool opAllowsDynamic = + policy == CommGlobalShapePolicy::AllowDynamicPartitionView; + bool isAllowedDynamicType = isa(ty); for (int64_t dim : shape) { - if (dim == ShapedType::kDynamic || dim <= 0) - return op->emitOpError() << "expects " << name - << " to have a positive static shape"; + if (dim == ShapedType::kDynamic) { + if (!opAllowsDynamic) + return op->emitOpError() + << "does not support dynamic dimensions on " << name; + if (!isAllowedDynamicType) + return op->emitOpError() << "allows dynamic dimensions on " << name + << " only for partition_tensor_view"; + continue; + } + if (dim <= 0) + return op->emitOpError() << "expects every static dimension of " << name + << " to be positive"; } return success(); } @@ -16299,8 +16318,12 @@ LogicalResult TGetAsyncOp::verify() { } LogicalResult TPutOp::verify() { - if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) || - failed(verifyCommGlobalLike(*this, getSrc(), "src")) || + if (failed(verifyCommGlobalLike( + *this, getDst(), "dst", + CommGlobalShapePolicy::AllowDynamicPartitionView)) || + failed(verifyCommGlobalLike( + *this, getSrc(), "src", + CommGlobalShapePolicy::AllowDynamicPartitionView)) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping", "pong"))) @@ -16308,15 +16331,20 @@ LogicalResult TPutOp::verify() { if (getElemTy(getDst().getType()) != getElemTy(getSrc().getType())) return emitOpError("expects src and dst to have the same element type"); if (getShapeVec(getDst().getType()) != getShapeVec(getSrc().getType())) - return emitOpError("expects src and dst to have the same static shape"); + return emitOpError( + "expects src and dst to have the same static/dynamic shape signature"); if (getElemTy(getPing().getType()) != getElemTy(getSrc().getType())) return emitOpError("expects staging tile element type to match src/dst"); return success(); } LogicalResult TGetOp::verify() { - if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) || - failed(verifyCommGlobalLike(*this, getSrc(), "src")) || + if (failed(verifyCommGlobalLike( + *this, getDst(), "dst", + CommGlobalShapePolicy::AllowDynamicPartitionView)) || + failed(verifyCommGlobalLike( + *this, getSrc(), "src", + CommGlobalShapePolicy::AllowDynamicPartitionView)) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping", "pong"))) @@ -16324,7 +16352,8 @@ LogicalResult TGetOp::verify() { if (getElemTy(getDst().getType()) != getElemTy(getSrc().getType())) return emitOpError("expects src and dst to have the same element type"); if (getShapeVec(getDst().getType()) != getShapeVec(getSrc().getType())) - return emitOpError("expects src and dst to have the same static shape"); + return emitOpError( + "expects src and dst to have the same static/dynamic shape signature"); if (getElemTy(getPing().getType()) != getElemTy(getSrc().getType())) return emitOpError("expects staging tile element type to match src/dst"); return success(); diff --git a/test/lit/pto/comm_dynamic_async_scope_invalid.pto b/test/lit/pto/comm_dynamic_async_scope_invalid.pto new file mode 100644 index 0000000000..00777cb678 --- /dev/null +++ b/test/lit/pto/comm_dynamic_async_scope_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @dynamic_async_p2p( + %dst: !pto.partition_tensor_view, + %src: !pto.partition_tensor_view, + %session: !pto.async_session) { + %event = pto.comm.tput_async(%dst, %src, %session : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.async_session) -> !pto.async_event + return + } +} + +// CHECK: error: 'pto.comm.tput_async' op expects dst to have a static shape diff --git a/test/lit/pto/comm_dynamic_collective_scope_invalid.pto b/test/lit/pto/comm_dynamic_collective_scope_invalid.pto new file mode 100644 index 0000000000..1b63d15220 --- /dev/null +++ b/test/lit/pto/comm_dynamic_collective_scope_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @dynamic_collective( + %src: !pto.partition_tensor_view, + %peer: !pto.partition_tensor_view, + %stage: !pto.tile_buf) { + pto.comm.tbroadcast(%src, recv(%stage), group(%peer) : !pto.partition_tensor_view, !pto.tile_buf, !pto.partition_tensor_view) {root = 0 : i32} + return + } +} + +// CHECK: error: 'pto.comm.tbroadcast' op does not support dynamic dimensions on src diff --git a/test/lit/pto/comm_dynamic_signal_scope_invalid.pto b/test/lit/pto/comm_dynamic_signal_scope_invalid.pto new file mode 100644 index 0000000000..47a130ca4b --- /dev/null +++ b/test/lit/pto/comm_dynamic_signal_scope_invalid.pto @@ -0,0 +1,19 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @dynamic_signal( + %signal: !pto.partition_tensor_view, %value: i32) { + pto.comm.tnotify(%signal, %value : !pto.partition_tensor_view, i32) {notifyOp = #pto} + return + } +} + +// CHECK: error: 'pto.comm.tnotify' op does not support dynamic dimensions on signal diff --git a/test/lit/pto/comm_p2p_dynamic_partition_emitc.pto b/test/lit/pto/comm_p2p_dynamic_partition_emitc.pto new file mode 100644 index 0000000000..419768f158 --- /dev/null +++ b/test/lit/pto/comm_p2p_dynamic_partition_emitc.pto @@ -0,0 +1,68 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: ptoas --pto-arch=a3 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-arch=a3 --pto-level=level3 --mlir-print-ir-after=pto-resolve-buffer-select %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE + +module { + func.func @comm_p2p_dynamic_partition( + %dst_ptr: !pto.ptr, %src_ptr: !pto.ptr, %runtime_rows: i32) { + %c0_i64 = arith.constant 0 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %c4096 = arith.constant 4096 : index + %rows = arith.index_cast %runtime_rows : i32 to index + %dst_view = pto.make_tensor_view %dst_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout} : !pto.tensor_view + %src_view = pto.make_tensor_view %src_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout} : !pto.tensor_view + %dst = pto.partition_view %dst_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view -> !pto.partition_tensor_view + %src = pto.partition_view %src_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view -> !pto.partition_tensor_view + %stage = pto.alloc_tile addr = %c0_i64 valid_row = %c1 valid_col = %c4096 : !pto.tile_buf + pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf) {atomicType = #pto} + pto.comm.tget(%dst, %src, buf(%stage) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf) + return + } + + // A static valid shape remains legal. At runtime its row/column values must + // divide the transfer shape whenever the transfer enters the chunked path. + func.func @comm_p2p_dynamic_partition_static_valid( + %dst_ptr: !pto.ptr, %src_ptr: !pto.ptr, %runtime_rows: i32) { + %c0_i64 = arith.constant 0 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %c4096 = arith.constant 4096 : index + %rows = arith.index_cast %runtime_rows : i32 to index + %dst_view = pto.make_tensor_view %dst_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout} : !pto.tensor_view + %src_view = pto.make_tensor_view %src_ptr, shape = [%c32, %c4096], strides = [%c4096, %c1] {layout = #pto.layout} : !pto.tensor_view + %dst = pto.partition_view %dst_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view -> !pto.partition_tensor_view + %src = pto.partition_view %src_view, offsets = [%c0, %c0], sizes = [%rows, %c4096] : !pto.tensor_view -> !pto.partition_tensor_view + %stage = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf) {atomicType = #pto} + return + } +} + +// EMITC: pto::Shape<1, 1, 1, -1, 4096> +// EMITC-SAME: ({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^)]+}}) +// EMITC: pto::Stride<-1, -1, -1, -1, -1> +// EMITC-SAME: ({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^)]+}}) +// EMITC: GlobalTensor, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND> +// EMITC: pto::comm::TPUT( +// EMITC: pto::comm::TGET( + +// NATIVE: IR Dump After PTOResolveBufferSelect +// NATIVE-LABEL: func.func @comm_p2p_dynamic_partition +// NATIVE: %[[ROWS:.*]] = arith.index_cast +// NATIVE: %[[DST:.*]] = pto.partition_view {{.*}} sizes = [%[[ROWS]], %c4096] +// NATIVE: %[[SRC:.*]] = pto.partition_view {{.*}} sizes = [%[[ROWS]], %c4096] +// NATIVE: pto.comm.tput(%[[DST]], %[[SRC]], buf( +// NATIVE-SAME: !pto.partition_tensor_view +// NATIVE: pto.comm.tget(%[[DST]], %[[SRC]], buf( +// NATIVE-SAME: !pto.partition_tensor_view +// NATIVE-NOT: memref.subview diff --git a/test/lit/pto/comm_p2p_dynamic_partition_verify_invalid.pto b/test/lit/pto/comm_p2p_dynamic_partition_verify_invalid.pto new file mode 100644 index 0000000000..512a010a52 --- /dev/null +++ b/test/lit/pto/comm_p2p_dynamic_partition_verify_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @shape_signature_mismatch( + %dst: !pto.partition_tensor_view, + %src: !pto.partition_tensor_view, + %stage: !pto.tile_buf) { + pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf) {atomicType = #pto} + return + } +} + +// CHECK: error: 'pto.comm.tput' op expects src and dst to have the same static/dynamic shape signature diff --git a/test/lit/pto/comm_p2p_dynamic_tensor_view_invalid.pto b/test/lit/pto/comm_p2p_dynamic_tensor_view_invalid.pto new file mode 100644 index 0000000000..7ca73a3677 --- /dev/null +++ b/test/lit/pto/comm_p2p_dynamic_tensor_view_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @direct_dynamic_tensor_view( + %dst: !pto.tensor_view, + %src: !pto.tensor_view, + %stage: !pto.tile_buf) { + pto.comm.tget(%dst, %src, buf(%stage) : !pto.tensor_view, !pto.tensor_view, !pto.tile_buf) + return + } +} + +// CHECK: error: 'pto.comm.tget' op allows dynamic dimensions on dst only for partition_tensor_view diff --git a/test/lit/pto/comm_p2p_nonpositive_static_invalid.pto b/test/lit/pto/comm_p2p_nonpositive_static_invalid.pto new file mode 100644 index 0000000000..5dec47d2ae --- /dev/null +++ b/test/lit/pto/comm_p2p_nonpositive_static_invalid.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @nonpositive_static_dimension( + %dst: !pto.partition_tensor_view<0x4096xi8>, + %src: !pto.partition_tensor_view<0x4096xi8>, + %stage: !pto.tile_buf) { + pto.comm.tput(%dst, %src, buf(%stage) : !pto.partition_tensor_view<0x4096xi8>, !pto.partition_tensor_view<0x4096xi8>, !pto.tile_buf) {atomicType = #pto} + return + } +} + +// CHECK: error: 'pto.comm.tput' op expects every static dimension of dst to be positive diff --git a/tools/ptobc/testdata/comm_p2p_dynamic_v0_roundtrip.pto b/tools/ptobc/testdata/comm_p2p_dynamic_v0_roundtrip.pto new file mode 100644 index 0000000000..776ff0bcca --- /dev/null +++ b/tools/ptobc/testdata/comm_p2p_dynamic_v0_roundtrip.pto @@ -0,0 +1,21 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// 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. + +module { + func.func @comm_p2p_dynamic_v0( + %dst: !pto.partition_tensor_view, + %src: !pto.partition_tensor_view, + %ping: !pto.tile_buf, + %pong: !pto.tile_buf) { + pto.comm.tput(%dst, %src, buf(%ping) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf) {atomicType = #pto} + pto.comm.tput(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf, !pto.tile_buf) {atomicType = #pto} + pto.comm.tget(%dst, %src, buf(%ping) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf) + pto.comm.tget(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view, !pto.partition_tensor_view, !pto.tile_buf, !pto.tile_buf) + return + } +} diff --git a/tools/ptobc/tests/CMakeLists.txt b/tools/ptobc/tests/CMakeLists.txt index 4c855c4a72..263be1fbee 100644 --- a/tools/ptobc/tests/CMakeLists.txt +++ b/tools/ptobc/tests/CMakeLists.txt @@ -113,6 +113,14 @@ add_test(NAME ptobc_tstore_fp_v0_encode ${CMAKE_CURRENT_LIST_DIR}/tstore_fp_v0_encode.sh ) +add_test(NAME ptobc_comm_p2p_dynamic_v0_encode + COMMAND ${CMAKE_COMMAND} -E env + PTOBC_BIN=$ + PTOAS_BIN=${CMAKE_BINARY_DIR}/tools/ptoas/ptoas + TESTDATA_DIR=${PTObc_TESTDATA_DIR} + ${CMAKE_CURRENT_LIST_DIR}/comm_p2p_dynamic_v0_encode.sh +) + add_test(NAME ptobc_tdequant_v0_encode COMMAND ${CMAKE_COMMAND} -E env PTOBC_BIN=$ diff --git a/tools/ptobc/tests/comm_p2p_dynamic_v0_encode.sh b/tools/ptobc/tests/comm_p2p_dynamic_v0_encode.sh new file mode 100755 index 0000000000..255346edcb --- /dev/null +++ b/tools/ptobc/tests/comm_p2p_dynamic_v0_encode.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# 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. + +set -euo pipefail + +PTOBC_BIN=${PTOBC_BIN:-} +if [[ -z "${PTOBC_BIN}" ]]; then + echo "error: PTOBC_BIN not set" >&2 + exit 2 +fi + +PTOAS_BIN=${PTOAS_BIN:-} +if [[ -z "${PTOAS_BIN}" ]]; then + echo "error: PTOAS_BIN not set" >&2 + exit 2 +fi + +TESTDATA_DIR=${TESTDATA_DIR:-} +if [[ -z "${TESTDATA_DIR}" ]]; then + echo "error: TESTDATA_DIR not set" >&2 + exit 2 +fi + +IN="${TESTDATA_DIR}/comm_p2p_dynamic_v0_roundtrip.pto" +OUT_DIR=${OUT_DIR:-"${PWD}/ptobc_comm_p2p_dynamic_out"} +mkdir -p "${OUT_DIR}" + +BC="${OUT_DIR}/comm_p2p_dynamic_v0_roundtrip.ptobc" +ROUNDTRIP="${OUT_DIR}/comm_p2p_dynamic_v0_roundtrip.roundtrip.pto" + +"${PTOBC_BIN}" encode "${IN}" -o "${BC}" +"${PTOBC_BIN}" decode "${BC}" -o "${ROUNDTRIP}" + +[[ $(grep -Fc "pto.comm.tput(" "${ROUNDTRIP}") -eq 2 ]] +[[ $(grep -Fc "pto.comm.tget(" "${ROUNDTRIP}") -eq 2 ]] +grep -F "!pto.partition_tensor_view" "${ROUNDTRIP}" >/dev/null +grep -F "atomic_add" "${ROUNDTRIP}" >/dev/null +grep -F "pto.comm.tput(%0, %1, buf(%2, %3)" "${ROUNDTRIP}" >/dev/null +grep -F "pto.comm.tget(%0, %1, buf(%2, %3)" "${ROUNDTRIP}" >/dev/null + +"${PTOAS_BIN}" --pto-arch=a3 --emit-pto-ir "${ROUNDTRIP}" -o /dev/null