AscendFHE 是面向同态加密计算的 Ascend NPU 硬件加速底层算子库，当前公开支持 Atlas A2 / Ascend 910B，SOC 配置为 ascend910b。仓库提供 FHE 常用模运算、RNS 基转换、内积累加、自同构重排和 NTT 等算子的 Ascend C kernel、host tiling 逻辑、 算子构建片段和 torch_npu 调用说明。

安装并注册 AscendFHE 算子包后，这些算子通过 torch.ops.ascend_npu_fhe.* 暴露给 PyTorch / torch_npu；未安装或未注册算子包的环境只能进行源码审查、环境检查和 Notebook 入口验证。

现在可以审查源码、准备 CANN 环境、使用 Notebook 做环境或注册检查；不能直接从根目录构建完整 扩展，也不能在未安装注册 AscendFHE 算子包的环境中直接调用 torch.ops.ascend_npu_fhe.*。

这条任务流用于最快完成“拿到源码 -> 确认内容 -> 准备环境 -> 打开交互入口”。真实算子执行 依赖目标 NPU 环境中已经安装并注册 AscendFHE 算子包。

1. 获取源码并查看算子：

git clone https://github.com/BJTUpupil/AscendFHE.git
cd AscendFHE

2. 确认当前仓库内容：

ls
rg --files kernels
rg --files doc notebooks

应能看到 kernels/、doc/ 和 notebooks/。其中 kernels/ 存放算子源码，doc/ 存放公开文档， notebooks/ 存放交互检查入口。

3. 准备 CANN 8.5.0 / Atlas A2 / Ascend 910B 环境：

conda create -n ascend-fhe python=3.10 -y
conda activate ascend-fhe
conda install ascend::cann-toolkit==8.5.0 -y
conda install ascend::cann-910b-ops==8.5.0 -y
source "${CONDA_PREFIX}/Ascend/ascend-toolkit/set_env.sh"
export ASCEND_HOME_PATH="${CONDA_PREFIX}/Ascend/ascend-toolkit"

驱动、固件、PyTorch 和 torch_npu 需要按目标 Atlas A2 / Ascend 910B 机器的实际版本安装。 完整安装路径见 CANN 8.5.0 环境安装方法。

4. 验证 NPU 运行环境：

echo "${ASCEND_HOME_PATH}"
npu-smi info
python -c "import torch, torch_npu; print(torch.npu.is_available())"

期望结果是 ASCEND_HOME_PATH 指向实际 CANN Toolkit 目录，npu-smi info 能识别设备， Python 命令输出 True。

5. 打开交互 Notebook：

jupyter notebook notebooks/ascend_npu_fhe_quickstart.ipynb

Notebook 提交版本默认不执行安装、环境检查或算子调用。进入目标 NPU 环境后，可先启用环境检查：

RUN_INSTALL = False
RUN_ENV_CHECK = True
RUN_OP_TESTS = False

6. 仅在完成算子包安装和注册后启用算子测试：

RUN_OP_TESTS = True

RUN_OP_TESTS=True 只适用于已经完成 AscendFHE 算子包安装和 torch.ops.ascend_npu_fhe.* 注册的环境。当前仓库最快可完成的是环境检查和接口入口验证；真正的 NPU 算子执行依赖目标环境中的 算子包安装与注册。

公开文档位于 doc/：

• 文档组织方式
• CANN 8.5.0 环境安装方法
• 第三方引用与可执行交互说明
• Notebook 一键安装、执行与测试脚本

docs/ 是本地忽略目录，不作为公开文档入口。

每个算子当前提供 Ascend C kernel、host tiling 逻辑和算子级 README。安装并注册 AscendFHE 算子包后，表中的调用形式即为公开 torch_npu API。

AscendFHE/
|-- README.md
|-- LICENSE
|-- doc/
|-- notebooks/
`-- kernels/
    |-- add_mod/
    |-- sub_mod/
    |-- mul_mod_shoup/
    |-- montgomery_mul/
    |-- bconv/
    |-- ip/
    |-- auto_op/
    `-- ntt/

每个算子目录采用固定结构：

