PTO-COMM-ISA-pto-isa-AtomGit

Adocs/isa/comm/README.md+91-0

@@ -0,0 +1,91 @@
	1	+# PTO Communication ISA Reference
	2	+
	3	+This directory contains the per-instruction reference for the PTO Communication ISA.
	4	+
	5	+- Source of truth (C++ intrinsics): `include/pto/comm/pto_comm_inst.hpp`
	6	+- Type definitions: `include/pto/comm/comm_types.hpp`
	7	+
	8	+## Point-to-Point Communication (Synchronous)
	9	+- `TPUT`: `docs/isa/comm/TPUT.md` - Remote write (GM → UB → GM)
	10	+- `TGET`: `docs/isa/comm/TGET.md` - Remote read (GM → UB → GM)
	11	+
	12	+## Signal-Based Synchronization
	13	+- `TNOTIFY`: `docs/isa/comm/TNOTIFY.md` - Send notification to remote NPU
	14	+- `TWAIT`: `docs/isa/comm/TWAIT.md` - Blocking wait for signal condition
	15	+- `TTEST`: `docs/isa/comm/TTEST.md` - Non-blocking test signal condition
	16	+
	17	+## Collective Communication
	18	+
	19	+- `TGATHER`: `docs/isa/comm/TGATHER.md` - Gather data from all ranks
	20	+- `TSCATTER`: `docs/isa/comm/TSCATTER.md` - Scatter data to all ranks
	21	+- `TREDUCE`: `docs/isa/comm/TREDUCE.md` - Reduce data from all ranks to local
	22	+- `TBROADCAST`: `docs/isa/comm/TBROADCAST.md` - Broadcast from current NPU to all ranks
	23	+
	24	+## Type Definitions
	25	+
	26	+### NotifyOp
	27	+
	28	+Operation type for `TNOTIFY`:
	29	+
	30	+\| Value \| Description \|
	31	+\|-------\|-------------\|
	32	+\| `NotifyOp::Set` \| Direct set (`signal = value`) \|
	33	+\| `NotifyOp::AtomicAdd` \| Atomic add (`signal += value`) \|
	34	+
	35	+### WaitCmp
	36	+
	37	+Comparison operators for `TWAIT` and `TTEST`:
	38	+
	39	+\| Value \| Description \|
	40	+\|-------\|-------------\|
	41	+\| `WaitCmp::EQ` \| Equal (`==`) \|
	42	+\| `WaitCmp::NE` \| Not equal (`!=`) \|
	43	+\| `WaitCmp::GT` \| Greater than (`>`) \|
	44	+\| `WaitCmp::GE` \| Greater or equal (`>=`) \|
	45	+\| `WaitCmp::LT` \| Less than (`<`) \|
	46	+\| `WaitCmp::LE` \| Less or equal (`<=`) \|
	47	+
	48	+```cpp
	49	+// Usage (unified runtime parameter style):
	50	+comm::TNOTIFY(signal, 1, comm::NotifyOp::Set);
	51	+comm::TWAIT(signal, 1, comm::WaitCmp::EQ);
	52	+comm::TTEST(signal, 1, comm::WaitCmp::GE);
	53	+```
	54	+
	55	+### ReduceOp
	56	+
	57	+Reduction operators for `TREDUCE`:
	58	+
	59	+\| Value \| Description \|
	60	+\|-------\|-------------\|
	61	+\| `ReduceOp::Sum` \| Element-wise sum \|
	62	+\| `ReduceOp::Max` \| Element-wise maximum \|
	63	+\| `ReduceOp::Min` \| Element-wise minimum \|
	64	+
	65	+### AtomicType
	66	+
	67	+Atomic operation type for `TPUT` (defined in `include/pto/common/constants.hpp`):
	68	+
	69	+\| Value \| Description \|
	70	+\|-------\|-------------\|
	71	+\| `AtomicType::AtomicNone` \| No atomic operation (default) \|
	72	+\| `AtomicType::AtomicAdd` \| Atomic add operation \|
	73	+
	74	+### ParallelGroup
	75	+
	76	+Wrapper for collective communication across multiple NPUs:
	77	+
	78	+```cpp
	79	+template <typename GlobalData>
	80	+struct ParallelGroup {
	81	+ // Pointer to an array of `GlobalData` objects (each wraps a GM address).
	82	+ // The array itself is local metadata; the wrapped addresses may refer to local or remote GM,
	83	+ // depending on the collective instruction.
	84	+ GlobalData *tensors;
	85	+ int nranks; // Number of ranks
	86	+ int rootIdx; // Root NPU's rank index
	87	+
	88	+ // Factory function (recommended): build from an existing tensor array.
	89	+ static ParallelGroup Create(GlobalData *tensorArray, int size, int rank_id);
	90	+};
	91	+```

Adocs/isa/comm/TBROADCAST.md+122-0

@@ -0,0 +1,122 @@
	1	+# TBROADCAST
	2	+
	3	+## Introduction
	4	+
	5	+Broadcast data from current NPU to all ranks in the parallel group. The calling NPU is the root and its data is copied to all other NPUs.
	6	+
	7	+Only the root needs to execute `TBROADCAST`. Non-root ranks only need to ensure their destination buffers are allocated and writable for the duration of the operation. Calling `TBROADCAST` on non-root ranks is undefined behavior.
	8	+
	9	+Large Tile Support: When the GlobalTensor exceeds the UB (Unified Buffer) tile capacity in rows and/or columns, the transfer is automatically chunked via 2D sliding.
	10	+
	11	+## Math Interpretation
	12	+
	13	+After the operation:
	14	+
	15	+$$ \mathrm{dst}^{(k)}_{i,j} = \mathrm{src}^{(\text{root})}_{i,j} \quad \forall k \in [0, N) $$
	16	+
	17	+where $N$ is the number of ranks and `root` is the calling NPU.
	18	+
	19	+## Assembly Syntax
	20	+
	21	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	22	+
	23	+Synchronous form:
	24	+
	25	+```text
	26	+tbroadcast %group, %src : (!pto.group<...>, !pto.memref<...>)
	27	+```
	28	+Lowering introduces UB staging tile(s) for the GM→UB→GM data path; the C++ intrinsic requires explicit `stagingTileData` (or `pingTile` / `pongTile`) operand(s).
	29	+
	30	+## C++ Intrinsic
	31	+
	32	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	33	+
	34	+```cpp
	35	+// Basic broadcast (single staging tile)
	36	+template <typename ParallelGroupType, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	37	+PTO_INST RecordEvent TBROADCAST(ParallelGroupType &parallelGroup, GlobalSrcData &srcGlobalData,
	38	+ TileData &stagingTileData, WaitEvents&... events);
	39	+
	40	+// Ping-pong broadcast (double buffering with two staging tiles)
	41	+template <typename ParallelGroupType, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	42	+PTO_INST RecordEvent TBROADCAST(ParallelGroupType &parallelGroup, GlobalSrcData &srcGlobalData,
	43	+ TileData &pingTile, TileData &pongTile, WaitEvents&... events);
	44	+```
	45	+
	46	+## Constraints
	47	+
	48	+- Type constraints:
	49	+ - `ParallelGroup::value_type::RawDType` must equal `GlobalSrcData::RawDType`.
	50	+ - `TileData::DType` must equal `GlobalSrcData::RawDType`.
	51	+- Memory constraints:
	52	+ - `srcGlobalData` must point to local memory (current NPU).
	53	+ - `stagingTileData` (or `pingTile` / `pongTile`) must be pre-allocated in UB.
	54	+- ParallelGroup constraints:
	55	+ - `parallelGroup.tensors[k]` must refer to rank `k`'s destination buffer (remote GM as seen by the root).
	56	+ - `parallelGroup.GetRootIdx()` identifies the calling NPU as the broadcast root.
	57	+ - All destination tensors are assumed to have the same shape and strides.
	58	+- Chunked mode constraints (when data exceeds a single UB tile):
	59	+ - If `TileData` has static `ValidRow`, `GetShape(DIM_3)` must be divisible by `ValidRow`. Use a Tile with `DYNAMIC` ValidRow for partial row support.
	60	+ - If `TileData` has static `ValidCol`, `GetShape(DIM_4)` must be divisible by `ValidCol`. Use a Tile with `DYNAMIC` ValidCol for partial column support.
	61	+
	62	+## Examples
	63	+
	64	+### Basic Broadcast
	65	+
	66	+```cpp
	67	+#include <pto/comm/pto_comm_inst.hpp>
	68	+
	69	+using namespace pto;
	70	+
	71	+template <typename T, int ROWS, int COLS, int TILE_ROWS, int TILE_COLS, int NRANKS>
	72	+void broadcast(__gm__ T* group_addrs[NRANKS], __gm__ T* my_data, int my_rank) {
	73	+ // Tile dimensions can differ from tensor dimensions.
	74	+ // The 2D sliding chunked path automatically tiles both row and column.
	75	+ using TileT = Tile<TileType::Vec, T, TILE_ROWS, TILE_COLS, BLayout::RowMajor, -1, -1>;
	76	+ using GTensor = GlobalTensor<T, Shape<1,1,1,ROWS,COLS>,
	77	+ BaseShape2D<T, ROWS, COLS, Layout::ND>, Layout::ND>;
	78	+
	79	+ GTensor tensors[NRANKS];
	80	+ for (int i = 0; i < NRANKS; ++i) {
	81	+ tensors[i] = GTensor(group_addrs[i]);
	82	+ }
	83	+
	84	+ comm::ParallelGroup<GTensor> group(tensors, NRANKS, my_rank);
	85	+ GTensor srcG(my_data);
	86	+ TileT stagingTile(TILE_ROWS, TILE_COLS);
	87	+
	88	+ // Current NPU broadcasts its data to all others
	89	+ comm::TBROADCAST(group, srcG, stagingTile);
	90	+}
	91	+```
	92	+
	93	+### Ping-Pong Broadcast (Double Buffering)
	94	+
	95	+Uses two UB tiles to overlap TLOAD of the next chunk with TSTORE of the current chunk.
	96	+
	97	+```cpp
	98	+#include <pto/comm/pto_comm_inst.hpp>
	99	+
	100	+using namespace pto;
	101	+
	102	+template <typename T, int ROWS, int COLS, int TILE_ROWS, int TILE_COLS, int NRANKS>
	103	+void broadcast_pingpong(__gm__ T* group_addrs[NRANKS], __gm__ T* my_data, int my_rank) {
	104	+
	105	+ using TileT = Tile<TileType::Vec, T, TILE_ROWS, TILE_COLS, BLayout::RowMajor, -1, -1>;
	106	+ using GPerRank = GlobalTensor<T, Shape<1,1,1,ROWS,COLS>,
	107	+ BaseShape2D<T, ROWS, COLS, Layout::ND>, Layout::ND>;
	108	+
	109	+ GPerRank tensors[NRANKS];
	110	+ for (int i = 0; i < NRANKS; ++i) {
	111	+ tensors[i] = GPerRank(group_addrs[i]);
	112	+ }
	113	+
	114	+ comm::ParallelGroup<GPerRank> group(tensors, NRANKS, my_rank);
	115	+ GPerRank srcG(my_data);
	116	+ TileT pingTile(TILE_ROWS, TILE_COLS);
	117	+ TileT pongTile(TILE_ROWS, TILE_COLS);
	118	+
	119	+ // Ping-pong: overlaps TLOAD and TSTORE for better throughput
	120	+ comm::TBROADCAST(group, srcG, pingTile, pongTile);
	121	+}
	122	+```

Adocs/isa/comm/TGATHER.md+128-0

@@ -0,0 +1,128 @@
	1	+# TGATHER
	2	+
	3	+## Introduction
	4	+
	5	+Gather operation: the calling NPU (root) collects data from all ranks in the parallel group and concatenates the results along DIM_3 (row dimension) into a local output buffer.
	6	+
	7	+
	8	+Only the root needs to execute `TGATHER`. Non-root ranks only need to ensure their source buffers are ready and remain valid for the duration of the operation. Calling `TGATHER` on non-root ranks is undefined behavior.
	9	+
	10	+Large Tile Support: When the GlobalTensor exceeds the UB tile capacity in rows and/or columns, the transfer is automatically chunked via 2D sliding — the same mechanism used by other PTO-COMM instructions.
	11	+
	12	+## Math Interpretation
	13	+
	14	+Each rank $r$ has source data of shape $(D_0, D_1, D_2, H, W)$. The gather concatenates all $N$ ranks along DIM_3:
	15	+
	16	+$$\mathrm{dst}_{d_0, d_1, d_2,\; r \cdot H + i,\; j} = \mathrm{src}^{(r)}_{d_0, d_1, d_2,\; i,\; j} \quad \forall\, r \in [0, N),\; i \in [0, H),\; j \in [0, W)$$
	17	+
	18	+The destination tensor has shape $(D_0, D_1, D_2, N \times H, W)$.
	19	+
	20	+## Assembly Syntax
	21	+
	22	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	23	+
	24	+Synchronous form:
	25	+
	26	+```text
	27	+tgather %group, %dst : (!pto.group<...>, !pto.memref<...>)
	28	+```
	29	+Lowering introduces UB staging tile(s) for the GM→UB→GM data path; the C++ intrinsic requires explicit `stagingTileData` (or `pingTile` / `pongTile`) operand(s).
	30	+
	31	+## C++ Intrinsic
	32	+
	33	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	34	+
	35	+```cpp
	36	+// Basic gather (single staging tile)
	37	+template <typename ParallelGroupType, typename GlobalDstData, typename TileData, typename... WaitEvents>
	38	+PTO_INST RecordEvent TGATHER(ParallelGroupType &parallelGroup, GlobalDstData &dstGlobalData,
	39	+ TileData &stagingTileData, WaitEvents&... events);
	40	+
	41	+// Ping-pong gather (double buffering with two staging tiles)
	42	+template <typename ParallelGroupType, typename GlobalDstData, typename TileData, typename... WaitEvents>
	43	+PTO_INST RecordEvent TGATHER(ParallelGroupType &parallelGroup, GlobalDstData &dstGlobalData,
	44	+ TileData &pingTile, TileData &pongTile, WaitEvents&... events);
	45	+```
	46	+
	47	+## Constraints
	48	+
	49	+- Type constraints:
	50	+ - `ParallelGroup::value_type::RawDType` must equal `GlobalDstData::RawDType`.
	51	+ - `TileData::DType` must equal `GlobalDstData::RawDType`.
	52	+- Memory constraints:
	53	+ - `dstGlobalData` must point to local memory (current NPU) and be large enough to hold the concatenated result from all ranks. Specifically, `dstGlobalData.GetShape(DIM_3)` must be $\geq N \times H$ where $H$ is each rank's `GetShape(DIM_3)`.
	54	+ - If `dstGlobalData.GetShape(DIM_3) > N × H`, only the first `N × H` rows are written; remaining rows are left unchanged.
	55	+ - `stagingTileData` (or `pingTile` / `pongTile`) must be pre-allocated in UB.
	56	+- ParallelGroup constraints:
	57	+ - `parallelGroup.tensors[r]` must refer to rank `r`'s source buffer (remote GM as seen by the root).
	58	+ - `parallelGroup.GetRootIdx()` identifies the calling NPU as the gather root.
	59	+ - All source tensors are assumed to have the same shape and strides; behavior is undefined if they differ.
	60	+- Chunked mode constraints (when source data exceeds a single UB tile):
	61	+ - If `TileData` has static `ValidRow`, `GetShape(DIM_3)` of each rank's source must be divisible by `ValidRow`. Use a Tile with `DYNAMIC` ValidRow for partial row support.
	62	+ - If `TileData` has static `ValidCol`, `GetShape(DIM_4)` must be divisible by `ValidCol`. Use a Tile with `DYNAMIC` ValidCol for partial column support.
	63	+
	64	+## Examples
	65	+
	66	+### Basic Gather (Single Staging Tile)
	67	+
	68	+Each rank contributes `ROWS × COLS` data. The root collects them into `NRANKS * ROWS` rows.
	69	+The tile size (`TILE_ROWS × TILE_COLS`) can be smaller than the per-rank data — when it is, the implementation automatically chunks the transfer along both DIM_3 and DIM_4 via 2D sliding.
	70	+
	71	+```cpp
	72	+#include <pto/comm/pto_comm_inst.hpp>
	73	+
	74	+using namespace pto;
	75	+
	76	+template <typename T, int ROWS, int COLS, int TILE_ROWS, int TILE_COLS, int NRANKS>
	77	+void gather(__gm__ T* group_addrs[NRANKS], __gm__ T* result, int my_rank) {
	78	+ using TileT = Tile<TileType::Vec, T, TILE_ROWS, TILE_COLS, BLayout::RowMajor, -1, -1>;
	79	+ using GPerRank = GlobalTensor<T, Shape<1,1,1,ROWS,COLS>,
	80	+ BaseShape2D<T, ROWS, COLS, Layout::ND>, Layout::ND>;
	81	+ using GResult = GlobalTensor<T, Shape<1,1,1,NRANKS*ROWS,COLS>,
	82	+ BaseShape2D<T, NRANKS*ROWS, COLS, Layout::ND>, Layout::ND>;
	83	+
	84	+ GPerRank tensors[NRANKS];
	85	+ for (int i = 0; i < NRANKS; ++i) {
	86	+ tensors[i] = GPerRank(group_addrs[i]);
	87	+ }
	88	+
	89	+ comm::ParallelGroup<GPerRank> group(tensors, NRANKS, my_rank);
	90	+ GResult dstG(result);
	91	+ TileT stagingTile(TILE_ROWS, TILE_COLS);
	92	+
	93	+ comm::TGATHER(group, dstG, stagingTile);
	94	+}
	95	+```
	96	+
	97	+### Ping-Pong Gather (Double Buffering)
	98	+
	99	+Uses two UB tiles to overlap TLOAD of the next chunk (MTE2) with TSTORE of the current chunk (MTE3).
	100	+
	101	+```cpp
	102	+#include <pto/comm/pto_comm_inst.hpp>
	103	+
	104	+using namespace pto;
	105	+
	106	+template <typename T, int ROWS, int COLS, int TILE_ROWS, int TILE_COLS, int NRANKS>
	107	+void gather_pingpong(__gm__ T* group_addrs[NRANKS], __gm__ T* result, int my_rank) {
	108	+ // Tile can be smaller than the data in both dimensions
	109	+ using TileT = Tile<TileType::Vec, T, TILE_ROWS, TILE_COLS, BLayout::RowMajor, -1, -1>;
	110	+ using GPerRank = GlobalTensor<T, Shape<1,1,1,ROWS,COLS>,
	111	+ BaseShape2D<T, ROWS, COLS, Layout::ND>, Layout::ND>;
	112	+ using GResult = GlobalTensor<T, Shape<1,1,1,NRANKS*ROWS,COLS>,
	113	+ BaseShape2D<T, NRANKS*ROWS, COLS, Layout::ND>, Layout::ND>;
	114	+
	115	+ GPerRank tensors[NRANKS];
	116	+ for (int i = 0; i < NRANKS; ++i) {
	117	+ tensors[i] = GPerRank(group_addrs[i]);
	118	+ }
	119	+
	120	+ comm::ParallelGroup<GPerRank> group(tensors, NRANKS, my_rank);
	121	+ GResult dstG(result);
	122	+ TileT pingTile(TILE_ROWS, TILE_COLS);
	123	+ TileT pongTile(TILE_ROWS, TILE_COLS);
	124	+
	125	+ // Ping-pong: overlaps TLOAD and TSTORE for better throughput
	126	+ comm::TGATHER(group, dstG, pingTile, pongTile);
	127	+}
	128	+```

