Kernel核函数#

Kernel是在AI Core上执行的函数，在函数签名中定义输入、输出和运行时参数，在函数体中组织数据搬运与计算。

定义Kernel#

使用JIT装饰器#

Kernel函数必须使用@pypto_pro.language.jit()装饰：

import pypto_pro.language as pl


@pl.jit(auto_mutex=True)
def add_kernel(
    x: pl.Tensor[[64, 64], pl.DT_FP16],
    y: pl.Tensor[[64, 64], pl.DT_FP16],
    out: pl.Tensor[[64, 64], pl.DT_FP16],
):
    ...

jit在首次启动Kernel时解析函数体并触发编译。JIT流程、编译签名和编译选项参考JIT编译。

声明Kernel参数#

Kernel参数需要通过类型标注明确数据类型和传递方式：

参数类型

用途

pypto_pro.language.Tensor

接收GM中的多维数据；在类型中声明shape、dtype和可选layout。

pypto_pro.language.Ptr

接收裸指针；通常使用pypto_pro.language.make_tensor构造Tensor视图。

pypto_pro.language.DT_*

接收整型、浮点型等运行时标量。

TilingData类

接收shape、stride、循环边界等结构化运行时参数。

Tensor参数适合直接使用调用侧Tensor的shape；Ptr参数适合由TilingData提供shape和stride：

from dataclasses import dataclass

import pypto_pro.language as pl


@dataclass
class AddTiling:
    rows: int
    cols: int


@pl.jit(auto_mutex=True)
def dynamic_kernel(
    x: pl.Ptr[pl.DT_FP16],
    out: pl.Ptr[pl.DT_FP16],
    scale: pl.DT_FP32,
    tiling: AddTiling,
):
    tensor_x = pl.make_tensor(x, [tiling.rows, tiling.cols])
    tensor_out = pl.make_tensor(out, [tiling.rows, tiling.cols])
    ...

TilingData必须位于Kernel形参列表和启动实参列表的末尾。完整字段和传输规则参考Tiling参数定义与传递。

Kernel不返回Python值。计算结果通过Tensor或Ptr对应的GM区域写回。

定义执行域#

使用pypto_pro.language.section_vector()和pypto_pro.language.section_cube()定义计算代码所在的执行域：

Kernel组成

执行方式

仅包含Vector执行域

启动Vector Kernel。

仅包含Cube执行域

启动Cube Kernel。

同时包含Cube和Vector执行域

启动混合Kernel。

 
import pypto_pro.language as pl


@pl.jit(auto_mutex=True)
def mixed_kernel(
    x: pl.Ptr[pl.DT_FP16],
    out: pl.Ptr[pl.DT_FP16],
):
    with pl.section_cube():
        # 矩阵计算代码
        ...

    with pl.section_vector():
        # 矢量计算代码
        ...

执行域决定可使用的指令、片上Buffer以及block_dim的含义。Vector计算参考Tile计算和Reg计算，矩阵计算参考Cube计算。

组织Kernel函数体#

Kernel函数体通常按照以下顺序组织：

  1. 定义Tile类型并绑定片上Buffer。

  2. 进入Vector或Cube执行域。

  3. 根据逻辑Block索引划分当前核的任务。

  4. 将数据从GM搬入片上Buffer。

  5. 执行矢量或矩阵计算。

  6. 将结果写回GM。

下面只展示Kernel结构，Tile创建和计算参数由相应章节说明：

import pypto_pro.language as pl


@pl.jit(auto_mutex=True)
def add_kernel(
    x: pl.Tensor[[64, 64], pl.DT_FP16],
    y: pl.Tensor[[64, 64], pl.DT_FP16],
    out: pl.Tensor[[64, 64], pl.DT_FP16],
):
    tile_type = pl.TileType(
        shape=[64, 64],
        dtype=pl.DT_FP16,
        target_memory=pl.MemorySpace.Vec,
    )
    tile_x = pl.make_tile_group(type=tile_type, addrs=0x0000, mutex_ids=[0])
    tile_y = pl.make_tile_group(type=tile_type, addrs=0x2000, mutex_ids=[1])
    tile_out = pl.make_tile_group(type=tile_type, addrs=0x4000, mutex_ids=[2])

    with pl.section_vector():
        cur_x = tile_x.current()
        cur_y = tile_y.current()
        cur_out = tile_out.current()
        pl.load(cur_x, x, [0, 0])
        pl.load(cur_y, y, [0, 0])
        pl.add(cur_out, cur_x, cur_y)
        pl.store(out, cur_out, [0, 0])

