muDNN 开发者指南

本文档介绍 muDNN 的架构、核心概念和使用方法。

注意

C++ API 已标记为弃用。新应用请使用 C API，C++ API 将在未来版本中移除。

安装与环境准备

muDNN 随 MUSA SDK 一起发布。安装步骤可参考 MUSA SDK 安装指南。

在使用本文档中的示例前，请先完成 Linux Driver 和 MUSA Toolkit 安装，并确保相关环境变量已生效。

概述

什么是 muDNN

muDNN（Moore Threads Deep Neural Network Library，深度神经网络库）是面向 MUSA 平台的深度神经网络加速库，提供张量、卷积、归一化、激活、注意力等常用神经网络算子接口。

关键特性

• 提供 C API 形式的神经网络算子接口，新应用推荐优先使用 C API。
• 支持张量算子、卷积神经网络算子和高级算子等多类能力。
• 通过 handle、descriptor 和 workspace 模型管理执行上下文与算子参数。
• 针对 MTGPU 架构优化，适用于深度学习训练和推理场景。

muDNN 架构

API 层次结构

muDNN C API 由以下头文件组成：

mudnn.h - 主头文件

聚合所有 C API 导出接口，包含：

• mudnn_version.h: 版本信息查询
• mudnn_ops.h: 张量算子和常用神经网络算子
• mudnn_cnn.h: 卷积相关接口
• mudnn_adv.h: 高级算子
• mudnn_graph.h: 图API

mudnn_ops.h - 基础算子

类型/函数	说明
mudnnHandle_t	muDNN 设备句柄，绑定到 MUSA 设备与流
mudnnTensorDescriptor_t	张量描述符，描述数据地址、类型、格式、维度
mudnnUnaryDescriptor_t / mudnnBinaryDescriptor_t	一元/二元张量算子描述符
mudnnLayerNormDescriptor_t / mudnnBatchNormDescriptor_t	归一化算子描述符
mudnnCreateXXXDescriptor	创建各类描述符
mudnnDestroyXXXDescriptor	销毁各类描述符

张量算子：

算子	函数前缀
一元算子	mudnnUnary
二元算子	mudnnBinary
三元算子	mudnnTernary
拼接	mudnnConcat
填充	mudnnFill
置换	mudnnPermute
归约	mudnnReduce
索引收集	mudnnGather
滑窗展开	mudnnUnfold
前缀扫描	mudnnScan, mudnnCum
权重归一化	mudnnWeightNorm
损失计算	mudnnL2Loss

mudnn_cnn.h - 卷积相关算子

算子类型	函数前缀
卷积兼容接口	mudnnConvolution（旧接口，使用前请查看 API 参考中的 deprecated 标记）
可变形卷积	mudnnDeformableConv

mudnn_adv.h - 高级算子

算子类型	函数前缀
注意力	mudnnScaledDotProductAttention（mudnnMultiHeadAttn 为旧接口）
损失	mudnnCTCLoss
矩阵乘法	mudnnBatchMatMul

快速开始

示例 1：LayerNorm（C API，使用非deprecated接口）

以下示例展示如何使用 muDNN C API 执行 LayerNorm 操作：

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <musa_runtime_api.h>
#include <mudnn.h>

