# metaprograming_in_hpc

**Repository Path**: zhengankun/metaprograming_in_hpc

## Basic Information

- **Project Name**: metaprograming_in_hpc
- **Description**: No description available
- **Primary Language**: C++
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-04-02
- **Last Updated**: 2025-04-02

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 介绍
在HPC（高性能计算）中，我们经常能看到模板元编程或者宏，通过模板参数调用函数，或者通过宏参数。最著名的是cutlass（几乎完全使用模板），以及openblas（很多地方使用了宏）。其实宏和模板都是元编程技术，它们发生在具体的编译之前，通过codegen技术将代码在编译期或者预处理期展开，避免了运行时动态的判断，从而得到更高的性能。


# 模板实现分块矩阵乘法和Activation融合
## 模板实现的kernel源码
1. 激活函数模板
```c++
#include <immintrin.h>
// 激活函数模板基类
template <typename T>
struct Activation {
    static T apply(T x) = delete; // 禁止直接实例化基类
};

// ReLU 特化
template <typename T>
struct ReLU : Activation<T> {
    static T apply(T x) {
        return x > T(0) ? x : T(0);
    }
};

// Sigmoid 特化
template <typename T>
struct Sigmoid : Activation<T> {
    static T apply(T x) {
        return T(1) / (T(1) + std::exp(-x));
    }
};


// AVX2 特化 ReLU
template <>
struct ReLU<float> : Activation<float> {
    static float apply(float x) {
        __m128 vec = _mm_set_ss(x);
        vec = _mm_max_ss(vec, _mm_setzero_ps());
        return _mm_cvtss_f32(vec);
    }
};

enum class ActType  {
    relu,
    sigmoid
};
```
2. 分块矩阵乘法和主体计算模板
```c++
#include <cmath>
#include <type_traits>
#include "activation.h"

// 基础内层kernel模板（需特化）
template <int BLOCK_M, int BLOCK_N, int BLOCK_K, typename T, template <typename> class ActivationFunc>
struct BlockGEMM
{
    static void compute(const T *A, const T *B, T *C,
                        int M, int N, int K,
                        int m_start, int n_start, int k_start)
    {
        {
            // 手动展开的SIMD优化计算逻辑
            for (int i = 0; i < BLOCK_M; ++i)
            {
                for (int j = 0; j < BLOCK_N; ++j)
                {
                    T sum = C[(m_start + i) * N + (n_start + j)];
                    for (int k = 0; k < BLOCK_N; ++k)
                    {
                        sum += A[(m_start + i) * K + (k_start + k)] *
                               B[(k_start + k) * N + (n_start + j)];
                    }
                    C[(m_start + i) * N + (n_start + j)] = ActivationFunc<T>::apply(sum);
                }
            }
        }
    }
};

// 特化16x16x4分块
template <typename T, template <typename> class ActivationFunc>
struct BlockGEMM<16, 16, 4, T, ActivationFunc>
{
    static void compute(const T *A, const T *B, T *C,
                        int M, int N, int K,
                        int m_start, int n_start, int k_start)
    {
        // 手动展开的SIMD优化计算逻辑
        for (int i = 0; i < 16; ++i)
        {
            for (int j = 0; j < 16; ++j)
            {
                T sum = C[(m_start + i) * N + (n_start + j)];
                for (int k = 0; k < 4; ++k)
                {
                    sum += A[(m_start + i) * K + (k_start + k)] *
                           B[(k_start + k) * N + (n_start + j)];
                }
                C[(m_start + i) * N + (n_start + j)] = ActivationFunc<T>::apply(sum);
            }
        }
    }
};

// 特化32x32x8分块（AVX512优化版本）
template <typename T, template <typename> class ActivationFunc>
struct BlockGEMM<32, 32, 8, T, ActivationFunc>
{
    static void compute(const T *A, const T *B, T *C,
                        int M, int N, int K,
                        int m_start, int n_start, int k_start)
    {
        // 使用AVX512指令的展开实现
        // ...（具体向量化实现逻辑）
    }
};

// 特化64x64x16分块，数据类型为double（AVX512优化版本）
template <template <typename> class ActivationFunc>
struct BlockGEMM<32, 32, 8, double, ActivationFunc>
{
    static void compute(const double *A, const double *B, double *C,
                        int M, int N, int K,
                        int m_start, int n_start, int k_start)
    {
        // 使用AVX512指令的展开实现
        // ...（具体向量化实现逻辑）
    }
};

// 调度器：根据分块策略选择实现
template <int BLOCK_M, int BLOCK_N, int BLOCK_K, typename T, template <typename> class ActivationFunc>
void gemm_impl(const T *A, const T *B, T *C, int M, int N, int K)
{
    static_assert(BLOCK_M > 0 && BLOCK_N > 0 && BLOCK_K > 0,
                  "Block size must be positive");

    for (int m = 0; m < M; m += BLOCK_M)
    {
        int m_end = std::min(m + BLOCK_M, M);
        for (int n = 0; n < N; n += BLOCK_N)
        {
            int n_end = std::min(n + BLOCK_N, N);
            for (int k = 0; k < K; k += BLOCK_K)
            {
                int k_end = std::min(k + BLOCK_K, K);

                BlockGEMM<BLOCK_M, BLOCK_N, BLOCK_K, T, ActivationFunc>::compute(
                    A, B, C, M, N, K, m, n, k);
            }
        }
    }
}

// 入口函数：自动选择最优分块策略
template <typename T, ActType act_type>
void gemm_with_activation(const T *A, const T *B, T *C, int M, int N, int K)
{
    if constexpr (std::is_same_v<T, float>)
    {
        if (M * N * K < 4096)
        {
            if constexpr (act_type == ActType::relu)
            {
                gemm_impl<16, 16, 4, float, ReLU>(A, B, C, M, N, K);
            }
            else
            {
                gemm_impl<16, 16, 4, float, Sigmoid<float>>(A, B, C, M, N, K);
            }
        }
        else
        {
            if constexpr (act_type == ActType::relu)
            {
                gemm_impl<32, 32, 8, float, ReLU>(A, B, C, M, N, K);
            }
            else
            {
                gemm_impl<32, 32, 8, float, Sigmoid<float>>(A, B, C, M, N, K);
            }
        }
    }
    else if constexpr (std::is_same_v<T, double>)
    {
        if constexpr (act_type == ActType::relu)
        {
            gemm_impl<64, 64, 16, double, ReLU<double>>(A, B, C, M, N, K);
        }
        else
        {
            gemm_impl<64, 64, 16, double, Sigmoid<double>>(A, B, C, M, N, K);
        }
    }
}
```
## 使用模板实现kernel的好处
1. **类型安全的泛型编程**
- 通过`template <typename T>`实现数据类型的泛化
- 编译器为每种类型（float/double）生成特化版本
- 示例：`BlockGEMM<32,32,8,double>`和`BlockGEMM<32,32,8,float>`会生成不同的机器指令
- 避免运行时类型判断的开销，如：`if constexpr(std::is_same_v<T, float>)`在编译期完成分支选择

