大模型部署框架 FastLLM 实现细节解析

class Data {
    public:
        bool lockInCPU = false; // 如果lock在CPU上，那么不允许移动到其余设备
        WeightType weightType = WeightType::NONE; // 权重类型，NONE代表非权重（或未知权重）

        DataType dataType = DataType::FLOAT32; // 数据类型
        int unitSize, unitSizeDiv = 1; // 单个元素的字节数 = unitSIze / unitSizeDiv

        std::vector <int> dims; // 数据形状
        std::vector <uint64_t> strides; // 跨度

        uint64_t expansionSize = 0; // 扩容后的尺寸
        uint64_t expansionBytes = 0; // 扩容后的字节数
        std::vector <int> expansionDims; // 预扩容的形状
        uint8_t *cpuData = nullptr; // 数据指针

     void *cudaData = nullptr;
        std::vector <void*> extraCudaData;

        void *deviceData = nullptr;
        std::vector <void*> extraDeviceData;

        DataDevice dataDevice = DataDevice::CPU;

        // 这两个参数用于量化，对FLOAT数据不适用
        int perChannelAxis = -1; // 沿哪个轴分通道量化，-1代表没有分通道
        std::vector <LowBitConfig> perChannelsConfigs; // perChannelsConfigs[i]代表第i个通道的min, max; 如果没有分通道，perChannelsConfigs[0]代表全局min, max
        std::vector <float> scales, mins;
        std::vector <int> zeros;
        std::vector <int> weightSum; // 作为权重时，有时候需要存一些和加速计算

        std::string fileName;
        long long filePos;
        std::shared_ptr<FileMmap> m_file;

        Data () {};

        Data (DataType type);

        Data (DataType type, const std::vector <int> &dims); // 构造函数

        // 构造函数，创建好之后从data复制数据
        // data中是原始数据，如果type不是float那么需要量化
        Data (DataType type, const std::vector <int> &dims, const std::vector <float> &data);

        ~Data(); // 析构函数

        Data (const Data &ori); // 深拷贝

        void CopyFrom(const Data &ori); // 复制

        uint64_t GetBytes() const; // 获取总字节数

        void Allocate(); // 分配内存

        void Allocate(float v); // 分配内存并初始化

        void Expansion(const std::vector <int> &dims); // 预扩容到相应尺寸

        void MallocSpace(uint64_t size); // 在设备上分配

        void FreeSpace(); // 回收设备上的内存

        void UpdateUnitSize(); // 更新unitSize

        void Resize(const std::vector <int> &dims); // 更改尺寸

        void Reshape(const std::vector <int> &dims); // 更改尺寸,但不修改数据

        uint64_t Count(int i) const; // dims[i] * strides[i]

        void PrintShape() const; // 输出形状

        void Print() const; // 输出

        void CalcWeightSum(); // 计算WeightSum

        void ToDevice(DataDevice device); // 移动到指定device

        void ToDevice(void *device);

        void set_file(std::shared_ptr<FileMmap> file) {
            m_file = file;
        }
    };

void Permute(const Data &input, const std::vector<int> &axis, Data &output) {
        Data axisData = Data(DataType::INT32PARAM, {(int)axis.size()});
        axisData.Allocate();
        for (int i = 0; i < axisData.Count(0); i++) {
            ((int32_t*)axisData.cpuData)[i] = axis[i];
        }
        curExecutor->Run("Permute", {
                {"input", (Data*)&input}, {"axis", &axisData}, {"output", (Data*)&output}
        }, {}, {});
    }

namespace fastllm {
    class Executor {
    private:
        std::vector <BaseDevice*> devices;
        std::map <std::string, float> profiler;

    public:
        Executor (); // 创建默认的Executor

        ~Executor(); // 析构

        void ClearDevices(); // 清空 devices

        void AddDevice(BaseDevice *device); // 增加一个device

        // 运行一个op
        void Run(const std::string &opType, const fastllm::DataDict &datas, const fastllm::FloatDict &floatParams,
                 const fastllm::IntDict &intParams);

        void ClearProfiler();

        void PrintProfiler();
    };
}

namespace fastllm {
    Executor::Executor() {
        this->devices.clear();
#ifdef USE_CUDA
        // 将一个指向 CudaDevice 类对象的指针插入到 devices 向量的末尾。
        // 这里通过 new 运算符创建了一个 CudaDevice 对象，并将返回的指针进行类型转换为 BaseDevice* 类型。
        this->devices.push_back((BaseDevice*) new CudaDevice());
#endif
        this->devices.push_back((BaseDevice*) new CpuDevice());
    }

    Executor::~Executor() {
        // 释放 devices 向量中的每个指针元素所占用的内存。
        for (int i = 0; i < devices.size(); i++) {
            delete devices[i];
        }
    }

    void Executor::ClearDevices() {
        // this->devices 指的是当前对象的 devices 成员，即指向 BaseDevice 类对象的指针向量。
        this->devices.clear();
    }
    
    // 该函数用于向 devices 向量中添加一个指向 BaseDevice 类对象的指针。
    void Executor::AddDevice(fastllm::BaseDevice *device) {
        this->devices.push_back(device);
    }

    void Executor::Run(const std::string &opType, const fastllm::DataDict &datas, const fastllm::FloatDict &floatParams,
                       const fastllm::IntDict &intParams) {
        // 创建一个 st 变量，用于记录函数开始执行的时间。
        auto st = std::chrono::system_clock::now();
        // 创建一个布尔变量 lockInCPU，用于记录是否将数据锁定在 CPU 上。
        bool lockInCPU = false;
        // 在第一个 for 循环中，遍历数据字典 datas，查找是否有 "___batch" 后缀的参数，
        // 并根据情况设置 lockInCPU 的值。it.first 是数据字典中的键（key），it.second 
        // 是对应的值（value）。如果存在 "___batch" 后缀的参数，则将 lockInCPU 设置为
        // 对应数据的 lockInCPU 属性（布尔值），否则设置为当前数据的 lockInCPU 属性。
        for (auto &it: datas) {
            if (intParams.find(it.first + "___batch") != intParams.end()) {
                int batch = intParams.find(it.first + "___batch")->second;
                for (int i = 0; i < batch; i++) {
                    lockInCPU |= ((Data**)it.second)[i]->lockInCPU;
                }
            } else {
                lockInCPU |= it.second->lockInCPU;
            }
        }
        // 第二个 for 循环遍历 devices 向量中的所有设备指针 device。
        // 在循环中，首先检查 lockInCPU 是否为真，并且当前设备的类型不是 "cpu"，
        // 如果是，则跳过当前设备（continue）。这个检查是为了保证数据锁定在 CPU 上时，只执行 CPU 设备上的操作。
        for (auto device: devices) {
            if (lockInCPU && device->deviceType != "cpu") {
                continue;
            }
            // 然后，通过调用 device->CanRun(opType, datas, floatParams, intParams) 
            // 检查当前设备是否可以运行指定的操作 opType。如果可以运行，则进行以下操作：
            if (device->CanRun(opType, datas, floatParams, intParams)) {
                // 第三个 for 循环遍历数据字典 datas，如果存在 "___batch" 后缀的参数，
                // 则将对应数据转移到当前设备上；否则，将当前数据转移到当前设备上。
                for (auto &it: datas) {
                    if (intParams.find(it.first + "___batch") != intParams.end()) {
                        int batch = intParams.find(it.first + "___batch")->second;
                        for (int i = 0; i < batch; i++) {
                            ((Data**)it.second)[i]->ToDevice((void *) device);
                        }
                    } else {
                        it.second->ToDevice((void *) device);
                    }
                }
                // 调用 device->Reshape(opType, datas, floatParams, intParams) 
                // 进行形状推导，device上的形状推导调用了opType对应的op的形状推导，
                // 并且被各个不同的op重写。
                device->Reshape(opType, datas, floatParams, intParams);
                // 对opType对应的这个算子进行推理。
                device->Run(opType, datas, floatParams, intParams);
                break;
            }
        }
        // 最后，计算操作运行时间，并将其加入 profiler 成员变量，用于性能分析。
        float spend = GetSpan(st, std::chrono::system_clock::now());
        profiler[opType] += spend;
    }
    
    // 清除profile的信息
    void Executor::ClearProfiler() {
        profiler.clear();
    }
    
    // 打印profile信息，也即输出每个层的运行时间和模型的总运行时间
    void Executor::PrintProfiler() {
        float sum = 0.0;
        for (auto &it : profiler) {
            printf("%s spend %f\n", it.first.c_str(), it.second);
            sum += it.second;
        }
        printf("total spend %f\n", sum);
    }
}

struct Tokenizer {
        struct TrieNode {
            int tokenId;
            std::map <int, TrieNode*> next;
            TrieNode();
        };
        TrieNode *root;

        std::unordered_map <int, std::string> tokenToStringDict;

        Tokenizer ();

        ~Tokenizer();

        void Clear(); // 清空分词器

