Fix memory consistency non-inline metadata helpers

docs/PTO_IR_manual.md

@@ -8713,6 +8713,89 @@ pto.barrier_sync [<TVEC>]
 
 ---
 
+##### `pto.cmo.cacheinvalid`
+
+**Summary:** Explicit GM cache-maintenance op and payload marker.
+
+**Semantics:**
+
+```text
+cacheinvalid(addr?, space)
+```
+
+**Forms:**
+
+- Whole-cache form: `pto.cmo.cacheinvalid all #pto.address_space<gm>`
+- Single-cache-line form: `pto.cmo.cacheinvalid %addr single_cache_line`
+
+**Arguments:**
+
+| Name | Type | Description |
+|------|------|-------------|
+| `addr` | optional GM address-like value | Optional payload address. In practice this can be a GM `!pto.ptr<...>` or a GM `!pto.partition_tensor_view<...>` value. Omit it for whole-cache form. |
+| `space` | `AddressSpaceAttr` | Memory space selector, usually `#pto.address_space<gm>` |
+
+**Results:** None.
+
+**Constraints & Verification:**
+
+- `space` must name a valid address space.
+- The whole-cache form is the conservative GM cache-maintenance form.
+- The single-cache-line form is used either as a real single-line CMO or as a payload marker for memory-consistency analysis.
+
+**Hardware Mapping:**
+
+- `pto.cmo.cacheinvalid all #pto.address_space<gm>` lowers to `dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE)`.
+- `pto.cmo.cacheinvalid %addr single_cache_line` lowers to `dcci((__gm__ void*)addr, cache_line_t::SINGLE_CACHE_LINE)` when the path is cacheable.
+- On non-cacheable payload paths, the single-line form may be consumed as marker-only by later passes.
+
+**Basic Examples:**
+
+```mlir
+pto.cmo.cacheinvalid all #pto.address_space<gm>
+pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr<f32>
+pto.cmo.cacheinvalid %payload_view single_cache_line : !pto.partition_tensor_view<128xf32>
+```
+
+---
+
+##### `pto.fence.barrier_all`
+
+**Summary:** Visibility fence for the selected memory scope.
+
+**Semantics:**
+
+```text
+barrier_all(scope)
+```
+
+**Arguments:**
+
+| Name | Type | Description |
+|------|------|-------------|
+| `scope` | `FenceScopeAttr` | Visibility scope: `local_memory`, `gm`, or `all` |
+
+**Results:** None.
+
+**Constraints & Verification:**
+
+- `local_memory` is reserved for local-buffer semantics.
+- `gm` and `all` are the scopes currently used for GM visibility ordering.
+
+**Hardware Mapping:**
+
+- `pto.fence.barrier_all #pto.fence_scope<gm>` lowers to `dsb(DSB_DDR)`.
+- `pto.fence.barrier_all #pto.fence_scope<all>` currently lowers to the same `dsb(DSB_DDR)` form.
+
+**Basic Examples:**
+
+```mlir
+pto.fence.barrier_all #pto.fence_scope<gm>
+pto.fence.barrier_all #pto.fence_scope<all>
+```
+
+---
+
 ##### `pto.record_event`
 
 **Summary:** Records an event for synchronization between producer and consumer operation classes.
@@ -10423,10 +10506,10 @@ pto.local_array_set %m[%i, %j], %v
 | Integer Sequence Generation | 1 | V |
 | Scalar Element Access | 2 | V |
 | MX Quantized | 3 | M/V |
-| Synchronization | 5 | - |
+| Synchronization | 7 | - |
 | CV-Related | 2 | - |
 | Runtime Intrinsics | 4 | - (Pure) |
 | Debug | 3 | - |
 | Stack-Local Array | 3 | - (Scalar / Host) |
 
-**Total: 110 operations**
+**Total: 112 operations**

docs/designs/ptoas-memory-consistency-design.md

@@ -266,10 +266,16 @@ PyPTO 生成规则：
 marker。PTOAS 不会尝试从所有前序 memory op 里推断“可能相关”的 payload。
 
 `pto.entry` launcher 可以调用多个 kernel 函数；每个 kernel 函数会被
