全同态加密在AI/机器学习中的应用：CipherFlow LattiSense 框架和 LattiAI 平台

# 1. 设置 FHE 参数
from frontend.custom_task import *
 
param = Param.create_bfv_default_param(n=16384)
set_fhe_param(param)
 
# 2. 定义输入密文节点
level = 2
x = BfvCiphertextNode('x', level)
y = BfvCiphertextNode('y', level)
 
# 3. 定义计算任务（乘法+重线性化）
z = mult_relin(x, y, 'z')
 
# 4. 生成执行指令
process_custom_task(
    input_args=[Argument('x', x), Argument('y', y)],
    output_args=[Argument('z', z)],
    output_instruction_path='examples/quick_start',
)

struct DatumNode {
    NodeIndex index;
    std::string id;
    std::vector<ComputeNode*> successors;
    bool is_output;
    DataType datum_type;  // TYPE_CIPHERTEXT, TYPE_PLAINTEXT
    int32_t level;        // 密文层级
    int32_t degree;       // 密文度数 (2 或 3)
    bool is_ntt;          // NTT 形式标志
    bool is_mform;        // 乘法形式标志
};
 
struct ComputeNode {
    NodeIndex index;
    std::string id;
    std::vector<DatumNode*> input_nodes;
    std::vector<DatumNode*> output_nodes;
    ExecutorFunc executor;  // 统一执行器函数
    OperationType op_type;  // ADD, SUB, MULTIPLY, ROTATE, etc.
};
 
struct MegaAG {
    std::unordered_map<NodeIndex, DatumNode> data;
    std::unordered_map<NodeIndex, ComputeNode> computes;
    std::vector<NodeIndex> inputs;
    std::vector<NodeIndex> outputs;
    Processor processor;  // CPU, GPU, FPGA
};

using ExecutorFunc = std::function<void(
    ExecutionContext& ctx,                                    // 硬件上下文
    const std::unordered_map<NodeIndex, std::any>& inputs,   // 输入数据
    std::any& output,                                        // 输出数据
    const ComputeNode& self                                  // 节点元信息
)>;

ExecutorFunc cpu_add_executor = [](ExecutionContext& ctx,
                                   const auto& inputs,
                                   std::any& output,
                                   const ComputeNode& self) {
    // 1. 提取 CPU 上下文
    auto& ckks_ctx = ctx.get_context<CkksContext>();
 
    // 2. 提取输入密文
    const auto& ct1 = std::any_cast<const CkksCiphertext&>(
        inputs.at(self.input_nodes[0]->index));
    const auto& ct2 = std::any_cast<const CkksCiphertext&>(
        inputs.at(self.input_nodes[1]->index));
 
    // 3. 执行 CPU 加法
    output = ckks_ctx.add(ct1, ct2);
};

ExecutorFunc gpu_add_executor = [](ExecutionContext& ctx,
                                   const auto& inputs,
                                   std::any& output,
                                   const ComputeNode& self) {
    // 1. 提取 GPU 算子上下文
    auto& he_op = ctx.get_context<HEOperator<Scheme::CKKS>>();
 
    // 2. 提取输入密文（GPU 设备内存）
    const auto& ct1 = std::any_cast<const Ciphertext<Scheme::CKKS>&>(
        inputs.at(self.input_nodes[0]->index));
    const auto& ct2 = std::any_cast<const Ciphertext<Scheme::CKKS>&>(
        inputs.at(self.input_nodes[1]->index));
 
    // 3. 提取执行选项（CUDA 流）
    auto* opts = ctx.get_other_arg<ExecutionOptions>(0);
 
    // 4. 执行 GPU 加法
    Ciphertext<Scheme::CKKS> result;
    he_op.add(ct1, ct2, result, opts ? *opts : ExecutionOptions());
 
    output = result;
};

# 在HEonGPU目录，配置构建
cmake -S . -B build -DCMAKE_CUDA_ARCHITECTURES=86 -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_INSTALL_PREFIX=$(pwd)/install
# 编译
cmake --build ./build/ -j$(nproc)
# 安装到HEonGPU/install下
cmake --install .

# 在主目录下，启用GPU加速配置构建
cmake -S . -B build -DCMAKE_INSTALL_PREFIX=$(pwd)/install \
  -DLATTISENSE_BUILD_EXAMPLES=ON \
  -DLATTISENSE_BUILD_TESTS=ON \
  -DLATTISENSE_ENABLE_GPU=ON -DLATTISENSE_CUDA_ARCH=86
# 编译
cmake --build ./build/ -j$(nproc)
# 运行build/bin/example和build/bin/benchmark下的样例和性能测试
# 部分测试需要运行相应的Python脚本构造测试数据

// 输入：H×H×C 的特征图，g = H/Ĥ
// 输出：g² 个 Ĥ×Ĥ×C 的子特征图
 
for (int c = 0; c < C; c++) {
    for (int i = 0; i < g; i++) {
        for (int j = 0; j < g; j++) {
            // 提取交错子通道
            for (int y = 0; y < Ĥ; y++) {
                for (int x = 0; x < Ĥ; x++) {
                    int src_y = i * Ĥ + y;
                    int src_x = j * Ĥ + x;
                    sub_channel[i*g+j][c][y*Ĥ+x] = input[c][src_y*H+src_x];
                }
            }
        }
    }
}

# CIFAR-10 ResNet20
python train.py \
    --epochs 200 \
    --batch-size 128 \
    --lr 0.1

