第18章 SIMD向量化与AI推理引擎¶

// SIMD数据类型
__m128    // 128-bit, 4×float
__m128d   // 128-bit, 2×double
__m128i   // 128-bit, 整数(8/16/32/64)

__m256    // 256-bit, 8×float
__m256d   // 256-bit, 4×double
__m256i   // 256-bit, 整数

__m512    // 512-bit, 16×float (AVX-512)

// 命名规则: _mm<width>_<operation>_<type>
// width: 空(128), 256, 512
// type: ps(packed single), pd(packed double),
//       epi32(packed int32), si128(128-bit integer)

// 示例
_mm256_add_ps    // 256-bit float加法
_mm256_mul_ps    // 256-bit float乘法
_mm256_fmadd_ps  // 256-bit fused multiply-add: a*b+c

#include <immintrin.h>
#include <cstdlib>
#include <cstdio>

int main() {
    // ===== 内存对齐分配 =====
    // AVX要求32字节对齐, AVX-512要求64字节对齐
    float* a = (float*)aligned_alloc(32, 8 * sizeof(float));
    float* b = (float*)aligned_alloc(32, 8 * sizeof(float));
    float* c = (float*)aligned_alloc(32, 8 * sizeof(float));

    for (int i = 0; i < 8; i++) {
        a[i] = (float)i;
        b[i] = (float)(i * 2);
    }

    // ===== 对齐加载 vs 非对齐加载 =====
    __m256 va = _mm256_load_ps(a);       // 对齐加载（地址必须32字节对齐）
    __m256 vb = _mm256_load_ps(b);       // 性能最优
    // __m256 vu = _mm256_loadu_ps(ptr); // 非对齐加载（安全但可能慢）

    // ===== 基本运算 =====
    __m256 vc = _mm256_add_ps(va, vb);   // c = a + b
    __m256 vd = _mm256_mul_ps(va, vb);   // d = a * b
    __m256 ve = _mm256_fmadd_ps(va, vb, vc); // e = a*b + c (FMA)

    // ===== 存储结果 =====
    _mm256_store_ps(c, vc);

    for (int i = 0; i < 8; i++) {
        printf("c[%d] = %.1f\n", i, c[i]);
    }

    free(a); free(b); free(c);
    return 0;
}

#include <immintrin.h>

// 标量版本
float dot_product_scalar(const float* a, const float* b, int n) {
    float sum = 0.0f;
    for (int i = 0; i < n; i++) {
        sum += a[i] * b[i];
    }
    return sum;
}

// AVX2 SIMD版本
float dot_product_avx2(const float* a, const float* b, int n) {
    __m256 sum_vec = _mm256_setzero_ps();  // 8个float的累加器

    int i = 0;
    // 主循环: 每次处理8个float
    for (; i + 7 < n; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        sum_vec = _mm256_fmadd_ps(va, vb, sum_vec); // sum += a*b
    }

    // 水平求和: 将8个float归约为1个
    // [s0,s1,s2,s3,s4,s5,s6,s7]
    __m128 hi = _mm256_extractf128_ps(sum_vec, 1);  // [s4,s5,s6,s7]
    __m128 lo = _mm256_castps256_ps128(sum_vec);     // [s0,s1,s2,s3]
    __m128 sum128 = _mm_add_ps(lo, hi);              // [s0+s4,s1+s5,s2+s6,s3+s7]
    sum128 = _mm_hadd_ps(sum128, sum128);             // 水平加
    sum128 = _mm_hadd_ps(sum128, sum128);
    float result = _mm_cvtss_f32(sum128);

    // 处理尾部元素
    for (; i < n; i++) {
        result += a[i] * b[i];
    }

    return result;
}

#include <immintrin.h>
#include <cstring>
#include <chrono>
#include <cstdio>

// ===== Naive GEMM: C = A × B =====
void gemm_naive(const float* A, const float* B, float* C,
                int M, int N, int K) {
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
}

