IndirectLoad 算子设计文档

简介

目的

本文档描述 Graph-autofusion Autofuse 组件中 IndirectLoad ASCIR 算子的设计方案。IndirectLoad 对应 Inductor 前端中 torch.gather(input, dim, index)、index_select、embedding 等可归一化为 GatherElements / elementwise index 的访存场景，面向 A5/v35 Autofuse 流程，目标是在垂直融合场景下支持多种执行模板。

目标读者包括 Autofuse 开发人员、代码检视人员、测试开发人员。

范围

包含：

• Inductor 前端算子 gather、index_select、embedding 的自动融合后端处理。
• 新增独立 IndirectLoad ASCIR op type。
• IndirectLoad 的输入要求、IR 属性、模板选择信息传递和 codegen 分派。
• SIMD、SIMT、SK-like 三类模板的适用场景、拓扑要求、schedule/tiling/codegen 交接约定。
• A5/v35 PlatformV2 下的 schedule case 生成、ATT/tiling 资源传递和 AscendC codegen 设计。

不包含：

• Inductor 前端方案。
• A2/A3 等非 A5 平台。
• 当前已支持的 TensorFlow 语义 Gather 算子。

背景与问题定义

PyTorch Gather 语义与前端场景

torch.gather(input, dim, index) 是逐元素间接读取：输出每个位置都从 index 的对应位置取索引值，并沿 input 的 dim 维读取元素。

index_select(input, dim, index) 是沿 dim 维按一维 index 选择整块 slice，输出 shape 等于将 input.shape[dim] 替换为 index.numel()。例如 input[[2, 0], :] 可表示为 index_select(input, 0, [2, 0])。

embedding(weight, indices) 是沿 weight 第 0 维查表，输出 shape 为 indices.shape + weight.shape[1:]。例如二维 weight 下，embedding(weight, [[2, 0]]) 会取出第 2 行和第 0 行，并在末尾保留 embedding 维度。

三者可归一到间接读取语义：gather 是最通用的 elementwise index load；index_select 是 index 仅沿一个维度变化、选择整块 slice 的特化；embedding 是 dim = 0 且输入通常为 embedding 表的特化。

在 A5 上，torch.gather 的语义对应内置算子 GatherElements，而其他两个算子对应 GatherV2。

gather 的核心公式为：

output.shape == index.shape
input.rank == index.rank == output.rank
output[i, j, k] = input[i, index[i, j, k], k] # i/j/k 是 index 的 shape 范围

对非 dim 维，index.size(d) <= input.size(d)。

以 dim 维为界，任何输入的 shape 都可以压成 [preDim, gatherDim, postDim]，下文简称为 外轴、dim 轴、内轴。

Inductor NPU 扩展侧会把部分间接访存模式降到 IndirectLoad，后端只消费统一的 ASCIR 表示，不区分前端来源。

前端归一化后，Autofuse 只依赖 input、index、dim 和 shape/stride 信息生成模板 case；模板选择不直接判断原始 PyTorch API 名称。

ASCIR 定义

REG_ASC_IR(IndirectLoad)
    .Input("x1", "T1")                  // input输入
    .Input("x2", "T2")                  // index输入
    .Output("y", "T1")                  // 输出
    .Attr<int64_t>("dim")               // 表示 gather 的轴
    .Attr<bool>("need_check_negative")  // 是否支持对 index 的运行值做负数归一化
    .Attr<bool>("need_check_bound")     // 是否支持对 index 的运行值做合法性校验
    .ComputeType(ComputeType::kComputeLoad)
    .Impl();

设计目标

1. 用独立 IndirectLoad op 表达 gather、index_select、embedding 语义，在指定场景下与某些算子融合。
2. 在语义不变的前提下，通过模板化方案覆盖不同的适用场景，在特定网络拿到与 triton 对比的目标性能提升比。

总体方案

支持的融合范围

%%{init: {'theme': 'base', 'themeVariables': { 'fontSize': '14px', 'nodeSpacing': 10, 'rankSpacing': 20}}}%%
flowchart LR
  XLoad(("load")):::normal --> XBrc(("broadcast?")):::normal --> XEle(("elementwise*")):::normal --> IL(("IndirectLoad")):::indirect
  ILoad(("load")):::normal --> IBrc(("broadcast?")):::normal --> IEle(("elementwise*")):::normal --> IL
  IL --> OEle(("elementwise*")):::normal --> R(("reduce?")):::normal --> Store(("store")):::normal

  classDef normal fill:#fff3bf,stroke:#1e1e1e,color:#000,stroke-width:2px
  classDef indirect fill:#f9a825,stroke:#f57f17,color:#000,stroke-width:2px

融合仅支持垂直方向，* 表示 0 个或多个，? 表示 0 个或 1 个。

Inductor场景总体架构

代码总体架构

前端入口
  torch.gather / index_select / embedding
    -> inductor_npu_ext lowering / NPU Kernel
      -> ascir.ops.IndirectLoad(input, index, dim)

Autofuse schedule
  fuser.schedule(graph)
    -> ASCIR graph contains IndirectLoad
    -> PlatformV2::GenerateTasks
      -> IndirectLoadScheduleCaseGenerator
        -> 派生 SIMD / SIMT / SK-like schedule case
        -> 写入 template_type、拓扑约束、调用边界/range/resource 信息
        -> 生成性能打分函数，运行时调用，影响模板优先级
    -> AutoScheduler / ATT / host tiling
      -> 对候选 case 求解 tiling、UB、workspace、dcache、usedCore 等信息

Autofuse codegen
  fuser.codegen(scheduled_graph)
    -> IndirectLoadRegApiCall 根据模板类型（内部属性）生成对应的代码

模块间流程交互

模板方案

模板	适用场景	融合方式
SIMD	输入输出可全载，单次搬运量、分核足够，reduce节点（如果有）的R轴处于 dim轴 内部	按 IndirectLoad 的 外轴 分核分块，前后继算子链复用 VFCall 能力，循环内保证流水不断
SIMT	通用兜底场景，直接读写 GM，逐元素运算	将 IndirectLoad 输出降成1维分核，逐元素通过 SIMT API 或者标量表达式进行计算，reduce 场景为 SIMT + SIMD 混合
SK-like	IndirectLoad 分核过低，上述方案会导致前后融合链性能受损	类似 superkernel 的思想，输入输出链复用 VFCall 能力，与 IndirectLoad 分别独立执行，通过 workspace 串接输入输出

