afe.ir.operation_functions ========================== .. py:module:: afe.ir.operation_functions Classes ------- .. autoapisummary:: afe.ir.operation_functions.RunMode Functions --------- .. autoapisummary:: afe.ir.operation_functions.placeholder afe.ir.operation_functions.constant afe.ir.operation_functions.float_convolution afe.ir.operation_functions.quantized_convolution afe.ir.operation_functions.prepare_weight_tensor afe.ir.operation_functions.conv_tensorflow afe.ir.operation_functions.float_add afe.ir.operation_functions.ml_kernels_add_sub afe.ir.operation_functions.relu afe.ir.operation_functions.clip afe.ir.operation_functions.prelu afe.ir.operation_functions.elu afe.ir.operation_functions.leaky_relu afe.ir.operation_functions.maxpool afe.ir.operation_functions.avgpool afe.ir.operation_functions.variance afe.ir.operation_functions.broadcast_to afe.ir.operation_functions.multiply_tensorflow afe.ir.operation_functions.ml_kernels_multiply afe.ir.operation_functions.pad afe.ir.operation_functions.mean afe.ir.operation_functions.squeeze afe.ir.operation_functions.argmax afe.ir.operation_functions.softmax afe.ir.operation_functions.lrn afe.ir.operation_functions.concatenate afe.ir.operation_functions.transpose afe.ir.operation_functions.depth_to_space afe.ir.operation_functions.reshape afe.ir.operation_functions.expand_dims afe.ir.operation_functions.batch_flatten afe.ir.operation_functions.min_op afe.ir.operation_functions.max_op afe.ir.operation_functions.sum_op afe.ir.operation_functions.prod afe.ir.operation_functions.tuple_get_item afe.ir.operation_functions.exp afe.ir.operation_functions.sqrt afe.ir.operation_functions.sigmoid afe.ir.operation_functions.swish afe.ir.operation_functions.quick_gelu afe.ir.operation_functions.hard_sigmoid afe.ir.operation_functions.hard_swish afe.ir.operation_functions.log afe.ir.operation_functions.softplus afe.ir.operation_functions.erf afe.ir.operation_functions.gelu afe.ir.operation_functions.log2 afe.ir.operation_functions.log10 afe.ir.operation_functions.subtract_tensorflow afe.ir.operation_functions.power afe.ir.operation_functions.divide afe.ir.operation_functions.reciprocal afe.ir.operation_functions.maximum afe.ir.operation_functions.minimum afe.ir.operation_functions.full afe.ir.operation_functions.tile afe.ir.operation_functions.split afe.ir.operation_functions.take afe.ir.operation_functions.expand_strided_slice_indices_to_shape_length afe.ir.operation_functions.get_strided_slice_out_shape afe.ir.operation_functions.strided_slice afe.ir.operation_functions.rsqrt afe.ir.operation_functions.tanh afe.ir.operation_functions.image_resize2d afe.ir.operation_functions.upsample afe.ir.operation_functions.gridsample afe.ir.operation_functions.layer_norm afe.ir.operation_functions.rms_norm afe.ir.operation_functions.instance_norm afe.ir.operation_functions.layout_transform afe.ir.operation_functions.calculate_tessellated_tensor_shape afe.ir.operation_functions.tessellation afe.ir.operation_functions.detessellation afe.ir.operation_functions.get_channel_aligned_shape afe.ir.operation_functions.get_mla_padded_2d_shape afe.ir.operation_functions.reshape_to_mla_padded_2d_shape afe.ir.operation_functions.reshape_from_mla_padded_2d_shape afe.ir.operation_functions.pack afe.ir.operation_functions.unpack afe.ir.operation_functions.normalization afe.ir.operation_functions.ev_quantize afe.ir.operation_functions.ev_dequantize afe.ir.operation_functions.ev_resize afe.ir.operation_functions.chroma_upsample afe.ir.operation_functions.yuv_rgb_conversion afe.ir.operation_functions.bgr_rgb_conversion afe.ir.operation_functions.ev_sigmoid afe.ir.operation_functions.nms_maxpool afe.ir.operation_functions.cast afe.ir.operation_functions.qnn_quantize afe.ir.operation_functions.qnn_dequantize afe.ir.operation_functions.do_requantize afe.ir.operation_functions.qnn_add afe.ir.operation_functions.qnn_mul afe.ir.operation_functions.external afe.ir.operation_functions.init_custom_op afe.ir.operation_functions.execute_custom_op afe.ir.operation_functions.batch_matmul Module Contents --------------- .. py:class:: RunMode(*args, **kwds) Supported run modes. MLA_MODE : use an implementation that exactly matches execution on the MLA. FAST_MODE : use a fast execution implementation that may not match MLA. JAX : use JAX implementation that matches MLA. .. py:attribute:: MLA_MODE :value: 1 .. py:attribute:: FAST_MODE :value: 2 .. py:attribute:: JAX :value: 3 .. py:method:: is_fast_mode() .. py:method:: is_jax() -> bool .. py:function:: placeholder(data: numpy.ndarray) -> numpy.ndarray .. py:function:: constant(data: numpy.ndarray) -> numpy.ndarray .. py:function:: float_convolution(attrs: afe.ir.attributes.ConvAddActivationAttrs, data: numpy.ndarray, mode: RunMode) -> numpy.ndarray Execute a floating-point convolution using an algorithm from ml_kernels. :param attrs: Attributes of the convolution operator :param data: Input activation data in NHWC layout :param mode: Mode of execution :return: Convolved result .. py:function:: quantized_convolution(attrs: afe.ir.attributes.ConvQuantAttrs, data: numpy.ndarray, mode: RunMode) -> numpy.ndarray Execute a quantized convolution using an algorithm from ml_kernels. :param attrs: Attributes of the convolution operator :param data: Input activation data in NHWC layout :param mode: Mode of execution :return: Convolved result .. py:function:: prepare_weight_tensor(weight: numpy.ndarray, do_transpose: bool) -> numpy.ndarray Helper function used to adapt weight tensor for use in ml_kernels convolution. :param weight: Weight tensor. :param do_transpose: Flag indicating whether weight needs to be transposed. Currently, Numpy and JAX ml_kernels convolution implementation observe weights using different layout. :returns: Modified weights tensor in order to adapt its shape to ml_kernels convolution implementations. .. py:function:: conv_tensorflow(attrs: afe.ir.attributes.ConvAttrs, data: numpy.ndarray, weight: numpy.ndarray) -> numpy.ndarray Compute a floating-point convolution by calling Tensorflow's convolution operator. This function may not exactly match MLA behavior. :param attrs: Attributes of the convolution :param data: Input tensor :param weight: Weight tensor Returns: Convolved tensor .. py:function:: float_add(lhs: numpy.ndarray, rhs: numpy.ndarray, axis: Optional[int] = None) -> numpy.ndarray .. py:function:: ml_kernels_add_sub(attrs: afe.ir.attributes.AddActivationAttrs | afe.ir.attributes.AddQuantAttrs | afe.ir.attributes.SubtractAttrs | afe.ir.attributes.SubtractQuantAttrs, lhs: numpy.ndarray, rhs: numpy.ndarray, use_jax: bool = False) -> numpy.ndarray Execute ml_kernels ideal_add_sub function for AddActivation and Subtract operators. Supports int8, int16 and bfloat16 data types. :param attrs: SiMaIR attributes for AddActivationOp or SubtractOp. :param lhs: The first operand of AddActivationOp or SubtractOp. :param rhs: The second operand of AddActivationOp or SubtractOp. :param use_jax: Flag indicating whether to use JAX implementation. :returns: The result of executing the AddActivationOp or SubtractOp. .. py:function:: relu(data: numpy.ndarray, zp: int = 0) -> numpy.ndarray .. py:function:: clip(attrs: afe.ir.attributes.ClipAttrs | afe.ir.attributes.ClipQuantAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: prelu(data: numpy.ndarray, alpha: Union[numpy.ndarray, float, int], axis: Optional[int] = None, zp: int = 0) -> numpy.ndarray .. py:function:: elu(data: numpy.ndarray) -> numpy.ndarray .. py:function:: leaky_relu(data: numpy.ndarray, alpha: Union[float, int]) -> numpy.ndarray .. py:function:: maxpool(attrs: afe.ir.attributes.PoolAttrs, data: numpy.ndarray, pad_value: Union[float, int], requantization: ml_kernels.requantization.BaseRequantization | None = None, mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray .. py:function:: avgpool(attrs: afe.ir.attributes.PoolAttrs, data: numpy.ndarray, pad_value: float | int, output_type: numpy.dtype, requant: ml_kernels.requantization.BaseRequantization | None = None, mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray .. py:function:: variance(data: numpy.ndarray, mean: numpy.ndarray, requant: ml_kernels.requantization.BaseRequantization | None = None, requant_var: ml_kernels.requantization.BaseRequantization | None = None) .. py:function:: broadcast_to(attrs: Union[afe.ir.attributes.BroadcastToAttrs, afe.ir.attributes.BroadcastToQuantAttrs], data: numpy.ndarray) .. py:function:: multiply_tensorflow(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray Floating-point multiplication. .. py:function:: ml_kernels_multiply(attrs: afe.ir.attributes.MultiplyAttrs | afe.ir.attributes.MultiplyQuantAttrs, lhs: numpy.ndarray, rhs: numpy.ndarray, mode: RunMode) -> numpy.ndarray Quantized multiplication. .. py:function:: pad(attrs: afe.ir.attributes.PadAttrs, data: numpy.ndarray, pad_value: numpy.ndarray) -> numpy.ndarray .. py:function:: mean(attrs: afe.ir.attributes.ReduceAttrs, data: numpy.ndarray, quantized=False) -> numpy.ndarray When in quantized set to True, using avg_pool2d to do mean along axis = * (1) * (2) * (1, 2) Parameters ---------- :param attrs: ReduceAttrs. Attributes needed to execute the mean operation. :param data: np.ndarray. Input data to the mean operation. :param quantized: bool. Default is False. Set to True if the mean operation is executed in a quantization domain. Return ------ :return: np.ndarray. Result of the mean operation. .. py:function:: squeeze(attrs: afe.ir.attributes.SqueezeAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: argmax(attrs: afe.ir.attributes.ArgMaxAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: softmax(attrs: Union[afe.ir.attributes.SoftmaxAttrs, afe.ir.attributes.SoftmaxQuantAttrs], data: numpy.ndarray, mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray .. py:function:: lrn(attrs: Union[afe.ir.attributes.LRNAttrs, afe.ir.attributes.LRNQuantAttrs], data: numpy.ndarray) -> numpy.ndarray .. py:function:: concatenate(attrs: afe.ir.attributes.ConcatenateAttrs | afe.ir.attributes.ConcatQuantAttrs, data_list: list[numpy.ndarray], run_mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray Execute concatenate operation. Used for both non-quantized and quantized inputs. :param attrs: Concatenate operation attributes. :param data_list: List of input tensors to be concatenated. :param run_mode: Execution mode. :returns: The result of concatenate operation. .. py:function:: transpose(attrs: afe.ir.attributes.TransposeAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: depth_to_space(attrs: afe.ir.attributes.DepthToSpaceAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: reshape(attrs: afe.ir.attributes.ReshapeAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: expand_dims(attrs: afe.ir.attributes.ExpandDimsAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: batch_flatten(data: numpy.ndarray) -> numpy.ndarray Flattens all the dimensions except for the batch dimension .. py:function:: min_op(attrs: afe.ir.attributes.ExtmAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: max_op(attrs: afe.ir.attributes.ExtmAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: sum_op(attrs: afe.ir.attributes.ReduceAttrs | afe.ir.attributes.ReduceQuantAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: prod(attrs: afe.ir.attributes.ProdAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: tuple_get_item(attrs: afe.ir.attributes.TupleGetItemAttrs, t: tuple) -> numpy.ndarray .. py:function:: exp(data: numpy.ndarray) -> numpy.ndarray .. py:function:: sqrt(data: numpy.ndarray) -> numpy.ndarray .. py:function:: sigmoid(data: numpy.ndarray, mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray .. py:function:: swish(data: numpy.ndarray) -> numpy.ndarray .. py:function:: quick_gelu(data: numpy.ndarray) -> numpy.ndarray .. py:function:: hard_sigmoid(data: numpy.ndarray, alpha: float, beta: float) -> numpy.ndarray .. py:function:: hard_swish(data: numpy.ndarray) -> numpy.ndarray .. py:function:: log(data: numpy.ndarray) -> numpy.ndarray .. py:function:: softplus(data: numpy.ndarray) -> numpy.ndarray .. py:function:: erf(data: numpy.ndarray) -> numpy.ndarray .. py:function:: gelu(x: numpy.ndarray) -> numpy.ndarray .. py:function:: log2(data: numpy.ndarray) -> numpy.ndarray .. py:function:: log10(data: numpy.ndarray) -> numpy.ndarray .. py:function:: subtract_tensorflow(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray .. py:function:: power(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray .. py:function:: divide(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray .. py:function:: reciprocal(data: numpy.ndarray) -> numpy.ndarray .. py:function:: maximum(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray .. py:function:: minimum(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray .. py:function:: full(attrs: afe.ir.attributes.FullAttrs, fill_value: numpy.ndarray) -> numpy.ndarray .. py:function:: tile(attrs: afe.ir.attributes.TileAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: split(attrs: afe.ir.attributes.SplitAttrs, data: numpy.ndarray) -> Tuple[numpy.ndarray, Ellipsis] .. py:function:: take(attrs: afe.ir.attributes.TakeAttrs, data: numpy.ndarray, indices: numpy.ndarray) -> numpy.ndarray .. py:function:: expand_strided_slice_indices_to_shape_length(begin: list[int], end: list[int], strides: list[int], axes: list[int] | None, input_shape: list[int]) -> tuple[list[int], list[int], list[int]] Helper function for expanding begin, end and strides to match the shape length. .. py:function:: get_strided_slice_out_shape(attrs: afe.ir.attributes.StridedSliceAttrs) -> tuple[int, Ellipsis] Get StridedSliceOp output shape. :param attrs: Strided slice operator attributes. :returns: Output shape of strided slice operator. .. py:function:: strided_slice(attrs: afe.ir.attributes.StridedSliceAttrs, data: numpy.ndarray, run_mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray Execute strided slice operator. :param attrs: Strided slice operator attributes. :param data: Input tensor. :param run_mode: Execution mode. :returns: The result of a strided slice operator. .. py:function:: rsqrt(data: numpy.ndarray) -> numpy.ndarray .. py:function:: tanh(data: numpy.ndarray) -> numpy.ndarray .. py:function:: image_resize2d(attrs: afe.ir.attributes.ImageResize2DAttrs, data: numpy.ndarray, rounding: Optional[str] = None) -> numpy.ndarray AFE and MLA does not have a way to support nearest_neighbor with asymmetric. However, the error should be ignorable during inference. .. py:function:: upsample(attrs: afe.ir.attributes.UpsamplingAttrs, data: numpy.ndarray, rounding: Optional[str] = None) -> numpy.ndarray Upsample the input tensor along H and/or W dimension .. py:function:: gridsample(attrs: afe.ir.attributes.GridSampleAttrs, data: numpy.ndarray, grid: numpy.ndarray) -> numpy.ndarray Image interpolation through GridSample .. py:function:: layer_norm(attrs: afe.ir.attributes.LayerNormAttrs | afe.ir.attributes.LayerNormQuantAttrs, data: numpy.ndarray, run_mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray .. py:function:: rms_norm(data: numpy.ndarray, attrs: Union[afe.ir.attributes.RMSNormAttrs, afe.ir.attributes.RMSNormQuantAttrs], mode: RunMode) -> numpy.ndarray .. py:function:: instance_norm(data: numpy.ndarray, mean: numpy.ndarray, variance: numpy.ndarray, attrs: afe.ir.attributes.InstanceNormAttrs | afe.ir.attributes.InstanceNormQuantAttrs) .. py:function:: layout_transform(attrs: afe.ir.attributes.LayoutTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: calculate_tessellated_tensor_shape(tensor_type: afe.ir.tensor_type.TensorType, slice_shape: Sequence[int], align_c16: bool) -> tuple[int, int] .. py:function:: tessellation(attrs: afe.ir.attributes.TessellationTransformAttrs, data: numpy.ndarray) -> numpy.ndarray Input tensor is 4D NHWC, int8 only Output tensor is 2D array .. py:function:: detessellation(attrs: afe.ir.attributes.DetessellationTransformAttrs, data: numpy.ndarray) -> numpy.ndarray Input tensor is 2D Output tensor is 4D: NHWC .. py:function:: get_channel_aligned_shape(tensor_shape: Sequence[int], elem_size: int) -> tuple[int, Ellipsis] Helper function to get a tensor shape where channel is aligned based on the element size. .. py:function:: get_mla_padded_2d_shape(tensor_shape: Sequence[int], elem_size: int) -> tuple[int, int] .. py:function:: reshape_to_mla_padded_2d_shape(tensor: numpy.ndarray) -> numpy.ndarray Reshape tensor to MLA 2D buffer shape (batch_size, data_size (prod(spatial dimensions) * C)) where data size must be multiple of 16. .. py:function:: reshape_from_mla_padded_2d_shape(tensor: numpy.ndarray, data_shape: Sequence[int], tensor_type: type) -> numpy.ndarray Reshape tensor from MLA 2D shape to 4D/5D shape. :param tensor: 2D tensor. :param data_shape: 4D/5D tensor shape. :return: Reshaped 4D/5D tensor. .. py:function:: pack(data: List[numpy.ndarray]) -> numpy.ndarray Multiple tensors are packed sequentially as a 2D array. Input data can be either a 2D tessellated tensor or a 5D or higher tensor that will be tessellated on the MLA. If there is the 5D or higher tensor, reshape it to MLA 2D shape. .. py:function:: unpack(attrs: afe.ir.attributes.UnpackTransformAttrs, data: numpy.ndarray) -> List[numpy.ndarray] A 2D array is unpacked to produce multiple 2D arrays .. py:function:: normalization(attrs: afe.ir.attributes.NormalizationTransformAttrs, data: numpy.ndarray) -> numpy.ndarray Normalization performs the following three steps: 1) Divide by a per-channel divisor 2) Subtract by per-channel mean values 3) Divide by per-channel standard deviation values .. py:function:: ev_quantize(attrs: afe.ir.attributes.QuantizationTransformAttrs, data: numpy.ndarray) -> numpy.ndarray Quantization transform. .. py:function:: ev_dequantize(attrs: afe.ir.attributes.DequantizationTransformAttrs, data: numpy.ndarray) -> numpy.ndarray Dequantization transform. .. py:function:: ev_resize(attrs: afe.ir.attributes.ResizeTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: chroma_upsample(attrs: afe.ir.attributes.ChromaUpsampleTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: yuv_rgb_conversion(attrs: afe.ir.attributes.YuvRgbConversionTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: bgr_rgb_conversion(attrs: afe.ir.attributes.BgrRgbConversionTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: ev_sigmoid(attrs: afe.ir.attributes.SigmoidTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: nms_maxpool(attrs: afe.ir.attributes.NmsMaxpoolTransformAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: cast(attrs: afe.ir.attributes.CastAttrs, data: numpy.ndarray) -> numpy.ndarray .. py:function:: qnn_quantize(attrs: afe.ir.attributes.QNNQuantizeAttrs, data: numpy.ndarray, output_scale: numpy.ndarray, output_zero_point: numpy.ndarray) -> numpy.ndarray For the rounding type used for this operator (away from 0), refer to: https://github.com/apache/tvm/pull/3512/commits/c089ebcdf4b13f98b776bb0213779f6783fa6743#diff-a47be721cf0f30d86d0f548a8cc5a1fe184d0827efd450c8446bfc05d962abf5R47 .. py:function:: qnn_dequantize(attrs: afe.ir.attributes.QNNDequantizeAttrs, data: numpy.ndarray, input_scale: numpy.ndarray, input_zero_point: numpy.ndarray) -> numpy.ndarray .. py:function:: do_requantize(in_scale, out_scale, in_zp, out_zp) -> bool .. py:function:: qnn_add(attrs: afe.ir.attributes.AwesomeAttributes, lhs: numpy.ndarray, rhs: numpy.ndarray, lhs_scale: float, lhs_zero_point: int, rhs_scale: float, rhs_zero_point: int, output_scale: float, output_zero_point: int, op: str = 'add') -> numpy.ndarray .. py:function:: qnn_mul(attrs: afe.ir.attributes.AwesomeAttributes, lhs: numpy.ndarray, rhs: numpy.ndarray, lhs_scale: float, lhs_zero_point: int, rhs_scale: float, rhs_zero_point: int, output_scale: float, output_zero_point: int) -> numpy.ndarray .. py:function:: external(attrs: afe.ir.attributes.ExternalAttrs, input_dict: Dict[str, numpy.ndarray]) -> Union[numpy.ndarray, Tuple[numpy.ndarray, Ellipsis]] .. py:function:: init_custom_op(attrs: afe.ir.attributes.CustomOpAttrs, input_dict: Dict[afe.ir.defines.InputName, numpy.ndarray], output_shape: Tuple[int, Ellipsis], force_compile: bool = True) -> None Initialize the custom op. Compile the custom op and put it into the CustomOpLibraryManager. Update the CustomOpAttrs with generated arguments list and function so it can be used at the execution time. :param attrs: CustomOpAttrs :param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor :param output_shape: Tuple[int, ...]. Output shape :param force_compile: bool. Default is True. If True, the init_custom_op will compile the custom op no matter the library is ready or not .. py:function:: execute_custom_op(attrs: afe.ir.attributes.CustomOpAttrs, input_dict: Dict[afe.ir.defines.InputName, numpy.ndarray]) -> numpy.ndarray Execute the custom op :param attrs: CustomOpAttrs :param input_dict: Dict[InputName, np.ndarray]. Input name to its tensor :return: np.ndarray .. py:function:: batch_matmul(lhs: numpy.ndarray, rhs: numpy.ndarray, attrs: afe.ir.attributes.BatchMatmulAttrs | afe.ir.attributes.BatchMatmulQuantAttrs, run_mode: RunMode) -> numpy.ndarray Execute batch matmul operation using ml_kernels implementation. :param lhs: Tensor representing lhs value of batch matmul operation. :param rhs: Tensor representing rhs value of batch matmul operation. :param attrs: BatchMatmul operator's attributes. :param run_mode: Execution mode. :returns: The result of batch matmul operation.