2. **编译时多态优化**
- 通过模板特化实现算法优化：
```cpp
template <>
struct BlockGEMM<32,32,8,double> { // AVX512特化版本
    static void compute(...) {
        // 使用_mm512指令的优化实现
    }
};
```
- 调度器在编译期静态绑定：
```cpp
gemm_impl<32,32,8,float,ReLU>(...);  // 直接调用AVX512优化版本
```

3. **零成本抽象**
- 激活函数的策略模式实现：
```cpp
template <template <typename> class ActivationFunc>
struct BlockGEMM {
    static void compute(...) {
        // 编译期绑定的激活函数调用
        C[...] = ActivationFunc<T>::apply(sum);
    }
};
```
- 对比运行时多态（虚函数）方案：
```cpp
// 伪代码对比
virtual float apply(float x) = 0;  // 会有虚函数调用开销
```

4. **硬件指令级优化**
- 通过特化利用SIMD指令：
```cpp
// AVX512特化版本示例
__m512d vecA = _mm512_load_pd(A + index);
__m512d vecB = _mm512_load_pd(B + index);
__m512d vecC = _mm512_fmadd_pd(vecA, vecB, vecC);
```
- 不同分块大小对应不同寄存器使用策略：
  - 16x16x4：适合L1 Cache
  - 32x32x8：匹配AVX512的8-wide向量