数据流与控制流

IndirectLoad 的数据流包含 input、index、output 三条链路：

• input 链路提供原始数据或 input pre 后的数据。
• index 链路提供每个 output 元素对应的 input dim 维坐标。
• output 链路消费 IndirectLoad 结果并执行 output post。

模块职责

模块	职责
ASCIR	定义 IndirectLoad(input, index, dim) op、IR attr 和 Python 构图入口。
Optimize / Schedule	识别 IndirectLoad 图，生成模板 graph case，改写轴和前后置链路边界。
ATT / Tiling	根据模板资源需求求解 UB、workspace、dcache、usedCore、tile size 和 actual size。
Codegen	按 template_type 分派，生成模板专用 kernel 片段并与普通 API 调用生成流程衔接。
测试	覆盖 ASCIR 构图、schedule case、codegen 文本、边界否决和设备侧正确性。

需求分析与设计

功能需求 1：ASCIR IndirectLoad 算子

介绍

新增 ASCIR IndirectLoad op type，表达 torch.gather(input, dim, index) 语义。

输入

端口	名称	说明
x[0]	input	源 tensor，rank 与 index 相同
x[1]	index	索引 tensor，shape 与 output 相同

处理

output[coord] = input[coord with coord[dim] = index[coord]]

输出

端口	名称	说明
y[0]	output	shape 与 index 相同，dtype 与 input 相同

IR 属性

属性	类型	说明
dim	int64	IndirectLoad 维度
need_check_negative	bool	是否对 index 运行值做负数归一化
need_check_bound	bool	是否对 index 运行值做合法性校验

ASCIR op 注册

IndirectLoad 需要在 autofuse/v35/ascir/generator/ascir_builtin_ops_v2.cpp 中通过 REG_ASC_IR(IndirectLoad) 注册，声明输入、输出、IR attr、ComputeType、dtype 约束和 v35 impl。基础注册绑定 IndirectLoadAscIrAttImplV2 和 IndirectLoadAscIrCodegenImplV2，使图中出现该 op 时可以进入 ATT 和 codegen 流程。

Python ASCIR 构图接口

pyascir 是 Inductor NPU 扩展生成 Autofuse ASCIR 图的 Python 接口。Inductor NPU 扩展在 codegen 阶段生成 ascir.HintGraph、ascir.ops.*、Autofuser.schedule() 和 Autofuser.codegen() 调用，因此 IndirectLoad 需要通过该接口暴露给上层图生成路径。

接入点包括：

• 在 autofuse/compiler/py_module/pyascir.h 的 REGISTERED_OPS 中注册 IndirectLoad。
• 在 autofuse/compiler/py_module/pyascir.cpp 中暴露 dim 等 IR attr 的 Python getter/setter。
• 在 Python 侧支持通过 ascir.ops.IndirectLoad("indirect_load") 创建节点，并设置 indirect_load.attr.ir_attr.dim。

前端 lowering 负责判断何时生成 IndirectLoad；Autofuse 侧保证生成后的 Python ASCIR 图可以被 schedule 和 codegen 消费。

ATT / Codegen impl 绑定

基础 op 注册需要在 autofuse/v35/ascir/generator/v2_ascir_att_impl.h 中声明 IndirectLoadAscIrAttImplV2，并在 autofuse/v35/ascir/generator/v2_ascir_codegen_impl.h 中提供 IndirectLoadAscIrCodegenImplV2。codegen impl 负责返回 API call 名称、API 名称、依赖头文件、节点合法性检查和必要的 tmp buffer 信息。

最简构图拓扑

Inductor NPU 扩展生成 IndirectLoad 时，最小 ASCIR 拓扑只包含 input/index 的 Data -> Load、IndirectLoad 和输出 Store -> Output：

x(Data)     -> input_load(Load)  -> IndirectLoad -> store(Store) -> y(Output)
index(Data) -> index_load(Load) -/

核心构图代码：

indirect_load = ascir.ops.IndirectLoad("indirect_load")
indirect_load.attr.sched.axis = [z4, z5, z6, z7]
indirect_load.attr.ir_attr.dim = 2
indirect_load.x1 = input_load
indirect_load.x2 = index_load

功能需求 2：模板候选生成

介绍

IndirectLoadScheduleCaseGenerator 在 PlatformV2::GenerateTasks 中识别 IndirectLoad 图，为 SIMD、SIMT、SK-like 生成对应 schedule task，并完成各模板需要的局部图改写和模板信息传递。

输入

• 原始 HintGraph / ImplGraph。
• IndirectLoad 节点及其 input、index、output 链路。
• 图中 axis、shape、stride、sched 信息。
• 候选 case 中重新定位后的 IndirectLoad 节点。

处理

• 候选识别：入口查找 IndirectLoad 节点；确认 IndirectLoad 的输入和融合拓扑满足基础要求。
• 模板枚举：为 SIMD、SIMT 复制 graph case，将模板类型、dcache_size 等信息写入属性；SK-like 需要额外生成 input pre、index pre、IndirectLoad、output post 四个逻辑阶段的边界和 workspace 串接信息。
• 改图优化：从语义上来说，IndirectLoad 是从输入中选取部分数据，理论上把 input 链路上的算子后移有可能减少运算量，但是 index 链路上无法适用。

SIMD 图改写

改图的核心是为了减小运算量，因此如果 IndirectLoad 把输入扩大，则不改图。例如下图中 dim 输出扩大时不前移：

%%{init: {'theme': 'base', 'themeVariables': { 'fontSize': '14px', 'nodeSpacing': 10, 'rankSpacing': 20}}}%%
graph TD
  A["input A [3, 4]"] --> B((Relu))
  B --> G((IndirectLoad)):::indirect
  I["index [3, 8], dim=1"]:::indirect --> G
  G --> S["netoutput [3, 8]"]
  classDef indirect fill:#f9a825,stroke:#f57f17,color:#000

