aclnnSinkhorn

产品支持情况

[object Object]undefined

功能说明

• 算子功能：

  计算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)$$

函数原型

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

• [object Object]
• [object Object]

aclnnSinkhornGetWorkspaceSize

• 参数说明：

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

• 返回值：

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

  [object Object]

aclnnSinkhorn

• 参数说明：

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

• 返回值：

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

约束说明

无。

调用示例

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