5. **编译期计算与优化**
- 分块参数静态指定：
```cpp
gemm_impl<64,64,16,double,...>(...); // 编译器可展开循环
```
- 静态断言保证参数合法性：
```cpp
static_assert(BLOCK_M > 0, "Block size must be positive");
```

6. **自动策略选择**
- 基于矩阵规模的编译期决策：
```cpp
if constexpr (M*N*K < 4096) { // 编译期计算
    // 选择小分块策略
}
```
- 类型特化的自动派发：
```cpp
else if constexpr (std::is_same_v<T, double>) {
    // 双精度特化路径
}
```

7. **代码生成效率**
- 循环展开模板示例：
```cpp
for(int i=0; i<16; ++i) {     // 固定迭代次数
    for(int j=0; j<16; ++j) { // 可被编译器完全展开
        // SIMD化计算
    }
}
```
- 对比动态分块代码：
```cpp
for(int i=0; i<block_m; ++i) {  // 运行时变量阻碍优化
```

8. **内存访问模式优化**
- 分块参数与内存层次结构的匹配：
  - 16x16x4分块：约4KB (float) 适合L1 Cache
  - 64x64x16分块：约128KB (double) 匹配L2 Cache
- 通过模板参数显式控制数据局部性：
```cpp
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
```
# 宏实现分块矩阵乘法和Activation融合
## 使用宏的源码
1. 激活函数的宏定义
```c
#include <immintrin.h>
#include <math.h>

// 激活函数类型枚举
typedef enum
{
    RELU,
    SIGMOID,
    TAHN
} act_type_t;

// 通用激活函数宏
#define DEFINE_ACTIVATION(TYPE, T, SUFFIX)   \
    static inline T activation_##SUFFIX(T x) \
    {                                        \
        return TYPE##_IMPL(x);               \
    }

// ReLU实现
#define RELU_IMPL(x) ((x) > 0 ? (x) : 0)
// DEFINE_ACTIVATION(RELU, float, relu_f32)
DEFINE_ACTIVATION(RELU, double, relu_f64)

// Sigmoid实现
#define SIGMOID_IMPL(x) (1.0 / (1.0 + exp(-(x))))
DEFINE_ACTIVATION(SIGMOID, float, sigmoid_f32)
DEFINE_ACTIVATION(SIGMOID, double, sigmoid_f64)

// Tanh实现
#define TANH_IMPL(x) tanh(x)
DEFINE_ACTIVATION(TANH, float, tanh_f32)
DEFINE_ACTIVATION(TANH, double, tanh_f64)

// AVX2优化的ReLU特化，使用宏实现，避免了代码可能无法被inline，直接被展开
#define ACTIVATION_RELU_F32(x) ({ \
    __m128 vec = _mm_set_ss(x); \
    vec = _mm_max_ss(vec, _mm_setzero_ps()); \
    _mm_cvtss_f32(vec); \
})
```
2. 分块矩阵乘法和主体计算模板
```c
#include <math.h>
#include <string.h>
#include "activation.h"

// 基础分块计算宏提前特化版本
#define DEFINE_BLOCK_F32_16x16x4(SUFFIX, ACT_FUNC)                                    \
    static void block_gemm_16x16x4_##SUFFIX(const float *A, const float *B, float *C, \
                                            int M, int N, int K,                      \
                                            int m_start, int n_start, int k_start)    \
    {                                                                                 \
        for (int i = 0; i < 16; ++i)                                                  \
        {                                                                             \
            for (int j = 0; j < 16; ++j)                                              \
            {                                                                         \
                float sum = C[(m_start + i) * N + (n_start + j)];                     \
                for (int k = 0; k < 4; ++k)                                           \
                {                                                                     \
                    sum += A[(m_start + i) * K + (k_start + k)] *                     \
                           B[(k_start + k) * N + (n_start + j)];                      \
                }                                                                     \
                C[(m_start + i) * N + (n_start + j)] = ACT_FUNC(sum);                 \
            }                                                                         \
        }                                                                             \
    }

DEFINE_BLOCK_F32_16x16x4(f32_sigmoid, activation_sigmoid_f32)
DEFINE_BLOCK_F32_16x16x4(f32_relu, ACTIVATION_RELU_F32)
// 基础分块计算宏一般版本
#define DEFINE_BLOCK_GEMM(DTYPE, SUFFIX, BLOCK_M, BLOCK_N, BLOCK_K, ACT_FUNC)                                     \
    static void block_gemm_##BLOCK_M##x##BLOCK_N##x##BLOCK_K##_##SUFFIX(const DTYPE *A, const DTYPE *B, DTYPE *C, \
                                                                        int M, int N, int K,                      \
                                                                        int m_start, int n_start, int k_start)    \
    {                                                                                                             \
        for (int i = 0; i < BLOCK_M; ++i)                                                                         \
        {                                                                                                         \
            for (int j = 0; j < BLOCK_N; ++j)                                                                     \
            {                                                                                                     \
                DTYPE sum = C[(m_start + i) * N + (n_start + j)];                                                 \
                for (int k = 0; k < BLOCK_K; ++k)                                                                 \
                {                                                                                                 \
                    sum += A[(m_start + i) * K + (k_start + k)] *                                                 \
                           B[(k_start + k) * N + (n_start + j)];                                                  \
                }                                                                                                 \
                C[(m_start + i) * N + (n_start + j)] = ACT_FUNC(sum);                                             \
            }                                                                                                     \
        }                                                                                                         \
    }

// 调度器宏
#define GEMM_IMPL(DTYPE, SUFFIX, BLOCK_M, BLOCK_N, BLOCK_K)                                               \
    void gemm_impl_##BLOCK_M##x##BLOCK_N##x##BLOCK_K##_##SUFFIX(const DTYPE *A, const DTYPE *B, DTYPE *C, \
                                                                int M, int N, int K)                      \
    {                                                                                                     \
        for (int m = 0; m < M; m += BLOCK_M)                                                              \
        {                                                                                                 \
            int m_end = m + BLOCK_M < M ? m + BLOCK_M : M;                                                \
            for (int n = 0; n < N; n += BLOCK_N)                                                          \
            {                                                                                             \
                int n_end = n + BLOCK_N < N ? n + BLOCK_N : N;                                            \
                for (int k = 0; k < K; k += BLOCK_K)                                                      \
                {                                                                                         \
                    int k_end = k + BLOCK_K < K ? k + BLOCK_K : K;                                        \
                    block_gemm_##BLOCK_M##x##BLOCK_N##x##BLOCK_K##_##SUFFIX(A, B, C, M, N, K, m, n, k);   \
                }                                                                                         \
            }                                                                                             \
        }                                                                                                 \
    }

// 组合宏：生成指定分块和激活函数的实现
#define GENERATE_IMPLEMENTATION(DTYPE, SUFFIX, BLOCK_M, BLOCK_N, BLOCK_K, ACT_FUNC) \
    DEFINE_BLOCK_GEMM(DTYPE, SUFFIX, BLOCK_M, BLOCK_N, BLOCK_K, ACT_FUNC)           \
    GEMM_IMPL(DTYPE, SUFFIX, BLOCK_M, BLOCK_N, BLOCK_K)

/******************** 具体实现实例化 ********************/
// float类型实现

// 提前特化版本
// 小矩阵: 16x16x4 + ReLU
GEMM_IMPL(float, f32_relu, 16, 16, 4)
// 小矩阵: 16x16x4 + Sigmoid
GEMM_IMPL(float, f32_sigmoid, 16, 16, 4)


// 大矩阵: 32x32x8 + ReLU
GENERATE_IMPLEMENTATION(float, f32_relu, 32, 32, 8, ACTIVATION_RELU_F32)
// 大矩阵: 32x32x8 + Sigmoid
GENERATE_IMPLEMENTATION(float, f32_sigmoid, 32, 32, 8, activation_sigmoid_f32)
// 超大矩阵: 64x64x16 + ReLU
GENERATE_IMPLEMENTATION(double, f64_relu, 64, 64, 16, activation_relu_f64)
// 超大矩阵: 64x64x16 + Sigmoid
GENERATE_IMPLEMENTATION(double, f64_sigmoid, 64, 64, 16, activation_sigmoid_f64)

    typedef enum dtype_e {
        F32,
        F64
    } dtype_t;

/******************** 入口函数选择逻辑 ********************/
void gemm_with_activation(dtype_t dtype, act_type_t act_type, const void *A,
                          const void *B,
                          void *C,
                          int M, int N, int K)
{
    if (dtype == F32)
    {
        if (M * N * K < 4096)
        {
            if (act_type == RELU)
            {
                gemm_impl_16x16x4_f32_relu((const float *)A, (const float *)B, (float *)C, M, N, K);
            }
            else
            {
                gemm_impl_16x16x4_f32_sigmoid((const float *)A, (const float *)B, (float *)C, M, N, K);
            }
        }
        else
        {
            if (act_type == RELU)
            {
                gemm_impl_32x32x8_f32_relu((const float *)A, (const float *)B, (float *)C, M, N, K);
            }
            else
            {
                gemm_impl_32x32x8_f32_sigmoid((const float *)A, (const float *)B, (float *)C, M, N, K);
            }
        }
    }
    else
    {
        if (act_type == RELU)
        {
            gemm_impl_64x64x16_f64_relu((const double *)A, (const double *)B, (double *)C, M, N, K);
        }
        else
        {
            gemm_impl_64x64x16_f64_sigmoid((const double *)A, (const double *)B, (double *)C, M, N, K);
        }
    }
}
```
---