kernels/<operator>/
|-- README.md
|-- operator.cmake
|-- op_host/
`-- op_kernel/

目录含义：

安装并注册 AscendFHE 算子包后，公开 API 如下：

torch.ops.ascend_npu_fhe.add_mod(a, b, q) -> y
torch.ops.ascend_npu_fhe.sub_mod(a, b, q) -> y
torch.ops.ascend_npu_fhe.mul_mod_shoup(a, W, Wprime, q) -> y
torch.ops.ascend_npu_fhe.montgomery_mul(x, y, q_mu, q) -> z
torch.ops.ascend_npu_fhe.bconv(
    x, lhs_dense, num_src_bases, num_dst_bases, poly_degree, fixed_shift,
    src_bases, dst_bases, q_hat_inv_mod_q, q_hat_inv_mod_q_shoup,
    src_bases_inv_fixed, error_comp_coefs, error_comp_coefs_shoup,
    merge_base_shoup
) -> y
torch.ops.ascend_npu_fhe.ip(x, key, key_shoup, num_rns, d, poly_degree, q_array) -> y
torch.ops.ascend_npu_fhe.auto_op(x, byte_offsets_32, sign_mask, num_rns, poly_degree, k, q_array) -> dst
torch.ops.ascend_npu_fhe.ntt(a_grid, w1_soa, w2_all, qs_tensor, q_mu_tensor, w3_mm2_mat_soa) -> stage5_out

统一接入链路如下：

Python API
  -> torch.ops.ascend_npu_fhe.<op_name>
  -> C++ extension wrapper
  -> 输入 dtype/shape/device 校验
  -> 输出 Tensor 创建
  -> host tiling 或等价 tiling 生成
  -> ACL / torch_npu 自定义算子 launch
  -> Ascend C kernel
  -> NPU 输出 Tensor

C++ wrapper 的统一职责：

1. 校验输入 Tensor 位于 NPU，且 dtype、shape 和属性满足算子要求。
2. 按算子 shape 推导规则创建输出 Tensor。
3. 调用现有 op_host/ 中的 tiling 逻辑，或实现等价 tiling 生成函数。
4. 调用对应 Ascend C kernel 入口。
5. 返回输出 Tensor 给 PyTorch。

建议新增的 torch_npu 适配层目录：

torch_npu_ext/
|-- CMakeLists.txt
|-- register_ops.cpp
|-- ops/
|   |-- add_mod.cpp
|   |-- sub_mod.cpp
|   |-- mul_mod_shoup.cpp
|   |-- montgomery_mul.cpp
|   |-- bconv.cpp
|   |-- ip.cpp
|   |-- auto_op.cpp
|   `-- ntt.cpp
`-- python/
    `-- ascend_npu_fhe/
        |-- __init__.py
        `-- ops.py

当前仓库不包含上述适配层源码；本 README 和各算子 README 描述的是接入实现时必须 遵守的接口和数据流。

本仓库公开支持的 CANN 版本固定为 CANN Community Edition 8.5.0。源码中的 AICore().AddConfig("ascend910b") 声明了硬件目标，公开支持路径固定为 Atlas A2 / Ascend 910B，并使用 910b ops 算子包。

本仓库不声明兼容 CANN Community Edition 8.5.0 以外的版本。安装方法见 CANN 8.5.0 环境安装方法。

实现 torch_npu 适配层后，每个算子按以下方式验证：

• AddMod：校验 y = (a + b) mod q。
• SubMod：校验 y = (a - b) mod q。
• MulModShoup：使用 W、Wprime、q 校验 Shoup 模乘结果。
• MontgomeryMul：使用 q_mu、q 校验 Montgomery 模乘结果。
• BConv：使用小规模源基和目标基校验基转换结果。
• IP：使用小规模 num_rns、d、poly_degree 校验内积累加结果。
• AutoOp：校验重排、符号处理和模归一化结果。
• NTT：使用固定旋转因子表和模数校验五阶段 NTT 输出。

验证记录必须包含：

算子名称
torch_npu API
输入形状
输入数据类型
属性值
参考实现输出
NPU 输出
校验和
运行命令
运行日志

添加新算子时按以下步骤修改仓库：

1. 新建 kernels/<name>/op_host/，添加 shape 推导、数据类型推导和 tiling 参数生成。
2. 新建 kernels/<name>/op_kernel/，添加 Ascend C 设备侧 kernel。
3. 新建 kernels/<name>/operator.cmake，声明构建所需源码和头文件。
4. 新建 kernels/<name>/README.md，说明独立功能和 torch_npu 接入方式。
5. 在后续 torch_npu 适配层中新增 torch_npu_ext/ops/<name>.cpp。
6. 更新根 README 的独立功能卡片和 torch_npu API。

本地检查命令：

rg -n "torch_npu|torch.ops.ascend_npu_fhe" README.md kernels/*/README.md
rg -n "__global__ __aicore__ void" kernels
rg --files kernels

仓库根目录的 LICENSE 为 GPL-2.0。修改和分发源码时必须保留源文件中的许可证声明。

公开发布前，在此处补充论文标题、作者、会议或期刊名称、年份和 BibTeX 条目。 实验记录中必须包含仓库提交哈希。