输入不扩大的时候，尽可能沿 input 路前移；如果前置算子是双输入，需要复制一次 IndirectLoad，对性能可能带来影响，则停止前移：

%%{init: {'theme': 'base', 'themeVariables': { 'fontSize': '14px', 'nodeSpacing': 10, 'rankSpacing': 20}}}%%
graph LR
  subgraph Before[变换前]
    direction TB
    A1[input A] --> D1((Add))
    E1[input B] --> D1
    D1 --> B1((Relu))
    B1 --> C1((Abs))
    C1 --> G1((IndirectLoad)):::indirect
    I1[index]:::indirect --> G1
    G1 --> S1[netoutput]
  end

  subgraph After[变换后]
    direction TB
    A2[input A] --> D2((Add))
    E2[input B] --> D2
    D2 --> G2((IndirectLoad)):::indirect
    I2[index]:::indirect --> G2
    G2 --> B2((Relu))
    B2 --> C2((Abs))
    C2 --> S2[netoutput]
  end

  Before -. IndirectLoad 前移 .-> After
  classDef indirect fill:#f9a825,stroke:#f57f17,color:#000

SIMT 图改写

SIMT 的分核策略是将输出 shape 降成1维处理，轴与 input 链路不一致，会导致 input 链路上的算子断流水，因此策略是在 input 链路上无条件前移。

%%{init: {'theme': 'base', 'themeVariables': { 'fontSize': '14px', 'nodeSpacing': 10, 'rankSpacing': 20}}}%%
graph LR
  subgraph Before[变换前]
    direction TB
    X1[input] --> A1((Abs))
    A1 --> D1((Add))
    S1[side input] --> D1
    D1 --> G1((IndirectLoad)):::indirect
    I1[index]:::indirect --> G1
    G1 --> C1((Relu))
    C1 --> O1[output]
  end

  subgraph After[变换后]
    direction TB
    X2[input] --> G2A
    I2A[index]:::indirect --> G2A((IndirectLoad)):::indirect
    S2[side input] --> G2B
    I2B[index]:::indirect --> G2B((IndirectLoad)):::indirect
    G2A --> A2((Abs))
    A2 --> D2((Add))
    G2B --> D2
    D2 --> C2((Relu))
    C2 --> O2[output]
  end

  Before -. IndirectLoad 前移 .-> After
  classDef indirect fill:#f9a825,stroke:#f57f17,color:#000

输出

• 一个或多个模板 schedule task。
• 每个 task 对应的 graph case、阶段边界、私有属性、模板校验结果和图改写结果。
• SIMT 模板需要将 dcache 大小 (32K) 写入图的属性中。

功能需求 3：SIMD 模板生成

Kernel伪码

以 input=[s0,s1,s2,s3]、index/output=[s4,s5,s6,s7]、dim=2 为例：

for each outer tile from output outer axes: # 按照 [s4,s5] 的合轴分核分块
    output_range = decode_outer_coords(outer tile)

    index_raw = Load(index[output_range, :, :])
    index_used = IndexPreVFCall(index_raw)

    input_outer = output_range # 按 output 的外轴取 input 切片
    input_raw = Load(input[input_outer, :, :])
    input_used = InputPreVFCall(input_raw)

    offset = BuildGatherOffset(index_used)
    gather_out = Gather(input_used, offset)

    output_used = OutputPostVFCall(gather_out)

    reduce_output = ReduceCall(output_used) # 生成模板的时候已确保 reduce 的数据依赖

    Store(output[output_range, :, :], reduce_output)

性能特征

SIMD 模板减少 MTE 搬运，在 tile 内所有计算实现向量化，流水表现较好，理论上是性能最优的模板。

根据 index_select、embedding 语义来看，前端将其转成 gather 之后，会有一个将低维度 input 做一个类似 broadcast 的动作，适合向量化运算，如果可以识别是这种场景，并且 cacheline 利用率较好，优选此模板。

基本思路

SIMD 模板按 output/index 的外轴调度。每个 tile 先用 output 外轴坐标定位 input 的一个 UB 输入切片，再在这个输入切片上执行 input pre，最后用 index 构造 offset 并调用 AscendC::Gather。

当前的机制要求一张 AscGraph 图上只能由一套循环轴，SIMD 模板按 output/index 的外轴调度，而 input 链与其不同，因此需要做一些特殊的处理。

input full space:       [s0, s1, s2, s3]
output/index space:     [s4, s5, s6, s7]
schedule outer:         [s4, s5] -> indirect_load_outer
input pre vector space: [s2, s3]
output vector space:    [s6, s7]

模板要求 s0 >= s4、s1 >= s5，只计算 input 外轴中 [0:s4, 0:s5] 对应的输入切片，不遍历 input 外轴剩余部分。

Schedule 处理

创建模板的时候，将外轴合并作为 schedule 的唯一轴，强制要求只能按外轴切分；并设置 input pre 和 output/index 保持两套坐标系：

input pre:
  axis       = [z0, z1, z2, z3]
  vectorized = [z2, z3]

index/output:
  axis       = [indirect_load_outer, z6, z7]
  vectorized = [z6, z7]

autoschedule 对 SIMD input pre process node 做特殊跳过：TilingGroup::GenTilingGroup()、Scheduler::TileSplit()、Scheduler::ApplyBlockSplit() 不再按普通 elementwise 节点处理这条 input pre 链。这样 input pre 的 input 坐标系不会参与普通 output/index 侧分组，但 codegen 仍可从图中收集这条链路生成 VF 函数。

Codegen 处理

最终 kernel 由普通 codegen 和 IndirectLoadRegApiCall::GenerateSimd() 配合生成：

1. 对于 index 链路，直接复用当前现有机制生成 VFCall。
2. codegen 识别到是 input 链路后，不生成普通 API body，只保留必要的 VectorFunc 函数定义。
3. codegen 进入 IndirectLoad 节点时，需要反推 input 链路，生成 input 链路上的 VFCall，并生成 AscendC::Gather 的调用。
4. 继续处理 output 链路，它的 shape 和轴属性都一致，可以直接复用当前现有机制生成 VFCall。
5. 当前端要求对 index 的运行值进行负数标准化、合法性校验的时候，需要插入 where 语句。

