diff --git a/docs/en/dev/ir/05-operators.md b/docs/en/dev/ir/05-operators.md index 424a866dfc..44c0109028 100644 --- a/docs/en/dev/ir/05-operators.md +++ b/docs/en/dev/ir/05-operators.md @@ -41,6 +41,113 @@ auto dynamic_dim = make_int(kDynamicDim); | `f_deduce_type(fn)` | Type deduction function | `.f_deduce_type(DeduceAddType)` | | `set_core_affinity(a)` | Which core executes the op (**placement**) | `.set_core_affinity(core_affinity::CoreAffinity::VECTOR)` | | `set_no_duplicate()` | Op must not run on a second core (**replication**) | `.set_no_duplicate()` | +| `set_arg_effect(i, e)` | What the op does to argument `i`'s buffer | `.set_arg_effect(2, ArgEffect::Write)` | +| `set_arg_effect(i, fn)` | Same, when a kwarg decides it | `.set_arg_effect(2, [](const auto& kw) { ... })` | +| `no_arg_writes()` | Classified: writes through no argument | `.no_arg_writes()` | +| `set_write_channel(c)` | Hardware path the op's writes take | `.set_write_channel(WriteChannel::Dma)` | + +### Argument effects + +An operator that updates one of its arguments in place must say so. Direction +inference, dependency analysis and the parameter-direction verifier all ask the +registry the same question — *does this call write the buffer this argument +names?* — and an operator that never answered reads as a pure consumer: + +```text +tile.mscatter writes output_tensor, but never declared it + → the parameter it writes keeps direction In + → no RAW edge is emitted against the kernel that reads it + → the scheduler is free to run the reader first + → stale data, or a deadlock waiting on a signal nobody wrote +``` + +| Effect | Meaning | Examples | +| ------ | ------- | -------- | +| `ArgEffect::Read` | Read, never written. Default for an unnamed argument | `tile.store`'s source tile | +| `ArgEffect::Write` | Overwritten without being read first | `tile.store`'s `output_tensor`, `pld.tile.get`'s `dst` | +| `ArgEffect::ReadWrite` | Read *and* written | `tile.matmul_acc`'s accumulator, an atomic store's destination | + +**Partial overwrite is still `Write`.** A store that lands on a sub-region does +not read the untouched remainder — nothing moves *into* the kernel — so its +destination is a pure write. Declaring it `ReadWrite` is not a harmless +approximation: it makes the enclosing parameter `InOut`, which stages the buffer +host→device and, across ranks, invents a dependency between two ranks writing +disjoint rows. + +Whether a destination is read is decided per operator, not per family: a +gather or exchange destination is pushed into and never loaded from, so it is +`Write`, while a reduce destination has its running value loaded back and is +`ReadWrite`. + +**`ReadWrite` is for operators that genuinely read the slot**: an accumulator +(`out += x` reads the running sum), an atomic store or assemble, or a +destination-passing operator whose untouched positions flow through to its SSA +result (`tile.scatter`, `array.update_element`). + +**Kwarg-dependent effects.** When a kwarg decides the answer, pass a resolver +instead of a constant. A kwarg can decide *whether* an argument is written at +all, not only how: `tile.mgather`'s third operand is a written GM scratch tensor +in Mat element mode and a read-only `valid_shape` in Mat row mode. The other +live cases are the `atomic` kwarg on the store family and the `op` kwarg on +`pld.system.notify`, whose default is atomic-add — so an unannotated notify +reads the slot it adds into: + +```cpp +REGISTER_OP("tile.store") + // ... arguments, memory spec ... + .set_arg_effect(2, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) +``` + +**Declared-read-only is not the same as unclassified.** `HasDeclaredArgEffects()` +distinguishes "a human decided this operator writes nothing" (`no_arg_writes()`, +e.g. `pld.system.wait`) from "nobody has looked at this operator yet". An +analysis that needs the answer can then refuse to guess instead of defaulting an +unclassified writer to read-only. + +**Enforcement.** `OpRegistry::ValidateArgEffects()` runs at import and rejects +two shapes, naming every offender and the fix rather than failing on first use: + +- an operator declaring `set_output_reuses_input(N)` — its SSA result *is* + argument N's buffer, so it writes through it — without a verdict about + argument N specifically. Classification is what is required, not a particular + answer: an operator whose in-place slot is metadata may declare it read-only. +- an operator declaring a write channel while writing through no argument. A + channel says *how* an operator writes, so one without a write is either a + stray declaration or a missing one. + +The second rule matters more than it looks. `set_write_channel()` creates the +effect spec as a side effect, so "the spec exists" cannot stand in for "a human +decided" — otherwise an operator that declared a channel and forgot its +`set_arg_effect` would pass the first rule with the argument it updates still +defaulting to `Read`. `no_arg_writes()` records the all-arguments verdict +explicitly, and combining it with `set_arg_effect` is rejected as +contradictory. + +`set_write_channel` records whether the writes travel the MTE3/DMA path or the +scalar D-cache path. PyPTO cannot order the two against one GM tensor, so a +function mixing them on one buffer is rejected; the channel lets that diagnostic +read the registry instead of re-listing operators. + +Declare it only for an operator whose writes really are one of those two paths, +and leave it unset otherwise — an unset channel keeps the operator out of that +diagnostic, which is where an operator belongs when neither path describes it: + +- `pld.system.notify` emits `pto.comm.tnotify`, a distinct comm instruction. + Claiming either channel would let the diagnostic reject a valid program. +- `system.set_ffts` hands the workspace *pointer* to the FFTS unit rather than + moving data; the hardware writes that region on its own schedule, which no + dependency edge models. It declares `no_arg_writes()`. +- A composite collective updates a data window and a signal through different + mechanisms, and one operator-level channel cannot describe both. Per-argument + channels would, but no case yet needs the distinction, and a wrong single + answer is worse than none. **`set_core_affinity` vs `set_no_duplicate`** — two orthogonal axes, and picking the wrong one makes a false claim about the ISA: diff --git a/docs/zh/dev/ir/05-operators.md b/docs/zh/dev/ir/05-operators.md index 8d5ea2511f..5863915dd9 100644 --- a/docs/zh/dev/ir/05-operators.md +++ b/docs/zh/dev/ir/05-operators.md @@ -41,6 +41,94 @@ auto dynamic_dim = make_int(kDynamicDim); | `f_deduce_type(fn)` | 类型推导函数 | `.f_deduce_type(DeduceAddType)` | | `set_core_affinity(a)` | 算子在哪个核上执行(**放置**) | `.set_core_affinity(core_affinity::CoreAffinity::VECTOR)` | | `set_no_duplicate()` | 算子不得在第二个核上运行(**复制**) | `.set_no_duplicate()` | +| `set_arg_effect(i, e)` | 算子对第 `i` 个参数缓冲区做了什么 | `.set_arg_effect(2, ArgEffect::Write)` | +| `set_arg_effect(i, fn)` | 同上,但由 kwarg 决定 | `.set_arg_effect(2, [](const auto& kw) { ... })` | +| `no_arg_writes()` | 已分类:不通过任何参数写入 | `.no_arg_writes()` | +| `set_write_channel(c)` | 算子写入所走的硬件通路 | `.set_write_channel(WriteChannel::Dma)` | + +### 参数效应(Argument effects) + +原地更新某个参数的算子必须显式声明。方向推导(direction inference)、依赖分析 +(dependency analysis)和参数方向验证器都向注册表询问同一个问题——*这次调用是否 +写入该参数所指的缓冲区?*——而从未回答过的算子会被读成纯消费者: + +```text +tile.mscatter 写 output_tensor,却从未声明 + → 它写入的参数方向停留在 In + → 不会对读取它的 kernel 发出 RAW 边 + → 调度器可以先运行读方 + → 读到陈旧数据,或等待一个无人写入的信号而死锁 +``` + +| 效应 | 含义 | 示例 | +| ---- | ---- | ---- | +| `ArgEffect::Read` | 只读,从不写入。未声明参数的默认值 | `tile.store` 的源 tile | +| `ArgEffect::Write` | 覆盖写,写前不读 | `tile.store` 的 `output_tensor`、`pld.tile.get` 的 `dst` | +| `ArgEffect::ReadWrite` | 既读*又*写 | `tile.matmul_acc` 的累加器、原子 store 的目的操作数 | + +**部分覆盖仍然是 `Write`。** 只写入子区域的 store 不会读取未触及的其余部分——没有 +任何数据流*进*内核——所以其目的操作数是纯写。把它声明成 `ReadWrite` 并非无害的保守 +近似:这会让外层参数变成 `InOut`,从而触发 host→device 搬运,并且在跨 rank 场景下 +为两个写入不相交行的 rank 凭空造出一条依赖。 + +目的操作数是否被读取按**算子**判定,而非按家族:gather / exchange 的目的窗口只被推入、 +从不被 load,因此是 `Write`;而 reduce 的目的窗口会把运行值 load 回来,因此是 `ReadWrite`。 + +**`ReadWrite` 留给真正会读取该槽位的算子**:累加器(`out += x` 会读取运行中的和)、 +原子 store/assemble,或那些未触及位置会流入 SSA 结果的 destination-passing 算子 +(`tile.scatter`、`array.update_element`)。 + +**由 kwarg 决定的效应。** 当答案取决于某个 kwarg 时,传入 resolver 而非常量。kwarg 不 +仅能决定*怎么*写,还能决定某个实参*是否*被写:`tile.mgather` 的第三个操作数在 Mat elem +模式下是被写的 GM scratch,在 Mat row 模式下则是只读的 `valid_shape`。另外两处是 store +家族的 `atomic` kwarg,以及 `pld.system.notify` 的 `op` kwarg——后者默认是 atomic-add, +因此未加标注的 notify 会读取它累加的槽位: + +```cpp +REGISTER_OP("tile.store") + // ... arguments, memory spec ... + .set_arg_effect(2, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) +``` + +**"已声明为只读"不等于"未分类"。** `HasDeclaredArgEffects()` 区分"有人判定该算子不 +写入任何参数"(`no_arg_writes()`,例如 `pld.system.wait`)与"还没有人看过这个算子"。 +需要答案的分析因而可以拒绝猜测,而不是把未分类的写者默认成只读。 + +**强制约束。** `OpRegistry::ValidateArgEffects()` 在 import 时运行,拒绝两种形态, +一次性列出所有违规算子和修复方式,而不是等到首次使用才失败: + +- 声明了 `set_output_reuses_input(N)`(其 SSA 结果*就是*第 N 个参数的缓冲区,因此会 + 通过它写入),却没有专门对第 N 个参数作出裁决。要求的是"做出分类"而非某个特定答案: + 原地槽位是元数据的算子可以声明为只读。 +- 声明了写通路却不通过任何参数写入。通路描述的是"怎么写",因此没有写的通路要么是多余 + 声明,要么是漏了声明。 + +第二条比看上去重要。`set_write_channel()` 会顺带创建效应 spec,因此"spec 存在"不能 +等同于"有人做过判断"——否则一个声明了通路却忘了 `set_arg_effect` 的算子会通过第一条 +检查,而它原地更新的那个参数仍然默认为 `Read`。`no_arg_writes()` 显式记录"对所有参数 +的裁决",且与 `set_arg_effect` 同时使用会被判为自相矛盾而拒绝。 + +`set_write_channel` 记录写入走的是 MTE3/DMA 通路还是标量 D-cache 通路。PyPTO 无法为 +同一个 GM tensor 排序这两者,因此会拒绝在同一缓冲区上混用两者的函数;有了通路声明, +该诊断可以查询注册表而不必再列一遍算子清单。 + +只为写入确实走这两条通路之一的算子声明它,其余一律留空——留空会把该算子排除在该诊断 +之外,而当两条通路都无法描述它时,这正是它该待的位置: + +- `pld.system.notify` 发出的是 `pto.comm.tnotify`,一条独立的 comm 指令。声明任一通路 + 都会让该诊断拒绝合法程序。 +- `system.set_ffts` 是把 workspace *指针*交给 FFTS 单元,而非搬运数据;该区域由硬件按 + 自己的节奏写入,没有任何依赖边能建模它。它声明 `no_arg_writes()`。 +- 复合集合通信通过不同机制更新数据窗口与 signal,单个算子级通路无法同时描述两者。 + 按参数记录通路可以做到,但目前没有任何用例需要这种区分,而一个错误的单一答案比没有 + 答案更糟。 **`set_core_affinity` 与 `set_no_duplicate`** —— 两个正交的维度,选错会对 ISA 做出错误的断言: diff --git a/include/pypto/ir/op_registry.h b/include/pypto/ir/op_registry.h index 842d62042e..ded1e2d1fc 100644 --- a/include/pypto/ir/op_registry.h +++ b/include/pypto/ir/op_registry.h @@ -21,10 +21,12 @@ #ifndef PYPTO_IR_OP_REGISTRY_H_ #define PYPTO_IR_OP_REGISTRY_H_ +#include #include #include #include #include +#include #include #include #include @@ -92,6 +94,133 @@ enum class ExecutionMemoryAccessEvidence : uint8_t { NoAccess, }; +/** + * @brief What executing an operation does to the buffer one argument names. + * + * This is an *effect* declaration, not a type. It answers the single question + * every direction and dependency analysis asks: does running this call read + * from, or write to, the memory this argument names? + * + * `Read` is the default for an argument the operator does not name, because + * the overwhelming majority of operators are functional — they consume their + * operands and produce a fresh SSA result. An operator that instead updates an + * operand in place must say so: a missing `Write` is not a conservative + * approximation, it silently erases a real dependency edge (the writer looks + * like a pure reader, so nothing is ordered against it). + */ +enum class ArgEffect : uint8_t { + Read = 0, ///< Read, never written. + Write = 1, ///< Overwritten without being read first (a destination operand). + ReadWrite = 2, ///< Read *and* written — accumulate, atomic update, partial in-place rewrite. +}; + +/// Merge two independent observations of one argument's effect. Each is a lower +/// bound on the accesses, so the merge is a union along Read < {Write} < ReadWrite. +[[nodiscard]] inline ArgEffect MergeArgEffect(ArgEffect lhs, ArgEffect rhs) { + if (lhs == rhs) return lhs; + return ArgEffect::ReadWrite; +} + +[[nodiscard]] inline bool ArgEffectReads(ArgEffect effect) { return effect != ArgEffect::Write; } +[[nodiscard]] inline bool ArgEffectWrites(ArgEffect effect) { return effect != ArgEffect::Read; } + +[[nodiscard]] inline std::string ArgEffectToString(ArgEffect effect) { + switch (effect) { + case ArgEffect::Read: + return "Read"; + case ArgEffect::Write: + return "Write"; + case ArgEffect::ReadWrite: + return "ReadWrite"; + } + return "Unknown"; +} + +/// Read an integer-valued kwarg out of a call's kwargs, or `fallback` when the +/// call does not carry it. Effect resolvers use this to branch on the enum-backed +/// int kwargs (`atomic`, `op`, ...) that decide whether a destination operand is +/// overwritten or accumulated into. +[[nodiscard]] inline int GetIntKwarg(const std::vector>& kwargs, + const std::string& key, int fallback) { + for (const auto& [k, v] : kwargs) { + if (k == key) return AnyCast(v, key); + } + return fallback; +} + +/// MemorySpace-valued counterpart of GetIntKwarg. A memory-space kwarg is stored +/// as a `MemorySpace`, not as an int (see `ConvertKwargsDict`), so it needs its +/// own accessor — `tile.mgather` selects which operand is its GM scratch from +/// `target_memory`. +[[nodiscard]] inline MemorySpace GetMemorySpaceKwarg( + const std::vector>& kwargs, const std::string& key, + MemorySpace fallback) { + for (const auto& [k, v] : kwargs) { + if (k == key) return AnyCast(v, key); + } + return fallback; +} + +/// String-valued counterpart of GetIntKwarg, for the kwargs that select an +/// operator mode by name (`system.syncall`'s hard/soft form). +[[nodiscard]] inline std::string GetStringKwarg(const std::vector>& kwargs, + const std::string& key, const std::string& fallback) { + for (const auto& [k, v] : kwargs) { + if (k == key) return AnyCast(v, key); + } + return fallback; +} + +/** + * @brief The hardware path an operation's writes travel. + * + * PyPTO cannot order an MTE3 (DMA) store against a scalar D-cache write to the + * same GM tensor, so a function that mixes both on one buffer is rejected. The + * channel is a property of the operator's lowering, declared once here rather + * than re-derived by the diagnostic. + */ +enum class WriteChannel : uint8_t { + Dma, ///< MTE3 / DMA store path (tile.store, tensor.assemble, cross-rank put/get). + Scalar, ///< Scalar D-cache write path (tensor.write). +}; + +/** + * @brief Per-argument execution effects declared by one operator. + * + * Absent (`std::nullopt` on the entry) means *nobody has classified this + * operator yet* — which is deliberately distinct from "declared read-only", so + * an analysis can refuse to guess instead of defaulting an unclassified writer + * to read-only. + */ +struct OpArgEffectSpec { + /// Effect resolved from the call's kwargs, for arguments whose effect is not + /// fixed by the operator alone (an atomic store reads its destination; a + /// `NotifyOp::kAtomicAdd` accumulates into the peer slot). + using Resolver = std::function>& kwargs)>; + + /// Effect per positional argument index. Indices past the end are `Read`. + std::vector per_arg; + + /// Kwarg-dependent effects, keyed by argument index. Overrides `per_arg`. + std::map kwarg_dependent; + + /// Which path this operator's writes take. Set only for operators that write. + std::optional write_channel; + + /// Argument indices a registration named explicitly. `per_arg` cannot answer + /// this on its own: it is resized to *cover* the highest declared index, so a + /// slot nobody named is indistinguishable there from one declared `Read`. + /// Validation needs the difference — an operator that classified the wrong + /// argument has not classified the one it updates in place. + std::set declared_args; + + /// True only when a registration called `no_arg_writes()`, which is a verdict + /// about every argument at once. The spec's mere existence cannot stand in for + /// this: `set_write_channel()` creates it too, and an operator that declared + /// only a channel has classified nothing. + bool declared_no_writes = false; +}; + /** * @brief Type-erased operator registration entry * @@ -556,6 +685,106 @@ class OpRegistryEntry { return cross_core_role_; } + /// Declare what executing this operator does to positional argument + /// `arg_index`. Every argument the operator does not name is `Read`; calling + /// this at all marks the operator *classified*, so a later analysis can tell + /// "reads everything" apart from "nobody looked yet". + /// + /// Declare the argument that carries the destination, not the one that + /// carries the data: `tile.store(tile, offsets, output_tensor)` writes + /// argument 2. + inline OpRegistryEntry& set_arg_effect(size_t arg_index, ArgEffect effect) { + auto& effects = EnsureArgEffects(); + CHECK(!effects.declared_no_writes) << "Operator '" << name_ << "' names argument " << arg_index + << " after declaring no_arg_writes(); the two are contradictory"; + if (effects.per_arg.size() <= arg_index) { + effects.per_arg.resize(arg_index + 1, ArgEffect::Read); + } + effects.per_arg[arg_index] = effect; + effects.declared_args.insert(arg_index); + return *this; + } + + /// Declare an argument whose effect the operator alone does not fix, because + /// a kwarg decides it — an atomic `tile.store` reads the accumulator it adds + /// into, while a plain one overwrites it. The resolver sees the call's kwargs + /// and must return the effect for that call. + inline OpRegistryEntry& set_arg_effect(size_t arg_index, OpArgEffectSpec::Resolver resolver) { + CHECK(resolver) << "Operator '" << name_ << "' argument " << arg_index + << " was given a null effect resolver"; + auto& effects = EnsureArgEffects(); + CHECK(!effects.declared_no_writes) << "Operator '" << name_ << "' names argument " << arg_index + << " after declaring no_arg_writes(); the two are contradictory"; + effects.kwarg_dependent[arg_index] = std::move(resolver); + effects.declared_args.insert(arg_index); + return *this; + } + + /// Declare that this operator writes through none of its arguments. Use it to + /// classify an operator whose name or side-effect-only signature would + /// otherwise leave a reader wondering — `pld.system.wait` polls a signal it + /// never writes. + inline OpRegistryEntry& no_arg_writes() { + auto& effects = EnsureArgEffects(); + CHECK(effects.declared_args.empty()) + << "Operator '" << name_ + << "' declares no_arg_writes() after naming an argument; the two are contradictory"; + effects.declared_no_writes = true; + return *this; + } + + /// Declare which hardware path this operator's writes take. Required for an + /// operator that writes a GM tensor, so the mixed-store diagnostic can tell + /// an MTE3 store from a scalar one without re-listing operators. + inline OpRegistryEntry& set_write_channel(WriteChannel channel) { + auto& effects = EnsureArgEffects(); + CHECK(!effects.write_channel.has_value()) << "Operator '" << name_ << "' write channel is already set"; + effects.write_channel = channel; + return *this; + } + + /// True when this operator declared its per-argument effects. False means the + /// operator has never been classified — an analysis that needs the answer must + /// say so loudly rather than assume read-only. + [[nodiscard]] bool HasDeclaredArgEffects() const { return arg_effects_.has_value(); } + + /// True when the registration reached a verdict about argument `arg_index` in + /// particular: it named that argument, or it declared with `no_arg_writes()` + /// that the operator writes through none of them. An operator that named some + /// *other* argument — or that only set a write channel, which creates the spec + /// as a side effect — has not decided about this one, and reading the resulting + /// `Read` as a decision is what this distinguishes. + [[nodiscard]] bool HasDeclaredArgEffect(size_t arg_index) const { + if (!arg_effects_.has_value()) return false; + if (arg_effects_->declared_no_writes) return true; + return arg_effects_->declared_args.count(arg_index) > 0; + } + + /// The effect on positional argument `arg_index` for a call carrying `kwargs`. + /// `Read` for any argument the operator did not name. + [[nodiscard]] ArgEffect GetArgEffect(size_t arg_index, + const std::vector>& kwargs) const { + if (!arg_effects_.has_value()) return ArgEffect::Read; + auto resolver = arg_effects_->kwarg_dependent.find(arg_index); + if (resolver != arg_effects_->kwarg_dependent.end()) return resolver->second(kwargs); + if (arg_index >= arg_effects_->per_arg.size()) return ArgEffect::Read; + return arg_effects_->per_arg[arg_index]; + } + + /// True when this operator writes through at least one argument under some + /// kwargs. Cheap pre-filter for analyses that only care about writers. + [[nodiscard]] bool WritesAnyArg() const { + if (!arg_effects_.has_value()) return false; + if (!arg_effects_->kwarg_dependent.empty()) return true; + return std::any_of(arg_effects_->per_arg.begin(), arg_effects_->per_arg.end(), ArgEffectWrites); + } + + /// The hardware path this operator's writes take, or nullopt when it declared + /// none (either it writes nothing, or its writes are not GM stores). + [[nodiscard]] std::optional GetWriteChannel() const { + return arg_effects_.has_value() ? arg_effects_->write_channel : std::nullopt; + } + inline OpRegistryEntry& set_internal_only(bool value = true) { internal_only_ = value; return *this; @@ -578,6 +807,15 @@ class OpRegistryEntry { } } + /// The effect spec, creating it when this is the operator's first declaration. + /// Returns a reference rather than leaving callers to dereference the optional: + /// engagement is obvious here and provable to a reader (and to clang-tidy's + /// unchecked-optional-access analysis) only at this one site. + OpArgEffectSpec& EnsureArgEffects() { + if (!arg_effects_.has_value()) return arg_effects_.emplace(); + return *arg_effects_; + } + /** * @brief Set the operator name * @@ -603,6 +841,7 @@ class OpRegistryEntry { const std::vector>&)>> deduce_type_; ///< Type deduction function std::optional memory_spec_; ///< Memory space specification + std::optional arg_effects_; ///< Per-argument execution effects; nullopt = unclassified bool is_inplace_safe_{true}; ///< Whether the op supports in-place execution (src == dst buffer) ExecutionMemoryAccessEvidence execution_memory_access_evidence_{ExecutionMemoryAccessEvidence::Unknown}; std::set forbid_output_alias_args_; ///< Input args whose buffer the output must not reuse @@ -754,6 +993,27 @@ class OpRegistry { */ void ValidateTileOps() const; + /** + * @brief Validate that every operator which updates an argument in place has + * declared its per-argument effects. + * + * An operator declaring `set_output_reuses_input(N)` writes through argument + * N — that is what reusing the buffer means. Direction inference, dependency + * analysis and the parameter-direction verifier all read those effects, and + * an undeclared operator reads as a pure consumer: the write vanishes, no + * dependency edge is emitted, and the failure surfaces on device as a race + * or a deadlock rather than at compile time. + * + * Classification, not a particular answer, is what is required: an operator + * whose in-place slot is metadata rather than data may declare it `Read` (via + * `no_arg_writes()`), which records that a human decided. + * + * Call at module init to catch an unclassified operator at import time. + * + * @throws ValueError listing every in-place operator with undeclared effects + */ + void ValidateArgEffects() const; + private: OpRegistry() = default; ~OpRegistry() = default; diff --git a/python/bindings/bindings.cpp b/python/bindings/bindings.cpp index 64d1c17199..1ebba5e4df 100644 --- a/python/bindings/bindings.cpp +++ b/python/bindings/bindings.cpp @@ -71,4 +71,8 @@ NB_MODULE(pypto_core, m) { // Validate that all tile.* ops have memory specs — fails at import time if any are missing pypto::ir::OpRegistry::GetInstance().ValidateTileOps(); + + // Validate that every in-place op declared what it does to the slot it updates — + // an undeclared writer reads as a pure consumer and its dependency edge vanishes + pypto::ir::OpRegistry::GetInstance().ValidateArgEffects(); } diff --git a/python/bindings/modules/ir.cpp b/python/bindings/modules/ir.cpp index 259668141d..683443ce5c 100644 --- a/python/bindings/modules/ir.cpp +++ b/python/bindings/modules/ir.cpp @@ -824,6 +824,56 @@ void BindIR(nb::module_& m) { }, nb::arg("op_name"), "Get memory space specification for a registered operator"); + nb::enum_(ir, "ArgEffect", "What executing an operator does to the buffer one argument names") + .value("Read", ArgEffect::Read, "Read, never written") + .value("Write", ArgEffect::Write, "Overwritten without being read first") + .value("ReadWrite", ArgEffect::ReadWrite, "Read and written (accumulate, atomic, in-place update)"); + + nb::enum_(ir, "WriteChannel", "The hardware path an operator's writes travel") + .value("Dma", WriteChannel::Dma, "MTE3 / DMA store path") + .value("Scalar", WriteChannel::Scalar, "Scalar D-cache write path"); + + ir.def( + "get_op_arg_effect", + [](const std::string& op_name, size_t arg_index, nb::kwargs kwargs) -> ArgEffect { + const auto& entry = OpRegistry::GetInstance().GetEntry(op_name); + // The same conversion every other kwarg-taking binding uses. Rolling a + // local int/str pair here rejected the enum-valued kwargs an operator + // legitimately carries — `tile.mgather`'s `target_memory` is a + // `MemorySpace`, and it is exactly what that operator's effect resolver + // reads. + nb::dict kwargs_dict; + for (auto [key, value] : kwargs) { + kwargs_dict[key] = value; + } + return entry.GetArgEffect(arg_index, ConvertKwargsDict(kwargs_dict)); + }, + nb::arg("op_name"), nb::arg("arg_index"), nb::arg("kwargs"), + "Effect an operator has on one positional argument, for a call carrying the given kwargs"); + + ir.def( + "op_has_declared_arg_effects", + [](const std::string& op_name) { + return OpRegistry::GetInstance().GetEntry(op_name).HasDeclaredArgEffects(); + }, + nb::arg("op_name"), "Whether an operator declared its per-argument effects (False = never classified)"); + + ir.def( + "op_has_declared_arg_effect", + [](const std::string& op_name, size_t arg_index) { + return OpRegistry::GetInstance().GetEntry(op_name).HasDeclaredArgEffect(arg_index); + }, + nb::arg("op_name"), nb::arg("arg_index"), + "Whether the registration reached a verdict about this argument in particular"); + + ir.def( + "get_op_write_channel", + [](const std::string& op_name) -> nb::object { + auto channel = OpRegistry::GetInstance().GetEntry(op_name).GetWriteChannel(); + return channel.has_value() ? nb::cast(*channel) : nb::none(); + }, + nb::arg("op_name"), "The hardware path an operator's writes travel, or None when it declared none"); + // Var - const shared_ptr auto var_class = nb::class_(ir, "Var", "Variable reference expression"); diff --git a/python/pypto/pypto_core/ir.pyi b/python/pypto/pypto_core/ir.pyi index 9cf79de699..d1ceeca380 100644 --- a/python/pypto/pypto_core/ir.pyi +++ b/python/pypto/pypto_core/ir.pyi @@ -3147,6 +3147,91 @@ def get_op_memory_spec(op_name: str) -> dict[str, Any] | None: * ``None`` — no resolver registered for this op. """ +class ArgEffect(enum.Enum): + """What executing an operator does to the buffer one argument names.""" + + Read = ... + """Read, never written. The default for an argument the operator does not name.""" + + Write = ... + """Overwritten without being read first (a destination operand).""" + + ReadWrite = ... + """Read and written — accumulate, atomic update, partial in-place rewrite.""" + +class WriteChannel(enum.Enum): + """The hardware path an operator's writes travel.""" + + Dma = ... + """MTE3 / DMA store path (tile.store, tensor.assemble, cross-rank put/get).""" + + Scalar = ... + """Scalar D-cache write path (tensor.write).""" + +def get_op_arg_effect(op_name: str, arg_index: int, **kwargs: Any) -> ArgEffect: + """Effect an operator has on one positional argument. + + Args: + op_name: Name of the operator + arg_index: Positional argument index + **kwargs: The kwargs a call would carry, for operators whose effect + depends on one (an atomic ``tile.store`` reads the accumulator it + adds into; ``pld.system.notify`` accumulates unless ``op`` selects + the set form) + + Returns: + The declared effect, or ``ArgEffect.Read`` for an argument the operator + did not name + + Raises: + Exception: If operator is not registered + """ + +def op_has_declared_arg_effects(op_name: str) -> bool: + """Whether an operator declared its per-argument effects. + + Args: + op_name: Name of the operator + + Returns: + False when the operator was never classified — distinct from a + declared read-only operator, so an analysis can refuse to guess + + Raises: + Exception: If operator is not registered + """ + +def op_has_declared_arg_effect(op_name: str, arg_index: int) -> bool: + """Whether the registration reached a verdict about one argument. + + Args: + op_name: Name of the operator + arg_index: Positional argument index + + Returns: + True when the operator named this argument, or declared with + ``no_arg_writes()`` that it writes through none of them. False when it + classified only *other* arguments — the resulting ``Read`` for this one + is a default, not a decision. + + Raises: + Exception: If operator is not registered + """ + +def get_op_write_channel(op_name: str) -> WriteChannel | None: + """The hardware path an operator's writes travel. + + Args: + op_name: Name of the operator + + Returns: + The declared channel, or None when the operator declared none (it + writes nothing, or its writes are not GM stores) + + Raises: + Exception: If operator is not registered + """ + # ========== Op Conversion Registry ========== def register_op_conversion(from_op: str, to_op: str) -> None: diff --git a/src/ir/op/array_ops/memory.cpp b/src/ir/op/array_ops/memory.cpp index d55e3312a0..4284ed8e3a 100644 --- a/src/ir/op/array_ops/memory.cpp +++ b/src/ir/op/array_ops/memory.cpp @@ -154,6 +154,9 @@ REGISTER_OP("array.update_element") .add_argument("array", "Source array (ArrayType)") .add_argument("index", "Element index (ScalarType, integer)") .add_argument("value", "Replacement value (ScalarType, dtype must match array)") + // Replaces one element and passes every other element of `array` through + // to the result, so the prior content is read. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceArrayUpdateElementType(args, kwargs); diff --git a/src/ir/op/distributed/allreduce.cpp b/src/ir/op/distributed/allreduce.cpp index fde18f6b5c..528d91f633 100644 --- a/src/ir/op/distributed/allreduce.cpp +++ b/src/ir/op/distributed/allreduce.cpp @@ -140,6 +140,10 @@ REGISTER_OP("pld.tensor.allreduce") .set_attr("mode") .set_attr("core_num") .no_memory_spec() + // Composite collective — target is reduced in place per chunk; signal is written by notify and read by + // wait. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceTensorAllReduceType); } // namespace ir diff --git a/src/ir/op/distributed/collective.cpp b/src/ir/op/distributed/collective.cpp index f467a3dd9c..8f4d055406 100644 --- a/src/ir/op/distributed/collective.cpp +++ b/src/ir/op/distributed/collective.cpp @@ -208,6 +208,11 @@ REGISTER_OP("builtin.tensor.allreduce") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.allreduce") + // Host-level collective: same read/write shape as the pld.tensor.* form + // it lowers from — the data window is updated in place and the signal is + // written by the notify phase and read by the wait phase. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorAllReduceType); REGISTER_OP("builtin.tensor.allreduce_ring") @@ -220,6 +225,11 @@ REGISTER_OP("builtin.tensor.allreduce_ring") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.allreduce_ring") + // Host-level collective: same read/write shape as the pld.tensor.* form + // it lowers from — the data window is updated in place and the signal is + // written by the notify phase and read by the wait phase. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorAllReduceRingType); // ============================================================================ @@ -258,6 +268,8 @@ REGISTER_OP("pld.tensor.barrier") .set_op_category("DistributedOp") .add_argument("signal", "Window-bound INT32 DistributedTensor used as cross-rank barrier (InOut)") .no_memory_spec() + // Composite collective — signal is written by the notify phase and read by the wait phase. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type(DeduceTensorBarrierType); // ============================================================================ @@ -311,6 +323,9 @@ REGISTER_OP("pld.tensor.broadcast") .add_argument("signal", "Window-bound INT32 DistributedTensor used as cross-rank barrier (InOut)") .set_attr("root") .no_memory_spec() + // Composite collective — target is read on the root and written on every rank; signal is notify+wait. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceTensorBroadcastType); // ============================================================================ @@ -430,6 +445,15 @@ REGISTER_OP("pld.tensor.allgather") .add_argument("target", "Window-bound DistributedTensor[NR, SIZE] — gathered result in-place (InOut)") .add_argument("signal", "Window-bound INT32 DistributedTensor used as cross-rank barrier (InOut)") .no_memory_spec() + // notify+wait. + // Composite collective — the data destination is overwritten, not updated: + // the lowering only pushes into it (`pld.tile.put`) and never loads from it, + // so nothing moves into the kernel through it. Declaring `ReadWrite` here + // would make the enclosing parameter `InOut`, stage the buffer host->device + // and invent a dependency on its incoming content. The signal is genuinely + // both: written by the notify phase and read by the wait phase. + .set_arg_effect(1, ArgEffect::Write) + .set_arg_effect(2, ArgEffect::ReadWrite) .f_deduce_type(DeduceTensorAllGatherType); // ============================================================================ @@ -526,6 +550,15 @@ REGISTER_OP("pld.tensor.all_to_all") "Window-bound DistributedTensor [NR, SIZE] — receives the result in-place (InOut)") .add_argument("signal", "Window-bound INT32 DistributedTensor used as cross-rank barrier (InOut)") .no_memory_spec() + // is notify+wait. + // Composite collective — the data destination is overwritten, not updated: + // the lowering only pushes into it (`pld.tile.put`) and never loads from it, + // so nothing moves into the kernel through it. Declaring `ReadWrite` here + // would make the enclosing parameter `InOut`, stage the buffer host->device + // and invent a dependency on its incoming content. The signal is genuinely + // both: written by the notify phase and read by the wait phase. + .set_arg_effect(1, ArgEffect::Write) + .set_arg_effect(2, ArgEffect::ReadWrite) .f_deduce_type(DeduceTensorAllToAllType); // ============================================================================ @@ -696,6 +729,16 @@ REGISTER_OP("pld.tensor.all_to_all_v") "Window-bound INT32 DistributedTensor [NR, 1] — after the barrier, " "recv_counts[src, 0] holds how many rows src sent to this rank (InOut)") .no_memory_spec() + // stays read-only. + // Composite collective — the data destination is overwritten, not updated: + // the lowering only pushes into it (`pld.tile.put`) and never loads from it, + // so nothing moves into the kernel through it. Declaring `ReadWrite` here + // would make the enclosing parameter `InOut`, stage the buffer host->device + // and invent a dependency on its incoming content. The signal is genuinely + // both: written by the notify phase and read by the wait phase. + .set_arg_effect(1, ArgEffect::Write) + .set_arg_effect(2, ArgEffect::ReadWrite) + .set_arg_effect(4, ArgEffect::Write) .f_deduce_type(DeduceTensorAllToAllVType); // ============================================================================ @@ -754,6 +797,9 @@ REGISTER_OP("pld.tensor.reduce_scatter") .add_argument("signal", "Window-bound INT32 DistributedTensor used as cross-rank barrier (InOut)") .set_attr("op") .no_memory_spec() + // Composite collective — same five-phase shape as allreduce. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceTensorReduceScatterType); // ============================================================================ @@ -782,6 +828,10 @@ REGISTER_OP("builtin.tensor.barrier") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.barrier") + // Host-level collective: same read/write shape as the pld.tensor.* form + // it lowers from — the data window is updated in place and the signal is + // written by the notify phase and read by the wait phase. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorBarrierType); // ============================================================================ @@ -824,6 +874,11 @@ REGISTER_OP("builtin.tensor.broadcast") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.broadcast") + // Host-level collective: same read/write shape as the pld.tensor.* form + // it lowers from — the data window is updated in place and the signal is + // written by the notify phase and read by the wait phase. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorBroadcastType); // ============================================================================ @@ -867,6 +922,11 @@ REGISTER_OP("builtin.tensor.reduce_scatter") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.reduce_scatter") + // Host-level collective: same read/write shape as the pld.tensor.* form + // it lowers from — the data window is updated in place and the signal is + // written by the notify phase and read by the wait phase. + .set_arg_effect(0, ArgEffect::ReadWrite) + .set_arg_effect(1, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorReduceScatterType); // ============================================================================ @@ -945,6 +1005,14 @@ REGISTER_OP("builtin.tensor.allgather") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.allgather") + // Composite collective — the data destination is overwritten, not updated: + // the lowering only pushes into it (`pld.tile.put`) and never loads from it, + // so nothing moves into the kernel through it. Declaring `ReadWrite` here + // would make the enclosing parameter `InOut`, stage the buffer host->device + // and invent a dependency on its incoming content. The signal is genuinely + // both: written by the notify phase and read by the wait phase. + .set_arg_effect(1, ArgEffect::Write) + .set_arg_effect(2, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorAllGatherType); // ============================================================================ @@ -1030,6 +1098,14 @@ REGISTER_OP("builtin.tensor.all_to_all") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.all_to_all") + // Composite collective — the data destination is overwritten, not updated: + // the lowering only pushes into it (`pld.tile.put`) and never loads from it, + // so nothing moves into the kernel through it. Declaring `ReadWrite` here + // would make the enclosing parameter `InOut`, stage the buffer host->device + // and invent a dependency on its incoming content. The signal is genuinely + // both: written by the notify phase and read by the wait phase. + .set_arg_effect(1, ArgEffect::Write) + .set_arg_effect(2, ArgEffect::ReadWrite) .f_deduce_type(DeduceBuiltinTensorAllToAllType); // ============================================================================ @@ -1184,6 +1260,15 @@ REGISTER_OP("builtin.tensor.all_to_all_v") .no_memory_spec() .set_internal_only(true) .set_template_dir(":pypto.runtime.builtins.collectives.all_to_all_v") + // Composite collective — the data destination is overwritten, not updated: + // the lowering only pushes into it (`pld.tile.put`) and never loads from it, + // so nothing moves into the kernel through it. Declaring `ReadWrite` here + // would make the enclosing parameter `InOut`, stage the buffer host->device + // and invent a dependency on its incoming content. The signal is genuinely + // both: written by the notify phase and read by the wait phase. + .set_arg_effect(1, ArgEffect::Write) + .set_arg_effect(2, ArgEffect::ReadWrite) + .set_arg_effect(4, ArgEffect::Write) .f_deduce_type(DeduceBuiltinTensorAllToAllVType); } // namespace ir diff --git a/src/ir/op/distributed/get.cpp b/src/ir/op/distributed/get.cpp index dd61a6141b..a0da82c737 100644 --- a/src/ir/op/distributed/get.cpp +++ b/src/ir/op/distributed/get.cpp @@ -223,6 +223,9 @@ REGISTER_OP("pld.tensor.get") .set_attr("pipeline") .set_core_affinity(core_affinity::CoreAffinity::VECTOR) .no_memory_spec() + // TGET lands the pulled bytes in the local destination `dst`. + .set_arg_effect(0, ArgEffect::Write) + .set_write_channel(WriteChannel::Dma) .f_deduce_type(DeduceGetType); // ============================================================================ @@ -250,6 +253,9 @@ REGISTER_OP("pld.tile.get") .add_argument("shape", "Optional per-dim transfer shape (MakeTuple); present only in the subregion form") .set_core_affinity(core_affinity::CoreAffinity::VECTOR) .no_memory_spec() + // TGET lands the pulled bytes in the local destination `dst`. + .set_arg_effect(0, ArgEffect::Write) + .set_write_channel(WriteChannel::Dma) .f_deduce_type(DeduceGetTileType); } // namespace ir diff --git a/src/ir/op/distributed/put.cpp b/src/ir/op/distributed/put.cpp index f76c069881..7ff59cb37f 100644 --- a/src/ir/op/distributed/put.cpp +++ b/src/ir/op/distributed/put.cpp @@ -78,6 +78,7 @@ #include "pypto/core/dtype.h" #include "pypto/core/logging.h" +#include "pypto/ir/comm.h" #include "pypto/ir/core_affinity_kind.h" #include "pypto/ir/expr.h" #include "pypto/ir/kind_traits.h" @@ -251,6 +252,16 @@ REGISTER_OP("pld.tensor.put") .set_attr("pipeline") .set_core_affinity(core_affinity::CoreAffinity::VECTOR) .no_memory_spec() + // A plain push overwrites the region it lands on; an atomic one accumulates + // into it, and accumulating reads the slot first. + .set_arg_effect(0, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type(DeducePutType); // ============================================================================ @@ -280,6 +291,16 @@ REGISTER_OP("pld.tile.put") .set_attr("atomic") .set_core_affinity(core_affinity::CoreAffinity::VECTOR) .no_memory_spec() + // A plain push overwrites the region it lands on; an atomic one accumulates + // into it, and accumulating reads the slot first. + .set_arg_effect(0, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type(DeducePutTileType); } // namespace ir diff --git a/src/ir/op/distributed/remote_load.cpp b/src/ir/op/distributed/remote_load.cpp index 46c4220766..91d6cd5bca 100644 --- a/src/ir/op/distributed/remote_load.cpp +++ b/src/ir/op/distributed/remote_load.cpp @@ -252,6 +252,8 @@ REGISTER_OP("pld.tile.remote_load") .add_argument("valid_shape", "Optional valid tile extent for ragged tails (MakeTuple of scalars)") .set_attr("allow_physical_tail_padding") .no_memory_spec() + // Pulls a peer's window into a fresh SSA tile: every operand is a read. + .no_arg_writes() .f_deduce_type(DeduceRemoteLoadType); } // namespace ir diff --git a/src/ir/op/distributed/remote_store.cpp b/src/ir/op/distributed/remote_store.cpp index 64451ed6d0..a047b211fd 100644 --- a/src/ir/op/distributed/remote_store.cpp +++ b/src/ir/op/distributed/remote_store.cpp @@ -76,6 +76,7 @@ #include "pypto/core/dtype.h" #include "pypto/core/logging.h" +#include "pypto/ir/comm.h" #include "pypto/ir/expr.h" #include "pypto/ir/kind_traits.h" #include "pypto/ir/memory_space.h" @@ -249,6 +250,16 @@ REGISTER_OP("pld.tile.remote_store") // InferTileMemorySpace pull a producer into a legal space instead of letting // an illegal one reach codegen. .set_input_memory(0, {MemorySpace::Vec, MemorySpace::Acc}) + // A plain push overwrites the region it lands on; an atomic one accumulates + // into it, and accumulating reads the slot first. + .set_arg_effect(1, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type(DeduceRemoteStoreType); // ============================================================================ @@ -270,6 +281,16 @@ REGISTER_OP("pld.tensor.remote_store") .add_argument("offsets", "Offsets in target tensor coordinates (MakeTuple of scalars)") .set_attr("atomic") .no_memory_spec() + // A plain push overwrites the region it lands on; an atomic one accumulates + // into it, and accumulating reads the slot first. + .set_arg_effect(1, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type(DeduceTensorRemoteStoreType); } // namespace ir diff --git a/src/ir/op/distributed/system.cpp b/src/ir/op/distributed/system.cpp index d981cf538d..8d8c713f91 100644 --- a/src/ir/op/distributed/system.cpp +++ b/src/ir/op/distributed/system.cpp @@ -221,6 +221,17 @@ REGISTER_OP("pld.system.notify") .set_attr("op") .set_no_duplicate() .no_memory_spec() + // TNOTIFY deposits `value` into the peer rank's slot of `target`. The set + // form overwrites the slot; the atomic-add form (the default) accumulates + // into it, which reads it. Without this the notifying task looks like a + // pure reader of the signal and the waiter carries no dependency on it. + .set_arg_effect(0, + [](const std::vector>& kwargs) { + return static_cast(GetIntKwarg( + kwargs, "op", static_cast(NotifyOp::kAtomicAdd))) == NotifyOp::kSet + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) .f_deduce_type(DeduceNotifyType); // ============================================================================ @@ -237,6 +248,9 @@ REGISTER_OP("pld.system.wait") .add_argument("expected", "Scalar threshold value") .set_attr("cmp") .no_memory_spec() + // Polls the local slot of `signal` until it satisfies the threshold; it + // never writes the signal, only the matching notify does. + .no_arg_writes() .f_deduce_type(DeduceWaitType); // ============================================================================ @@ -264,6 +278,8 @@ REGISTER_OP("pld.system.defer_wait") .add_argument("expected", "Integer or index scalar threshold value") .set_attr("cmp") .no_memory_spec() + // Registers a completion condition on `signal`; a read, like wait. + .no_arg_writes() .f_deduce_type(DeduceDeferWaitType); } // namespace ir diff --git a/src/ir/op/sync_ops/cross_core.cpp b/src/ir/op/sync_ops/cross_core.cpp index 053b74d270..5e48cf66ca 100644 --- a/src/ir/op/sync_ops/cross_core.cpp +++ b/src/ir/op/sync_ops/cross_core.cpp @@ -118,6 +118,11 @@ REGISTER_OP("system.set_ffts") .set_description("Declare the A3 FFTS setup operand for explicit cross-core synchronization") .set_op_category("CrossCoreOp") .add_argument("workspace", "One-dimensional INT64 FFTS workspace") + // Hands the workspace *pointer* to the FFTS unit (`pto.set_ffts %ws : + // !pto.ptr`) — it declares where the hardware's scratch lives rather + // than moving any data itself. The FFTS unit writes that region on its own + // schedule, which no PyPTO dependency edge models or could usefully order. + .no_arg_writes() .f_deduce_type(DeduceSetFFTSType); REGISTER_OP("system.sync_set") diff --git a/src/ir/op/sync_ops/sync.cpp b/src/ir/op/sync_ops/sync.cpp index c3befe5aee..5124ddbc82 100644 --- a/src/ir/op/sync_ops/sync.cpp +++ b/src/ir/op/sync_ops/sync.cpp @@ -133,6 +133,14 @@ REGISTER_OP("system.syncall") .add_argument("used_cores", "Soft form: optional participant core count (i32; omitted = auto)") .set_attr("core_type") .set_attr("mode") + // Soft form: every core writes its arrival counter into `gm_workspace` and + // polls it, so the workspace is read and written. The hard form is an FFTS + // barrier that touches no workspace at all. + .set_arg_effect(0, + [](const std::vector>& kwargs) { + return GetStringKwarg(kwargs, "mode", "hard") == "soft" ? ArgEffect::ReadWrite + : ArgEffect::Read; + }) .f_deduce_type(DeduceUnknownType); } // namespace ir diff --git a/src/ir/op/tensor_ops/memory.cpp b/src/ir/op/tensor_ops/memory.cpp index 5441cfc4fd..de9c68a850 100644 --- a/src/ir/op/tensor_ops/memory.cpp +++ b/src/ir/op/tensor_ops/memory.cpp @@ -571,6 +571,16 @@ REGISTER_OP("tensor.assemble") // buffer, not a new allocation. Declaring it here keeps param/buffer lineage // analyses off a hardcoded op list. .set_output_reuses_input(0) + // A plain push overwrites the region it lands on; an atomic one accumulates + // into it, and accumulating reads the slot first. + .set_arg_effect(0, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTensorAssembleType(args, kwargs); @@ -869,6 +879,11 @@ REGISTER_OP("tensor.write") .add_argument("tensor", "Destination tensor (TensorType)") .add_argument("indices", "Index dimensions (TupleType of ScalarType)") .add_argument("value", "Value to write (ScalarType)") + // Writes one element of `tensor` through the scalar D-cache path. The + // channel matters: PyPTO cannot order a scalar write against an MTE3 + // store to the same GM tensor, and rejects a function that mixes them. + .set_arg_effect(0, ArgEffect::Write) + .set_write_channel(WriteChannel::Scalar) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTensorWriteType(args, kwargs); diff --git a/src/ir/op/tensor_ops/transform.cpp b/src/ir/op/tensor_ops/transform.cpp index d4c1b39e54..7d86a70fd2 100644 --- a/src/ir/op/tensor_ops/transform.cpp +++ b/src/ir/op/tensor_ops/transform.cpp @@ -817,6 +817,11 @@ REGISTER_OP("tensor.set_validshape") // a lineage walk that cannot see the aliasing reports the result as a fresh // kernel allocation. .set_output_reuses_input(0) + // The in-place slot is metadata, not data: this op rebinds the valid extent + // and moves nothing, so no dependency edge should order against it. A + // verdict on record — the gate requires one, and "writes nothing" is the + // honest answer rather than an omission. + .no_arg_writes() .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTensorSetValidShapeType(args, kwargs); diff --git a/src/ir/op/tile_ops/batch_matmul.cpp b/src/ir/op/tile_ops/batch_matmul.cpp index 2cca86cb1d..1fb2a4cbe5 100644 --- a/src/ir/op/tile_ops/batch_matmul.cpp +++ b/src/ir/op/tile_ops/batch_matmul.cpp @@ -303,6 +303,8 @@ REGISTER_OP("tile.batch_matmul_acc") .set_input_memory(2, MemorySpace::Right) .set_output_memory(MemorySpace::Acc) .set_output_reuses_input(0) + // Accumulates into `acc`, same as tile.matmul_acc. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileBatchMatMulAccType(args, kwargs, "tile.batch_matmul_acc"); diff --git a/src/ir/op/tile_ops/elementwise.cpp b/src/ir/op/tile_ops/elementwise.cpp index d2c6b7c6eb..772d141221 100644 --- a/src/ir/op/tile_ops/elementwise.cpp +++ b/src/ir/op/tile_ops/elementwise.cpp @@ -1315,6 +1315,8 @@ REGISTER_OP("tile.fillpad_inplace") .set_input_memory(0, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) .set_output_reuses_input(0) + // Rewrites only the padding elements; the data region passes through. + .set_arg_effect(0, ArgEffect::ReadWrite) .set_attr("pad_value") .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { diff --git a/src/ir/op/tile_ops/matmul.cpp b/src/ir/op/tile_ops/matmul.cpp index 7bf5339c0e..4b804c5909 100644 --- a/src/ir/op/tile_ops/matmul.cpp +++ b/src/ir/op/tile_ops/matmul.cpp @@ -430,6 +430,8 @@ REGISTER_OP("tile.matmul_acc") .set_input_memory(2, MemorySpace::Right) .set_output_memory(MemorySpace::Acc) .set_output_reuses_input(0) + // Accumulates into `acc`: C += A@B reads the running sum it adds to. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileMatMulAccType(args, kwargs, "tile.matmul_acc"); @@ -479,6 +481,8 @@ REGISTER_OP("tile.gemv_acc") .set_input_memory(2, MemorySpace::Right) .set_output_memory(MemorySpace::Acc) .set_output_reuses_input(0) + // Accumulates into `acc`, same as tile.matmul_acc. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileGemvAccType(args, kwargs, "tile.gemv_acc"); diff --git a/src/ir/op/tile_ops/matmul_mx.cpp b/src/ir/op/tile_ops/matmul_mx.cpp index 252711951c..9d89457c82 100644 --- a/src/ir/op/tile_ops/matmul_mx.cpp +++ b/src/ir/op/tile_ops/matmul_mx.cpp @@ -349,6 +349,8 @@ REGISTER_OP("tile.matmul_mx_acc") .set_input_memory(4, MemorySpace::RightScale) .set_output_memory(MemorySpace::Acc) .set_output_reuses_input(0) + // Accumulates into `acc`, same as tile.matmul_acc. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileMatMulMxAccType(args, kwargs, "tile.matmul_mx_acc"); @@ -460,6 +462,11 @@ REGISTER_OP("tile.tget_scale_addr") .add_argument("src", "Resolved Left/Right MX data tile (FP8E4M3FN) whose address is scaled") .set_output_memory_inherit_input() .set_output_reuses_input(0) + // Binds a derived address into the shared physical scale buffer `dst_scale` + // names — a mutation of that buffer, which is why InsertMxScaleAddr never + // reuses one binding across two MX matmul consumers. Declaring the write + // keeps that hazard visible to any analysis that orders accesses. + .set_arg_effect(0, ArgEffect::Write) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileTGetScaleAddrType(args, kwargs, "tile.tget_scale_addr"); diff --git a/src/ir/op/tile_ops/memory.cpp b/src/ir/op/tile_ops/memory.cpp index dc29e318c0..1aad11e01e 100644 --- a/src/ir/op/tile_ops/memory.cpp +++ b/src/ir/op/tile_ops/memory.cpp @@ -976,6 +976,11 @@ REGISTER_OP("tile.write") .add_argument("tile", "Destination tile (TileType)") .add_argument("indices", "Index dimensions (TupleType of ScalarType)") .add_argument("value", "Scalar value to write (ScalarType)") + // Rewrites one element and passes every other element of the tile through + // to the result, so the prior content is read. No write channel: this is a + // tile-local write, not one of the GM store paths the mixed-store + // diagnostic orders against each other. + .set_arg_effect(0, ArgEffect::ReadWrite) .set_input_memory(0, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) .f_deduce_type([](const std::vector& args, @@ -1087,6 +1092,18 @@ REGISTER_OP("tile.store") .set_attr("atomic") .set_input_memory(0, {MemorySpace::Vec, MemorySpace::Acc}) .set_output_reuses_input(2) + // A plain store overwrites the region it lands on: the untouched remainder + // is neither loaded nor re-stored, so nothing moves *into* the kernel and + // the destination is a pure write. An atomic store is not an overwrite at + // all — `out += x` reads the accumulator it adds to. + .set_arg_effect(2, + [](const std::vector>& kwargs) { + return GetIntKwarg(kwargs, "atomic", static_cast(AtomicType::kNone)) == + static_cast(AtomicType::kNone) + ? ArgEffect::Write + : ArgEffect::ReadWrite; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileStoreType(args, kwargs, "tile.store"); @@ -1163,6 +1180,10 @@ REGISTER_OP("tile.mscatter") .set_input_memory(0, MemorySpace::Vec) .set_input_memory(1, MemorySpace::Vec) .set_output_reuses_input(2) + // Scatters `src` into the indexed cells of `output_tensor` without reading + // any of it — the same pure-write destination contract as tile.store. + .set_arg_effect(2, ArgEffect::Write) + .set_write_channel(WriteChannel::Dma) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileMscatterType(args, kwargs, "tile.mscatter"); @@ -1381,6 +1402,21 @@ REGISTER_OP("tile.mgather") .set_attr("target_memory") .set_output_memory_from_kwarg("target_memory", MemorySpace::Vec) .not_inplace_safe() + // Argument 2 is the GM `scratch` tensor only in Mat *elem* mode, where the + // gathered elements are staged through it. In Mat row mode that position + // holds `valid_shape`, and in Vec mode it is absent — declaring an + // unconditional write there would claim a tuple operand is a written + // buffer and could promote a read-only parameter to an output. + .set_arg_effect(2, + [](const std::vector>& kwargs) { + const bool mat_output = + GetMemorySpaceKwarg(kwargs, "target_memory", MemorySpace::Vec) == MemorySpace::Mat; + const bool elem_mode = + GetIntKwarg(kwargs, "coalesce", static_cast(MgatherCoalesceMode::kRow)) == + static_cast(MgatherCoalesceMode::kElem); + return mat_output && elem_mode ? ArgEffect::Write : ArgEffect::Read; + }) + .set_write_channel(WriteChannel::Dma) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileMgatherType(args, kwargs, "tile.mgather"); diff --git a/src/ir/op/tile_ops/paged_gather.cpp b/src/ir/op/tile_ops/paged_gather.cpp index 2ae647f2e6..ee6acee007 100644 --- a/src/ir/op/tile_ops/paged_gather.cpp +++ b/src/ir/op/tile_ops/paged_gather.cpp @@ -93,6 +93,9 @@ REGISTER_OP("tile.gather_row") "May be a runtime Scalar[INDEX]; defaults to shapes.") .set_attr("transpose") .set_output_reuses_input(0) + // DPS: loads one GM row into a sub-region of `dst`; the rest of the + // accumulator tile passes through, so the prior content is read. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileGatherRowType(args, kwargs, "tile.gather_row"); diff --git a/src/ir/op/tile_ops/scatter.cpp b/src/ir/op/tile_ops/scatter.cpp index e3f9521611..d12310632f 100644 --- a/src/ir/op/tile_ops/scatter.cpp +++ b/src/ir/op/tile_ops/scatter.cpp @@ -196,6 +196,9 @@ REGISTER_OP("tile.scatter") .set_input_memory(2, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) .set_output_reuses_input(0) + // DPS: rewrites the indexed positions of `dst` and passes every other + // position through, so the prior content reaches the result — a read. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileScatterType(args, kwargs, "tile.scatter"); @@ -290,6 +293,8 @@ REGISTER_OP("tile.scatter_mask") .set_input_memory(1, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) .set_output_reuses_input(0) + // DPS: mask-selected columns are rewritten, the rest pass through. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileScatterMaskType(args, kwargs, "tile.scatter_mask"); diff --git a/src/ir/op/tile_ops/transform.cpp b/src/ir/op/tile_ops/transform.cpp index c559aa477b..684281c5e8 100644 --- a/src/ir/op/tile_ops/transform.cpp +++ b/src/ir/op/tile_ops/transform.cpp @@ -662,6 +662,9 @@ REGISTER_OP("tile.assemble") .add_argument("target", "Target tile (TileType)") .add_argument("source", "Source tile to write (TileType)") .add_argument("offset", "Offset dimensions (TupleType of ScalarType(INT64/UINT64/INDEX))") + // Rewrites the offset sub-region of `target` and passes the rest through to + // the result, so the prior content is read. Tile-local, hence no channel. + .set_arg_effect(0, ArgEffect::ReadWrite) .set_output_memory_inherit_input() .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { @@ -874,6 +877,8 @@ REGISTER_OP("tile.scatter_update") .set_input_memory(2, MemorySpace::Vec) .set_output_memory(MemorySpace::Vec) .set_output_reuses_input(0) + // DPS: the indexed rows are rewritten, every other row passes through. + .set_arg_effect(0, ArgEffect::ReadWrite) .f_deduce_type([](const std::vector& args, const std::vector>& kwargs) { return DeduceTileScatterUpdateType(args, kwargs); diff --git a/src/ir/op_registry.cpp b/src/ir/op_registry.cpp index b63a21bc23..bee2aa4f90 100644 --- a/src/ir/op_registry.cpp +++ b/src/ir/op_registry.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -279,5 +280,53 @@ void OpRegistry::ValidateTileOps() const { } } +void OpRegistry::ValidateArgEffects() const { + std::vector unclassified; + std::vector channel_without_write; + for (const auto& [name, entry] : registry_) { + // A write channel describes *how* an operator writes, so declaring one + // while writing nothing is incoherent — and it is the shape that hides a + // missing classification, since `set_write_channel()` creates the effect + // spec as a side effect and would otherwise make the operator look + // classified. + if (entry.GetWriteChannel().has_value() && !entry.WritesAnyArg()) { + channel_without_write.push_back(name); + } + const auto& spec = entry.GetMemorySpec(); + if (!spec.has_value() || !spec->output_reuses_input_arg.has_value()) continue; + const size_t reused = *spec->output_reuses_input_arg; + // Ask about the reused argument specifically. A registration that named a + // different argument still leaves this one defaulting to `Read`, and the + // whole point of the gate is that such a default is a decision nobody made. + if (entry.HasDeclaredArgEffect(reused)) continue; + unclassified.push_back(name + " (in-place on argument " + std::to_string(reused) + ")"); + } + if (!channel_without_write.empty()) { + std::sort(channel_without_write.begin(), channel_without_write.end()); + std::string msg = + "The following ops declare a write channel but write through no argument. A channel says " + "how an op writes, so one without a write is either a stray declaration or a missing " + "one — add the .set_arg_effect(, ...) that was meant to accompany it, or drop the " + ".set_write_channel(...):"; + for (const auto& name : channel_without_write) { + msg += "\n - " + name; + } + throw ValueError(msg); + } + if (!unclassified.empty()) { + std::sort(unclassified.begin(), unclassified.end()); + std::string msg = + "The following ops update an argument in place but never declared what they do to it. " + "Direction inference reads an undeclared operator as a pure consumer, so the write is " + "silently dropped. Add .set_arg_effect(, ArgEffect::Write) — ArgEffect::ReadWrite " + "when the op accumulates into the slot — or .no_arg_writes() when the slot is metadata " + "rather than data:"; + for (const auto& name : unclassified) { + msg += "\n - " + name; + } + throw ValueError(msg); + } +} + } // namespace ir } // namespace pypto diff --git a/tests/ut/ir/operators/test_op_registry.py b/tests/ut/ir/operators/test_op_registry.py index 33182b794c..5a3be4a11c 100644 --- a/tests/ut/ir/operators/test_op_registry.py +++ b/tests/ut/ir/operators/test_op_registry.py @@ -987,5 +987,303 @@ def test_value_error_still_surfaces_as_value_error(self): ir.create_op_call("tile.cast", [self._arg()], ir.Span.unknown()) +class TestArgEffects: + """Per-argument read/write effects declared on the operator registry. + + Every direction and dependency analysis needs one answer to "does this call + write the buffer this argument names". These tests pin that answer at its + source, so a new operator cannot quietly join the set of writers nobody + models — which is how a written parameter keeps direction ``In``, loses its + RAW edge, and deadlocks or races on device. + """ + + def test_unnamed_argument_defaults_to_read(self): + """The tile a store copies *from* is read, not written.""" + assert ir.get_op_arg_effect("tile.store", 0) == ir.ArgEffect.Read + + def test_index_past_the_argument_list_is_read(self): + assert ir.get_op_arg_effect("tile.store", 99) == ir.ArgEffect.Read + + def test_functional_op_is_unclassified(self): + """`tensor.add` writes through no argument and was never classified; + `False` here is what lets an analysis tell that apart from a declared + read-only operator.""" + assert ir.op_has_declared_arg_effects("tensor.add") is False + assert ir.get_op_arg_effect("tensor.add", 0) == ir.ArgEffect.Read + + def test_declared_read_only_op_is_classified(self): + """`pld.system.wait` polls a signal it never writes — classified, but + with no write. That is a decision on record, not an omission.""" + assert ir.op_has_declared_arg_effects("pld.system.wait") is True + assert ir.get_op_arg_effect("pld.system.wait", 0) == ir.ArgEffect.Read + + def test_unknown_op_raises(self): + with pytest.raises(ValueError): + ir.get_op_arg_effect("tile.not_an_op", 0) + + @pytest.mark.parametrize( + ("op_name", "arg_index", "expected"), + [ + # A store overwrites the region it lands on; the untouched remainder + # is neither loaded nor re-stored, so nothing moves into the kernel. + ("tile.store", 2, ir.ArgEffect.Write), + # Same contract, and the one that was missing: a GM tensor written + # only by a scatter used to read as a pure input. + ("tile.mscatter", 2, ir.ArgEffect.Write), + ("tensor.assemble", 0, ir.ArgEffect.Write), + ("tensor.write", 0, ir.ArgEffect.Write), + # Cross-rank pushes and pulls land in their destination operand. + ("pld.tile.remote_store", 1, ir.ArgEffect.Write), + ("pld.tensor.remote_store", 1, ir.ArgEffect.Write), + ("pld.tile.put", 0, ir.ArgEffect.Write), + ("pld.tile.get", 0, ir.ArgEffect.Write), + # Accumulators read the running sum they add to. + ("tile.matmul_acc", 0, ir.ArgEffect.ReadWrite), + ("tile.gemv_acc", 0, ir.ArgEffect.ReadWrite), + # Destination-passing style: the positions the op does not rewrite + # pass through to the result, so the prior content is read. + ("tile.scatter", 0, ir.ArgEffect.ReadWrite), + ("tile.scatter_update", 0, ir.ArgEffect.ReadWrite), + ("array.update_element", 0, ir.ArgEffect.ReadWrite), + ("tile.write", 0, ir.ArgEffect.ReadWrite), + ("tile.assemble", 0, ir.ArgEffect.ReadWrite), + # Composite collectives update their window and signal in place. + ("pld.tensor.allreduce", 0, ir.ArgEffect.ReadWrite), + ("pld.tensor.allreduce", 1, ir.ArgEffect.ReadWrite), + # A gather/exchange destination is overwritten, not updated: the + # lowering only pushes into it and never loads from it. `recv_counts` + # is deposited with NotifyOp::Set, so it is not an accumulate either. + ("pld.tensor.allgather", 0, ir.ArgEffect.Read), + ("pld.tensor.allgather", 1, ir.ArgEffect.Write), + ("pld.tensor.allgather", 2, ir.ArgEffect.ReadWrite), + ("pld.tensor.all_to_all", 1, ir.ArgEffect.Write), + ("pld.tensor.all_to_all_v", 1, ir.ArgEffect.Write), + ("pld.tensor.all_to_all_v", 3, ir.ArgEffect.Read), + ("pld.tensor.all_to_all_v", 4, ir.ArgEffect.Write), + # A reduce destination *is* read — its lowering loads the running + # value back — so the distinction is per operator, not per family. + ("pld.tensor.allreduce", 0, ir.ArgEffect.ReadWrite), + ("pld.tensor.reduce_scatter", 0, ir.ArgEffect.ReadWrite), + ], + ) + def test_declared_effects(self, op_name, arg_index, expected): + assert ir.get_op_arg_effect(op_name, arg_index) == expected + + def test_atomic_store_reads_its_destination(self): + """`out += x` is not an overwrite: the accumulate reads the slot first. + Declaring it `Write` would let the runtime skip staging the buffer, and + the sum would start from allocator garbage.""" + plain = ir.get_op_arg_effect("tile.store", 2) + atomic = ir.get_op_arg_effect("tile.store", 2, atomic=int(ir.AtomicType.Add)) + assert plain == ir.ArgEffect.Write + assert atomic == ir.ArgEffect.ReadWrite + + def test_atomic_assemble_reads_its_destination(self): + assert ir.get_op_arg_effect("tensor.assemble", 0) == ir.ArgEffect.Write + assert ( + ir.get_op_arg_effect("tensor.assemble", 0, atomic=int(ir.AtomicType.Add)) + == ir.ArgEffect.ReadWrite + ) + + def test_notify_defaults_to_accumulating(self): + """`pld.system.notify`'s `op` kwarg defaults to atomic-add, so an + unannotated notify reads the slot it adds into; only the set form is a + pure overwrite.""" + assert ir.get_op_arg_effect("pld.system.notify", 0) == ir.ArgEffect.ReadWrite + assert ir.get_op_arg_effect("pld.system.notify", 0, op=int(ir.NotifyOp.Set)) == ir.ArgEffect.Write + + def test_mgather_scratch_only_in_mat_elem_mode(self): + """`tile.mgather`'s argument 2 is a written GM scratch tensor only when + the gather stages through one. + + `DeduceTileMgatherType` puts `scratch` at that position for Mat *elem* + mode; Mat row mode holds `valid_shape` there and Vec mode has no third + operand at all. Declaring the write unconditionally would claim a tuple + operand is a written buffer, and could promote a read-only parameter to + an output. + """ + mat = ir.MemorySpace.Mat + # `MgatherCoalesceMode` (include/pypto/ir/comm.h) is not bound to Python; + # the DSL passes the same ints, and the op deducer validates the range. + elem, row = 1, 0 + assert ir.get_op_arg_effect("tile.mgather", 2, target_memory=mat, coalesce=elem) == ( + ir.ArgEffect.Write + ) + assert ir.get_op_arg_effect("tile.mgather", 2, target_memory=mat, coalesce=row) == (ir.ArgEffect.Read) + # Vec is the default output space and carries no third operand. + assert ir.get_op_arg_effect("tile.mgather", 2) == ir.ArgEffect.Read + + def test_enum_valued_kwargs_reach_the_resolver(self): + """A resolver may key on any kwarg the operator declares, including an + enum-valued one. The query converts kwargs the same way every other + binding does, so a `MemorySpace` argument resolves instead of raising.""" + assert ( + ir.get_op_arg_effect("tile.mgather", 2, target_memory=ir.MemorySpace.Mat, coalesce=1) + == ir.ArgEffect.Write + ) + + def test_set_ffts_declares_no_write(self): + """`system.set_ffts` hands the workspace *pointer* to the FFTS unit + (`pto.set_ffts %ws : !pto.ptr`); it declares where the hardware's + scratch lives rather than moving any data. The FFTS unit writes that + region on its own schedule, which no PyPTO dependency edge models.""" + assert ir.op_has_declared_arg_effects("system.set_ffts") is True + assert ir.get_op_arg_effect("system.set_ffts", 0) == ir.ArgEffect.Read + assert ir.get_op_write_channel("system.set_ffts") is None + + def test_in_place_gate_asks_about_the_reused_argument(self): + """The import-time gate must ask about the argument the operator updates + in place, not merely whether *some* argument was classified. + + `per_arg` cannot answer that on its own — it is resized to cover the + highest declared index, so a slot nobody named looks like a declared + `Read`. Without the distinction, an operator declaring + `set_output_reuses_input(2)` while classifying argument 1 would pass the + gate with argument 2 still defaulting to `Read`. + """ + # tile.store declares set_output_reuses_input(2) and classifies 2. + assert ir.op_has_declared_arg_effect("tile.store", 2) is True + # Argument 0 is covered by `per_arg` (it was resized past it) but was + # never named, so no verdict was reached about it. + assert ir.op_has_declared_arg_effect("tile.store", 0) is False + # `no_arg_writes()` is a verdict about every argument at once. + assert ir.op_has_declared_arg_effect("pld.system.wait", 0) is True + assert ir.op_has_declared_arg_effect("pld.system.wait", 7) is True + # An operator nobody classified reaches no verdict about any argument. + assert ir.op_has_declared_arg_effect("tensor.add", 0) is False + + def test_a_write_channel_alone_is_not_a_verdict(self): + """Declaring only a write channel must not make an operator look classified. + + `set_write_channel()` creates the effect spec as a side effect, so + "the spec exists" cannot stand in for "a human decided". Were it allowed + to, an operator that declared a channel and forgot its `set_arg_effect` + would pass the in-place gate with the argument it updates still + defaulting to `Read` — the exact silent default this registry exists to + remove. `no_arg_writes()` records the verdict explicitly instead. + + Every operator that declares a channel therefore also writes something, + which `ValidateArgEffects()` enforces at import; this pins the invariant + that check maintains. + """ + for op_name, written_index in _CHANNEL_OPS.items(): + assert ir.get_op_write_channel(op_name) is not None, op_name + assert ir.op_has_declared_arg_effect(op_name, written_index), ( + f"{op_name} declares a write channel but reached no verdict about " + f"argument {written_index}, the one that channel describes" + ) + + def test_composite_collectives_declare_no_write_channel(self): + """A composite collective updates a data window and a signal through + different mechanisms, and one operator-level channel cannot describe + both. Declaring `Dma` for the pair would let the mixed-store diagnostic + pair a collective's signal write against a scalar `tensor.write` on the + same buffer and reject a program that is fine. Recording no channel + keeps them out of that diagnostic, exactly as before this API existed. + """ + for op_name in ( + ir.get_op("pld.tensor.allreduce").name, + ir.get_op("pld.tensor.barrier").name, + ir.get_op("pld.tensor.allgather").name, + ir.get_op("builtin.tensor.broadcast").name, + ): + assert ir.get_op_write_channel(op_name) is None, op_name + + def test_notify_declares_no_write_channel(self): + """`pld.system.notify` emits `pto.comm.tnotify`, which is neither the + MTE3 store path nor the scalar D-cache path the mixed-store diagnostic + orders against each other. Claiming either would make that diagnostic + reject a valid program, so it declares the write without a channel.""" + assert ir.get_op_arg_effect("pld.system.notify", 0) == ir.ArgEffect.ReadWrite + assert ir.get_op_write_channel("pld.system.notify") is None + + def test_hard_syncall_does_not_touch_the_workspace(self): + """The soft form counts arrivals in the GM workspace; the hard form is + an FFTS barrier that never reads or writes it.""" + assert ir.get_op_arg_effect("system.syncall", 0) == ir.ArgEffect.Read + assert ir.get_op_arg_effect("system.syncall", 0, mode="soft") == ir.ArgEffect.ReadWrite + + @pytest.mark.parametrize( + ("op_name", "expected"), + [ + ("tile.store", ir.WriteChannel.Dma), + ("tensor.assemble", ir.WriteChannel.Dma), + ("tile.mscatter", ir.WriteChannel.Dma), + # The one scalar D-cache writer. PyPTO cannot order a scalar write + # against an MTE3 store to the same GM tensor, and rejects a + # function that mixes them. + ("tensor.write", ir.WriteChannel.Scalar), + # Declared classified, writes nothing, so no channel. + ("pld.system.wait", None), + ], + ) + def test_write_channel(self, op_name, expected): + assert ir.get_op_write_channel(op_name) == expected + + def test_every_in_place_op_is_classified(self): + """An operator whose result reuses an input's buffer writes through that + argument. Leaving the effect undeclared is what let `tile.mscatter` + write a GM output while every direction analysis read it as an input. + + `pypto` fails at import when this is violated (see + `OpRegistry::ValidateArgEffects`); asserting it here names the operator + and the fix instead of failing the whole test session on import. + """ + for op_name in _IN_PLACE_OPS: + assert ir.op_has_declared_arg_effects(op_name), ( + f"{op_name} updates an argument in place but never declared what it does to it. " + f"Add .set_arg_effect(, ArgEffect::Write) to its REGISTER_OP block — " + f"ArgEffect::ReadWrite when it accumulates, or .no_arg_writes() when the slot " + f"is metadata rather than data." + ) + + +#: Operators declaring a write channel, mapped to the argument that channel +#: describes. Each must also declare a write there — a channel says *how* an +#: operator writes, so one without a write is either a stray declaration or a +#: missing one. `tile.mgather` reaches its verdict through a kwarg resolver, +#: which still counts: the registration named the argument. +_CHANNEL_OPS = { + ir.get_op(name).name: index + for name, index in ( + ("tile.store", 2), + ("tile.mscatter", 2), + ("tile.mgather", 2), + ("tensor.write", 0), + ("tensor.assemble", 0), + ("pld.tile.put", 0), + ("pld.tile.get", 0), + ("pld.tile.remote_store", 1), + ) +} + +#: Operators declaring ``set_output_reuses_input``: their SSA result IS an +#: argument's buffer, so they write through it and must classify that argument. +#: Routed through ``get_op`` so a renamed operator fails at import rather than +#: silently dropping out of the coverage this list asserts. +_IN_PLACE_OPS = [ + ir.get_op(name).name + for name in ( + "array.update_element", + # main declared these in-place after this series began; the import gate + # is what surfaced them, so pin them here too. + "tensor.assemble", + "tensor.set_validshape", + "tile.batch_matmul_acc", + "tile.fillpad_inplace", + "tile.gather_row", + "tile.gemv_acc", + "tile.matmul_acc", + "tile.matmul_mx_acc", + "tile.mscatter", + "tile.scatter", + "tile.scatter_mask", + "tile.scatter_update", + "tile.store", + "tile.tget_scale_addr", + ) +] + + if __name__ == "__main__": pytest.main([__file__, "-v"])