        void Insert(const std::string &s, int tokenId); // 插入一个token

        Data Encode(const std::string &s); // 编码

        std::string Decode(const Data &data); // 解码

        std::string DecodeTokens(const std::vector <int> &tokens); // 解码
    };

   // 这是 Tokenizer 类的嵌套结构 TrieNode 的构造函数的实现。
   // 在构造函数中，将 tokenId 成员变量的值初始化为 -999999。
   // 这个值在构造函数中被硬编码，它是作为一个特殊标记来使用的。
  Tokenizer::TrieNode::TrieNode() {
        this->tokenId = -999999;
    }
    
    // Tokenizer 类的构造函数的实现。
    // 在构造函数中，通过 new 运算符创建一个新的 TrieNode 对象，
    // 并将其指针赋值给 root 成员变量。这样，构造函数创建了一个空的字典树，
    // 并将其根节点指针存储在 root 中。
    Tokenizer::Tokenizer() {
        root = new TrieNode();
    }
    
    // Tokenizer 类的析构函数的实现。
    // 在析构函数中，首先调用 Clear() 函数，用于释放动态分配的资源和清空数据。
    // 然后，调用 delete 运算符释放通过 new 运算符创建的 root 对象的内存，从而释放整个字典树的内存。
    Tokenizer::~Tokenizer() {
        Clear();
        delete root;
    }
    
    // 这是 Tokenizer 类的成员函数 Clear() 的定义，用于清空分词器并释放动态分配的资源。
    void Tokenizer::Clear() {
        // 创建一个指向 TrieNode 的指针向量 q，用于辅助遍历字典树。
        std::vector <TrieNode*> q;
        // 将字典树的根节点 root 加入 q 向量，作为遍历的起始点。
        q.push_back(root);
        // 开始遍历 q 向量中的节点，这是一个广度优先搜索（BFS）的过程。
        for (int i = 0; i < q.size(); i++) {
            // 取出当前遍历到的节点 now。
            TrieNode *now = q[i];
            // 对当前节点 now 的所有子节点进行遍历。
            for (auto it : now->next) {
                // 将当前节点 now 的子节点加入 q 向量中，以便继续遍历子节点的子节点。
                q.push_back(it.second);
            }
        }
        // 当遍历完成后，q 向量中包含了字典树中的所有节点。
        // 创建一个新的 TrieNode 对象，并将其指针赋值给 root 成员变量，表示创建了一个空的字典树。
        root = new TrieNode();
        //  清空 tokenToStringDict 映射表，以确保所有 token 的映射被清空。
        tokenToStringDict.clear();
    }
    
    // 这是 Tokenizer 类的成员函数 Insert 的定义，用于向分词器中插入一个 token。
    void Tokenizer::Insert(const std::string &s, int tokenId) {
        // 创建一个指向 TrieNode 的指针 now，并将其初始化为指向字典树的根节点 root。
        TrieNode *now = this->root;
        // 开始遍历输入的字符串 s 中的每个字符。
        for (int i = 0; i < s.size(); i++) {
            // 检查当前字符 s[i] 是否已经存在于当前节点 now 的 next 映射表中。
            // 如果当前字符 s[i] 不存在于当前节点 now 的子节点中，
            // 在 now->next 中添加新的子节点，该子节点的键为当前字符 s[i] 的编码值，
            // 值为指向新创建的 TrieNode 对象的指针。这表示在字典树中添加了一个新的字符节点。
            if (now->next.find(s[i]) == now->next.end()) {
                now->next[s[i]] = new TrieNode();
            }
            // 将 now 移动到下一个字符 s[i] 对应的节点，以便继续处理下一个字符。
            now = now->next[s[i]];
        }
        // 遍历完成后，now 将指向字典树中最后一个字符的节点。
        // 设置当前节点的 tokenId 成员变量，表示当前节点代表一个 token，
        // 并使用传入的 tokenId 值来标识该 token。
        now->tokenId = tokenId;
        // 将传入的 tokenId 和对应的字符串 s 添加到 tokenToStringDict 
        // 映射表中，用于后续的解码过程。
        tokenToStringDict[tokenId] = s;
    }
    
    // 这是 Tokenizer 类的成员函数 Encode 的定义，用于对输入的字符串 s 进行编码。
    Data Tokenizer::Encode(const std::string &s) {
        // 创建一个浮点数向量 v，用于存储编码结果。该向量将存储找到的 token 对应的 tokenId 值。
        std::vector <float> v;
        // 开始遍历输入的字符串 s 中的每个字符。
        for (int i = 0; i < s.size(); i++) {
            // 创建两个整数变量 tokenId 和 pos，
            // 用于记录找到的 token 的 tokenId 值和 token 的结束位置。
            int tokenId = -999999, pos = i - 1;
            // 创建一个指向 TrieNode 的指针 now，并将其初始化为指向字典树的根节点 root。
            TrieNode *now = this->root;
            // 从当前字符 s[i] 开始继续遍历字符串 s。
            for (int j = i; j < s.size(); j++) {
              // 检查当前字符 s[j] 是否存在于当前节点 now 的 next 映射表中。
              // 如果存在，表示当前字符构成了一个 token 的一部分，继续遍历子节点。
                if (now->next.find(s[j]) != now->next.end()) {
                    // 将 now 移动到下一个字符 s[j] 对应的节点。
                    now = now->next[s[j]];
                    // 检查当前节点 now 是否代表一个 token，即它的 tokenId 是否有效。
                    if (now->tokenId != -999999) {
                        // 如果当前节点代表一个 token，将 tokenId 和当前位置 j 存储到 
                        // tokenId 和 pos 变量中，以便记录找到的 token 的信息。 
                        tokenId = now->tokenId;
                        pos = j;
                    }
                } else { // 如果当前字符不再是 token 的一部分，退出内层循环，继续外层循环。
                    break;
                }
            }
            // 如果 pos 大于等于当前位置 i，表示找到了一个 token。
            // 这里 pos 存储了找到的 token 的结束位置，i 移动到 pos 处，以便继续遍历下一个字符。
            if (pos >= i) {
                i = pos;
                v.push_back(tokenId);
                //printf("%d ", tokenId);
            }
        }
        //printf("\n");
        // 遍历完成后，v 向量中存储了输入字符串中所有找到的 token 对应的 tokenId 值。
        // 创建一个 Data 对象并返回，表示编码的结果。这里 Data 是一个数据结构，
        // 用于存储数据及其相关信息。编码结果是一个一维浮点数数组，
        // 表示输入字符串中所有找到的 token 对应的 tokenId 值。
        return Data (DataType::FLOAT32, {1, (int)v.size()}, v);
    }
    
    // 这是 Tokenizer 类的成员函数 DecodeTokens 的定义，
    // 用于对输入的 token 数组进行解码，将 token 转换回原始的字符串。
    std::string Tokenizer::DecodeTokens(const std::vector<int> &tokens) {
        // 创建一个空字符串 ret，用于存储解码结果。
        std::string ret = "";
        // 开始遍历输入的 token 数组 tokens。
        for (int i = 0; i < tokens.size(); i++) {
            // 获取当前 token 对应的原始字符串 s，通过查询 tokenToStringDict 映射表，
            // 将 tokens[i] 转换回字符串。
            std::string s = tokenToStringDict[tokens[i]];
            // 判断当前 token 是否需要特殊处理：
            // 如果 s 是类似 "<0xHH>" 格式的 token（其中 HH 表示十六进制数），
            // 则需要将其转换为对应的字符。首先，提取 HH，然后将其转换为对应的字符，
            // 并用空格代替原始的 token。
            if (s.size() == 6 && s.substr(0, 3) == "<0x" && s.back() == '>') {
                int c = 0;
                for (int i = 3; i < 5; i++) {
                    c *= 16;
                    if (s[i] >= '0' && s[i] <= '9') {
                        c += (s[i] - '0');
                    } else {
                        c += (s[i] - 'A' + 10);
                    }
                }

                s = " ";
                s[0] = c;
            }
            // 根据不同的 token 进行解码：
            if (s == "<n>") {
                ret += "\n";
            } else if (s == "<|tab|>") {
                ret += "\t";
            } else {
                ret += s;
            }
        }
        
        // 将特殊字符 "\xE2\x96\x81"（UTF-8 编码）替换为空格 " "，这是用于表示空格的特殊字符。
        std::string blank = "";
        blank += 226, blank += 150, blank += 129;
        while (true) {
            std::string::size_type pos(0);
            if ((pos = ret.find(blank)) != std::string::npos)
                ret.replace(pos, blank.length(), " ");
            else break;
        }
        // 检查是否有 "<|blank_数字>" 格式的特殊 token，如果有，将其解码成对应数量的空格字符。
        int pos = ret.find("<|blank_");
        if (pos != -1) {
            int space_num = atoi(ret.substr(8, ret.size() - 10).c_str());
            return std::string(space_num, ' ');
        }

        return ret;
    }

    std::string Tokenizer::Decode(const Data &data) {
        std::vector <int> tokens;
        for (int i = 0; i < data.Count(0); i++) {
            tokens.push_back((int) ((float *) data.cpuData)[i]);
        }
        return DecodeTokens(tokens);
    }