Adocs/isa/comm/TGET.md+109-0

@@ -0,0 +1,109 @@
	1	+# TGET
	2	+
	3	+## Introduction
	4	+
	5	+Remote read operation: read remote NPU's data to local memory. Data is transferred via a UB tile as intermediate staging buffer.
	6	+
	7	+When the GlobalTensor exceeds the UB tile capacity, TGET automatically performs 2D sliding — chunking rows (DIM_3) and columns (DIM_4) to fit each chunk into the tile, iterating over all outer dimensions (DIM_0, DIM_1, DIM_2).
	8	+
	9	+## Math Interpretation
	10	+
	11	+For each element `(i, j)` in the valid region:
	12	+
	13	+$$ \mathrm{dst}^{\mathrm{local}}_{i,j} = \mathrm{src}^{\mathrm{remote}}_{i,j} $$
	14	+
	15	+Data flow: `srcGlobalData (remote GM)` → `stagingTileData (UB)` → `dstGlobalData (local GM)`
	16	+
	17	+## Assembly Syntax
	18	+
	19	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	20	+
	21	+Synchronous form:
	22	+
	23	+```text
	24	+tget %dst_local, %src_remote : (!pto.memref<...>, !pto.memref<...>)
	25	+```
	26	+Lowering introduces UB staging tile(s) for the GM→UB→GM data path; the C++ intrinsic requires explicit `stagingTileData` (or `pingTile` / `pongTile`) operand(s).
	27	+
	28	+## C++ Intrinsic
	29	+
	30	+Declared in `include/pto/comm/pto_comm_inst.hpp`
	31	+
	32	+### Single-tile (auto-chunking)
	33	+
	34	+```cpp
	35	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	36	+PTO_INST RecordEvent TGET(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
	37	+ TileData &stagingTileData, WaitEvents&... events);
	38	+```
	39	+
	40	+### Ping-pong double buffering
	41	+
	42	+Uses two staging tiles to overlap TLOAD and TSTORE for adjacent chunks, hiding one DMA transfer behind the other.
	43	+
	44	+```cpp
	45	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	46	+PTO_INST RecordEvent TGET(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
	47	+ TileData &pingTile, TileData &pongTile, WaitEvents&... events);
	48	+```
	49	+
	50	+## Constraints
	51	+
	52	+- Type constraints:
	53	+ - `GlobalSrcData::RawDType` must equal `GlobalDstData::RawDType`.
	54	+ - `TileData::DType` must equal `GlobalSrcData::RawDType`.
	55	+ - `GlobalSrcData::layout` must equal `GlobalDstData::layout`.
	56	+- Memory constraints:
	57	+ - `srcGlobalData` must point to remote address (on source NPU).
	58	+ - `dstGlobalData` must point to local address (on current NPU).
	59	+ - `stagingTileData` / `pingTile` / `pongTile` must be pre-allocated in Unified Buffer.
	60	+- Valid region:
	61	+ - Transfer size is determined by `GlobalTensor` shape (auto-chunked to fit tile).
	62	+- Ping-pong:
	63	+ - `pingTile` and `pongTile` must have the same type and dimensions.
	64	+ - Must reside at non-overlapping UB offsets.
	65	+
	66	+## Examples
	67	+
	68	+### Basic Usage
	69	+
	70	+```cpp
	71	+#include <pto/comm/pto_comm_inst.hpp>
	72	+#include <pto/pto-inst.hpp>
	73	+
	74	+using namespace pto;
	75	+
	76	+template <typename T>
	77	+void example_tget(__gm__ T* local_data, __gm__ T* remote_addr) {
	78	+ using TileT = Tile<TileType::Vec, T, 16, 16>;
	79	+ using GShape = Shape<1, 1, 1, 16, 16>;
	80	+ using GStride = BaseShape2D<T, 16, 16, Layout::ND>;
	81	+ /*
	82	+ If the globalTensor is larger than UB Tile, TGET will perform 2D sliding automatically.
	83	+ using GShape = Shape<1, 1, 1, 4096, 4096>;
	84	+ using GStride = BaseShape2D<T, 4096, 4096, Layout::ND>;
	85	+ */
	86	+ using GTensor = GlobalTensor<T, GShape, GStride, Layout::ND>;
	87	+
	88	+ GTensor srcG(remote_addr);
	89	+ GTensor dstG(local_data);
	90	+ TileT stagingTile;
	91	+ TASSIGN(stagingTile, 0);
	92	+
	93	+ // Basic remote read
	94	+ comm::TGET(dstG, srcG, stagingTile);
	95	+}
	96	+```
	97	+
	98	+### Ping-pong Double Buffering
	99	+
	100	+```cpp
	101	+constexpr size_t tileUBBytes = ((64 * 64 * sizeof(float) + 1023) / 1024) * 1024;
	102	+TileT pingTile(64, 64);
	103	+TileT pongTile(64, 64);
	104	+TASSIGN(pingTile, 0);
	105	+TASSIGN(pongTile, tileUBBytes); // Non-overlapping UB region
	106	+
	107	+// Overlaps TLOAD[i+1] with TSTORE[i] for better pipeline utilization
	108	+comm::TGET(dstG, srcG, pingTile, pongTile);
	109	+```

Adocs/isa/comm/TNOTIFY.md+100-0

@@ -0,0 +1,100 @@
	1	+# TNOTIFY
	2	+
	3	+## Introduction
	4	+
	5	+Send flag notification to remote NPU. Used for lightweight synchronization between NPUs without transferring bulk data.
	6	+
	7	+## Math Interpretation
	8	+
	9	+For `NotifyOp::Set`:
	10	+
	11	+$$ \mathrm{signal}^{\mathrm{remote}} = \mathrm{value} $$
	12	+
	13	+For `NotifyOp::AtomicAdd`:
	14	+
	15	+$$ \mathrm{signal}^{\mathrm{remote}} \mathrel{+}= \mathrm{value} \quad (\text{atomic}) $$
	16	+
	17	+## Assembly Syntax
	18	+
	19	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	20	+
	21	+```text
	22	+tnotify %signal_remote, %value {op = #pto.notify_op<Set>} : (!pto.memref<i32>, i32)
	23	+tnotify %signal_remote, %value {op = #pto.notify_op<AtomicAdd>} : (!pto.memref<i32>, i32)
	24	+```
	25	+
	26	+## C++ Intrinsic
	27	+
	28	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	29	+
	30	+```cpp
	31	+template <typename GlobalSignalData, typename... WaitEvents>
	32	+PTO_INST void TNOTIFY(GlobalSignalData &dstSignalData, int32_t value, NotifyOp op, WaitEvents&... events);
	33	+```
	34	+
	35	+## Constraints
	36	+
	37	+- Type constraints:
	38	+ - `GlobalSignalData::DType` must be `int32_t` (32-bit signal).
	39	+- Memory constraints:
	40	+ - `dstSignalData` must point to remote address (on target NPU).
	41	+ - `dstSignalData` should be 4-byte aligned.
	42	+- Operation semantics:
	43	+ - `NotifyOp::Set`: Direct store to remote memory.
	44	+ - `NotifyOp::AtomicAdd`: Hardware atomic add using `st_atomic` instruction.
	45	+
	46	+## Examples
	47	+
	48	+### Basic Set Notification
	49	+
	50	+```cpp
	51	+#include <pto/comm/pto_comm_inst.hpp>
	52	+
	53	+using namespace pto;
	54	+
	55	+void notify_set(__gm__ int32_t* remote_signal) {
	56	+ comm::Signal sig(remote_signal);
	57	+
	58	+ // Set remote signal to 1
	59	+ comm::TNOTIFY(sig, 1, comm::NotifyOp::Set);
	60	+}
	61	+```
	62	+
	63	+### Atomic Counter Increment
	64	+
	65	+```cpp
	66	+#include <pto/comm/pto_comm_inst.hpp>
	67	+
	68	+using namespace pto;
	69	+
	70	+void atomic_increment(__gm__ int32_t* remote_counter) {
	71	+ comm::Signal counter(remote_counter);
	72	+
	73	+ // Atomically add 1 to remote counter
	74	+ comm::TNOTIFY(counter, 1, comm::NotifyOp::AtomicAdd);
	75	+}
	76	+```
	77	+
	78	+### Producer-Consumer Pattern
	79	+
	80	+```cpp
	81	+#include <pto/comm/pto_comm_inst.hpp>
	82	+
	83	+using namespace pto;
	84	+
	85	+// Producer: notify when data is ready
	86	+void producer(__gm__ int32_t* remote_flag) {
	87	+ // ... produce data ...
	88	+
	89	+ comm::Signal flag(remote_flag);
	90	+ comm::TNOTIFY(flag, 1, comm::NotifyOp::Set);
	91	+}
	92	+
	93	+// Consumer: wait for data
	94	+void consumer(__gm__ int32_t* local_flag) {
	95	+ comm::Signal flag(local_flag);
	96	+ comm::TWAIT(flag, 1, comm::WaitCmp::EQ);
	97	+
	98	+ // ... consume data ...
	99	+}
	100	+```

Adocs/isa/comm/TPUT.md+131-0

@@ -0,0 +1,131 @@
	1	+# TPUT
	2	+
	3	+## Introduction
	4	+
	5	+Remote write operation: write local data to remote NPU's memory. Data is transferred via a UB tile as intermediate staging buffer.
	6	+
	7	+When the GlobalTensor exceeds the UB tile capacity, TPUT automatically performs 2D sliding — chunking rows (DIM_3) and columns (DIM_4) to fit each chunk into the tile, iterating over all outer dimensions (DIM_0, DIM_1, DIM_2).
	8	+
	9	+## Math Interpretation
	10	+
	11	+For each element `(i, j)` in the valid region:
	12	+
	13	+$$ \mathrm{dst}^{\mathrm{remote}}_{i,j} = \mathrm{src}^{\mathrm{local}}_{i,j} $$
	14	+
	15	+Data flow: `srcGlobalData (local GM)` → `stagingTileData (UB)` → `dstGlobalData (remote GM)`
	16	+
	17	+## Assembly Syntax
	18	+
	19	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	20	+
	21	+Synchronous form:
	22	+
	23	+```text
	24	+tput %dst_remote, %src_local : (!pto.memref<...>, !pto.memref<...>)
	25	+```
	26	+Lowering introduces UB staging tile(s) for the GM→UB→GM data path; the C++ intrinsic requires explicit `stagingTileData` (or `pingTile` / `pongTile`) operand(s).
	27	+
	28	+## C++ Intrinsic
	29	+
	30	+Declared in `include/pto/comm/pto_comm_inst.hpp`
	31	+
	32	+### Single-tile (auto-chunking)
	33	+
	34	+```cpp
	35	+template <AtomicType atomicType = AtomicType::AtomicNone,
	36	+ typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	37	+PTO_INST RecordEvent TPUT(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
	38	+ TileData &stagingTileData, WaitEvents&... events);
	39	+```
	40	+
	41	+### Ping-pong double buffering
	42	+
	43	+Uses two staging tiles to overlap TLOAD and TSTORE for adjacent chunks, hiding one DMA transfer behind the other.
	44	+
	45	+```cpp
	46	+template <AtomicType atomicType = AtomicType::AtomicNone,
	47	+ typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	48	+PTO_INST RecordEvent TPUT(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
	49	+ TileData &pingTile, TileData &pongTile, WaitEvents&... events);
	50	+```
	51	+
	52	+### Runtime atomic type
	53	+
	54	+```cpp
	55	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	56	+PTO_INST RecordEvent TPUT(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData,
	57	+ TileData &stagingTileData, AtomicType atomicType, WaitEvents&... events);
	58	+```
	59	+
	60	+## Constraints
	61	+
	62	+- Type constraints:
	63	+ - `GlobalSrcData::RawDType` must equal `GlobalDstData::RawDType`.
	64	+ - `TileData::DType` must equal `GlobalSrcData::RawDType`.
	65	+ - `GlobalSrcData::layout` must equal `GlobalDstData::layout`.
	66	+- Memory constraints:
	67	+ - `dstGlobalData` must point to remote address (on target NPU).
	68	+ - `srcGlobalData` must point to local address (on current NPU).
	69	+ - `stagingTileData` / `pingTile` / `pongTile` must be pre-allocated in Unified Buffer.
	70	+- Valid region:
	71	+ - Transfer size is determined by `GlobalTensor` shape (auto-chunked to fit tile).
	72	+- Atomic operation:
	73	+ - `atomicType` supports `AtomicNone` and `AtomicAdd`.
	74	+- Ping-pong:
	75	+ - `pingTile` and `pongTile` must have the same type and dimensions.
	76	+ - Must reside at non-overlapping UB offsets.
	77	+
	78	+## Examples
	79	+
	80	+### Basic Usage
	81	+
	82	+```cpp
	83	+#include <pto/comm/pto_comm_inst.hpp>
	84	+#include <pto/pto-inst.hpp>
	85	+
	86	+using namespace pto;
	87	+
	88	+template <typename T>
	89	+void example_tput(__gm__ T* local_data, __gm__ T* remote_addr) {
	90	+ using TileT = Tile<TileType::Vec, T, 16, 16>;
	91	+ using GShape = Shape<1, 1, 1, 16, 16>;
	92	+ using GStride = BaseShape2D<T, 16, 16, Layout::ND>;
	93	+ /*
	94	+ If the globalTensor is larger than UB Tile, TPUT will perform 2D sliding automatically.
	95	+ using GShape = Shape<1, 1, 1, 4096, 4096>;
	96	+ using GStride = BaseShape2D<T, 4096, 4096, Layout::ND>;
	97	+ */
	98	+ using GTensor = GlobalTensor<T, GShape, GStride, Layout::ND>;
	99	+
	100	+ GTensor srcG(local_data);
	101	+ GTensor dstG(remote_addr);
	102	+ TileT stagingTile;
	103	+ TASSIGN(stagingTile, 0);
	104	+
	105	+ // Basic remote write
	106	+ comm::TPUT(dstG, srcG, stagingTile);
	107	+
	108	+ // Remote write with atomic add
	109	+ comm::TPUT<AtomicType::AtomicAdd>(dstG, srcG, stagingTile);
	110	+}
	111	+```
	112	+
	113	+### Ping-pong Double Buffering
	114	+
	115	+```cpp
	116	+constexpr size_t tileUBBytes = ((64 * 64 * sizeof(float) + 1023) / 1024) * 1024;
	117	+TileT pingTile(64, 64);
	118	+TileT pongTile(64, 64);
	119	+TASSIGN(pingTile, 0);
	120	+TASSIGN(pongTile, tileUBBytes); // Non-overlapping UB region
	121	+
	122	+// Overlaps TLOAD[i+1] with TSTORE[i] for better pipeline utilization
	123	+comm::TPUT(dstG, srcG, pingTile, pongTile);
	124	+```
	125	+
	126	+### Runtime Atomic Type
	127	+
	128	+```cpp
	129	+// Select atomic type at runtime instead of compile-time template parameter
	130	+comm::TPUT(dstG, srcG, stagingTile, AtomicType::AtomicAdd);
	131	+```

