[Feature Request] LowerTileOp for NoC Intrinsic

[xiaoyao-NKU]

Required prerequisites

• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

Motivation

We need to downgrade the NoC primitives in the frontend to lower-level IR in order to better interface with the Broadcast primitives provided by the hardware.

Solution

We need to implement the NoC primitives in the form of TileOp, and write lowering functions. In the LowerTileOp pass, we will downgrade them into lower-level broadcast operations.

Alternatives

No response

Additional context

No response

[xiaoyao-NKU]

PR: Add inter-core communication primitives for TileLang

Summary

• What: Add inter-core communication primitives and their TileLang bindings and lowering.
• Why: Enable expressive, validated mesh communication patterns (broadcast, put, all-gather, barriers, fences, core-id helpers) for SUNMMIO-style device meshes.

Key Changes

• Python API: tilelang/language/comm.py — high-level wrappers:

  • Added get_target_mesh_shape, core_tuple_to_id, core_id_to_tuple, and CoreId.
  • Implemented current_core(), broadcast(), put(), all_gather(), barrier(), fence().
  • Input validation for dtypes, shapes, sizes; mapping of direction strings via DIRECTION_MAP; conversion to buffer regions and emission of TIR intrinsics via tir.call_intrin.

• C++ operators & lowering: src/op/comm.h, src/op/comm.cc:

  • Registered new TIR ops: tl.comm_barrier, tl.comm_fence, tl.CoreId, tl.comm_current_core, tl.comm_is_current_core, tl.tileop.comm_broadcast, tl.tileop.comm_put, tl.tileop.comm_allgather.
  • Implemented BroadcastNode, PutNode, AllgatherNode with:
    • Robust validation (mesh dims from target attributes device_mesh_ncol_* / device_mesh_nrow_*).
    • Lowering that emits Call/Evaluate sequences for 1D/2D broadcasts, masked broadcasts, multi-step puts (via intermediate core), and expanded all-gather sequences.
    • InferLayout() implementations (reuse Copy layout logic; compute linear layout for allgather recv).

• Tests: exercised by testing/python/language/test_tilelang_language_comm.py (unit tests compare generated TIR/TIR script output).

Behavior & Constraints

• Target support: SUNMMIO targets only (code asserts TargetIsSunmmio).
• Direction encoding: 0 = horizontal, 1 = vertical, 2 = all.
• API expectations: all_gather requires the recv buffer to have a leading dimension matching participant count (row/col/all).
• Mesh detection: relies on target attributes like device_mesh_ncol_4, device_mesh_nrow_4. Non-conforming targets must provide equivalent attributes.

Motivation

• Provide a consistent, tested surface for writing inter-core communication in TileLang programs and a matching lowering path to generate correct TIR calls for the runtime.

Testing

• Unit tests under testing/python/language/test_tilelang_language_comm.py validate:
  • Python API script generation.
  • Lowering behavior for broadcast, put, and all-gather with a 4×4 mesh target.

Run locally with the project Python environment:

python -m pytest testing/python/language/test_tilelang_language_comm.py

[xiaoyao-NKU]

tl.comm_allreduce — Allreduce 说明

概述

tl.comm_allreduce 提供跨 mesh 的 all-reduce 功能，支持按行（horizontal）、按列（vertical）或全局（all）方向的规约操作。该功能在 TileLang 的 Python API 与 C++ Lowering 层均有实现，适用于 SUNMMIO 目标。

主要位置

• C++ 实现与注册：src/op/comm.cc, src/op/comm.h
• Python 高层封装：tilelang/language/comm.py
• 测试与脚本：testing/python/language/test_tilelang_language_comm.py

Python API（高层）

函数签名：

all_reduce(buffer, out, reduce_type, direction, dim=-1, clear=True)

要点：

• buffer：待规约缓冲（T.Buffer）。
• out：规约结果缓冲（可为去除规约维后的形状，或在规约维上为 1）。
• reduce_type：规约类型，见下方支持列表。
• direction："horizontal" / "vertical" / "all"（或缩写）。
• dim：规约维，默认为 -1（最后一维）。
• clear：是否在规约前清空 out，clear=False 时会分配临时 dst_copy 保存中间结果。

函数会分配中间缓冲 row_allgather 与 col_allgather（通过 T.alloc_fragment）并将其传递至下层 intrinsic。

C++ 侧（Lower 与布局推断）

• 在 AllreduceOpNode 中定义了字段：src, dst, row_allgather, col_allgather, type, direction, dim, clear, dst_copy 等。
• InferLayout / ComputeLayout：提供对 local.fragment scope 的布局推断逻辑（部分逻辑标注为“待完善”）。
• Lower：仅在 SUNMMIO 目标有效，会从 target mattr 中读取 mesh 配置（mesh_x / mesh_y），对参数与缓冲尺寸做严格检查，然后生成由 allgather/broadcast/reduce/copy 组合而成的规约实现。
• Op 在注册时为 tl.comm_allreduce，并设为 opaque call effect（.set_num_inputs(-1)）。

支持的规约类型

Python 侧定义的 REDUCE_TYPE_LIST 包括：

• sum, abssum, max, min, absmax, bitand, bitor, bitxor。

参数与约束

• dim 必须在合法范围内（或 -1 表示最后一维）；
• out 的形状必须是去掉被规约维后的形状，或在规约维上为 1；
• direction 必须为支持的方向；
• reduce_type 必须在支持列表中；
• Lower 阶段强依赖 SUNMMIO target 的 mattr（如 device_mesh_nrow_* / device_mesh_ncol_*）。

中间缓冲与 clear 行为

• 默认 clear=True：直接在 out 上写入规约结果；
• clear=False：会分配 dst_copy 作增量规约的临时拷贝，保证 out 的初始内容在规约中被正确使用/保留。