功能需求 4：SIMT 模板生成

Kernel伪码

以 input=[s0,s1,s2,s3]、index/output=[s4,s5,s6,s7]、dim=2 为例：

for each output/index linear tile: # 按照 [s4,s5,s6,s7] 拉平后的轴分核分块
    # 无 reduce 时，tile 覆盖一段 output/index 线性空间
    # 有 reduce 时，tile 覆盖完整 reduce 内轴
    # IndexPreScalar/GatherScalar/OutputPostScalar 为 SIMT API 调用或是 scalar 运算

    AscendC::Simt::VF_CALL {
        for each lane element in tile:
            index_value = IndexPreScalar(index_gm[output_pos])
            value = GatherScalar[input_pos]
            value = OutputPostScalar(value)

            if has_reduce:
                reduce_input[lane_id] = value
            else:
                output_gm[output_pos] = value
    }

    if has_reduce:
        sync()
        output_gm[reduced_pos] = ReduceCall(reduce_input)

性能特征

SIMT 对 GM 直接逐元素访问，通过 DCache 机制保证读写性能，同样可以减少 MTE 的搬运，与 gather 离散访问的特性更加契合。

但是当子图中有 reduce 时需要 SIMT/SIMD 混合编程，此外某些 elementwise 算子尚未提供 SIMT API，需要通过 scalar 表达式直接运算。

基本思路

SIMT 模板直接逐元素读写 GM，有如下特点需要额外处理：

1. 改图处理已保证 input 链路上没有 input pre 算子链，因此整个流程直接按照 output 的合轴切分。
2. 需要旁路 Load / Store 节点，避免 UB 的引入。
3. 其他算子需要生成 SIMT API 或者标量计算表达式。
4. reduce 场景切分的时候需要考虑 reduce 内轴完整并且全载，而且 kernel 代码是 SIMT + SIMD reduce api 混合。

input full space:        [s0, s1, s2, s3]
output/index full space: [s4, s5, s6, s7]
schedule linear:         [s4, s5, s6, s7] -> indirect_load_linear
index/output scalar:     one lane -> one output element

Schedule 处理

创建模板的时候，将 output/index 的全部轴合并为 indirect_load_linear，并作为唯一 schedule 轴：

input:
  axis       = [z0, z1, z2, z3]
  vectorized = [z0, z1, z2, z3]

index/output:
  axis       = [indirect_load_linear]
  vectorized = [indirect_load_linear]

input Load 保留在图中作为 input GM 边界，但类似 SIMD 流程，autoschedule 也对 SIMT input Load 做一些特殊跳过处理，后续由 codegen 仍可从 IndirectLoad 追溯到 input GM tensor。

Codegen 处理

最终 kernel 由 IndirectLoadRegApiCall::GenerateSimt() 生成单个 SIMT helper 调用：

1. 从 IndirectLoad 输入追溯 input Load/Data，取得 input GM tensor。
2. 从 index 链路追溯 index Load/Data，取得 index GM tensor，并收集 index pre 标量表达式。
3. 从 output 链路追溯 Store，取得 output GM tensor，并收集 output post 标量表达式。
4. 生成 IndirectLoadIndexTransform 和 IndirectLoadOutputTransform 两个 __simt_callee__ functor。
5. 调用 IndirectLoadSimtExtend，由 helper 完成 index 读取、input 地址计算、input GM 读取和 output GM 写回。
6. 当前端要求对 index 的运行值进行负数标准化、合法性校验的时候，直接通过 scalar 语句完成。

功能需求 5：SK-like 模板生成

Kernel伪码

以 input=[s0,s1,s2,s3]、index/output=[s4,s5,s6,s7]、dim=2 为例：

# input pre 阶段：复用普通 VF 能力，按 input pre 自己的分核分块结果生成
input_gm[input_pre_range] -> workspace_input[input_pre_range]

# index pre 阶段：复用普通 VF 能力，按 index pre 自己的分核分块结果生成
index_gm[index_pre_range] -> workspace_index[index_pre_range]

# IndirectLoad 阶段：按 output 外轴合轴后的范围分核分块
workspace_input[input_pos] + workspace_index[index_pos] -> workspace_gather[output_pos]

# output post 阶段：复用普通 VF/Reduce 能力，按 output post 自己的规则生成
workspace_gather[output_post_range] -> output_gm[output_post_range]

性能特征

该模板适合 IndirectLoad 本体分核明显不足，且前后链路独立执行后能显著提高 core 利用率的场景。对搬运的优化效果有限，收益主要来自多个小算子的头开销合并。

基本思路

SK-like 模板按 input pre、index pre、IndirectLoad、output post 四个逻辑阶段组织，在 graph case 中为阶段边界插入 workspace，并沿用现有 ScheduleGroup / ATT / Codegen 流程完成分组、tiling 和代码生成。除 IndirectLoad 阶段外，其余阶段不新增专用调度语义，复用现有普通 VF/Reduce 能力。workspace 串接方式可参考 autofuse/optimize/fused_graph/fused_graph_modifier.cpp 中 SubgraphConnectionsToWorkspace、ChangeStartingOutputToWorkspace 的处理。

input pre     : input GM -> workspace_input
index pre     : index GM -> workspace_index
IndirectLoad  : workspace_input + workspace_index -> workspace_gather
output post   : workspace_gather -> output GM

input pre、index pre、output post 不内联到 IndirectLoad 阶段内部，而是分别作为独立逻辑阶段处理，并使用各自的分核分块原则。IndirectLoad 阶段只消费 workspace，按 output 外轴合轴后的范围分核分块，构造 offset 并调用 AscendC::Gather。这样 IndirectLoad 阶段的调度范围不会限制前后链路复用现有能力，物理 kernel 和 ScheduleGroup 组织由后续既有流程决定。

Schedule 处理