Adocs/isa/comm/TREDUCE.md+118-0

@@ -0,0 +1,118 @@
	1	+# TREDUCE
	2	+
	3	+## Introduction
	4	+
	5	+Reduce operation: gather data from multiple remote NPUs and perform element-wise reduction locally.
	6	+
	7	+
	8	+Only the root needs to execute `TREDUCE`. Non-root ranks only need to ensure their source buffers are ready and remain valid for the duration of the operation. Calling `TREDUCE` on non-root ranks is undefined behavior.
	9	+
	10	+Large Tile Support: When the GlobalTensor exceeds the UB tile capacity in rows and/or columns, the reduction is automatically chunked via 2D sliding.
	11	+
	12	+## Math Interpretation
	13	+
	14	+For each element `(i, j)` in the valid region:
	15	+
	16	+$$ \mathrm{dst}^{\mathrm{local}}_{i,j} = \bigoplus_{r=0}^{N-1} \mathrm{src}^{(r)}_{i,j} $$
	17	+
	18	+where $N$ is the number of ranks and $\oplus$ is the reduction operation (sum, max, min, etc.).
	19	+
	20	+## Assembly Syntax
	21	+
	22	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	23	+
	24	+Synchronous form:
	25	+
	26	+```text
	27	+treduce %group, %dst {op = #pto.reduce_op<Sum>} : (!pto.group<...>, !pto.memref<...>)
	28	+treduce %group, %dst {op = #pto.reduce_op<Max>} : (!pto.group<...>, !pto.memref<...>)
	29	+```
	30	+Lowering introduces internal accumulator and receive tiles for the reduce pipeline; the C++ intrinsic requires explicit `accTileData`, `recvTileData` (or `accTileData`, `pingTileData`, `pongTileData`) operand(s).
	31	+
	32	+## C++ Intrinsic
	33	+
	34	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	35	+
	36	+```cpp
	37	+// Basic reduce (accumulator + receive tile)
	38	+template <typename ParallelGroupType, typename GlobalDstData, typename TileData, typename... WaitEvents>
	39	+PTO_INST RecordEvent TREDUCE(ParallelGroupType &parallelGroup, GlobalDstData &dstGlobalData,
	40	+ TileData &accTileData, TileData &recvTileData, ReduceOp op, WaitEvents&... events);
	41	+
	42	+// Ping-pong reduce (accumulator + ping + pong tiles for double buffering)
	43	+template <typename ParallelGroupType, typename GlobalDstData, typename TileData, typename... WaitEvents>
	44	+PTO_INST RecordEvent TREDUCE(ParallelGroupType &parallelGroup, GlobalDstData &dstGlobalData,
	45	+ TileData &accTileData, TileData &pingTileData, TileData &pongTileData,
	46	+ ReduceOp op, WaitEvents&... events);
	47	+```
	48	+
	49	+## Constraints
	50	+
	51	+- Type constraints:
	52	+ - `ParallelGroup::value_type::RawDType` must equal `GlobalDstData::RawDType`.
	53	+ - `TileData::DType` must equal `GlobalDstData::RawDType`.
	54	+- Memory constraints:
	55	+ - `dstGlobalData` must point to local address (on current NPU).
	56	+ - `accTileData`, `recvTileData` (or `accTileData`, `pingTileData`, `pongTileData`) must be pre-allocated UB tiles.
	57	+- ParallelGroup constraints:
	58	+ - `parallelGroup.tensors[r]` must refer to rank `r`'s source buffer (remote GM as seen by the root).
	59	+ - `parallelGroup.GetRootIdx()` identifies the calling NPU as the reduce root.
	60	+ - All source tensors are assumed to have the same shape and strides.
	61	+- Chunked mode constraints (when data exceeds a single UB tile):
	62	+ - If `TileData` has static `ValidRow`, `GetShape(DIM_3)` must be divisible by `ValidRow`. Use a Tile with `DYNAMIC` ValidRow for partial row support.
	63	+ - If `TileData` has static `ValidCol`, `GetShape(DIM_4)` must be divisible by `ValidCol`. Use a Tile with `DYNAMIC` ValidCol for partial column support.
	64	+
	65	+## Examples
	66	+
	67	+### Basic Reduce Sum
	68	+
	69	+```cpp
	70	+#include <pto/comm/pto_comm_inst.hpp>
	71	+
	72	+using namespace pto;
	73	+
	74	+template <typename T, int SIZE, int NRANKS>
	75	+void reduce_sum(__gm__ T* group_addrs[NRANKS], __gm__ T* result, int my_rank) {
	76	+ using TileT = Tile<TileType::Vec, T, 1, SIZE>;
	77	+ using GTensor = GlobalTensor<T, Shape<1,1,1,1,SIZE>,
	78	+ BaseShape2D<T, 1, SIZE, Layout::ND>, Layout::ND>;
	79	+
	80	+ // Stack-allocated tensors
	81	+ GTensor tensors[NRANKS];
	82	+ for (int i = 0; i < NRANKS; ++i) {
	83	+ tensors[i] = GTensor(group_addrs[i]);
	84	+ }
	85	+
	86	+ comm::ParallelGroup<GTensor> group(tensors, NRANKS, my_rank);
	87	+ GTensor dstG(result);
	88	+ TileT accTile, recvTile;
	89	+
	90	+ comm::TREDUCE(group, dstG, accTile, recvTile, comm::ReduceOp::Sum);
	91	+}
	92	+```
	93	+
	94	+### Max Reduce
	95	+
	96	+```cpp
	97	+#include <pto/comm/pto_comm_inst.hpp>
	98	+
	99	+using namespace pto;
	100	+
	101	+template <typename T, int SIZE, int NRANKS>
	102	+void reduce_max(__gm__ T* group_addrs[NRANKS], __gm__ T* result, int my_rank) {
	103	+ using TileT = Tile<TileType::Vec, T, 1, SIZE>;
	104	+ using GTensor = GlobalTensor<T, Shape<1,1,1,1,SIZE>,
	105	+ BaseShape2D<T, 1, SIZE, Layout::ND>, Layout::ND>;
	106	+
	107	+ GTensor tensors[NRANKS];
	108	+ for (int i = 0; i < NRANKS; ++i) {
	109	+ tensors[i] = GTensor(group_addrs[i]);
	110	+ }
	111	+
	112	+ comm::ParallelGroup<GTensor> group(tensors, NRANKS, my_rank);
	113	+ GTensor dstG(result);
	114	+ TileT accTile, recvTile;
	115	+
	116	+ comm::TREDUCE(group, dstG, accTile, recvTile, comm::ReduceOp::Max);
	117	+}
	118	+```

Adocs/isa/comm/TSCATTER.md+126-0

@@ -0,0 +1,126 @@
	1	+# TSCATTER
	2	+
	3	+## Introduction
	4	+
	5	+Scatter operation: the calling NPU (root) distributes data to all ranks in the parallel group by splitting the local source tensor along DIM_3 (row dimension). This is the inverse of `TGATHER`.
	6	+
	7	+
	8	+Only the root needs to execute `TSCATTER`. Non-root ranks only need to ensure their destination buffers are allocated and writable for the duration of the operation. Calling `TSCATTER` on non-root ranks is undefined behavior.
	9	+
	10	+Large Tile Support: When the per-rank data exceeds the UB tile capacity in rows and/or columns, the transfer is automatically chunked via 2D sliding.
	11	+
	12	+## Math Interpretation
	13	+
	14	+The local source tensor has shape $(D_0, D_1, D_2, N \times H, W)$, where $N$ is the number of ranks and each rank receives $H$ rows. After the operation:
	15	+
	16	+$$\mathrm{dst}^{(r)}_{d_0, d_1, d_2,\; i,\; j} = \mathrm{src}^{\mathrm{local}}_{d_0, d_1, d_2,\; r \cdot H + i,\; j} \quad \forall\, r \in [0, N),\; i \in [0, H),\; j \in [0, W)$$
	17	+
	18	+## Assembly Syntax
	19	+
	20	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	21	+
	22	+Synchronous form:
	23	+
	24	+```text
	25	+tscatter %group, %src : (!pto.group<...>, !pto.memref<...>)
	26	+```
	27	+Lowering introduces UB staging tile(s) for the GM→UB→GM data path; the C++ intrinsic requires explicit `stagingTileData` (or `pingTile` / `pongTile`) operand(s).
	28	+
	29	+## C++ Intrinsic
	30	+
	31	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	32	+
	33	+```cpp
	34	+// Basic scatter (single staging tile)
	35	+template <typename ParallelGroupType, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	36	+PTO_INST RecordEvent TSCATTER(ParallelGroupType &parallelGroup, GlobalSrcData &srcGlobalData,
	37	+ TileData &stagingTileData, WaitEvents&... events);
	38	+
	39	+// Ping-pong scatter (double buffering with two staging tiles)
	40	+template <typename ParallelGroupType, typename GlobalSrcData, typename TileData, typename... WaitEvents>
	41	+PTO_INST RecordEvent TSCATTER(ParallelGroupType &parallelGroup, GlobalSrcData &srcGlobalData,
	42	+ TileData &pingTile, TileData &pongTile, WaitEvents&... events);
	43	+```
	44	+
	45	+## Constraints
	46	+
	47	+- Type constraints:
	48	+ - `ParallelGroup::value_type::RawDType` must equal `GlobalSrcData::RawDType`.
	49	+ - `TileData::DType` must equal `GlobalSrcData::RawDType`.
	50	+- Memory constraints:
	51	+ - `srcGlobalData` must point to local memory (current NPU) and be large enough to hold data for all ranks. Specifically, `srcGlobalData.GetShape(DIM_3)` must be $\geq N \times H$ where $H$ is each rank's `GetShape(DIM_3)`.
	52	+ - If `srcGlobalData.GetShape(DIM_3) > N × H`, only the first `N × H` rows are read; remaining rows are ignored.
	53	+ - `stagingTileData` (or `pingTile` / `pongTile`) must be pre-allocated in UB.
	54	+- ParallelGroup constraints:
	55	+ - `parallelGroup.tensors[r]` must refer to rank `r`'s destination buffer (remote GM as seen by the root).
	56	+ - `parallelGroup.GetRootIdx()` identifies the calling NPU as the scatter root.
	57	+ - All destination tensors are assumed to have the same shape and strides; behavior is undefined if they differ.
	58	+- Chunked mode constraints (when per-rank data exceeds a single UB tile):
	59	+ - If `TileData` has static `ValidRow`, `GetShape(DIM_3)` of each rank's destination must be divisible by `ValidRow`. Use a Tile with `DYNAMIC` ValidRow for partial row support.
	60	+ - If `TileData` has static `ValidCol`, `GetShape(DIM_4)` must be divisible by `ValidCol`. Use a Tile with `DYNAMIC` ValidCol for partial column support.
	61	+
	62	+## Examples
	63	+
	64	+### Basic Scatter (Single Staging Tile)
	65	+
	66	+Root has `NRANKS * ROWS` rows of width `COLS`. Each rank receives `ROWS × COLS`, split along DIM_3.
	67	+The tile size (`TILE_ROWS × TILE_COLS`) can be smaller than the per-rank data — when it is, the implementation automatically chunks the transfer along both DIM_3 and DIM_4 via 2D sliding.
	68	+
	69	+```cpp
	70	+#include <pto/comm/pto_comm_inst.hpp>
	71	+
	72	+using namespace pto;
	73	+
	74	+template <typename T, int ROWS, int COLS, int TILE_ROWS, int TILE_COLS, int NRANKS>
	75	+void scatter(__gm__ T* local_data, __gm__ T* group_addrs[NRANKS], int my_rank) {
	76	+ using TileT = Tile<TileType::Vec, T, TILE_ROWS, TILE_COLS, BLayout::RowMajor, -1, -1>;
	77	+ using GPerRank = GlobalTensor<T, Shape<1,1,1,ROWS,COLS>,
	78	+ BaseShape2D<T, ROWS, COLS, Layout::ND>, Layout::ND>;
	79	+ using GSource = GlobalTensor<T, Shape<1,1,1,NRANKS*ROWS,COLS>,
	80	+ BaseShape2D<T, NRANKS*ROWS, COLS, Layout::ND>, Layout::ND>;
	81	+
	82	+ GPerRank tensors[NRANKS];
	83	+ for (int i = 0; i < NRANKS; ++i) {
	84	+ tensors[i] = GPerRank(group_addrs[i]);
	85	+ }
	86	+
	87	+ comm::ParallelGroup<GPerRank> group(tensors, NRANKS, my_rank);
	88	+ GSource srcG(local_data);
	89	+ TileT stagingTile(TILE_ROWS, TILE_COLS);
	90	+
	91	+ comm::TSCATTER(group, srcG, stagingTile);
	92	+}
	93	+```
	94	+
	95	+### Ping-Pong Scatter (Double Buffering)
	96	+
	97	+Uses two UB tiles to overlap TLOAD of the next chunk (MTE2) with TSTORE of the current chunk (MTE3).
	98	+
	99	+```cpp
	100	+#include <pto/comm/pto_comm_inst.hpp>
	101	+
	102	+using namespace pto;
	103	+
	104	+template <typename T, int ROWS, int COLS, int TILE_ROWS, int TILE_COLS, int NRANKS>
	105	+void scatter_pingpong(__gm__ T* local_data, __gm__ T* group_addrs[NRANKS], int my_rank) {
	106	+ // Tile can be smaller than the data in both dimensions
	107	+ using TileT = Tile<TileType::Vec, T, TILE_ROWS, TILE_COLS, BLayout::RowMajor, -1, -1>;
	108	+ using GPerRank = GlobalTensor<T, Shape<1,1,1,ROWS,COLS>,
	109	+ BaseShape2D<T, ROWS, COLS, Layout::ND>, Layout::ND>;
	110	+ using GSource = GlobalTensor<T, Shape<1,1,1,NRANKS*ROWS,COLS>,
	111	+ BaseShape2D<T, NRANKS*ROWS, COLS, Layout::ND>, Layout::ND>;
	112	+
	113	+ GPerRank tensors[NRANKS];
	114	+ for (int i = 0; i < NRANKS; ++i) {
	115	+ tensors[i] = GPerRank(group_addrs[i]);
	116	+ }
	117	+
	118	+ comm::ParallelGroup<GPerRank> group(tensors, NRANKS, my_rank);
	119	+ GSource srcG(local_data);
	120	+ TileT pingTile(TILE_ROWS, TILE_COLS);
	121	+ TileT pongTile(TILE_ROWS, TILE_COLS);
	122	+
	123	+ // Ping-pong: overlaps TLOAD and TSTORE for better throughput
	124	+ comm::TSCATTER(group, srcG, pingTile, pongTile);
	125	+}
	126	+```

Adocs/isa/comm/TTEST.md+150-0

