Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

KLDiv 算子设计描述

一、背景信息 （必填）

PyTorch 的 kl_div / KLDivLoss 是训练侧常用损失函数。CANN 存量 v1 原型算子 KLDiv（REG_OP，reduction 必选属性、无 log_target）已大量存在于旧 GE 图中。新算子体系迁移到 ops-math 仓后，存量旧图（IR 图迁移、OM 重编译）要求 KLDiv 这个 op type 仍可精确匹配并正确路由执行。本需求即在 ops-math 仓实现 KLDiv 壳路由算子：算子原型对齐 v1，计算内核纯转发 KLDivV2 的 log_target=False DAG，恒按 log_target=False 语义计算 KL 散度损失（Kullback-Leibler Divergence）。

二、价值/作用 （必填）

• 旧图零改造迁移：存量 GE 图中的 KLDiv 节点无需改写即可在新版本 CANN 上运行，保障存量业务连续性。
• 应用场景：知识蒸馏（teacher-student 分布对齐）、变分推断/VAE、策略优化等一切使用 KL 散度损失的训练场景；语义兼容 PyTorch kl_div（log_target=False 分支）。
• 低成本维护：计算内核完全复用 KLDivV2 成熟的 Reduce DAG 与 tiling 模板（atvoss/reduce 调度框架），壳层无独立计算逻辑，无重复开发。

三、设计方案 （必填）

壳路由设计总体思路："原型对齐 v1 + 计算纯转发 V2"。

旧GE图 KLDiv节点 ──op type精确匹配──> 壳原型(ops::KLDiv, 对齐v1)
      │                                   │
      │                          Infershape/InferDtype(壳自有)
      │                                   │
      └── atc图编译 ──> tiling(壳自有, 复用ReduceOp模板)
                                          │
                                  kernel(kl_div_apt) 纯转发
                                          │
                          KLDivV2 log_target=False DAG (Sum/Mean两个变体)

3.1 使能方式（涉及哪些框架：如Aclnn直调、Pytorch训练等）

3.2 总体设计

3.2.1 算子支持的数据类型

• FLOAT16 / BFLOAT16 输入在 kernel 内部提升为 FLOAT32 计算，结果按 round-to-nearest-even（RINT）回落原 dtype。

3.2.2 host侧设计

• 原型注册（op_host/kl_div_def.cpp）：

  • 双输入 x/target（REQUIRED、ND 格式）、输出 y；reduction 为 REQUIRED String 属性（对齐 v1，不声明 log_target，旧图携带的该属性按未知属性忽略）；
  • 开启动态编译（DynamicCompileStaticFlag）、动态 rank（DynamicRankSupportFlag）、动态 shape（DynamicShapeSupportFlag）；
  • opFile 指向自有 kl_div_apt；AICore 配置注册 ascend950。

• Infershape（op_host/kl_div_infershape.cpp）：

  • reduction 白名单校验（与 tiling 侧 STR_2_INT 报错同语义）；
  • 静态 shape 严格校验 x 与 target 同 shape（动态 shape 不做相等性比较）；
  • 推导规则：none → y 与 x 同 shape；mean/sum/batchmean → y 为标量（SetDimNum(0)）。

• InferDtype（op_graph/kl_div_graph_infer.cpp）：y 与 x 同 dtype，校验 x/target 同 dtype。

• Tiling（op_host/arch35/kl_div_tiling_arch35.cpp，编译信息 KLDivCompileInfo 复用 ReduceOpCompileInfo）：

  读reduction → GetInputParam(none: axes.clear()逐元素; 其余: 全轴归约)
    → CheckDtype(x/target同dtype) → CheckShape(x/target同shape, none路径防越界读)
    → DoTiling按(reduction, dtype size)分派V2 DAG模板:
         sum/none  → Tiling4ReduceOp<KLDivDagSumLogFalse<float|half, float>>
         mean/batchmean → Tiling4ReduceOp<KLDivDagMeanLogFalse<float|half, float>>
         (SIZE4→float, SIZE2→half; 其余dtype报错 → 等效拒绝double)
    → SetEmptyValue(空tensor短路值) → GEN_REDUCE_TILING_KEY(key, reduction)  // 四参数, 无logTarget维
  

  • SetEmptyValue：空 tensor 短路输出值——sum→0、mean→NaN、batchmean→dim0 为 0 时 NaN 否则 0、none→哨兵值 -0.1；batchmean 且 dim0≠0 时装载 meanVar = 1/dim0。
  • TilingParse（KLDivParse）预取平台信息（AIV 核数、UB 大小、cacheline、ubBlockSize、vRegSize），0 值校验。

