external-cannbot-ops-triton-latency-optimizer

Latency Optimizer Skill

Input and Output

Input Parameters

• code_file_path: input code file path (Triton Ascend operator code)
• output_path: output code file path (the optimized code must be written to this path)
• npu: NPU device ID
• arch: hardware architecture

Output Requirements

Must produce：

1. output_path the specified optimized code file
2. optimization strategy description (in code comments or return information)
3. functional consistency description
4. precision consistency description

If there are no more optimization points：

• still must produce output_path(content identical to or slightly adjusted from the input)
• clearly state "no more optimization points" in the return information

Optimization Point Execution Order

The Agent must check the optimization points one by one in strict accordance with the following order,only one optimization point may be attempted at a time; after a hit, refer to the corresponding document。

⚠️ Prerequisites: the "hit condition" of an optimization point must first be hit (the code features satisfy one of the typical code features and the applicable conditions hold) before the corresponding reference document can be loaded. If not hit, skip it and loading the reference document is prohibited.

---

Optimization Point 1: Input Parameter Staticization Optimization

Applicable conditions: there exist fixed parameters in the code that can be declared as tl.constexpr fixed parameters

Typical code features：

@triton.jit
def kernel(A, B, C, M, N,
            stride_am, stride_an,  # 运行时不变化的固定值，但未声明为 constexpr
            BLOCK_SIZE_M: tl.constexpr,
            BLOCK_SIZE_K: tl.constexpr):

Judgment logic：

• If there exist fixed parameters in the code that do not change at runtime (such as stride, fixed values, BLOCK_SIZE, etc.) that are not declared as tl.constexpr → involved
• If all fixed parameters have been correctly declared as tl.constexpr → not involved, skip

Hit condition: the code features satisfy one of the above typical code features, and the applicable conditions hold

Reference document：references/constexpr_parameters.md

---

Optimization Point 2: Tiling Optimization (Contiguous Axis Vectorization)

Applicable conditions: reduction-type or normalization operators that process multi-dimensional tensors (3D and above), and the reduction axis is not the most contiguous axis in the memory layout

Typical code features：

@triton.jit
def kernel(input_ptr, output_ptr, dim1, dim2, ...):
    # 特征 1：向量化偏移 tl.arange 作用在非连续轴（如 dim1/M 轴）
    m_offsets = tl.arange(0, BLOCK_SIZE_M)
    # 特征 2：访存偏移计算中，向量化部分乘上了较大的 stride
    input_offset = m_offsets * stride_m + n_idx * stride_n
    # 特征 3：循环内部频繁进行还原操作（如 tl.sum）将向量压缩为标量
    acc = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
    ...
    total_sum = tl.sum(acc, axis=0)

Judgment logic：

• Check tl.load : if the vector offset produced by tl.arange acts on the axis of stride > 1 , while the axis of stride = 1 is only treated as a scalar index → involved
• Check the loop accumulator: if the accumulator is chunked on the reduction axis, but the memory access pattern causes non-contiguous memory reads → involved
• If tl.arange already acts on the most contiguous axis of memory (usually the last dimension of the last tensor), and merged memory access is implemented → not involved, skip

Hit condition: the code logic is intended to reduce over a certain dimension, but its chunking strategy causes the hardware to perform strided memory access

Reference document：references/tiling_optimization.md

---

Optimization Point 3: Core Partitioning Optimization

Applicable conditions: the Grid size setting in the code is unreasonable, or the NPU hardware resources are not fully utilized

Typical code features：

# 特征 1：Grid 远大于物理核数
grid = (batch_size,)  # 如果 batch_size=128，远超 48 核

