嵌入式平台点云处理加速

// 完整TensorRT部署流程 (C++)

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime_api.h>

class PointCloudTensorRT {
public:
    // 1. 构建引擎 (离线执行一次，保存.plan文件)
    bool buildEngine(const std::string& onnxPath, 
                     const std::string& enginePath,
                     bool useFP16 = true,
                     bool useINT8 = false) {
        
        nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger_);
        const auto explicitBatch = 1U << static_cast<uint32_t>(
            nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
        
        nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
        nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger_);
        
        // 解析ONNX
        if (!parser->parseFromFile(onnxPath.c_str(), 
                                   static_cast<int>(nvinfer1::ILogger::Severity::kWARNING))) {
            return false;
        }
        
        // 配置builder
        nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
        config->setMaxWorkspaceSize(1ULL << 30);  // 1GB工作空间
        
        if (useFP16) {
            config->setFlag(nvinfer1::BuilderFlag::kFP16);
        }
        
        if (useINT8) {
            config->setFlag(nvinfer1::BuilderFlag::kINT8);
            // 需要设置INT8校准器
            config->setInt8Calibrator(calibrator_);
        }
        
        // 针对点云模型的动态shape优化
        auto profile = builder->createOptimizationProfile();
        // 点云数量动态: [min, opt, max]
        profile->setDimensions("points", nvinfer1::OptProfileSelector::kMIN, 
                               nvinfer1::Dims2(1000, 4));
        profile->setDimensions("points", nvinfer1::OptProfileSelector::kOPT, 
                               nvinfer1::Dims2(20000, 4));
        profile->setDimensions("points", nvinfer1::OptProfileSelector::kMAX, 
                               nvinfer1::Dims2(40000, 4));
        config->addOptimizationProfile(profile);
        
        // 构建引擎
        nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
        
        // 序列化保存
        nvinfer1::IHostMemory* serialized = engine->serialize();
        std::ofstream file(enginePath, std::ios::binary);
        file.write(reinterpret_cast<const char*>(serialized->data()), serialized->size());
        
        // 清理
        serialized->destroy();
        engine->destroy();
        config->destroy();
        parser->destroy();
        network->destroy();
        builder->destroy();
        
        return true;
    }
    
    // 2. 加载引擎并推理
    bool init(const std::string& enginePath) {
        // 反序列化引擎
        std::ifstream file(enginePath, std::ios::binary);
        std::vector<char> engineData((std::istreambuf_iterator<char>(file)),
                                      std::istreambuf_iterator<char>());
        
        runtime_ = nvinfer1::createInferRuntime(logger_);
        engine_ = runtime_->deserializeCudaEngine(engineData.data(), engineData.size());
        context_ = engine_->createExecutionContext();
        
        // 分配GPU内存
        cudaMalloc(&d_points_, maxPoints_ * 4 * sizeof(float));
        cudaMalloc(&d_voxels_, maxVoxels_ * 10 * 4 * sizeof(float));
        cudaMalloc(&d_features_, maxVoxels_ * 64 * sizeof(float));
        cudaMalloc(&d_output_, maxDetections_ * 7 * sizeof(float));
        
        // 创建CUDA流
        cudaStreamCreate(&stream_);
        
        return true;
    }
    
    // 3. 执行推理 (PointPillars示例)
    std::vector<Box3D> infer(const std::vector<float>& points) {
        // 预处理: 点云体素化 (GPU加速)
        int numPoints = points.size() / 4;
        cudaMemcpyAsync(d_points_, points.data(), 
                        numPoints * 4 * sizeof(float), 
                        cudaMemcpyHostToDevice, stream_);
        
        // 启动自定义CUDA kernel进行体素化
        launchVoxelizationKernel(d_points_, numPoints, 
                                 d_voxels_, d_voxel_coords_, 
                                 numVoxels_, stream_);
        
        // TensorRT推理
        void* bindings[] = {d_voxels_, d_voxel_coords_, d_features_, d_output_};
        context_->enqueueV2(bindings, stream_, nullptr);
        
        // 后处理: NMS (GPU或CPU)
        std::vector<Box3D> detections;
        launchNMSKernel(d_output_, d_scores_, d_boxes_, 
                        numDetections_, nmsThreshold_, stream_);
        
        // 拷贝回CPU
        cudaMemcpyAsync(outputHost_, d_output_, 
                        numDetections_ * 7 * sizeof(float),
                        cudaMemcpyDeviceToHost, stream_);
        cudaStreamSynchronize(stream_);
        
        // 解析输出
        for (int i = 0; i < numDetections_; ++i) {
            if (outputHost_[i * 7 + 6] > scoreThreshold_) {  // 置信度
                Box3D box;
                box.cx = outputHost_[i * 7 + 0];
                box.cy = outputHost_[i * 7 + 1];
                box.cz = outputHost_[i * 7 + 2];
                box.w = outputHost_[i * 7 + 3];
                box.l = outputHost_[i * 7 + 4];
                box.h = outputHost_[i * 7 + 5];
                box.score = outputHost_[i * 7 + 6];
                detections.push_back(box);
            }
        }
        
        return detections;
    }

