CINN 基础算子体系

1. 常见张量算子（12 个）

以下算子可替换为硬件厂商库函数调用：

• batch_norm_grad：批正则化操作的反向计算逻辑。

• batch_norm_inference：用于前向推理过程中的批正则化计算逻辑。

• batch_norm_training：用于训练过程中的批正则化前向计算逻辑。

• cholesky：对输入的对称正定矩阵计算 Cholesky 分解。

• conv：对给定输入张量进行卷积计算，包括二维卷积和三维卷积。

• dot：对两个输入张量进行矩阵乘法操作。

• fft：对输入复数张量进行快速傅里叶变换。

• rfft：对输入实数张量进行快速傅里叶变换，输出为复数张量。

• pool：对输入张量进行池化操作。

• sort：对输入张量沿给定轴进行排序，输出排序好的数据，其形状和输入相同。

• triangular_solve：计算具有唯一解的线性方程组，其中系数方阵 A 为上（下）三角系数矩阵。若系数方阵 A 不可逆，则线性方程不可解。

• argmax：沿给定轴计算输入张量在该轴上的最大元素的索引值张量。

2. 激活函数（5 个）

• sin：对输入张量逐元素求正弦值。计算公式为$output = \sin(input)$。

• cos：对输入张量逐元素求余弦值。计算公式为$output = \cos(input)$。

• tanh：对输入张量逐元素求双曲正切值。计算公式为$output = \frac{e^x - e^{-x}}{e^x + e^{-x}}$。

• atan2：对输入张量 x、y 逐元素求原点(0,0)至点(x,y)的方位角，即与横坐标轴 x 轴的夹角，并将各个位置的输出元素保存到返回结果中。 计算公式如下： $$ output = \begin{cases} \arctan(\frac{x}{y}) & y > 0 \ \arctan(\frac{x}{y}) + \pi & x \ge 0, y < 0 \ \arctan(\frac{x}{y}) - \pi & x < 0, y < 0 \ +\frac{\pi}{2} & x > 0, y = 0 \ - \frac{\pi}{2}& x < 0, y = 0 \ undefined & x = 0, y = 0 \end{cases} $$

• clamp：对输入张量逐元素进行裁剪，使得输出元素限制在指定的[min_value, max_value]范围内，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \min(\max(input, min_value), max_value)$。

3. elementwise 一元算子（22 个）

• abs：对输入张量逐元素求绝对值，并将各个位置的输出元素保存到返回结果中。计算公式为$output = |input|$。

• ceil：对输入张量逐元素向上取整，并将各个位置的输出元素保存到返回结果中。 计算公式为$output = \lceil input \rceil$。

• floor：对输入张量逐元素向下取整，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \lfloor input \rfloor$。

• round：对输入张量逐元素四舍五入到最接近的整数数值，并将各个位置的输出元素保存到返回结果中。

• clz：clz 是 count leading zeros 的缩写，主要用于对输入张量逐元素统计前导零的数目。即在二进制表达下，第一个 1 比特位前的 0 比特位的数目（均包含符号位），并将各个位置的统计数目保存到返回结果中。

• population_count：在二进制表达下，对输入张量逐元素统计 1 比特位的数目（包含符号位），并将各个位置的统计数目保存到返回结果中。

• exp：对输入张量逐元素求以自然数 e 为底的指数值，并将各个位置的输出元素保存到返回结果中。计算公式为$output = e^{input}$。

• expm1：对输入张量逐元素求以自然数 e 为底的指数值并减 1，并将各个位置的输出元素保存到返回结果中。计算公式为$output = e^{input} - 1$。

• is_finite：对输入张量逐元素判断是否 Finite，换而言之，是否既非+/-INF，亦非+/-NaN，并将各个位置的判断结果保存到返回结果中。计算公式为$output = (input \ne +\inf) && (input \ne -\inf) && (input \ne \text{nan})$。

• log：对输入张量逐元素求自然对数，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \ln(input)$。

• log1p：对输入张量逐元素求加一的自然对数，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \ln(input + 1)$。

• logistic：对输入张量逐元素求标准逻辑斯谛函数（又称 sigmoid 函数）值，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \frac{1}{1 + e^{-input}}$。

• sign：对输入张量逐元素进行正负判断，并将各个位置的正负判断值保存到返回结果中。计算公式如下： $$ output = \begin{cases} -1 & input \lt 0 \ -0 & input = -0 \ NaN & input = NaN \ +0 & input = +0 \ 1 & input \gt 0 \end{cases} $$

• bitwise_not：对输入张量逐元素按位取反，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \sim input$。

• logical_not：对输入张量逐元素进行逻辑非运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = ! input$。

