aclnnSinkhorn

产品支持情况

• [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]。
• [object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]。

功能说明

• 算子功能：

  计算Sinkhorn距离，可以用于MoE模型中的专家路由。

• 计算公式：

  $$p=Sinkhorn(cost, tol)$$

  ---

  输入： cost(R, C): 2维成本矩阵 tol: 误差

  初始化:

  $$cost = exp(cost) \\ d0 = ones(R) \\ d1 = ones(C) \\ eps = 0.00000001 \\ error = 1e9 \\ d1\_old= d1 \\$$

  重复执行:

  $$d0 = \frac{1}{R * (sum(d1 * cost, 1) + eps)} \\ d1 = \frac{1}{C * (sum(d0.unsqueeze(1) * cost, 0) + eps)} \\ error = mean(abs(d1\_old - d1)) \\ d1\_old = d1$$

  直至:

  $$error <= tol$$

  输出:

  $$p = d1 * cost * d0.unsqueeze(1)$$

函数原型

每个算子分为undefined，必须先调用“aclnnSinkhornGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnSinkhorn”接口执行计算。

• aclnnStatus aclnnSinkhornGetWorkspaceSize(const aclTensor *cost, const aclScalar *tol, aclTensor *p, uint64_t *workspaceSize, aclOpExecutor** executor)
• aclnnStatus aclnnSinkhorn(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnSinkhornGetWorkspaceSize

• 参数说明：

  • cost（aclTensor*，计算输入）：表示成本张量，公式中的cost，Device侧的aclTensor。数据类型支持BFLOAT16、FLOAT16、FLOAT。undefined支持ND，输入为二维矩阵且列数不超过4096。支持undefined。
  • tol (aclScalar*, 入参) ：表示计算Sinkhorn的误差，数据类型支持FLOAT。如果传入空指针，则tol取0.0001。
  • p（aclTensor*，计算输出）：表示最优传输张量，公式中的p，Device侧的aclTensor。数据类型支持BFLOAT16、FLOAT16、FLOAT。undefined支持ND。shape维度为2。不支持undefined。数据类型和shape与入参cost的数据类型和shape一致。
  • workspaceSize（uint64_t*，出参）：返回用户需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

  aclnnStatus：返回状态码，具体参见undefined。

  [object Object]

aclnnSinkhorn

• 参数说明：

  • workspace（void*，入参）：在Device侧申请的workspace内存地址。
  • workspaceSize（uint64_t，入参）：在Device侧申请的workspace大小，由第一段接口aclnnSinkhornGetWorkspaceSize获取。
  • executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
  • stream（aclrtStream，入参）：指定执行任务的Stream。

• 返回值：

  aclnnStatus：返回状态码，具体参见undefined。

约束说明

无。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考undefined。