# 特征 2：Grid 远小于物理核数
grid = (batch_size // 64,)  # 如果 batch_size=128，只有 2 核

# 特征 3：每个 program 只处理 1 行数据
row_idx = tl.program_id(0)
x = tl.load(ptr + row_idx * stride + cols, mask=mask)

# 特征 4：未使用编译优化选项（multibuffer、unit_flag）
kernel[grid](...)  # 未传入 multibuffer、unit_flag

Judgment logic：

• Check whether the Grid size is close to the number of physical cores (40-48)
  • If Grid >> 48 or Grid << 48, or the Grid value cannot be determined → involved
• Check the amount of data processed by each program
  • If each program processes only a small amount of data (e.g., 1 row) → involved
• Check whether compilation optimization options are used
  • If multibuffer is not used and it is a memory-intensive operator → involved
• If the Grid is reasonable and optimization options are used → not involved, skip

Hit condition: The Grid size setting in the code is unreasonable, or NPU hardware resources are not fully utilized

Reference document：references/vector_core_partition.md

---

Optimization point 4: Discrete memory access optimization

Applicable conditions: The code contains access to global memory through random/unpredictable indices

Typical code characteristics：

# 索引来源于 tl.load 加载的值（随机性）
idx = tl.load(indices_ptr + offset)  # idx 是运行时确定的随机值
val = tl.load(data_ptr + idx)        # 通过随机索引访问

# 或者索引来源于 kernel 入参（可能是随机值）
val = tl.load(ptr + random_index)

Judgment logic：

• Check tl.load index source:
  • If the index is tl.program_id linear transformation → deterministic and continuous, not involved
  • If the index is a linear transformation of the loop variable → deterministic stride, not involved
  • If the index comes from tl.load the loaded value or kernel input parameter → potentially random, involved
• If all memory access indices are deterministic continuous/stride patterns → not involved, skip

Hit condition: The code contains access to global memory through random/unpredictable indices

Reference document：references/discrete_memory_access.md

---

Optimization point 5: Scalar to Vector optimization

Applicable conditions: The code contains scalar operations that can be converted to vector operations to fully utilize the NPU Vector computation unit

Typical code characteristics：

# 特征 1：标量广播操作
scalar_val = 0.5  # Python 标量
result = x * scalar_val  # scalar 广播，无法启用 vector 加速

# 特征 2：标量规约操作
sum_val = 0.0  # 标量累加器
for n in range(N):
    val = tl.load(x_ptr + n)
    sum_val += val  # 标量加法

# 特征 3：标量控制流
if x > 0:  # 标量条件，导致 SIMD 分支分化
    result = tl.exp(x)
else:
    result = tl.cos(x)

# 特征 4：int 类型比较/除法/取余
is_invalid = tok < 0  # int 类型比较，退化为标量
c = a // b  # int 类型除法，退化为标量
d = a % b   # int 类型取余，退化为标量

Judgment logic：

• Check whether there are computations between Python scalars and vector data (scalar broadcasting)
• Check whether there is a scalar accumulator (e.g., sum_val = 0.0）
• Check whether there is if-else control flow processing vector data
• Check whether there are int32/int64 -type comparison, division, and modulo operations
• If any of the above conditions exists → involved
• If all operations already use vector form → not involved, skip

Hit condition: The code contains scalar operations that can be converted to vector operations

Reference document：references/scalar_to_vector.md

---

Optimization point 6: Avoid vector API scalar downgrade

Applicable conditions: The code contains vector operations that may be downgraded by the compiler to scalar loops, including general arithmetic operations, comparison operations, extended multiplication, cumulative operations (cumsum/cumprod), or reduction operations (reduce)

Typical code characteristics：

# 特征 1：通用算术操作使用 i64，或者满足降级条件
z = x + y  # x/y 为 i64
z = x % y  # x/y 为 i32且执行取余计算

# 特征 2：整数比较操作（非 i32 EQ/NE，或非浮点比较）
mask = x < y  # i8/i16/i32/i64 的 LT/GT/LE/GE 比较

# 特征 3：扩展乘法
z = x * y  # 触发 vmulext，始终降级

# 特征 4：cumsum/cumprod 在最后一个维度上操作
x_cumsum = tl.cumsum(x_1d, axis=0)  # 一维张量，或 cumDim 是 lastDim

# 特征 5：reduce 操作在特定条件下
# i64 类型的 sum/prod/max/min
# 整数类型的 argmax/argmin
# 浮点类型 argmax/argmin 且 flatten 后维度 > 2

Judgment logic：

• Check general arithmetic operations (add/sub/mul/min/max/abs/shl/shr/interleave/deinterleave): if the data type is i64
• Check comparison operations: if the data type is i8/i16/i64 (all comparisons), or i32 LT/GT/LE/GE → involved
• Check modulo operations: if the data type is any int type → involved
• Check extended multiplication (vmulext): any extended multiplication → involved
• Check cumsum/cumprod: if the cumulative dimension is the last dimension of the input tensor (when one-dimensional, axis=0 is the last dimension), or the data type is i64 → involved
• Check reduce operations: if they are i64-type sum/prod/max/min; integer-type argmax/argmin; floating-point-type argmax/argmin with flattened dimension > 2 → involved
• If none of the above cases exist → not involved, skip

Hit condition: there is any of the above vector operations in the code, and the corresponding scalar fallback condition is satisfied

Reference document：references/avoid_scalar_lowering.md

---

Optimization point 7: Pass merging optimization

Applicable condition: the code contains multiple traversals of the same data to compute different statistics

Typical code characteristics：

# Pass 1: 计算 mean
for ...:
    data = tl.load(...)
    mean += tl.sum(data)

# Pass 2: 计算 variance（再次遍历！）
for ...:
    data = tl.load(...)  # 重复加载
    var += tl.sum((data - mean) ** 2)

# Pass 3: 归一化（第三次遍历！）
for ...:
    data = tl.load(...)  # 第三次加载
    tl.store(...)

Decision logic：

• Check whether there are multiple independent loops traversing the same data
• Check whether multiple statistics can be computed simultaneously (e.g., sum + sum_sq can compute mean + var simultaneously)
• If there are multiple traversals and they can be merged → involved
• If there is only a single traversal, or the statistics have dependencies and cannot be merged → not involved, skip

Hit condition: the code contains multiple traversals of the same data, and the computations can be merged

Reference document：references/pass-merge.md

---

Optimization point 8: Dimension merging optimization

Applicable condition: the code contains multiple levels of nested loops processing contiguous dimensions, and there are no dependencies between dimensions

Typical code characteristics：

# 问题代码：3层循环处理 NCHW 布局
for n in range(N):           # 64 次
    for h in range(H):       # 512 次
        for w_start in range(0, W, BLOCK_SIZE):  # 循环层数过多
            base_offset = n * stride_n + c * stride_c + h * stride_h
            data = tl.load(input_ptr + base_offset + ...)

Decision logic：

• Check whether there are multiple levels of nested loops (3 levels or more)
• Check whether the loop dimensions are in a contiguous memory layout (e.g., H×W of NCHW)
• Check whether there are dependencies between dimensions
• If there are multiple levels of loops and the dimensions are contiguous with no dependencies → involved
• If there are few loop levels, or there are dependencies between dimensions → not involved, skip

Hit condition: the code contains multiple levels of nested loops processing contiguous dimensions, and they can be merged

Reference document：references/dimension-merge.md

---

Optimization point 9: Libdevice function usage

Applicable condition: the code contains manually implemented mathematical functions, while tl.extra.cann.libdevice already has optimized versions

Typical code characteristics：

# 手动实现 round
return (x + 0.5).to(tl.int8)

# 手动实现 relu
out = tl.maximum(x, 0.0)

# 手动实现 tanh、sinh、pow 等数学函数

Decision logic：

• Check whether the following functions are manually implemented in the code: round, trunc, relu, tanh, sinh, cosh, pow, atan, acos, asin, expm1, log1p, hypot, etc.
• If there are manual implementations and tl.extra.cann.libdevice has corresponding functions → involved
• If there are no mathematical function implementations in the code, or the libdevice versions are already used → not involved, skip

Hit condition: the code contains manually implemented mathematical functions, and libdevice has optimized versions

Reference document：references/libdevice-usage.md

---

Optimization point 10: Loop invariant hoisting

Applicable condition: the code contains nested loops, and the inner loop contains a tl.load

Typical code characteristics：

# 问题代码：内层循环重复加载相同值
for outer_idx in range(outer_size):
    for inner_idx in range(inner_size):
        param_idx = outer_idx  # 只依赖外层变量
        val = tl.load(param_ptr + param_idx)  # 重复加载相同值
        ...

# 或者通过整除映射到更粗粒度
for block in range(num_blocks):
    offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    channel = offsets // spatial_size
    w = tl.load(weight_ptr + channel)  # 相同 channel 重复加载

Decision logic：

• Check whether a nested loop structure exists
• Check whether there is a tl.load(param_ptr + index_expr)
• Check whether index_expr depends only on the outer loop variable, not on the inner loop variable
• If it exists and the inner loop iteration count >> the outer loop iteration count → involved
• If there are no nested loops, or all loads depend on the inner variable → not involved, skip

Hit condition：There are nested loops in the code, and the inner loop contains variables that depend only on the outer loop tl.load

Reference documentation：references/loop-invariant-hoisting.md

---

Optimization point 11: Load instruction reordering

Applicable conditions：There are loops in the code, and within the loops there are multiple tl.load and tl.store, and there is blocking caused by data dependencies

Typical code characteristics：

for i in range(HEAD_NUM):
    # load B 在前，会等待上一次循环的 store B
    idx_B = tl.load(p_B_index)
    b_B = tl.load(p_B)
    
    # load A 在后，必须等 load B 完成
    b_A = tl.load(p_A)
    
    # calculation
    b_O = b_A * b_B
    
    # store
    tl.store(p_O, b_O)
    tl.store(p_B, b_B)  # store B 会阻塞下一次循环的 load B

Judgment logic：

• Check whether a loop structure exists
• Check whether there are multiple tl.load and tl.store
• Check whether there is a case where load A and store B have no data dependency between them but are blocked by other dependencies
• If there are reorderable load instructions → involved
• If there is only one load in the loop, or all loads have dependency relationships → not involved, skip

Hit conditions：There are loops in the code, and there are load instructions that can be issued earlier through reordering

Reference documentation：references/load-order.md

---

Optimization point 12: Autotune automatic tuning

Applicable conditions：There are one or more tunable parameters in the code (for example BLOCK_SIZE, BLOCK_M, etc.), and these parameters have not been fully tuned; considering that other optimization points may introduce tunable hyperparameters, optimize this optimization point last

Typical code characteristics：

# 未使用 autotune，手动指定固定参数
@triton.jit
def kernel(..., BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    ...

# 调用时固定参数
kernel[grid](..., BLOCK_M=128, BLOCK_N=128)

Judgment logic：

• Check whether @triton.autotune decorator is already used
• Check whether there are multiple tunable tl.constexpr parameters
• If autotune is not used and there are tunable parameters → involved
• If autotune is already used → not involved, skip

Hit conditions：There are multiple tunable parameters in the code, and autotune is not used

Reference documentation：references/autotune.md

---

Optimization workflow

1. 按顺序检查优化点 1 → 2 → 3 → ... → 12
2. 对于当前优化点，先判断是否命中（代码特征满足 + 适用条件成立）：
   - 未命中 → 跳过，检查下一优化点
   - 命中 → 参考对应文档，应用优化策略
3. 应用优化后，必须加载 references/checklist.md 检查代码规范
4. 如果代码规范不满足 → 修改代码直到满足规范
5. 代码规范满足后 → 返回优化后的代码，回到1继续检查优化点

Important constraints：

• ⚠️ Only the optimization methods specified by this skill may be used; any optimization method beyond this skill is prohibited
• ⚠️ The "hit conditions" of an optimization point must be hit first before the reference documentation can be loaded; if not hit, skip
• Only one optimization point can be used per optimization iteration; multiple rounds of optimization are allowed, example:

  第一轮：检查 1→2→3→...，命中优化点 X，应用后验证
  第二轮：检查 1→2→...，命中优化点 Y，应用后验证
  第三轮：检查 1→2→...，命中优化点 Z，应用后验证
  ...
  直到所有优化点都不命中

• Only one document can be referenced at a time

Optimization verification rules

⚠️ Mandatory requirement: before performing any accuracy verification or performance verification, a checklist check must be executed first to ensure all code specifications are satisfied. The verification workflow is as follows:

1. Checklist check：Load references/checklist.md, and check item by item whether the code satisfies all specification requirements
2. Does not satisfy the specifications → modify the code until all specification requirements are satisfied, then re-execute the checklist check to confirm
3. After the specifications are satisfied → execute accuracy verification and performance verification

• Success：The optimized performance does not degrade (speedup ≥ 1.0); this optimization result is used as the baseline for the next optimization iteration
• Failure：The optimized performance degrades (speedup < 1.0); discard this optimization result and use the code before optimization as the baseline for the next optimization iteration

Reference material index

Document type	Document path
Input argument staticization optimization	references/constexpr_parameters.md
Tiling optimization	references/tiling_optimization.md
Core partitioning optimization	references/vector_core_partition.md
Discrete memory access optimization	references/discrete_memory_access.md
Scalar to Vector optimization	references/scalar_to_vector.md
Avoid vector API scalar fallback	references/avoid_scalar_lowering.md
Pass merging optimization	references/pass-merge.md
Dimension merging optimization	references/dimension-merge.md
Libdevice function usage	references/libdevice-usage.md
Loop invariant hoisting	references/loop-invariant-hoisting.md
Load instruction reordering	references/load-order.md
Autotune automatic tuning	references/autotune.md
Code specification check	references/checklist.md

Final step: Block Size Scaling

After all the above instruction-level optimization strategies have been fully applied,must load references/block_size_scaling.md, execute Block Size Scaling as the final optimization step.

Execution flow

1. Read code_file_path code
2. Analyze code characteristics, load the corresponding optimization strategy document
3. Apply optimization strategies, generate optimized code
4. Write the optimized code to output_path
5. Return optimization description and status information

Key constraints

• Must produce a code file: even if there are no optimization points, still write out output_path
• Functional consistency: the computation results before and after optimization must be consistent
• Precision consistency: numerical precision must not be reduced
• Executability: the output code must be directly verifiable and benchmarkable by triton-op-verifier