• negate：对输入张量逐元素取相反数，并将各个位置的输出元素保存到返回结果中。计算公式为$output = - input$。

• reciprocal：对输入张量逐元素取倒数，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \frac{1}{input}$。

• imag：对输入张量逐元素取复数的虚部数值，并将各个位置的输出元素保存到返回结果中。

• real：对输入张量逐元素取复数的实部数值，并将各个位置的输出元素保存到返回结果中。

• rsqrt：对输入张量逐元素取平方根的倒数，并将各个位置的输出元素保存到返回结果中。计算公式为$output = \frac{1}{\sqrt{input}}$。

• sqrt：对输入张量逐元素取平方根，并将各个位置的输出元素保存到返回结果中。计算公式为$output =\sqrt{input}$。

• cbrt：对输入张量逐元素取立方根，并将各个位置的输出元素保存到返回结果中。计算公式为$output =\sqrt[3]{input}$。

4. elementwise 二元算子（20 个）

• add：逐元素相加算子，对两个输入张量逐元素相加，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x + y$。

• sub：逐元素相减算子，对两个输入张量逐元素相减，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x - y$。

• mul：逐元素相乘算子，对两个输入张量逐元素相乘，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x \odot y$。

• div：逐元素相除算子，对两个输入张量逐元素相除，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x / y$。

• complex：给定实部输入张量和虚部输入张量，逐元素转换为复数，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x + yi$。

• mod：逐元素取模算子，对两个输入张量逐元素取模，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x % y$。

• remainder：逐元素取余算子。不同于取模在取余数过程中尽可能让商向负无穷靠近，取余尽可能让商向零靠近。

• arithmetic_right_shift：逐元素算术右移算子，逐元素对第一个输入张量 x 右移 y 个位置，空出的高位用最高位（符号位）填补，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x \gg y$。

• logical_right_shift：逐元素逻辑右移算子，逐元素对第一个输入张量 x 右移 y 个位置，空出的高位用 0 填补，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x \ggg y$。

• logical_left_shift：逐元素逻辑左移算子，逐元素对第一个输入张量 x 左移 y 个位置，空出的低位用 0 填补，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x \ll y$。

• bitwise_and：逐元素按位与算子，对两个输入张量逐元素进行按位与运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x & y$。

• bitwise_or：逐元素按位或算子，对两个输入张量逐元素进行按位或运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x | y$。

• bitwise_xor：逐元素按位异或算子，对两个输入张量逐元素进行按位异或运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x \verb!^! y$。

• logical_and：逐元素逻辑与算子，对两个输入张量逐元素进行逻辑与运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x && y$。

• logical_or：逐元素逻辑或算子，对两个输入张量逐元素进行逻辑或运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = x || y$。

• logical_xor：逐元素逻辑异或算子，对两个输入张量逐元素进行逻辑异或运算，并将各个位置的输出元素保存到返回结果中。计算公式为$output = (x || y) && !(x && y)$。

• maximum：逐元素对比两个输入张量，并且把各个位置更大的元素保存到返回结果中。计算公式为$output = \max(x,y)$。

• minimum：逐元素对比两个输入张量，并且把各个位置更小的元素保存到返回结果中。计算公式为$output = \min(x,y)$。

• pow：逐元素计算输入张量 x 的 y 次幂的值，并且把各个位置的输出元素保存到返回结果中。计算公式为$output = x^y$。

• table：表查找操作，主要用于实现嵌入层(Embedding Layer)。

5. elementwise 三元算子（1 个）

• select：逐元素判断condition输入张量中的值，并且根据各个位置的判断结果把true_value输入张量或false_value输入张量中的元素保存到返回结果中。计算公式如下： $$ output = condition ? true_value : false_value $$

6. 比较算子（3 个）

• equal：逐元素判断两个输入 Tensor 是否相等，并且把各个位置的判断结果保存到返回结果中。计算公式为$output = (x == y)$。

• greater：逐元素判断第一个输入 Tensor 是否大于第二个输入 Tensor，并且把各个位置的判断结果保存到返回结果中。计算公式为$output = (x > y)$。

• greater_equal：逐元素判断第一个输入 Tensor 是否大于等于第二个输入 Tensor，并且把各个位置的判断结果保存到返回结果中。计算公式为$output = (x \ge y)$。

7. reduction 算子（6 个）

• reduce_all：沿给定的坐标轴对输入张量元素进行逻辑与运算。

• reduce_any：沿给定的坐标轴对输入张量元素进行逻辑或运算。

• reduce_max：沿给定的坐标轴对输入张量元素取最大值。

• reduce_min：沿给定的坐标轴对输入张量元素取最小值。

• reduce_product：沿给定的坐标轴对输入张量元素进行规约求乘积操作。

