ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

CANN ops-math 中 aclnnIm2col 算子接口详解:两段式调用、参数语义与 NPU 源码实现

CANN ops-math 中 aclnnIm2col 算子接口详解:两段式调用、参数语义与 NPU 源码实现 CANN ops-math 中 aclnnIm2col 算子接口详解两段式调用、参数语义与 NPU 源码实现【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读aclnnIm2col是 CANN ops-math 仓库中用于把三维/四维 NCHW 输入中的滑动窗口展开为列矩阵的算子 API是卷积计算如 im2col GEMM前数据重排的经典原语。本文以 experimental/conversion/im2col/docs/aclnnIm2col.md 为骨架完整讲解其功能语义、两段式调用方式、全部入参约束与返回码并结合experimental/conversion/im2col下的 op_api、op_host、op_kernel 与单测源码深入剖析其连续化处理、输出 shape 推导、tiling 多路径分派与 AscendC 内核实现。读完本文你将掌握如何在 Atlas A2/A3 系列产品上正确调用该接口并理解其底层执行链路。功能说明aclnnIm2col将三维[C, H, W]或四维[N, C, H, W]的 NCHW 输入按滑动窗口展开为列矩阵展开后每一列对应一个卷积核位置上的局部感受野数据常用于卷积计算前的数据重排为后续矩阵乘GEMM形式的卷积提供排布良好的输入。输出维度计算公式H、W 两个空间方向各自独立地按如下公式计算输出尺寸D 代表 H 或 W$$ outD \left\lfloor\frac{inD 2 \times paddingD - dilationD \times (kernelD - 1) - 1}{strideD}\right\rfloor 1 $$其中kernelD为卷积核尺寸dilationD为膨胀系数paddingD为单侧对称填充strideD为滑动步长。公式与常规卷积输出尺寸推导一致等效卷积核尺寸为dilationD * (kernelD - 1) 1分子为输入加两侧填充后减去等效核尺寸向下取整除以步长再加 1。输出 shape 语义三维输入[C, H, W]输出 shape 为[C * kernelH * kernelW, outH * outW]即“通道数 × 核面积”行、outH * outW列四维输入[N, C, H, W]输出 shape 为[N, C * kernelH * kernelW, outH * outW]batch 维被保留在最外层。产品支持情况该接口在当前仓库中的支持范围如下Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持。在源码层面算子定义 im2col_def.cpp 中通过this-AICore().AddConfig(ascend910b)与this-AICore().AddConfig(ascend910_93)注册了 Atlas A2Ascend 910B与 Atlas A3Ascend 910_93两个 AICore 配置与文档声明一致。注意目录 README 明确指出experimental/conversion/im2col的实现注册在 Atlas A2、Atlas A3 产品不复用仓库中 conversion/im2col 目录下的 Ascend 950Arch35实现。两段式调用模式aclnnIm2col采用 CANN 单算子 API 通用的两段式调用方式这与仓库文档 docs/zh/context/two_phase_api.md 描述的通用模式一致先调用第一段aclnnIm2colGetWorkspaceSize完成入参校验、输出 shape 推导与执行器构建并返回本次计算所需的 Device 侧 workspace 大小与aclOpExecutor执行器按返回的workspaceSize申请 NPU 内存再调用第二段aclnnIm2col真正下发计算。其中 workspace 是指除输入/输出外算子在 NPU 上完成计算所需的临时内存。需要注意第二段接口不可重复调用同一executor只应执行一次。函数原型aclnnStatus aclnnIm2colGetWorkspaceSize( const aclTensor *self, const aclIntArray *kernelSize, const aclIntArray *dilation, const aclIntArray *padding, const aclIntArray *stride, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnIm2col( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnIm2colGetWorkspaceSize 参数说明第一段接口的完整参数语义如下表所示。参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续Tensorselfconst aclTensor*输入待展开的输入Tensor。三维时按CHW解释四维时按NCHW解释。三维输入不支持空Tensor四维输入仅支持N为0的空Tensor。FLOAT16、FLOAT、BFLOAT16、BOOLND3或4√kernelSizeconst aclIntArray*输入卷积核大小[kernelH, kernelW]。长度为2所有元素必须大于0。INT64-2-dilationconst aclIntArray*输入卷积核膨胀系数[dilationH, dilationW]。长度为2所有元素必须大于0。INT64-2-paddingconst aclIntArray*输入H、W方向两侧的对称填充值[paddingH, paddingW]。长度为2所有元素必须大于等于0。INT64-2-strideconst aclIntArray*输入H、W方向的滑动步长[strideH, strideW]。长度为2所有元素必须大于0。INT64-2-outconst aclTensor*输出滑动窗口展开后的输出Tensor。数据类型必须与self相同shape必须与公式推导结果相同。与self一致NDself为三维时是2维self为四维时是3维√workspaceSizeuint64_t*输出返回Device侧workspace大小。不可为空指针。UINT64---executoraclOpExecutor**输出返回包含算子计算流程的执行器。不可为空指针。----参数校验的源码实现上述校验规则在 op_api/aclnn_im2col.cpp 的CheckParams中逐条落实空指针检查CheckNotNull对self、四个aclIntArray、out、workspaceSize、executor逐一判空数据类型检查DTYPE_SUPPORT_LIST仅允许DT_FLOAT16、DT_FLOAT、DT_BF16、DT_BOOL且out的数据类型必须与self完全一致维度检查CheckInputDims要求self维度为 3 或 4四维时 N 必须非负C/H/W 必须大于 0属性数组检查CheckArray要求kernelSize、dilation、padding、stride长度均为 2kernelSize/dilation/stride元素大于 0padding元素大于等于 0输出 shape 校验CheckOutputDims依据前述公式计算outH、outW再与kernelSize推导输出通道数C * kernelH * kernelW与空间尺寸outH * outW要求与传入out的 shape 完全一致同时对 INT64 乘法做了溢出防护SafeMul使用int64_t上限检查。返回值aclnnStatus返回状态码具体说明可参见 docs/zh/context/aclnn_return_code.md。第一段接口完成入参校验出现以下场景时报错161001ACLNN_ERR_PARAM_NULLPTRself、kernelSize、dilation、padding、stride、out、workspaceSize或executor为空指针。 161002ACLNN_ERR_PARAM_INVALIDself或out的数据类型、数据格式、维度或shape不满足要求属性长度或取值不满足 要求计算得到的输出shape无效或发生溢出。aclnnIm2col 参数说明第二段接口的参数如下。参数名输入/输出描述使用说明workspacevoid*输入Device侧workspace内存地址。由第一段接口返回的workspaceSize申请workspaceSize为0时可传空指针。workspaceSizeuint64_t输入Device侧workspace大小。由aclnnIm2colGetWorkspaceSize获取。executoraclOpExecutor*输入包含算子计算流程的执行器。由aclnnIm2colGetWorkspaceSize获取。streamaclrtStream输入指定执行任务的Stream。不可为空指针。第二段接口的返回值同为aclnnStatus具体参见 docs/zh/context/aclnn_return_code.md。从源码看aclnnIm2col的实现非常简洁仅通过CommonOpExecutorRun(workspace, workspaceSize, executor, stream)将第一段构建好的执行器在指定 Stream 上运行这也是 CANN aclnn API 的标准执行入口见 op_api/aclnn_im2col.cpp。约束说明self的 C、H、W 必须大于 0四维输入的 N 必须大于等于 0。计算得到的outH、outW必须大于 0且所有输入、输出 shape 乘积均不得发生 INT64 溢出。支持非连续输入与输出 Tensor接口内部完成连续化和结果回写。确定性计算aclnnIm2col默认确定性实现。空 Tensor 的边界语义源码 op_api/aclnn_im2col.cpp 中对空 Tensor 做了专门处理当self-IsEmpty()为真时接口直接返回成功并释放执行器不进入计算图构建流程。这与参数表中“四维输入仅支持 N 为 0 的空 Tensor”的限制对应——三维输入不允许空 Tensor。源码级实现剖析第一段接口内部的图构建流程在通过CheckParams校验后aclnnIm2colGetWorkspaceSize会利用 L0低阶算子图 API 组装实际的计算流程其关键步骤为见 op_api/aclnn_im2col.cpp连续化l0op::Contiguous(self, ...)将可能的非连续输入转换为连续 Tensor——这就是文档“接口内部完成连续化”的实现位置维度对齐三维输入CHW会通过CreateView以 batch1 的方式升维为四维[1, C, H, W]统一交给四维语义的算子处理输出侧再通过CreateView将三维输出[C*kh*kw, outH*outW]还原为二维格式归一l0op::ReFormat将输入显式转为FORMAT_NCHW保证后续算子按 NCHW 语义工作输出侧再ReFormat回用户要求的 view formatpadding 展开接口层面的padding是[paddingH, paddingW]的对称填充长度 2内部将其展开为四边填充[padTop, padBottom, padLeft, padRight] [paddingH, paddingH, paddingW, paddingW]长度 4后传入底层算子底层算子构建l0op::Im2col(selfReFormat, kernelSize, dilation, newPadding, stride, executor)创建实际的计算节点其PADDING_MODE固定为CALCULATED见 op_api/im2col.cpp结果回写通过l0op::Cast保证输出数据类型一致、l0op::ViewCopy将计算结果写回用户提供的outTensor——这是“非连续输出回写”的实现位置workspace 汇总*workspaceSize uniqueExecutor-GetWorkspaceSize()汇总整个图流程所需的临时内存大小。输出 shape 推导InferShape底层Im2col算子的 shape 推导在 op_host/im2col_infershape.cpp 中实现首先通过Ops::Math::GetImgDataDimsByNCHWOrder将输入按 NCHW 顺序拆出N, C, H, W从运行时属性中解包ksizes固定 2 个元素且必须 0、strides、dilations同样要求 0、pads4 个元素且不允许为负padding_mode仅支持CALCULATED其余取值直接判失败输出维度固定为 3[N, C*kernelH*kernelW, outH*outW]其中outH/outW按前述公式计算动态 shape维度为 -1会原样透传乘法路径全程有 INT64 溢出防护。Host 侧 tiling 与多路径分派Host 侧 op_host/im2col_tiling.cpp 承担路径选择与核数规划。从模板参数声明 op_kernel/im2col_tiling_key.h 可以看出内核按数据类型 × 路径组合实例化共 5 条计算路径IM2COL_PATH_CONTIGUOUS_WW 方向连续适合 stride1 等连续排布场景直接DataCopyPad搬运IM2COL_PATH_GATHER_WW 方向需按步长 gather先搬一段连续数据再Gather采样IM2COL_PATH_GATHER_BOOLBOOL 专用路径BOOL 无原生计算单元先按 half 拓宽再 gather 回写IM2COL_PATH_CHANNEL_TEMPLATE通道模板路径索引模板只建一次、多通道复用IM2COL_PATH_CHANNEL_TRANSPOSE通道转置路径适用于需要跨通道重排的场景。其中 FLOAT16/FLOAT/BF16 支持 CONTIGUOUS_W、GATHER_W、CHANNEL_TEMPLATE、CHANNEL_TRANSPOSE 四条路径而 BOOL 仅支持 CONTIGUOUS_W 与 GATHER_BOOL 两条路径。tiling 阶段还包含多级调度策略见 op_kernel/im2col_kernel.h 中Init的分支逻辑优先按“通道维”切分fastChannel如 kernel1×1 时退化为 channelIdentity 的纯搬运其次按“kernel 组”切分fastGroup最后退化到按行分 tilebatchRows并在各 AICore 间按base/extra工作项均衡分配。单测 tests/ut/op_host/test_im2col_tiling.cpp 对这几条分支均有覆盖例如selects_channel_identity验证 1×1 核时走fastChannel channelIdentity零 workspace 路径falls_back_to_row_tiles验证大宽度场景回退到行 tile。内核侧执行内核入口 op_kernel/im2col.cpp 为模板化__global__ __aicore__函数按数据类型选择存储类型1 字节 →int8_t2 字节 →uint16_t4 字节 →uint32_tif constexpr在编译期分派到Im2colChannelTransposeKernel或Im2colKernelStorageT, PATH。后者使用 AscendC 的TPipe流水在 MTE2搬入、V向量计算、MTE3搬出之间通过硬事件同步如SyncMte2ToV、SyncVToMte3并对 padding 区域先补零、gather 索引越界处写入无效索引后再Maxs截断到 0保证填充位置输出为 0。调用示例仓库提供了可直接编译运行的完整示例 examples/test_aclnn_im2col.cpp覆盖三维、四维输入以及 FLOAT、FLOAT16、BFLOAT16、BOOL 四种数据类型。示例覆盖的用例主函数依次运行 4 个用例用例名dtype输入 shapekerneldilationpaddingstridefp32_rank3FLOAT{2, 5, 6}三维{3, 2}{1, 1}{1, 0}{2, 1}fp16_contiguousFLOAT16{2, 3, 7, 8}{3, 3}{1, 1}{1, 1}{1, 1}bf16_dilationBF16{1, 2, 8, 9}{3, 2}{2, 3}{2, 1}{2, 2}bool_gatherBOOL{1, 2, 33, 35}{3, 5}{1, 2}{2, 4}{2, 1}其中fp32_rank3验证三维 CHW 语义与 squeeze/unsqueeze 路径bf16_dilation验证非 1 膨胀系数bool_gather验证 BOOL 的 gather 实现与大 padding。关键调用流程示例代码展示了完整的两段式调用与资源管理流程环境初始化aclInit→aclrtSetDevice(0)→aclrtCreateStream输出 shape 推导示例自带BuildOutputShape用与文档公式一致的CalculateOutputDim内部使用__int128防止中间乘法溢出计算outH/outW与输出 shape三维输入推导为 2 维[C*kh*kw, outH*outW]四维推导为 3 维[N, C*kh*kw, outH*outW]Tensor 创建通过aclCreateTensor创建 ND 格式的输入输出并借助ContiguousStrides构造连续 strides属性数组创建aclCreateIntArray分别创建kernelSize、dilation、padding、stride四个 INT64 数组第一段调用uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; aclnnStatus ret aclnnIm2colGetWorkspaceSize(input.tensor, kernelArray, dilationArray, paddingArray, strideArray, output.tensor, workspaceSize, executor);workspace 申请与第二段调用workspaceSize 0时用aclrtMalloc申请ACL_MEM_MALLOC_HUGE_FIRST然后调用aclnnIm2col(workspace, workspaceSize, executor, stream)随后aclrtSynchronizeStream同步等待完成结果校验将输出从 Device 拷回 Host按字节计算 checksum 打印output_elements与checksum资源释放依次aclrtFree(workspace)、aclDestroyIntArray、aclDestroyTensor、aclrtDestroyStream、aclrtResetDevice、aclFinalize。值得注意的细节是当ret ACL_SUCCESS workspaceSize 0时才申请 workspace即 workspace 为 0如 1×1 核的纯搬运场景时跳过申请、直接向第二段接口传空指针这与参数表中“workspaceSize 为 0 时可传空指针”的说明完全对应。编译与运行示例代码编译与运行方式遵循 CANN 单算子 API 样例的通用流程详见 docs/zh/context/compile_and_run_sample.md要点如下前提已搭建好驱动、固件、CANN 软件包、ops 包等基础环境头文件需要#include acl/acl.h与#include aclnnop/aclnn_im2col.hCMake 链接以 CANN 安装路径默认/usr/local/Ascend/cann可用环境变量ASCEND_CUSTOM_PATH覆盖为基准链接libascendcl.so、libnnopbase.so、libopapi_math.so运行编译产物在bin目录下直接执行即可在默认设备 0 上运行全部用例并打印各用例的 checksum。测试验证仓库为该算子提供了三级测试覆盖可作为接口正确性的验证依据op_api 层tests/ut/op_api/test_aclnn_im2col.cpp 走完整 aclnn 两段式调用op_host 层tests/ut/op_host/test_im2col_infershape.cpp 验证 shape 推导tests/ut/op_host/test_im2col_tiling.cpp 验证路径选择其中rejects_unsupported_dtypeINT32 拒绝、rejects_negative_padding负 padding 拒绝、rejects_shape_product_overflowshape 乘积溢出拒绝直接对应文档约束说明中的校验规则op_kernel 层tests/ut/op_kernel/test_im2col.cpp 与实例化文件 im2col_kernel_inst.cpp 验证内核数值正确性。总结aclnnIm2col以两段式 aclnn 接口封装了 im2col 数据重排逻辑第一段接口完成参数校验、shape 推导、非连续输入连续化、padding 展开与 L0 图构建第二段接口在指定 Stream 上执行。其底层在 Atlas A2/A3 产品上通过多路径 AscendC 内核实现兼顾了连续搬移、gather、BOOL 特化与通道模板复用等多种场景。开发者只需按照参数表构造合法的self、四个 INT64 属性数组与推导一致的out按 workspaceSize 申请内存后即可完成调用并可通过仓库中的示例与单测快速验证。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表