Tile声明、地址和TileGroup操作参考Tile创建和操作。auto_mutex=True会将通过pypto_pro.language.make_tile_group创建的Tile对象自动插入mutex同步，确保核内Tile的读写时序；需要显式同步的场景参考同步API。

调用Kernel#

Kernel使用方括号指定启动配置，使用圆括号传入函数实参：

调用形式

含义

kernel(args…)

使用当前Stream，block_dim=1。

kernel[block_dim](args…)

使用当前Stream，并指定实际核数。

kernel[stream, block_dim](args…)

指定Stream和实际核数。

kernel[stream, block_dim, tiling_key](args…)

选择TilingKey对应的编译实例。

kernel[stream, block_dim, tiling_key, datatype](args…)

同时选择TilingKey和datatype特化实例。

使用TilingKey或datatype特化时，必须通过方括号传入相应字典。 仅使用datatype特化时，datatype字典位于第三项；同时使用TilingKey和datatype时，两者分别位于第三项和第四项。

import os

import torch
import torch_npu


device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0))
device = f"npu:{device_id}"
torch.npu.set_device(device)

x = torch.rand(64, 64, device=device, dtype=torch.float16)
y = torch.rand(64, 64, device=device, dtype=torch.float16)
out = torch.empty_like(x)

add_kernel[None, 1](x, y, out)
torch.npu.synchronize()

Kernel启动相对于Host异步执行。在Host读取结果、检查精度或统计完整执行时间之前，需要同步对应的Stream。

stream的含义与设置#

stream指定Kernel下发的NPU执行流，用于维护一些异步操作的执行顺序，确保按照应用程序中的代码调用顺序在device上执行。传入None时使用当前Stream；显式传入Stream时，可以只等待该Stream上的任务：

import torch
import torch_npu


stream = torch.npu.Stream()
add_kernel[stream, num_cores](x, y, out)
stream.synchronize()

同一Stream中的任务按照下发顺序执行。不同Stream之间存在数据依赖时，需要通过Stream同步机制显式建立依赖。

blockDim的含义与设置#

block_dim表示实际可用核数，取值必须是正整数，且不得超过最大可用核数。Kernel可通过pypto_pro.language.get_block_num()读取实际生效的核数。

Kernel类型

block_dim的含义

Vector Kernel

用于设置启动多少个Vector（AIV）实例执行，比如某款AI处理器上有40个Vector核，建议设置为40。

Cube Kernel

用于设置启动多少个Cube（AIC）实例执行，比如某款AI处理器上有20个Cube核，建议设置为20。

Cube与Vector混合Kernel

按照AIV和AIC组合启动，用于设置启动多少个组合执行，比如某款AI处理器上有40个Vector核和20个Cube核，一个组合是2个Vector核和1个Cube核，建议设置为20，此时会启动20个组合，即40个Vector核和20个Cube核。此时Vector核数的统计还需要乘以pypto_pro.language.get_subblock_num()。

使用TilingKey和datatype#

TilingKey用于选择有限的编译期模式，datatype用于根据输入或输出数据类型生成专用实例：

import pypto_pro.language as pl


key = {"UseScale": 1, "BlockM": 128}
datatype = {"x": pl.DT_FP16, "out": pl.DT_FP16}

kernel[None, block_dim, key, datatype](x, out, tiling)

TilingKey的声明、编码和启动规则参考Tiling参数定义与传递。datatype字段由@pypto_pro.language.jit(datatype=...)声明，具体编译行为参考JIT编译。