aclnnRotaryPositionEmbedding

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	×
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	×
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能：执行单路旋转位置编码计算。
• 计算公式：

（1）half模式（mode等于0）： $$ x1 = x[..., : x.shape[-1] // 2] $$

$$x2 = x[..., x.shape[-1] // 2 :]$$ $$x\_rotate = torch.cat((-x2, x1), dim=-1)$$ $$y = x * cos + x\_rotate * sin$$

（2）interleave模式（mode等于1）： $$ x1 = x[..., ::2].view(-1, 1) $$

$$x2 = x[..., 1::2].view(-1, 1)$$ $$x\_rotate = torch.cat((-x2, x1), dim=-1).view(x.shape[0], x.shape[1], x.shape[2], x.shape[3])$$ $$y = x * cos + x\_rotate * sin$$

（3）quarter模式（mode等于2）： $$ x1 = x[..., : x.shape[-1] // 4] $$

$$x2 = x[..., x.shape[-1] // 4 : x.shape[-1] // 2]$$ $$x3 = x[..., x.shape[-1] // 2 : x.shape[-1] // 4 * 3]$$ $$x4 = x[..., x.shape[-1] // 4 * 3 :]$$ $$x\_rotate = torch.cat((-x2, x1, -x4, x3), dim=-1)$$ $$y = x * cos + x\_rotate * sin$$

（4）interleave-half模式（mode等于3），该模式会先将奇数位的输入抽取到前半部分，将偶数位的输入抽取到后半部分，再进行half处理： $$ x1 = x[..., ::2] $$

$$x2 = x[..., 1::2]$$ $$x\_part1 = torch.cat((x1, x2), dim=-1)$$ $$x\_part2 = torch.cat((-x2, x1), dim=-1)$$

$$ y = x_part1 * cos + x_part2 * sin $$

函数原型

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

• aclnnStatus aclnnRotaryPositionEmbeddingGetWorkspaceSize(const aclTensor *x, const aclTensor *cos, const aclTensor *sin, int64_t mode, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnRotaryPositionEmbedding(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnRotaryPositionEmbeddingGetWorkspaceSize

• 参数说明：

  • x（aclTensor*，计算输入）：表示待执行旋转位置编码的张量，公式中的x，Device侧的aclTensor。支持undefined，undefined支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32。
  • cos（aclTensor*，计算输入）：表示参与计算的位置编码张量，公式中的cos，Device侧的aclTensor。支持undefined，undefined支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32，且必须与x的数据类型一致。
  • sin（aclTensor*，计算输入）：表示参与计算的位置编码张量，公式中的sin，Device侧的aclTensor，支持undefined，undefined支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32，且必须与x的数据类型一致。
  • mode（int64_t，计算输入）：表示公式中的旋转模式，数据类型为int64，支持的枚举值如下：
    • 0：表示旋转模式为half模式。
    • 1：表示旋转模式为interleave模式。
    • 2：表示旋转模式为quarter模式。
    • 3：表示旋转模式为interleave-half模式。
  • out（aclTensor*，计算输出）：表示旋转位置编码的计算结果张量，公式中的y，Device侧的aclTensor。 不支持undefined，undefined支持ND，维度为4维，数据类型支持BFLOAT16、FLOAT16、FLOAT32，且必须与x的数据类型一致。
  • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
  • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnRotaryPositionEmbedding

• 参数说明：

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

• 返回值：

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

约束说明

用(B, S, N, D)表示四维输入x的shape，在该表示下，各参数的shape约束可以描述如下：

• 输入张量x、cos、sin及输出张量y的D维度大小必须相同，且小于等于1024。对于half、interleave和interleave-half模式，D必须能被2整除，对于quarter模式，D必须能被4整除。
• 输入张量x和输出张量y的shape必须完全相同。
• 输入张量cos和sin的shape必须完全相同，且必须满足下列条件之一：
  • 前三维大小都为1，即shape为(1, 1, 1, D)。
  • 前三维的大小和x前三维的大小完全相等，即shape为(B, S, N, D)。
  • 前三维中，第二维和第三维中的一个大小为1，剩余的维度及第一维大小与x的对应维度相等，即shape为(B, 1, N, D)或(B, S, 1, D)。
  • 前三维中，两个维度大小为1，剩余的一个维度大小与x的对应维度相等，即shape为(1, 1, N, D)，(1, S, 1, D)或(B, 1, 1, D)。 当x为空tensor时，输出也为空tensor，且不受上述shape约束限制。

调用示例

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