@@ -0,0 +1,150 @@
	1	+# TTEST
	2	+
	3	+## Introduction
	4	+
	5	+Non-blocking test if signal(s) meet comparison condition. Returns `true` if condition is satisfied, `false` otherwise. Used for polling-based synchronization with timeout or interleaved work.
	6	+
	7	+Supports single signal or multi-dimensional signal tensor (up to 5-D, shape derived from GlobalTensor). For tensor, returns `true` only if ALL signals meet the condition.
	8	+
	9	+## Math Interpretation
	10	+
	11	+Test and return result:
	12	+
	13	+Single signal:
	14	+
	15	+$$ \mathrm{result} = (\mathrm{signal} \;\mathtt{cmp}\; \mathrm{cmpValue}) $$
	16	+
	17	+Signal tensor (all must satisfy):
	18	+
	19	+$$ \mathrm{result} = \bigwedge_{d_0, d_1, d_2, d_3, d_4} (\mathrm{signal}_{d_0, d_1, d_2, d_3, d_4} \;\mathtt{cmp}\; \mathrm{cmpValue}) $$
	20	+
	21	+where `cmp` ∈ {`EQ`, `NE`, `GT`, `GE`, `LT`, `LE`}
	22	+
	23	+## Assembly Syntax
	24	+
	25	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	26	+
	27	+```text
	28	+%result = ttest %signal, %cmp_value {cmp = #pto.cmp<EQ>} : (!pto.memref<i32>, i32) -> i1
	29	+%result = ttest %signal_matrix, %cmp_value {cmp = #pto.cmp<GE>} : (!pto.memref<i32, MxN>, i32) -> i1
	30	+```
	31	+
	32	+## C++ Intrinsic
	33	+
	34	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	35	+
	36	+```cpp
	37	+template <typename GlobalSignalData, typename... WaitEvents>
	38	+PTO_INST bool TTEST(GlobalSignalData &signalData, int32_t cmpValue, WaitCmp cmp, WaitEvents&... events);
	39	+```
	40	+
	41	+## Constraints
	42	+
	43	+- Type constraints:
	44	+ - `GlobalSignalData::DType` must be `int32_t` (32-bit signal).
	45	+- Memory constraints:
	46	+ - `signalData` must point to local address (on current NPU).
	47	+- Return value:
	48	+ - Returns `true` if condition is satisfied, `false` otherwise.
	49	+ - For signal tensor, returns `true` only if ALL signals satisfy the condition.
	50	+- Shape semantics:
	51	+ - For single signal: Shape is `<1,1,1,1,1>`.
	52	+ - For signal tensor: Shape determines the multi-dimensional region (up to 5-D) to test.
	53	+- Comparison operators (WaitCmp):
	54	+ \| Value \| Condition \|
	55	+ \|-------\|-----------\|
	56	+ \| `EQ` \| `signal == cmpValue` \|
	57	+ \| `NE` \| `signal != cmpValue` \|
	58	+ \| `GT` \| `signal > cmpValue` \|
	59	+ \| `GE` \| `signal >= cmpValue` \|
	60	+ \| `LT` \| `signal < cmpValue` \|
	61	+ \| `LE` \| `signal <= cmpValue` \|
	62	+
	63	+## Examples
	64	+
	65	+### Basic Test
	66	+
	67	+```cpp
	68	+#include <pto/comm/pto_comm_inst.hpp>
	69	+
	70	+using namespace pto;
	71	+
	72	+bool check_ready(__gm__ int32_t* local_signal) {
	73	+ comm::Signal sig(local_signal);
	74	+
	75	+ // Check if signal == 1
	76	+ return comm::TTEST(sig, 1, comm::WaitCmp::EQ);
	77	+}
	78	+```
	79	+
	80	+### Test Signal Matrix
	81	+
	82	+```cpp
	83	+#include <pto/comm/pto_comm_inst.hpp>
	84	+
	85	+using namespace pto;
	86	+
	87	+// Test if all signals from a 4x8 dense grid of workers are ready
	88	+bool check_worker_grid(__gm__ int32_t* signal_matrix) {
	89	+ comm::Signal2D<4, 8> grid(signal_matrix);
	90	+
	91	+ // Returns true only if all 32 signals == 1
	92	+ return comm::TTEST(grid, 1, comm::WaitCmp::EQ);
	93	+}
	94	+```
	95	+
	96	+### Polling with Timeout
	97	+
	98	+```cpp
	99	+#include <pto/comm/pto_comm_inst.hpp>
	100	+
	101	+using namespace pto;
	102	+
	103	+bool poll_with_timeout(__gm__ int32_t* local_signal, int max_iterations) {
	104	+ comm::Signal sig(local_signal);
	105	+
	106	+ for (int i = 0; i < max_iterations; ++i) {
	107	+ if (comm::TTEST(sig, 1, comm::WaitCmp::EQ)) {
	108	+ return true; // Signal received
	109	+ }
	110	+ // Could do other work here between polls
	111	+ }
	112	+ return false; // Timeout
	113	+}
	114	+```
	115	+
	116	+### Progress-Based Polling
	117	+
	118	+```cpp
	119	+#include <pto/comm/pto_comm_inst.hpp>
	120	+
	121	+using namespace pto;
	122	+
	123	+void process_with_progress(__gm__ int32_t* local_counter, int expected_count) {
	124	+ comm::Signal counter(local_counter);
	125	+
	126	+ while (!comm::TTEST(counter, expected_count, comm::WaitCmp::GE)) {
	127	+ // Do some useful work while waiting
	128	+ // ...
	129	+ }
	130	+ // All expected signals received
	131	+}
	132	+```
	133	+
	134	+### Compare TWAIT vs TTEST
	135	+
	136	+```cpp
	137	+#include <pto/comm/pto_comm_inst.hpp>
	138	+
	139	+using namespace pto;
	140	+
	141	+void compare_wait_test(__gm__ int32_t* local_signal) {
	142	+ comm::Signal sig(local_signal);
	143	+
	144	+ // Blocking: spins until signal == 1
	145	+ comm::TWAIT(sig, 1, comm::WaitCmp::EQ);
	146	+
	147	+ // Non-blocking: returns immediately with result
	148	+ bool ready = comm::TTEST(sig, 1, comm::WaitCmp::EQ);
	149	+}
	150	+```

Adocs/isa/comm/TWAIT.md+131-0

@@ -0,0 +1,131 @@
	1	+# TWAIT
	2	+
	3	+## Introduction
	4	+
	5	+Blocking wait until signal(s) meet comparison condition. Used in conjunction with `TNOTIFY` for flag-based synchronization.
	6	+
	7	+Supports single signal or multi-dimensional signal tensor (up to 5-D, shape derived from GlobalTensor).
	8	+
	9	+
	10	+## Math Interpretation
	11	+
	12	+Wait (spin) until the following condition is satisfied:
	13	+
	14	+Single signal:
	15	+
	16	+$$ \mathrm{signal} \;\mathtt{cmp}\; \mathrm{cmpValue} $$
	17	+
	18	+Signal tensor (all elements must satisfy):
	19	+
	20	+$$ \forall d_0, d_1, d_2, d_3, d_4: \mathrm{signal}_{d_0, d_1, d_2, d_3, d_4} \;\mathtt{cmp}\; \mathrm{cmpValue} $$
	21	+
	22	+where `cmp` ∈ {`EQ`, `NE`, `GT`, `GE`, `LT`, `LE`}
	23	+
	24	+## Assembly Syntax
	25	+
	26	+PTO-AS form: see `docs/grammar/PTO-AS.md`.
	27	+
	28	+```text
	29	+twait %signal, %cmp_value {cmp = #pto.cmp<EQ>} : (!pto.memref<i32>, i32)
	30	+twait %signal_matrix, %cmp_value {cmp = #pto.cmp<GE>} : (!pto.memref<i32, MxN>, i32)
	31	+```
	32	+
	33	+## C++ Intrinsic
	34	+
	35	+Declared in `include/pto/comm/pto_comm_inst.hpp`:
	36	+
	37	+```cpp
	38	+template <typename GlobalSignalData, typename... WaitEvents>
	39	+PTO_INST void TWAIT(GlobalSignalData &signalData, int32_t cmpValue, WaitCmp cmp, WaitEvents&... events);
	40	+```
	41	+
	42	+## Constraints
	43	+
	44	+- Type constraints:
	45	+ - `GlobalSignalData::DType` must be `int32_t` (32-bit signal).
	46	+- Memory constraints:
	47	+ - `signalData` must point to local address (on current NPU).
	48	+- Shape semantics:
	49	+ - For single signal: Shape is `<1,1,1,1,1>`.
	50	+ - For signal tensor: Shape determines the multi-dimensional region (up to 5-D) to wait on. All signals in the tensor must satisfy the condition.
	51	+- Comparison operators (WaitCmp):
	52	+ \| Value \| Condition \|
	53	+ \|-------\|-----------\|
	54	+ \| `EQ` \| `signal == cmpValue` \|
	55	+ \| `NE` \| `signal != cmpValue` \|
	56	+ \| `GT` \| `signal > cmpValue` \|
	57	+ \| `GE` \| `signal >= cmpValue` \|
	58	+ \| `LT` \| `signal < cmpValue` \|
	59	+ \| `LE` \| `signal <= cmpValue` \|
	60	+
	61	+## Examples
	62	+
	63	+### Wait for Single Signal
	64	+
	65	+```cpp
	66	+#include <pto/comm/pto_comm_inst.hpp>
	67	+
	68	+using namespace pto;
	69	+
	70	+void wait_for_ready(__gm__ int32_t* local_signal) {
	71	+ comm::Signal sig(local_signal);
	72	+
	73	+ // Wait until signal == 1
	74	+ comm::TWAIT(sig, 1, comm::WaitCmp::EQ);
	75	+}
	76	+```
	77	+
	78	+### Wait for Signal Matrix
	79	+
	80	+```cpp
	81	+#include <pto/comm/pto_comm_inst.hpp>
	82	+
	83	+using namespace pto;
	84	+
	85	+// Wait for signals from a 4x8 dense grid of workers
	86	+void wait_worker_grid(__gm__ int32_t* signal_matrix) {
	87	+ comm::Signal2D<4, 8> grid(signal_matrix);
	88	+
	89	+ // Wait until all 32 signals == 1
	90	+ comm::TWAIT(grid, 1, comm::WaitCmp::EQ);
	91	+}
	92	+```
	93	+
	94	+### Wait for Counter Threshold
	95	+
	96	+```cpp
	97	+#include <pto/comm/pto_comm_inst.hpp>
	98	+
	99	+using namespace pto;
	100	+
	101	+void wait_for_count(__gm__ int32_t* local_counter, int expected_count) {
	102	+ comm::Signal counter(local_counter);
	103	+
	104	+ // Wait until counter >= expected_count
	105	+ comm::TWAIT(counter, expected_count, comm::WaitCmp::GE);
	106	+}
	107	+```
	108	+
	109	+### Producer-Consumer Pattern
	110	+
	111	+```cpp
	112	+#include <pto/comm/pto_comm_inst.hpp>
	113	+
	114	+using namespace pto;
	115	+
	116	+// Producer: notify when data is ready
	117	+void producer(__gm__ int32_t* remote_flag) {
	118	+ // ... produce data ...
	119	+
	120	+ comm::Signal flag(remote_flag);
	121	+ comm::TNOTIFY(flag, 1, comm::NotifyOp::Set);
	122	+}
	123	+
	124	+// Consumer: wait for data
	125	+void consumer(__gm__ int32_t* local_flag) {
	126	+ comm::Signal flag(local_flag);
	127	+ comm::TWAIT(flag, 1, comm::WaitCmp::EQ);
	128	+
	129	+ // ... consume data ...
	130	+}
	131	+```

Ainclude/pto/comm/TBroadCast.hpp+421-0