// ===== AVX2 GEMM: 向量化内循环 =====
void gemm_avx2(const float* A, const float* B, float* C,
               int M, int N, int K) {
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j += 8) {  // 每次处理8列
            __m256 c_vec = _mm256_setzero_ps();

            for (int k = 0; k < K; k++) {
                // 广播A[i][k]到8个lane
                __m256 a_broadcast = _mm256_set1_ps(A[i * K + k]);
                // 加载B[k][j:j+8]
                __m256 b_vec = _mm256_loadu_ps(&B[k * N + j]);
                // FMA: c += a * b
                c_vec = _mm256_fmadd_ps(a_broadcast, b_vec, c_vec);
            }

            _mm256_storeu_ps(&C[i * N + j], c_vec);
        }
    }
}

// ===== 分块(Tiling) + AVX2 GEMM =====
void gemm_tiled_avx2(const float* A, const float* B, float* C,
                     int M, int N, int K) {
    constexpr int TILE_M = 4;
    constexpr int TILE_N = 16;  // 2个AVX寄存器
    constexpr int TILE_K = 64;

    memset(C, 0, M * N * sizeof(float));

    for (int i0 = 0; i0 < M; i0 += TILE_M) {
        for (int k0 = 0; k0 < K; k0 += TILE_K) {
            for (int j0 = 0; j0 < N; j0 += TILE_N) {
                int i_end = std::min(i0 + TILE_M, M);
                int k_end = std::min(k0 + TILE_K, K);
                int j_end = std::min(j0 + TILE_N, N);

                for (int i = i0; i < i_end; i++) {
                    // 加载C的当前块到寄存器
                    __m256 c0 = _mm256_loadu_ps(&C[i * N + j0]);
                    __m256 c1 = _mm256_loadu_ps(&C[i * N + j0 + 8]);

                    for (int k = k0; k < k_end; k++) {
                        __m256 a_val = _mm256_set1_ps(A[i * K + k]);
                        __m256 b0 = _mm256_loadu_ps(&B[k * N + j0]);
                        __m256 b1 = _mm256_loadu_ps(&B[k * N + j0 + 8]);

                        c0 = _mm256_fmadd_ps(a_val, b0, c0);
                        c1 = _mm256_fmadd_ps(a_val, b1, c1);
                    }

                    _mm256_storeu_ps(&C[i * N + j0], c0);
                    _mm256_storeu_ps(&C[i * N + j0 + 8], c1);
                }
            }
        }
    }
}

// ===== 性能对比 =====
int main() {
    constexpr int M = 512, N = 512, K = 512;

    float* A = (float*)aligned_alloc(32, M * K * sizeof(float));
    float* B = (float*)aligned_alloc(32, K * N * sizeof(float));
    float* C_naive = (float*)aligned_alloc(32, M * N * sizeof(float));
    float* C_avx = (float*)aligned_alloc(32, M * N * sizeof(float));
    float* C_tiled = (float*)aligned_alloc(32, M * N * sizeof(float));

    // 初始化
    for (int i = 0; i < M * K; i++) A[i] = (float)(i % 100) / 100.0f;
    for (int i = 0; i < K * N; i++) B[i] = (float)(i % 100) / 100.0f;

    auto bench = [](auto fn, const float* A, const float* B, float* C,
                    int M, int N, int K, const char* name) {
        auto start = std::chrono::high_resolution_clock::now();
        fn(A, B, C, M, N, K);
        auto end = std::chrono::high_resolution_clock::now();
        double ms = std::chrono::duration<double, std::milli>(end - start).count();
        double gflops = 2.0 * M * N * K / (ms * 1e6);
        printf("%s: %.2f ms, %.2f GFLOPS\n", name, ms, gflops);
    };

    bench(gemm_naive, A, B, C_naive, M, N, K, "Naive");
    bench(gemm_avx2, A, B, C_avx, M, N, K, "AVX2");
    bench(gemm_tiled_avx2, A, B, C_tiled, M, N, K, "Tiled+AVX2");

    free(A); free(B); free(C_naive); free(C_avx); free(C_tiled);
    return 0;
}