#define MUDNN_CHECK(cmd)                                              \
  do {                                                                \
    mudnnStatus_t status = cmd;                                       \
    if (status != MUDNN_STATUS_SUCCESS) {                             \
      fprintf(stderr, "muDNN error %d at %s:%d\n", status,            \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

#define MUSACHECK(cmd)                                                \
  do {                                                                \
    musaError_t err = cmd;                                            \
    if (err != musaSuccess) {                                         \
      fprintf(stderr, "MUSA error %d at %s:%d\n", err,                \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

int main() {
    mudnnHandle_t handle;
    mudnnTensorDescriptor_t input_desc, output_desc;
    mudnnTensorDescriptor_t mean_desc, variance_desc, gamma_desc, beta_desc;
    mudnnLayerNormDescriptor_t layernorm_desc;
    
    // 1. 创建 muDNN 句柄
    MUDNN_CHECK(mudnnCreate(&handle));
    
    // 2. 创建描述符
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&input_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&output_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&mean_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&variance_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&gamma_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&beta_desc));
    MUDNN_CHECK(mudnnCreateLayerNormDescriptor(&layernorm_desc));
    
    // 3. 设置张量维度 (N, C, H, W) - NCHW 格式
    const int N = 4, C = 64, H = 32, W = 32;
    
    // 设置输入张量描述符
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        input_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    
    // output/gamma/beta 使用输入张量维度
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        output_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        gamma_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        beta_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));

    // 对 C/H/W 归一化时，每个 N 对应一个 mean/variance
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        mean_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, 1, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        variance_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, 1, 1, 1));
    
    // 设置 LayerNorm 描述符
    int axes[] = {1, 2, 3};  // 对 C, H, W 维度归一化
    MUDNN_CHECK(mudnnSetLayerNormDescriptor(
        layernorm_desc, 
        1e-5,  // epsilon
        MUDNN_LAYER_NORM_MODE_DIRECT,  // 模式
        axes,  // 归一化轴
        3      // 轴数量
    ));
    
    // 4. 分配设备内存
    float *input_data, *output_data, *mean_data, *variance_data;
    float *gamma_data, *beta_data;
    size_t size = N * C * H * W * sizeof(float);
    size_t mean_var_size = N * sizeof(float);
    
    MUSACHECK(musaMalloc((void**)&input_data, size));
    MUSACHECK(musaMalloc((void**)&output_data, size));
    MUSACHECK(musaMalloc((void**)&mean_data, mean_var_size));
    MUSACHECK(musaMalloc((void**)&variance_data, mean_var_size));
    MUSACHECK(musaMalloc((void**)&gamma_data, size));
    MUSACHECK(musaMalloc((void**)&beta_data, size));
    MUSACHECK(musaMemset(input_data, 0, size));
    MUSACHECK(musaMemset(output_data, 0, size));
    MUSACHECK(musaMemset(mean_data, 0, mean_var_size));
    MUSACHECK(musaMemset(variance_data, 0, mean_var_size));
    MUSACHECK(musaMemset(gamma_data, 0, size));
    MUSACHECK(musaMemset(beta_data, 0, size));
    
    // 5. 查询工作区大小
    size_t workspace_size;
    MUDNN_CHECK(mudnnGetLayerNormForwardWorkspaceSize(
        handle, layernorm_desc, output_desc, mean_desc, variance_desc,
        input_desc, gamma_desc, beta_desc, &workspace_size
    ));
    
    void* workspace = NULL;
    if (workspace_size > 0) {
        MUSACHECK(musaMalloc(&workspace, workspace_size));
    }
    
    // 6. 执行 LayerNorm
    MUDNN_CHECK(mudnnLayerNormForward(
        handle, layernorm_desc,
        output_desc, output_data,
        mean_desc, mean_data,
        variance_desc, variance_data,
        input_desc, input_data,
        gamma_desc, gamma_data,
        beta_desc, beta_data,
        workspace, workspace_size
    ));
    
    // 7. 同步流
    MUSACHECK(musaStreamSynchronize(musaStreamPerThread));
    
    printf("LayerNorm executed successfully!\n");
    
    // 8. 清理资源
    MUSACHECK(musaFree(input_data));
    MUSACHECK(musaFree(output_data));
    MUSACHECK(musaFree(mean_data));
    MUSACHECK(musaFree(variance_data));
    MUSACHECK(musaFree(gamma_data));
    MUSACHECK(musaFree(beta_data));
    if (workspace) MUSACHECK(musaFree(workspace));
    
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(input_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(output_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(mean_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(variance_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(gamma_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(beta_desc));
    MUDNN_CHECK(mudnnDestroyLayerNormDescriptor(layernorm_desc));
    MUDNN_CHECK(mudnnDestroy(handle));
    
    return 0;
}

编译命令:

mcc -x musa layernorm_example.cc --offload-arch=mp_31 \
    -I${MUSA_SDK_PATH}/include \
    -L${MUSA_SDK_PATH}/lib \
    -lmudnn -lmudnn_ops -lmudnn_base -lmudnn_tensor \
    -lmudnn_tensor_unary -lmudnn_tensor_binary -lmudnn_xmma -lmusart \
    -Wl,-rpath,${MUSA_SDK_PATH}/lib \
    -o layernorm_example

运行:

./layernorm_example

示例 2：GEMM（使用 muDNN C API）

以下示例展示如何使用 muDNN 执行矩阵乘法（GEMM）操作：

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <musa_runtime_api.h>
#include <mudnn.h>

#define MUDNN_CHECK(cmd)                                              \
  do {                                                                \
    mudnnStatus_t status = cmd;                                       \
    if (status != MUDNN_STATUS_SUCCESS) {                             \
      fprintf(stderr, "muDNN error %d at %s:%d\n", status,            \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

#define MUSACHECK(cmd)                                                \
  do {                                                                \
    musaError_t err = cmd;                                            \
    if (err != musaSuccess) {                                         \
      fprintf(stderr, "MUSA error %d at %s:%d\n", err,                \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

static void contiguous_strides(const int* shape, int nb_dims, int* strides) {
  strides[nb_dims - 1] = 1;
  for (int i = nb_dims - 2; i >= 0; --i) {
    strides[i] = strides[i + 1] * shape[i + 1];
  }
}

static void set_tensor_desc_3d(mudnnTensorDescriptor_t desc, mudnnDataType_t dtype,
                               const int shape[3]) {
  int strides[3];
  contiguous_strides(shape, 3, strides);
  MUDNN_CHECK(mudnnSetTensorNdDescriptor(desc, dtype, 3, shape, strides));
}

int main() {
  printf("muDNN version: %zu\n", mudnnGetVersion());
  MUSACHECK(musaSetDevice(0));

  const int batch = 1;
  const int m = 2;
  const int k = 3;
  const int n = 4;
  const int a_shape[3] = {batch, m, k};
  const int b_shape[3] = {batch, k, n};
  const int d_shape[3] = {batch, m, n};

  const float h_a_f32[] = {
      1.0f, 2.0f, 3.0f,
      4.0f, 5.0f, 6.0f,
  };
  const float h_b_f32[] = {
      1.0f, 2.0f, 3.0f, 4.0f,
      5.0f, 6.0f, 7.0f, 8.0f,
      9.0f, 10.0f, 11.0f, 12.0f,
  };
  float h_d[1 * 2 * 4] = {0.0f};

  void* d_a = NULL;
  void* d_b = NULL;
  void* d_d = NULL;

  MUSACHECK(musaMalloc(&d_a, sizeof(h_a_f32)));
  MUSACHECK(musaMalloc(&d_b, sizeof(h_b_f32)));
  MUSACHECK(musaMalloc(&d_d, sizeof(h_d)));

  MUSACHECK(musaMemcpy(d_a, h_a_f32, sizeof(h_a_f32), musaMemcpyHostToDevice));
  MUSACHECK(musaMemcpy(d_b, h_b_f32, sizeof(h_b_f32), musaMemcpyHostToDevice));
  MUSACHECK(musaMemset(d_d, 0, sizeof(h_d)));

  mudnnHandle_t handle = NULL;
  mudnnTensorDescriptor_t a_desc = NULL;
  mudnnTensorDescriptor_t b_desc = NULL;
  mudnnTensorDescriptor_t d_desc = NULL;
  mudnnBatchMatMulDescriptor_t bmm_desc = NULL;

  MUDNN_CHECK(mudnnCreate(&handle));
  MUDNN_CHECK(mudnnCreateTensorDescriptor(&a_desc));
  MUDNN_CHECK(mudnnCreateTensorDescriptor(&b_desc));
  MUDNN_CHECK(mudnnCreateTensorDescriptor(&d_desc));
  MUDNN_CHECK(mudnnCreateBatchMatMulDescriptor(&bmm_desc));

  set_tensor_desc_3d(a_desc, MUDNN_DATA_FLOAT, a_shape);
  set_tensor_desc_3d(b_desc, MUDNN_DATA_FLOAT, b_shape);
  set_tensor_desc_3d(d_desc, MUDNN_DATA_FLOAT, d_shape);

  MUDNN_CHECK(mudnnSetBatchMatMulDescriptor(
        bmm_desc, MUDNN_TENSOR_OP_MATH, false, false));

  size_t workspace_size = 0;
  MUDNN_CHECK(mudnnBatchMatmulGetWorkspaceSize(handle, bmm_desc, a_desc, b_desc,
                                         NULL, d_desc,
                                         &workspace_size));

  void* workspace = NULL;
  if (workspace_size > 0) {
    MUSACHECK(musaMalloc(&workspace, workspace_size));
  }

  MUDNN_CHECK(mudnnBatchMatMulBias(handle, bmm_desc, a_desc, d_a, b_desc, d_b,
                                     NULL, NULL, d_desc, d_d, NULL,
                                     NULL, workspace, workspace_size));

  MUSACHECK(musaMemcpy(h_d, d_d, sizeof(h_d), musaMemcpyDeviceToHost));

  MUSACHECK(musaDeviceSynchronize());

  if (workspace != NULL) {
    MUSACHECK(musaFree(workspace));
  }

  MUDNN_CHECK(mudnnDestroyBatchMatMulDescriptor(bmm_desc));
  MUDNN_CHECK(mudnnDestroyTensorDescriptor(a_desc));
  MUDNN_CHECK(mudnnDestroyTensorDescriptor(b_desc));
  MUDNN_CHECK(mudnnDestroyTensorDescriptor(d_desc));
  MUDNN_CHECK(mudnnDestroy(handle));

  MUSACHECK(musaFree(d_a));
  MUSACHECK(musaFree(d_b));
  MUSACHECK(musaFree(d_d));

  return 0;
}

示例 3：BatchNorm（C API）

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <musa_runtime_api.h>
#include <mudnn.h>

#define MUDNN_CHECK(cmd)                                              \
  do {                                                                \
    mudnnStatus_t status = cmd;                                       \
    if (status != MUDNN_STATUS_SUCCESS) {                             \
      fprintf(stderr, "muDNN error %d at %s:%d\n", status,            \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

#define MUSACHECK(cmd)                                                \
  do {                                                                \
    musaError_t err = cmd;                                            \
    if (err != musaSuccess) {                                         \
      fprintf(stderr, "MUSA error %d at %s:%d\n", err,                \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

int main() {
    mudnnHandle_t handle;
    mudnnTensorDescriptor_t input_desc, output_desc;
    mudnnTensorDescriptor_t scale_desc, bias_desc;
    mudnnTensorDescriptor_t acc_mean_desc, acc_var_desc;
    mudnnTensorDescriptor_t fresh_mean_desc, fresh_var_desc;
    mudnnBatchNormDescriptor_t bn_desc;
    
    // 1. 创建句柄和描述符
    MUDNN_CHECK(mudnnCreate(&handle));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&input_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&output_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&scale_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&bias_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&acc_mean_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&acc_var_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&fresh_mean_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&fresh_var_desc));
    MUDNN_CHECK(mudnnCreateBatchNormDescriptor(&bn_desc));
    
    // 2. 设置张量维度 (N, C, H, W) NCHW 格式
    const int N = 4, C = 64, H = 32, W = 32;
    
    // 输入/输出张量描述符
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        input_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        output_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    
    // scale/bias、acc mean/var 和 fresh mean/var 描述符（每通道）
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        scale_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        bias_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        acc_mean_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        acc_var_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        fresh_mean_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        fresh_var_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    
    // 3. 分配设备内存
    float *input_data, *output_data, *scale_data, *bias_data;
    float *acc_mean_data, *acc_var_data, *fresh_mean_data, *fresh_var_data;
    size_t input_size = N * C * H * W * sizeof(float);
    size_t channel_size = C * sizeof(float);
    
    MUSACHECK(musaMalloc((void**)&input_data, input_size));
    MUSACHECK(musaMalloc((void**)&output_data, input_size));
    MUSACHECK(musaMalloc((void**)&scale_data, channel_size));
    MUSACHECK(musaMalloc((void**)&bias_data, channel_size));
    MUSACHECK(musaMalloc((void**)&acc_mean_data, channel_size));
    MUSACHECK(musaMalloc((void**)&acc_var_data, channel_size));
    MUSACHECK(musaMalloc((void**)&fresh_mean_data, channel_size));
    MUSACHECK(musaMalloc((void**)&fresh_var_data, channel_size));
    MUSACHECK(musaMemset(input_data, 0, input_size));
    MUSACHECK(musaMemset(output_data, 0, input_size));
    MUSACHECK(musaMemset(scale_data, 0, channel_size));
    MUSACHECK(musaMemset(bias_data, 0, channel_size));
    MUSACHECK(musaMemset(acc_mean_data, 0, channel_size));
    MUSACHECK(musaMemset(acc_var_data, 0, channel_size));
    MUSACHECK(musaMemset(fresh_mean_data, 0, channel_size));
    MUSACHECK(musaMemset(fresh_var_data, 0, channel_size));
    
    // 4. 设置 BatchNorm 参数
    MUDNN_CHECK(mudnnSetBatchNormDescriptor(
        bn_desc, 
        MUDNN_BATCHNORM_PER_ACTIVATION, 
        1e-5,  // epsilon
        true   // isTraining = true for training mode
    ));
    
    // 5. 查询工作区大小
    size_t workspace_size;
    MUDNN_CHECK(mudnnGetBatchNormForwardTrainingWorkspaceSize(
        handle, bn_desc,
        output_desc, input_desc,
        acc_mean_desc, acc_var_desc,
        fresh_mean_desc, fresh_var_desc,
        scale_desc, bias_desc,
        &workspace_size));
    
    void* workspace = NULL;
    if (workspace_size > 0) {
        MUSACHECK(musaMalloc(&workspace, workspace_size));
    }
    
    // 6. 执行 BatchNorm（训练模式）
    // 注意：参数顺序为 outputDesc, output, inputDesc, input, ...
    MUDNN_CHECK(mudnnBatchNormForwardTraining(
        handle, bn_desc,
        output_desc, output_data,
        input_desc, input_data,
        acc_mean_desc, acc_mean_data,
        acc_var_desc, acc_var_data,
        fresh_mean_desc, fresh_mean_data,
        fresh_var_desc, fresh_var_data,
        scale_desc, scale_data,
        bias_desc, bias_data,
        0.9,                        // momentum
        workspace, workspace_size
    ));
    
    // 7. 同步流
    MUSACHECK(musaStreamSynchronize(musaStreamPerThread));
    
    printf("BatchNorm executed successfully!\n");
    
    // 8. 清理资源
    MUSACHECK(musaFree(input_data));
    MUSACHECK(musaFree(output_data));
    MUSACHECK(musaFree(scale_data));
    MUSACHECK(musaFree(bias_data));
    MUSACHECK(musaFree(acc_mean_data));
    MUSACHECK(musaFree(acc_var_data));
    MUSACHECK(musaFree(fresh_mean_data));
    MUSACHECK(musaFree(fresh_var_data));
    if (workspace) MUSACHECK(musaFree(workspace));
    
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(input_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(output_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(scale_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(bias_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(acc_mean_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(acc_var_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(fresh_mean_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(fresh_var_desc));
    MUDNN_CHECK(mudnnDestroyBatchNormDescriptor(bn_desc));
    MUDNN_CHECK(mudnnDestroy(handle));
    
    return 0;
}

核心概念

Handle（设备句柄）

mudnnHandle_t 是 muDNN 与 MUSA 设备交互的核心接口：

mudnnHandle_t handle;
mudnnCreate(&handle);

// 设置流
musaStream_t stream;
musaStreamCreate(&stream);
mudnnSetStream(handle, stream);

// 清理
mudnnDestroy(handle);
musaStreamDestroy(stream);

Tensor Descriptor（张量描述符）

张量描述符描述多维数组的元数据：

mudnnTensorDescriptor_t tensor_desc;
mudnnCreateTensorDescriptor(&tensor_desc);

// 4D 张量 (NCHW)
mudnnSetTensor4dDescriptor(
    tensor_desc,
    MUDNN_TENSOR_NCHW,      // 格式
    MUDNN_DATA_FLOAT,       // 数据类型
    N, C, H, W              // 维度
);

// 2D 张量
const int dims[2] = {rows, cols};
const int strides[2] = {cols, 1};
mudnnSetTensorNdDescriptor(
    tensor_desc,
    MUDNN_DATA_FLOAT,       // 数据类型
    2,                      // 维度数量
    dims,
    strides
);

// 清理
mudnnDestroyTensorDescriptor(tensor_desc);

Workspace（工作区）

muDNN 算子需要临时工作区内存：

// 以 LayerNorm 为例：查询工作区大小
size_t workspace_size;
mudnnGetLayerNormForwardWorkspaceSize(
    handle, layernorm_desc, 
    output_desc, mean_desc, variance_desc,
    input_desc, gamma_desc, beta_desc,
    &workspace_size
);

// 2. 分配工作区
void* workspace;
if (workspace_size > 0) {
    musaMalloc(&workspace, workspace_size);
}

// 3. 使用工作区执行算子
mudnnLayerNormForward(
    handle, layernorm_desc,
    output_desc, output_data,
    mean_desc, mean_data,
    variance_desc, variance_data,
    input_desc, input_data,
    gamma_desc, gamma_data,
    beta_desc, beta_data,
    workspace, workspace_size
);

最佳实践

错误处理

始终检查 muDNN 函数返回状态：

#define MUDNN_CHECK(cmd)                                              \
  do {                                                                \
    mudnnStatus_t status = cmd;                                       \
    if (status != MUDNN_STATUS_SUCCESS) {                             \
      fprintf(stderr, "muDNN error %d at %s:%d\n", status,            \
              __FILE__, __LINE__);                                    \
      exit(EXIT_FAILURE);                                             \
    }                                                                 \
  } while (0)

资源管理

使用宏简化错误检查和资源清理：

// 在函数开始处创建所有资源
mudnnHandle_t handle;
mudnnTensorDescriptor_t input_desc, output_desc;
void* workspace = NULL;

MUDNN_CHECK(mudnnCreate(&handle));
MUDNN_CHECK(mudnnCreateTensorDescriptor(&input_desc));
MUDNN_CHECK(mudnnCreateTensorDescriptor(&output_desc));

// 使用 goto 统一清理
cleanup:
  if (input_desc) mudnnDestroyTensorDescriptor(input_desc);
  if (output_desc) mudnnDestroyTensorDescriptor(output_desc);
  if (workspace) musaFree(workspace);
  if (handle) mudnnDestroy(handle);

性能优化

提示

• 优先使用 API 参考中未标记 deprecated 的接口，并按对应 workspace 查询函数分配临时内存
• 重用工作区内存，避免重复分配
• 批处理操作以减少内核启动开销
• 使用 mudnnSetStream 绑定 MUSA 流以实现并发执行

与 AI 框架对接扩展

muDNN 可以与主流 AI 框架（PyTorch、TensorFlow）对接，使框架能够使用 muDNN 算子库在 MTGPU 上运行 AI 模型。

本节代码用于说明框架扩展中调用 muDNN C API 的基本方式，框架侧设备注册、内存分配、流管理和构建脚本需要结合实际工程适配。

torch_musa - muDNN 对接扩展

本节介绍如何将 MTGPU 设备通过 muDNN 库和 torch_musa 集成到 PyTorch 框架中。

算子派发机制（Dispatch）

PyTorch 中的 dispatcher 是一个分发器，当执行一个 operator 时，dispatcher 会根据 tensor 的输入信息和其他信息（参数个数、返回值类型等）计算得到一个 dispatch key，然后根据 dispatch key 从 dispatch table 中找到对应的 kernel 函数指针，最后回调执行。

dispatch key 不仅有硬件后端，还有一些更抽象的概念如 autograd、tracing 等。当同一个算子有不同 dispatch key 的实现时，PyTorch 会调用优先级最高的 dispatch key 对应的 kernel 实现。

设备注册流程

要让 PyTorch 支持 MTGPU 设备，需要完成以下步骤：

1. 添加 backend

   • 在 torch/c10/core/Backend.h 的 enum class Backend 中添加 MUSA
   • 在 backendToDispatchKey() 中将 backend 与 DispatchKey(PrivateUse1) 绑定
   • 在 backendToDeviceType() 中将 backend 与 DeviceType(MTGPU) 绑定

2. 添加 device

   • 在 torch/c10/core/DeviceType.h 的 enum class DeviceType 中添加 MTGPU
   • 在 torch/c10/core/Device.h 的 struct Device 中添加 is_mtgpu()
   • 在 torch/c10/core/DispatchKey.cpp 的 toBackendComponent() 中添加 MTGPU 对应的 BackendComponent
   • 在 torch/c10/core/TensorOptions.h 的 computeDispatchKey() 中根据 MTGPU 获得对应的 DispatchKey
   • 在 torch/torch/library.h 的 dispatch 中将 Device_type MTGPU 与 DispatchKey::PrivateUse1 绑定

算子对接流程

以 relu 算子为例，介绍算子库对接的基本流程：

1. 查看 muDNN 是否支持该算子

   • relu 属于一元算子，支持 MUDNN_UNARY_RELU 模式

2. 查找 PyTorch 中的 dispatch 定义

   • 在 torch/include/ATen/native/native_functions.yaml 中搜索 relu：

   - func: relu(Tensor self) -> Tensor
     device_check: NoCheck
     variants: function, method
     dispatch:
       CPU, CUDA: relu
       MPS: relu_mps

3. 获取函数原型

   • 在 torch/include/ATen/ops/relu.h 中找到：Tensor relu(const at::Tensor & self)

4. 注册 MTGPU 上的算子

   #include <ATen/ATen.h>
   #include <mudnn.h>
   #include <vector>
   
   at::Tensor musa_relu(const at::Tensor& input) {
       // 步骤 1: 创建输出 tensor
       at::Tensor result = at::native::empty_mtgpu(input.sizes(), c10::DeviceType::MTGPU, ...);
        
       // 步骤 2: 创建 muDNN 描述符
       mudnnHandle_t handle;
       mudnnTensorDescriptor_t input_desc, output_desc;
       mudnnUnaryDescriptor_t unary_desc;
       
       mudnnCreate(&handle);
       mudnnCreateTensorDescriptor(&input_desc);
       mudnnCreateTensorDescriptor(&output_desc);
       mudnnCreateUnaryDescriptor(&unary_desc);
       
       // 设置张量描述符
       std::vector<int> input_dims(input.dim());
       std::vector<int> input_strides(input.dim());
       std::vector<int> output_dims(result.dim());
       std::vector<int> output_strides(result.dim());
       const int nb_dims = static_cast<int>(input.dim());
       for (int i = 0; i < nb_dims; ++i) {
           input_dims[i] = static_cast<int>(input.size(i));
           input_strides[i] = static_cast<int>(input.stride(i));
           output_dims[i] = static_cast<int>(result.size(i));
           output_strides[i] = static_cast<int>(result.stride(i));
       }
       mudnnSetTensorNdDescriptor(input_desc, MUDNN_DATA_FLOAT, nb_dims,
                                  input_dims.data(), input_strides.data());
       mudnnSetTensorNdDescriptor(output_desc, MUDNN_DATA_FLOAT, nb_dims,
                                  output_dims.data(), output_strides.data());
       
       // 设置一元算子参数
       mudnnSetUnaryDescriptor(unary_desc, MUDNN_UNARY_RELU,
                               MUDNN_PARAM_DATA_NONE, nullptr,
                               MUDNN_PARAM_DATA_NONE, nullptr);
       
       // 步骤 3: 获取数据指针
       void* input_ptr = input.data_ptr();
       void* output_ptr = result.data_ptr();
       
       // 步骤 4: 执行算子
       mudnnUnary(handle, unary_desc, output_desc, output_ptr, input_desc, input_ptr);
       
       // 清理
       mudnnDestroyUnaryDescriptor(unary_desc);
       mudnnDestroyTensorDescriptor(input_desc);
       mudnnDestroyTensorDescriptor(output_desc);
       mudnnDestroy(handle);
   
       return result;
   }
   
   // 注册算子
   TORCH_LIBRARY_IMPL(aten, MTGPU, m) {
       m.impl("relu", &musa_relu);
   }

注意

由于 PyTorch 的 tensor 的 storage 具有 offset 属性，当多个 tensor 共用一块 storage 时，tensor 的 offset 可能不为零。此外，muDNN 目前不支持不连续的 tensor，因此在自定义函数中需要进行连续性检测和处理。

算子测试

算子注册完成后，需要进行正确性测试：

单算子测试:

import torch
import torch_musa

data = torch.randn(3, 4)
out = torch.relu(data)

with torch.autograd.inference_mode(mode=True):
    data_mt = data.to("musa")
    out_mt = torch.relu(data_mt)
    out_c = out_mt.to("cpu")
    
# 比较 out 和 out_c 是否一致

模型测试:

import torch
import torch_musa
from torchvision.models import alexnet

model = alexnet(weights=None).eval()
model_mt = model.to("musa")

x = torch.rand((1, 3, 224, 224))
x_mt = x.to("musa")

with torch.autograd.inference_mode(mode=True):
    y = model(x)
    y_mt = model_mt(x_mt)
    y_c = y_mt.to("cpu")
    
# 比较 y 和 y_c 是否一致

muTensorFlow - muDNN 对接扩展

Kernel 注册

在 TensorFlow 中，Op 类似于函数声明，Kernel 类似于函数实现。自定义 Kernel 需要继承从 OpKernel 类并重写 Compute 成员函数：

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include <mudnn.h>
#include <vector>

class VecAddOp : public tensorflow::OpKernel {
 public:
  explicit VecAddOp(tensorflow::OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(tensorflow::OpKernelContext* context) override {
    const tensorflow::Tensor& input_0 = context->input(0);
    const tensorflow::Tensor& input_1 = context->input(1);
    
    // 检查输入数据类型是否相同
    OP_REQUIRES(context, input_0.dtype() == tensorflow::DT_FLOAT &&
                         input_1.dtype() == tensorflow::DT_FLOAT,
                tensorflow::errors::InvalidArgument("input data type should be float"));
    
    tensorflow::Tensor* output = nullptr;
    // 步骤 1: 给 output 分配内存空间
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_0.shape(), &output));
    
    // 步骤 2: 创建 muDNN 描述符
    mudnnHandle_t handle;
    mudnnTensorDescriptor_t in1_desc, in2_desc, out_desc;
    mudnnBinaryDescriptor_t binary_desc;
    
    mudnnCreate(&handle);
    mudnnCreateTensorDescriptor(&in1_desc);
    mudnnCreateTensorDescriptor(&in2_desc);
    mudnnCreateTensorDescriptor(&out_desc);
    mudnnCreateBinaryDescriptor(&binary_desc);
    
    // 设置张量描述符
    auto set_tensor_desc = [](mudnnTensorDescriptor_t desc, const tensorflow::Tensor& tensor) {
      const int nb_dims = tensor.dims();
      std::vector<int> dims(nb_dims);
      std::vector<int> strides(nb_dims, 1);
      for (int i = 0; i < nb_dims; ++i) {
        dims[i] = static_cast<int>(tensor.dim_size(i));
      }
      for (int i = nb_dims - 2; i >= 0; --i) {
        strides[i] = strides[i + 1] * dims[i + 1];
      }
      mudnnSetTensorNdDescriptor(desc, MUDNN_DATA_FLOAT, nb_dims,
                                 dims.data(), strides.data());
    };
    set_tensor_desc(in1_desc, input_0);
    set_tensor_desc(in2_desc, input_1);
    set_tensor_desc(out_desc, *output);
    
    // 设置二元算子参数
    mudnnSetBinaryDescriptor(binary_desc, MUDNN_BINARY_ADD,
                             MUDNN_PARAM_DATA_NONE, nullptr,
                             MUDNN_PARAM_DATA_NONE, nullptr);
    
    // 步骤 3: 获取数据指针
    void* in1_ptr = input_0.flat<float>().data();
    void* in2_ptr = input_1.flat<float>().data();
    void* out_ptr = output->flat<float>().data();
    
    // 步骤 4: 执行二元加法
    mudnnBinary(handle, binary_desc, out_desc, out_ptr, in1_desc, in1_ptr, in2_desc, in2_ptr);
    
    // 清理
    mudnnDestroyBinaryDescriptor(binary_desc);
    mudnnDestroyTensorDescriptor(in1_desc);
    mudnnDestroyTensorDescriptor(in2_desc);
    mudnnDestroyTensorDescriptor(out_desc);
    mudnnDestroy(handle);
  }
};

// 注册 Kernel
MUSA_KERNEL_REGISTER(VecAdd) {
  REGISTER_KERNEL_BUILDER(tensorflow::Name("VecAdd").Device(tensorflow::DEVICE_MTGPU), VecAddOp);
}

对于自定义 kernel 的代码，需要将其放在 muTensorFlow/musa-plugin/src/kernels 目录下，随后在 muTensorFlow/musa-plugin 目录下调用 python setup.py install 完成整个 musa-plugin 的安装。

Op 注册

自定义 Op 和设备没有关系，参考 TensorFlow 的官方教程即可：

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("VecAdd")
    .Input("in1: float")   // 输入 0，类型为 float
    .Input("in2: float")   // 输入 1，类型为 float
    .Output("out: float")  // 输出 0，类型为 float
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));  // 设置输出 0 的 shape
      return tensorflow::Status::OK();
    });

编译为动态库：

TF_CFLAGS=$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))')
TF_LFLAGS=$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))')
g++ -std=c++14 -shared vec_add_ops.cc -o _vec_add_ops.so -fPIC \
    ${TF_CFLAGS} ${TF_LFLAGS} \
    -I${MUSA_SDK_PATH}/include -L${MUSA_SDK_PATH}/lib \
    -lmudnn -lmudnn_ops -lmudnn_base -lmudnn_tensor \
    -lmudnn_tensor_binary -lmusart -Wl,-rpath,${MUSA_SDK_PATH}/lib -O2

加载动态库并使用：

import tensorflow as tf
vec_add_module = tf.load_op_library('./_vec_add_ops.so')

result = vec_add_module.vec_add([[1, 2], [3, 4]], [[1, 2], [3, 4]])
print(result)

图像相关算子

muDNN 提供部分可用于图像处理流程的算子，例如：

池化

• 最大池化：取邻域内的最大值
• 平均池化：取邻域内的平均值

插值

• 最邻近插值：赋予离目标点最近的原像素值
• 双线性插值：根据相邻的四个点进行插值
• 双三次插值：用 16 个相邻点做插值，效果更平滑

坐标变换

• 平移：将输入图像坐标平移
• 旋转：旋转图像坐标
• 缩放：缩放图像坐标

具体支持范围、参数约束和 deprecated 状态以 muDNN C API 参考 为准。

API 参考

详细的 API 文档请参阅：

• muDNN C API 参考 - 推荐，长期支持
• muDNN C++ API 参考 - 已弃用，仅供参考