创建模板的时候，需要在复制出的 graph case 上保留四个逻辑阶段的边界、轴和 workspace layout，并将阶段间 Data/Output 边界改写为 Workspace 边界。input pre、index pre、output post 阶段复用现有 schedule/tiling 机制；output post 阶段如果包含 reduce，reduce 的 case 生成、是否进一步拆分、tiling 和 codegen 均交由现有 reduce 能力处理。IndirectLoad 阶段作为专用阶段，单独使用 output 外轴合轴后的 schedule 轴。

input pre/index pre/output post:
  schedule/range/output   = 复用现有能力生成的结果

IndirectLoad:
  axis       = [indirect_load_outer, z6, z7] # [s4,s5] -> indirect_load_outer
  vectorized = [z6, z7]
  input      = workspace_input, workspace_index
  output     = workspace_gather

input pre、index pre 和 output post 分别按自己的坐标系独立调度；各自 range 由现有 schedule/tiling 结果决定。三个阶段分别保留自己的 actual size 和 usedCore。

IndirectLoad 阶段不复用其它阶段的切分结果，只按 [s4,s5] -> indirect_load_outer 后的外轴合轴分核分块。schedule 需要保证四个逻辑阶段之间的依赖顺序，并将 workspace 大小、各阶段 actual size、usedCore 和 group 间变量依赖写入模板信息。最终物理 ScheduleGroup 数量由复用能力决定，不要求与四个逻辑阶段一一对应。

Codegen 处理

Codegen 不新增一套固定四段式 kernel 组织。graph case 中的 workspace 边界进入现有 ScheduleGroup 后，input pre、index pre、output post 继续复用现有 codegen 能力，IndirectLoad 阶段生成专用调用：

1. 按 input pre 的独立 schedule 生成 input pre 阶段，将 input_pre_range 写入 workspace_input。
2. 按 index pre 的独立 schedule 生成 index pre 阶段，将 index_pre_range 写入 workspace_index。
3. 按 output 外轴合轴 schedule 生成 IndirectLoad 阶段，从 workspace_input 和 workspace_index 读取数据，调用 AscendC::Gather 后写入 workspace_gather。
4. 按 output post 的独立规则生成 output post 阶段，从 workspace_gather 的 output_post_range 读取数据并写回 output GM。
5. 如果 output post 包含 reduce，则 reduce 仍由现有 output 侧能力处理，不并入 IndirectLoad 阶段。

功能需求 6：模板选择

介绍

IndirectLoad 同时生成 SIMD、SIMT、SK-like 三类候选模板，生成对应打分函数，打分函数随 ScheduleTask::score_func 传递给 ATT/tiling。ATT/tiling 阶段执行打分函数，根据当前实际 shape、blockDim 和资源信息决定候选模板优先级。

判定依据

模板选择使用两个判定量：

判定量	计算方式	作用
input 连续访问字节数	从 input 尾轴向前分析连续轴，计算连续轴乘积，再乘以 dtype size	判断 SIMD/SK-like 是否具备有效连续搬运
index 外轴分核规模	计算 index 中 dim 之前各轴的 repeat 乘积	判断 IndirectLoad 本体是否有足够分核

input 连续轴判定不依赖 dim 位置，只看 input stride 是否连续：

continuous_product = 1
for i in range(rank - 1, 0, -1):
    if strides[i - 1] == repeats[i] * strides[i]:
        continuous_product *= repeats[i]
    else:
        break

continuous_bytes = continuous_product * dtype_size

index 外轴分核规模按 dim 之前的轴计算：

index_outer_product = 1
for i in range(0, dim):
    index_outer_product *= index_repeats[i]

模板选择规则

条件	模板
continuous_bytes >= cacheline_size * Cache_thres 且 index_outer_product >= block_dim * Core_thres	SIMD
continuous_bytes >= cacheline_size * Cache_thres 且 index_outer_product < block_dim * Core_thres	SK-like
其他情况	SIMT

SIMD 适合 input 连续搬运收益足够、IndirectLoad 本体分核也足够的场景。SK-like 适合 input 连续搬运收益足够，但 IndirectLoad 本体分核不足、前后链路独立分核更有收益的场景。SIMT 作为兜底模板，覆盖连续访问不足、shape 不规则或资源条件不满足的场景。

打分函数生成

IndirectLoadScheduleCaseGenerator 在生成 SIMD、SIMT、SK-like 候选 graph case 时，同时生成对应 score_func。每个候选 case 对应一个打分函数：

模板	打分策略
SIMD	两个判定条件都满足时返回正分，否则返回负分
SIMT	兜底模板，返回中性分或低优先级正分
SK-like	连续访问条件满足但分核条件不满足时返回正分，否则返回负分

打分函数通过 ScheduleTask::score_func 传递给 ATT/tiling。ATT/tiling 根据实际 tiling_data.block_dim、shape 和模板资源情况执行打分，再选择候选模板。

数据流

IndirectLoadScheduleCaseGenerator
  -> 生成 SIMD / SIMT / SK-like graph case
  -> 计算 input 连续轴乘积、index 外轴乘积
  -> 生成每个 case 的 score_func
  -> 写入 ScheduleTask::score_func

ATT / tiling
  -> 执行 score_func
  -> 读取 tiling_data.block_dim
  -> 结合 UB / workspace / dcache 资源约束
  -> 选择模板优先级

Codegen
  -> 根据 template_type 生成对应模板代码

阈值暂定如下：

参数	含义	初始值
Cache_thres	cacheline 利用率阈值	0.5
Core_thres	分核利用率阈值	0.5

软件设计

设计原则

IndirectLoad 的语义信息和模板实现信息需要解耦。前端图只表达 gather(input, dim, index) 语义，不携带 SIMD、SIMT、SK-like 的执行策略；模板生成阶段根据 shape、stride、拓扑和资源约束生成候选模板，并把模板内部信息传递给后续 schedule、ATT、codegen 和 runtime 路径。

三类模板之间保持候选独立性。任一模板因拓扑、资源或 codegen 约束被否决时，不影响其它模板继续参与选择。模板选择不在图构造阶段固化，而是在 tiling 和运行期资源信息可用后，通过 score func、资源公式和 tiling case 共同决定。