private:
    nvinfer1::IRuntime* runtime_ = nullptr;
    nvinfer1::ICudaEngine* engine_ = nullptr;
    nvinfer1::IExecutionContext* context_ = nullptr;
    cudaStream_t stream_;
    
    // GPU缓冲区
    void* d_points_ = nullptr;
    void* d_voxels_ = nullptr;
    void* d_voxel_coords_ = nullptr;
    void* d_features_ = nullptr;
    void* d_output_ = nullptr;
    float* outputHost_ = nullptr;
    
    static constexpr int maxPoints_ = 40000;
    static constexpr int maxVoxels_ = 20000;
    static constexpr int maxDetections_ = 500;
    static constexpr float scoreThreshold_ = 0.3f;
    static constexpr float nmsThreshold_ = 0.5f;
    
    int numVoxels_ = 0;
    
    Logger logger_;
};

// Voxelization Plugin (TensorRT插件)

#include <NvInferPlugin.h>

class VoxelizationPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
    // 输出描述
    int getNbOutputs() const override { return 2; }  // voxels, coords
    
    nvinfer1::DimsExprs getOutputDimensions(
        int outputIndex,
        const nvinfer1::DimsExprs* inputs,
        int nbInputs,
        nvinfer1::IExprBuilder& exprBuilder) override {
        
        // inputs[0]: [N, 4] points
        // output 0: [maxVoxels, maxPointsPerVoxel, 4] voxels
        // output 1: [maxVoxels, 4] coords (x,y,z, num_points)
        
        nvinfer1::DimsExprs output;
        if (outputIndex == 0) {
            output.nbDims = 3;
            output.d[0] = exprBuilder.constant(maxVoxels_);
            output.d[1] = exprBuilder.constant(maxPointsPerVoxel_);
            output.d[2] = exprBuilder.constant(4);
        } else {
            output.nbDims = 2;
            output.d[0] = exprBuilder.constant(maxVoxels_);
            output.d[1] = exprBuilder.constant(4);
        }
        return output;
    }
    
    // 核心CUDA kernel配置
    int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
                const nvinfer1::PluginTensorDesc* outputDesc,
                const void* const* inputs,
                void* const* outputs,
                void* workspace,
                cudaStream_t stream) override {
        
        const float* points = static_cast<const float*>(inputs[0]);
        float* voxels = static_cast<float*>(outputs[0]);
        int* coords = static_cast<int*>(outputs[1]);
        
        // 启动体素化kernel
        int numPoints = inputDesc[0].dims.d[0];
        voxelizeKernel<<<blocks, threads, 0, stream>>>(
            points, numPoints,
            voxelSizeX_, voxelSizeY_, voxelSizeZ_,
            gridMinX_, gridMinY_, gridMinZ_,
            gridSizeX_, gridSizeY_, gridSizeZ_,
            maxPointsPerVoxel_, maxVoxels_,
            voxels, coords, numVoxels_
        );
        
        return 0;
    }