// ✅ 编译器可以自动向量化的循环
void add_arrays(float* c, const float* a, const float* b, int n) {
    for (int i = 0; i < n; i++) {
        c[i] = a[i] + b[i];  // 简单逐元素操作
    }
}

// ❌ 无法自动向量化的情况

// 1. 循环依赖
void cumsum(float* a, int n) {
    for (int i = 1; i < n; i++) {
        a[i] += a[i-1];  // 依赖前一个元素
    }
}

// 2. 条件分支
void conditional(float* c, const float* a, int n) {
    for (int i = 0; i < n; i++) {
        if (a[i] > 0) c[i] = a[i];
        else c[i] = -a[i];
        // 可用 _mm256_blendv_ps 手动向量化
    }
}

// 3. 函数调用
void with_call(float* c, const float* a, int n) {
    for (int i = 0; i < n; i++) {
        c[i] = my_func(a[i]);  // 未内联的函数调用
    }
}

// 使用#pragma提示编译器
#pragma GCC ivdep           // 告诉GCC忽略依赖（程序员保证无依赖）
#pragma clang loop vectorize(enable)
#pragma omp simd            // OpenMP SIMD指令

// __restrict__ 消除指针别名
void add_restrict(float* __restrict__ c,
                  const float* __restrict__ a,
                  const float* __restrict__ b, int n) {
    for (int i = 0; i < n; i++) {
        c[i] = a[i] + b[i];
    }
}

#include <arm_neon.h>

// NEON数据类型
// float32x4_t  : 4个float32
// int32x4_t    : 4个int32
// uint8x16_t   : 16个uint8
// float16x8_t  : 8个float16 (ARMv8.2-A)

// NEON向量点积
float dot_product_neon(const float* a, const float* b, int n) {
    float32x4_t sum_vec = vdupq_n_f32(0.0f);

    int i = 0;
    for (; i + 3 < n; i += 4) {
        float32x4_t va = vld1q_f32(a + i);
        float32x4_t vb = vld1q_f32(b + i);
        sum_vec = vfmaq_f32(sum_vec, va, vb);  // FMA
    }

    // 水平求和: NEON有专门的指令
    float result = vaddvq_f32(sum_vec);  // ARMv8.1: 直接水平加

    for (; i < n; i++) {
        result += a[i] * b[i];
    }
    return result;
}

#include <arm_neon.h>

// INT8矩阵乘法（量化推理核心）
// C[M×N] = A[M×K] × B[K×N]，结果为int32防止int8溢出
// 算法：广播A元素，向量化沿B的列方向（N维度）
void gemm_int8_neon(const int8_t* A, const int8_t* B, int32_t* C,
                    int M, int N, int K) {
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j += 8) {
            // 8路int32累加器，对应C[i][j:j+8]
            int32x4_t sum0 = vdupq_n_s32(0);  // C[i][j+0..j+3]
            int32x4_t sum1 = vdupq_n_s32(0);  // C[i][j+4..j+7]

            for (int k = 0; k < K; k++) {
                // 广播 A[i][k] 为int16向量（标量→8路复制）
                int16x8_t va = vdupq_n_s16((int16_t)A[i * K + k]);

                // 加载 B[k][j:j+8]（B的第k行，第j~j+7列）
                int8x8_t vb = vld1_s8(B + k * N + j);
                // int8 → int16 扩展（防止乘法溢出）
                int16x8_t vb_wide = vmovl_s8(vb);

                // int16 × int16 → int32 乘加累积
                // sum0 += A[i][k] * B[k][j+0..j+3]
                sum0 = vmlal_s16(sum0, vget_low_s16(va), vget_low_s16(vb_wide));
                // sum1 += A[i][k] * B[k][j+4..j+7]
                sum1 = vmlal_s16(sum1, vget_high_s16(va), vget_high_s16(vb_wide));
            }

            // 数学验证：sum0[x] = Σ_k A[i][k]·B[k][j+x] = C[i][j+x] ✓
            vst1q_s32(C + i * N + j, sum0);
            vst1q_s32(C + i * N + j + 4, sum1);
        }
    }
}

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime.h>
#include <fstream>
#include <iostream>
#include <vector>
#include <memory>

