CANN/cann-ops-competitions Im2col算子开发设计文档

Im2col算子开发设计文档

【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions

需求背景

需求来源

本需求来源于 CANN 社区任务 2026 年 5 月发放的Im2col算子开发任务。任务要求参考 CANN 内置aclnnIm2col的 TBE 实现，使用 Ascend C 实现功能一致的算子，并新增bool数据类型支持。

适配硬件为 Atlas A2 训练系列产品和 Atlas A3 系列产品。交付内容包括算子代码、README、 aclnn 调用测试、自验证报告和通过评审的设计文档。

背景介绍

Im2col算子实现优化

Im2col将输入图像的滑动窗口展开为列矩阵，常用于将卷积转换为矩阵乘法。该算子只进行 数据重排和越界补零，不改变有效输入元素的数值。

本任务需要保持aclnnIm2col的接口语义，覆盖原有浮点类型，并为bool增加一条 Ascend C 实现路径。

Im2col算子（TBE）实现路径和相关API路径

以下路径基于 CANN 8.5.2 安装环境、ascend910b算子信息库和ops-math接口源码核验。安装路径中的{ASCEND_OPP_PATH}在本次调查环境中为/home/developer/Ascend/cann-8.5.2/opp。

接口层完成以下映射：

1. self为3维(C,H,W)时，在第0维补batch维，统一为4维NCHW输入。
2. 非连续输入先执行Contiguous。
3. padding=[paddingH,paddingW]扩展为pads=[paddingH,paddingH,paddingW,paddingW]。
4. l0op将kernelSize、stride、dilation、"CALCULATED"、pads下发给内部Im2col。
5. 3维输入计算完成后去除batch维，恢复为2维输出。

Im2col算子TBE实现现状分析

TBE支持的数据类型和数据格式

910B 的内置 Im2col TBE 配置覆盖 FP16/FP32/BF16 和CALCULATED/SAME/VALID三种 padding 模式，总共 9 个动态 NCHW 条目；当前公开aclnnIm2colAPI 使用显式padding参数，并固定映射到CALCULATED模式。本文以 ascend910b 注册的动态 NCHW TBE 路径作为基线；legacy 静态实现用于补充说明历史调度结构。

动态TBE实现描述

1.dynamic/im2col.py入口

1. @register_operator("Im2col", pattern="ExtractImagePatches")注册动态入口。
2. 校验输入format必须为NCHW，dtype必须为float16/float32/bfloat16。
3. 通过None in (ksizes, strides, dilations, pads)判断属性是否为binary形式。
4. 属性已知时，prepare_params将长度1或2的ksizes/strides/dilations扩展为4维形式，并将pads扩展为四向。
5. 创建ExtractImagePatchesNCHW对象并调用build(kernel_name)。

2.ExtractImagePatchesNCHW初始化

1. 归一化dtype；其中bfloat16在TVM内部按16位承载类型处理。
2. 根据shape中是否包含-1/-2区分常量shape和动态shape。
3. 解析kh/kw、sh/sw、rh/rw、pt/pb/pl/pr。
4. CALCULATED使用显式公式计算Ho/Wo；SAME/VALID通过_calc_output_size_and_pad计算输出尺寸和padding。
5. 常量shape调用op_tiling.do_op_tiling获取tiling key及n/c/kh/kw/ho/wo/ub七个factor；动态shape使用TVM变量表示。

3._compute()计算图

1. 创建输入placeholder，并通过x_ub=identity(x)标记UB计算节点。
2. 将逻辑索引从NCHW转换为NHWC。
3. 非VALID模式调用_padding()，通过四向tvm.select生成补零张量。
4. 生成6维滑窗结果：y_nhwc[n,kh,kw,oh,ow,c] = x[n,oh*sh+kh*rh,ow*sw+kw*rw,c]。
5. 将索引顺序转换为(N,C,KH,KW,Ho,Wo)。
6. 生成输出tensor；其flatten语义为(N,C*KH*KW,Ho*Wo)。

4.build()构建过程

build()依次执行tbe.compute()、_compute()、注册compile info、tbe.auto_schedule(y)和tbe.build(...)，最终生成CCE内核二进制。

动态TBE源码构建流程图

Im2col算子功能分析

设输入为(N,C,H,W)，属性为(kernelH,kernelW)、(dilationH,dilationW)、(paddingH,paddingW)和(strideH,strideW)：

outH = floor((H + 2*paddingH - dilationH*(kernelH-1) - 1) / strideH + 1) outW = floor((W + 2*paddingW - dilationW*(kernelW-1) - 1) / strideW + 1) L = outH * outW

输出shape为：

4维输入: (N, C*kernelH*kernelW, L) 3维输入: (C*kernelH*kernelW, L)

输出元素映射为：

row = c*kernelH*kernelW + kh*kernelW + kw col = oh*outW + ow ih = oh*strideH - paddingH + kh*dilationH iw = ow*strideW - paddingW + kw*dilationW

(ih,iw)有效时复制输入值，否则写0；bool的越界填充值为false。当outH<=0或outW<=0时按非法参数报错。

需求分析

需求描述

使用 Ascend C 实现aclnnIm2col，对齐原 TBE 的功能和公开接口语义，支持float16、float32、bfloat16、bool，并满足泛化、精度和性能要求。

需求拆解