模板元数据只作为 Autofuse 内部信息流转，不改变 IndirectLoad 的 ASCIR 语义。template_type、SIMT dcache reserve、SK-like 阶段边界和 workspace layout 等信息用于候选模板分派、资源计算和代码生成，不作为用户可感知的算子语义。

关键数据模型

IndirectLoad 语义信息

语义信息包含 input、index、dim、rank、shape、stride 和 dtype，用于描述 gather 的数学含义。该信息在三类模板中保持一致，模板只能改变执行方式，不能改变输出 shape、索引解释和异常边界。

模板选择信息

模板选择信息描述候选模板类型、适用条件和打分函数。SIMD、SIMT、SK-like 分别生成独立候选；每个候选保留自己的拓扑约束、schedule 轴、资源需求和否决条件。打分函数随候选模板进入 ATT/host tiling，用于在实际 shape、blockDim 和资源公式可用时选择优先级。

模板资源信息

模板资源信息描述 UB、workspace、tmp buffer、dcache reserve、actual size 和 usedCore 等约束。静态阶段只生成资源公式或资源上界，动态 shape 下的真实资源值由 host tiling 代入实际 shape 和 tiling case 后计算。资源不足时应否决对应 tiling case 或候选模板，而不是生成越界 kernel。

SK-like 阶段信息

SK-like 阶段信息描述 input pre、index pre、IndirectLoad、output post 四个逻辑阶段的边界、workspace layout、阶段间数据依赖和范围映射。四个阶段是逻辑执行边界，不要求与物理 ScheduleGroup 或 kernel 组织一一对应。最终执行顺序、workspace 分配和数据可见性由现有 ScheduleGroup / ATT / Codegen / Runtime 流程共同保证。

模块职责

模块	职责
ASCIR	表达 IndirectLoad(input, index, dim) 语义，保留输入、索引和输出关系。
Schedule case	识别 IndirectLoad 融合子图，生成 SIMD、SIMT、SK-like 候选模板及其内部元数据。
AutoScheduler	基于候选 graph case 枚举 tiling case，形成后续 ATT 和 codegen 可消费的 schedule 结果。
ATT/host tiling	承载 score func、资源公式、tiling key、workspace size 和实际 tiling data，完成候选选择和资源合法性判断。
Codegen	根据模板类型和 schedule 结果生成对应 kernel 代码，保证三类模板的语义输出一致。
Runtime	按 host tiling 写入的 tiling data、workspace size 和 launch 信息执行 kernel。

模板设计

SIMD 模板

SIMD 模板以 output/index 外轴作为主 schedule 轴。每个 tile 根据 output 外轴坐标确定 input 输入切片，将 input slice 搬入 UB 后执行 input pre，再由 index 构造 AscendC::Gather 所需 offset，最后执行 output post 和可选 reduce。

该模板的边界是 input slice、offset tmp 和中间结果必须满足 UB 资源约束。连续 input slice 搬运、tile 内复用和 AscendC::Gather 能被流水覆盖时，SIMD 优先级最高；若 input slice 或 tmp buffer 超资源，或者外轴分核不足导致 core 利用率较低，则对应 tiling case 应被否决或降级。

SIMT 模板

SIMT 模板将 output/index 线性化后按 lane 执行，每个 lane 读取 index、计算 input GM 地址、读取 input 值并写回 output。index pre 和 output post 中可标量化的表达式在 lane 内执行；input 链路不承载 input pre 算子链，避免在逐元素 GM 访问路径中引入不匹配的向量化语义。

该模板依赖 SIMT GM 访问、DCache reserve 和访存合并能力改善离散访问性能。index 分布具有局部性、lane 访问能落到较少 cacheline 时收益较好；若访问完全离散或 SIMT API 无法覆盖子图表达式，则只能作为通用兜底模板或被对应约束否决。reduce 场景需要与现有 SIMD reduce 能力协同，tiling 时必须保证 reduce 内轴完整性。

SK-like 模板

SK-like 模板面向 IndirectLoad 本体分核不足、但前后链路独立执行后可以提升 core 利用率的场景。模板将融合子图拆成 input pre、index pre、IndirectLoad、output post 四个逻辑阶段，通过 workspace 边界串接阶段间数据，并复用现有 VF/Reduce 的 schedule、ATT 和 codegen 能力。

该模板的收益来自阶段独立分核分块后提升前后链路并行度，代价是额外 workspace 写读和阶段间依赖管理。workspace 搬运进入 critical path 或 workspace size 超资源时，不应优先于 SIMD/SIMT。SK-like 不引入固定四段式 kernel 组织，物理 ScheduleGroup 数量和 kernel 组织由现有流程根据 graph case 和 schedule 结果决定。

数据流

ASCIR graph
  -> IndirectLoad 语义输入、dim、shape、stride、dtype

Schedule case
  -> 生成 SIMD / SIMT / SK-like 候选模板
  -> 写入模板类型、资源需求、score func、阶段边界和 workspace layout

AutoScheduler
  -> 枚举 tiling case
  -> 形成 schedule 结果和 group 间依赖

ATT / host tiling
  -> 计算 UB / workspace / dcache 资源
  -> 执行 score func 和 perf 评估
  -> 选择 tiling case，写入 tiling key、tiling data 和 workspace size

Codegen
  -> 根据模板类型生成 SIMD / SIMT / SK-like kernel 代码
  -> 校验 tiling data、tmp buffer、workspace 和模板元数据一致性

Runtime
  -> 按 host tiling 结果 launch kernel

错误处理和否决策略

Schedule 阶段

Schedule 阶段负责否决语义或拓扑不满足的候选模板。rank、dim、shape、stride、dtype 不合法时，整个 IndirectLoad 模板生成失败；单个模板的拓扑约束不满足时，只否决该模板，不影响其它模板继续生成。

SIMD 需要确认 input slice 能由 output 外轴唯一定位，且 input pre、index pre、output post 的执行空间边界清晰。SIMT 需要确认 input 链路不存在必须在 input 空间执行的 pre 算子链，index pre 和 output post 可转成 SIMT 表达式或标量表达式。SK-like 需要确认四个逻辑阶段可以通过 workspace 边界串接，并能保持数据依赖顺序。