private:
    // 体素化参数
    float voxelSizeX_ = 0.16f, voxelSizeY_ = 0.16f, voxelSizeZ_ = 4.0f;
    float gridMinX_ = -75.2f, gridMinY_ = -75.2f, gridMinZ_ = -2.0f;
    int gridSizeX_ = 940, gridSizeY_ = 940, gridSizeZ_ = 1;
    int maxPointsPerVoxel_ = 100;
    int maxVoxels_ = 20000;
};

// 注册插件
REGISTER_TENSORRT_PLUGIN(VoxelizationPluginCreator);

# PyTorch → ONNX → TensorRT INT8 (Python校准示例)
# 实际部署需转换为C++校准器

import torch
import tensorrt as trt
import numpy as np

class PointCloudCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file="calibration.cache"):
        super().__init__()
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.batch_size = 1
        self.current_index = 0
        
        # 分配GPU内存
        self.device_input = cuda.mem_alloc(
            40000 * 4 * 4)  # max_points * 4 coords * float32
        
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        if self.current_index >= len(self.data_loader):
            return None
            
        # 获取点云数据
        points = self.data_loader[self.current_index]
        self.current_index += 1
        
        # 填充固定长度
        points_padded = np.zeros((40000, 4), dtype=np.float32)
        num_points = min(len(points), 40000)
        points_padded[:num_points] = points[:num_points]
        
        # 拷贝到GPU
        cuda.memcpy_htod(self.device_input, points_padded.ravel())
        return [int(self.device_input)]
    
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
        return None
    
    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

# 使用校准器构建INT8引擎
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = PointCloudCalibrator(data_loader)

// 通道剪枝 (针对PointPillars的PFN层)

#include <torch/torch.h>

struct PrunedPFN : torch::nn::Module {
    // 原始: 线性层 64 → 64
    // 剪枝后: 64 → 32 (移除50%通道)
    
    torch::nn::Linear linear_{nullptr};
    torch::nn::BatchNorm1d bn_{nullptr};
    
    PrunedPFN(int in_channels = 64, int out_channels = 32) {
        linear_ = register_module("linear", 
            torch::nn::Linear(in_channels, out_channels));
        bn_ = register_module("bn", 
            torch::nn::BatchNorm1d(out_channels));
    }
    
    torch::Tensor forward(torch::Tensor x) {
        // x: [N, 64]
        x = torch::relu(bn_->forward(linear_->forward(x)));
        // 使用max pooling聚合点特征
        x = torch::max(x, /*dim=*/0).values;
        return x;
    }
};

// 剪枝策略
class ChannelPruner {
public:
    // 基于L1范数的通道重要性评估
    std::vector<int> getImportantChannels(
        const torch::Tensor& weight, 
        float prune_ratio) {
        
        // weight: [out_ch, in_ch]
        auto l1_norm = weight.abs().sum(1);  // 每个输出通道的L1范数
        auto sorted = std::get<1>(l1_norm.sort(/*descending=*/true));
        
        int num_keep = weight.size(0) * (1 - prune_ratio);
        std::vector<int> indices(sorted.data_ptr<int64_t>(), 
                                 sorted.data_ptr<int64_t>() + num_keep);
        return indices;
    }
    
    // 剪枝并微调
    void pruneAndFinetune(
        std::shared_ptr<torch::nn::Module> model,
        float prune_ratio,
        torch::data::DataLoader<>& dataloader) {
        
        // 1. 评估各层重要性
        // 2. 移除低重要性通道
        // 3. 微调恢复精度
        // 4. 导出新的ONNX
    }
};

// 教师-学生模型蒸馏 (大模型→小模型)

