AT_DISPATCH_FLOATING_TYPES 宏主要用于以下目的：

1. 数据类型调度：根据输入张量的数据类型选择合适的数据类型进行处理。
2. 模板编程：结合 C++ 模板编程，根据不同的数据类型生成不同的代码路径。
3. 代码简化：减少手动写类型检查和类型转换代码的繁琐过程。

使用示例

假设我们有一个简单的 CUDA 内核函数 example_kernel，它对输入张量进行某种操作。我们希望这个内核函数可以处理 float 和 double 类型的数据。以下是如何使用 AT_DISPATCH_FLOATING_TYPES 来实现这个目标。

1. 定义 CUDA 内核

首先，我们定义一个简单的 CUDA 内核函数：

template <typename scalar_t>
__global__ void example_kernel(scalar_t* data, int64_t size) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < size) {
    data[index] *= 2; // 例如，简单地将每个元素乘以2
  }
}

2. 定义 C++ 函数并使用 AT_DISPATCH_FLOATING_TYPES

接下来，我们定义一个 C++ 函数，使用 AT_DISPATCH_FLOATING_TYPES 来调度数据类型，并调用相应的 CUDA 内核：

#include <torch/extension.h>
#include <vector>

std::vector<torch::Tensor> example_forward(torch::Tensor input) {
  const auto size = input.size(0);
  auto output = torch::zeros_like(input);

  const int threads = 1024;
  const int blocks = (size + threads - 1) / threads;

  AT_DISPATCH_FLOATING_TYPES(input.type(), "example_forward_cuda", ([&] {
    example_kernel<scalar_t><<<blocks, threads>>>(
        input.data<scalar_t>(), size);
  }));

  return {output};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &example_forward, "Example forward");
}

在这段代码中：

• AT_DISPATCH_FLOATING_TYPES(input.type(), "example_forward_cuda", ([&] { ... }))：
  • input.type()：获取输入张量的数据类型。
  • "example_forward_cuda"：操作的名称，用于错误信息。
  • [&]：捕获外部变量的 lambda 表达式。
  • example_kernel<scalar_t><<<blocks, threads>>>(input.data<scalar_t>(), size)：根据调度的数据类型调用相应的 CUDA 内核。

3. 在 Python 中调用

最后，我们在 Python 中加载和调用这个扩展：

import torch
from torch.utils.cpp_extension import load

# JIT 编译并加载 C++ 扩展
example_cpp = load(name="example_cpp", sources=["example.cpp"], verbose=True)

# 创建输入张量
input = torch.randn(1024, device='cuda', dtype=torch.float32)

# 调用前向传播函数
output = example_cpp.forward(input)

print(output)

另一个示例：支持多种数据类型

假设我们希望实现一个支持 float、double 和 half 类型的前向传播函数。我们可以使用 AT_DISPATCH_FLOATING_TYPES_AND_HALF 宏：

#include <torch/extension.h>
#include <vector>

template <typename scalar_t>
__global__ void example_kernel(scalar_t* data, int64_t size) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < size) {
    data[index] *= 2; // 简单地将每个元素乘以2
  }
}

std::vector<torch::Tensor> example_forward(torch::Tensor input) {
  const auto size = input.size(0);
  auto output = torch::zeros_like(input);

  const int threads = 1024;
  const int blocks = (size + threads - 1) / threads;

  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "example_forward_cuda", ([&] {
    example_kernel<scalar_t><<<blocks, threads>>>(
        input.data<scalar_t>(), size);
  }));

  return {output};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &example_forward, "Example forward");
}

AT_DISPATCH_ALL_TYPES 是 PyTorch 提供的一个宏，用于在编译时调度所有基础数据类型，并执行相应的数据类型特定的代码。这对于编写可以处理多种数据类型的高效代码特别有用，尤其是在需要调用 CUDA 内核时。

让我们通过一个完整的例子来展示如何使用 AT_DISPATCH_ALL_TYPES。

使用示例

假设我们有一个简单的 CUDA 内核函数 example_kernel，它对输入张量进行某种操作。我们希望这个内核函数可以处理所有基础数据类型（例如 int, float, double 等）。以下是如何使用 AT_DISPATCH_ALL_TYPES 来实现这个目标。

1. 定义 CUDA 内核

首先，我们定义一个简单的 CUDA 内核函数：

// example_kernel.cu
template <typename scalar_t>
__global__ void example_kernel(scalar_t* data, int64_t size) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index < size) {
    data[index] *= 2; // 例如，简单地将每个元素乘以2
  }
}

2. 定义 C++ 函数并使用 AT_DISPATCH_ALL_TYPES

接下来，我们定义一个 C++ 函数，使用 AT_DISPATCH_ALL_TYPES 来调度数据类型，并调用相应的 CUDA 内核：

// example.cpp
#include <torch/extension.h>
#include <vector>

std::vector<torch::Tensor> example_forward(torch::Tensor input) {
  const auto size = input.size(0);
  auto output = torch::zeros_like(input);

  const int threads = 1024;
  const int blocks = (size + threads - 1) / threads;

  AT_DISPATCH_ALL_TYPES(input.scalar_type(), "example_forward_cuda", ([&] {
    example_kernel<scalar_t><<<blocks, threads>>>(
        input.data<scalar_t>(), size);
  }));

  return {output};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &example_forward, "Example forward");
}

在这段代码中：

• AT_DISPATCH_ALL_TYPES(input.scalar_type(), "example_forward_cuda", ([&] { ... }))：
  • input.scalar_type()：获取输入张量的数据类型。
  • "example_forward_cuda"：操作的名称，用于错误信息。
  • [&]：捕获外部变量的 lambda 表达式。
  • example_kernel<scalar_t><<<blocks, threads>>>(input.data<scalar_t>(), size)：根据调度的数据类型调用相应的 CUDA 内核。

3. 在 Python 中调用

最后，我们在 Python 中加载和调用这个扩展：

import torch
from torch.utils.cpp_extension import load

# JIT 编译并加载 C++ 扩展
example_cpp = load(name="example_cpp", sources=["example.cpp", "example_kernel.cu"], verbose=True)

# 创建输入张量
input = torch.randn(1024, device='cuda', dtype=torch.float32)

# 调用前向传播函数
output = example_cpp.forward(input)

print(output)

总结

• 数据类型调度：AT_DISPATCH_ALL_TYPES 宏用于调度所有基础数据类型，包括 int, float, double 等。
• 模板编程：结合 C++ 模板编程，根据不同的数据类型生成不同的代码路径。
• 代码简化：减少手动写类型检查和类型转换代码的繁琐过程。