ATT/host tiling 阶段

ATT/host tiling 阶段负责资源合法性判断。SIMD 的 input slice、offset tmp 和中间 UB buffer 超资源时，否决对应 tiling case。SIMT 的 dcache reserve、lane 访问范围和 reduce 内轴完整性不满足时，否决对应 tiling case。SK-like 的 workspace size、阶段 actual size、usedCore 或 group 间依赖不一致时，否决对应 tiling case。

动态 shape 场景下，静态阶段不直接使用未知 shape 做最终判断，只生成可代入的资源公式和 score func；host tiling 使用真实 shape 和 tiling case 计算实际资源，保证 kernel launch 前完成合法性判断。

Codegen 阶段

Codegen 阶段负责校验模板元数据与 schedule 结果一致。若模板类型缺失、tiling data 缺失、workspace layout 不一致、tmp buffer size 与资源公式不一致，或 SIMT helper 无法表达目标子图，应返回错误并记录清晰日志。Codegen 不应在发现资源不合法时生成依赖未定义行为的 kernel。

非功能需求

可维护性

• 模板行为通过 template_type 分派，避免在普通 Gather/Load 路径中散落 IndirectLoad 特判。
• SIMD、SIMT、SK-like 的拓扑约束、tiling 资源和 codegen 逻辑分章节描述，便于独立扩展。
• 普通 VF、普通 Gather 和普通 Load 路径不承担 IndirectLoad 的语义判断。

可测试性

• dump 文件保证可用性。
• Python smoke 覆盖 SIMD/SIMT/SK-like codegen 分支。
• UT 覆盖 schedule case 生成、template type 写入、拓扑否决和调用边界/range 信息。
• ST/E2E 覆盖 kernel 生成和设备侧精度。

可移植性

非 A5 平台不进入该路径。

可靠性

• shape、rank、dim、stride、dtype、offset 使用前必须校验。
• workspace、UB、dcache 由 tiling 阶段计算并传给 codegen，codegen 做一致性检查。
• 图改写只在模板局部 graph case 中执行，避免污染其它模板。

性能

编译时长

新增 schedule task、ATT 和 codegen 分支会增加编译期图遍历、tiling 代码生成和字符串生成成本。应限制在包含 IndirectLoad 的图中触发，不影响普通图路径。

执行性能

• SIMD 目标是通过 UB 内输入切片全量搬入和 AscendC::Gather 减少逐元素 GM 读取。
• SIMT 是通用执行路径，性能取决于 index 分布和 GM 访问局部性。
• SK-like 目标是避免低分核 IndirectLoad 拖累前后 V 算子并行度。

内存和产物大小

SIMD 增加 UB offset tmp；SK-like 增加 workspace；SIMT 增加 dcache reserve。SK-like 的 workspace size 由阶段串接边界决定，host tiling 需要写入 workspace size，并保证后续 codegen 和 kernel 消费一致。新增 codegen/helper 会增加少量动态库和头文件产物大小。

特性交叉影响

场景	适用性	分析说明
SuperKernel Python 接口	不适用	不修改 SuperKernel Python 包和选项。
SuperKernel C++/AOT 接口	不适用	不修改 libascendsk.so 和 AOT 接口。
Autofuse 图优化	适用	新增 IndirectLoad schedule case、模板局部 graph rewrite 和 VF 隔离规则。
Autofuse Codegen/Backend	适用	新增 IndirectLoad 专用 codegen、helper 和 tiling 消费逻辑。
AscendC API / Runtime 交互	适用	SIMD/SK-like 使用 AscendC::Gather，SIMT 逐元素计算 input GM 地址并读取 input。
Python/C++ 混合绑定	适用	通过 pyautofuse 暴露 ASCIR op 构造能力。
构建与打包	适用	新增源码和头文件需进入构建/安装路径。
测试与覆盖率	适用	需要新增 UT/ST/smoke 覆盖三模板。
性能与日志	适用	影响融合收益、UB/workspace、GM 访存和日志路径。
兼容性	适用	不改变现有 Gather 行为；新增 op type 和模板枚举需评估 ABI/API。

安全与兼容性

安全检查

• 不硬编码敏感信息、公网地址或芯片型号判断。
• dim、rank、shape、offset、workspace size 使用前必须校验。
• byte offset 和 shape 乘积必须防止溢出。
• 图改写保持数据边和控制边等价。
• AscendC 接口调用满足内存生命周期和同步约束。
• 测试临时产物不进入版本库。

ABI/API 兼容性

• 新增 ASCIR op type 和 IR attr 属于新增能力。
• 不改变现有 Gather、Load、GatherToLoadPass 行为。
• Python/C++ 构图接口新增 IndirectLoad，不改变已有接口调用方式。

图改写确定性

模板 graph case 从复制图中重新获取 IndirectLoad 节点，避免复用原图节点指针。图改写依赖明确的 axis、node 拓扑和端口关系，不依赖无序容器遍历结果生成最终结构。

Runtime 生命周期约束

SIMD 路径遵循普通 UB queue 生命周期。SIMT 路径跳过普通 UB output 生命周期，但 input/index/output GM 指针必须在 helper 调用期间保持有效。SK-like 路径通过 workspace 边界串接四个逻辑阶段，workspace 生命周期、group 间变量依赖和数据可见性由现有 ScheduleGroup / ATT / Runtime 约束保证。

测试设计

测试边界

测试覆盖 ASCIR 构图、schedule case、template type 写入、tiling/codegen 分派、SIMD/SIMT/SK-like kernel 文本和设备侧正确性。

最简拓扑用于模拟前端直接构造 IndirectLoad：没有 input pre、index pre 和 output post，只保留 input/index 的 Data -> Load、IndirectLoad、Store -> Output。

from pyautofuse import ascir, Autofuser, AutofuserOptions

ascir.utils.set_platform("3510", 1, 245760)