python train.py \
    --poly_model_convert \
    --pretrained train_baseline.pth \
    --degree 4 \
    --upper-bound 3.0 \
    --poly-module RangeNormPoly2d

class PolyReLU2D {
public:
    // 使用 Chebyshev 多项式近似 ReLU
    Feature2DEncrypted run(CkksContext& ctx,
                          const Feature2DEncrypted& input) {
        // ReLU(x) ≈ Σ c_k * T_k(x)，其中 T_k 是 Chebyshev 多项式
        // 使用 Clenshaw 算法数值稳定求值
        return evaluate_clenshaw(ctx, input, relu_coeffs_);
    }
 
private:
    std::vector<double> relu_coeffs_;  // 预计算的多项式系数
};

class Simple_Polyrelu(nn.Module):
    """Polynomial activation approximating ReLU or SiLU via Hermite expansion.
 
    Exported as a single ``nn_tools::Simple_Polyrelu`` custom op in ONNX.
 
    Args:
        scale_before: Input scaling factor.
        scale_after:  Output scaling factor.
        degree:       Polynomial degree (2 or 4).
        activation:   Target activation ('relu' or 'silu').
    """
 
    # Hermite coefficients for ReLU approximation
    _RELU_COEFF = {
        2: (0.39894228, 0.50000000, 0.28209479 / np.sqrt(2), 0.0, 0.0),
        4: (0.39894228, 0.50000000, 0.28209479 / np.sqrt(2), 0.0, -0.08143375 / np.sqrt(24)),
    }
 
    # Coefficients for SiLU approximation
    _SILU_COEFF = {
        2: (0.20662096, 0.50000000, 0.24808519 / np.sqrt(2), 0.0, 0.0),
        4: (0.20662096, 0.50000000, 0.24808519 / np.sqrt(2), 0.0, -0.03780501 / np.sqrt(24)),
    }

        if degree == 2:
            return a0 + (a1 + a2 * x) * x - a2
        elif degree == 4:
            return a0 + a1 * x + a2 * (x**2 - 1) + a4 * (x**4 - 6 * x**2 + 3)
        else:
            raise ValueError(f'Unsupported degree: {degree}')

python run_compile.py \
    --input trained_poly.onnx \
    --output ./runs/cifar10/ \
    --style=multiplexed

// 客户端：加密输入
InferenceClient client("./task/client");
auto encrypted_input = client.encrypt("img.csv");
 
// 服务端：推理（选择硬件后端）
InferenceServer server("./task/server", use_gpu);
auto encrypted_output = server.evaluate(encrypted_input);
 
// 客户端：解密结果
auto result = client.decrypt(encrypted_output);

Feature2DEncrypted run(CkksContext& ctx,
                      const Feature2DEncrypted& input) {
    for (int i = 0; i < kernel_h; i++) {
        for (int j = 0; j < kernel_w; j++) {
            // 1. 旋转输入密文
            int step = (i - center_y) * width + (j - center_x);
            CkksCiphertext rotated = ctx.rotate(input, step);
 
            // 2. 元素级乘法（密文 × 明文权重）
            CkksCiphertext prod = ctx.mult_plain(rotated, weight[i][j]);
 
            // 3. 累加求和
            output = ctx.add(output, prod);
        }
    }
    return output;
}

Feature2DEncrypted run(CkksContext& ctx,
                      const Feature2DEncrypted& input) {
    // 每个输入通道独立执行基本卷积
    for (int c = 0; c < n_channels; c++) {
        auto channel_result = basic_convolution(input.get_channel(c), kernel[c]);
 
        // 通过旋转聚合到指定位置
        channel_result = ctx.rotate(channel_result, c * stride);
 
        // 重新打包为 Multiplexed 格式
        output = multiplex_pack(output, channel_result);
    }
    return output;
}

Feature2DEncrypted run(CkksContext& ctx,
                      const Feature2DEncrypted& input) {
    // 1. 将每个通道分解为 g² 个交错子通道
    std::vector<Feature2DEncrypted> sub_channels;
    for (int i = 0; i < g; i++) {
        for (int j = 0; j < g; j++) {
            sub_channels.push_back(extract_sub_channel(input, i, j, g));
        }
    }
 
    // 2. 对每个子通道执行步长为 g 的卷积
    std::vector<Feature2DEncrypted> partial_results;
    for (int idx = 0; idx < g * g; idx++) {
        auto result = strided_convolution(sub_channels[idx], kernel, g);
        partial_results.push_back(result);
    }
 
    // 3. 累加所有子通道结果
    Feature2DEncrypted output = partial_results[0];
    for (int idx = 1; idx < g * g; idx++) {
        output = ctx.add(output, partial_results[idx]);
    }
 
    return output;
}

// 权重矩阵按对角线方向打包为明文向量
// 主对角线：[w₀₀, w₁₁, w₂₂, w₃₃]
// 次对角线1：[w₀₁, w₁₂, w₂₃, w₃₀]
// ...
 
Feature2DEncrypted run(CkksContext& ctx,
                      const Feature2DEncrypted& input,
                      const WeightMatrix& weight) {
    Feature2DEncrypted output = ctx.encrypt_zero();
 
    for (int d = 0; d < n_diagonals; d++) {
        // 1. 旋转输入密文
        CkksCiphertext rotated = ctx.rotate(input, d);
 
        // 2. 密文 × 明文对角线
        CkksPlaintext diag_pt = encode_diagonal(weight, d);
        CkksCiphertext prod = ctx.mult_plain(rotated, diag_pt);
 
        // 3. 累加到输出
        output = ctx.add(output, prod);
    }
 
    return output;
}