class DistillationLoss {
public:
    torch::Tensor compute(
        torch::Tensor student_logits,
        torch::Tensor teacher_logits,
        torch::Tensor labels,
        float temperature = 4.0f,
        float alpha = 0.5f) {
        
        // 软目标损失 (KL散度)
        auto soft_loss = torch::kl_div(
            torch::log_softmax(student_logits / temperature, 1),
            torch::softmax(teacher_logits / temperature, 1),
            /*reduction=*/torch::kBatchMean
        ) * (temperature * temperature);
        
        // 硬目标损失 (交叉熵)
        auto hard_loss = torch::cross_entropy_loss(student_logits, labels);
        
        return alpha * soft_loss + (1 - alpha) * hard_loss;
    }
};

// 训练循环
void trainWithDistillation(
    StudentModel& student,
    TeacherModel& teacher,
    DataLoader& loader) {
    
    teacher.eval();  // 教师固定
    
    for (auto& batch : loader) {
        auto points = batch.points;
        auto labels = batch.labels;
        
        // 前向传播
        auto student_out = student.forward(points);
        auto teacher_out = teacher.forward(points);
        
        // 计算蒸馏损失
        auto loss = distillation_loss.compute(
            student_out, teacher_out, labels);
        
        // 反向传播
        optimizer.zero_grad();
        loss.backward();
        optimizer.step();
    }
}

// RK3588点云检测适配方案
// 由于NPU限制，采用CPU预处理 + NPU推理 + CPU后处理

#include "rknn_api.h"