@@ -0,0 +1,421 @@
	1	+/**
	2	+Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+CANN Open Software License Agreement Version 2.0 (the "License").
	5	+Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+See LICENSE in the root of the software repository for the full text of the License.
	9	+*/
	10	+
	11	+#ifndef PTO_COMM_TBROADCAST_HPP
	12	+#define PTO_COMM_TBROADCAST_HPP
	13	+
	14	+#include <type_traits>
	15	+
	16	+#include "pto/common/debug.h"
	17	+#include "pto/common/type.hpp"
	18	+#include "pto/common/constants.hpp"
	19	+#include "pto/common/pto_instr.hpp"
	20	+#include "pto/comm/comm_types.hpp"
	21	+
	22	+namespace pto {
	23	+namespace comm {
	24	+
	25	+// ============================================================================
	26	+// TBROADCAST_IMPL: Broadcast data from root NPU to all ranks
	27	+//
	28	+// The root loads srcGlobalData and stores it to every rank's buffer in the
	29	+// ParallelGroup.
	30	+//
	31	+// When the GlobalTensor exceeds the UB tile capacity in rows and/or columns,
	32	+// the transfer is automatically chunked via 2D sliding:
	33	+// - Outer dimensions (DIM_0, DIM_1, DIM_2) are iterated explicitly.
	34	+// - DIM_3 (rows) is split into tileValidRow-sized chunks.
	35	+// - DIM_4 (cols) is split into tileValidCol-sized chunks.
	36	+//
	37	+// Constraints for chunked mode:
	38	+// - If TileData has static ValidRow, shape3 must be divisible by ValidRow.
	39	+// - If TileData has static ValidCol, shape4 must be divisible by ValidCol.
	40	+// - All ranks in the ParallelGroup are assumed to have the same shape/strides.
	41	+// ============================================================================
	42	+
	43	+template <typename ParallelGroupType, typename GlobalSrcData, typename TileData>
	44	+PTO_INTERNAL void TBROADCAST_IMPL(ParallelGroupType &parallelGroup, GlobalSrcData &srcGlobalData,
	45	+ TileData &stagingTileData)
	46	+{
	47	+ using GlobalDstData = typename ParallelGroupTraits<ParallelGroupType>::GlobalDataType;
	48	+ using T = typename GlobalSrcData::RawDType;
	49	+
	50	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	51	+ "TBROADCAST: TileData element type must match GlobalData element type");
	52	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>,
	53	+ "TBROADCAST: ParallelGroup element type must match source element type");
	54	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TBROADCAST: src/dst layout mismatch");
	55	+
	56	+ const int nranks = parallelGroup.GetSize();
	57	+ const int rootIdx = parallelGroup.GetRootIdx();
	58	+
	59	+ PTO_ASSERT(nranks > 0, "ParallelGroup size must be greater than 0!");
	60	+ PTO_ASSERT(rootIdx >= 0 && rootIdx < nranks, "rootIdx must be in range [0, nranks)!");
	61	+
	62	+ // Get GlobalTensor dimensions from source
	63	+ const int gShape0 = srcGlobalData.GetShape(GlobalTensorDim::DIM_0);
	64	+ const int gShape1 = srcGlobalData.GetShape(GlobalTensorDim::DIM_1);
	65	+ const int gShape2 = srcGlobalData.GetShape(GlobalTensorDim::DIM_2);
	66	+ const int gShape3 = srcGlobalData.GetShape(GlobalTensorDim::DIM_3);
	67	+ const int gShape4 = srcGlobalData.GetShape(GlobalTensorDim::DIM_4);
	68	+
	69	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	70	+ const int tileValidRow = stagingTileData.GetValidRow();
	71	+ const int tileValidCol = stagingTileData.GetValidCol();
	72	+
	73	+ PTO_ASSERT(tileValidRow > 0, "TBROADCAST: tileValidRow must be greater than 0");
	74	+ PTO_ASSERT(tileValidCol > 0, "TBROADCAST: tileValidCol must be greater than 0");
	75	+
	76	+ if (totalRows == 0 \|\| gShape4 == 0) {
	77	+ return;
	78	+ }
	79	+
	80	+ // ---- Single rank: copy src to dst[0] ----
	81	+ if (nranks == 1) {
	82	+ TLOAD(stagingTileData, srcGlobalData);
	83	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	84	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	85	+ TSTORE(parallelGroup[rootIdx], stagingTileData);
	86	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	87	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	88	+ return;
	89	+ }
	90	+
	91	+ // ---- Simple path: data fits in UB tile in both dimensions ----
	92	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	93	+ // Root loads data to UB once
	94	+ TLOAD(stagingTileData, srcGlobalData);
	95	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	96	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	97	+
	98	+ // Broadcast to all ranks
	99	+ for (int r = 0; r < nranks; ++r) {
	100	+ TSTORE(parallelGroup[r], stagingTileData);
	101	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	102	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	103	+ }
	104	+ return;
	105	+ }
	106	+
	107	+ // ---- 2D sliding chunked path ----
	108	+ //
	109	+ // Strategy: for each chunk, TLOAD from srcGlobalData into UB tile,
	110	+ // then TSTORE to every rank's destination at the corresponding offset.
	111	+
	112	+ PTO_ASSERT(tileValidRow > 0, "TBROADCAST: tile ValidRow must be greater than 0 for chunked transfer");
	113	+ PTO_ASSERT(tileValidCol > 0, "TBROADCAST: tile ValidCol must be greater than 0 for chunked transfer");
	114	+
	115	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	116	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	117	+
	118	+ if constexpr (!isDynamicRow) {
	119	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	120	+ "TBROADCAST chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. "
	121	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	122	+ }
	123	+ if constexpr (!isDynamicCol) {
	124	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	125	+ "TBROADCAST chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. "
	126	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	127	+ }
	128	+
	129	+ // Source strides
	130	+ const int srcStride0 = srcGlobalData.GetStride(GlobalTensorDim::DIM_0);
	131	+ const int srcStride1 = srcGlobalData.GetStride(GlobalTensorDim::DIM_1);
	132	+ const int srcStride2 = srcGlobalData.GetStride(GlobalTensorDim::DIM_2);
	133	+ const int srcStride3 = srcGlobalData.GetStride(GlobalTensorDim::DIM_3);
	134	+ const int srcStride4 = srcGlobalData.GetStride(GlobalTensorDim::DIM_4);
	135	+
	136	+ // Destination strides (from first rank, assumed same for all)
	137	+ const int dstStride0 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_0);
	138	+ const int dstStride1 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_1);
	139	+ const int dstStride2 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_2);
	140	+ const int dstStride3 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_3);
	141	+ const int dstStride4 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_4);
	142	+
	143	+ // View types with fully dynamic shape/stride
	144	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	145	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	146	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	147	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	148	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	149	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	150	+
	151	+ // 2D sliding: iterate outer dims, then chunk rows (dim3) and columns (dim4)
	152	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	153	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	154	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	155	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	156	+ static_cast<int64_t>(i2) * srcStride2;
	157	+ int64_t dstBase = static_cast<int64_t>(i0) * dstStride0 + static_cast<int64_t>(i1) * dstStride1 +
	158	+ static_cast<int64_t>(i2) * dstStride2;
	159	+
	160	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	161	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	162	+
	163	+ if constexpr (isDynamicRow) {
	164	+ stagingTileData.RowMaskInternal = currentRows;
	165	+ }
	166	+
	167	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	168	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	169	+
	170	+ if constexpr (isDynamicCol) {
	171	+ stagingTileData.ColMaskInternal = currentCols;
	172	+ }
	173	+
	174	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	175	+ static_cast<int64_t>(colOff) * srcStride4;
	176	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	177	+ static_cast<int64_t>(colOff) * dstStride4;
	178	+
	179	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	180	+
	181	+ // TLOAD source chunk into UB
	182	+ SrcViewT srcView(srcGlobalData.data() + srcOffset, chunkShape, srcChunkStride);
	183	+ TLOAD(stagingTileData, srcView);
	184	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	185	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	186	+
	187	+ // TSTORE to all ranks
	188	+ // MTE3 guarantees in-order execution, so no inter-TSTORE sync needed.
	189	+ for (int r = 0; r < nranks; ++r) {
	190	+ DstViewT dstView(parallelGroup[r].data() + dstOffset, chunkShape, dstChunkStride);
	191	+ TSTORE(dstView, stagingTileData);
	192	+ }
	193	+
	194	+ // Sync before next chunk's TLOAD
	195	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	196	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	197	+ }
	198	+ }
	199	+ }
	200	+ }
	201	+ }
	202	+}
	203	+
	204	+// ============================================================================
	205	+// TBROADCAST_IMPL (ping-pong): Broadcast with double buffering
	206	+//
	207	+// Uses two staging tiles (pingTile, pongTile) to overlap TLOAD of the next
	208	+// chunk (MTE2) with TSTORE of the current chunk to all ranks (MTE3).
	209	+//
	210	+// Timeline without ping-pong:
	211	+// [TLOAD chunk0] -> [N×TSTORE chunk0] -> [TLOAD chunk1] -> [N×TSTORE chunk1] -> ...
	212	+//
	213	+// Timeline with ping-pong:
	214	+// [TLOAD chunk0] -> [N×TSTORE chunk0 \| TLOAD chunk1] -> [N×TSTORE chunk1 \| TLOAD chunk2] -> ...
	215	+//
	216	+// Constraints: same as TBROADCAST_IMPL for chunked mode.
	217	+// ============================================================================
	218	+
	219	+template <typename ParallelGroupType, typename GlobalSrcData, typename TileData>
	220	+PTO_INTERNAL void TBROADCAST_IMPL(ParallelGroupType &parallelGroup, GlobalSrcData &srcGlobalData, TileData &pingTile,
	221	+ TileData &pongTile)
	222	+{
	223	+ using GlobalDstData = typename ParallelGroupTraits<ParallelGroupType>::GlobalDataType;
	224	+ using T = typename GlobalSrcData::RawDType;
	225	+
	226	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	227	+ "TBROADCAST: TileData element type must match GlobalData element type");
	228	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>,
	229	+ "TBROADCAST: ParallelGroup element type must match source element type");
	230	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TBROADCAST: src/dst layout mismatch");
	231	+
	232	+ const int nranks = parallelGroup.GetSize();
	233	+ const int rootIdx = parallelGroup.GetRootIdx();
	234	+
	235	+ PTO_ASSERT(nranks > 0, "ParallelGroup size must be greater than 0!");
	236	+ PTO_ASSERT(rootIdx >= 0 && rootIdx < nranks, "rootIdx must be in range [0, nranks)!");
	237	+
	238	+ // Get GlobalTensor dimensions from source
	239	+ const int gShape0 = srcGlobalData.GetShape(GlobalTensorDim::DIM_0);
	240	+ const int gShape1 = srcGlobalData.GetShape(GlobalTensorDim::DIM_1);
	241	+ const int gShape2 = srcGlobalData.GetShape(GlobalTensorDim::DIM_2);
	242	+ const int gShape3 = srcGlobalData.GetShape(GlobalTensorDim::DIM_3);
	243	+ const int gShape4 = srcGlobalData.GetShape(GlobalTensorDim::DIM_4);
	244	+
	245	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	246	+ const int tileValidRow = pingTile.GetValidRow();
	247	+ const int tileValidCol = pingTile.GetValidCol();
	248	+
	249	+ PTO_ASSERT(tileValidRow > 0, "TBROADCAST: tileValidRow must be greater than 0");
	250	+ PTO_ASSERT(tileValidCol > 0, "TBROADCAST: tileValidCol must be greater than 0");
	251	+
	252	+ if (totalRows == 0 \|\| gShape4 == 0) {
	253	+ return;
	254	+ }
	255	+
	256	+ // ---- Single rank: copy src to dst[0] ----
	257	+ if (nranks == 1) {
	258	+ TLOAD(pingTile, srcGlobalData);
	259	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	260	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	261	+ TSTORE(parallelGroup[rootIdx], pingTile);
	262	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	263	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	264	+ return;
	265	+ }
	266	+
	267	+ // ---- Simple path: single chunk, no ping-pong benefit ----
	268	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	269	+ TLOAD(pingTile, srcGlobalData);
	270	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	271	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	272	+
	273	+ for (int r = 0; r < nranks; ++r) {
	274	+ TSTORE(parallelGroup[r], pingTile);
	275	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	276	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	277	+ }
	278	+ return;
	279	+ }
	280	+
	281	+ // ---- 2D sliding chunked path with ping-pong double buffering ----
	282	+
	283	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	284	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	285	+
	286	+ if constexpr (!isDynamicRow) {
	287	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	288	+ "TBROADCAST chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. "
	289	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	290	+ }
	291	+ if constexpr (!isDynamicCol) {
	292	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	293	+ "TBROADCAST chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. "
	294	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	295	+ }
	296	+
	297	+ const int srcStride0 = srcGlobalData.GetStride(GlobalTensorDim::DIM_0);
	298	+ const int srcStride1 = srcGlobalData.GetStride(GlobalTensorDim::DIM_1);
	299	+ const int srcStride2 = srcGlobalData.GetStride(GlobalTensorDim::DIM_2);
	300	+ const int srcStride3 = srcGlobalData.GetStride(GlobalTensorDim::DIM_3);
	301	+ const int srcStride4 = srcGlobalData.GetStride(GlobalTensorDim::DIM_4);
	302	+
	303	+ const int dstStride0 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_0);
	304	+ const int dstStride1 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_1);
	305	+ const int dstStride2 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_2);
	306	+ const int dstStride3 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_3);
	307	+ const int dstStride4 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_4);
	308	+
	309	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	310	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	311	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	312	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	313	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	314	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	315	+
	316	+ // Ping-pong state
	317	+ bool usePing = true;
	318	+ bool hasPending = false;
	319	+ int64_t pendingDstOffset = 0;
	320	+ int pendingRows = 0;
	321	+ int pendingCols = 0;
	322	+
	323	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	324	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	325	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	326	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	327	+ static_cast<int64_t>(i2) * srcStride2;
	328	+ int64_t dstBase = static_cast<int64_t>(i0) * dstStride0 + static_cast<int64_t>(i1) * dstStride1 +
	329	+ static_cast<int64_t>(i2) * dstStride2;
	330	+
	331	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	332	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	333	+
	334	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	335	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	336	+
	337	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	338	+ static_cast<int64_t>(colOff) * srcStride4;
	339	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	340	+ static_cast<int64_t>(colOff) * dstStride4;
	341	+
	342	+ // Select load tile for this iteration
	343	+ TileData &loadTile = usePing ? pingTile : pongTile;
	344	+ event_t curEvent = usePing ? EVENT_ID0 : EVENT_ID1;
	345	+
	346	+ // Configure masks on the load tile
	347	+ if constexpr (isDynamicRow)
	348	+ loadTile.RowMaskInternal = currentRows;
	349	+ if constexpr (isDynamicCol)
	350	+ loadTile.ColMaskInternal = currentCols;
	351	+
	352	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	353	+ SrcViewT srcView(srcGlobalData.data() + srcOffset, chunkShape, srcChunkStride);
	354	+
	355	+ if (hasPending) {
	356	+ // The other tile holds data from the previous TLOAD
	357	+ TileData &storeTile = usePing ? pongTile : pingTile;
	358	+ event_t prevEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	359	+
	360	+ // Wait for previous TLOAD to finish
	361	+ wait_flag(PIPE_MTE2, PIPE_MTE3, prevEvent);
	362	+
	363	+ // Build view for the deferred TSTOREs
	364	+ DynShape pendShape(1, 1, 1, pendingRows, pendingCols);
	365	+
	366	+ // Issue N TSTOREs + TLOAD concurrently (MTE3 and MTE2 in parallel).
	367	+ // MTE3 guarantees in-order execution, so no inter-TSTORE sync needed.
	368	+ for (int r = 0; r < nranks; ++r) {
	369	+ DstViewT dstView(parallelGroup[r].data() + pendingDstOffset, pendShape, dstChunkStride);
	370	+ TSTORE(dstView, storeTile);
	371	+ }
	372	+ TLOAD(loadTile, srcView);
	373	+
	374	+ set_flag(PIPE_MTE3, PIPE_MTE2, prevEvent); // all stores done
	375	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent); // load done
	376	+
	377	+ // Ensure storeTile UB is safe before it can be overwritten
	378	+ wait_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);
	379	+ } else {
	380	+ // First chunk: just issue TLOAD
	381	+ TLOAD(loadTile, srcView);
	382	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent);
	383	+ }
	384	+
	385	+ // Record this chunk as pending
	386	+ pendingDstOffset = dstOffset;
	387	+ pendingRows = currentRows;
	388	+ pendingCols = currentCols;
	389	+ hasPending = true;
	390	+ usePing = !usePing;
	391	+ }
	392	+ }
	393	+ }
	394	+ }
	395	+ }
	396	+
	397	+ // Epilogue: drain the last pending chunk
	398	+ if (hasPending) {
	399	+ TileData &lastTile = usePing ? pongTile : pingTile;
	400	+ event_t lastEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	401	+
	402	+ wait_flag(PIPE_MTE2, PIPE_MTE3, lastEvent);
	403	+
	404	+ DynShape lastShape(1, 1, 1, pendingRows, pendingCols);
	405	+
	406	+ // MTE3 guarantees that consecutive TSTOREs on the same pipe are executed
	407	+ // in order, so no additional inter-TSTORE synchronization is needed here.
	408	+ for (int r = 0; r < nranks; ++r) {
	409	+ DstViewT dstView(parallelGroup[r].data() + pendingDstOffset, lastShape, dstChunkStride);
	410	+ TSTORE(dstView, lastTile);
	411	+ }
	412	+
	413	+ set_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	414	+ wait_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	415	+ }
	416	+}
	417	+
	418	+} // namespace comm
	419	+} // namespace pto
	420	+
	421	+#endif // PTO_COMM_TBROADCAST_HPP

Ainclude/pto/comm/TGather.hpp+441-0

