企业官网建设流程全解析

CANN ops-math SliceLastDim 算子深度解析：最后一维切片原理、参数配置与图模式调用实战

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

导读

本文以 conversion/slice_last_dim/README.md 为核心，系统讲解 CANN ops-math 数学算子库中 SliceLastDim 算子的完整技术细节：它专用于对输入张量的最后一维执行切片（支持 start/end 指定范围、stride 指定步长），是处理连续内存布局下"尾维裁剪"场景的高效基础算子。读完本文，你将掌握该算子的支持范围、参数语义、shape 推导公式、底层 Tiling 与 Kernel 实现原理，以及基于 GE IR 图模式构建算子图的完整可运行样例代码。

产品支持情况

SliceLastDim 算子在不同昇腾产品上的支持情况如下表所示（以仓库 README.md 为准）：

需要说明的是，该算子的 kernel 与 tiling 实现位于arch35目录（对应 Ascend 950 系列等新架构），而Atlas 推理系列产品（非 A2/A3 系列）暂不支持，使用时请先核对目标设备型号。

功能说明

SliceLastDim 的核心语义非常聚焦：仅对输入张量的最后一维进行切片，其余维度原样透传。切片范围由start、end指定，采样步长由stride指定，其行为与 NumPy 的x[..., start:end:stride]高度一致。

计算示例

仓库 README 给出如下示例：

• 输入xshape =(2, 3, 10)，start = 0，end = 10，stride = 2
• 输出yshape =(2, 3, 5)
• 元素对应关系：y[i, j, k] = x[i, j, start + k * stride] = x[i, j, 2 * k]

即在最后一维上以步长 2 抽取 10 个元素中的 5 个，前两维(2, 3)保持不变。

计算公式

输出y的 shape 与输入x相同，仅最后一维发生变化：

y.shape[-1] = ceil((end - start) / stride) y[..., k] = x[..., start + k * stride]

关于start/end的语义，README 明确了两点关键规则：

1. 支持负数索引：负数时相对最后一维长度进行折算（-1表示最后一个元素）；
2. 自动截断：start被截断到[0, lastDim]，end同样被截断到[0, lastDim]，保证索引不会越界。

该规则在 slice_last_dim_infershape.cpp 与 slice_last_dim_tiling_arch35.cpp 中均有完整实现，Host 侧（infershape 与 tiling）对负索引的处理逻辑保持一致：先加lastDim转正，再做上下界截断。

参数说明

将 README 参数表与源码定义交叉核对，可以确认以下几点实现细节：

• 算子原型定义位于 slice_last_dim_proto.h，通过REG_OP(SliceLastDim)声明：x的 TensorType 为{DT_FLOAT16, DT_FLOAT, DT_DOUBLE, DT_INT8, DT_INT16, DT_INT32, DT_INT64}，start、end为REQUIRED_ATTR，stride为带默认值 1 的ATTR；
• 算子算子定义（OpDef）位于 slice_last_dim_def.cpp，这里登记的输入数据类型为{INT8, INT16, INT32, INT64, FLOAT16, FLOAT, BF16}，数据格式均为ND，且显式声明了DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)等能力标志，与 README"支持动态 Shape、动态 Rank、动态编译静态"的约束说明一一对应；
• Format 固定为 ND：README 要求输入输出均为 ND 格式，这与 OpDef 中baseFormat全部为FORMAT_ND一致，说明该算子不参与 Format 转换，仅在标准连续内存布局上工作。

约束说明

使用 SliceLastDim 时需满足以下约束（来自 README.md）：

• 输入张量x的维度范围为[1, 8]，该约束在 tiling 侧同样有校验：见 slice_last_dim_tiling_arch35.cpp，当rank_不在[1, MAX_RANK]（MAX_RANK 为 8）时直接返回GRAPH_FAILED并记录错误日志；
• stride必须 ≥ 1，tiling 侧在 GetParams 中同样做了强校验；
• start和end支持负数索引，负数时加上最后一维长度转换为正索引；start截断到[0, lastDim]，end截断到[0, lastDim]；
• 当end <= start时，输出最后一维长度为 0（此时整个算子退化为空输出，Kernel 的Process()会直接提前返回，见 slice_last_dim.h）；
• 支持动态 Shape、动态 Rank、动态编译静态。

