手写 CUDA 到昇腾迁移：CANN custom-op 迁移指南

你手上有一段跑得飞快的 CUDA 算子，现在老板说："这套代码要上昇腾，用 CANN 跑起来。"

你可能会想：是不是要把所有 cudaMalloc 改成 aclrtMalloc，把 <<<>>> 改成 <<<...>>> 的新语法，再啃几百页的硬件手册？

其实没那么夸张。大多数情况下，迁移可以分成两大部分：

脚本层迁移：把 PyTorch/TensorFlow 的 GPU 代码迁移到 NPU。
自定义算子迁移：把你手写的 CUDA kernel 用 CANN 的 custom-op 框架重新实现一遍。

这篇文章就是一份"从 CUDA 到昇腾迁移"的实战笔记，重点讲解第二种情况：如何把手写 CUDA kernel 迁移为 CANN 的自定义算子（custom-op）。

🗺️ 迁移总体思路

我们可以把迁移过程想象成一次"搬家"：

旧家 (CUDA)：一栋已经装修好的房子，家具（业务逻辑）都是按 NVIDIA 的规矩摆的。
新家 (昇腾)：一栋结构相似但插座、水管位置都不同的房子。
任务：把家具搬过去，并调整插座位置、重接水电，让它们在新家也能正常工作。

迁移的核心思路可以总结为以下四步：

分析现状：梳理现有 CUDA 算子，明确其功能和接口。
脚本迁移：将上层的训练/推理脚本从 CUDA 环境迁移到 NPU 环境。
算子迁移：为 CUDA kernel 在昇腾上找到对应的实现方式（如 Ascend C），并重新实现其核心逻辑。
验证调优：确保功能正确、精度对齐，并进行性能调优。

流程解读：

分析现状 ：搞清楚模型里哪些是标准算子，哪些是你自己写的 .cu 文件里的自定义算子。
脚本迁移：让模型脚本在昇腾环境下能"跑起来"，即使某些自定义算子暂时还是"空壳"。
算子迁移：为你的 CUDA kernel 在昇腾上找到一个"新家"，用 Ascend C 或 TIK 重新实现其核心逻辑。
验证调优：确保功能正确、精度对齐，并进行性能调优。

📝 第一步：分析现状，明确迁移范围

在动手前，先搞清楚两件事：

算子类型 ：你的代码里，哪些是 PyTorch/TensorFlow 自带的标准算子，哪些是你自己写的 .cu 文件里的自定义算子？
依赖关系：自定义算子都用在模型的哪些地方？输入输出是什么类型、什么维度？有没有控制流（if/for）等复杂逻辑？

核心原则：

标准算子 ：优先使用 CANN 或框架（如 torch_npu）已提供的实现，不要自己重写。
自定义算子：这是迁移工作的核心，需要逐一攻克。

对于 PyTorch 用户，可以先用 transfer_to_npu 等工具自动迁移脚本，它会自动处理大部分 torch.cuda 到 torch.npu 的替换，并生成一个"不支持算子列表"，这个列表就是你后续需要手动迁移的 custom-op 清单。

🚚 第二步：脚本迁移，为算子"腾地方"

脚本迁移的目标，是让你的模型脚本在昇腾环境下能"跑起来"，即使某些自定义算子暂时还是"空壳"。

1. 环境准备

安装 CANN 软件栈。
安装适配了昇腾的 PyTorch 版本（如 2.1.0+）及 torch_npu 插件。
通过 python3 -c "import torch; import torch_npu; print(torch_npu.npu.is_available())" 验证环境是否就绪。

2. 代码修改

替换设备 ：将 torch.device("cuda:0") 改为 torch.device("npu:0")。
导入插件 ：在脚本开头添加 import torch_npu 和 from torch_npu.contrib import transfer_to_npu，后者可帮助自动映射部分 CUDA API。
修改分布式 ：将 nccl 相关的初始化替换为 hccl。

完成这一步后，大部分模型应该能在 NPU 上跑起来，但调用到自定义算子时会报错，这是正常现象，接下来就是解决这些报错。

🏗️ 第三步：算子迁移，核心实现

这是整个迁移过程中技术含量最高的一步：为你的 CUDA kernel 在昇腾上找到一个"新家"。

1. 理解 CANN 自定义算子实现方式

CANN 提供了多种自定义算子开发路径，你可以根据算子的复杂度和性能要求选择：

