neuronxcc包介绍及示例代码

介绍

neuronxcc（或 neuronx-cc）是 AWS Neuron 编译器，这是 AWS 专为其定制机器学习芯片开发的编译工具包的核心组件。

主要功能

neuronxcc 接受各种格式的机器学习模型（TensorFlow、MXNet、PyTorch、XLA HLO）并将它们优化为在 Neuron 设备上运行 Neuron Compiler — AWS Neuron Documentation。

技术特点

AOT 编译器：Neuron 编译器是一个提前编译器（Ahead-of-Time compiler），可以加速模型在 NeuronCores 上的执行 Neuron Compiler (neuronx-cc) release notes — AWS Neuron Documentation
自动类型转换：编译器自动将 FP32 转换为内部支持的数据类型，如 FP16 或 BF16 Neuron Compiler FAQ (neuronx-cc) — AWS Neuron Documentation
输出格式：编译后生成 NEFF（Neuron Executable File Format）归档文件

支持的硬件平台

AWS Inferentia：用于推理工作负载
AWS Trainium：支持在 Trn1 实例上编译训练模型 Neuron Compiler (neuronx-cc) release notes — AWS Neuron Documentation

版本区别

存在两个版本：

neuron-cc：用于 Inf1 实例
neuronx-cc：用于更新的 Inf2 和 Trn1 实例

使用场景

这个编译器主要用于：

• 将深度学习模型优化部署到 AWS 的专用 AI 芯片上
• 提高模型推理和训练的性能
• 降低在 AWS AI 实例上运行 ML 工作负载的成本

代码示例

