triton-ascend-case-vector-mask-i32

name	triton-ascend-case-vector-mask-i32
description	比较两侧本身多为 `tl.int32`（offset / attn_arg 等），`arith.cmpi` 产出的是 `i1`；多段结果在 **`i1` 张量上做 `&`/`|`** 时， lowering 会在每条逻辑附近插入 **`extui`/`trunci`** 与 `select` 对齐；**每段比较后立刻 `.to(tl.int32)`**，让整条 mask 在 **`i32` 0/1 上 `&`/`|`**，后端更易连续处理 **`vand.i32`/`vor.i32`**。
category	improvement
version	1.0.0
metadata	{"case_type":"improvement","backend":"ascend","dsl":"triton_ascend","hardware":"Atlas A5"}

任务特征

• 多段 tl.int32 参与比较（如 q_off/k_off、q_attn_arg/k_attn_arg），得到的 bool 再用 & / | 拼成 mask，最后交给 tl.where。

原因

让 arith.cmpi 的 i1 结果在tensor<…xi1> 上反复 andi/ori，Ascend 后端会在每条逻辑与最后的 arith.select 之间额外插入arith.extui / arith.trunci等，把向量的宽度对齐前后使用的tl.int32，显式 .to(tl.int32)中间 不再在 i1 上对齐宽度，vector 上更易使用指令vand.i32 / vor.i32。

参考写法

@triton.jit
def mask(...):
    triu = (q_off[:, None] <= k_off[None, :]).to(tl.int32)
    return (
        (triu & ((q_arg[:, None] == k_arg[None, :]).to(tl.int32)
                 | (k_arg[None, :] == 0).to(tl.int32)))
        | (q_off[:, None] == k_off[None, :]).to(tl.int32))

注意

• mask 语义只需 0/1，tl.int32 的数据类型来源于比较操作符两边的数据类型。