@@ -0,0 +1,441 @@
	1	+/**
	2	+Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+CANN Open Software License Agreement Version 2.0 (the "License").
	5	+Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+See LICENSE in the root of the software repository for the full text of the License.
	9	+*/
	10	+
	11	+#ifndef PTO_COMM_TGATHER_HPP
	12	+#define PTO_COMM_TGATHER_HPP
	13	+
	14	+#include <type_traits>
	15	+
	16	+#include "pto/common/debug.h"
	17	+#include "pto/common/type.hpp"
	18	+#include "pto/common/constants.hpp"
	19	+#include "pto/common/pto_instr.hpp"
	20	+#include "pto/comm/comm_types.hpp"
	21	+
	22	+namespace pto {
	23	+namespace comm {
	24	+
	25	+// ============================================================================
	26	+// TGATHER_IMPL: Gather operation - root collects data from all ranks
	27	+//
	28	+// The calling NPU is the root and gathers data from all ranks, concatenating
	29	+// the results along DIM_3 (row dimension) into a local output buffer.
	30	+//
	31	+// Each rank r contributes data of shape (D0, D1, D2, H, W). The destination
	32	+// tensor has shape (D0, D1, D2, N*H, W), where rank r's data is placed at
	33	+// rows [rH, (r+1)H).
	34	+//
	35	+// When the per-rank GlobalTensor exceeds the UB tile capacity in rows and/or
	36	+// columns, the transfer is automatically chunked via 2D sliding:
	37	+// - Outer dimensions (DIM_0, DIM_1, DIM_2) are iterated explicitly.
	38	+// - DIM_3 (rows) is split into tileValidRow-sized chunks.
	39	+// - DIM_4 (cols) is split into tileValidCol-sized chunks.
	40	+//
	41	+// Constraints for chunked mode:
	42	+// - If TileData has static ValidRow, per-rank DIM_3 must be divisible by ValidRow.
	43	+// - If TileData has static ValidCol, DIM_4 must be divisible by ValidCol.
	44	+// - All source tensors in the ParallelGroup are assumed to have the same shape/strides.
	45	+// ============================================================================
	46	+
	47	+template <typename ParallelGroupType, typename GlobalDstData, typename TileData>
	48	+PTO_INTERNAL void TGATHER_IMPL(ParallelGroupType &parallelGroup, GlobalDstData &dstGlobalData,
	49	+ TileData &stagingTileData)
	50	+{
	51	+ using GlobalSrcData = typename ParallelGroupTraits<ParallelGroupType>::GlobalDataType;
	52	+ using T = typename GlobalSrcData::RawDType;
	53	+
	54	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>, "TGATHER: GlobalData type mismatch!");
	55	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	56	+ "TGATHER: TileData element type must match GlobalData element type");
	57	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TGATHER: src/dst layout mismatch");
	58	+
	59	+ const int nranks = parallelGroup.GetSize();
	60	+ const int rootIdx = parallelGroup.GetRootIdx();
	61	+
	62	+ PTO_ASSERT(nranks > 0, "ParallelGroup size must be greater than 0!");
	63	+ PTO_ASSERT(rootIdx >= 0 && rootIdx < nranks, "rootIdx must be in range [0, nranks)!");
	64	+
	65	+ // Get per-rank dimensions (from first rank, all assumed same)
	66	+ const int gShape0 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_0);
	67	+ const int gShape1 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_1);
	68	+ const int gShape2 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_2);
	69	+ const int gShape3 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_3); // H (per-rank rows)
	70	+ const int gShape4 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_4); // W
	71	+
	72	+ const int perRankRows = gShape3;
	73	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	74	+ const int tileValidRow = stagingTileData.GetValidRow();
	75	+ const int tileValidCol = stagingTileData.GetValidCol();
	76	+
	77	+ PTO_ASSERT(tileValidRow > 0, "TGATHER: tileValidRow must be greater than 0");
	78	+ PTO_ASSERT(tileValidCol > 0, "TGATHER: tileValidCol must be greater than 0");
	79	+
	80	+ if (totalRows == 0 \|\| gShape4 == 0) {
	81	+ return;
	82	+ }
	83	+
	84	+ // ---- Simple path: per-rank data fits in UB tile ----
	85	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	86	+ if (nranks == 1) {
	87	+ // Single rank: direct copy, no offset needed
	88	+ TLOAD(stagingTileData, parallelGroup[0]);
	89	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	90	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	91	+ TSTORE(dstGlobalData, stagingTileData);
	92	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	93	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	94	+ return;
	95	+ }
	96	+
	97	+ // Multiple ranks: need destination views with per-rank offset
	98	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	99	+
	100	+ using DynShape5D = Shape<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	101	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	102	+ using DstViewT = GlobalTensor<T, DynShape5D, DynStride, GlobalDstData::layout>;
	103	+
	104	+ DynShape5D perRankShape(gShape0, gShape1, gShape2, gShape3, gShape4);
	105	+ DynStride dstViewStride(dstGlobalData.GetStride(GlobalTensorDim::DIM_0),
	106	+ dstGlobalData.GetStride(GlobalTensorDim::DIM_1),
	107	+ dstGlobalData.GetStride(GlobalTensorDim::DIM_2), dstStride3,
	108	+ dstGlobalData.GetStride(GlobalTensorDim::DIM_4));
	109	+
	110	+ for (int r = 0; r < nranks; ++r) {
	111	+ TLOAD(stagingTileData, parallelGroup[r]);
	112	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	113	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	114	+
	115	+ int64_t dstOffset = static_cast<int64_t>(r) * perRankRows * dstStride3;
	116	+ DstViewT dstView(dstGlobalData.data() + dstOffset, perRankShape, dstViewStride);
	117	+ TSTORE(dstView, stagingTileData);
	118	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	119	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	120	+ }
	121	+ return;
	122	+ }
	123	+
	124	+ // ---- 2D sliding chunked path ----
	125	+ //
	126	+ // For each rank r, iterate outer dims and chunk rows/cols.
	127	+ // Source: parallelGroup[r] at chunk offset
	128	+ // Destination: dstGlobalData at (rank base + chunk offset)
	129	+
	130	+ PTO_ASSERT(tileValidRow > 0, "TGATHER: tile ValidRow must be greater than 0 for chunked transfer");
	131	+ PTO_ASSERT(tileValidCol > 0, "TGATHER: tile ValidCol must be greater than 0 for chunked transfer");
	132	+
	133	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	134	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	135	+
	136	+ if constexpr (!isDynamicRow) {
	137	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	138	+ "TGATHER chunked: per-rank DIM_3 must be divisible by tile ValidRow when static. "
	139	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	140	+ }
	141	+ if constexpr (!isDynamicCol) {
	142	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	143	+ "TGATHER chunked: DIM_4 must be divisible by tile ValidCol when static. "
	144	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	145	+ }
	146	+
	147	+ // Source strides (from first rank, all assumed same)
	148	+ const int srcStride0 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_0);
	149	+ const int srcStride1 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_1);
	150	+ const int srcStride2 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_2);
	151	+ const int srcStride3 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_3);
	152	+ const int srcStride4 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_4);
	153	+
	154	+ // Destination strides
	155	+ const int dstStride0 = dstGlobalData.GetStride(GlobalTensorDim::DIM_0);
	156	+ const int dstStride1 = dstGlobalData.GetStride(GlobalTensorDim::DIM_1);
	157	+ const int dstStride2 = dstGlobalData.GetStride(GlobalTensorDim::DIM_2);
	158	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	159	+ const int dstStride4 = dstGlobalData.GetStride(GlobalTensorDim::DIM_4);
	160	+
	161	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	162	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	163	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	164	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	165	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	166	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	167	+
	168	+ for (int r = 0; r < nranks; ++r) {
	169	+ int64_t rankDstBase = static_cast<int64_t>(r) * perRankRows * dstStride3;
	170	+
	171	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	172	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	173	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	174	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	175	+ static_cast<int64_t>(i2) * srcStride2;
	176	+ int64_t dstBase = rankDstBase + static_cast<int64_t>(i0) * dstStride0 +
	177	+ static_cast<int64_t>(i1) * dstStride1 + static_cast<int64_t>(i2) * dstStride2;
	178	+
	179	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	180	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	181	+
	182	+ if constexpr (isDynamicRow) {
	183	+ stagingTileData.RowMaskInternal = currentRows;
	184	+ }
	185	+
	186	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	187	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	188	+
	189	+ if constexpr (isDynamicCol) {
	190	+ stagingTileData.ColMaskInternal = currentCols;
	191	+ }
	192	+
	193	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	194	+ static_cast<int64_t>(colOff) * srcStride4;
	195	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	196	+ static_cast<int64_t>(colOff) * dstStride4;
	197	+
	198	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	199	+
	200	+ // TLOAD from rank r's source at chunk position
	201	+ SrcViewT srcView(parallelGroup[r].data() + srcOffset, chunkShape, srcChunkStride);
	202	+ TLOAD(stagingTileData, srcView);
	203	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	204	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	205	+
	206	+ // TSTORE to local destination at rank + chunk position
	207	+ DstViewT dstView(dstGlobalData.data() + dstOffset, chunkShape, dstChunkStride);
	208	+ TSTORE(dstView, stagingTileData);
	209	+
	210	+ // Sync before next chunk's TLOAD
	211	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	212	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	213	+ }
	214	+ }
	215	+ }
	216	+ }
	217	+ }
	218	+ }
	219	+}
	220	+
	221	+// ============================================================================
	222	+// TGATHER_IMPL (ping-pong): Gather with double buffering
	223	+//
	224	+// Uses two staging tiles (pingTile, pongTile) to overlap TLOAD of the next
	225	+// chunk (MTE2) with TSTORE of the current chunk (MTE3).
	226	+//
	227	+// Timeline without ping-pong:
	228	+// [TLOAD chunk0] -> [TSTORE chunk0] -> [TLOAD chunk1] -> [TSTORE chunk1] -> ...
	229	+//
	230	+// Timeline with ping-pong:
	231	+// [TLOAD chunk0] -> [TSTORE chunk0 \| TLOAD chunk1] -> [TSTORE chunk1 \| TLOAD chunk2] -> ...
	232	+//
	233	+// Constraints: same as TGATHER_IMPL for chunked mode.
	234	+// ============================================================================
	235	+
	236	+template <typename ParallelGroupType, typename GlobalDstData, typename TileData>
	237	+PTO_INTERNAL void TGATHER_IMPL(ParallelGroupType &parallelGroup, GlobalDstData &dstGlobalData, TileData &pingTile,
	238	+ TileData &pongTile)
	239	+{
	240	+ using GlobalSrcData = typename ParallelGroupTraits<ParallelGroupType>::GlobalDataType;
	241	+ using T = typename GlobalSrcData::RawDType;
	242	+
	243	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>, "TGATHER: GlobalData type mismatch!");
	244	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	245	+ "TGATHER: TileData element type must match GlobalData element type");
	246	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TGATHER: src/dst layout mismatch");
	247	+
	248	+ const int nranks = parallelGroup.GetSize();
	249	+ const int rootIdx = parallelGroup.GetRootIdx();
	250	+
	251	+ PTO_ASSERT(nranks > 0, "ParallelGroup size must be greater than 0!");
	252	+ PTO_ASSERT(rootIdx >= 0 && rootIdx < nranks, "rootIdx must be in range [0, nranks)!");
	253	+
	254	+ // Get per-rank dimensions
	255	+ const int gShape0 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_0);
	256	+ const int gShape1 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_1);
	257	+ const int gShape2 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_2);
	258	+ const int gShape3 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_3);
	259	+ const int gShape4 = parallelGroup[0].GetShape(GlobalTensorDim::DIM_4);
	260	+
	261	+ const int perRankRows = gShape3;
	262	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	263	+ const int tileValidRow = pingTile.GetValidRow();
	264	+ const int tileValidCol = pingTile.GetValidCol();
	265	+
	266	+ PTO_ASSERT(tileValidRow > 0, "TGATHER: tileValidRow must be greater than 0");
	267	+ PTO_ASSERT(tileValidCol > 0, "TGATHER: tileValidCol must be greater than 0");
	268	+
	269	+ if (totalRows == 0 \|\| gShape4 == 0) {
	270	+ return;
	271	+ }
	272	+
	273	+ // ---- Simple path: per-rank data fits in UB tile, no ping-pong benefit ----
	274	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	275	+ if (nranks == 1) {
	276	+ TLOAD(pingTile, parallelGroup[0]);
	277	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	278	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	279	+ TSTORE(dstGlobalData, pingTile);
	280	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	281	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	282	+ return;
	283	+ }
	284	+
	285	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	286	+
	287	+ using DynShape5D = Shape<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	288	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	289	+ using DstViewT = GlobalTensor<T, DynShape5D, DynStride, GlobalDstData::layout>;
	290	+
	291	+ DynShape5D perRankShape(gShape0, gShape1, gShape2, gShape3, gShape4);
	292	+ DynStride dstViewStride(dstGlobalData.GetStride(GlobalTensorDim::DIM_0),
	293	+ dstGlobalData.GetStride(GlobalTensorDim::DIM_1),
	294	+ dstGlobalData.GetStride(GlobalTensorDim::DIM_2), dstStride3,
	295	+ dstGlobalData.GetStride(GlobalTensorDim::DIM_4));
	296	+
	297	+ for (int r = 0; r < nranks; ++r) {
	298	+ TLOAD(pingTile, parallelGroup[r]);
	299	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	300	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	301	+
	302	+ int64_t dstOffset = static_cast<int64_t>(r) * perRankRows * dstStride3;
	303	+ DstViewT dstView(dstGlobalData.data() + dstOffset, perRankShape, dstViewStride);
	304	+ TSTORE(dstView, pingTile);
	305	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	306	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	307	+ }
	308	+ return;
	309	+ }
	310	+
	311	+ // ---- 2D sliding chunked path with ping-pong double buffering ----
	312	+
	313	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	314	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	315	+
	316	+ if constexpr (!isDynamicRow) {
	317	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	318	+ "TGATHER chunked: per-rank DIM_3 must be divisible by tile ValidRow when static.");
	319	+ }
	320	+ if constexpr (!isDynamicCol) {
	321	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	322	+ "TGATHER chunked: DIM_4 must be divisible by tile ValidCol when static.");
	323	+ }
	324	+
	325	+ const int srcStride0 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_0);
	326	+ const int srcStride1 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_1);
	327	+ const int srcStride2 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_2);
	328	+ const int srcStride3 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_3);
	329	+ const int srcStride4 = parallelGroup[0].GetStride(GlobalTensorDim::DIM_4);
	330	+
	331	+ const int dstStride0 = dstGlobalData.GetStride(GlobalTensorDim::DIM_0);
	332	+ const int dstStride1 = dstGlobalData.GetStride(GlobalTensorDim::DIM_1);
	333	+ const int dstStride2 = dstGlobalData.GetStride(GlobalTensorDim::DIM_2);
	334	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	335	+ const int dstStride4 = dstGlobalData.GetStride(GlobalTensorDim::DIM_4);
	336	+
	337	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	338	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	339	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	340	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	341	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	342	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	343	+
	344	+ // Ping-pong state
	345	+ bool usePing = true;
	346	+ bool hasPending = false;
	347	+ int64_t pendingDstOffset = 0;
	348	+ int pendingRows = 0;
	349	+ int pendingCols = 0;
	350	+
	351	+ for (int r = 0; r < nranks; ++r) {
	352	+ int64_t rankDstBase = static_cast<int64_t>(r) * perRankRows * dstStride3;
	353	+
	354	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	355	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	356	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	357	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	358	+ static_cast<int64_t>(i2) * srcStride2;
	359	+ int64_t dstBase = rankDstBase + static_cast<int64_t>(i0) * dstStride0 +
	360	+ static_cast<int64_t>(i1) * dstStride1 + static_cast<int64_t>(i2) * dstStride2;
	361	+
	362	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	363	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	364	+
	365	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	366	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	367	+
	368	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	369	+ static_cast<int64_t>(colOff) * srcStride4;
	370	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	371	+ static_cast<int64_t>(colOff) * dstStride4;
	372	+
	373	+ // Select load tile
	374	+ TileData &loadTile = usePing ? pingTile : pongTile;
	375	+ event_t curEvent = usePing ? EVENT_ID0 : EVENT_ID1;
	376	+
	377	+ if constexpr (isDynamicRow)
	378	+ loadTile.RowMaskInternal = currentRows;
	379	+ if constexpr (isDynamicCol)
	380	+ loadTile.ColMaskInternal = currentCols;
	381	+
	382	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	383	+ SrcViewT srcView(parallelGroup[r].data() + srcOffset, chunkShape, srcChunkStride);
	384	+
	385	+ if (hasPending) {
	386	+ TileData &storeTile = usePing ? pongTile : pingTile;
	387	+ event_t prevEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	388	+
	389	+ // Wait for previous TLOAD to finish
	390	+ wait_flag(PIPE_MTE2, PIPE_MTE3, prevEvent);
	391	+
	392	+ DynShape pendShape(1, 1, 1, pendingRows, pendingCols);
	393	+ DstViewT dstView(dstGlobalData.data() + pendingDstOffset, pendShape, dstChunkStride);
	394	+
	395	+ // Issue TSTORE + TLOAD concurrently (MTE3 and MTE2 in parallel)
	396	+ TSTORE(dstView, storeTile);
	397	+ TLOAD(loadTile, srcView);
	398	+
	399	+ set_flag(PIPE_MTE3, PIPE_MTE2, prevEvent); // store done
	400	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent); // load done
	401	+
	402	+ // Ensure storeTile UB is safe before overwrite
	403	+ wait_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);
	404	+ } else {
	405	+ // First chunk: just issue TLOAD
	406	+ TLOAD(loadTile, srcView);
	407	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent);
	408	+ }
	409	+
	410	+ pendingDstOffset = dstOffset;
	411	+ pendingRows = currentRows;
	412	+ pendingCols = currentCols;
	413	+ hasPending = true;
	414	+ usePing = !usePing;
	415	+ }
	416	+ }
	417	+ }
	418	+ }
	419	+ }
	420	+ }
	421	+
	422	+ // Epilogue: drain the last pending chunk
	423	+ if (hasPending) {
	424	+ TileData &lastTile = usePing ? pongTile : pingTile;
	425	+ event_t lastEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	426	+
	427	+ wait_flag(PIPE_MTE2, PIPE_MTE3, lastEvent);
	428	+
	429	+ DynShape lastShape(1, 1, 1, pendingRows, pendingCols);
	430	+ DstViewT dstView(dstGlobalData.data() + pendingDstOffset, lastShape, dstChunkStride);
	431	+ TSTORE(dstView, lastTile);
	432	+
	433	+ set_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	434	+ wait_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	435	+ }
	436	+}
	437	+
	438	+} // namespace comm
	439	+} // namespace pto
	440	+
	441	+#endif // PTO_COMM_TGATHER_HPP

Ainclude/pto/comm/TGet.hpp+364-0

