kernel_define.h 代码结构与功能实现分析

kernel_define.h 是 InferLLM 框架中的核心头文件，定义了内核计算所需的基本数据结构、枚举类型和宏定义。这个文件为整个计算内核系统提供了基础设施，使得不同平台的优化实现可以共享相同的接口和数据结构。

1. 文件结构概览

文件结构可以分为以下几个部分：

1. 头文件引入和宏定义
2. 命名空间和类型定义
3. 内核相关枚举类型
4. 任务调度相关结构
5. 量化计算相关结构
6. 宏定义工具

2. 头文件引入和宏定义

#pragma once

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <functional>
#include <vector>

#define PI                           (3.1415)
#define PGELU                        (0.044715)
#define INFER_ATTRIBUTE_TARGET(simd) __attribute__((target(simd)))

这部分引入了基本的 C/C++ 标准库头文件，并定义了一些常量和宏：

• PI 和 PGELU：数学常量，其中 PGELU 是 GELU 激活函数中使用的常数
• INFER_ATTRIBUTE_TARGET：用于指定函数应该使用特定的 SIMD 指令集编译

3. 命名空间和类型定义

namespace inferllm {

template <class Dtype>
using InData = std::vector<const Dtype*>;

所有代码都在 inferllm 命名空间中，InData 是一个模板类型别名，表示指向常量数据的指针数组，用于传递多个输入数据。

4. 内核相关枚举类型

enum class KernelID {
    EmbeddingGetInt4Float = 0,
    EmbeddingGetInt8Float,
    // ... 其他内核ID
    MatmulInt4WeightReorder,
};

enum class KernelOptMethod {
    MatmulInt4Reorder = 0,
};

enum class ElemMode {
    Add = 0,
    Mul,
    Silu,
    Gelu,
};

enum class RotMode {
    Mode0 = 0,
    Mode1,
    ModelRotHalf,
};

enum class KernelType { Naive = 0, Arm = 1, X86 = 2, GPU = 3 };

这部分定义了几个重要的枚举类型：

• KernelID：标识不同类型的计算操作，包括：

