跳到主要内容

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_tmuDNN 设备句柄，绑定到 MUSA 设备与流
mudnnTensorDescriptor_t张量描述符，描述数据地址、类型、格式、维度
mudnnFilterDescriptor_t滤波器/权重描述符
mudnnCreateXXXDescriptor创建各类描述符
mudnnDestroyXXXDescriptor销毁各类描述符

张量算子：

算子函数前缀
一元算子mudnnUnary
二元算子mudnnBinary
三元算子mudnnTernary
拼接mudnnConcat
填充mudnnFill
置换mudnnPermute
归约mudnnReduce

mudnn_cnn.h - 神经网络算子​

算子类型函数前缀
卷积mudnnConvolution
池化mudnnPooling
归一化mudnnBatchNorm, mudnnLayerNorm, mudnnGroupNorm
激活mudnnActivation, mudnnSoftmax
DropoutmudnnDropout

mudnn_adv.h - 高级算子​

算子类型函数前缀
注意力mudnnMultiHeadAttention, mudnnScaledDotProductAttention
损失mudnnCTCLoss
矩阵乘法mudnnGemm

快速开始​

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

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

#include <stdio.h>
#include <stdlib.h>
#include <musa_runtime_api.h>
#include <mudnn/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));
    
    // 输出/../gamma/beta 使用相同维度
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        output_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        mean_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, N, C, H, W));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        variance_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));
    
    // 设置 LayerNorm 描述符
    int axes[] = {1, 2, 3};  // 对 C, H, W 维度归一化
    MUDNN_CHECK(mudnnSetLayerNormDescriptor(
        layernorm_desc, 
        1e-5,  // epsilon
        MUDNN_LAYERNORM_MODE_ELEMENTWISE,  // 模式
        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 channel_size = C * sizeof(float);
    
    MUSACHECK(musaMalloc((void**)&input_data, size));
    MUSACHECK(musaMalloc((void**)&output_data, size));
    MUSACHECK(musaMalloc((void**)&mean_data, size));
    MUSACHECK(musaMalloc((void**)&variance_data, size));
    MUSACHECK(musaMalloc((void**)&gamma_data, channel_size));
    MUSACHECK(musaMalloc((void**)&beta_data, channel_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;
}

编译命令:

说明

编译依赖于 MUSA SDK 的 musa_asm 组件，请确保已正确安装。

mcc layernorm_example.c \
    -I${MUSA_SDK_PATH}/include \
    -L${MUSA_SDK_PATH}/lib \
    -lmudnn -lmusa_runtime \
    -Wl,-rpath,${MUSA_SDK_PATH}/lib \
    -o layernorm_example

运行:

./layernorm_example

示例 2：GEMM（使用 muDNN）​

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

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <vector>
#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)

namespace {

std::vector<int> ContiguousStrides(const std::vector<int>& shape) {
  std::vector<int> stride(shape.size(), 1);
  for (int i = static_cast<int>(shape.size()) - 2; i >= 0; --i) {
    stride[static_cast<size_t>(i)] =
        stride[static_cast<size_t>(i + 1)] * shape[static_cast<size_t>(i + 1)];
  }
  return stride;
}

void SetTensorDesc3D(mudnnTensorDescriptor_t desc, mudnnDataType_t dtype,
                    const std::vector<int>& shape) {
  const std::vector<int> stride = ContiguousStrides(shape);
  MUDNN_CHECK(mudnnSetTensorNdDescriptor(desc, dtype, static_cast<int>(shape.size()),
                                         shape.data(), stride.data()));
}

}  // namespace

int main() {
  std::cout << "muDNN version: " << mudnnGetVersion() << std::endl;
  MUSACHECK(musaSetDevice(0));

  constexpr int batch = 1;
  constexpr int m = 2;
  constexpr int k = 3;
  constexpr int n = 4;

  const std::vector<float> h_a_f32 = {
      1.0f, 2.0f, 3.0f,
      4.0f, 5.0f, 6.0f,
  };
  const std::vector<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,
  };
  std::vector<float> h_d(static_cast<size_t>(batch * m * n), 0.0f);

  void* d_a = nullptr;
  void* d_b = nullptr;
  void* d_d = nullptr;

  MUSACHECK(musaMalloc(&d_a, h_a_f32.size() * sizeof(float)));
  MUSACHECK(musaMalloc(&d_b, h_b_f32.size() * sizeof(float)));
  MUSACHECK(musaMalloc(&d_d, h_d.size() * sizeof(float)));

  MUSACHECK(musaMemcpy(d_a, h_a_f32.data(), h_a_f32.size() * sizeof(float),
                      musaMemcpyHostToDevice));
  MUSACHECK(musaMemcpy(d_b, h_b_f32.data(), h_b_f32.size() * sizeof(float),
                      musaMemcpyHostToDevice));
  MUSACHECK(musaMemset(d_d, 0, h_d.size() * sizeof(float)));

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

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

  SetTensorDesc3D(a_desc, MUDNN_DATA_FLOAT, {batch, m, k});
  SetTensorDesc3D(b_desc, MUDNN_DATA_FLOAT, {batch, k, n});
  SetTensorDesc3D(d_desc, MUDNN_DATA_FLOAT, {batch, m, n});

  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,
                                         nullptr, d_desc,
                                         &workspace_size));

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

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

  MUSACHECK(musaMemcpy(h_d.data(), d_d, h_d.size() * sizeof(float),
                      musaMemcpyDeviceToHost));

  MUSACHECK(musaDeviceSynchronize());

  if (workspace != nullptr) {
    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 <stdio.h>
#include <stdlib.h>
#include <musa_runtime_api.h>
#include <mudnn/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 mean_desc, variance_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(&mean_desc));
    MUDNN_CHECK(mudnnCreateTensorDescriptor(&variance_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 和 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(
        mean_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    MUDNN_CHECK(mudnnSetTensor4dDescriptor(
        variance_desc, MUDNN_TENSOR_NCHW, MUDNN_DATA_FLOAT, 1, C, 1, 1));
    
    // 3. 分配设备内存
    float *input_data, *output_data, *scale_data, *bias_data, *mean_data, *variance_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**)&mean_data, channel_size));
    MUSACHECK(musaMalloc((void**)&variance_data, channel_size));
    
    // 4. 设置 BatchNorm 参数（使用 v5.1.0 API）
    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, input_desc, &workspace_size));
    
    void* workspace = NULL;
    if (workspace_size > 0) {
        MUSACHECK(musaMalloc(&workspace, workspace_size));
    }
    
    // 6. 执行 BatchNorm（训练模式）- 使用 v5.1.0 API
    // 注意：参数顺序为 outputDesc, output, inputDesc, input, ...
    MUDNN_CHECK(mudnnBatchNormForwardTraining(
        handle, bn_desc,
        output_desc, output_data,
        input_desc, input_data,
        mean_desc, mean_data,       // accumulated mean
        variance_desc, variance_data,  // accumulated variance
        mean_desc, mean_data,       // fresh mean (same buffer for simplicity)
        variance_desc, variance_data,  // fresh variance (same buffer for simplicity)
        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(mean_data));
    MUSACHECK(musaFree(variance_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(mean_desc));
    MUDNN_CHECK(mudnnDestroyTensorDescriptor(variance_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 张量
mudnnSetTensor2dDescriptor(
    tensor_desc,
    MUDNN_DATA_FLOAT,       // 数据类型
    rows, cols              // 维度
);

// 清理
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);

性能优化​

提示
  • 选择合适的卷积算法（使用 mudnnFindConvolutionForwardAlgorithm）
  • 重用工作区内存，避免重复分配
  • 批处理操作以减少内核启动开销
  • 使用 mudnnSetStream 绑定 MUSA 流以实现并发执行

与 AI 框架对接扩展​

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

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_ACTIVATION_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/mudnn.h>

    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;
        mudnnActivationDescriptor_t activ_desc;
        
        mudnnCreate(&handle);
        mudnnCreateTensorDescriptor(&input_desc);
        mudnnCreateTensorDescriptor(&output_desc);
        mudnnCreateActivationDescriptor(&activ_desc);
        
        // 设置张量描述符
        mudnnSetTensorDescriptor(input_desc, MUDNN_DATA_FLOAT, input.dim(), input.sizes().data());
        mudnnSetTensorDescriptor(output_desc, MUDNN_DATA_FLOAT, result.dim(), result.sizes().data());
        
        // 设置激活函数参数
        mudnnSetActivationDescriptor(activ_desc, MUDNN_ACTIVATION_RELU, 0.0, 0.0);
        
        // 步骤 3: 获取数据指针
        void* input_ptr = input.data_ptr();
        void* output_ptr = result.data_ptr();
        
        // 步骤 4: 执行算子
        float alpha = 1.0f;
        mudnnActivationForward(handle, activ_desc, &alpha, input_desc, input_ptr, &alpha, output_desc, output_ptr);
        
        // 清理
        mudnnDestroyActivationDescriptor(activ_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(pretrained=True).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/mudnn.h>

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() == input_1.dtype(),
                tensorflow::errors::InvalidArgument("input data type should be the same!"));
    
    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);
    
    // 设置张量描述符
    mudnnSetTensorDescriptor(in1_desc, MUDNN_DATA_FLOAT, input_0.dims(), input_0.shape().data());
    mudnnSetTensorDescriptor(in2_desc, MUDNN_DATA_FLOAT, input_1.dims(), input_1.shape().data());
    mudnnSetTensorDescriptor(out_desc, MUDNN_DATA_FLOAT, output->dims(), output->shape().data());
    
    // 设置二元算子参数
    mudnnSetBinaryDescriptor(binary_desc, MUDNN_BINARY_ADD, MUDNN_DATA_FLOAT);
    
    // 步骤 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: 执行二元加法
    float alpha = 1.0f;
    mudnnBinary(handle, binary_desc, &alpha, in1_desc, in1_ptr, in2_desc, in2_ptr, &alpha, out_desc, out_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: int32")   // 输入 0，类型为 int32
    .Input("in2: int32")   // 输入 1，类型为 int32
    .Output("out: int32")  // 输出 0，类型为 int32
    .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[@]} -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 个相邻点做插值，效果更平滑

坐标变换​

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

muDNN 图像预处理算子仅支持池化、插值、坐标变换，不支持其他操作。

API 参考​

详细的 API 文档请参阅：

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。