### **使用宏的主要好处**
该代码通过宏（Macro）实现了分块矩阵乘法与激活函数组合的代码生成，这种方式在特定场景下具有显著优势，但也需权衡可读性。以下是详细分析：

---

### **1. 减少重复代码，提升开发效率**
- **场景**：需要为不同分块大小（16x16x4、32x32x8、64x64x16）、不同数据类型（`float`/`double`）、不同激活函数（ReLU/Sigmoid）生成相似的函数实现。
- **宏的作用**：
  - 通过 `DEFINE_BLOCK_GEMM` 和 `GEMM_IMPL` 宏定义分块计算和调度逻辑的模板。
  - 使用 `GENERATE_IMPLEMENTATION` 组合宏，一次性生成所有需要的函数变体（如 `gemm_impl_32x32x8_f32_relu`）。
- **优势**：
  - 避免手动编写每个分块和激活函数组合的代码，减少重复劳动。
  - 修改分块逻辑时，只需调整宏定义，所有生成的函数自动同步更新。

---

### **2. 灵活支持多参数组合**
- **场景**：需要支持多种分块大小、数据类型和激活函数的任意组合。
- **宏的作用**：
  - 宏参数化分块维度（`BLOCK_M`, `BLOCK_N`, `BLOCK_K`）、数据类型（`DTYPE`）、激活函数（`ACT_FUNC`）。
  - 例如，`GENERATE_IMPLEMENTATION(float, f32_relu, 32, 32, 8, ACTIVATION_RELU_F32)` 生成 `float` 类型、32x32x8分块、ReLU激活的具体实现。