-`pto-memory-consistency` 独立分析。kernel body 内部若通过 `func.call` 调用包含 payload
-访问、CMO、fence 或 signal op 的 helper，PyPTO 应在 `pto-memory-consistency` 前将 helper
-inline，或者把 payload、CMO、fence 和 signal 保持在同一个 caller 中。否则 pass 会报错，
-避免 caller 侧 `TNotify` 或 `TWait` 看不到 callee 内部的 memory-consistency 状态。
+`pto-memory-consistency` 独立分析。kernel body 内部可以通过 `func.call` 调用只包含
+payload 访问的 helper，PTOAS 会把 callee 中的 payload producer 摘要回 caller，并继续由
+caller 侧的 `cmo.cacheinvalid` marker、`fence.barrier_all` 和 signal op 决定 release
+边界。如果 helper 内部已经插入 `pto.barrier`，callee 摘要也会消费对应 pipe drain，避免
+caller 侧重复补同一个 pipe barrier。
+
+但 CMO、fence、`TNotify`、`TWait` 和 `TTest` 这类 memory-consistency boundary op 不能藏在
+non-inline helper 中。PyPTO 应在 `pto-memory-consistency` 前将包含 boundary op 的 helper
+inline，或者把这些 boundary op 保持在 caller 中；否则 pass 会报错，避免 caller 侧看不到
+callee 内部的 release 或 acquire 边界。
 
 ### 6.1 Issue #872：TPUT 发布 Signal
 

lib/PTO/Transforms/PTOMemoryConsistency.cpp

@@ -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 {
@@ -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 {
@@ -408,10 +420,67 @@ 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 TNotifyReleaseState
+getTNotifyReleaseExitStateForBlock(Block &block,
+                                   TNotifyReleaseState pendingState,
+                                   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 =
+      callee.getBody().hasOneBlock()
+          ? getTNotifyReleaseExitStateForBlock(callee.getBody().front(),
+                                               TNotifyReleaseState{},
+                                               activeCallees)
+          : 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());
@@ -442,23 +511,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 &&
@@ -474,28 +557,26 @@ static void applyFenceBarrierAllForSummary(pto::FenceBarrierAllOp fence,
   state.applyFenceBarrierAll(fence.getScope().getScope());
 }
 
-static TNotifyReleaseState getTNotifyReleaseExitStateForBlock(
-    Block &block, TNotifyReleaseState pendingState);
-
 static TNotifyReleaseState
-getTNotifyReleaseExitState(Operation *op,
-                           TNotifyReleaseState pendingState = {}) {
+getTNotifyReleaseExitState(Operation *op, TNotifyReleaseState pendingState,
+                           llvm::DenseSet<Operation *> &activeCallees) {
   if (isa<pto::TNotifyOp>(op))
     pendingState.clear();
 
-  pendingState.merge(getDirectTNotifyReleaseState(op));
+  pendingState.merge(getDirectTNotifyReleaseState(op, activeCallees));
 
   TNotifyReleaseState regionEntryState = pendingState;
   TNotifyReleaseState combinedRegionExitState;
   for (Region &region : op->getRegions()) {
     if (region.hasOneBlock()) {
       combinedRegionExitState.merge(
-          getTNotifyReleaseExitStateForBlock(region.front(), regionEntryState));
+          getTNotifyReleaseExitStateForBlock(region.front(), regionEntryState,
+                                             activeCallees));
       continue;
     }
 
     TNotifyReleaseState regionExitState = regionEntryState;
-    regionExitState.merge(collectTNotifyReleaseState(region));
+    regionExitState.merge(collectTNotifyReleaseState(region, activeCallees));
     combinedRegionExitState.merge(regionExitState);
   }
   pendingState.merge(combinedRegionExitState);
@@ -509,40 +590,31 @@ getTNotifyReleaseExitState(Operation *op,
   return pendingState;
 }
 
-static TNotifyReleaseState getTNotifyReleaseExitStateForBlock(
-    Block &block, TNotifyReleaseState pendingState) {
+static TNotifyReleaseState
+getTNotifyReleaseExitState(Operation *op, TNotifyReleaseState pendingState = {}) {
+  llvm::DenseSet<Operation *> activeCallees;
+  return getTNotifyReleaseExitState(op, pendingState, activeCallees);
+}
+
+static TNotifyReleaseState
+getTNotifyReleaseExitStateForBlock(Block &block,
+                                   TNotifyReleaseState pendingState,
+                                   llvm::DenseSet<Operation *> &activeCallees) {
   for (Operation &op : block)
-    pendingState = getTNotifyReleaseExitState(&op, pendingState);
+    pendingState = getTNotifyReleaseExitState(&op, pendingState, activeCallees);
   return pendingState;
 }
 
 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) {
+  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;
@@ -555,13 +627,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();
   });
@@ -590,14 +662,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;
     });
   }