"""
NeuronX Compiler (NKI) 示例代码集合
展示如何使用 neuronxcc.nki 编写自定义内核
"""import neuronxcc.nki.isa as nisa
import neuronxcc.nki.language as nl
from neuronxcc import nki
from neuronxcc.nki.language import par_dim
import numpy as np# ============================================================
# 示例 1: 简单的向量加法内核
# ============================================================@nki.jit
def vector_add_kernel(a_tensor, b_tensor, result_tensor):"""简单的向量加法内核"""# 获取张量的形状信息N = a_tensor.shape[0]# 使用并行维度进行计算i = nl.arange(N)[:]# 从HBM加载数据到SBUFa_tile = nl.load(a_tensor[i])b_tile = nl.load(b_tensor[i])# 执行加法运算result_tile = nisa.tensor_tensor(a_tile, b_tile, op=nl.add)# 将结果存储回HBMnl.store(result_tensor[i], result_tile)# ============================================================
# 示例 2: 矩阵乘法内核
# ============================================================@nki.jit
def matmul_kernel(a_tensor, b_tensor, result_tensor):"""简单的矩阵乘法内核M x K @ K x N = M x N"""M, K = a_tensor.shapeK_b, N = b_tensor.shape# 分块大小（根据硬件限制调整）tile_m = 128tile_n = 128tile_k = 512# 使用并行维度i_m = par_dim(M // tile_m)i_n = par_dim(N // tile_n)for i in nl.affine_range(i_m):for j in nl.affine_range(i_n):# 初始化累加器acc = nl.zeros((tile_m, tile_n), dtype=nl.float32)# K维度的分块循环for k in nl.affine_range(K // tile_k):# 加载A和B的分块a_tile = nl.load(a_tensor[i*tile_m:(i+1)*tile_m, k*tile_k:(k+1)*tile_k])b_tile = nl.load(b_tensor[k*tile_k:(k+1)*tile_k, j*tile_n:(j+1)*tile_n])# 执行矩阵乘法并累加acc = nisa.nc_matmul(a_tile, b_tile, acc)# 存储结果nl.store(result_tensor[i*tile_m:(i+1)*tile_m, j*tile_n:(j+1)*tile_n], acc)# ============================================================
# 示例 3: 激活函数内核 (ReLU)
# ============================================================@nki.jit
def relu_kernel(input_tensor, output_tensor):"""ReLU激活函数内核"""# 获取张量形状shape = input_tensor.shapetotal_elements = np.prod(shape)# 分块处理tile_size = 1024num_tiles = (total_elements + tile_size - 1) // tile_sizefor tile_idx in nl.affine_range(num_tiles):start_idx = tile_idx * tile_sizeend_idx = nl.minimum(start_idx + tile_size, total_elements)# 加载数据input_tile = nl.load(input_tensor.reshape(-1)[start_idx:end_idx])# 创建零张量用于比较zeros = nl.zeros_like(input_tile)# 执行 ReLU: max(0, x)relu_result = nisa.tensor_tensor(input_tile, zeros, op=nl.maximum)# 存储结果nl.store(output_tensor.reshape(-1)[start_idx:end_idx], relu_result)# ============================================================
# 示例 4: Softmax 内核
# ============================================================@nki.jit
def softmax_kernel(input_tensor, output_tensor):"""Softmax 激活函数内核假设输入是 [batch_size, features] 的二维张量"""batch_size, features = input_tensor.shape# 并行处理每个batchi = par_dim(batch_size)for batch_idx in nl.affine_range(i):# 加载当前batch的数据x = nl.load(input_tensor[batch_idx, :])# 步骤1: 找到最大值 (数值稳定性)x_max = nisa.tensor_reduce(x, op=nl.max, axis=0, keepdims=True)# 步骤2: 减去最大值x_shifted = nisa.tensor_scalar(x, x_max, op=nl.subtract)# 步骤3: 计算指数exp_x = nisa.activation(x_shifted, op=nl.exp)# 步骤4: 计算和exp_sum = nisa.tensor_reduce(exp_x, op=nl.add, axis=0, keepdims=True)# 步骤5: 除法得到最终结果softmax_result = nisa.tensor_tensor(exp_x, exp_sum, op=nl.divide)# 存储结果nl.store(output_tensor[batch_idx, :], softmax_result)# ============================================================
# 示例 5: 卷积内核（简化版）
# ============================================================@nki.jit
def conv2d_kernel(input_tensor, weight_tensor, output_tensor, stride=1, padding=0):"""简化的2D卷积内核input_tensor: [batch, in_channels, height, width]weight_tensor: [out_channels, in_channels, kernel_h, kernel_w]"""batch, in_channels, in_h, in_w = input_tensor.shapeout_channels, _, kernel_h, kernel_w = weight_tensor.shape# 计算输出尺寸out_h = (in_h + 2 * padding - kernel_h) // stride + 1out_w = (in_w + 2 * padding - kernel_w) // stride + 1# 并行化batch和output channel维度b = par_dim(batch)oc = par_dim(out_channels)for batch_idx in nl.affine_range(b):for out_ch in nl.affine_range(oc):# 为每个输出通道初始化结果output_slice = nl.zeros((out_h, out_w), dtype=nl.float32)# 遍历输入通道for in_ch in nl.affine_range(in_channels):# 加载权重和输入特征图weight = nl.load(weight_tensor[out_ch, in_ch, :, :])# 执行卷积操作（简化实现）for h in nl.affine_range(out_h):for w in nl.affine_range(out_w):h_start = h * stridew_start = w * stride# 提取输入窗口input_window = nl.load(input_tensor[batch_idx, in_ch,h_start:h_start+kernel_h,w_start:w_start+kernel_w])# 计算点积conv_result = nisa.tensor_reduce(nisa.tensor_tensor(input_window, weight, op=nl.multiply),op=nl.add)# 累加到输出output_slice = nisa.tensor_scalar(output_slice, conv_result, op=nl.add, mask=(h, w))# 存储最终结果nl.store(output_tensor[batch_idx, out_ch, :, :], output_slice)# ============================================================
# 示例 6: 内存拷贝和转置
# ============================================================@nki.jit
def transpose_kernel(input_tensor, output_tensor):"""矩阵转置内核input: [M, N] -> output: [N, M]"""M, N = input_tensor.shape# 分块大小tile_size = 128# 分块处理for i in nl.affine_range((M + tile_size - 1) // tile_size):for j in nl.affine_range((N + tile_size - 1) // tile_size):i_start = i * tile_sizei_end = nl.minimum(i_start + tile_size, M)j_start = j * tile_sizej_end = nl.minimum(j_start + tile_size, N)# 加载输入分块input_tile = nl.load(input_tensor[i_start:i_end, j_start:j_end])# 执行转置transposed_tile = nisa.nc_transpose(input_tile, axes=(1, 0))# 存储到转置位置nl.store(output_tensor[j_start:j_end, i_start:i_end], transposed_tile)# ============================================================
# 示例 7: 归约操作 (Sum Reduction)
# ============================================================@nki.jit
def sum_reduction_kernel(input_tensor, output_tensor, axis):"""沿指定轴进行求和归约"""if axis == 0:# 沿第一个维度求和M, N = input_tensor.shapefor j in nl.affine_range(N):# 加载列数据col_data = nl.load(input_tensor[:, j])# 执行归约sum_result = nisa.tensor_reduce(col_data, op=nl.add, axis=0)# 存储结果nl.store(output_tensor[j], sum_result)elif axis == 1:# 沿第二个维度求和M, N = input_tensor.shapefor i in nl.affine_range(M):# 加载行数据row_data = nl.load(input_tensor[i, :])# 执行归约sum_result = nisa.tensor_reduce(row_data, op=nl.add, axis=0)# 存储结果nl.store(output_tensor[i], sum_result)

这些示例展示了使用 neuronxcc.nki (Neuron Kernel Interface) 编写自定义内核的各种模式：# 关键概念说明@nki.jit 装饰器：将Python函数编译为Neuron内核
par_dim()：定义并行维度，利用多核心并行计算
nl.load() / nl.store()：在HBM和SBUF之间传输数据
nisa.* 指令：底层的Neuron ISA指令集
nl.affine_range()：创建循环迭代器# 内存层次结构HBM（高带宽内存）：主内存
SBUF（静态缓冲区）：片上缓存
PSUM：部分求和缓冲区（用于矩阵乘法）# 优化要点分块处理：将大张量分成小块以适应片上内存
并行化：使用 par_dim 在多个NeuronCore上并行执行
内存访问模式：优化数据加载/存储模式以最大化带宽利用
数值稳定性：在Softmax等操作中考虑数值稳定性