- **优势**：
  - 新增分块策略（如 8x8x2）或激活函数（如 Tanh）时，只需调用宏生成新函数，无需修改核心逻辑。
  - 代码扩展性强，符合 **开闭原则**（对扩展开放，对修改关闭）。

---

### **3. 编译时优化与零开销抽象**
- **场景**：需要生成高性能的底层计算代码，避免运行时分支判断。
- **宏的作用**：
  - 宏在预处理阶段展开为具体的函数实现，生成**完全展开的代码**（如 `block_gemm_32x32x8_f32_relu`）。
  - 激活函数通过宏参数内联到计算逻辑中，例如 `ACT_FUNC(sum)` 直接展开为 `activation_relu_f32(sum)`。
- **优势**：
  - **消除函数调用开销**：激活函数直接内联，无需通过函数指针或条件判断。
  - **编译器优化友好**：展开后的代码允许编译器进行循环展开、向量化（SIMD）等优化。
  - **零运行时抽象成本**：所有逻辑在编译时确定，无动态派发开销。

---

### **4. 类型安全的代码生成**
- **场景**：需要为不同数据类型（`float`/`double`）生成类型特化的代码。
- **宏的作用**：
  - 通过 `DTYPE` 参数生成特定类型的代码（如 `float *A` 或 `double *A`）。
  - 结合 `activation.h` 中的类型安全函数（如 `activation_relu_f32` 和 `activation_relu_f64`）。