using namespace nvinfer1;

// ===== 1. 日志类（必须实现） =====
class Logger : public ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING) {
            std::cout << "[TRT] " << msg << std::endl;
        }
    }
};

// ===== 智能指针删除器 =====
struct TRTDeleter {
    template <typename T>
    void operator()(T* obj) const {
        delete obj;
    }
};
template <typename T>
using TRTUniquePtr = std::unique_ptr<T, TRTDeleter>;

// ===== 2. 从ONNX构建TensorRT Engine =====
std::vector<char> buildEngine(const std::string& onnx_path,
                               Logger& logger,
                               bool use_fp16 = true) {
    // 创建Builder
    auto builder = TRTUniquePtr<IBuilder>(createInferBuilder(logger));

    // 创建Network (显式batch)
    const auto explicitBatch = 1U <<
        static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = TRTUniquePtr<INetworkDefinition>(
        builder->createNetworkV2(explicitBatch));

    // 创建ONNX Parser
    auto parser = TRTUniquePtr<nvonnxparser::IParser>(
        nvonnxparser::createParser(*network, logger));

    // 解析ONNX模型
    if (!parser->parseFromFile(onnx_path.c_str(),
            static_cast<int>(ILogger::Severity::kWARNING))) {
        std::cerr << "Failed to parse ONNX file: " << onnx_path << std::endl;
        return {};
    }

    // 构建配置
    auto config = TRTUniquePtr<IBuilderConfig>(builder->createBuilderConfig());
    config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1ULL << 30); // 1GB

    if (use_fp16 && builder->platformHasFastFp16()) {
        config->setFlag(BuilderFlag::kFP16);
        std::cout << "FP16 mode enabled" << std::endl;
    }

    // 构建序列化引擎
    auto serialized = TRTUniquePtr<IHostMemory>(
        builder->buildSerializedNetwork(*network, *config));

    if (!serialized) {
        std::cerr << "Failed to build engine" << std::endl;
        return {};
    }

    // 返回序列化数据
    std::vector<char> engine_data(
        static_cast<char*>(serialized->data()),
        static_cast<char*>(serialized->data()) + serialized->size());

    return engine_data;
}

// ===== 3. 保存/加载Engine =====
void saveEngine(const std::vector<char>& engine_data,
                const std::string& path) {
    std::ofstream file(path, std::ios::binary);
    file.write(engine_data.data(), engine_data.size());
}

std::vector<char> loadEngine(const std::string& path) {
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<char> data(size);
    file.read(data.data(), size);
    return data;
}

// ===== 4. 推理类 =====
class TRTInference {
public:
    TRTInference(const std::vector<char>& engine_data, Logger& logger) {
        runtime_.reset(createInferRuntime(logger));
        engine_.reset(runtime_->deserializeCudaEngine(
            engine_data.data(), engine_data.size()));
        context_.reset(engine_->createExecutionContext());

        cudaStreamCreate(&stream_);

        // 分配GPU内存
        allocateBuffers();
    }

    ~TRTInference() {
        for (auto& buf : device_buffers_) {
            cudaFree(buf);
        }
        cudaStreamDestroy(stream_);
    }

