已合并
feat(linearize): triton-linearize integrate with master #615
candyhong创建于 2025年10月24日
feat(linearize): triton-linearize integrate with master #615
已合并
candyhong创建于 2025年10月24日
candyhong成员
2025年10月24日

相关的Issue

原因(目的、解决的问题等)

when there're division and modulo operator in addptr or mask expression, current triton-ascend will regard the correspondig load/store as unstructured , triton-to-unstructured pass will turn the unstructured dimensions to several levels of loops and perform scalar data movement . That is functional solution and the performence is not good . Triton-linerize pass will analyze addptr or mask expression with division/modulo operator and turn them into structured representation and the load/store could be executed as normal NDA instruction to guarantee the performence . TritonLinearize work as memory optimization solution described in the below diagram .
image.png

描述(做了什么,变更了什么)

linearize

  1. rewrite AddPtr, rewriteMask, RewriteLoad - kaixin, done

  2. UT - kaixin done, but just simpler one, need add more cases . - baiqing

  3. rewriteStoreOp -kaixin done

  4. rewriteAtomic -yihan

  5. use I64 as index, instead of I32 -yiming

  6. 随路转置 -ziqi

  7. rewrite Loop - Kaixin, done partially, while not added ? jiangzhen

  8. error handling and PR to main stream -kaixin done partially, there're other cases
    fallback indirect addptr - kaixin done
    fallback ptr loop - kaixin done
    fallback indirect mask - xuejie
    fallback undivisble offset -kaixin, done

  9. optimize mask generation to avoid selectOp -jiangzheng

  10. handle dynamic stride

  11. axis split : 16 -> 8 * 2, handle axis cat : 16 -> 10 + 6 ?

测试用例(新增、改动、可能影响的功能)

test_linearize.py

局限性

  • the current implemention assume offsets are divisible by the divisor , the low epression will be fallback to triton-to-unstructured .
    (2 + tl.arange(0,4) ) % 4
  • beow is the most common case that triton-linearize to handle , when offset are made by block_size, the block_size must be divisor or multile of M
    offset = s * block_size
    (offset + tl.arange(0, block_size) ) // M * S1 + (offset + tl.arange(0, block_size) ) % M * S2
likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 candyhong 的贡献)
Ccandyhong成员
2025年10月24日 创建了 pull request,commit 6ab81975
Ccandyhong成员
2025年10月24日 预合并成功(commit_id: c223d1b29fd5ae4f053b0387e8f5b6cb4a267a45)
Ccandyhong成员
2025年10月24日 推送  1 个提交:ef2037a2-fix: resolve conflicts with master
Ccandyhong成员
2025年10月24日 预合并成功(commit_id: dddb26c92f12df748cef6d53b3e0331ba6e9b57f)
ascend-robot
ascend-robot成员
2025年10月24日 评论:

以下是根据您提交的修改文件推荐的Reviewer和Committer序列,需各模块评审通过后方可合入

Module List Reviewers Committers
repo-Ascend/triton-ascend fengrui886, kang-ingu, shijingchang, zhucehw, kpxing weizhan4, ccdedreams, wcleungaj, wangtao489, liuzhuheng
likedislike
此处折叠了86条消息 查看更多
ascend-robotascend-robot成员
2025年10月24日 添加了label:approved
wangzhanpeng5
2025年10月24日 评论:

/lgtm

likedislike
ascend-robotascend-robot成员
2025年10月24日 添加了label:lgtm
ascend-robot
ascend-robot成员
2025年10月24日 评论:

Review Guide

This Pull-Request Passes Review.
Committers who writed a comment of /approve are: zhang-chunli01.
Reviewers who writed a comment of /lgtm are: zhang-chunli01, wangzhanpeng5.

likedislike
ascend-robotascend-robot成员
2025年10月24日 合入了pull request