其中"动态编译静态"能力在 OpDef 中通过aicoreConfig.DynamicCompileStaticFlag(true)开启，配合 Tiling 阶段通过GET_TPL_TILING_KEY生成的模板参数（详见下文），使得同一份编译产物可以覆盖运行时才知道的具体 shape。

源码级原理剖析

Shape 推导（InferShape）

算子输出 shape 的推导逻辑集中在 slice_last_dim_infershape.cpp：

1. 复制输入x的前rank - 1维到输出y；
2. 对最后一维：先处理负索引（s += lastDim、e += lastDim），再做截断（s < 0归零、e > lastDim置为lastDim）；
3. 计算切片长度length = (e > s) ? (e - s + stride - 1) / stride : 0，即向上取整的ceil((end - start) / stride)；
4. 若最后一维为未知维度（动态 shape，值为UNKNOWN_DIM = -1），则输出最后一维也保持未知。

InferDataType则将输入数据类型原样传递给输出。两者通过IMPL_OP_INFERSHAPE(SliceLastDim)注册到算子实现注册表。

Tiling 策略（Host 侧）

Tiling 是决定算子在 NPU 上执行效率的关键环节，实现在 slice_last_dim_tiling_arch35.cpp 中，整体流程为GetSocInfo → GetParams → ComputeAndSetTiling → SetTilingData：

• GetSocInfo：读取 AIV 核数coreNum_、UB 内存大小、CacheLine 大小与 UB Block 大小等 SoC 硬件参数；
• GetParams：解析输入 shape、start/end/stride属性，做负索引归一化与截断，并计算lastDimOut = ceil(end - start, stride)、outerSize（前rank-1维的乘积）；同时根据stride == 1与否设置copyMode（0 表示连续拷贝、1 表示跨步拷贝）；
• DoCacheLineTiling / DoTilingSearch：以 UB 缓冲区（默认上限64 * 1024字节、双缓冲BUFFER_NUM = 2）为约束，从"最后一维"与"外维"两个候选切分轴中选择能最大化核利用率的ubAxis与ubFactor，并保证实际使用核数不低于总核数的 80%（MIN_CORE_UTIL_RATIO = 0.8）；
• SetTilingData：将outerSize、lastDimIn、lastDimOut、start、stride、totalCount、perCoreCount、ubFactor、bufferSize写入 SliceLastDimTilingData 结构体，并通过context_->SetTilingKey(GET_TPL_TILING_KEY(copyMode_, ubAxis_))与context_->SetBlockDim(realCoreNum)设置编译模板参数与实际启用的核数。

copyMode与ubAxis的候选组合在 slice_last_dim_tiling_key.h 中通过ASCENDC_TPL_ARGS_DECL声明：copyMode ∈ {连续(0), 跨步(1)}，ubAxis ∈ {0, 1}，共 4 种模板实例。

Kernel 实现（Device 侧）

Kernel 入口 slice_last_dim_apt.cpp 是一个模板化的__global__ __aicore__函数，模板参数为<DTYPE_X, copyMode, ubAxis>。核心计算类 SliceLastDim 采用标准的CopyIn → 同步 → CopyOut → 同步流水线模型：

• 双缓冲：UB 缓冲区按bufferSize * BUFFER_NUM分配，Process()中通过(i - beginIdx) % BUFFER_NUM轮换使用两个缓冲区，配合MTE2_MTE3/MTE3_MTE2硬件事件同步实现数据搬移与计算的重叠；
• 连续模式（CopyMode = 0）：当stride == 1时，直接使用DataCopyPad按"块长度 × 块数 + 源/目的 stride"的参数形式批量搬移整行连续数据；
• 跨步模式（CopyMode = 1）：当stride > 1时，使用NdDmaDci + NdDmaLoopInfo描述循环搬移，loopSrcStride[0] = stride、loopDstStride[0] = 1即表达"源端每步跳 stride 个元素、目的端连续写入"的语义；
• 空输出保护：outerSize == 0 || lastDimOut == 0时直接返回，避免对空张量发起无意义的 DMA。