if (pos != -1) {
            int space_num = atoi(ret.substr(8, ret.size() - 10).c_str());
            return std::string(space_num, ' ');
        }

   // 这是 CpuDevice 类的成员函数 Malloc 的定义，用于在 CPU 上分配一块内存空间。
    bool CpuDevice::Malloc(void **ret, size_t size) {
        *ret = (void*)new uint8_t [size];
        return true;
    }
    
    // 这是 CpuDevice 类的成员函数 Free 的定义，用于在 CPU 上释放之前分配的内存。
    bool CpuDevice::Free(void *ret) {
        delete[] (uint8_t*)ret;
        return true;
    }
    
    // 这是 CpuDevice 类的成员函数 CopyDataFromCPU 的定义，用于将数据从 CPU 拷贝到指定的设备上。
    // 这里什么都不做，直接返回true。
    bool CpuDevice::CopyDataFromCPU(void *dst, void *src, size_t size) {
        return true;
    }
    
    // 这是 CpuDevice 类的成员函数 CopyDataToCPU 的定义，用于将数据从指定的设备拷贝到 CPU 上。
    bool CpuDevice::CopyDataToCPU(void *dst, void *src, size_t size) {
        return true;
    }

// 如果定义了 __AVX__ 和 __AVX2__，那么会启用第一个 DotU8U8 函数和 DotU4U8 函数。
// 如果只定义了 __AVX__，但没有定义 __AVX2__，那么会启用第二个 DotU8U8 函数和 DotU4U8 函数。

#ifdef __AVX__
#ifdef __AVX2__
    // 这是一段使用了 Intel AVX2 指令集（Advanced Vector Extensions 2）的代码，
    // 用于计算两个8位无符号整数数组的点积。
    // 定义了一个函数 DotU8U8，它接受两个指向 8 位无符号整数的指针 a 和 b，
    // 以及一个整数 n。这个函数的目的是计算数组 a 和 b 的点积，其中数组的长度为 n。
    int DotU8U8(uint8_t *a, uint8_t *b, int n) {
        // 初始化一个 256 位的整数向量 acc，所有位都设置为零。这个向量用于存储点积的累加值。
        __m256i acc = _mm256_setzero_si256();
        //  初始化两个变量，i 用于循环计数，ans 用于存储最后的结果。
        int i = 0;
        int ans = 0;
        // 等这几行代码初始化了一些常量向量
        const __m256i lowMask = _mm256_set1_epi8(0xf);
        const __m256i ones = _mm256_set1_epi16(1);
        const __m256i ones8 = _mm256_set1_epi8(1);
        const __m256i xors = _mm256_set1_epi8(-128);
        // 这是一个循环，每次处理 32 个元素。这是因为 AVX2 可以同时处理 32 个 8 位整数。
        for (; i + 31 < n; i += 32) {
            // 这两行代码从数组 a 和 b 中加载数据到 256 位的向量 bx 和 by。
            __m256i bx = _mm256_loadu_si256((const __m256i *) (a + i));
            __m256i by = _mm256_loadu_si256((const __m256i *) (b + i));
            
            // 这行代码将 by 中的每个元素减去 128，这对应于上面表达式中的 ((int)b[i] - 128)。
            by = _mm256_xor_si256(by, xors);
            // 这行代码对于那些原本是 0 的元素（在减去 128 后变为 -128 的元素）加 1，
            // 以避免后续乘法操作时的溢出。
            by = _mm256_add_epi8(by, _mm256_and_si256(_mm256_cmpeq_epi8(by, xors), ones8));
            
            //  这行代码将 bx 中的符号应用到 by 中，对应于上面表达式中的 ((int8_t*)a)[i]。
            by = _mm256_sign_epi8(by, bx);
            // 这行代码将 bx 中的所有非零元素变为 1，这是为了在后续的乘法操作中保持 by 中元素的原值。
            bx = _mm256_sign_epi8(bx, bx);
            
            // 这行代码先对 bx 和 by 进行乘法运算（这对应于上面表达式中的 * 操作），
            // 然后再与 acc 进行加法操作（这对应于上面表达式中的 += 操作）。
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(bx, by), ones));
        }
        // 这是另一个循环，用于处理数组中剩余的元素（数量小于 32）。
        // 这些元素通过常规的方式计算点积，然后累加到 ans 中。
        for (; i < n; i++) {
            ans += ((int8_t*)a)[i] * ((int)b[i] - 128);
        }
        
        // 最后，将 acc 中的所有元素相加，然后再加上 ans，返回最终的结果。
        return ans + I32sum(acc);
    };
#else
    // 定义了一个函数 DotU8U8，它接受两个指向 8 位无符号整数的指针 a 和 b，
    // 以及一个整数 n。这个函数的目的是计算数组 a 和 b 的点积，其中数组的长度为 n。
    int DotU8U8(uint8_t *a, uint8_t *b, int n) {
        // 初始化一个 256 位的整数向量 acc，所有位都设置为零。这个向量用于存储点积的累加值。
        __m256i acc = _mm256_setzero_si256();

        int i = 0;
        int ans = 0;
        // 这是一个循环，每次处理 32 个元素。这是因为 AVX 可以同时处理 32 个 8 位整数。
        for (; i + 31 < n; i += 32) {
            // 这两行代码从数组 a 和 b 中加载数据到 256 位的向量 bx 和 by。
            __m256i bx = _mm256_loadu_si256((const __m256i *) (a + i));
            __m256i by = _mm256_loadu_si256((const __m256i *) (b + i));
            
            // 接下来的四行代码将 bx 和 by 中的 8 位整数扩展为 16 位整数。
            // 这是因为在后续的乘法和累加操作中，如果仍然使用 8 位整数，可能会发生溢出。
            __m256i mx0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(bx, 0));
            __m256i mx1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(bx, 1));

            __m256i my0 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(by, 0));
            __m256i my1 = _mm256_cvtepu8_epi16(_mm256_extractf128_si256(by, 1));
            
            // 这两行代码首先对 mx0 和 my0，以及 mx1 和 my1 进行乘法累加操作，
            // 然后再与 acc 进行加法操作，结果存储在 acc 中。
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(mx0, my0));
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(mx1, my1));
        }
        //  这是另一个循环，用于处理数组中剩余的元素（数量小于 32）。
        // 这些元素通过常规的方式计算点积，然后累加到 ans 中。
        for (; i < n; i++) {
            ans += a[i] * b[i];
        }
        
        // 最后，将 acc 中的所有元素相加，然后再加上 ans，返回最终的结果。
        return ans + I32sum(acc);
    };
#endif
    // 它接受两个指向 8 位无符号整数的指针 a 和 b，以及一个整数 n。
    // 这个函数的目的是计算数组 a 和 b 的点积，其中数组的长度为 n。
    int DotU4U8(uint8_t *a, uint8_t *b, int n) {
        // 初始化一个 256 位的整数向量 acc，所有位都设置为零。这个向量用于存储点积的累加值。
        __m256i acc = _mm256_setzero_si256();

        int i = 0;
        int ans = 0;
        // 初始化两个常量向量，lowMask 中的每个元素都是 0xf，ones 中的每个元素都是 1。
        const __m256i lowMask = _mm256_set1_epi8(0xf);
        const __m256i ones = _mm256_set1_epi16(1);
        for (; i + 31 < n; i += 32) {
            // 从数组 a 中加载 16 个元素到 128 位的向量 orix 中。
            // 这里 i / 2 的原因是每个元素实际上只有 4 位。
            __m128i orix = _mm_loadu_si128((const __m128i *) (a + i / 2));
            // 将 orix 中的元素分成高 4 位和低 4 位，然后将它们合并成一个 256 位的向量 bytex。
            __m256i bytex = _mm256_set_m128i(_mm_srli_epi16(orix, 4), orix);
            // 使用按位与操作，取 bytex 中的每个元素的低 4 位，结果存储在 bx 中。
            __m256i bx = _mm256_and_si256(lowMask, bytex);
            // 从数组 b 中加载数据到 256 位的向量 by。
            __m256i by = _mm256_loadu_si256((const __m256i *) (b + i));
            // 这行代码首先进行了两个向量的乘法累加操作，然后再与 acc 进行加法操作，结果存储在 acc 中。
            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(by, bx), ones));
        }
        for (; i < n; i++) {
            ans += a[i] * b[i];
        }

        return ans + I32sum(acc);
    };
#endif

