Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
123 changes: 90 additions & 33 deletions lib/PTO/Transforms/PTOMemoryConsistency.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLFunctionalExtras.h"

namespace mlir {
namespace pto {
Expand Down Expand Up @@ -318,6 +319,17 @@ struct TNotifyReleaseState {
UnitAttr::get(cmo.getContext()));
}
}

void remapPayloads(llvm::function_ref<Value(Value)> mapper) {
for (PendingReleaseAccess &access : pendingAccesses)
if (access.payload)
access.payload = mapper(access.payload);
for (Value &payload : addressedCmoPayloads)
if (payload)
payload = mapper(payload);
if (hasReleaseCmoMarker())
recomputeAddressedState();
}
};

struct SignalAcquireState {
Expand Down Expand Up @@ -408,10 +420,58 @@ static TNotifyReleaseState getReleaseStateForMacroModel(Operation *op) {
return state;
}

static TNotifyReleaseState getDirectTNotifyReleaseState(Operation *op) {
static func::FuncOp lookupCallee(func::CallOp call) {
return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
call.getOperation(), call.getCalleeAttr());
}

static TNotifyReleaseState
collectTNotifyReleaseState(Region &region,
llvm::DenseSet<Operation *> &activeCallees);

static Value remapCalleePayloadToCallOperand(Value payload, func::FuncOp callee,
func::CallOp call) {
if (!payload)
return payload;

Value root = getPayloadAliasRoot(payload);
auto arg = dyn_cast<BlockArgument>(root);
if (!arg)
return payload;
if (arg.getOwner() != &callee.getBody().front())
return payload;
if (arg.getArgNumber() >= call.getNumOperands())
return payload;
return call.getOperand(arg.getArgNumber());
}

static TNotifyReleaseState
getReleaseStateForCall(func::CallOp call,
llvm::DenseSet<Operation *> &activeCallees) {
func::FuncOp callee = lookupCallee(call);
if (!callee || callee.isExternal())
return {};
if (!activeCallees.insert(callee.getOperation()).second)
return {};

TNotifyReleaseState state =
collectTNotifyReleaseState(callee.getBody(), activeCallees);
activeCallees.erase(callee.getOperation());
state.remapPayloads([&](Value payload) {
return remapCalleePayloadToCallOperand(payload, callee, call);
});
return state;
}

static TNotifyReleaseState
getDirectTNotifyReleaseState(Operation *op,
llvm::DenseSet<Operation *> &activeCallees) {
if (isa<pto::BarrierOp, pto::CmoCacheInvalidOp, pto::FenceBarrierAllOp>(op))
return {};

if (auto call = dyn_cast<func::CallOp>(op))
return getReleaseStateForCall(call, activeCallees);

if (auto store = dyn_cast<pto::StoreScalarOp>(op)) {
if (isGmScalarMemory(store.getPtr().getType()))
return getCacheableGmStoreReleaseState(store.getPtr());
Expand Down Expand Up @@ -442,23 +502,37 @@ static TNotifyReleaseState getDirectTNotifyReleaseState(Operation *op) {
return {};
}

static TNotifyReleaseState collectTNotifyReleaseState(Operation *op) {
TNotifyReleaseState state = getDirectTNotifyReleaseState(op);
static TNotifyReleaseState getDirectTNotifyReleaseState(Operation *op) {
llvm::DenseSet<Operation *> activeCallees;
return getDirectTNotifyReleaseState(op, activeCallees);
}

static TNotifyReleaseState
collectTNotifyReleaseState(Operation *op,
llvm::DenseSet<Operation *> &activeCallees) {
TNotifyReleaseState state = getDirectTNotifyReleaseState(op, activeCallees);
for (Region &region : op->getRegions())
for (Block &block : region)
for (Operation &nested : block)
state.merge(collectTNotifyReleaseState(&nested));
state.merge(collectTNotifyReleaseState(&nested, activeCallees));
return state;
}

static TNotifyReleaseState collectTNotifyReleaseState(Region &region) {
static TNotifyReleaseState
collectTNotifyReleaseState(Region &region,
llvm::DenseSet<Operation *> &activeCallees) {
TNotifyReleaseState state;
for (Block &block : region)
for (Operation &nested : block)
state.merge(collectTNotifyReleaseState(&nested));
state.merge(collectTNotifyReleaseState(&nested, activeCallees));
return state;
}

static TNotifyReleaseState collectTNotifyReleaseState(Region &region) {
llvm::DenseSet<Operation *> activeCallees;
return collectTNotifyReleaseState(region, activeCallees);
}

static void applyFenceBarrierAllForSummary(pto::FenceBarrierAllOp fence,
TNotifyReleaseState &state) {
if (fence.getScope().getScope() != pto::FenceScope::GM &&
Expand Down Expand Up @@ -520,29 +594,12 @@ static bool isLoopLikeOp(Operation *op) {
return isa<scf::ForOp, scf::WhileOp, scf::ParallelOp, scf::ForallOp>(op);
}

static func::FuncOp lookupCallee(func::CallOp call) {
return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
call.getOperation(), call.getCalleeAttr());
}

static bool isMemoryConsistencyRelevantDirectOp(Operation *op) {
if (isa<pto::BarrierOp, pto::CmoCacheInvalidOp, pto::FenceBarrierAllOp, pto::TNotifyOp,
pto::TWaitOp, pto::TTestOp, pto::TLoadOp, pto::TPrefetchOp,
pto::TStoreOp, pto::TStoreFPOp>(op))
return true;

if (auto load = dyn_cast<pto::LoadScalarOp>(op))
return isGmScalarMemory(load.getPtr().getType());
if (auto store = dyn_cast<pto::StoreScalarOp>(op))
return isGmScalarMemory(store.getPtr().getType());

TNotifyReleaseState macroState = getReleaseStateForMacroModel(op);
return macroState.drainMte2 || macroState.drainMte3 ||
macroState.drainFix || macroState.needsDsbDdr ||
macroState.needsGmCacheCmo;
static bool isMemoryConsistencyBoundaryOp(Operation *op) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Narrowing the fail-fast to boundary-ops-only introduces a release/acquire asymmetry that can silently drop an acquire diagnostic.

Before this change, a non-inline callee containing a GM load_scalar/store_scalar was rejected here (the old isMemoryConsistencyRelevantDirectOp returned true for GM scalar load/store), which forced the helper to be inlined so the acquire walk could see its loads. Now such a callee is accepted, but only the release path was made interprocedural (getReleaseStateForCall summarizes callee payload producers into the caller). The acquire path was not: annotateSignalAcquire / collectSignalAcquireState / annotateSignalAcquireForBlock have no func.call handling, and every function is analyzed independently starting from an empty SignalAcquireState.

Consequence:

%ok = pto.comm.ttest ...      // caller acquire -> pendingInvalidateGmCache
func.call @helper(%ctx)        // @helper does a cacheable GM load_scalar

The cacheable load inside @helper is now reached after an acquire, but diagnoseAcquireLoad never fires for it: @helper is analyzed with an empty acquire state, and the caller's acquire state does not cross the call. Previously this program was rejected (forcing inline, after which the load would be diagnosed as needing pto.cmo.cacheinvalid all #pto.address_space<gm> before the cacheable load). So a program that should error now compiles silently -> potential stale GM read after acquire.

The motivating metadata helpers (load_scalar on CommContext) are legitimately exempt, but nothing constrains a now-allowed non-inline callee to metadata-only loads. Is this asymmetry intentional? If so, it would be good to pin it with a regression test and note it in the design doc; otherwise the acquire path likely needs the same interprocedural treatment as the release path. Either way it is currently untested.

return isa<pto::CmoCacheInvalidOp, pto::FenceBarrierAllOp, pto::TNotifyOp,
pto::TWaitOp, pto::TTestOp>(op);
}

static bool calleeContainsMemoryConsistencyRelevantOps(
static bool calleeContainsMemoryConsistencyBoundaryOps(
func::FuncOp callee, llvm::DenseSet<Operation *> &activeCallees) {
if (!callee || callee.isExternal())
return false;
Expand All @@ -555,13 +612,13 @@ static bool calleeContainsMemoryConsistencyRelevantOps(

if (auto nestedCall = dyn_cast<func::CallOp>(op)) {
func::FuncOp nestedCallee = lookupCallee(nestedCall);
if (calleeContainsMemoryConsistencyRelevantOps(nestedCallee,
activeCallees))
if (calleeContainsMemoryConsistencyBoundaryOps(nestedCallee,
activeCallees))
return WalkResult::interrupt();
return WalkResult::advance();
}

if (isMemoryConsistencyRelevantDirectOp(op))
if (isMemoryConsistencyBoundaryOp(op))
return WalkResult::interrupt();
return WalkResult::advance();
});
Expand Down Expand Up @@ -590,14 +647,14 @@ static bool diagnoseNonInlinedMemoryConsistencyCalls(ModuleOp module) {
return;

llvm::DenseSet<Operation *> activeCallees;
if (!calleeContainsMemoryConsistencyRelevantOps(callee, activeCallees))
if (!calleeContainsMemoryConsistencyBoundaryOps(callee, activeCallees))
return;

call.emitOpError()
<< "calls @" << callee.getSymName()
<< ", which contains PTO memory consistency relevant operations; "
<< ", which contains PTO memory consistency boundary operations; "
"inline the callee before `pto-memory-consistency` or keep "
"payload, CMO, fence, and signal operations in the caller";
"CMO, fence, and signal operations in the caller";
hasFailure = true;
});
}
Expand Down
49 changes: 9 additions & 40 deletions test/lit/pto/memory_consistency_noninline_call_invalid.pto
Original file line number Diff line number Diff line change
Expand Up @@ -8,54 +8,23 @@

// RUN: not ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s

// Non-inlined calls are not context-sensitive: a caller-side TNotify cannot
// safely observe release-relevant payload writes hidden in a callee body. Such
// callees must be inlined before the memory consistency pass.
// Non-inlined calls must not hide memory-consistency boundary operations. A
// caller-side release/acquire analysis cannot safely reason about CMO, fence,
// and signal ordering when those ops are hidden in a callee body.

module {
func.func private @producer(
%tile: !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
%dst_ptr: !pto.ptr<f32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
%dst_view = pto.make_tensor_view %dst_ptr,
shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view<?x?xf32>
%dst = pto.partition_view %dst_view,
offsets = [%c0, %c0], sizes = [%c1, %c32]
: !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<1x32xf32>
pto.tstore ins(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>)
outs(%dst : !pto.partition_tensor_view<1x32xf32>)
func.func private @hidden_boundary(%payload_ptr: !pto.ptr<i32>) {
pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr<i32>
pto.fence.barrier_all #pto.fence_scope<gm>
return
}

func.func @call_hidden_payload_write(
%dst_ptr: !pto.ptr<f32>,
%signal_ptr: !pto.ptr<i32>)
func.func @call_hidden_boundary(%payload_ptr: !pto.ptr<i32>)
attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
%v_i32 = arith.constant 1 : i32

%tile = pto.alloc_tile :
!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>

func.call @producer(%tile, %dst_ptr) :
(!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
!pto.ptr<f32>) -> ()
pto.fence.barrier_all #pto.fence_scope<gm>

%sig_view = pto.make_tensor_view %signal_ptr,
shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32>
%sig = pto.partition_view %sig_view,
offsets = [%c0], sizes = [%c1]
: !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32>
pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32)
{notifyOp = #pto<notify_op set>}
func.call @hidden_boundary(%payload_ptr) : (!pto.ptr<i32>) -> ()
return
}
}

// CHECK: calls @producer, which contains PTO memory consistency relevant operations
// CHECK: calls @hidden_boundary, which contains PTO memory consistency boundary operations
// CHECK: inline the callee before `pto-memory-consistency`
31 changes: 31 additions & 0 deletions test/lit/pto/memory_consistency_noninline_metadata_call.pto
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// 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 ANY 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.

// A non-inline helper that only reads communication metadata should not be
// treated as hidden payload/signal state.

// RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s

module {
func.func private @metadata_offset(%ctx: !pto.ptr<i32>) -> i32 {
%c0 = arith.constant 0 : index
%off = pto.load_scalar %ctx[%c0] : !pto.ptr<i32> -> i32
return %off : i32
}

func.func @metadata_helper_call_is_allowed(%ctx: !pto.ptr<i32>)
attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
%c0 = arith.constant 0 : index
%off = func.call @metadata_offset(%ctx) : (!pto.ptr<i32>) -> i32
pto.store_scalar %off, %ctx[%c0] : !pto.ptr<i32>, i32
return
}
}

// CHECK-NOT: error:
// CHECK: AICORE void metadata_helper_call_is_allowed(
64 changes: 64 additions & 0 deletions test/lit/pto/memory_consistency_noninline_release_call.pto
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
// 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 ANY 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.

// The callee hides a payload write, but the caller still provides the payload
// marker explicitly. The memory consistency pass should keep the boundary and
// infer the required drain from the marker instead of blindly rejecting the
// call because of the hidden load/store body.

// RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s

module {
func.func private @producer(
%tile: !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
%dst_ptr: !pto.ptr<f32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
%dst_view = pto.make_tensor_view %dst_ptr,
shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view<?x?xf32>
%dst = pto.partition_view %dst_view,
offsets = [%c0, %c0], sizes = [%c1, %c32]
: !pto.tensor_view<?x?xf32> -> !pto.partition_tensor_view<1x32xf32>
pto.tstore ins(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>)
outs(%dst : !pto.partition_tensor_view<1x32xf32>)
return
}

func.func @call_hidden_payload_write_with_marker(
%dst_ptr: !pto.ptr<f32>,
%signal_ptr: !pto.ptr<i32>)
attributes {pto.kernel_kind = #pto.kernel_kind<vector>} {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%v_i32 = arith.constant 1 : i32

%tile = pto.alloc_tile :
!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>

func.call @producer(%tile, %dst_ptr) :
(!pto.tile_buf<loc=vec, dtype=f32, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>,
!pto.ptr<f32>) -> ()

pto.cmo.cacheinvalid %dst_ptr single_cache_line : !pto.ptr<f32>
pto.fence.barrier_all #pto.fence_scope<gm>

%sig_view = pto.make_tensor_view %signal_ptr,
shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32>
%sig = pto.partition_view %sig_view,
offsets = [%c0], sizes = [%c1]
: !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32>
pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32)
{notifyOp = #pto<notify_op set>}
return
}
}

// CHECK: pipe_barrier(PIPE_MTE3)
// CHECK: dsb(
// CHECK: pto::comm::TNOTIFY(
Loading