    // 执行推理
    std::vector<float> infer(const std::vector<float>& input) {
        // 拷贝输入到GPU
        cudaMemcpyAsync(device_buffers_[input_idx_], input.data(),
                        input.size() * sizeof(float),
                        cudaMemcpyHostToDevice, stream_);

        // 设置tensor地址
        for (int i = 0; i < engine_->getNbIOTensors(); i++) {
            const char* name = engine_->getIOTensorName(i);
            context_->setTensorAddress(name, device_buffers_[i]);
        }

        // 执行推理
        context_->enqueueV3(stream_);

        // 拷贝输出到CPU
        std::vector<float> output(output_size_);
        cudaMemcpyAsync(output.data(), device_buffers_[output_idx_],
                        output_size_ * sizeof(float),
                        cudaMemcpyDeviceToHost, stream_);

        cudaStreamSynchronize(stream_);
        return output;
    }

private:
    void allocateBuffers() {
        int nb_tensors = engine_->getNbIOTensors();
        device_buffers_.resize(nb_tensors);

        for (int i = 0; i < nb_tensors; i++) {
            const char* name = engine_->getIOTensorName(i);
            auto dims = engine_->getTensorShape(name);
            auto dtype = engine_->getTensorDataType(name);

            size_t vol = 1;
            for (int d = 0; d < dims.nbDims; d++) vol *= dims.d[d];

            size_t bytes = vol * sizeof(float); // 简化：假设float
            cudaMalloc(&device_buffers_[i], bytes);

            if (engine_->getTensorIOMode(name) == TensorIOMode::kINPUT) {
                input_idx_ = i;
            } else {
                output_idx_ = i;
                output_size_ = vol;
            }
        }
    }

    TRTUniquePtr<IRuntime> runtime_;
    TRTUniquePtr<ICudaEngine> engine_;
    TRTUniquePtr<IExecutionContext> context_;
    cudaStream_t stream_;
    std::vector<void*> device_buffers_;
    int input_idx_ = 0, output_idx_ = 1;
    size_t output_size_ = 0;
};

// ===== 5. 完整使用流程 =====
int main() {
    Logger logger;

    // 构建Engine（首次）
    std::string engine_path = "model.engine";
    std::vector<char> engine_data;

    std::ifstream check(engine_path);
    if (check.good()) {
        std::cout << "Loading cached engine..." << std::endl;
        engine_data = loadEngine(engine_path);
    } else {
        std::cout << "Building engine from ONNX..." << std::endl;
        engine_data = buildEngine("model.onnx", logger, true);
        saveEngine(engine_data, engine_path);
    }

    // 创建推理器
    TRTInference infer(engine_data, logger);

    // 准备输入 (假设 [1, 3, 224, 224])
    std::vector<float> input(1 * 3 * 224 * 224, 1.0f);

    // 推理
    auto output = infer.infer(input);

    std::cout << "Output size: " << output.size() << std::endl;
    std::cout << "Top-1 class: "
              << std::distance(output.begin(),
                   std::max_element(output.begin(), output.end()))
              << std::endl;

    return 0;
}

#include <NvInfer.h>
#include <vector>
#include <string>
#include <fstream>

class Int8Calibrator : public nvinfer1::IInt8EntropyCalibrator2 {
public:
    Int8Calibrator(int batch_size, int input_h, int input_w,
                   const std::vector<std::string>& image_files,
                   const std::string& cache_file)
        : batch_size_(batch_size),
          input_h_(input_h), input_w_(input_w),
          image_files_(image_files),
          cache_file_(cache_file),
          current_batch_(0) {

        input_count_ = batch_size * 3 * input_h * input_w;
        cudaMalloc(&device_input_, input_count_ * sizeof(float));
    }

    ~Int8Calibrator() {
        cudaFree(device_input_);
    }

    int getBatchSize() const noexcept override {
        return batch_size_;
    }

