Skip to content

[Fix] Harden row-expand layouts and clean up legacy tooling 🤖 - #1495

Open
platelett wants to merge 3 commits into
tile-ai:ascendc_ptofrom
platelett:fix/flags-rowexpand-env
Open

[Fix] Harden row-expand layouts and clean up legacy tooling 🤖#1495
platelett wants to merge 3 commits into
tile-ai:ascendc_ptofrom
platelett:fix/flags-rowexpand-env

Conversation

@platelett

@platelett platelett commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

背景

本 PR 处理三个“调用看起来有效,但实际行为已经偏离接口表面含义”的问题:

  • row-expand 的访问指针覆盖多 stage,frontend 却可能只记录一个 stage 的行数;
  • T.init_flag / T.clear_flag 在 AscendC 生成 flag 指令,在 PTO 则被静默忽略;
  • set_env.sh 在 Bash 首次 source 时正常,但换成 Zsh 或重复 source 后结果会变化。

三项修改相互独立,因此分别放在三个 commit 中 review。

1. 修正 row-expand 的行数和布局

Row-expand 的最后一维固定为 256 字节,即 fp32 的 64 列或 fp16 的 128 列;前面的维度都应合并为行数。

例如,一个 fp32 Buffer 的 shape 是 [3, 8, 64]。完整 Buffer 应被视为 [24, 64],但旧逻辑只看到 [8, 64]。类似地,[1:3, :, :] 的指针已经覆盖两个 stage,src1 的 shape 计算却可能丢掉 stage 维,导致行标量数量与实际数据不一致。

本 PR 让 BufferBufferRegion 都按同一规则折叠前导维度。只有实际连续的 region 才能折叠;例如从 [2, 128, 8] 中选择 [0:2, 0:64, :] 会跨过未选择的内存,因此现在直接报错。

这不改变乘、减、除的计算,只修正 shape,并在进入 codegen 前拒绝当前后端无法表示的布局。

在下面列出的支持范围内,AscendC 和 PTO 现在都会使用同一个折叠后的 rows,并通过了相同的单 stage 和多 stage 数值用例。

具体接受哪些布局
  • dst、src0、src1 使用相同的 fp16/fp32 dtype;
  • rows 是静态的 8 的倍数,范围为 8..248
  • dst/src0 的物理行 stride 相同,并按 32 字节对齐;
  • 不带 tmp 时,src1 每行包含一个重复标量的 32 字节 block;
  • tmp 时,src1 是连续标量数组,tmp 每行提供一个 32 字节 block;
  • chunk * 64 等符号化 offset 在 TIR 能证明其对齐时仍可使用。

AscendC codegen 也会检查 uint8_t 能表示的 row 数和 stride,防止手工构造 intrinsic 时发生截断。

2. 删除只在 AscendC 生效的旧 flag 接口

T.init_flag / T.clear_flag 随 2025 年的实验性 AscendC backend加入。名字像通用同步 API,实际却是在 Python 端拼接 AscendC::SetFlag/WaitFlag 源码,再由 AscendC codegen 原样输出。

PTO 只支持结构化的 T.set_flag / T.wait_flag,不会处理这种 AscendC 字符串 AttrStmt。调用旧接口时 PTO 不报错,但也不生成 flag 指令,因而可能静默丢失同步。

当前代码和历史中没有找到这两个公开函数的调用、测试或文档;示例中的同名函数是各自定义的局部 T.macro,与它们无关。经 ChaoyangJi 确认,本 PR 直接删除这组公开函数及 AscendC 特例;结构化 flag API 保持不变。

3. set_env.sh 在 Bash/Zsh 中保持稳定

新脚本分别用 BASH_SOURCE[0]${(%):-%N} 找到被 source 的文件,并只在需要时把 canonical TL_ROOT 放到 PYTHONPATH 首位:

场景 旧结果 新结果
Bash 首次 source,原值为 /opt/lib $TL_ROOT:/opt/lib 相同
PYTHONPATH 为空 $TL_ROOT:,包含当前目录 $TL_ROOT
连续 source 两次 $TL_ROOT:$TL_ROOT:$old $TL_ROOT:$old
从 Zsh source 不能可靠定位脚本 正确得到同一个 TL_ROOT

因此,常见的 Bash 首次 source 保持原搜索顺序;空环境和重复 source 的差异则是有意删除空 path 和重复前缀。

测试

仓库中只保留最小回归,大矩阵仅在本地执行:

  • TileLang 增量重建通过;
  • row-expand pytest 15 passed,本地 AscendC/PTO 多 stage NPU 矩阵 24/24 passed
  • 现有 xAttention 完整编译通过,exp_experiment 回归 5 passed
  • legacy flag 回归 1 passed,Bash/Zsh 环境矩阵 12/12 passed
  • Ruff、clang-format、shell syntax 和 git diff --check 均通过。

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀

@platelett
platelett force-pushed the fix/flags-rowexpand-env branch from 89013a0 to 0f2fb10 Compare July 29, 2026 13:20
Fold contiguous leading Buffer and BufferRegion dimensions into the logical row count instead of dropping them.

Reject gapped regions and enforce the 256-byte row, packed src1, scalar-plus-tmp, dtype, count, and workspace contracts. Add focused frontend coverage and a two-backend multistage numerical regression.
Remove the public init_flag and clear_flag helpers together with the Ascend codegen branch that injected their prebuilt SetFlag/WaitFlag strings.

Keep the structured set_flag/wait_flag APIs unchanged and add a regression that confirms the legacy helpers are no longer exported.
Resolve the sourced script path with the active shell, export the canonical TL_ROOT, and prepend it to PYTHONPATH only when needed.

Handle unset and empty PYTHONPATH without an implicit current-directory entry, remain safe under nounset, and avoid duplication across repeated sourcing.
@platelett
platelett force-pushed the fix/flags-rowexpand-env branch from 0f2fb10 to d079907 Compare July 29, 2026 16:51
@platelett
platelett marked this pull request as ready for review July 31, 2026 03:50
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@LLMZhangYC LLMZhangYC left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

lgtm

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants