From 4b5395dac4087899549acd40185a03d28b2f9546 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Fri, 10 Jul 2026 17:09:28 +0800 Subject: [PATCH] fix: align MAT fill and expand sync pipes --- include/PTO/IR/PTOOps.td | 27 ++++++++++++- test/lit/pto/texpands_mat_pipe_selection.pto | 39 +++++++++++++++++++ .../pto/tfillpad_a5_mat_pipe_selection.pto | 21 +++++++--- 3 files changed, 79 insertions(+), 8 deletions(-) create mode 100644 test/lit/pto/texpands_mat_pipe_selection.pto diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 1bca2965e..0cff3bc63 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -4535,7 +4535,30 @@ def TExpandsOp : PTO_TOp<"texpands", [ }]; let extraClassDeclaration = [{ - ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } + ::mlir::pto::PIPE getPipe() { + auto getASFromType = [](Type ty) + -> std::optional<::mlir::pto::AddressSpace> { + if (auto tb = ::mlir::dyn_cast<::mlir::pto::TileBufType>(ty)) { + if (auto as = ::mlir::dyn_cast_or_null<::mlir::pto::AddressSpaceAttr>( + tb.getMemorySpace())) + return as.getAddressSpace(); + return std::nullopt; + } + if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { + if (auto ms = mr.getMemorySpace()) { + if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) + return as.getAddressSpace(); + } + return std::nullopt; + } + return std::nullopt; + }; + + auto dstAS = getASFromType(getDst().getType()); + if (dstAS && *dstAS == ::mlir::pto::AddressSpace::MAT) + return ::mlir::pto::PIPE::PIPE_MTE2; + return ::mlir::pto::PIPE::PIPE_V; + } ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } }]; } @@ -4837,7 +4860,7 @@ def TFillPadOp : PTO_TOp<"tfillpad", [ if (sOpt.has_value() && dOpt.has_value() && sOpt.value() == ::mlir::pto::AddressSpace::MAT && dOpt.value() == ::mlir::pto::AddressSpace::MAT) - return ::mlir::pto::PIPE::PIPE_MTE1; + return ::mlir::pto::PIPE::PIPE_MTE2; return ::mlir::pto::PIPE::PIPE_V; } ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } diff --git a/test/lit/pto/texpands_mat_pipe_selection.pto b/test/lit/pto/texpands_mat_pipe_selection.pto new file mode 100644 index 000000000..dadb993e9 --- /dev/null +++ b/test/lit/pto/texpands_mat_pipe_selection.pto @@ -0,0 +1,39 @@ +// 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 a5 --enable-insert-sync %s | FileCheck %s +// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + // MAT TEXPANDS runs on MTE2. The following MAT->RIGHT TMOV runs on MTE1, + // so sync insertion must order TEXPANDS before TMOV across those pipes. + func.func @texpands_mat_sync(%scalar: f16) + attributes {pto.kernel_kind = #pto.kernel_kind} { + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.texpands ins(%scalar : f16) + outs(%src : !pto.tile_buf) + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: AICORE void texpands_mat_sync( +// CHECK: TEXPANDS( +// CHECK: set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID{{[0-9]+}}); +// CHECK: wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID{{[0-9]+}}); +// CHECK: TMOV( diff --git a/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto b/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto index 866559923..e523366cc 100644 --- a/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto +++ b/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto @@ -1,9 +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: ptoas --pto-arch a5 --enable-insert-sync %s | FileCheck %s +// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s module attributes {"pto.target_arch" = "a5"} { - // A5: mat->mat TFILLPAD follows the cube data-movement path and must not be - // treated as PIPE_V, otherwise sync insertion emits invalid V pipe events. - func.func @tfillpad_mat_sync(%in: memref<32x16xf16, #pto.address_space>) { + // A5: mat->mat TFILLPAD runs on MTE2. The following MAT->RIGHT TMOV runs on + // MTE1, so sync insertion must order TFILLPAD before TMOV across those pipes. + func.func @tfillpad_mat_sync(%in: memref<32x16xf16, #pto.address_space>) + attributes {pto.kernel_kind = #pto.kernel_kind} { %src = pto.alloc_tile : !pto.tile_buf @@ -27,8 +37,7 @@ module attributes {"pto.target_arch" = "a5"} { } // CHECK-LABEL: AICORE void tfillpad_mat_sync( +// CHECK: TFILLPAD( // CHECK: set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID{{[0-9]+}}); // CHECK: wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID{{[0-9]+}}); -// CHECK: TFILLPAD( -// CHECK-NOT: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID{{[0-9]+}}); -// CHECK-NOT: wait_flag(PIPE_V, PIPE_MTE1, EVENT_ID{{[0-9]+}}); +// CHECK: TMOV(