Ascend C + msOpGen ：官方主推方式，性能控制力强。通过 msOpGen 工具从 JSON 描述文件生成工程模板，开发者只需填充核函数实现、Tiling 策略和框架适配代码即可。
TIK (Tensor Iterator Kernel)：基于 Python 的 DSL，适合快速实现原型。通过 TIK 提供的 API 描述数据搬运和计算流程，由编译器生成内核，上手快但性能调优空间相对较小。
AOL (Ascend Operator Library)：用于封装已用 ACL (Ascend Computing Language) 写好的高性能 kernel，使其能被框架调用。

对于从 CUDA 迁移的场景，Ascend C + msOpGen 是最推荐的路线，因为它能让你对性能有更精细的控制。

2. 建立映射关系：CUDA 概念 vs. CANN 概念

要将 CUDA kernel 翻译成 Ascend C，首先需要理解两者核心概念的对应关系：

CUDA 概念	CANN (Ascend C) 对应概念
`cudaMalloc` / `cudaFree`	`aclrtMalloc` / `aclrtFree` (设备内存管理)
`__global__` 函数	`__aicore__ __global__` 函数 (核函数入口)
`blockIdx`, `threadIdx`	`GetBlockIdx()`, `GetThreadIdx()` (逻辑坐标)
`<<<grid, block>>>`	Tiling 策略 (在 Host 侧计算分片参数)
`cudaMemcpyAsync`	`DataCopy` + 双缓冲 (实现流水)
`threadIdx.x` 上的循环	`vec_add` 等向量指令 (实现并行)

3. 迁移流程：从 Kernel 到 Operator

以迁移一个简单的向量加法 kernel 为例，迁移流程如下：

编写算子原型 JSON：定义算子的名称、输入输出、数据类型和格式。
生成工程 ：使用 msOpGen 工具，根据 JSON 文件生成 Ascend C 算子工程。
实现核函数 ：在生成的工程模板中，填充 __aicore__ 核函数的实现，核心是编写"搬入(CopyIn) -> 计算(Compute) -> 搬出(CopyOut)"的三级流水逻辑。
实现 Host 侧逻辑：完成 Tiling 策略计算、算子原型注册、Shape 推导函数等，使框架能够正确调用你的算子。
编译与集成：编译生成算子插件（.so 文件），并通过 PyTorch 的 C++/Python 扩展机制将其集成到你的模型中。

4. 关键细节处理

内存布局：确保输入/输出 Tensor 在昇腾上的内存布局（如 NCHW）与你 kernel 的实现预期一致。
数据类型：昇腾对 FP16/FP32 有专门优化，优先使用 FP16 以获得更好的性能。
边界处理：在 Tiling 和 kernel 实现中，要特别注意处理数据长度不能被 Tiling 大小整除的边界情况。

✅ 第四步：验证与调优

算子实现完成后，工作并未结束，还需要进行严格的验证和性能调优。

1. 功能与精度验证

单元测试：构造简单的输入数据，对比 CANN 算子与原始 CUDA 算子的输出结果，确保数值一致（允许微小的浮点误差）。
端到端验证 ：将算子集成回完整的模型，在 NPU 上运行训练或推理，观察 Loss 曲线和最终精度是否与 GPU 基线对齐。如果出现精度偏差，可使用 msprobe 等工具定位问题算子。

2. 性能调优

Profiling：使用 CANN 的 profiling 工具分析算子的执行时间、内存带宽占用和计算单元利用率。
优化方向：根据 profiling 结果，重点优化数据搬运（如调整 Tiling 大小、使用双缓冲）和计算（如使用向量指令、循环展开）的效率。

💡 小白迁移心法

先跑通，再优化：不要一开始就追求极致的性能，先让算子功能正确、精度对齐。
善用工具 ：充分利用 msOpGen、transfer_to_npu 等官方工具，能极大减少重复劳动。
学会看报错：昇腾的错误信息通常很详细，仔细阅读能帮你快速定位问题。
参考官方示例：CANN 社区提供了大量自定义算子示例，看懂它们比自己瞎猜要高效得多。

总而言之，从手写 CUDA 迁移到昇腾 CANN custom-op，更像是一场有图纸、有工具的"搬家"工程，而非从零开始的"造房"。只要理清思路，步步为营，你就能成功将你的算法高效地运行在昇腾平台上。