  • 嵌入查找操作（如 EmbeddingGetInt4Float）
  • 元素级操作（如 ElemwiseFloat）
  • 矩阵乘法操作（如 MatmulInt4Float）
  • 注意力计算相关操作（如 HeadBatchedMatmulFloat）
  • 位置编码操作（如 RopeFloat）

• KernelOptMethod：定义内核优化方法，目前只有 MatmulInt4Reorder

• ElemMode：定义元素级操作的模式，包括加法、乘法、Silu 和 Gelu 激活函数

• RotMode：定义旋转位置编码的模式，用于实现 RoPE（Rotary Position Embedding）

• KernelType：定义内核类型，包括朴素实现、ARM 优化、X86 优化和 GPU 实现

5. 任务调度相关结构

struct TaskId {
    uint32_t start;
    uint32_t end;
    uint32_t thread_id;
};

//! the task, the first parameter is the task start id, the second parameter is
//! the task end if, the third parameter is the thread id
using MultiThreadingTask = std::function<void(TaskId)>;

//! the task pair, the first parameter is the task, the second parameter is the
//! number of sub task, some kernel may need to split the task into several
using TaskSet = std::vector<std::pair<MultiThreadingTask, uint32_t>>;

这部分定义了多线程任务调度相关的结构：

• TaskId：表示任务的范围和线程 ID
• MultiThreadingTask：表示可以在多线程环境中执行的任务函数
• TaskSet：表示一组任务及其子任务数量

这些结构使得内核计算可以在多线程环境中高效执行，通过将大型计算任务分解为多个子任务，并分配给不同的线程处理。

6. 量化计算相关结构

#define QK40 32
struct BlockQ40 {
    float d;               // delta
    uint8_t qs[QK40 / 2];  // nibbles / quants
};
static_assert(sizeof(BlockQ40) == 20, "BlockQ40 size error");

struct BlockQ40X8 {
    uint8_t qs[QK40 / 2 * 8];     // nibbles / quants
    float scale[8];               // delta
};
static_assert(sizeof(BlockQ40X8) == 160, "BlockQ40X8 size error");

#define QK80 32
struct BlockQ80 {
    float d;          // delta
    int8_t qs[QK80];  // nibbles
};
static_assert(sizeof(BlockQ80) == 36, "BlockQ80 size error");

这部分定义了量化计算相关的数据结构：

• BlockQ40：4 位量化块，包含一个缩放因子 d 和 16 个字节的量化值（每个字节存储两个 4 位值）
• BlockQ40X8：8 组 4 位量化块，包含 8 个缩放因子和 128 个字节的量化值
• BlockQ80：8 位量化块，包含一个缩放因子 d 和 32 个字节的量化值

这些结构使得模型可以使用低精度整数（4 位或 8 位）进行计算，从而减少内存占用和计算量，同时通过缩放因子保持计算精度。

7. 宏定义工具

#define PartialImplementKernel(kernel_id, fun)       \
    template <typename... Args>                      \
    struct Comp<KernelID::kernel_id, Args...> {      \
        static TaskSet get_all_task(Args... args) {  \
            return fun(std::forward<Args>(args)...); \
        }                                            \
    };

#define PartialImplementSpace(kernel_id, fun)        \
    template <typename... Args>                      \
    struct Space<KernelID::kernel_id, Args...> {     \
        static size_t get(Args... args) {            \
            return fun(std::forward<Args>(args)...); \
        }                                            \
    };

这部分定义了两个宏，用于简化内核实现：

• PartialImplementKernel：为特定的内核 ID 实现 Comp 模板类，该类提供 get_all_task 静态方法，返回任务集
• PartialImplementSpace：为特定的内核 ID 实现 Space 模板类，该类提供 get 静态方法，返回工作空间大小

这些宏使得不同平台的优化实现可以方便地注册到内核系统中，而不需要修改核心代码。

8. 功能分析

通过分析 kernel_define.h 文件，可以看出它实现了以下功能：

8.1 内核标识与分类

通过 KernelID 枚举，将不同类型的计算操作进行分类和标识，使得内核系统可以根据操作类型选择合适的实现。

8.2 多线程任务调度

通过 TaskId、MultiThreadingTask 和 TaskSet 结构，提供了多线程任务调度的基础设施，使得计算任务可以在多线程环境中高效执行。

8.3 量化计算支持

通过 BlockQ40、BlockQ40X8 和 BlockQ80 结构，提供了 4 位和 8 位量化计算的支持，使得模型可以使用低精度整数进行计算，从而减少内存占用和计算量。

8.4 内核实现注册

通过 PartialImplementKernel 和 PartialImplementSpace 宏，提供了内核实现注册的机制，使得不同平台的优化实现可以方便地注册到内核系统中。

9. 设计模式分析

kernel_define.h 文件采用了多种设计模式：

9.1 策略模式

通过 KernelType 枚举，可以在运行时选择不同的计算策略（朴素实现、ARM 优化、X86 优化、GPU 实现）。

9.2 模板方法模式

通过 PartialImplementKernel 和 PartialImplementSpace 宏定义的模板类，提供了统一的接口，而具体实现由不同平台的优化代码提供。

9.3 命令模式

通过 MultiThreadingTask 和 TaskSet，将计算任务封装为命令对象，由线程池执行。

总结

kernel_define.h 是 InferLLM 框架中的核心头文件，定义了内核计算所需的基本数据结构、枚举类型和宏定义。它为整个计算内核系统提供了基础设施，使得不同平台的优化实现可以共享相同的接口和数据结构。通过多线程任务调度、量化计算支持和内核实现注册等机制，使得 InferLLM 框架能够在不同硬件平台上高效运行，同时保持代码的可维护性和扩展性。