// CpuEmbedding 算子的形状推导函数，这个函数接受四个参数：
// 一个 std::string 类型的 opType，两个字典类型的 datas 和 floatParams，以及一个 intParams。
void CpuEmbedding::Reshape(const std::string &opType, const fastllm::DataDict &datas,
                               const fastllm::FloatDict &floatParams, const fastllm::IntDict &intParams) {
        // 这三行代码从 datas 字典中查找键为 "input"、"output" 和 "weight" 的元素，
        // 并将找到的元素的值赋给 input、output 和 weight。
        // 这里的 "input"、"output" 和 "weight" 可以理解为嵌入层的输入、输出和权重。
        Data &input = *(datas.find("input")->second);
        Data &output = *(datas.find("output")->second);
        Data &weight = *(datas.find("weight")->second);
        
        // 这行代码检查 weight 的维度数量是否为 2。如果不是，就会抛出一个错误。
        AssertInFastLLM(weight.dims.size() == 2, "Embedding's weight's dim should be 2.\n");
        // 这行代码检查 weight 的数据类型是否为 FLOAT32 或 BFLOAT16。如果不是，就会抛出一个错误。
        AssertInFastLLM(weight.dataType == DataType::FLOAT32 ||
                        weight.dataType == DataType::BFLOAT16, "Embedding's weight's type should be float32 or bfloat16.\n");
        // 这行代码检查 input 的数据类型是否为 FLOAT32。如果不是，就会抛出一个错误。
        AssertInFastLLM(input.dataType == DataType::FLOAT32, "Embedding's input's type should be float32.\n");
        
        // 这行代码将 weight 的 weightType 属性设置为 EMBEDDING。
        weight.weightType = WeightType::EMBEDDING;
        // 这行代码从 weight 的维度中提取词汇大小（vocabSize）和嵌入大小（embSize）。
        int vocabSize = weight.dims[0], embSize = weight.dims[1];
        // 这两行代码将 embSize 添加到 input 的维度中，形成一个新的维度。
        std::vector <int> dims = input.dims;
        dims.push_back(embSize);
        
        // 这两行代码将 output 的数据类型设置为 FLOAT32，并重新调整其维度。
        output.dataType = DataType::FLOAT32;
        output.Resize(dims);
    }
    
    // 这是一个名为 CpuEmbedding::Run 的函数，它在某个名为 CpuEmbedding 的类中被定义。
    // 这个函数接受四个参数：一个 std::string 类型的 opType，
    // 两个字典类型的 datas 和 floatParams，以及一个 intParams。
    // 这个函数的主要任务是执行嵌入层（Embedding layer）的运算。
    // 嵌入层通常用于将离散型特征（例如词汇）转换为连续的向量表示。
    // 具体的实现方法是，对于每个输入的索引，从权重矩阵中查找对应的行，
    // 然后将其复制到输出矩阵的对应位置。
    void CpuEmbedding::Run(const std::string &opType, const fastllm::DataDict &datas,
                               const fastllm::FloatDict &floatParams, const fastllm::IntDict &intParams) {
        // 这三行代码从 datas 字典中查找键为 "input"、"output" 和 "weight" 的元素，
        // 并将找到的元素的值赋给 input、output 和 weight。
        // 这里的 "input"、"output" 和 "weight" 可以理解为嵌入层的输入、输出和权重。
        Data &input = *(datas.find("input")->second);
        Data &output = *(datas.find("output")->second);
        Data &weight = *(datas.find("weight")->second);;

        output.Allocate(); // 这行代码为 output 分配内存。
        
        // 这行代码从 weight 的维度中提取词汇大小（vocabSize）和嵌入大小（embSize）。
        int vocabSize = weight.dims[0], embSize = weight.dims[1];
        // 这行代码计算 input 的长度。
        uint64_t inputLen = input.Count(0);
        // 这行代码获取 input 的数据，并将其转换为浮点数的指针。
        float *inputData = (float*)input.cpuData;
        
        // 接下来的代码根据内存模式和权重的数据类型的不同，分别处理了四种情况。
        // 这四种情况可以归纳为两个大类：内存模式和权重的数据类型。
        // 内存模式：如果 GetLowMemMode() 返回 true，则表示处于低内存模式。
        // 在这种模式下，权重数据不会一次性全部加载到内存中，而是每次只加载需要的部分。
        // 否则，权重数据会全部加载到内存中。
        if (GetLowMemMode()) {
            FILE *fi = fopen(weight.fileName.c_str(), "rb");
            // 权重的数据类型：如果权重的数据类型为 FLOAT32，则使用浮点数进行计算。
            // 如果权重的数据类型为 BFLOAT16，则使用 16 位浮点数进行计算。
            if (weight.dataType == DataType::FLOAT32) {
                float *outputData = (float *) output.cpuData;
                for (int i = 0; i < inputLen; i++) {
                    // 这行代码从 inputData 中取出第 i 个元素，并将其四舍五入到最近的整数。
                    int token = (int) (inputData[i] + 1e-9);
                    // 这两行代码将文件指针移动到第 token 行的开始位置。
#if defined(_WIN32) or defined(_WIN64)
                    _fseeki64(fi, (long long)token * embSize * sizeof(float) + weight.filePos, 0);
#else
                    fseek(fi, (long long)token * embSize * sizeof(float) + weight.filePos, 0);
#endif
                    // 这行代码从文件中读取 embSize 个浮点数，并将它们存储在 outputData 的对应位置。
                    int ret = fread(outputData + i * embSize, sizeof(float), embSize, fi);
                }
            } else {
                // 如果权重的数据类型为 BFLOAT16，则使用 16 位浮点数进行计算。
                // 这部分代码的逻辑与 FLOAT32 部分的逻辑类似，只是多了一个步骤：
                // 将 16 位的浮点数转换为 32 位的浮点数。
                uint16_t *outputData = (uint16_t *) output.cpuData;
                uint16_t *weightData = new uint16_t[embSize];
                for (int i = 0; i < inputLen; i++) {
                    int token = (int) (inputData[i] + 1e-9);
#if defined(_WIN32) or defined(_WIN64)
                    _fseeki64(fi, (long long)token * embSize * sizeof(uint16_t) + weight.filePos, 0);
#else
                    fseek(fi, (long long)token * embSize * sizeof(uint16_t) + weight.filePos, 0);
#endif
                    int ret = fread(weightData, sizeof(uint16_t), embSize, fi);
                    for (int j = 0; j < embSize; j++) {
                        outputData[i * embSize * 2 + j * 2] = 0;
                        outputData[i * embSize * 2 + j * 2 + 1] = weightData[j];
                    }
                }
                delete[] weightData;
            }
            // 最后，fclose(fi); 这行代码关闭了文件。
            fclose(fi);
        } else {
            if (weight.dataType == DataType::FLOAT32) {
                // 这两行代码获取 output 和 weight 的数据，并将它们转换为浮点数的指针。
                float *outputData = (float *) output.cpuData;
                float *weightData = (float *) weight.cpuData;
                for (int i = 0; i < inputLen; i++) {
                    int token = (int) (inputData[i] + 1e-9);
                    // 这行代码从 weightData 中复制 embSize 个浮点数到 outputData 的对应位置。
                    // 这里的 token 是索引，embSize 是嵌入向量的长度。
                    memcpy(outputData + i * embSize, weightData + token * embSize, embSize * sizeof(float));
                }
            } else {
                uint16_t *outputData = (uint16_t *) output.cpuData;
                uint16_t *weightData = (uint16_t *) weight.cpuData;
                for (int i = 0; i < inputLen; i++) {
                    int token = (int) (inputData[i] + 1e-9);
                    for (int j = 0; j < embSize; j++) {
                        outputData[i * embSize * 2 + j * 2] = 0;
                        outputData[i * embSize * 2 + j * 2 + 1] = weightData[token * embSize + j];
                    }
                }
            }
        }
    }