@@ -0,0 +1,364 @@
	1	+/**
	2	+Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+CANN Open Software License Agreement Version 2.0 (the "License").
	5	+Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+See LICENSE in the root of the software repository for the full text of the License.
	9	+*/
	10	+
	11	+#ifndef PTO_COMM_TGET_HPP
	12	+#define PTO_COMM_TGET_HPP
	13	+
	14	+#include <type_traits>
	15	+
	16	+#include "pto/common/debug.h"
	17	+#include "pto/common/type.hpp"
	18	+#include "pto/common/constants.hpp"
	19	+#include "pto/common/pto_instr.hpp"
	20	+#include "pto/comm/comm_types.hpp"
	21	+
	22	+namespace pto {
	23	+namespace comm {
	24	+
	25	+// ============================================================================
	26	+// TGET_IMPL: Remote read operation implementation
	27	+//
	28	+// Data flow: srcGlobalData (remote GM) → stagingTileData (UB) → dstGlobalData (local GM)
	29	+//
	30	+// When the GlobalTensor exceeds the UB tile capacity in rows and/or columns,
	31	+// the transfer is automatically chunked via 2D sliding:
	32	+// - Outer dimensions (DIM_0, DIM_1, DIM_2) are iterated explicitly.
	33	+// - DIM_3 (rows) is split into tileValidRow-sized chunks.
	34	+// - DIM_4 (cols) is split into tileValidCol-sized chunks.
	35	+//
	36	+// Constraints for chunked mode:
	37	+// - If TileData has static ValidRow, shape3 must be divisible by ValidRow.
	38	+// Use DYNAMIC ValidRow for partial row chunk support.
	39	+// - If TileData has static ValidCol, shape4 must be divisible by ValidCol.
	40	+// Use DYNAMIC ValidCol for partial column chunk support.
	41	+// ============================================================================
	42	+
	43	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData>
	44	+PTO_INTERNAL void TGET_IMPL(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData, TileData &stagingTileData)
	45	+{
	46	+ using T = typename GlobalSrcData::RawDType;
	47	+
	48	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>, "TGET: src/dst element type mismatch");
	49	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TGET: src/dst layout mismatch");
	50	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	51	+ "TGET: TileData element type must match GlobalData element type");
	52	+
	53	+ // Get GlobalTensor dimensions
	54	+ const int gShape0 = srcGlobalData.GetShape(GlobalTensorDim::DIM_0);
	55	+ const int gShape1 = srcGlobalData.GetShape(GlobalTensorDim::DIM_1);
	56	+ const int gShape2 = srcGlobalData.GetShape(GlobalTensorDim::DIM_2);
	57	+ const int gShape3 = srcGlobalData.GetShape(GlobalTensorDim::DIM_3);
	58	+ const int gShape4 = srcGlobalData.GetShape(GlobalTensorDim::DIM_4);
	59	+
	60	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	61	+ const int tileValidRow = stagingTileData.GetValidRow();
	62	+ const int tileValidCol = stagingTileData.GetValidCol();
	63	+
	64	+ PTO_ASSERT(tileValidRow > 0, "TGET: tileValidRow must be greater than 0");
	65	+ PTO_ASSERT(tileValidCol > 0, "TGET: tileValidCol must be greater than 0");
	66	+
	67	+ if (totalRows == 0 \|\| gShape4 == 0) {
	68	+ return;
	69	+ }
	70	+
	71	+ // ---- Simple path: data fits in UB tile in both dimensions ----
	72	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	73	+ TLOAD(stagingTileData, srcGlobalData);
	74	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	75	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	76	+ TSTORE(dstGlobalData, stagingTileData);
	77	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	78	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	79	+ return;
	80	+ }
	81	+
	82	+ // ---- 2D sliding chunked path ----
	83	+ //
	84	+ // Strategy (ND layout):
	85	+ // - Iterate over outer dimensions (dim0, dim1, dim2) explicitly.
	86	+ // - Within each (i0, i1, i2) block, slide a (tileValidRow × tileValidCol)
	87	+ // window over the (dim3 × dim4) plane.
	88	+ // - For each chunk, create a view: shape = (1, 1, 1, curRows, curCols),
	89	+ // preserving the original strides for correct GM addressing.
	90	+ // - TLOAD the chunk view into UB, then TSTORE from UB to local GM.
	91	+
	92	+ PTO_ASSERT(tileValidRow > 0, "TGET: tile ValidRow must be greater than 0 for chunked transfer");
	93	+ PTO_ASSERT(tileValidCol > 0, "TGET: tile ValidCol must be greater than 0 for chunked transfer");
	94	+
	95	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	96	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	97	+
	98	+ // Row validation: static ValidRow requires shape3 to be exactly divisible
	99	+ if constexpr (!isDynamicRow) {
	100	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	101	+ "TGET chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. "
	102	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	103	+ }
	104	+ // Column validation: static ValidCol requires shape4 to be exactly divisible
	105	+ if constexpr (!isDynamicCol) {
	106	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	107	+ "TGET chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. "
	108	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	109	+ }
	110	+
	111	+ // Get strides for offset calculation
	112	+ const int srcStride0 = srcGlobalData.GetStride(GlobalTensorDim::DIM_0);
	113	+ const int srcStride1 = srcGlobalData.GetStride(GlobalTensorDim::DIM_1);
	114	+ const int srcStride2 = srcGlobalData.GetStride(GlobalTensorDim::DIM_2);
	115	+ const int srcStride3 = srcGlobalData.GetStride(GlobalTensorDim::DIM_3);
	116	+ const int srcStride4 = srcGlobalData.GetStride(GlobalTensorDim::DIM_4);
	117	+ const int dstStride0 = dstGlobalData.GetStride(GlobalTensorDim::DIM_0);
	118	+ const int dstStride1 = dstGlobalData.GetStride(GlobalTensorDim::DIM_1);
	119	+ const int dstStride2 = dstGlobalData.GetStride(GlobalTensorDim::DIM_2);
	120	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	121	+ const int dstStride4 = dstGlobalData.GetStride(GlobalTensorDim::DIM_4);
	122	+
	123	+ // View types with fully dynamic shape/stride for chunk GlobalTensors
	124	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	125	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	126	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	127	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	128	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	129	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	130	+
	131	+ // 2D sliding: iterate outer dims, then chunk rows (dim3) and columns (dim4)
	132	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	133	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	134	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	135	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	136	+ static_cast<int64_t>(i2) * srcStride2;
	137	+ int64_t dstBase = static_cast<int64_t>(i0) * dstStride0 + static_cast<int64_t>(i1) * dstStride1 +
	138	+ static_cast<int64_t>(i2) * dstStride2;
	139	+
	140	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	141	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	142	+
	143	+ if constexpr (isDynamicRow) {
	144	+ stagingTileData.RowMaskInternal = currentRows;
	145	+ }
	146	+
	147	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	148	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	149	+
	150	+ if constexpr (isDynamicCol) {
	151	+ stagingTileData.ColMaskInternal = currentCols;
	152	+ }
	153	+
	154	+ // Compute element offsets
	155	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	156	+ static_cast<int64_t>(colOff) * srcStride4;
	157	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	158	+ static_cast<int64_t>(colOff) * dstStride4;
	159	+
	160	+ // Create chunk views with adjusted shape
	161	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	162	+
	163	+ SrcViewT srcView(srcGlobalData.data() + srcOffset, chunkShape, srcChunkStride);
	164	+ DstViewT dstView(dstGlobalData.data() + dstOffset, chunkShape, dstChunkStride);
	165	+
	166	+ // Transfer: remote GM → UB → local GM
	167	+ TLOAD(stagingTileData, srcView);
	168	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	169	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	170	+ TSTORE(dstView, stagingTileData);
	171	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	172	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	173	+ }
	174	+ }
	175	+ }
	176	+ }
	177	+ }
	178	+}
	179	+
	180	+// ============================================================================
	181	+// TGET_IMPL (ping-pong): Remote read with double buffering
	182	+//
	183	+// Uses two staging tiles (pingTile, pongTile) to overlap TLOAD (MTE2) and
	184	+// TSTORE (MTE3) for adjacent chunks, effectively hiding one DMA transfer
	185	+// behind the other.
	186	+//
	187	+// Timeline without ping-pong:
	188	+// [TLOAD chunk0] -> [TSTORE chunk0] -> [TLOAD chunk1] -> [TSTORE chunk1] -> ...
	189	+//
	190	+// Timeline with ping-pong (overlap TSTORE[i] with TLOAD[i+1]):
	191	+// [TLOAD chunk0] -> [TSTORE chunk0 \| TLOAD chunk1] -> [TSTORE chunk1 \| TLOAD chunk2] -> ...
	192	+//
	193	+// Requirements:
	194	+// - pingTile and pongTile must have the same type and dimensions.
	195	+// - Uses EVENT_ID0 (pingTile) and EVENT_ID1 (pongTile) for synchronization.
	196	+// ============================================================================
	197	+
	198	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData>
	199	+PTO_INTERNAL void TGET_IMPL(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData, TileData &pingTile,
	200	+ TileData &pongTile)
	201	+{
	202	+ using T = typename GlobalSrcData::RawDType;
	203	+
	204	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>, "TGET: src/dst element type mismatch");
	205	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TGET: src/dst layout mismatch");
	206	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	207	+ "TGET: TileData element type must match GlobalData element type");
	208	+
	209	+ const int gShape0 = srcGlobalData.GetShape(GlobalTensorDim::DIM_0);
	210	+ const int gShape1 = srcGlobalData.GetShape(GlobalTensorDim::DIM_1);
	211	+ const int gShape2 = srcGlobalData.GetShape(GlobalTensorDim::DIM_2);
	212	+ const int gShape3 = srcGlobalData.GetShape(GlobalTensorDim::DIM_3);
	213	+ const int gShape4 = srcGlobalData.GetShape(GlobalTensorDim::DIM_4);
	214	+
	215	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	216	+ const int tileValidRow = pingTile.GetValidRow();
	217	+ const int tileValidCol = pingTile.GetValidCol();
	218	+
	219	+ PTO_ASSERT(tileValidRow > 0, "TGET: tileValidRow must be greater than 0");
	220	+ PTO_ASSERT(tileValidCol > 0, "TGET: tileValidCol must be greater than 0");
	221	+
	222	+ if (totalRows == 0 \|\| gShape4 == 0) {
	223	+ return;
	224	+ }
	225	+
	226	+ // ---- Simple path: single chunk, no ping-pong benefit ----
	227	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	228	+ TLOAD(pingTile, srcGlobalData);
	229	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	230	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	231	+ TSTORE(dstGlobalData, pingTile);
	232	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	233	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	234	+ return;
	235	+ }
	236	+
	237	+ // ---- 2D sliding chunked path with ping-pong double buffering ----
	238	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	239	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	240	+
	241	+ if constexpr (!isDynamicRow) {
	242	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	243	+ "TGET chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. "
	244	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	245	+ }
	246	+ if constexpr (!isDynamicCol) {
	247	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	248	+ "TGET chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. "
	249	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	250	+ }
	251	+
	252	+ const int srcStride0 = srcGlobalData.GetStride(GlobalTensorDim::DIM_0);
	253	+ const int srcStride1 = srcGlobalData.GetStride(GlobalTensorDim::DIM_1);
	254	+ const int srcStride2 = srcGlobalData.GetStride(GlobalTensorDim::DIM_2);
	255	+ const int srcStride3 = srcGlobalData.GetStride(GlobalTensorDim::DIM_3);
	256	+ const int srcStride4 = srcGlobalData.GetStride(GlobalTensorDim::DIM_4);
	257	+ const int dstStride0 = dstGlobalData.GetStride(GlobalTensorDim::DIM_0);
	258	+ const int dstStride1 = dstGlobalData.GetStride(GlobalTensorDim::DIM_1);
	259	+ const int dstStride2 = dstGlobalData.GetStride(GlobalTensorDim::DIM_2);
	260	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	261	+ const int dstStride4 = dstGlobalData.GetStride(GlobalTensorDim::DIM_4);
	262	+
	263	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	264	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	265	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	266	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	267	+
	268	+ // Precompute strides (identical for all chunk views)
	269	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	270	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	271	+
	272	+ // Ping-pong state (same as TPUT_IMPL ping-pong)
	273	+ // See TPUT_IMPL comments for detailed pipeline analysis.
	274	+ bool usePing = true;
	275	+ bool hasPending = false;
	276	+ int64_t pendingDstOffset = 0;
	277	+ int pendingRows = 0;
	278	+ int pendingCols = 0;
	279	+
	280	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	281	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	282	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	283	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	284	+ static_cast<int64_t>(i2) * srcStride2;
	285	+ int64_t dstBase = static_cast<int64_t>(i0) * dstStride0 + static_cast<int64_t>(i1) * dstStride1 +
	286	+ static_cast<int64_t>(i2) * dstStride2;
	287	+
	288	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	289	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	290	+
	291	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	292	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	293	+
	294	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	295	+ static_cast<int64_t>(colOff) * srcStride4;
	296	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	297	+ static_cast<int64_t>(colOff) * dstStride4;
	298	+
	299	+ TileData &loadTile = usePing ? pingTile : pongTile;
	300	+ event_t curEvent = usePing ? EVENT_ID0 : EVENT_ID1;
	301	+
	302	+ if constexpr (isDynamicRow)
	303	+ loadTile.RowMaskInternal = currentRows;
	304	+ if constexpr (isDynamicCol)
	305	+ loadTile.ColMaskInternal = currentCols;
	306	+
	307	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	308	+ SrcViewT srcView(srcGlobalData.data() + srcOffset, chunkShape, srcChunkStride);
	309	+
	310	+ if (hasPending) {
	311	+ TileData &storeTile = usePing ? pongTile : pingTile;
	312	+ event_t prevEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	313	+
	314	+ // Wait for previous TLOAD to finish (data in storeTile is ready)
	315	+ wait_flag(PIPE_MTE2, PIPE_MTE3, prevEvent);
	316	+
	317	+ DynShape pendShape(1, 1, 1, pendingRows, pendingCols);
	318	+ DstViewT pendView(dstGlobalData.data() + pendingDstOffset, pendShape, dstChunkStride);
	319	+
	320	+ // Issue TSTORE + TLOAD concurrently (MTE3 and MTE2 in parallel)
	321	+ TSTORE(pendView, storeTile);
	322	+ TLOAD(loadTile, srcView);
	323	+
	324	+ set_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);
	325	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent);
	326	+
	327	+ // Ensure storeTile's UB has been fully read by MTE3
	328	+ wait_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);
	329	+ } else {
	330	+ TLOAD(loadTile, srcView);
	331	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent);
	332	+ }
	333	+
	334	+ pendingDstOffset = dstOffset;
	335	+ pendingRows = currentRows;
	336	+ pendingCols = currentCols;
	337	+ hasPending = true;
	338	+ usePing = !usePing;
	339	+ }
	340	+ }
	341	+ }
	342	+ }
	343	+ }
	344	+
	345	+ // Epilogue: drain the last pending TSTORE
	346	+ if (hasPending) {
	347	+ TileData &lastTile = usePing ? pongTile : pingTile;
	348	+ event_t lastEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	349	+
	350	+ wait_flag(PIPE_MTE2, PIPE_MTE3, lastEvent);
	351	+
	352	+ DynShape lastShape(1, 1, 1, pendingRows, pendingCols);
	353	+ DstViewT lastView(dstGlobalData.data() + pendingDstOffset, lastShape, dstChunkStride);
	354	+
	355	+ TSTORE(lastView, lastTile);
	356	+ set_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	357	+ wait_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	358	+ }
	359	+}
	360	+
	361	+} // namespace comm
	362	+} // namespace pto
	363	+
	364	+#endif // PTO_COMM_TGET_HPP

Ainclude/pto/comm/TNotify.hpp+66-0

@@ -0,0 +1,66 @@
	1	+/**
	2	+Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+CANN Open Software License Agreement Version 2.0 (the "License").
	5	+Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+See LICENSE in the root of the software repository for the full text of the License.
	9	+*/
	10	+
	11	+#ifndef PTO_COMM_TNOTIFY_HPP
	12	+#define PTO_COMM_TNOTIFY_HPP
	13	+
	14	+#include "pto/common/type.hpp"
	15	+#include "pto/common/utils.hpp"
	16	+#include "pto/comm/comm_types.hpp"
	17	+
	18	+namespace pto {
	19	+namespace comm {
	20	+
	21	+namespace detail {
	22	+PTO_INTERNAL void DcciSignal(__gm__ int32_t *ptr)
	23	+{
	24	+ __asm__ __volatile__("");
	25	+ dcci(ptr, SINGLE_CACHE_LINE);
	26	+ __asm__ __volatile__("");
	27	+}
	28	+} // namespace detail
	29	+
	30	+// ============================================================================
	31	+// TNOTIFY_IMPL: Send flag notification to remote NPU
	32	+//
	33	+// Signal type must be int32_t.
	34	+// dstSignalData should be 4-byte aligned.
	35	+// ============================================================================
	36	+
	37	+template <typename GlobalSignalData>
	38	+PTO_INTERNAL void TNOTIFY_IMPL(GlobalSignalData &dstSignalData, int32_t value, NotifyOp op)
	39	+{
	40	+ static_assert(std::is_same_v<typename GlobalSignalData::RawDType, int32_t>, "TNOTIFY: signal type must be int32_t");
	41	+
	42	+ volatile __gm__ int32_t sigPtr = (volatile __gm__ int32_t )dstSignalData.data();
	43	+
	44	+ if (op == NotifyOp::AtomicAdd) {
	45	+ // Atomic add using hardware atomic instruction
	46	+ set_st_atomic_cfg(ATOMIC_S32, ATOMIC_SUM);
	47	+ detail::DcciSignal((__gm__ int32_t *)sigPtr);
	48	+ st_atomic<int32_t>(value, (__gm__ int32_t *)sigPtr);
	49	+ detail::DcciSignal((__gm__ int32_t *)sigPtr);
	50	+ dsb(DSB_DDR);
	51	+ } else {
	52	+ // Set operation - direct store to remote memory
	53	+ // Invalidate cache first to prevent stale cached data from overwriting the new value
	54	+ detail::DcciSignal((__gm__ int32_t *)sigPtr);
	55	+ *sigPtr = value;
	56	+ detail::DcciSignal((__gm__ int32_t *)sigPtr);
	57	+ dsb(DSB_DDR);
	58	+ }
	59	+
	60	+ pipe_barrier(PIPE_ALL);
	61	+}
	62	+
	63	+} // namespace comm
	64	+} // namespace pto
	65	+
	66	+#endif // PTO_COMM_TNOTIFY_HPP

Ainclude/pto/comm/TPut.hpp+393-0

