Skip to content

Ansor Transform Steps

Annotation

AnnotationStep: AN

相关代码：AnnotationStepNode，AnnotationStepNode::WriteToRecord

输出的 json 格式：

["AN", stage_id, iter_id, int(annotation)]

annotation 是 IteratorAnnotation 类型的变量，这是一个枚举类型，其值分别代表：

值 含义 值 含义
1 循环展开 2 向量化（vectorized）
3 并行（paralleled） 4 绑定到 vthread
5 绑定到 blockIdx.x 6 绑定到 threadIdx.x
7 绑定到 blockIdx.y 8 绑定到 threadIdx.y
9 绑定到 blockIdx.z 10 绑定到 threadIdx.z
11 张量化（tensorize） 0 无

Fuse

FuseStep: FU

相关代码：FuseStepNode，FuseStepNode::WriteToRecord

json 格式：

["FU", stage_id, fused_ids];
  • 合并对应 Stage 中 fused_ids 对应的 iterator。

Pragma

PragmaStep: PR

相关代码：PragmaStepNode，PragmaStepNode::WriteToRecord

该 step 对应 pragma 的预处理指令。

输出的 json 格式：

["PR", stage_id, iter_id, pragma_type]
  • pragma_type：如 "auto_unroll_max_step$16" 即设置该 iterator 下自动进行循环展开，且展开的最大步数为 16。

Reorder

ReorderStep: RE

相关代码：ReorderStepNode ReorderStepNode::WriteToRecord

["RE", stage_id, after_ids]

after_ids：在 reorder 之后，所有 iterator 的顺序。

Split

SplitStep: SP

相关代码：SplitStepNode，SplitStepNode::WriteToRecord

输出的 json 格式：

["SP", stage_id, iter_id, extent, lengths, int(inner_to_outer)]

extent：该轴的 extent 长度。

lengths：split 因子数组。

inner_to_outer：若为 1，表示 lengths 指示了从内层到外层的迭代长度

FollowSplitStep: FSP

相关代码：FollowSplitStepNode FollowSplitStepNode::WriteToRecord FollowSplitStepNode::ExtractSplitLengths

与 SplitStep 相似，不过，FollowSplitStep 使用来自于其他 step 的 factors（即，Follow another split step）。

输出的 json 格式：

["FSP", stage_id, iter_id, src_step_id, n_split];
  • src_step_id：要 follow 的 split step 的索引
  • n_split：split 级别的数量
    • 取索引为 src_step_id 的 Transformer step，该 step 一定是 Split 类别的。
    • 本次 Split 的 lengths 数组有 nsplit 个 factor，它们由源 Split Step 的 lengths 数组前 n_split - 1 个 factor，以及剩余 factor 的累积组成。

FollowFusedSplitStep: FFSP

相关代码：FollowFusedSplitStepNode FollowFusedSplitStepNode::WriteToRecord

与 FollowSplitStep 类似，但是 FollowFusedSplitStep 使用的 split factors 来自于多个 step。

输出的 json 格式：

["FFSP", stage_id, iter_id, src_step_ids, level, int(factor_or_nparts)];

  • src_step_ids：要遵循的 split step 的索引的数组

    • 数组中每个元素是一个 Transformer step，如果该 step 是一个 split step，那么就获取其第 level 层的长度（从 0 开始）
    • 将这些长度乘起来，作为该 FFSP 步骤的 split factor。

  • factor_or_nparts：如果为真（1），使用 factor；若为假（0），使用 nparts。

Compute At

ComputeAtStep: CA

ComputeAtStepNode ComputeAtStepNode::WriteToRecord

["CA", stage_id, target_stage_id, target_iter_id]

stage_id 对应 node 的计算将在 target_stage_id stage 的 target_iter_id 迭代下进行。

Compute Inline

ComputeInlineStep: CI

ComputeInlineStepNode ComputeInlineStepNode::WriteToRecord

["CI", stage_id]

该 step 表明 stage_id 对应的 compute 将在其下游 stage 节点 inline 展开。

Compute Cache

CacheReadStep: CHR

CacheReadStepNode CacheReadStepNode::WriteToRecord

["CHR", stage_id, scope_name, reader_stage_ids]
  • scope_name：字符串，表示新增的 read stage 的 scope，如 local、shared、global。
  • reader_stage_ids：read stage 的 stage_id 数组（即该 read cache node 的 consumers node）。

CacheWriteStep: CHW

CacheWriteStepNode CacheWriteStepNode::WriteToRecord

["CHR", stage_id, scope_name]
  • scope_name：新 compute stage 的 scope name，如 local、shared、global。

Others

StorageAlignStep: SA

StorageAlignStepNode

暂时没遇到过

RfactorStep: RF

RfactorStepNode

和将 reduce iterator 进行 parallel 优化相关，暂时没遇到过。

ComputeRootStep: CR

ComputeRootStepNode