class RK3588PointDetector {
public:
    bool init(const std::string& rknnPath) {
        // 加载RKNN模型
        FILE* fp = fopen(rknnPath.c_str(), "rb");
        fseek(fp, 0, SEEK_END);
        int model_size = ftell(fp);
        rewind(fp);
        
        model_data_.resize(model_size);
        fread(model_data_.data(), 1, model_size, fp);
        fclose(fp);
        
        // 初始化RKNN
        int ret = rknn_init(&ctx_, model_data_.data(), model_size, 0, NULL);
        if (ret < 0) return false;
        
        // 查询输入输出
        rknn_input_output_num io_num;
        rknn_query(ctx_, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
        
        // 设置输入参数 (BEV伪图像)
        input_attrs_.resize(io_num.n_input);
        for (int i = 0; i < io_num.n_input; i++) {
            input_attrs_[i].index = i;
            rknn_query(ctx_, RKNN_QUERY_INPUT_ATTR, 
                      &input_attrs_[i], sizeof(rknn_tensor_attr));
        }
        
        return true;
    }
    
    std::vector<Box3D> detect(const std::vector<float>& points) {
        // Step 1: CPU端体素化 (RK3588无高效GPU)
        auto bev_image = cpuVoxelize(points);  // [H, W, C]
        
        // Step 2: NPU推理
        rknn_input inputs[1];
        inputs[0].index = 0;
        inputs[0].type = RKNN_TENSOR_UINT8;  // INT8量化输入
        inputs[0].size = bev_image.size();
        inputs[0].fmt = RKNN_TENSOR_NHWC;
        inputs[0].buf = bev_image.data();
        
        rknn_inputs_set(ctx_, 1, inputs);
        rknn_run(ctx_, NULL);
        
        // 获取输出 (检测头)
        rknn_output outputs[2];  // 分类 + 回归
        outputs[0].want_float = true;
        outputs[1].want_float = true;
        rknn_outputs_get(ctx_, 2, outputs, NULL);
        
        // Step 3: CPU后处理 (NMS)
        auto detections = cpuNMS(outputs[0].buf, outputs[1].buf);
        
        rknn_outputs_release(ctx_, 2, outputs);
        
        return detections;
    }

private:
    // CPU优化体素化 (NEON加速)
    std::vector<uint8_t> cpuVoxelize(const std::vector<float>& points) {
        const int H = 400, W = 400, C = 64;  // BEV尺寸
        std::vector<uint8_t> bev(H * W * C, 0);
        
        // 使用ARM NEON加速
        #ifdef __ARM_NEON
        // NEON实现的点云投影
        #else
        // 标量实现
        for (const auto& pt : points) {
            int x = (pt.x - min_x) / voxel_size;
            int y = (pt.y - min_y) / voxel_size;
            if (x >= 0 && x < W && y >= 0 && y < H) {
                // 填充特征通道
                int idx = (y * W + x) * C;
                // 简化的特征编码
            }
        }
        #endif
        
        return bev;
    }
    
    // CPU NMS优化
    std::vector<Box3D> cpuNMS(void* cls_output, void* reg_output) {
        // 解码检测框
        // 使用多线程加速NMS
        std::vector<Box3D> candidates;
        // ... 解码逻辑
        
        // 按分数排序
        std::sort(candidates.begin(), candidates.end(),
                  [](const Box3D& a, const Box3D& b) {
                      return a.score > b.score;
                  });
        
        std::vector<Box3D> result;
        std::vector<bool> suppressed(candidates.size(), false);
        
        for (size_t i = 0; i < candidates.size(); ++i) {
            if (suppressed[i]) continue;
            result.push_back(candidates[i]);
            
            #pragma omp parallel for
            for (size_t j = i + 1; j < candidates.size(); ++j) {
                if (suppressed[j]) continue;
                float iou = computeIoU(candidates[i], candidates[j]);
                if (iou > nms_threshold_) {
                    suppressed[j] = true;
                }
            }
        }
        
        return result;
    }
    
    rknn_context ctx_;
    std::vector<char> model_data_;
    std::vector<rknn_tensor_attr> input_attrs_;
    std::vector<rknn_tensor_attr> output_attrs_;
    
    float min_x_ = -50.0f, min_y_ = -50.0f;
    float voxel_size_ = 0.25f;
    float nms_threshold_ = 0.5f;
};

// 优化1: 共享内存体素化
__global__ void voxelizeKernelShared(
    const float* points,
    int num_points,
    float voxel_size_x, float voxel_size_y, float voxel_size_z,
    int grid_size_x, int grid_size_y, int grid_size_z,
    float min_x, float min_y, float min_z,
    float* voxels,
    int* voxel_coords,
    int* num_voxels,
    int max_points_per_voxel,
    int max_voxels) {
    
    __shared__ int voxel_map[4096];  // 共享内存哈希表
    __shared__ int voxel_count[4096];
    
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    int local_tid = threadIdx.x;
    
    // 初始化共享内存
    if (local_tid < 4096) {
        voxel_map[local_tid] = -1;
        voxel_count[local_tid] = 0;
    }
    __syncthreads();
    
    // 处理点
    if (tid < num_points) {
        const float* pt = points + tid * 4;
        
        // 计算体素索引
        int vx = floor((pt[0] - min_x) / voxel_size_x);
        int vy = floor((pt[1] - min_y) / voxel_size_y);
        int vz = floor((pt[2] - min_z) / voxel_size_z);
        
        if (vx >= 0 && vx < grid_size_x && 
            vy >= 0 && vy < grid_size_y && 
            vz >= 0 && vz < grid_size_z) {
            
            // 哈希映射
            int hash = (vx * 73856093 ^ vy * 19349663 ^ vz * 83492791) % 4096;
            
            // 原子操作插入
            int voxel_idx = atomicCAS(&voxel_map[hash], -1, *num_voxels);
            if (voxel_idx == -1) {
                voxel_idx = atomicAdd(num_voxels, 1);
                voxel_map[hash] = voxel_idx;
                // 记录坐标
                voxel_coords[voxel_idx * 4 + 0] = vx;
                voxel_coords[voxel_idx * 4 + 1] = vy;
                voxel_coords[voxel_idx * 4 + 2] = vz;
            }
            
            // 填充点
            int point_idx = atomicAdd(&voxel_count[hash], 1);
            if (point_idx < max_points_per_voxel) {
                int offset = (voxel_idx * max_points_per_voxel + point_idx) * 4;
                voxels[offset + 0] = pt[0];
                voxels[offset + 1] = pt[1];
                voxels[offset + 2] = pt[2];
                voxels[offset + 3] = pt[3];  // intensity
            }
        }
    }
}

// 优化2: 旋转NMS (CUDA)
__global__ void rotateNMSKernel(
    const float* boxes,      // [N, 7] (x,y,z,w,l,h,theta)
    const float* scores,     // [N]
    int num_boxes,
    float nms_threshold,
    bool* keep) {
    
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= num_boxes) return;
    