@@ -0,0 +1,393 @@
	1	+/**
	2	+Copyright (c) 2025 Huawei Technologies Co., Ltd.
	3	+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+CANN Open Software License Agreement Version 2.0 (the "License").
	5	+Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+See LICENSE in the root of the software repository for the full text of the License.
	9	+*/
	10	+
	11	+#ifndef PTO_COMM_TPUT_HPP
	12	+#define PTO_COMM_TPUT_HPP
	13	+
	14	+#include <type_traits>
	15	+
	16	+#include "pto/common/debug.h"
	17	+#include "pto/common/type.hpp"
	18	+#include "pto/common/constants.hpp"
	19	+#include "pto/common/pto_instr.hpp"
	20	+#include "pto/comm/comm_types.hpp"
	21	+
	22	+namespace pto {
	23	+namespace comm {
	24	+
	25	+// ============================================================================
	26	+// TPUT_IMPL: Remote write operation implementation
	27	+//
	28	+// Data flow: srcGlobalData (local GM) → stagingTileData (UB) → dstGlobalData (remote GM)
	29	+// - atomicType: Atomic operation type (AtomicNone or AtomicAdd)
	30	+//
	31	+// When the GlobalTensor exceeds the UB tile capacity in rows and/or columns,
	32	+// the transfer is automatically chunked via 2D sliding:
	33	+// - Outer dimensions (DIM_0, DIM_1, DIM_2) are iterated explicitly.
	34	+// - DIM_3 (rows) is split into tileValidRow-sized chunks.
	35	+// - DIM_4 (cols) is split into tileValidCol-sized chunks.
	36	+//
	37	+// Constraints for chunked mode:
	38	+// - If TileData has static ValidRow, shape3 must be divisible by ValidRow.
	39	+// Use DYNAMIC ValidRow for partial row chunk support.
	40	+// - If TileData has static ValidCol, shape4 must be divisible by ValidCol.
	41	+// Use DYNAMIC ValidCol for partial column chunk support.
	42	+// ============================================================================
	43	+
	44	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData,
	45	+ AtomicType atomicType = AtomicType::AtomicNone>
	46	+PTO_INTERNAL void TPUT_IMPL(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData, TileData &stagingTileData)
	47	+{
	48	+ using T = typename GlobalSrcData::RawDType;
	49	+
	50	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>, "TPUT: src/dst element type mismatch");
	51	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TPUT: src/dst layout mismatch");
	52	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	53	+ "TPUT: TileData element type must match GlobalData element type");
	54	+
	55	+ // Get GlobalTensor dimensions
	56	+ const int gShape0 = srcGlobalData.GetShape(GlobalTensorDim::DIM_0);
	57	+ const int gShape1 = srcGlobalData.GetShape(GlobalTensorDim::DIM_1);
	58	+ const int gShape2 = srcGlobalData.GetShape(GlobalTensorDim::DIM_2);
	59	+ const int gShape3 = srcGlobalData.GetShape(GlobalTensorDim::DIM_3);
	60	+ const int gShape4 = srcGlobalData.GetShape(GlobalTensorDim::DIM_4);
	61	+
	62	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	63	+ const int tileValidRow = stagingTileData.GetValidRow();
	64	+ const int tileValidCol = stagingTileData.GetValidCol();
	65	+
	66	+ PTO_ASSERT(tileValidRow > 0, "TPUT: tileValidRow must be greater than 0");
	67	+ PTO_ASSERT(tileValidCol > 0, "TPUT: tileValidCol must be greater than 0");
	68	+
	69	+ if (totalRows == 0 \|\| gShape4 == 0) {
	70	+ return;
	71	+ }
	72	+
	73	+ // ---- Simple path: data fits in UB tile in both dimensions ----
	74	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	75	+ TLOAD(stagingTileData, srcGlobalData);
	76	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	77	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	78	+ TSTORE<TileData, GlobalDstData, atomicType>(dstGlobalData, stagingTileData);
	79	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	80	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	81	+ return;
	82	+ }
	83	+
	84	+ // ---- 2D sliding chunked path ----
	85	+ //
	86	+ // Strategy (ND layout):
	87	+ // - Iterate over outer dimensions (dim0, dim1, dim2) explicitly.
	88	+ // - Within each (i0, i1, i2) block, slide a (tileValidRow × tileValidCol)
	89	+ // window over the (dim3 × dim4) plane.
	90	+ // - For each chunk, create a view: shape = (1, 1, 1, curRows, curCols),
	91	+ // preserving the original strides for correct GM addressing.
	92	+ // - TLOAD the chunk view into UB, then TSTORE from UB to remote GM.
	93	+
	94	+ PTO_ASSERT(tileValidRow > 0, "TPUT: tile ValidRow must be greater than 0 for chunked transfer");
	95	+ PTO_ASSERT(tileValidCol > 0, "TPUT: tile ValidCol must be greater than 0 for chunked transfer");
	96	+
	97	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	98	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	99	+
	100	+ // Row validation: static ValidRow requires shape3 to be exactly divisible
	101	+ if constexpr (!isDynamicRow) {
	102	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	103	+ "TPUT chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. "
	104	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	105	+ }
	106	+ // Column validation: static ValidCol requires shape4 to be exactly divisible
	107	+ if constexpr (!isDynamicCol) {
	108	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	109	+ "TPUT chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. "
	110	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	111	+ }
	112	+
	113	+ // Get strides for offset calculation
	114	+ const int srcStride0 = srcGlobalData.GetStride(GlobalTensorDim::DIM_0);
	115	+ const int srcStride1 = srcGlobalData.GetStride(GlobalTensorDim::DIM_1);
	116	+ const int srcStride2 = srcGlobalData.GetStride(GlobalTensorDim::DIM_2);
	117	+ const int srcStride3 = srcGlobalData.GetStride(GlobalTensorDim::DIM_3);
	118	+ const int srcStride4 = srcGlobalData.GetStride(GlobalTensorDim::DIM_4);
	119	+
	120	+ const int dstStride0 = dstGlobalData.GetStride(GlobalTensorDim::DIM_0);
	121	+ const int dstStride1 = dstGlobalData.GetStride(GlobalTensorDim::DIM_1);
	122	+ const int dstStride2 = dstGlobalData.GetStride(GlobalTensorDim::DIM_2);
	123	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	124	+ const int dstStride4 = dstGlobalData.GetStride(GlobalTensorDim::DIM_4);
	125	+
	126	+ // View types with fully dynamic shape/stride for chunk GlobalTensors
	127	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	128	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	129	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	130	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	131	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	132	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	133	+
	134	+ // 2D sliding: iterate outer dims, then chunk rows (dim3) and columns (dim4)
	135	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	136	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	137	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	138	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	139	+ static_cast<int64_t>(i2) * srcStride2;
	140	+ int64_t dstBase = static_cast<int64_t>(i0) * dstStride0 + static_cast<int64_t>(i1) * dstStride1 +
	141	+ static_cast<int64_t>(i2) * dstStride2;
	142	+
	143	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	144	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	145	+
	146	+ if constexpr (isDynamicRow) {
	147	+ stagingTileData.RowMaskInternal = currentRows;
	148	+ }
	149	+
	150	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	151	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	152	+
	153	+ if constexpr (isDynamicCol) {
	154	+ stagingTileData.ColMaskInternal = currentCols;
	155	+ }
	156	+
	157	+ // Compute element offsets
	158	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	159	+ static_cast<int64_t>(colOff) * srcStride4;
	160	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	161	+ static_cast<int64_t>(colOff) * dstStride4;
	162	+
	163	+ // Create chunk views with adjusted shape
	164	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	165	+
	166	+ SrcViewT srcView(srcGlobalData.data() + srcOffset, chunkShape, srcChunkStride);
	167	+ DstViewT dstView(dstGlobalData.data() + dstOffset, chunkShape, dstChunkStride);
	168	+
	169	+ // Transfer: local GM → UB → remote GM
	170	+ TLOAD(stagingTileData, srcView);
	171	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	172	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	173	+ TSTORE<TileData, DstViewT, atomicType>(dstView, stagingTileData);
	174	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	175	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	176	+ }
	177	+ }
	178	+ }
	179	+ }
	180	+ }
	181	+}
	182	+
	183	+// ============================================================================
	184	+// TPUT_IMPL (ping-pong): Remote write with double buffering
	185	+//
	186	+// Uses two staging tiles (pingTile, pongTile) to overlap TLOAD (MTE2) and
	187	+// TSTORE (MTE3) for adjacent chunks, effectively hiding one DMA transfer
	188	+// behind the other.
	189	+//
	190	+// Timeline without ping-pong:
	191	+// [TLOAD chunk0] -> [TSTORE chunk0] -> [TLOAD chunk1] -> [TSTORE chunk1] -> ...
	192	+//
	193	+// Timeline with ping-pong (overlap TSTORE[i] with TLOAD[i+1]):
	194	+// [TLOAD chunk0] -> [TSTORE chunk0 \| TLOAD chunk1] -> [TSTORE chunk1 \| TLOAD chunk2] -> ...
	195	+//
	196	+// Requirements:
	197	+// - pingTile and pongTile must have the same type and dimensions.
	198	+// - Uses EVENT_ID0 (pingTile) and EVENT_ID1 (pongTile) for synchronization.
	199	+// ============================================================================
	200	+
	201	+template <typename GlobalDstData, typename GlobalSrcData, typename TileData,
	202	+ AtomicType atomicType = AtomicType::AtomicNone>
	203	+PTO_INTERNAL void TPUT_IMPL(GlobalDstData &dstGlobalData, GlobalSrcData &srcGlobalData, TileData &pingTile,
	204	+ TileData &pongTile)
	205	+{
	206	+ using T = typename GlobalSrcData::RawDType;
	207	+
	208	+ static_assert(std::is_same_v<T, typename GlobalDstData::RawDType>, "TPUT: src/dst element type mismatch");
	209	+ static_assert(GlobalSrcData::layout == GlobalDstData::layout, "TPUT: src/dst layout mismatch");
	210	+ static_assert(std::is_same_v<T, typename TileData::DType>,
	211	+ "TPUT: TileData element type must match GlobalData element type");
	212	+
	213	+ const int gShape0 = srcGlobalData.GetShape(GlobalTensorDim::DIM_0);
	214	+ const int gShape1 = srcGlobalData.GetShape(GlobalTensorDim::DIM_1);
	215	+ const int gShape2 = srcGlobalData.GetShape(GlobalTensorDim::DIM_2);
	216	+ const int gShape3 = srcGlobalData.GetShape(GlobalTensorDim::DIM_3);
	217	+ const int gShape4 = srcGlobalData.GetShape(GlobalTensorDim::DIM_4);
	218	+
	219	+ const int64_t totalRows = static_cast<int64_t>(gShape0) * gShape1 * gShape2 * gShape3;
	220	+ const int tileValidRow = pingTile.GetValidRow();
	221	+ const int tileValidCol = pingTile.GetValidCol();
	222	+
	223	+ PTO_ASSERT(tileValidRow > 0, "TPUT: tileValidRow must be greater than 0");
	224	+ PTO_ASSERT(tileValidCol > 0, "TPUT: tileValidCol must be greater than 0");
	225	+
	226	+ if (totalRows == 0 \|\| gShape4 == 0) {
	227	+ return;
	228	+ }
	229	+
	230	+ // ---- Simple path: single chunk, no ping-pong benefit ----
	231	+ if (totalRows <= tileValidRow && gShape4 <= tileValidCol) {
	232	+ TLOAD(pingTile, srcGlobalData);
	233	+ set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	234	+ wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
	235	+ TSTORE<TileData, GlobalDstData, atomicType>(dstGlobalData, pingTile);
	236	+ set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	237	+ wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
	238	+ return;
	239	+ }
	240	+
	241	+ // ---- 2D sliding chunked path with ping-pong double buffering ----
	242	+ constexpr bool isDynamicRow = (TileData::ValidRow == DYNAMIC);
	243	+ constexpr bool isDynamicCol = (TileData::ValidCol == DYNAMIC);
	244	+
	245	+ if constexpr (!isDynamicRow) {
	246	+ PTO_ASSERT(gShape3 % tileValidRow == 0,
	247	+ "TPUT chunked: shape3 must be divisible by tile ValidRow when ValidRow is static. "
	248	+ "Use a Tile with DYNAMIC ValidRow for partial row chunk support.");
	249	+ }
	250	+ if constexpr (!isDynamicCol) {
	251	+ PTO_ASSERT(gShape4 % tileValidCol == 0,
	252	+ "TPUT chunked: shape4 must be divisible by tile ValidCol when ValidCol is static. "
	253	+ "Use a Tile with DYNAMIC ValidCol for partial column chunk support.");
	254	+ }
	255	+
	256	+ const int srcStride0 = srcGlobalData.GetStride(GlobalTensorDim::DIM_0);
	257	+ const int srcStride1 = srcGlobalData.GetStride(GlobalTensorDim::DIM_1);
	258	+ const int srcStride2 = srcGlobalData.GetStride(GlobalTensorDim::DIM_2);
	259	+ const int srcStride3 = srcGlobalData.GetStride(GlobalTensorDim::DIM_3);
	260	+ const int srcStride4 = srcGlobalData.GetStride(GlobalTensorDim::DIM_4);
	261	+
	262	+ const int dstStride0 = dstGlobalData.GetStride(GlobalTensorDim::DIM_0);
	263	+ const int dstStride1 = dstGlobalData.GetStride(GlobalTensorDim::DIM_1);
	264	+ const int dstStride2 = dstGlobalData.GetStride(GlobalTensorDim::DIM_2);
	265	+ const int dstStride3 = dstGlobalData.GetStride(GlobalTensorDim::DIM_3);
	266	+ const int dstStride4 = dstGlobalData.GetStride(GlobalTensorDim::DIM_4);
	267	+
	268	+ using DynShape = Shape<1, 1, 1, DYNAMIC, DYNAMIC>;
	269	+ using DynStride = Stride<DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC, DYNAMIC>;
	270	+ using SrcViewT = GlobalTensor<T, DynShape, DynStride, GlobalSrcData::layout>;
	271	+ using DstViewT = GlobalTensor<T, DynShape, DynStride, GlobalDstData::layout>;
	272	+
	273	+ // Precompute strides (identical for all chunk views)
	274	+ DynStride srcChunkStride(srcStride0, srcStride1, srcStride2, srcStride3, srcStride4);
	275	+ DynStride dstChunkStride(dstStride0, dstStride1, dstStride2, dstStride3, dstStride4);
	276	+
	277	+ // Ping-pong state: tracks the deferred TSTORE from the previous iteration
	278	+ // usePing: true → next TLOAD goes into pingTile (EVENT_ID0)
	279	+ // false → next TLOAD goes into pongTile (EVENT_ID1)
	280	+ // hasPending: whether there is a deferred TSTORE waiting to be issued
	281	+ //
	282	+ // Pipeline overlap: TSTORE and TLOAD are dispatched to separate HW engines
	283	+ // (MTE3 and MTE2). Within each iteration, they run concurrently. The
	284	+ // wait_flag at the end ensures storeTile's UB is safe to reuse before
	285	+ // the NEXT iteration's TLOAD can overwrite it.
	286	+ //
	287	+ // MTE2 queue: [..., TLOAD(loadTile), set_flag(curEvent), wait_flag(prevEvent), ...]
	288	+ // MTE3 queue: [..., wait_flag(prevEvent), TSTORE(storeTile), set_flag(prevEvent), ...]
	289	+ //
	290	+ // The TLOAD and TSTORE in the same iteration execute in parallel on their
	291	+ // respective engines. The wait_flag(MTE3→MTE2) is AFTER the TLOAD in MTE2's
	292	+ // queue, so it doesn't block the current TLOAD — only the NEXT one.
	293	+ bool usePing = true;
	294	+ bool hasPending = false;
	295	+ int64_t pendingDstOffset = 0;
	296	+ int pendingRows = 0;
	297	+ int pendingCols = 0;
	298	+
	299	+ for (int i0 = 0; i0 < gShape0; ++i0) {
	300	+ for (int i1 = 0; i1 < gShape1; ++i1) {
	301	+ for (int i2 = 0; i2 < gShape2; ++i2) {
	302	+ int64_t srcBase = static_cast<int64_t>(i0) * srcStride0 + static_cast<int64_t>(i1) * srcStride1 +
	303	+ static_cast<int64_t>(i2) * srcStride2;
	304	+ int64_t dstBase = static_cast<int64_t>(i0) * dstStride0 + static_cast<int64_t>(i1) * dstStride1 +
	305	+ static_cast<int64_t>(i2) * dstStride2;
	306	+
	307	+ for (int rowOff = 0; rowOff < gShape3; rowOff += tileValidRow) {
	308	+ int currentRows = (rowOff + tileValidRow <= gShape3) ? tileValidRow : (gShape3 - rowOff);
	309	+
	310	+ for (int colOff = 0; colOff < gShape4; colOff += tileValidCol) {
	311	+ int currentCols = (colOff + tileValidCol <= gShape4) ? tileValidCol : (gShape4 - colOff);
	312	+
	313	+ int64_t srcOffset = srcBase + static_cast<int64_t>(rowOff) * srcStride3 +
	314	+ static_cast<int64_t>(colOff) * srcStride4;
	315	+ int64_t dstOffset = dstBase + static_cast<int64_t>(rowOff) * dstStride3 +
	316	+ static_cast<int64_t>(colOff) * dstStride4;
	317	+
	318	+ // Select the tile for this iteration's TLOAD
	319	+ TileData &loadTile = usePing ? pingTile : pongTile;
	320	+ event_t curEvent = usePing ? EVENT_ID0 : EVENT_ID1;
	321	+
	322	+ // Configure masks on the load tile
	323	+ if constexpr (isDynamicRow)
	324	+ loadTile.RowMaskInternal = currentRows;
	325	+ if constexpr (isDynamicCol)
	326	+ loadTile.ColMaskInternal = currentCols;
	327	+
	328	+ DynShape chunkShape(1, 1, 1, currentRows, currentCols);
	329	+ SrcViewT srcView(srcGlobalData.data() + srcOffset, chunkShape, srcChunkStride);
	330	+
	331	+ if (hasPending) {
	332	+ // The other tile holds data from the previous TLOAD
	333	+ TileData &storeTile = usePing ? pongTile : pingTile;
	334	+ event_t prevEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	335	+
	336	+ // Wait for previous TLOAD to finish (data in storeTile is ready)
	337	+ wait_flag(PIPE_MTE2, PIPE_MTE3, prevEvent);
	338	+
	339	+ // Build view for the deferred TSTORE
	340	+ DynShape pendShape(1, 1, 1, pendingRows, pendingCols);
	341	+ DstViewT pendView(dstGlobalData.data() + pendingDstOffset, pendShape, dstChunkStride);
	342	+
	343	+ // Issue TSTORE + TLOAD concurrently (MTE3 and MTE2 in parallel)
	344	+ TSTORE<TileData, DstViewT, atomicType>(pendView, storeTile);
	345	+ TLOAD(loadTile, srcView);
	346	+
	347	+ set_flag(PIPE_MTE3, PIPE_MTE2, prevEvent); // storeTile TSTORE done
	348	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent); // loadTile TLOAD done
	349	+
	350	+ // Ensure storeTile's UB has been fully read by MTE3 before
	351	+ // it can be overwritten by a future TLOAD.
	352	+ // Note: this wait is AFTER the TLOAD in MTE2's queue, so the
	353	+ // current TLOAD already runs in parallel with TSTORE.
	354	+ wait_flag(PIPE_MTE3, PIPE_MTE2, prevEvent);
	355	+ } else {
	356	+ // First chunk: just issue TLOAD (no pending TSTORE yet)
	357	+ TLOAD(loadTile, srcView);
	358	+ set_flag(PIPE_MTE2, PIPE_MTE3, curEvent);
	359	+ }
	360	+
	361	+ // Record this chunk as pending for the next iteration
	362	+ pendingDstOffset = dstOffset;
	363	+ pendingRows = currentRows;
	364	+ pendingCols = currentCols;
	365	+ hasPending = true;
	366	+ usePing = !usePing;
	367	+ }
	368	+ }
	369	+ }
	370	+ }
	371	+ }
	372	+
	373	+ // Epilogue: drain the last pending TSTORE
	374	+ if (hasPending) {
	375	+ // After the last flip, the tile holding the final data is the opposite of usePing
	376	+ TileData &lastTile = usePing ? pongTile : pingTile;
	377	+ event_t lastEvent = usePing ? EVENT_ID1 : EVENT_ID0;
	378	+
	379	+ wait_flag(PIPE_MTE2, PIPE_MTE3, lastEvent);
	380	+
	381	+ DynShape lastShape(1, 1, 1, pendingRows, pendingCols);
	382	+ DstViewT lastView(dstGlobalData.data() + pendingDstOffset, lastShape, dstChunkStride);
	383	+
	384	+ TSTORE<TileData, DstViewT, atomicType>(lastView, lastTile);
	385	+ set_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	386	+ wait_flag(PIPE_MTE3, PIPE_MTE2, lastEvent);
	387	+ }
	388	+}
	389	+
	390	+} // namespace comm
	391	+} // namespace pto
	392	+
	393	+#endif // PTO_COMM_TPUT_HPP