本教程介绍如何使用 torch_mlu.utils.cpp_extension.load_inline 在 Python 中直接编译和调用 BangC kernel。
Torch-Inline 允许你在 Python 中直接编写和编译 BangC 代码，无需单独的编译脚本。它的优势：
- 实时编译：首次运行时自动编译，后续使用缓存
- 无缝集成：编译后的函数可以直接从 Python 调用
- 类型安全：利用 PyTorch 的 C++ API 进行类型检查
| 方法 | 优点 | 缺点 |
|---|---|---|
| 纯 PyTorch | 简单易用 | 性能可能不够优化 |
| 独立 BangC 程序 | 性能最优 | 开发复杂，需要手动编译 |
| Torch-Inline | 平衡性能和开发效率 | 需要了解 C/C++ 和 BangC |
load_inline 是 Torch-Inline 的核心函数：
load_inline(
name, # 扩展名称
cpp_sources, # C++ 函数声明
bang_sources, # BangC kernel 实现
functions, # 暴露给 Python 的函数列表
verbose=False, # 显示编译过程
extra_cflags=[],# C++ 编译选项
extra_ldflags=[]# 链接选项
)返回值： Python 模块对象，可通过 module.function_name() 调用。
BangC 是寒武纪 MLU 的编程语言。关键概念：
- GDRAM: 全局内存（存储输入输出数据）
- NRAM: 片上存储器（高速但容量小）
- Kernel: 运行在 MLU 上的计算函数
- Core: MLU 的计算核心，可并行执行
使用 torch::Tensor 在 C++ 中操作 PyTorch 张量：
torch::Tensor vec_add_bang(torch::Tensor A, torch::Tensor B) {
auto out = torch::empty_like(A); // 创建输出张量
float* ptr_A = A.data_ptr<float>(); // 获取数据指针
// 调用 BangC kernel
return out;
}创建 vec_add_inline.py：
import torch
import torch_mlu
from torch_mlu.utils.cpp_extension import load_inline
# =============================================================================
# BangC 源代码
# =============================================================================
source = r"""
#include <bang.h>
#include <torch/extension.h>
#include <cnrt.h>
#include <cstdint>
// 配置参数
#define NRAM_SIZE_BYTES (384 * 1024) // NRAM 容量
#define TILE_SIZE (NRAM_SIZE_BYTES / sizeof(float) / 8)
// CNRT 错误检查宏
#define CNRT_CHECK(cmd) do { cnrtRet_t __ret = (cmd); (void)__ret; } while(0)
// =============================================================================
// BangC Kernel: 向量加法 C = A + B
// =============================================================================
__mlu_global__ void vec_add_kernel(const float* __restrict__ A,
const float* __restrict__ B,
float* __restrict__ C,
int32_t n) {
// 获取核心信息
int32_t core_num = taskDimX * taskDimY * taskDimZ;
int32_t core_id = taskIdX + taskDimX * (taskIdY + taskDimY * taskIdZ);
// 计算当前核心的数据范围
int32_t data_per_core = n / core_num;
int32_t data_last_core = data_per_core + n % core_num;
int32_t start_idx = core_id * data_per_core;
int32_t work_elems = (core_id == core_num - 1) ? data_last_core : data_per_core;
// NRAM 缓冲区
__nram__ float nram_A[TILE_SIZE];
__nram__ float nram_B[TILE_SIZE];
__nram__ float nram_C[TILE_SIZE];
// 分块处理
int32_t processed = 0;
while (processed < work_elems) {
int32_t remain = work_elems - processed;
int32_t cur = remain > TILE_SIZE ? TILE_SIZE : remain;
// 加载数据到 NRAM
__memcpy(nram_A, A + start_idx + processed, cur * sizeof(float), GDRAM2NRAM);
__memcpy(nram_B, B + start_idx + processed, cur * sizeof(float), GDRAM2NRAM);
// 在 NRAM 中执行向量加法
__bang_add(nram_C, nram_A, nram_B, cur);
// 将结果写回 GDRAM
__memcpy(C + start_idx + processed, nram_C, cur * sizeof(float), NRAM2GDRAM);
processed += cur;
}
__sync_all();
}
// =============================================================================
// C++ 包装函数
// =============================================================================
torch::Tensor vec_add_bang(torch::Tensor A, torch::Tensor B) {
// 创建输出张量
auto out = torch::empty_like(A);
int32_t numel = static_cast<int32_t>(A.numel());
// 创建 MLU 队列
cnrtQueue_t queue;
CNRT_CHECK(cnrtSetDevice(0));
CNRT_CHECK(cnrtQueueCreate(&queue));
// Kernel 启动配置
cnrtDim3_t dim = {4, 1, 1};
cnrtFunctionType_t ktype = cnrtFuncTypeUnion1;
// 启动 BangC kernel
vec_add_kernel<<<dim, ktype, queue>>>(
A.data_ptr<float>(),
B.data_ptr<float>(),
out.data_ptr<float>(),
numel
);
// 同步队列
cnrtQueueSync(queue);
CNRT_CHECK(cnrtQueueDestroy(queue));
return out;
}
"""
# =============================================================================
# C++ 函数声明
# =============================================================================
cpp_src = r"torch::Tensor vec_add_bang(torch::Tensor A, torch::Tensor B);"
# =============================================================================
# 编译扩展
# =============================================================================
vec_add_extension = load_inline(
name="vec_add_bang_ext",
cpp_sources=cpp_src,
bang_sources=source,
functions=["vec_add_bang"],
verbose=True,
extra_cflags=["-O3"],
extra_ldflags=["-lcnrt"],
)
# =============================================================================
# 测试代码
# =============================================================================
print("=" * 60)
print("Torch-Inline 向量加法示例")
print("=" * 60)
# 创建测试数据
size = 1024 * 1024
A = torch.randn(size)
B = torch.randn(size)
print(f"输入向量 A: shape={A.shape}, dtype={A.dtype}")
print(f"输入向量 B: shape={B.shape}, dtype={B.dtype}")
# CPU 参考计算
ref_result = A + B
print(f"\nCPU 参考结果: min={ref_result.min():.4f}, max={ref_result.max():.4f}")
# MLU 计算（如果可用）
if torch.mlu.is_available():
print("\nMLU 设备可用，在 MLU 上运行...")
A_mlu = A.to('mlu')
B_mlu = B.to('mlu')
# 调用编译好的 BangC 函数
custom_result = vec_add_extension.vec_add_bang(A_mlu, B_mlu)
custom_result = custom_result.cpu()
print(f"MLU 自定义结果: min={custom_result.min():.4f}, max={custom_result.max():.4f}")
# 验证结果
diff = torch.abs(ref_result - custom_result)
max_diff = diff.max().item()
print(f"\n最大差异: {max_diff:.8f}")
if max_diff < 1e-6:
print("✓ 结果验证通过！")
else:
print("✗ 结果验证失败！")
else:
print("\nMLU 设备不可用，跳过 MLU 测试")
print("=" * 60)cd /workspace/volume/ict/Experiments/08_torchinline
python vec_add_inline.py预期输出：
============================================================
Torch-Inline 向量加法示例
============================================================
输入向量 A: shape=torch.Size([1048576]), dtype=torch.float32
输入向量 B: shape=torch.Size([1048576]), dtype=torch.float32
CPU 参考结果: min=-0.1234, max=2.3456
MLU 设备可用，在 MLU 上运行...
MLU 自定义结果: min=-0.1234, max=2.3456
最大差异: 0.00000000
✓ 结果验证通过！
============================================================
┌─────────────────────────────────────────────────────────────┐
│ 1. 导入模块 │
│ import torch, torch_mlu │
│ from torch_mlu.utils.cpp_extension import load_inline│
├─────────────────────────────────────────────────────────────┤
│ 2. BangC 源代码 (source = r""" ... """) │
│ - 头文件包含 │
│ - 宏定义 │
│ - BangC kernel 实现 │
│ - C++ 包装函数 │
├─────────────────────────────────────────────────────────────┤
│ 3. C++ 接口声明 (cpp_src = r""" ... """) │
│ - 函数签名声明 │
├─────────────────────────────────────────────────────────────┤
│ 4. 编译扩展 (load_inline(...)) │
│ - name: 扩展名称 │
│ - cpp_sources: C++ 源代码 │
│ - bang_sources: BangC 源代码 │
│ - functions: 函数列表 │
├─────────────────────────────────────────────────────────────┤
│ 5. 测试和使用 │
│ - 创建测试数据 │
│ - 调用编译好的函数 │
│ - 验证结果 │
└─────────────────────────────────────────────────────────────┘
__mlu_global__ void vec_add_kernel(const float* __restrict__ A,
const float* __restrict__ B,
float* __restrict__ C,
int32_t n) {
// Kernel 实现
}关键字解释：
__mlu_global__: 声明这是 MLU kernel 函数__restrict__: 告诉编译器指针不会重叠，可优化const: 表示输入参数是只读的
int32_t core_num = taskDimX * taskDimY * taskDimZ;
int32_t core_id = taskIdX + taskDimX * (taskIdY + taskDimY * taskIdZ);解释：
taskDimX,taskDimY,taskDimZ: kernel 启动时的核心网格维度taskIdX,taskIdY,taskIdZ: 当前核心在网格中的坐标core_id: 当前核心的唯一 ID（0 到 core_num-1）
__nram__ float nram_A[TILE_SIZE];
__memcpy(nram_A, A + start_idx, cur * sizeof(float), GDRAM2NRAM);
__bang_add(nram_C, nram_A, nram_B, cur);
__memcpy(C + start_idx, nram_C, cur * sizeof(float), NRAM2GDRAM);内存层次：
GDRAM (慢，大容量)
↓ GDRAM2NRAM
NRAM (快，小容量)
↓ 计算 (__bang_add)
NRAM (快，小容量)
↓ NRAM2GDRAM
GDRAM (慢，大容量)
cnrtDim3_t dim = {4, 1, 1}; // 使用 4 个核心
vec_add_kernel<<<dim, cnrtFuncTypeUnion1, queue>>>(...);参数说明：
dim: 核心网格维度 (x=4, y=1, z=1) → 共 4 个核心cnrtFuncTypeUnion1: 核心类型queue: CNRT 队列，管理 kernel 执行
在编写代码之前，你需要理解以下概念：
- 内存层次：GDRAM（全局）↔ NRAM（片上）→ 计算
- 并行执行：多个核心同时处理不同数据
- 分块策略：大数据分成小块适应 NRAM 容量
- 数据指针：通过
.data_ptr<type>()获取张量的原始指针
关键点：
__mlu_global__ void kernel_name(const float* A, float* C, int32_t n) {
// 1. 获取核心信息
int core_id = taskIdX;
int core_num = taskDimX;
// 2. 计算数据范围
int start = core_id * n / core_num;
int work = n / core_num;
// 3. 声明 NRAM 缓冲区
__nram__ float nram_buf[TILE_SIZE];
// 4. 分块处理
int processed = 0;
while (processed < work) {
int cur = min(TILE_SIZE, work - processed);
// 5. 加载到 NRAM
__memcpy(nram_buf, A + start + processed, cur * sizeof(float), GDRAM2NRAM);
// 6. 计算
// __bang_... 操作
// 7. 写回 GDRAM
__memcpy(C + start + processed, nram_buf, cur * sizeof(float), NRAM2GDRAM);
processed += cur;
}
__sync_all();
}torch::Tensor func_name(torch::Tensor A) {
// 1. 创建输出张量
auto out = torch::empty_like(A);
// 2. 获取张量信息
int32_t n = static_cast<int32_t>(A.numel());
// 3. 创建队列
cnrtQueue_t queue;
cnrtSetDevice(0);
cnrtQueueCreate(&queue);
// 4. 启动 kernel
cnrtDim3_t dim = {4, 1, 1};
kernel_name<<<dim, cnrtFuncTypeUnion1, queue>>>(
A.data_ptr<float>(),
out.data_ptr<float>(),
n
);
// 5. 同步和清理
cnrtQueueSync(queue);
cnrtQueueDestroy(queue);
return out;
}# 1. 编译扩展
extension = load_inline(
name="ext_name",
cpp_sources=cpp_src, # C++ 函数声明
bang_sources=source, # BangC 源代码
functions=["func_name"], # 函数列表
verbose=True, # 显示编译过程
)
# 2. 在 Python 中调用
A = torch.randn(1024).to('mlu')
result = extension.func_name(A)# 1. 小规模测试
A = torch.randn(16)
B = torch.randn(16)
# 2. CPU 参考结果
ref = A + B
# 3. MLU 计算
A_mlu = A.to('mlu')
B_mlu = B.to('mlu')
result = extension.func_name(A_mlu, B_mlu).cpu()
# 4. 验证
if torch.allclose(ref, result, rtol=1e-5):
print("✓ 结果正确")
else:
print("✗ 结果错误")问题： error: 'torch' namespace not found
解决： 确保包含正确头文件：
#include <torch/extension.h>问题： error: '__bang_add' was not declared
解决： 确保包含 BangC 头文件：
#include <bang.h>问题： RuntimeError: CNRT error: ...
解决： 检查：
- 设备是否可用：
cnrtSetDevice(0) - 队列是否创建成功
- 内存是否足够分配
- Kernel 参数是否正确
问题： CPU 和 MLU 结果不一致
解决：
- 使用小规模数据测试
- 检查数据类型
- 检查矩阵布局（行主序 vs 列主序）
- 对比参考实现
方法：
# 1. 使用小数据
A = torch.randn(16)
# 2. 打印中间结果
print("Input:", A)
print("Output:", result)
# 3. 逐步验证
# 先在 CPU 上验证逻辑
# 再迁移到 MLU# 实时监控
cnmon -smi# 清理 PyTorch 扩展缓存
rm -rf ~/.cache/torch_extensions/-
基础 BangC 编程
- 参考
../01_vecadd/了解更多 BangC 基础 - 学习内存管理、并行策略等
- 参考
-
高级算子实现
- 矩阵乘法（
../05_matmul/） - Softmax（
../03_softmax/） - 卷积（
../04_maxpooling/和../06_conv/）
- 矩阵乘法（
-
Torch-MLU API
-
从简单开始
- 先实现简单的逐元素操作
- 再尝试归约操作
- 最后实现复杂算子
-
充分测试
- 小规模数据验证正确性
- 大规模数据测试性能
- 对比 CPU 和 MLU 结果
-
性能优化
- 使用
cnmon -smi监控设备 - 调整分块大小
- 优化 NRAM 利用率
- 使用
-
代码复用
- 封装常用 kernel
- 建立函数库
- 参考现有实现
Torch-Inline 的核心要点：
- load_inline：在 Python 中编译 BangC 代码
- mlu_global：声明 BangC kernel
- nram：声明 NRAM 缓冲区
- __memcpy：GDRAM ↔ NRAM 传输
- _bang*：BangC 算子执行计算
初级：理解概念 → 运行示例
↓
中级：修改代码 → 实现简单算子
↓
高级：优化性能 → 实现复杂算子
↓
专家：深入研究 → 性能调优
# 运行示例
python vec_add_inline.py
# 监控 MLU
cnmon
# 清理缓存
rm -rf ~/.cache/torch_extensions/