    if (!keep[tid]) return;
    
    // 预计算当前box的角点
    float cx = boxes[tid * 7 + 0];
    float cy = boxes[tid * 7 + 1];
    float w = boxes[tid * 7 + 3];
    float l = boxes[tid * 7 + 4];
    float theta = boxes[tid * 7 + 6];
    
    float cos_t = cos(theta), sin_t = sin(theta);
    
    // 与其他box比较
    for (int i = tid + 1; i < num_boxes; ++i) {
        if (!keep[i]) continue;
        
        // 快速距离检查
        float dx = boxes[i * 7 + 0] - cx;
        float dy = boxes[i * 7 + 1] - cy;
        if (dx * dx + dy * dy > (w + l) * (w + l)) continue;
        
        // 精确旋转IoU计算
        float iou = computeRotateIoU(boxes + tid * 7, boxes + i * 7);
        if (iou > nms_threshold) {
            keep[i] = false;
        }
    }
}

// 内存池管理 (避免频繁malloc/free)

class CudaMemoryPool {
public:
    static CudaMemoryPool& getInstance() {
        static CudaMemoryPool instance;
        return instance;
    }
    
    void* allocate(size_t size) {
        std::lock_guard<std::mutex> lock(mutex_);
        
        // 查找合适块
        for (auto it = free_blocks_.begin(); it != free_blocks_.end(); ++it) {
            if (it->size >= size) {
                void* ptr = it->ptr;
                allocated_blocks_[ptr] = *it;
                free_blocks_.erase(it);
                return ptr;
            }
        }
        
        // 分配新块
        void* ptr;
        cudaMalloc(&ptr, size);
        allocated_blocks_[ptr] = {ptr, size};
        return ptr;
    }
    
    void deallocate(void* ptr) {
        std::lock_guard<std::mutex> lock(mutex_);
        
        auto it = allocated_blocks_.find(ptr);
        if (it != allocated_blocks_.end()) {
            free_blocks_.push_back(it->second);
            allocated_blocks_.erase(it);
        }
    }
    
    // 预分配常用缓冲区
    void preallocate() {
        // 点云缓冲区
        allocate(40000 * 4 * sizeof(float));
        // 体素缓冲区
        allocate(20000 * 100 * 4 * sizeof(float));
        // 特征缓冲区
        allocate(20000 * 64 * sizeof(float));
        // 输出缓冲区
        allocate(500 * 7 * sizeof(float));
    }

private:
    struct Block {
        void* ptr;
        size_t size;
    };
    
    std::unordered_map<void*, Block> allocated_blocks_;
    std::vector<Block> free_blocks_;
    std::mutex mutex_;
};

// 使用RAII管理
class CudaBuffer {
public:
    CudaBuffer(size_t size) {
        ptr_ = CudaMemoryPool::getInstance().allocate(size);
        size_ = size;
    }
    
    ~CudaBuffer() {
        CudaMemoryPool::getInstance().deallocate(ptr_);
    }
    
    void* data() { return ptr_; }
    
private:
    void* ptr_;
    size_t size_;
};

// 生产环境部署检查清单

□ 模型优化
  □ ONNX导出 (opset 11+)
  □ TensorRT引擎构建 (FP16/INT8)
  □ 动态shape支持 (点云数量可变)
  □ 层融合验证

□ 内存管理
  □ 预分配所有缓冲区
  □ 内存池管理
  □ 零拷贝数据传输 (pinned memory)

□ 并发处理
  □ 多流并行 (预处理/推理/后处理)
  □ 双缓冲机制 (ping-pong buffer)
  □ 线程池管理

□ 稳定性
  □ 异常处理 (CUDA错误恢复)
  □ 温度监控 (降频保护)
  □ 模型热更新

□ 性能监控
  □ 各阶段耗时统计
  □ GPU利用率监控
  □ 内存使用追踪