- **优势**：
  - 避免类型强制转换导致的潜在错误。
  - 为不同数据类型生成最优化的内存访问和计算指令。

---

### **5. 代码一致性保证**
- **场景**：确保所有分块实现遵循相同的逻辑，避免人为编写错误。
- **宏的作用**：
  - 所有分块函数的核心逻辑由宏统一生成（如三重循环结构）。
  - 例如，`DEFINE_BLOCK_GEMM` 定义了统一的分块计算模式。
- **优势**：
  - 不同分块版本的代码行为一致，降低因手动编写导致的逻辑差异风险。
  - 便于集中修复问题（如修复边界检查只需修改宏）。

---

### **权衡与局限性**
尽管宏提供了上述优势，但也需注意以下问题：
1. **可读性降低**：宏展开后的代码难以调试，函数名冗长（如 `gemm_impl_64x64x16_f64_sigmoid`）。
2. **编译错误信息模糊**：宏参数错误时，编译器报错可能指向宏定义而非调用位置。
3. **灵活性受限**：分块策略和激活函数必须在编译时确定，无法动态调整。

---
# 对比总结

性能对比示例（假设AVX512优化）：
| 实现方式       | FP32性能 (TFLOPS) | 代码大小(KB) | 编译时间(s) |
|----------------|-------------------|-------------|------------|
| 模板特化版     | 2.1               | 128         | 8.2        |
| 动态派发版     | 1.3               | 84          | 3.1        |
| 纯C宏实现      | 1.8               | 156         | 2.4        |

关键优化点说明：
1. **寄存器阻塞**：通过固定分块大小（如32x32x8），编译器可最大化寄存器利用率
2. **指令流水**：静态循环展开使CPU能更好预测分支
3. **内存预取**：固定访问模式允许编译器插入prefetch指令
4. **SIMD对齐**：模板参数确保数据对齐要求（如AVX512需要64字节对齐）

使用宏或模板而非if-else/虚函数，主要是为了**速度和硬件适配**：  
1. **性能**：宏/模板在编译时生成代码，消除虚函数查表、if-else分支预测的开销，直接优化为高效指令；  
2. **灵活性**：模板按类型生成专用代码（如float/double），宏适配不同硬件指令（如AVX/NEON），避免运行时判断拖慢速度；  
3. **零额外消耗**：直接展开代码，比函数调用或条件分支快10倍以上，适合计算密集型任务。