NVIDIA cuDNN

• 加速学习：cuDNN 提供了针对 Tensor Core 优化的内核，在计算密集型操作中实现最佳性能，并为不同问题规模挑选合适的内核提供启发式算法。

• 融合支持：cuDNN 支持将计算密集型和内存密集型操作进行融合。常见的通用融合模式通常通过运行时内核生成实现，特殊的融合模式则使用预编写的优化内核。

• 表达性算子图 API：用户可以将计算定义为张量上的操作图。cuDNN 库既有直接的 C API，也有开源的 C++ 前端以提升使用便利性。大多数用户选择前端作为使用 cuDNN 的入口。

cuDNN API 代码示例

该代码使用 cuDNN 与 PyTorch 集成，实现了带有偏置的批量矩阵乘法。

import torch
import cudnn

# Prepare sample input data. nvmath-python accepts input tensors from pytorch, cupy, and
# numpy.
b, m, n, k = 1, 1024, 1024, 512
A = torch.randn(b, m, k, dtype=torch.float32, device="cuda")
B = torch.randn(b, k, n, dtype=torch.float32, device="cuda")
bias = torch.randn(b, m, 1, dtype=torch.float32, device="cuda")

result = torch.empty(b, m, n, dtype=torch.float32, device="cuda")

# Use the stateful Graph object in order to perform multiple matrix multiplications
# without replanning. The cudnn API allows us to fine-tune our operations by, for
# example, selecting a mixed-precision compute type.
graph = cudnn.pygraph(
   intermediate_data_type=cudnn.data_type.FLOAT,
   compute_data_type=cudnn.data_type.FLOAT,
)

a_cudnn_tensor    = graph.tensor_like(A)
b_cudnn_tensor    = graph.tensor_like(B)
bias_cudnn_tensor = graph.tensor_like(bias)

c_cudnn_tensor = graph.matmul(name="matmul", A=a_cudnn_tensor, B=b_cudnn_tensor)
d_cudnn_tensor = graph.bias(name="bias", input=c_cudnn_tensor, bias=bias_cudnn_tensor)

# Build the matrix multiplication. Building returns a sequence of algorithms that can be
# configured. Each algorithm is a JIT generated function that can be executed on the GPU.

graph.build([cudnn.heur_mode.A])
workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)

# Execute the matrix multiplication.
graph.execute(
   {
       a_cudnn_tensor: A,
       b_cudnn_tensor: B,
       bias_cudnn_tensor: bias,
       d_cudnn_tensor: result,
   },
   workspace
)

cuDNN Graph API 描述的操作图示例

深度神经网络

深度学习神经网络广泛应用于计算机视觉、对话式 AI 以及推荐系统，并推动了诸如智能驾驶与智能语音助手等突破性进展。NVIDIA 的 GPU 加速深度学习框架显著缩短了这些技术的训练时间，将数天的训练过程缩短到数小时。

cuDNN 为云端、嵌入式设备和智能驾驶汽车中的深度神经网络提供高性能、低延迟的推理基础库。

• 加速 attention 训练／预填、卷积和矩阵乘法（matmul）等计算密集型操作
  

• 优化如 attention 解码、池化、softmax、归一化、激活、逐点操作、张量变换等内存密集型操作
  

• 支持计算密集型与内存密集型操作的融合
  

• 提供运行时融合引擎，可为常见融合模式在运行时生成内核
  

• 针对如融合 attention 等重要专用模式进行优化

• 根据具体问题规模应用启发式算法，选择合适的实现