    bool getBatch(void* bindings[], const char* names[],
                  int nbBindings) noexcept override {
        if (current_batch_ * batch_size_ >= (int)image_files_.size()) {
            return false;  // 校准完成
        }

        // 加载一个batch的图片并预处理
        std::vector<float> host_data(input_count_);
        for (int i = 0; i < batch_size_; i++) {
            int idx = current_batch_ * batch_size_ + i;
            if (idx < (int)image_files_.size()) {
                loadAndPreprocess(image_files_[idx],
                    host_data.data() + i * 3 * input_h_ * input_w_,
                    input_h_, input_w_);
            }
        }

        cudaMemcpy(device_input_, host_data.data(),
                   input_count_ * sizeof(float), cudaMemcpyHostToDevice);
        bindings[0] = device_input_;

        current_batch_++;
        return true;
    }

    // 读取/写入校准缓存
    const void* readCalibrationCache(size_t& length) noexcept override {
        cache_.clear();
        std::ifstream input(cache_file_, std::ios::binary);
        if (input.good()) {
            input.seekg(0, std::ios::end);
            length = input.tellg();
            input.seekg(0, std::ios::beg);
            cache_.resize(length);
            input.read(cache_.data(), length);
            return cache_.data();
        }
        return nullptr;
    }

    void writeCalibrationCache(const void* cache, size_t length) noexcept override {
        std::ofstream output(cache_file_, std::ios::binary);
        output.write(static_cast<const char*>(cache), length);
    }

private:
    void loadAndPreprocess(const std::string& path, float* buffer,
                           int h, int w) {
        // 实际项目中用OpenCV读取并resize、normalize
        // 这里简化为随机数据
        for (int i = 0; i < 3 * h * w; i++) {
            buffer[i] = static_cast<float>(rand()) / RAND_MAX;
        }
    }

    int batch_size_;
    int input_h_, input_w_;
    int input_count_;
    std::vector<std::string> image_files_;
    std::string cache_file_;
    int current_batch_;
    void* device_input_;
    std::vector<char> cache_;
};

// 使用INT8校准构建Engine
std::vector<char> buildInt8Engine(const std::string& onnx_path,
                                  Logger& logger) {
    auto builder = TRTUniquePtr<IBuilder>(createInferBuilder(logger));
    auto network = TRTUniquePtr<INetworkDefinition>(
        builder->createNetworkV2(1U << static_cast<uint32_t>(
            NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));

    auto parser = TRTUniquePtr<nvonnxparser::IParser>(
        nvonnxparser::createParser(*network, logger));
    parser->parseFromFile(onnx_path.c_str(),
        static_cast<int>(ILogger::Severity::kWARNING));

    auto config = TRTUniquePtr<IBuilderConfig>(builder->createBuilderConfig());
    config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1ULL << 30);
    config->setFlag(BuilderFlag::kINT8);

    // 准备校准图片
    std::vector<std::string> calib_images;
    for (int i = 0; i < 500; i++) {
        calib_images.push_back("calib_data/img_" + std::to_string(i) + ".jpg");
    }

    auto calibrator = std::make_unique<Int8Calibrator>(
        8, 224, 224, calib_images, "calibration.cache");
    config->setInt8Calibrator(calibrator.get());

    auto serialized = TRTUniquePtr<IHostMemory>(
        builder->buildSerializedNetwork(*network, *config));

    return std::vector<char>(
        static_cast<char*>(serialized->data()),
        static_cast<char*>(serialized->data()) + serialized->size());
}

