Skip to content

fix: align MAT fill and expand sync pipes #930

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
HecreReed wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: align MAT fill and expand sync pipes #930

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 25 additions & 2 deletions include/PTO/IR/PTOOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
};
Comment on lines +4539 to +4555

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The helper lambda getASFromType can be simplified to reduce boilerplate and improve readability by extracting the memory space attribute first, and then performing a single dyn_cast_or_null check.

      auto getASFromType = [](::mlir::Type ty) -> std::optional<::mlir::pto::AddressSpace> {
        ::mlir::Attribute ms;
        if (auto tb = ::mlir::dyn_cast<::mlir::pto::TileBufType>(ty))
          ms = tb.getMemorySpace();
        else if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty))
          ms = mr.getMemorySpace();
        if (auto as = ::mlir::dyn_cast_or_null<::mlir::pto::AddressSpaceAttr>(ms))
          return as.getAddressSpace();
        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(); }
}];
}
Expand Down Expand Up @@ -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(); }
Expand Down
39 changes: 39 additions & 0 deletions test/lit/pto/texpands_mat_pipe_selection.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<cube>} {
%src = pto.alloc_tile
: !pto.tile_buf<loc=mat, dtype=f16, rows=32, cols=32, v_row=32, v_col=32,
blayout=col_major, slayout=row_major, fractal=512, pad=0>
%dst = pto.alloc_tile
: !pto.tile_buf<loc=right, dtype=f16, rows=32, cols=32, v_row=32, v_col=32,
blayout=row_major, slayout=col_major, fractal=512, pad=0>

pto.texpands ins(%scalar : f16)
outs(%src : !pto.tile_buf<loc=mat, dtype=f16, rows=32, cols=32, v_row=32, v_col=32,
blayout=col_major, slayout=row_major, fractal=512, pad=0>)
pto.tmov ins(%src : !pto.tile_buf<loc=mat, dtype=f16, rows=32, cols=32, v_row=32, v_col=32,
blayout=col_major, slayout=row_major, fractal=512, pad=0>)
outs(%dst : !pto.tile_buf<loc=right, dtype=f16, rows=32, cols=32, v_row=32, v_col=32,
blayout=row_major, slayout=col_major, fractal=512, pad=0>)
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(
21 changes: 15 additions & 6 deletions test/lit/pto/tfillpad_a5_mat_pipe_selection.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<gm>>) {
// 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<gm>>)
attributes {pto.kernel_kind = #pto.kernel_kind<cube>} {
%src = pto.alloc_tile
: !pto.tile_buf<loc=mat, dtype=f16, rows=32, cols=32, v_row=32, v_col=16,
blayout=col_major, slayout=row_major, fractal=512, pad=1>
Expand All @@ -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(
Loading