• TilingKey（op_kernel/kl_div_tiling_key.h）：REDUCE_TPL_KEY + reduction(0~3) 两维模板参数；选择子覆盖 empty（空 tensor 短路）/ A（全连续 pattern）/ AR normal（外连续内归约）/ AR group（分组归约混合核型）。

• Binary config（op_host/config/ascend950/kl_div_binary.json）：3 dtype × 4 reduction 共 12 组预编译二进制，shape 为 -2 任意，format_match_mode 为 FormatAgnostic。

3.2.3 kernel侧设计

op_kernel/kl_div_apt.cpp 为纯转发层，无独立计算逻辑：模板参数 <REDUCE_TPL_PARAM, int32_t reduction>，TilingData 结构复用 KLDivV2TilingData，由 ReduceSch 统一调度：

kernel kl_div(x, target, y, workspace, tiling):
  if reduction == mean/batchmean:
      op = ReduceSch<KLDivV2::KLDivDagMeanLogFalse<DTYPE_X, float>>   // meanVar装载
      op.SetVar<float,0>(tilingData.reduceTiling.meanVar)
  else (sum/none):
      op = ReduceSch<KLDivV2::KLDivDagSumLogFalse<DTYPE_X, float>>
  op.Init(&pipe, x, target, y, workspace)
  op.Process(tilingData.emptyValue)      // 空tensor短路输出emptyValue

DAG 计算流程（恒 log_target=False，提升至 float32）：

x, target ─CopyIn→ Cast(f32)
mask = Compare(target≠0.0)
t'   = Select(mask, target, 1.0)        // 防log(0); target<0时保留target
loss = target * (Log(t') - x)           // target=0→0; target<0→log负数=NaN
     → ReduceSum
     → [Muls(meanVar)]                  // 仅mean/batchmean: mean或1/dim0
     → Cast(原dtype, RINT) → CopyOut

3.3 支持硬件

Ascend950

本仓实现：CMake SUPPORT_COMPUTE_UNIT "ascend950"，tiling 走 arch35 目录，binary config 仅 ascend950 目录。

3.4 算子约束限制

• x 与 target 必须同 shape、同 dtype；不支持广播（"none" 路径双输入逐元素 CopyIn，shape 不同会越界读，tiling/infershape 双重拦截）。
• reduction 为必选属性，取值仅支持 none/mean/sum/batchmean，无默认值，构图必须显式设置。
• 原型纸面声明支持 DT_DOUBLE，但 tiling 阶段 dtype 白名单拒绝，实际不支持 double。
• 恒按 log_target=False 语义计算；旧图中携带的 log_target 属性会被作为未知属性忽略（v1 原型即无此属性）。
• 值域约束：target=0 的位置输出 0；target<0 的位置输出 NaN（tiling/infershape 不校验，调用方需自行保证输入语义合法）。
• FLOAT16/BFLOAT16 输入在 kernel 内部提升 FLOAT32 计算，结果按 round-to-nearest-even 回落原 dtype。
• 输出 shape：reduction="none" 时 y 与 x 同 shape；mean/sum/batchmean 时 y 为标量。
• 本算子为壳路由算子，无 aclnn 接口：旧图迁移走 atc 图编译路径；PyTorch 适配请使用 kl_div_v2 目录的 aclnnKlDiv 接口。
• 支持 1D ~ 8D 输入；空 tensor 短路：sum→0、mean→NaN、batchmean→dim0 为 0 时 NaN 否则 0。

💡 备注（选填）

• KLDiv 为壳路由算子，后续不再增强；如有新计算需求（如 log_target=True），请使用 KLDivV2。
• 图模式调用示例：examples/test_geir_kl_div.cpp（通过 op::KLDiv 构图），编译验证流程参见 docs/zh/invocation/quick_op_invocation.md。