自定义算子¶
概述¶
AscendNPU IR已为上游模型支持丰富的算子集合。但部分业务场景下用户仍需自定义算子实现专属计算逻辑，典型场景包括：
现有算子的组合无法满足所需的计算需求。
厂商希望自定义算子保持私有。
多个算子的组合无法达到最优性能。
自定义算子允许用户自由使用AscendNPU IR提供的接口，提供能与其他算子一起编译的自有算子。
接口说明¶
参数说明：
参数 |
描述 |
|---|---|
|
唯一算子名称 |
|
输入参数 |
|
输出结果，可指定为 |
属性说明：
属性名 |
说明 |
备注/示例 |
|---|---|---|
|
在哪种核类型上执行 |
参见 |
|
在哪个Pipe上执行（用于 |
参见 |
|
Macro自定义算子的输入Pipe |
参见 |
|
Macro自定义算子的输出Pipe |
参见 |
|
在向量单元上的运行模式 |
参见 |
|
实现函数的名称 |
- |
|
Macro自定义算子的同步槽位元信息。 |
列表长度必须与macro实现体内、同一Pipe对上的 |
|
逐操作数的迭代器语义，供结构化降级与flatten类Pass使用（ |
可选；设置时长度应覆盖参与tiling的输入与输出 |
|
逐操作数的仿射索引映射（与Linalg的 |
可选，类型为 |
|
flatten/布局类Pass支持的最大张量秩 |
可选 |
|
逐操作数的维度对齐提示 |
通过 |
|
逐操作数的字典属性数组 |
类型为 |
|
临时buffer的元素类型 |
由 |
|
临时buffer的一维大小（元素个数） |
同上 |
|
传给设备实现的临时 |
通常由 |
|
表示算子无副作用 |
- |
|
实现产物路径及可选编译命令 |
通常由Triton前端设置 |
降级流程¶
┌─────────────────────────────────────────────────────────────────┐
│ CustomOp │
│ hivm.hir.custom "name" { attrs... } ins(..) outs(...) │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ HIVMToStandard │
│ ───────────────────────────────────────────────────────────────│
│ • 内置算子 │
│ -> 调用内置库 │
│ • 用户提供的实现 -> │
| -> 调用用户提供的函数名 |
| -> bishengir-compile 使用用户提供的链接命令进行链接 |
└─────────────────────────────────────────────────────────────────┘
│
▼
毕昇编译器将其编译为 Object Files
支持能力¶
特性 |
说明 |
|---|---|
CoreType |
自定义算子执行核 |
Pipe |
自定义算子执行pipe（ |
InPipe / OutPipe |
Macro自定义算子的输入/输出pipe（ |
VFMode |
自定义算子在向量核上的运行模式：SIMT/SIMD/MIX |
Symbol |
使用者提供的函数名称 |
sync_event_slots |
Macro同步槽位声明，用于GraphSyncSolver集成 |
iterator_types |
Tiling / flatten的逐操作数迭代器语义 |
indexing_map |
结构化降级的逐操作数仿射映射 |
max_rank |
布局Pass支持的最大张量秩（默认5） |
align_dim / arg_attrs |
对齐调整Pass的逐操作数对齐提示 |
extra_buffers_* / tmps |
临时buffer声明与分配 |
no_side_effect |
纯算子标记，便于优化 |
内置算子 |
一组内置算子（名称预留） |
限制约束¶
当前自定义算子体系存在两项待完善能力，具体如下：
用户侧实现适配不完善
支持将自定义算子降级为用户提供的实现，相关流程包含：HIVM IR链接到用户提供的源码或目标文件、向bishengir-compile注册特定链接命令，目前该功能仍在开发迭代中。
Pass交互适配不完善
各类自定义算子的变换Pass，包含Flatten优化、对齐调整、内存规划、布局变换等，仍有多项适配逻辑待补充，相关工作正在推进。
MLIR示例¶
算子声明示例¶
内置算子：
%0 = hivm.hir.custom
"__builtin_gather_load"
ins(%arg0, %arg1, %c4_i64, %c0_i32, %c2_i64, %c1_i64, %c2_i32, %c2_i32, %c0_i32, %c0_i32
: memref<?xf32>, tensor<3x3xi64>, i64, i32, i64, i64, i32, i32, i32, i32)
outs(%empty : tensor<3x3xf32>) -> tensor<3x3xf32>
自定义算子：
%0 = hivm.hir.custom
{ hivm.tcore_type = #hivm.tcore_type<VECTOR>, hivm.pipe = #hivm.pipe<PIPE_V>, hivm.vf_mode = #hivm.vf_mode<SIMD>,
symbol = "my_custom" }
"my_custom_op"
ins(%arg0, %arg1, %c4_i64, %c0_i32, %c2_i64, %c1_i64, %c2_i32, %c2_i32, %c0_i32, %c0_i32
: memref<?xf32>, tensor<3x3xi64>, i64, i32, i64, i64, i32, i32, i32, i32)
outs(%empty : tensor<3x3xf32>) -> tensor<3x3xf32>
自定义Macro算子：
%0 = hivm.hir.custom_macro
{ hivm.tcore_type = #hivm.tcore_type<VECTOR>,
hivm.pipe_in = #hivm.pipe<PIPE_MTE2>,
hivm.pipe_out = #hivm.pipe<PIPE_V>,
hivm.vf_mode = #hivm.vf_mode<SIMD>,
symbol = "custom_macro_add_int32",
sync_event_slots = [
#hivm.sync_event_slot<#hivm.pipe<PIPE_MTE2>, #hivm.pipe<PIPE_V>>
] }
"macro_add"
ins(%arg0, %arg1 : memref<?xi32>, memref<?xi32>)
outs(%dst : memref<32xi32, #hivm.address_space<ub>>)
核心属性用法示例¶
Macro算子同步槽位（sync_event_slots）¶
对于在InPipe（如PIPE_MTE2）搬运数据，在OutPipe（如PIPE_V）执行计算的Macro自定义算子，通常需要在设备实现中插入set_flag与wait_flag操作，确保MTE2将数据写入UB完成后，向量Pipe再开始读取。
在Triton侧通过sync_event_slots与SYNC_HINT.INTERNAL声明这些宏内部的同步关系。
GraphSyncSolver会为每个槽位分配独立的event ID，并在需要时在macro算子前后补充set/wait操作，保证其与上下游HIVM算子正确衔接。
规则：macro实现体内，每一对同方向Pipe的set_flag/wait_flag，需要在Python/MLIR中对应声明一个sync_event_slots条目。
模式 |
Python / MLIR槽位数 |
典型设备代码 |
适用场景 |
|---|---|---|---|
单槽位 |
1 × |
两次GM→UB搬运完成后，一次 |
默认场景：两次load连续执行，向量计算开始前只需一次MTE2到V的同步 |
双槽位 |
2 × |
|
每次GM→UB需要独立event的场景，如流水重叠、或需按次确认搬运完成 |
单槽位说明：
两个操作数都通过MTE2搬入UB后，向量Pipe才开始计算。单个event即可标识本macro所需的MTE2输入搬运已全部完成，PIPE_V可安全读取两块UB暂存空间，对应伪代码如下：
load_gm_to_ubuf(src0)
load_gm_to_ubuf(src1)
set_flag(MTE2 → V, event 0)
wait_flag(MTE2 → V, event 0)
vector_vadd(...)
双槽位说明：
第一次load后立即执行set_flag，可在等待首次搬运完成的同时启动第二次GM到UB的搬运。每个槽位对应GraphSyncSolver分配的一个独立event ID，对应伪代码如下：
load_gm_to_ubuf(src0)
set_flag(MTE2 → V, event 0)
load_gm_to_ubuf(src1)
set_flag(MTE2 → V, event 1)
wait_flag(MTE2 → V, event 0)
wait_flag(MTE2 → V, event 1)
vector_vadd(...)
双槽位对应的MLIR：
sync_event_slots = [
#hivm.sync_event_slot<#hivm.pipe<PIPE_MTE2>, #hivm.pipe<PIPE_V>>,
#hivm.sync_event_slot<#hivm.pipe<PIPE_MTE2>, #hivm.pipe<PIPE_V>>
]
若Python/MLIR中声明的槽位数量，与C++ 实现中set_flag/wait_flag的配对数量不一致，会导致GraphSyncSolver分配的event ID与kernel不匹配，这是自定义Macro算子集成时的常见错误。
Tiling属性（iterator_types、indexing_map、max_rank）¶
通过配置这些属性，CustomOp可实现HIVM结构化算子接口（getIteratorTypesArray、getIndexingMaps），使flatten、broadcast及布局Pass能够像处理原生结构化算子一样，对自定义算子的操作数进行分析与优化。
iterator_types支持的取值包括：parallel、broadcast、transpose、reduction、interleave、deinterleave、inverse、pad、concat、gather、cumulative、opaque。
MLIR示例（摘自custom-op-attrs.mlir）：
#map2d = affine_map<(d0, d1) -> (d0, d1)>
%0 = hivm.hir.custom
{hivm.tcore_type = #hivm.tcore_type<VECTOR>,
hivm.pipe = #hivm.pipe<PIPE_V>,
hivm.vf_mode = #hivm.vf_mode<SIMD>,
symbol = "k_named_maps",
max_rank = 5 : i64,
iterator_types = [#hivm.iterator_type<parallel>, #hivm.iterator_type<parallel>],
indexing_map = [#map2d, #map2d, #map2d]}
"user.named_maps"
ins(%arg0, %arg1 : memref<2x2xf32>, tensor<2x2xf32>)
outs(%empty : tensor<2x2xf32>) -> tensor<2x2xf32>
Triton注册示例：
@al.register_custom_op
class tiled_custom_op:
core = al.CORE.VECTOR
pipe = al.PIPE.PIPE_V
mode = al.MODE.SIMT
symbol = "my_tiled_func"
bitcode = "/path/to/kernel.bc"
iterator_types = [
al.IteratorType.Parallel,
al.IteratorType.Broadcast,
]
def __init__(self, x, y, out=None):
# 每个结构化操作数（输入 + 输出）对应一个仿射映射。
self.indexing_map = [
al.affine_map.get_identity(2),
al.affine_map.get_identity(2),
al.affine_map.get_identity(2),
]
操作数对齐属性（align_dim、arg_attrs）¶
align_dim用于标记指定操作数需要对齐的维度。在MLIR中，该属性通过arg_attrs挂载到对应索引的操作数上，示例如下：
%0 = hivm.hir.custom { ... }
"my_custom_op"
ins(%arg0 {align_dim = 1 : i64}, %arg1 {align_dim = 0 : i64} : memref<?xf32>, memref<?xf32>)
outs(%dst : tensor<?xf32>) -> tensor<?xf32>
Triton示例（在__init__中设置，键为参数名或位置索引）：
def __init__(self, x, ptr1, ptr2, out=None):
self.align_dim = {"ptr2": 1, 1: 0} # ptr2 第 1 维；第 2 个参数第 0 维
临时缓冲区属性（extra_buffers_*、temp_buffers）¶
通过相关属性可在算子上声明临时缓冲区的类型与大小，hivm-alloc-extra-buffer会自动分配一维memref，并接入算子的tmps操作数段，示例如下：
// 经 hivm-alloc-extra-buffer 之后：
%alloc0 = memref.alloc() : memref<512xf32>
%alloc1 = memref.alloc() : memref<128xf16>
%0 = hivm.hir.custom
{ extra_buffers_types = [f32, f16],
extra_buffers_sizes = [512 : i64, 128 : i64], ... }
"my_custom_op"
ins(%arg0 : memref<2x2xf32>)
tmps(%alloc0, %alloc1 : memref<512xf32>, memref<128xf16>)
outs(%empty : tensor<2x2xf32>) -> tensor<2x2xf32>
Triton示例（单个buffer用元组，多个用(dtype, size)列表）：
@al.register_custom_op
class my_custom_op_extra_buf:
core = al.CORE.VECTOR
pipe = al.PIPE.PIPE_V
mode = al.MODE.SIMT
symbol = "my_extra_buf_func"
bitcode = "/path/to/kernel.bc"
def __init__(self, x, out=None):
self.extra_buffers = [
(tl.bfloat16, 256),
(tl.float32, 512),
]
# 或单个 buffer：self.extra_buffers = (tl.float16, 128)
设备kernel会按照extra_buffers的声明顺序，通过tmps操作数接收分配的scratch memref。
no_side_effect¶
该属性用于标记仅读取输入、仅写入声明输出的纯算子，示例如下：
%0 = hivm.hir.custom {no_side_effect, symbol = "pure_kernel", ...}
"pure_op" ins(...) outs(...) -> ...
Triton自定义算子示例¶
标准自定义算子示例¶
Python脚本：test_custom_op.py
# 更多关于 Triton 自定义算子设计的详情，请参考
# https://gitcode.com/Ascend/triton-ascend/pull/988
import triton
import triton.language as tl
import triton.language.extra.cann.extension as al
import torch
import torch_npu
import pytest
def torch_add(a, b):
return a + b
@al.register_custom_op
class add:
core = al.CORE.VECTOR
pipe = al.PIPE.PIPE_V
mode = al.MODE.SIMD
def __init__(self, a, b, out=None):
assert out, "out is required"
self.symbol = "custom_add_" + str(a.dtype)
self.bitcode = 'add.bc'
@triton.jit
def triton_custom_add(
output_ptr,
a_ptr,
b_ptr,
L: tl.constexpr
):
idx = tl.arange(0, L)
a = tl.load(a_ptr + idx)
b = tl.load(b_ptr + idx)
buf = tl.full([L], 0, a.dtype)
res = al.custom("add", a, b, out=buf)
tl.store(output_ptr + idx, res)
testlist = [
(32)
]
typelist = [torch.int32]
@pytest.mark.parametrize("DT", typelist)
@pytest.mark.parametrize("L", testlist)
def test_custom(DT, L):
a = torch.ones(L, dtype=DT).npu()
b = torch.ones(L, dtype=DT).npu()
ref = torch_add(a, b)
out = torch.zeros(L, dtype=DT).npu()
triton_custom_add[1, 1, 1](out, a, b, L)
torch.testing.assert_close(out, ref)
C++ API定义：add.cpp
#define __aiv__ [aicore]
#define INTRINSIC_NO_ARGS(NAME) NAME()
#define INTRINSIC(NAME, ...) NAME(__VA_ARGS__)
template <typename T, size_t Dim>
struct memref_t {
T *allocated;
T *aligned;
int64_t offset;
int64_t sizes[Dim];
int64_t strides[Dim];
};
template <size_t OPERANUM, typename SRC_T, typename DST_T = SRC_T>
struct intrin_args {
__ubuf__ DST_T *dst;
__ubuf__ SRC_T *src[OPERANUM];
SRC_T scalar;
uint64_t repeat;
uint16_t dst_block_stride;
uint16_t src_block_stride[OPERANUM];
uint16_t dst_repeat_stride;
uint16_t src_repeat_stride[OPERANUM];
};
template <typename SRC_TYPE, typename DST_TYPE = SRC_TYPE>
__aiv__ __attribute__((always_inline)) void
vector_eltwise_vadd_intrin(intrin_args<2, SRC_TYPE, DST_TYPE> args) {
#define ELTWISE_VV_ARGS \
args.dst, args.src[0], args.src[1], args.repeat, args.dst_block_stride, \
args.src_block_stride[0], args.src_block_stride[1], \
args.dst_repeat_stride, args.src_repeat_stride[0], \
args.src_repeat_stride[1]
INTRINSIC(vadd, ELTWISE_VV_ARGS);
}
extern "C" {
__aiv__ __attribute__((always_inline)) void _mlir_ciface_custom_add_int32(
memref_t<__ubuf__ int32_t, 1> *src0, memref_t<__ubuf__ int32_t, 1> *src1,
memref_t<__ubuf__ int32_t, 1> *dst) {
uint16_t src0_block_stride = 1;
uint16_t src1_block_stride = 1;
uint16_t src0_repeat_stride = 8;
uint16_t src1_repeat_stride = 8;
auto new_src0_ptr = src0->aligned + src0->offset;
auto new_src1_ptr = src1->aligned + src1->offset;
auto dst_ptr = dst->aligned + dst->offset;
INTRINSIC_NO_ARGS(set_mask_count);
const int64_t n = dst->sizes[0];
INTRINSIC(set_vector_mask, 0, n);
vector_eltwise_vadd_intrin<int32_t>(
intrin_args<2, int32_t>{dst_ptr,
{new_src0_ptr, new_src1_ptr},
0,
1,
1,
{src0_block_stride, src1_block_stride},
8,
{src0_repeat_stride, src1_repeat_stride}});
INTRINSIC_NO_ARGS(set_mask_norm);
}
}
编译.bc文件指令：
ccec -x cce --cce-aicore-arch=dav-c220-vec --cce-aicore-only -c -emit-llvm ./add.cpp -o ./add.bc
Python脚本执行命令：
python -m pytest -sv test_custom_op.py
Macro自定义算子示例¶
Python脚本：test_custom_macro_op.py
import triton
import triton.language as tl
import triton.language.extra.cann.extension as al
import torch
import torch_npu
import pytest
@al.register_custom_op
class macro_add:
core = al.CORE.VECTOR
pipe = (al.PIPE.PIPE_MTE2, al.PIPE.PIPE_V)
mode = al.MODE.SIMD
sync_event_slots = [
(al.PIPE.PIPE_MTE2, al.PIPE.PIPE_V, al.SYNC_HINT.INTERNAL),
]
def __init__(self, a, b, out=None):
assert out is not None, "out is required"
self.symbol = "custom_macro_add_int32"
self.bitcode = "macro_add.bc"
@triton.jit
def triton_custom_macro_add(output_ptr, a_ptr, b_ptr, L: tl.constexpr):
idx = tl.arange(0, L)
a = tl.load(a_ptr + idx)
b = tl.load(b_ptr + idx)
buf = tl.full([L], 0, tl.int32)
res = al.custom("macro_add", a, b, out=buf)
tl.store(output_ptr + idx, res)
@pytest.mark.parametrize("L", [32])
def test_custom_macro(L):
a = torch.ones(L, dtype=torch.int32).npu()
b = torch.ones(L, dtype=torch.int32).npu()
out = torch.zeros(L, dtype=torch.int32).npu()
triton_custom_macro_add[(1,)](out, a, b, L)
torch.testing.assert_close(out.cpu(), (a + b).cpu())
C++ API定义：macro_add.cpp
#define __aiv__ [aicore]
#define INTRINSIC_NO_ARGS(NAME) NAME()
#define INTRINSIC(NAME, ...) NAME(__VA_ARGS__)
constexpr int64_t UB_SCRATCH_SRC0_BYTES = 256;
constexpr int64_t UB_SCRATCH_SRC1_BYTES = 384;
constexpr int64_t kGmToUbChunkElems = 16;
template <typename T, size_t Dim>
struct memref_t {
T *allocated;
T *aligned;
int64_t offset;
int64_t sizes[Dim];
int64_t strides[Dim];
};
extern "C" {
__aiv__ __attribute__((always_inline)) void _mlir_ciface_custom_macro_add_int32(
memref_t<__gm__ int32_t, 1> *src0, memref_t<__gm__ int32_t, 1> *src1,
memref_t<__ubuf__ int32_t, 1> *dst) {
const int64_t n = dst->sizes[0];
auto ub_src0 = ub_scratch_at_bytes(UB_SCRATCH_SRC0_BYTES, n);
auto ub_src1 = ub_scratch_at_bytes(UB_SCRATCH_SRC1_BYTES, n);
memref_t<__gm__ int32_t, 1> gm_src0 = {src0->allocated, src0->aligned,
src0->offset, {n}, {1}};
memref_t<__gm__ int32_t, 1> gm_src1 = {src1->allocated, src1->aligned,
src1->offset, {n}, {1}};
// 单内部 event：两次 GM→UB 后一次 MTE2→V 握手。
load_gm_to_ubuf_1d(&gm_src0, &ub_src0);
load_gm_to_ubuf_1d(&gm_src1, &ub_src1);
INTRINSIC(set_flag, PIPE_MTE2, PIPE_V, 0);
INTRINSIC(wait_flag, PIPE_MTE2, PIPE_V, 0);
vector_vadd_ub(&ub_src0, &ub_src1, dst);
}
}
编译.bc文件指令：
ccec -x cce --cce-aicore-arch=dav-c220-vec --cce-aicore-only -c -emit-llvm ./macro_add.cpp -o ./macro_add.bc
Python脚本执行命令：
python -m pytest -sv test_custom_macro_op.py
降级为MLIR：
module attributes {hacc.target = #hacc.target<"Ascend910B3">} {
func.func @triton_custom_add(%arg0: memref<?xi8>, %arg1: memref<?xi8>, %arg2: memref<?xi32> {tt.divisibility = 16 : i32, tt.tensor_kind = 1 : i32}, %arg3: memref<?xi32> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg4: memref<?xi32> {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32}, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: i32) attributes {SyncBlockLockArgIdx = 0 : i64, WorkspaceArgIdx = 1 : i64, global_kernel = "local", mix_mode = "aiv", parallel_mode = "simd"} {
%c0_i32 = arith.constant 0 : i32
%0 = tensor.empty() : tensor<32xi32>
%1 = linalg.fill ins(%c0_i32 : i32) outs(%0 : tensor<32xi32>) -> tensor<32xi32>
%reinterpret_cast = memref.reinterpret_cast %arg3 to offset: [0], sizes: [32], strides: [1] : memref<?xi32> to memref<32xi32, strided<[1]>>
%alloc = memref.alloc() : memref<32xi32>
memref.copy %reinterpret_cast, %alloc : memref<32xi32, strided<[1]>> to memref<32xi32>
%2 = bufferization.to_tensor %alloc restrict writable : memref<32xi32>
%reinterpret_cast_0 = memref.reinterpret_cast %arg4 to offset: [0], sizes: [32], strides: [1] : memref<?xi32> to memref<32xi32, strided<[1]>>
%alloc_1 = memref.alloc() : memref<32xi32>
memref.copy %reinterpret_cast_0, %alloc_1 : memref<32xi32, strided<[1]>> to memref<32xi32>
%3 = bufferization.to_tensor %alloc_1 restrict writable : memref<32xi32>
%4 = hivm.hir.custom {bitcode = "/home/test/add.bc", hivm.pipe = #hivm.pipe<PIPE_V>, hivm.tcore_type = #hivm.tcore_type<VECTOR>, hivm.vf_mode = #hivm.vf_mode<SIMD>, symbol = "custom_add_int32"} "add" ins(%2, %3 : tensor<32xi32>, tensor<32xi32>) outs(%1 : tensor<32xi32>) -> tensor<32xi32>
%reinterpret_cast_2 = memref.reinterpret_cast %arg2 to offset: [0], sizes: [32], strides: [1] : memref<?xi32> to memref<32xi32, strided<[1]>>
bufferization.materialize_in_destination %4 in writable %reinterpret_cast_2 : (tensor<32xi32>, memref<32xi32, strided<[1]>>) -> ()
return
}
}