void CpuLayerNormOp::Run(const std::string &opType, const fastllm::DataDict &datas,
                             const fastllm::FloatDict &floatParams, const fastllm::IntDict &intParams) {
        // 这四行代码从 datas 字典中查找键为 "input"、"output"、"gamma" 和 "beta" 的元素，
        // 并将找到的元素的值赋给 input、output、gamma 和 beta。
        // 这里的 "input" 是层归一化的输入，"output" 是输出，
        // "gamma" 和 "beta" 是用于对归一化后的结果进行缩放和移位的可学习参数。
        Data &input = *(datas.find("input")->second);
        Data &output = *(datas.find("output")->second);
        Data &gamma = *(datas.find("gamma")->second);
        Data &beta = *(datas.find("beta")->second);
        
        // 这行代码为 output 分配内存。
        output.Allocate();
        
        // 这行代码从 intParams 字典中查找键为 "axis" 的元素。
        // 如果找到，则使用找到的值作为归一化的轴；否则，使用默认值 -1。在层归一化中，轴通常是特征维度。
        int axis = intParams.find("axis") != intParams.end() ? intParams.find("axis")->second : -1;
        // 这两行代码计算 input 的维度数，并将 axis 转换为非负数。
        // 这是为了处理负数的轴值，因为在 Python 中，轴可以是负数，表示从后向前数的位置。
        int dimsLen = input.dims.size();
        axis = (axis % dimsLen + dimsLen) % dimsLen;
        
        // 这三行代码计算 outer、channels 和 inner。
        // outer 是归一化操作的外部维度的元素总数，channels 是归一化操作的轴的大小，
        // inner 是归一化操作的内部维度的元素总数。
        int outer = input.Count(0) / input.Count(axis);
        int channels = input.dims[axis];
        int inner = input.strides[axis];
        
        // 这行代码为 mean 和 var 分配内存，它们用于存储每个归一化组的均值和方差。
        float *mean = new float[inner], *var = new float[inner];
        float *inputData = (float *) input.cpuData;
        float *outputData = (float *) output.cpuData;
        float *gammaData = (float *) gamma.cpuData;
        float *betaData = (float *) beta.cpuData;
        
        // 在这个条件下，每个通道只有一个元素，所以可以并行地对每个通道进行层归一化。
        if (inner == 1) {
            // 这是一个循环，对 input 中的每一个外部元素进行处理。
            for (int i = 0; i < outer; i++) {
                // 这行代码定义了三个浮点数变量，分别用于存储均值、平方和和方差。
                float mean = 0.f, s2 = 0.f, var = 0.f;
                int j = 0;
                // 这是一段条件编译的代码，只有在目标平台为 ARM 架构时才会编译和执行。
                // 这段代码使用了 ARM 架构的 SIMD 指令来加速计算。
#ifdef __aarch64__
                float32x4_t sums = vdupq_n_f32(0.0);
                    float32x4_t sums2 = vdupq_n_f32(0.0);
                    for (; j + 3 < channels; j += 4) {
                        float32x4_t vi = vld1q_f32(inputData + j);
                        sums = vaddq_f32(sums, vi);
                        sums2 = vaddq_f32(sums2, vmulq_f32(vi, vi));
                    }
                    mean = sums[0] + sums[1] + sums[2] + sums[3];
                    s2 = sums2[0] + sums2[1] + sums2[2] + sums2[3];
#endif
#ifdef __AVX2__
                // 这是另一段条件编译的代码，只有在目标平台支持 AVX2 指令集时才会编译和执行。
                // 这段代码使用了 AVX2 的 SIMD 指令来加速计算。
                __m256 sum_vec = _mm256_setzero_ps();
                __m256 squared_sum_vec = _mm256_setzero_ps();

                for (; j < channels - 7; j += 8) {
                    __m256 data_vec = _mm256_loadu_ps(inputData + j);
                    sum_vec = _mm256_add_ps(sum_vec, data_vec);

                    __m256 squared_data_vec = _mm256_mul_ps(data_vec, data_vec);
                    squared_sum_vec = _mm256_add_ps(squared_sum_vec, squared_data_vec);
                }

                float sum_array[8];
                _mm256_storeu_ps(sum_array, sum_vec);
                mean = sum_array[0] + sum_array[1] + sum_array[2] + sum_array[3] +
                            sum_array[4] + sum_array[5] + sum_array[6] + sum_array[7];

                float squared_sum_array[8];
                _mm256_storeu_ps(squared_sum_array, squared_sum_vec);
                s2 = squared_sum_array[0] + squared_sum_array[1] +
                                    squared_sum_array[2] + squared_sum_array[3] +
                                    squared_sum_array[4] + squared_sum_array[5] +
                                    squared_sum_array[6] + squared_sum_array[7];
#endif
                // 这是一个循环，对 input 中剩余的每一个通道进行处理。
                for (; j < channels; j++) {
                    mean += inputData[j];
                    s2 += inputData[j] * inputData[j];
                }
                // 这两行代码计算了均值和方差。
                mean /= channels;
                var = sqrt(s2 / channels - mean*mean + 1e-10);
                // 接下来是对output的每一个通道进行并行处理
                j = 0;
#ifdef __aarch64__
                float32x4_t means = vdupq_n_f32(mean);
                    float32x4_t vars = vdupq_n_f32(1.0 / var);
                    for (; j + 3 < channels; j += 4) {
                        float32x4_t va = vld1q_f32(gammaData + j), vb = vld1q_f32(betaData + j);
                        float32x4_t vi = vld1q_f32(inputData + j);
                        float32x4_t vo = vaddq_f32(vmulq_f32(vmulq_f32(vsubq_f32(vi, means), vars), va), vb);
                        vst1q_f32(outputData + j, vo);
                    }
#endif
                for (; j < channels; j++) {
                    float a = gammaData[j], b = betaData[j];
                    outputData[j] = (inputData[j] - mean) / var * a + b;
                }
                
                // 这两行代码更新了 inputData 和 outputData 的指针位置，
                // 以便在下一轮循环中处理下一个外部元素。
                inputData += channels;
                outputData += channels;
            }
            return;
        } else {
            // 这段代码同样是执行层归一化（Layer Normalization）操作，但这次的操作更为通用，
            // 能处理 inner 不等于 1 的情况，即每个通道有多个元素的情况。
            // 这是一个循环，对 input 中的每一个外部元素进行处理。
            for (int i = 0; i < outer; i++) {
                // 这两行代码将 mean 和 var 数组的所有元素初始化为 0。
                std::fill(mean, mean + inner, 0.f);
                std::fill(var, var + inner, 0.f);
                // 这行代码定义了一个指针 inputWalk，指向 inputData。
                float *inputWalk = inputData;
                // 这是一个循环，对每个通道进行处理。
                for (int j = 0; j < channels; j++) {
                   // 这是一个嵌套循环，对每个通道内的每个元素进行处理。
                    for (int k = 0; k < inner; k++) {
                        // 这行代码将当前元素的值加到对应的 mean 中，然后 inputWalk 指针向后移动。
                        mean[k] += *inputWalk++; 
                    }
                }
                // 这是另一个循环，计算每个通道的均值。
                for (int k = 0; k < inner; k++) {
                    mean[k] /= channels;
                }
                // 方差类似
                inputWalk = inputData;
                for (int j = 0; j < channels; j++) {
                    for (int k = 0; k < inner; k++) {
                        float x = (*inputWalk++) - mean[k];
                        var[k] += x * x;
                    }
                }
                for (int k = 0; k < inner; k++) {
                    var[k] = sqrt(var[k] / channels + 1e-5);
                }
                
                // 计算输出也是类似
                inputWalk = inputData;
                float *outputWalk = outputData;
                for (int j = 0; j < channels; j++) {
                    float a = gammaData[j], b = betaData[j];
                    for (int k = 0; k < inner; k++) {
                        *outputWalk++ = ((*inputWalk++) - mean[k]) / var[k] * a + b;
                    }
                }

                inputData += channels * inner;
                outputData += channels * inner;
            }
            delete[] mean;
            delete[] var;
        }
    }