• reduce_sum：沿给定的坐标轴对输入张量元素进行规约求和操作。

8. broadcast 算子（2 个）

• broadcast：将输入张量广播到指定的形状。

• broadcast_in_dim：将输入张量广播到指定的形状。相比 broadcast，broadcast_in_dim 通过给定轴映射列表，支持了广播维度不连续情况下的广播。

9. datalayout 变换算子（6 个）

• concat：沿给定的坐标轴将所有输入张量合并为一个张量。

• pad：根据padding_mode对输入张量进行扩充，扩充区域填充值value。

• reshape：在保持输入数据不变的情况下，改变张量的形状。

• reverse：沿给定的坐标轴反转输入张量中的数据，数据类型和形状均保持不变。

• slice：根据给定的起点（包含）、终点（不包含）和步长值，对输入张量进行切片，并将切片后的张量保存到返回结果中。

• transpose：根据给定重排轴列表，对输入张量进行数据重排。

10. 图像操作算子（1 个）

• resize：将输入图片通过指定插值方法调整为指定大小，输入图片应该是 4-D 张量，且形状为[N, C, H, W]，注意调整仅适用于 H、W 对应维度。

11. 集合[通信]算子（11 个）

• gather：根据指定的偏移量数组将多个操作数片段进行拼接。

• scatter：根据输入的updates张量和indices数组，对输入张量进行相应切片的更新。

• select_and_scatter：select和scatter的复合操作。即在给定窗口大小中从输入张量中选择目标元素，然后将数值scatter到输出张量中。

• send：发送指定操作数的值到另一个共享相同通道的recv指令。

• recv：从另一个共享相同通道的send指令中接收指定 shape 的数据。

• all_gather：将来自不同分布式节点的数据进行concat操作。

• all_reduce：将来自不同分布式节点的数据进行reduce操作。

• reduce_scatter：在执行完all_reduce操作后执行scatter操作。

• all_to_all：将数据从所有分布式节点发送到所有分布式节点。

• after_all：接收可变数量的令牌，生成单一令牌。主要用于控制特定操作的执行在一组操作完成之后进行。

• collective_permute：在不同分布式节点之间发送和接收数据。

12. 数据节点算子（3 个）

• const：将输入的常量数据保存到输出张量中。

• identity：拷贝输入张量中的所有元素到输出张量中，类似copy操作。

• iota：创建一个一维数组，其值从 0 开始并以 1 递增。

13. 类型转换算子（2 个）

• cast：将输入张量中的所有元素转换为指定类型并保存到输出张量中，亦可称为convert。

• bitcast_convert：在不改变底层存储的情况下，强制转换数据类型。若转换前后数据类型的字节大小不相同，则形状会改变。比如一个 shape 为[10]的 float32 类型数据被强制转换为 float16 类型后，其 shape 应为[10, 2]。

14. 量化操作算子（3 个）

• quantize：量化操作，即将浮点数量化为定点数，会生成scale（比例系数）和zero_point（零点偏移值）。

• dequantize：反量化操作，即根据scale和zero_point的值将定点数反量化为浮点数。

• rescale：使用给定的multiplier（系数）和shift（偏移）将量化值类型转换为其他定点数类型。如将int8_t量化值转为int16_t类型。

15. 随机数生成算子（3 个）

• rng_bit_generator：使用指定的随机数生成算法生成随机数，用于填充给定 shape 的张量。

• rng_normal：生成符合正态分布的随机数。

• rng_uniform：生成符合均匀分布的随机数。

16. 控制流算子（3 个）

• if：if条件算子。

• case：switch-case条件算子。

• while：while循环算子。

17. 动态 shape 算子（8 个）

• dynamic_conv：具有动态 shape 语义的卷积操作。

• dynamic_gather：具有动态 shape 语义的gather操作。

• dynamic_iota：具有动态 shape 语义的iota操作。

• dynamic_pad：具有动态 shape 语义的padding操作。

• dynamic_reshape：具有动态 shape 语义的reshape操作。

• dynamic_slice：具有动态 shape 语义的切片操作。

• dynamic_update_slice：具有动态 shape 语义的update_slice操作，即先进行切片操作在根据update输入张量对切片进行更新。

• dynamic_broadcast_in_dim：具有动态 shape 语义的broadcast_in_dim操作。

动态 shape 语义主要是将对应静态 shape 算子中编译时确定的属性输入替换为运行时可变的张量输入。

18. 爱因斯坦求和约定（2 个）

• einsum：二元爱因斯坦求和约定算子。

• unary_einsum：一元爱因斯坦求和约定算子。

19. 自定义算子（1 个）

• custom_call：custom_call主要用于调用设备相关的外部函数。