1. 支持3维(C,H,W)和4维(N,C,H,W)输入。
2. 支持float16、float32、bfloat16、bool，输入输出dtype一致。
3. 支持长度为2的kernelSize、dilation、padding、stride及其合法组合。
4. 支持非连续输入、非32B对齐尾块、小shape和大shape。
5. 有效位置保持二进制一致，padding区域写0或false。
6. 大shape使用多核并行；bool性能按任务书与TBE float16基线比较。

详细设计

算子分析

算子原型、支持数据类型和形状

aclnnIm2col接口层接收self、kernelSize、dilation、padding、stride、out。其中padding在接口层扩展为四边pads，padding_mode固定为CALCULATED后下发给Im2colCustom。Im2col不执行数值运算，浮点类型直接搬运原始元素；bool按1字节类型搬运。

算子实现

实现方案

对外接口保持aclnnIm2col，内部AICore算子注册为Im2colCustom，与系统内置Im2col区分；aclnn层将参数规整后加入该算子的AICore launcher。实现采用以下分层：

1. aclnn层校验参数、扩展四向padding、连续化输入，并通过view将3维输入统一为4维。
2. Host tiling按64B cache line对输出线性区间分核，填充统一的Im2colCustomTilingData，使用一个固定tiling key。
3. Kernel入口按sizeof(DTYPE_X)分派：1字节类型进入bool专用实现，2/4字节类型进入Im2colCustomGeneric<T>。
4. 两类Kernel均在Process()内根据实际shape选择连续拷贝、特化3x3/stride=2/padding=1或通用按行处理路径。

Ascend C实现流程图

host侧设计

参数处理

1. 校验非空指针、输入维度、属性长度和属性取值。
2. 校验输入输出dtype一致，输出shape与公式一致。
3. 非连续输入执行Contiguous。
4. 3维输入补batch维；padding从两元素扩展为四边。
5. outH/outW<1时返回参数错误，不启动kernel。

分核和tiling

1. 根据dtype计算元素字节数：bool为1B，float16/bfloat16为2B，float32为4B。
2. 将总输出元素数换算为64B cache line数量，按平台AIV核数均分cache line。
3. elementsPerCore按cache line元素数对齐，使每个核写回独立的cache line区间。。
4. 实际核数不超过平台AIV核数和输出cache line数量。
5. tiling数据记录输入shape、卷积属性、输出shape、总输出元素数和elementsPerCore。
6. 当前Host只设置IM2COL_CUSTOM_TILING_KEY，具体优化路径由Kernel根据dtype和shape判断。

kernel侧设计

连续拷贝快路径

输入输出元素数量和线性顺序一致。每个核处理连续区间，以对齐tile执行DataCopyPad或 等价DMA搬运；尾tile按真实长度写回。

特化3x3 stride2路径

当kernel=3x3、stride=2、dilation=1、padding=1，且输入输出宽度、plane大小和 tile容量满足对齐条件时，Kernel按(n,c)plane分配任务，将输入plane搬入UB后完成 stride2抽取和9个kernel位置的输出组织。

通用按行路径

固定(n,c,kh,kw,oh)，计算：

ih = oh*strideH - paddingH + kh*dilationH

Kernel由输出线性区间得到flatOutputRow和outputWStart。当从完整输出行开始时：

1. strideW=1尝试批量处理连续输出行。
2. strideW!=1尝试批量处理strided输出行。
3. 无法批量处理时，按当前行剩余元素和tile容量生成segment。
4. 每个segment计算有效W区间；有效位置读取输入，无效位置写0。

浮点类型使用Im2colCustomGeneric<T>，bool使用1字节专用实现。两者共享相同输出索引语义，但UB组织和可用指令不同。

bool和尾块

1. bool使用uint8_t等1字节承载类型，值0表示false。
2. 输出优先在UB形成连续块后DMA写回。
3. 非32B对齐尾块使用按真实长度的安全搬运方式，避免越界覆盖。

Ascend C实现与TBE动态路径的差异点和原因

支持硬件

算子约束限制

1. 输入维度为3或4。
2. kernelSize、dilation、padding、stride长度均为2。
3. kernelSize、dilation、stride各元素大于0，padding各元素大于等于0。
4. 输入输出dtype一致。
5. 推导的outH/outW必须大于0。
6. 非连续输入在aclnn层转连续后进入kernel。

可维可测分析

精度标准/性能标准

小shape的TBE耗时低于10us且差距不超过3us时，按任务书补充流水图和瓶颈分析。

自验证用例设计

正确性测试使用CPU参考实现生成golden。浮点按任务书精度标准验收，并对数据重排场景增加严格相等检查；bool逐元素严格相等，padding区域单独检查0值。性能使用msprof的内核Task Duration(us)，每个shape独立运行多次并保留原始日志。

兼容性分析

1. 继承aclnnIm2col的函数签名、3维/4维输入语义和输出shape。
2. 在原有浮点类型行为基础上扩展bool数据类型。
3. 对外保持ND Tensor语义，3维输入view变换和内部tiling封装在实现内部。
4. aclnn层完成非连续输入规整，kernel接收连续NCHW数据。

参考资料

1. 04_tasks/01_community-task-2026/docs/202605/im2col_task_doc.md
2. 04_tasks/01_community-task-2026/resources/design_template.md
3. ops-math/conversion/im2col
4. CANN 8.5.2 内置 TBEIm2col源码、Proto和ascend910b信息库
5. aclnnIm2colAPI文档： https://www.hiascend.com/document/detail/zh/canncommercial/850/API/aolapi/context/ops-math/aclnnIm2col.md

【免费下载链接】cann-ops-competitions本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-ops-competitions