图模式调用实战

README 中该算子的官方调用方式为图模式调用：通过 算子 IR 构图后交给 GE 执行。完整可运行的样例位于 examples/test_geir_slice_last_dim.cpp，其核心流程如下。

1. 构建算子节点

auto node = op::SliceLastDim("slice_last_dim_1"); vector<int64_t> xShape = {4, 8}; TensorDesc xDesc(ge::Shape(xShape), FORMAT_ND, inDtype); xDesc.SetPlacement(ge::kPlacementHost); xDesc.SetFormat(FORMAT_ND); // 构造输入数据 Data 节点 auto xData = op::Data("x_data").set_attr_index(0); xData.update_input_desc_x(xDesc); xData.update_output_desc_y(xDesc); node.set_input_x(xData); node.SetAttr("start", static_cast<int64_t>(0)); node.SetAttr("end", static_cast<int64_t>(4)); node.SetAttr("stride", static_cast<int64_t>(1)); vector<int64_t> yShape = {4, 4}; TensorDesc yDesc(ge::Shape(yShape), FORMAT_ND, inDtype); node.update_output_desc_y(yDesc); graph.AddOp(xData); graph.AddOp(node);

样例中以x[4, 8], start=0, end=4, stride=1 → y[4, 4]为例（见文件头注释），属性通过node.SetAttr(...)显式设置，输出 shape 由用户预先声明。

2. 初始化 GE 并执行图

std::map<AscendString, AscendString> global_options = { {"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}}; Status ret = ge::GEInitialize(global_options); // ... 校验 ret ... Graph graph("tc_ge_irrun_test"); // ... CreateOppInGraph：构图并填充 inputs / outputs ... graph.SetInputs(inputs).SetOutputs(outputs); ge::Session* session = new Session(build_options); uint32_t graph_id = 0; session->AddGraph(graph_id, graph, graph_options); vector<Tensor> output; ret = session->RunGraph(graph_id, input, output); // ... 校验 ret ... ge::GEFinalize();

完整流程为：GEInitialize（配置设备号与运行模式）→ 构建Graph→Session::AddGraph注册图 →Session::RunGraph执行 →GEFinalize收尾。

3. 数据落盘

执行完毕后，样例通过ProcessInputData/ProcessOutputData将输入、输出张量按原始字节写为二进制文件，便于后续比对验证：

./tc_ge_irrun_test_0008_npu_input_0.bin // 输入数据 ./tc_ge_irrun_test_0008_npu_output_0.bin // 输出数据

其中输入数据由GenData按value + i（从 1.0 起递增）填充，支持的 dtype 包括 FLOAT、INT64、INT32、INT16、INT8 以及各类 UINT 类型，覆盖了 README 参数表中声明的多数数据类型。

单元测试验证

仓库在 tests/ut/op_host/test_slice_last_dim_infershape.cpp 中提供了 5 组 InferShape 单元测试，可直接印证 README 中的功能语义：

其中slice_last_dim_3d_with_stride与 README 的计算示例完全一致（(2,3,10) + stride=2 → (2,3,5)），slice_last_dim_dynamic验证了动态 shape 场景下未知维度（-1）的透传行为。此外，Tiling 侧的单元测试位于 tests/ut/op_host/arch35/test_slice_last_dim_tiling.cpp，用于验证切分策略在 Ascend 950 系列架构上的正确性。

总结

SliceLastDim 是 CANN ops-math 中一个"小而精"的基础算子，其设计目标明确：在 ND 连续布局下高效完成最后一维的定长/跨步切片。通过本文可以看到，一个看似简单的切片语义在工程实现上覆盖了算子原型声明（Proto/OpDef）、shape 推导（InferShape）、硬件感知的 Tiling 切分（含模板参数枚举）、双缓冲 DMA Kernel 流水线以及完整的 GE IR 图模式调用样例与单元测试。开发者若需在昇腾 NPU 上实现"尾维裁剪"类操作，可直接复用该算子，或参考其实现模式完成同类算子的开发。

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math