void CpuLinearOp::Run(const std::string &opType, const fastllm::DataDict &datas,
                          const fastllm::FloatDict &floatParams, const fastllm::IntDict &intParams) {
//auto st = std::chrono::system_clock::now();
        Data &input = *(datas.find("input")->second);
        Data &output = *(datas.find("output")->second);
        Data &weight = *(datas.find("weight")->second);
        Data &bias = *(datas.find("bias")->second);

        output.Allocate(0.0f);
        int n = input.Count(0) / input.dims.back();
        int m = input.dims.back();
        int k = output.dims.back();
        
        // 这段代码处理权重数据类型为FLOAT32的情况。首先，它将输入、权重、输出和
        // 偏置数据的指针分别转换为 float* 类型的指针。对于偏置数据，如果其维度长度大于0，
        // 则获取其数据指针，否则设为nullptr。
        if (weight.dataType == DataType::FLOAT32) {
            float *inputData = (float *) input.cpuData;
            float *weightData = (float *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() > 0 ? (float *) bias.cpuData : nullptr;
            
            // 接下来，计算需要的线程数（threadNum）。这里用的是用户设定的线程数
            //（通过 GetThreads() 获得）。然后，每个线程负责的任务数（per）
            // 为 k（输出数据的最后一个维度）除以线程数。cur 用来表示当前任务的起始位置。
            int threadNum = GetThreads();
            int per = k / threadNum;
            int cur = 0;
            // 接着，创建线程池（通过 GetPool() 获取）和用于保存线程任务的std::future数组。
            // 对于每个线程，确定其需要处理的任务范围（从 cur 到 end），然后提交线程任务。
            // 线程任务是通过调用 FloatLinearPart 函数来执行的，该函数需要输入数据、
            // 权重数据、偏置数据、输出数据、输入维度（n）、权重维度（m）、输出维度（k）
            // 以及任务范围（从 cur 到 end）作为参数。
            auto pool = GetPool();
            std::vector <std::future <void> > futures;
            for (int i = 0; i < threadNum - 1; i++) {
                int end = cur + per + (cur + per * (threadNum - i) < k);
                futures.push_back(pool->Submit(FloatLinearPart, inputData, weightData, biasData, outputData,
                                                  n, m, k, cur, end));
                cur = end;
            }
            
            // 然后，主线程也执行一部分任务，处理范围为从 cur 到 k。
            FloatLinearPart(inputData, weightData, biasData, outputData, n, m, k, cur, k);
            // 最后，主线程等待所有子线程完成工作。通过调用 std::future::get() 
            // 方法来阻塞主线程，直到对应的子线程完成任务。
            // 这样，可以保证所有的线程任务都完成后，主线程才继续执行。
            for (int i = 0; i < futures.size(); i++) {
                futures[i].get();
            }
        } else if (weight.dataType == DataType::FLOAT16) {
            float *inputData = (float *) input.cpuData;
            uint16_t *weightData = (uint16_t *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() > 0 ? (float *) bias.cpuData : nullptr;
#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
            uint16_t *temp = new uint16_t[n * m];
            for (int i = 0; i < n * m; i++) {
                temp[i] = float_to_half(inputData[i]);
            }
            inputData = (float*)temp;
#endif
            int threadNum = GetThreads();
            int per = k / threadNum;
            int cur = 0;
            auto pool = GetPool();
            std::vector <std::future <void> > futures;
            for (int i = 0; i < threadNum - 1; i++) {
                int end = cur + per + (cur + per * (threadNum - i) < k);
                futures.push_back(pool->Submit(Float16LinearPart, inputData, weightData, biasData, outputData,
                                                  n, m, k, cur, end));
                cur = end;
            }

            Float16LinearPart(inputData, weightData, biasData, outputData, n, m, k, cur, k);
            for (int i = 0; i < futures.size(); i++) {
                futures[i].get();
            }
#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
            delete[] temp;
#endif
        } else if (weight.dataType == DataType::INT8) { // 这段代码处理权重数据类型为 INT8 的情况。
            // 这段代码首先对输入、权重、输出和偏置数据的指针进行类型转换，
            // 并根据偏置数据的维度是否大于0来决定是否获取偏置数据的指针。然后，它计算了权重数据的总和。
            float *inputData = (float *) input.cpuData;
            uint8_t *weightData = (uint8_t *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() > 0 ? (float *) bias.cpuData : nullptr;
            weight.CalcWeightSum();
            
            // 之后，代码创建一个std::vector<LowBitConfig>对象，
            // LowBitConfig是一个用于存储数据量化信息的类，包括最小值、最大值、位宽和零点。
            // 这些信息是通过遍历输入数据获得的。
            std::vector <LowBitConfig> inputConfigs;
            for (int i = 0; i < n; i++) {
                float minValue = 1e9, maxValue = -1e9;
                for (int j = 0; j < m; j++) {
                    minValue = std::min(minValue, inputData[i * m + j]);
                    maxValue = std::max(maxValue, inputData[i * m + j]);
                }
                inputConfigs.push_back(LowBitConfig(minValue, maxValue, 8, 0));
            }
            // 接着，创建一个std::vector<uint8_t>对象uinput，并将其大小设置为输入数据的大小（n * m）。
            // uinput中的每个元素都是输入数据元素经过inputConfigs中对应配置信息量化后的结果。
            // 注意这里的量化过程可能会根据是否定义了__AVX2__进行不同的处理。
            std::vector <uint8_t> uinput;
            uinput.resize(n * m);
            for (int i = 0; i < n * m; i++) {
#ifdef __AVX2__
                uinput[i] = inputConfigs[i / m].quantization(inputData[i]);
                uinput[i] = (uinput[i] + !uinput[i]) ^ 128;
#else
                uinput[i] = inputConfigs[i / m].quantization(inputData[i]);
#endif
            }
            
            // 随后，调用MultiplyMultiThread函数，使用多线程并行计算uinput和weightData的乘积，
            // 并将结果存储在outputData中。
            MultiplyMultiThread(uinput.data(), weightData, (int32_t*)outputData, n, m, k, GetThreads());
            // 这段代码的目的是把在使用INT8进行量化计算时由于量化造成的误差进行修正，
            // 使得结果更接近于使用浮点数进行计算的结果。也就是反量化过程。
            for (int i = 0; i < n; i++) {
                // 这一步中，对于每一个输入向量（i从0到n），代码首先初始化inputSum为0，
                // 然后遍历输入向量的每个元素（j从0到m），将元素值加到inputSum上。
                // 如果定义了__AVX2__，则在加到inputSum之前，元素值会先与128进行异或操作。
                uint32_t inputSum = 0;
                for (int j = 0; j < m; j++) {
#ifdef __AVX2__
                    inputSum += uinput[i * m + j] ^ 128;
#else
                    inputSum += uinput[i * m + j];
#endif
                }
                
                // 接下来，代码遍历每个输出元素（j从0到k），并按照以下步骤进行调整和缩放：
                for (int j = 0; j < k; j++) {
                    // 首先，获取输出元素的原始值value。
                    int value = ((int32_t*)outputData)[i * k + j];
#ifdef __AVX2__
                    // 如果定义了__AVX2__，则value会增加128 * weight.weightSum[j]、
                    // 128 * inputSum，并减去m * 128 * 128。
                    value += (128 * weight.weightSum[j]);
                    value += (128 * inputSum);
                    value -= m * 128 * 128;
#endif
                    value -= weight.weightSum[j] * inputConfigs[i].zeroPoint;
                    value -= inputSum * weight.perChannelsConfigs[j].zeroPoint;
                    value += (int)inputConfigs[i].zeroPoint * weight.perChannelsConfigs[j].zeroPoint * m;
                    outputData[i * k + j] = weight.perChannelsConfigs[j].scale * inputConfigs[i].scale * value +
                                            (biasData == nullptr ? 0.0 : biasData[j]);
                }
            }
        } else if (weight.dataType == DataType::INT4 || weight.dataType == DataType::INT4_NOZERO) {
            float *inputData = (float *) input.cpuData;
            uint8_t *weightData = (uint8_t *) weight.cpuData;
            float *outputData = (float *) output.cpuData;
            float *biasData = bias.dims.size() > 0 ? (float *) bias.cpuData : nullptr;
            weight.CalcWeightSum();

            std::vector <LowBitConfig> inputConfigs;
            for (int i = 0; i < n; i++) {
                float minValue = 1e9, maxValue = -1e9;
                for (int j = 0; j < m; j++) {
                    minValue = std::min(minValue, inputData[i * m + j]);
                    maxValue = std::max(maxValue, inputData[i * m + j]);
                }
                inputConfigs.push_back(LowBitConfig(minValue, maxValue, 8, 0));
            }
            std::vector <uint8_t> uinput;
            uinput.resize(n * m);
            for (int i = 0; i < n * m; i++) {
                uinput[i] = inputConfigs[i / m].quantization(inputData[i]);
            }
#ifdef __AVX__
            uint8_t *temp = new uint8_t[32];
            for (int i = 0; i < n; i++) {
                for (int j = 0; j + 31 < m; j += 32) {
                    memcpy(temp, uinput.data() + i * m + j, 32);
                    for (int k = 0; k < 16; k++) {
                        uinput[i * m + j + k] = temp[k * 2 + 1];
                        uinput[i * m + j + k + 16] = temp[k * 2];
                    }
                }
            }
            delete[] temp;
#endif
            if (weight.dataType == DataType::INT4) {
                MultiplyInt4MultiThread(uinput.data(), weightData, (int32_t *) outputData, n, m, k,
                                        weight.weightSum.data(), weight.zeros.data(), weight.scales.data(), biasData,
                                        inputConfigs, GetThreads());
            } else {
                MultiplyInt4NoZeroMultiThread(uinput.data(), weightData, (int32_t *) outputData, n, m, k,
                                        weight.weightSum.data(), weight.mins.data(), weight.scales.data(), biasData,
                                        inputConfigs, GetThreads());
            }
        } else {
            ErrorInFastLLM("Linear error: unsupport weight's dataType.\n");
        }
//float spend = GetSpan(st, std::chrono::system_clock::now());
//float gops = (float)n * m * k / spend / 1e9;
// printf("n = %d, m = %d, k = %d, spend %f s, gops = %f\n", n, m, k, spend, gops);
    }

//a = [n, m], b = [k, m], c = aT(b') = [n, k]
    void MultiplyMultiThread(uint8_t *a, uint8_t *b, int32_t *c, int n, int m, int k, int threadNum) {
        int per = k / threadNum;
        int cur = 0;
        if (threadNum == 1) {
            Multiply(a, b + cur * m, c + cur, n, m, k - cur, k);
        } else {
            auto pool = GetPool();
            std::vector<std::future<void> > futures;
            for (int i = 0; i < threadNum; i++) {
                int end = cur + per + (cur + per * (threadNum - i) < k);
                if (i == threadNum - 1) {
                    end = k;
                }
                futures.push_back(pool->Submit(Multiply, a, b + cur * m, c + cur, n, m, end - cur, k));
                cur = end;
            }
            for (int i = 0; i < futures.size(); i++) {
                futures[i].get();
            }
        }
    }

   //a = [n, m], b = [k, m], c = aT(b') = [n, k]
    void Multiply(uint8_t *a, uint8_t *b, int32_t *c, int n, int m, int k, int kstride) {
#ifdef __ARM_FEATURE_DOTPROD
        int block = 0;
        for (; block < n; block++) {
            uint8_t *weightWalk = b;
            uint8_t *inputStart = a + block * m;

            for (int i = 0; i < k; i++) {
                int value = 0;
                uint8_t *inputWalk = inputStart;
                int j = 0;
                uint32x4_t sum0 = {0, 0, 0, 0};
                for (; j + 31 < m; j += 32) {
                    uint8x16_t vi = vld1q_u8(inputWalk);
                    uint8x16_t vi0 = vld1q_u8(inputWalk + 16);
                    uint8x16_t vw = vld1q_u8(weightWalk);
                    uint8x16_t vw0 = vld1q_u8(weightWalk + 16);
                    sum0 = vdotq_u32(sum0, vi, vw);
                    sum0 = vdotq_u32(sum0, vi0, vw0);
                    inputWalk += 32;
                    weightWalk += 32;
                }

                value += sum0[0] + sum0[1] + sum0[2] + sum0[3];
                for (; j < m; j++) {
        value += (int)(*(weightWalk++)) * (*(inputWalk++));
       }
                c[block * kstride + i] = value;
            }
        }
#elif defined(__aarch64__)
        int block = 0;
        for (; block < n; block++) {
            uint8_t *weightWalk = b;
            uint8_t *inputStart = a + block * m;

            for (int i = 0; i < k; i++) {
                int value = 0;
                uint8_t *inputWalk = inputStart;

                int per = 64;
                int cnt = m / per;
                int sur = m % per;

                uint32x4_t sum = {0};
                uint16x8_t temp = {0};
                uint16x8_t temp1 = {0};
                uint16x8_t temp2 = {0};
                uint16x8_t temp3 = {0};
                uint16x8_t temp4 = {0};
                uint16x8_t temp5 = {0};
                uint16x8_t temp6 = {0};
                uint16x8_t temp7 = {0};

                while (cnt--) {
                    temp = vmull_u8(vld1_u8(inputWalk), vld1_u8(weightWalk));
                    temp1 = vmull_u8(vld1_u8(inputWalk + 8), vld1_u8(weightWalk + 8));
                    temp2 = vmull_u8(vld1_u8(inputWalk + 16), vld1_u8(weightWalk + 16));
                    temp3 = vmull_u8(vld1_u8(inputWalk + 24), vld1_u8(weightWalk + 24));
                    temp4 = vmull_u8(vld1_u8(inputWalk + 32), vld1_u8(weightWalk + 32));
                    temp5 = vmull_u8(vld1_u8(inputWalk + 40), vld1_u8(weightWalk + 40));
                    temp6 = vmull_u8(vld1_u8(inputWalk + 48), vld1_u8(weightWalk + 48));
                    temp7 = vmull_u8(vld1_u8(inputWalk + 56), vld1_u8(weightWalk + 56));

                    sum = vpadalq_u16(sum, temp);
                    sum = vpadalq_u16(sum, temp1);
                    sum = vpadalq_u16(sum, temp2);
                    sum = vpadalq_u16(sum, temp3);
                    sum = vpadalq_u16(sum, temp4);
                    sum = vpadalq_u16(sum, temp5);
                    sum = vpadalq_u16(sum, temp6);
                    sum = vpadalq_u16(sum, temp7);

                    inputWalk += per;
                    weightWalk += per;
                }

                value += (sum[0] + sum[1] + sum[2] + sum[3]);
                while (sur--) {
                    value += (int)(*(weightWalk++)) * (*(inputWalk++));
                }

                c[block * kstride + i] = value;
            }
        }
#elif defined(__AVX__)
        int block = 0;
        for (; block < n; block++) {
            uint8_t *weightWalk = b;
            uint8_t *inputStart = a + block * m;

            for (int i = 0; i < k; i++) {
                uint8_t *inputWalk = inputStart;

                c[block * kstride + i] = DotU8U8(inputWalk, weightWalk, m);
                weightWalk += m;
            }
        }
#else
        int block = 0;
     for (; block < n; block++) {
      uint8_t *weightWalk = b;
      uint8_t *inputStart = a + block * m;

      for (int i = 0; i < k; i++) {
       int value = 0;
       uint8_t *inputWalk = inputStart;
       for (int j = 0; j < m; j++) {
        value += (int)(*(weightWalk++)) * (*(inputWalk++));
       }

       c[block * kstride + i] = value;
      }
     }
#endif
    }

# 这个类是用来创建旋转位置编码（Rotary Position Embedding）的。
# Llama模型引入了旋转位置编码，以改进长序列处理的性能。
class LlamaRotaryEmbedding(torch.nn.Module):
    # 这是类的初始化方法，接收四个参数：dim（嵌入的维度），max_position_embeddings
    # （最大的位置嵌入长度，默认为2048），base（基数，默认为10000）和device（设备类型，例如CPU或GPU）。
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        self.dim = dim # 将输入的dim参数保存到self.dim属性中。
        # # 将输入的max_position_embeddings参数保存到self.max_position_embeddings属性中。
        self.max_position_embeddings = max_position_embeddings
        # 将输入的base参数保存到self.base属性中。
        self.base = base
        # 计算逆频率并保存到变量inv_freq中。逆频率是一种用于位置编码的技巧，
        # 它可以帮助模型更好地捕捉位置信息。
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        # 将inv_freq保存到模型的缓存中。register_buffer是PyTorch nn.Module的一个方法，
        # 它用于保存一些不需要计算梯度的变量。
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        # 调用_set_cos_sin_cache方法，预先计算并保存正弦和余弦的缓存值。
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )
    
    # 这是一个私有方法，接收三个参数：seq_len（序列长度），device（设备类型）和dtype（数据类型）
    def _set_cos_sin_cache(self, seq_len, device, dtype):
        # 将输入的seq_len参数保存到self.max_seq_len_cached属性中。
        self.max_seq_len_cached = seq_len
        # 生成一个长度为max_seq_len_cached的序列，并保存到变量t中。
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
        
        # 使用外积计算频率和t的乘积，结果保存到变量freqs中。
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        # 将频率的两份副本拼接在一起，结果保存到变量emb中。
        emb = torch.cat((freqs, freqs), dim=-1)
        # 计算emb的余弦值，然后将结果保存到模型的缓存中。
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
        # 计算emb的正弦值，然后将结果保存到模型的缓存中。
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
    
    # 这是模型的前向传播方法，接收两个参数：x（输入数据）和seq_len（序列长度）。
    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        # 如果输入的序列长度大于缓存的最大序列长度，那么调用_set_cos_sin_cache方法，更新缓存。
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
        
        # 返回对应输入位置的正弦和余弦值。这些值将用于旋转位置编码。
        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )

def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

void CudaLlamaRotatePosition2DOp::Run(const std::string &opType, const fastllm::DataDict &datas,
                                     const fastllm::FloatDict &floatParams, const fastllm::IntDict &intParams) {
        Data &data = *(datas.find("input")->second);
        Data &positionIds = *(datas.find("positionIds")->second);
        Data &sinData = *(datas.find("sin")->second);
        Data &cosData = *(datas.find("cos")->second);
        int rotaryDim = intParams.find("rotaryDim") != intParams.end() ? intParams.find("rotaryDim")->second : 128;

        FastllmCudaLlamaRotatePosition2D(data, positionIds, sinData, cosData, rotaryDim);
    }

// 这是一个在 GPU 上运行的 CUDA 函数，用于执行 Llama 模型的位置编码旋转操作。
// data：输入的数据，这个数据将会被旋转。
// positionIds：位置编码的数据。
// sinData，cosData：用于旋转的 sin 和 cos 值。
// rotaryDim：旋转的维度。
bool FastllmCudaLlamaRotatePosition2D(fastllm::Data &data, const fastllm::Data &positionIds,
                                      const fastllm::Data &sinData, const fastllm::Data &cosData, int rotaryDim) {
    // 使用 FastllmCudaPrepareInput 函数将输入的数据从 CPU 复制到 GPU。
    // 这个函数会返回一个指向 GPU 内存的指针。                                  
    float *cudaData = (float *) FastllmCudaPrepareInput(data);
    float *cudaPositionIds = (float *) FastllmCudaPrepareInput(positionIds);
    float *cudaSin = (float *) FastllmCudaPrepareInput(sinData);
    float *cudaCos = (float *) FastllmCudaPrepareInput(cosData);
    
    // 计算旋转操作需要的一些参数，包括 outer，spatial，bs，len，n 和 m。
    // 这些参数是用于确定 CUDA 核函数的执行配置和一些数据操作的。
    int outer = data.dims[0] * data.dims[1];
    int spatial = data.Count(2);
    int bs = data.dims[0], len = data.dims[1];
    int n = data.dims[2], m = data.dims[3];
    // 调用 CUDA 核函数 FastllmLlamaRotatePosition2DKernel 来在 GPU 上执行位置编码的旋转操作。
    // <<<outer * n, min(rotaryDim, m / 2)>>> 是 CUDA 中定义并行线程块和线程的语法，
    // outer * n 是线程块的数量，min(rotaryDim, m / 2) 是每个线程块中的线程数量。
    // 核函数的参数包括之前准备的数据和一些计算参数。
    FastllmLlamaRotatePosition2DKernel <<< outer * n, min(rotaryDim, m / 2) >>> (cudaData, cudaPositionIds, cudaSin, cudaCos,
                                                                                 len, bs, spatial, n, m,
                                                                                 (int)positionIds.dims.back(), (int)sinData.dims[1], rotaryDim);

    // 使用 FastllmCudaFinishInput 函数释放 positionIds，sinData 和 cosData 在 GPU 上的内存。
    // 这些数据在这个函数中不再需要。
    FastllmCudaFinishInput(positionIds, cudaPositionIds);
    FastllmCudaFinishInput(sinData, cudaSin);
    FastllmCudaFinishInput(cosData, cudaCos);
    // 使用 FastllmCudaFinishOutput 函数将旋转后的数据从 GPU 复制回 CPU。
    // 这个函数也会释放 data 在 GPU 上的内存。
    FastllmCudaFinishOutput(data, cudaData);
    return true;
}

// float *data：输入数据，大小为 [bs, len, n, m]，其中 bs 是批量大小，
// len 是序列长度，n 是头的数量，m 是每个头的维度。
// float *positionIds：位置编码的索引，大小为 [bs, len]。
// float *sin 和 float *cos：预先计算的正弦和余弦值，用于旋转编码。
// int len, int bs, int spatial, int n, int m：输入数据的各个维度大小。
// int partStride 和 int sinCosStride：用于索引 positionIds 和 sin/cos 的步长。
// int rotateDim：旋转维度。
__global__ void FastllmLlamaRotatePosition2DKernel(float *data, float *positionIds, float *sin, float *cos,
                                                   int len, int bs, int spatial, int n, int m, int partStride, int sinCosStride, int rotateDim) {
    // 首先，计算出当前线程应处理的位置 o，长度 l 和批次 b。
    int o = (blockIdx.x / n);
    int l = o % len;
    int b = o / len;
    int j = threadIdx.x;
    // 然后，根据 positionIds 获取对应的旋转角度的正弦值 curSin 和余弦值 curCos。
    int index = (int) (positionIds[b * partStride + l]);

    float curSin = sin[index * sinCosStride + j];
    float curCos = cos[index * sinCosStride + j];
    float *d = (float *) data + o * spatial + j;
    int i = blockIdx.x % n;
    // 接着，获取输入数据对应位置的值 va 和 vb。
    float va = d[i * m], vb = d[i * m + m / 2];
    // 最后，根据旋转矩阵的公式，计算旋转后的值，并将结果写回输入数据中。
    d[i * m] = va * curCos - vb * curSin;
    d[i * m + m / 2] = va * curSin + vb * curCos;
}

void CpuLlamaRotatePosition2DOp::Run(const std::string &opType, const fastllm::DataDict &datas,
                                    const fastllm::FloatDict &floatParams, const fastllm::IntDict &intParams) {
        Data &data = *(datas.find("input")->second);
        Data &positionIds = *(datas.find("positionIds")->second);
        Data &sinData = *(datas.find("sin")->second);
        Data &cosData = *(datas.find("cos")->second);
        int rotaryDim = intParams.find("rotaryDim") != intParams.end() ? intParams.find("rotaryDim")->second : 128;

        int bs = data.dims[0], len = data.dims[1];
        int spatial = data.Count(2);
        int n = data.dims[2], m = data.dims[3];
        int stride = (int)sinData.dims[1];
        for (int b = 0; b < bs; b++) {
            for (int l = 0; l < len; l++) {
                int index = (int) ((float *) positionIds.cpuData)[b * positionIds.dims.back() + l];
                float *sin = ((float *) sinData.cpuData) + stride * index;
                float *cos = ((float *) cosData.cpuData) + stride * index;
                float *d = (float *) data.cpuData + (b * len + l) * spatial;
                for (int i = 0; i < n; i++) {
                    for (int j = 0; j < rotaryDim && j < m / 2; j++) {
                        float a = d[j], b = d[j + m / 2];
                        d[j] = a * cos[j] - b * sin[j];
                        d[j + m / 2] = a * sin[j] + b * cos[j];
                    }

                    d += m;
                }
            }
        }
    }

if (generationConfig.IsSimpleGreedy()) {
            // 对 logits 进行 TopK 操作，将结果存储在 topk 中。
            // 这里的 TopK 操作是找到 logits 中最大的 K 个值，这里 K=1，所以是找到最大值。
            TopK(logits, topk, 1); 
            topk.ToDevice(DataDevice::CPU);
            for (int b = 0; b < batch; b++) {
                int base = (maxLen - 1) * batch + b; // 计算基础索引值 base。
                // 将 topk 中对应索引的值取整并添加到 lastRet 中。
                lastRet.push_back((int) (((float *) topk.cpuData)[base * 2] + 1e-3));
            }
        } else {
            for (int b = 0; b < batch; b++) {
                int base = (maxLen - 1) * batch + b; // 计算基础索引值 base。
                // 使用 LLMSampling 方法进行抽样，将结果添加到 lastRet 中。
                lastRet.push_back(LLMSampling(logits, base, generationConfig, lastTokens.units[b]));
            }
        }

// 这段代码是一个用于从给定的 logits（通常表示预测的概率分布）进行采样的函数，
// 采样策略主要受 GenerationConfig 和 LastTokensUnit 参数的影响。
int LLMSampling(Data &logits, int outerOffset,
                    const GenerationConfig &config, const LastTokensUnit &tokens) {
        // 将 logits 数据从当前设备转移到 CPU。
        logits.ToDevice(DataDevice::CPU);
        // 从 logits 的维度中获取词汇量 vocabSize。
        int vocabSize = logits.dims.back();
        // 计算 base 指针，指向要处理的 logits 的开始位置。
        float *base = ((float*)logits.cpuData) + outerOffset * vocabSize;
        
        // 判断 config.repeat_penalty 是否不等于1，如果不等于1，
        // 则对 tokens.tokenSet 中每个 id 对应的 base[id] 值进行修改。
        if (fabs(config.repeat_penalty - 1.0) > 1e-6) {
            for (int id : tokens.tokenSet) {
                base[id] = (base[id] < 0 ? base[id] * config.repeat_penalty : base[id] / config.repeat_penalty);
            }
        }
        // 计算温度的倒数 invTemp。
        float invTemp = 1.0f / config.temperature;
        // 定义一个向量 v，用于存储 <logit值，索引>。
        std::vector <std::pair <float, int> > v;
        // 遍历每个 logit，将其值乘以 invTemp，并存入 v 中。
        for (int i = 0; i < vocabSize; i++) {
            v.push_back(std::make_pair(-base[i] * invTemp, i));
        }
        // 计算 topk，它是词汇量 vocabSize 和 config.top_k 中的较小值。
        int topk = std::min(vocabSize, config.top_k);
        // 对 v 中的前 topk 个元素进行排序。
        std::partial_sort(v.begin(), v.begin() + topk, v.end());
        // 初始化 psum 和 maxValue，maxValue 是 v 中最大的元素。
        float psum = 0.0, maxValue = -v.begin()->first;
        // 定义一个向量 ps，用于存储处理后的概率。
        std::vector <float> ps;
        // 遍历 v 中的前 topk 个元素，将其值取 exp 并减去 maxValue，存入 ps，同时更新 psum。
        for (int i = 0; i < topk; i++) {
            ps.push_back(expf(-v[i].first - maxValue));
            psum += ps.back();
        }
        float curSum = 0.0;
        // 遍历 ps，将其每个元素除以 psum 并更新 curSum，
        // 当 curSum 大于 config.top_p 时，更新 topk 并退出循环。
        for (int i = 0; i < topk; i++) {
            ps[i] /= psum;
            curSum += ps[i];
            if (curSum > config.top_p) {
                topk = i + 1;
                break;
            }
        }
        // 生成一个随机数 rnd。
        float rnd = fastllmRandom.randP();
        curSum = 0.0;
        // 遍历 ps 中的前 topk 个元素，将其累加到 curSum，
        // 当 curSum 大于 rnd 或者达到最后一个元素时，
        // 返回对应 v[i].second，也就是返回采样得到的 id。
        for (int i = 0; i < topk; i++) {
            curSum += ps[i];
            if (curSum > rnd || i == topk - 1) {
                return v[i].second;
            }
        }
        // 如果以上步骤都没有返回，那么返回 -1。
        return -1;
    }