// 动态shape允许一个Engine支持多种输入尺寸
std::vector<char> buildDynamicEngine(const std::string& onnx_path,
                                     Logger& logger) {
    auto builder = TRTUniquePtr<IBuilder>(createInferBuilder(logger));
    auto network = TRTUniquePtr<INetworkDefinition>(
        builder->createNetworkV2(1U << static_cast<uint32_t>(
            NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));

    auto parser = TRTUniquePtr<nvonnxparser::IParser>(
        nvonnxparser::createParser(*network, logger));
    parser->parseFromFile(onnx_path.c_str(),
        static_cast<int>(ILogger::Severity::kWARNING));

    auto config = TRTUniquePtr<IBuilderConfig>(builder->createBuilderConfig());
    config->setFlag(BuilderFlag::kFP16);

    // 创建Optimization Profile
    auto profile = builder->createOptimizationProfile();

    // 设置输入的动态范围: [min, opt, max]
    // 假设输入名称为"input", NCHW格式
    profile->setDimensions("input", OptProfileSelector::kMIN,
                           Dims4{1, 3, 224, 224});
    profile->setDimensions("input", OptProfileSelector::kOPT,
                           Dims4{8, 3, 224, 224});    // 优化目标batch
    profile->setDimensions("input", OptProfileSelector::kMAX,
                           Dims4{32, 3, 224, 224});

    config->addOptimizationProfile(profile);

    auto serialized = TRTUniquePtr<IHostMemory>(
        builder->buildSerializedNetwork(*network, *config));

    return std::vector<char>(
        static_cast<char*>(serialized->data()),
        static_cast<char*>(serialized->data()) + serialized->size());
}

// 推理时设置实际shape
void inferDynamic(IExecutionContext* context, int actual_batch) {
    // 设置实际输入shape
    context->setInputShape("input", Dims4{actual_batch, 3, 224, 224});

    // 分配对应大小的buffer并推理
    // ...
}

#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <vector>
#include <string>
#include <cstring>

using namespace nvinfer1;

// 自定义GELU Plugin
class GELUPlugin : public IPluginV2DynamicExt {
public:
    GELUPlugin() = default;

    // ===== 必须实现的接口 =====

    // 返回plugin名称
    const char* getPluginType() const noexcept override {
        return "CustomGELU";
    }

    const char* getPluginVersion() const noexcept override {
        return "1";
    }

    int getNbOutputs() const noexcept override {
        return 1;
    }

    // 输出数据类型
    DataType getOutputDataType(int index, const DataType* inputTypes,
                                int nbInputs) const noexcept override {
        return inputTypes[0];
    }

    // 输出维度
    DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs,
                                  int nbInputs,
                                  IExprBuilder& exprBuilder) noexcept override {
        return inputs[0];  // GELU不改变shape
    }

    // 工作空间需求
    size_t getWorkspaceSize(const PluginTensorDesc* inputs, int nbInputs,
                            const PluginTensorDesc* outputs,
                            int nbOutputs) const noexcept override {
        return 0;
    }

    // 核心: CUDA核函数执行
    int enqueue(const PluginTensorDesc* inputDesc,
                const PluginTensorDesc* outputDesc,
                const void* const* inputs, void* const* outputs,
                void* workspace, cudaStream_t stream) noexcept override {

        int n = 1;
        for (int i = 0; i < inputDesc[0].dims.nbDims; i++) {
            n *= inputDesc[0].dims.d[i];
        }

        // 调用CUDA kernel
        geluKernel(static_cast<const float*>(inputs[0]),
                   static_cast<float*>(outputs[0]),
                   n, stream);
        return 0;
    }

    // 序列化
    size_t getSerializationSize() const noexcept override { return 0; }
    void serialize(void* buffer) const noexcept override {}

    // 克隆
    IPluginV2DynamicExt* clone() const noexcept override {
        return new GELUPlugin();
    }

    // 配置
    bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut,
                                    int nbInputs,
                                    int nbOutputs) noexcept override {
        return inOut[pos].type == DataType::kFLOAT &&
               inOut[pos].format == TensorFormat::kLINEAR;
    }

    void configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs,
                         const DynamicPluginTensorDesc* out,
                         int nbOutputs) noexcept override {}

    int initialize() noexcept override { return 0; }
    void terminate() noexcept override {}
    void destroy() noexcept override { delete this; }
    void setPluginNamespace(const char* ns) noexcept override { namespace_ = ns; }
    const char* getPluginNamespace() const noexcept override { return namespace_.c_str(); }

private:
    std::string namespace_;

    // GELU CUDA kernel声明
    void geluKernel(const float* input, float* output, int n, cudaStream_t stream);
};