graph = ascir.HintGraph("indirect_load_minimal_graph")
s0, s1, s2, s3 = [graph.create_size(f"s{i}") for i in range(4)]
s4, s5, s6, s7 = [graph.create_size(f"s{i}") for i in range(4, 8)]
z0 = graph.create_axis("z0", s0)
z1 = graph.create_axis("z1", s1)
z2 = graph.create_axis("z2", s2)
z3 = graph.create_axis("z3", s3)
z4 = graph.create_axis("z4", s4)
z5 = graph.create_axis("z5", s5)
z6 = graph.create_axis("z6", s6)
z7 = graph.create_axis("z7", s7)

x = ascir.ops.Data("x", graph)
index = ascir.ops.Data("index", graph)
y = ascir.ops.Output("y", graph)

x.attr.ir_attr.index = 0
x.y.dtype = ascir.dtypes.float16
x.y.axis = [z0, z1, z2, z3]
x.y.size = [s0, s1, s2, s3]
x.y.strides = [s1 * s2 * s3, s2 * s3, s3, ascir.SizeExpr(1)]

index.attr.ir_attr.index = 1
index.y.dtype = ascir.dtypes.int32
index.y.axis = [z4, z5, z6, z7]
index.y.size = [s4, s5, s6, s7]
index.y.strides = [s5 * s6 * s7, s6 * s7, s7, ascir.SizeExpr(1)]

input_load = ascir.ops.Load("input_load")
input_load.attr.ir_attr.offset = ascir.SizeExpr(0)
input_load.attr.sched.axis = [z0, z1, z2, z3]
input_load.x = x
input_load.y.dtype = ascir.dtypes.float16
input_load.y.axis = [z0, z1, z2, z3]
input_load.y.size = [s0, s1, s2, s3]
input_load.y.strides = [s1 * s2 * s3, s2 * s3, s3, ascir.SizeExpr(1)]

index_load = ascir.ops.Load("index_load")
index_load.attr.ir_attr.offset = ascir.SizeExpr(0)
index_load.attr.sched.axis = [z4, z5, z6, z7]
index_load.x = index
index_load.y.dtype = ascir.dtypes.int32
index_load.y.axis = [z4, z5, z6, z7]
index_load.y.size = [s4, s5, s6, s7]
index_load.y.strides = [s5 * s6 * s7, s6 * s7, s7, ascir.SizeExpr(1)]

indirect_load = ascir.ops.IndirectLoad("indirect_load")
indirect_load.attr.sched.axis = [z4, z5, z6, z7]
indirect_load.attr.ir_attr.dim = 2
indirect_load.x1 = input_load
indirect_load.x2 = index_load
indirect_load.y.dtype = ascir.dtypes.float16
indirect_load.y.axis = [z4, z5, z6, z7]
indirect_load.y.size = [s4, s5, s6, s7]
indirect_load.y.strides = [s5 * s6 * s7, s6 * s7, s7, ascir.SizeExpr(1)]

store = ascir.ops.Store("store")
store.attr.sched.axis = [z4, z5, z6, z7]
store.attr.ir_attr.offset = ascir.SizeExpr(0)
store.x = indirect_load
store.y.dtype = ascir.dtypes.float16
store.y.axis = [z4, z5, z6, z7]
store.y.size = [s4, s5, s6, s7]
store.y.strides = [s5 * s6 * s7, s6 * s7, s7, ascir.SizeExpr(1)]

y.x = store
y.attr.ir_attr.index = 0
y.attr.sched.axis = [z4, z5, z6, z7]
y.y.dtype = ascir.dtypes.float16
y.y.axis = [z4, z5, z6, z7]
y.y.size = [s4, s5, s6, s7]
y.y.strides = [s5 * s6 * s7, s6 * s7, s7, ascir.SizeExpr(1)]

options = AutofuserOptions(tiling_lib_path="", tiling_lib_codegen_symbol="")
fuser = Autofuser(options)
impl_graphs = fuser.schedule(graph)
assert impl_graphs

tiling_def, host_tiling, op_kernel = fuser.codegen(impl_graphs)
assert tiling_def
assert host_tiling
assert op_kernel

测试用例设计

测试类别	关键测试项	测试方法	用例类型
功能	IndirectLoad op 构图和 attr	Python/C++ 构图	UT
功能	SIMD codegen	检查输入切片全量搬入、offset、AscendC::Gather 代码	UT/ST
功能	SK-like 阶段边界信息	检查四个逻辑阶段、workspace 和同步点信息	UT
功能	SIMT codegen	检查 helper、线性化逐元素 GM 读取、内联 elementwise	UT/ST
异常	非法 dim/rank/shape	构造非法图	UT
兼容	现有 Gather 不受影响	现有 Gather 用例	UT/ST
性能	三模板收益边界	benchmark	Benchmark

测试命令

cmake --build build --target aihac_codegen -j 8
PYTHONPATH=build/autofuse/compiler/py_module:$PYTHONPATH python3 autofuse/tests/ut/ascendc/compile/indirect_load_simd.py
PYTHONPATH=build/autofuse/compiler/py_module:$PYTHONPATH python3 autofuse/tests/ut/ascendc/compile/indirect_load_simt.py
sh build.sh -u --module=autofuse_framework -j 8
sh build.sh -s --module=autofuse_e2e -j 8

验收标准

1. 新增代码通过 UT/ST 验证，满足覆盖率要求，充分构造各种异常场景。
2. 泛化用例功能正常，无 AIC ERROR，精度达标。
3. 泛化用例性能，典型融合片段对比单算子模式不劣化。
4. 关键客户网络(ETA/SIM 等)性能达标，对比 triton 性能提升挑战 10%。

设计文档检查结果

• 跨特性交叉影响：已按 cross_feature_check.md 检查 SuperKernel、Autofuse optimize/codegen、AscendC、Python/C++ 绑定、构建打包和测试场景影响。
• 编码红线：已检查敏感信息、外部输入校验、整数溢出、资源生命周期、ABI/API、图改写确定性和 AscendC 生命周期约束。
• 测试设计：已说明 UT/ST/benchmark 覆盖策略，覆盖 ASCIR、schedule、tiling、codegen 和兼容性。
• 性能评估：已说明编译期影响、三模板执行性能边界、UB/workspace/dcache 资源影响。