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.conv_tensorflow afe.ir.operation_functions.add 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.adaptive_avgpool2d afe.ir.operation_functions.broadcast_to afe.ir.operation_functions.multiply afe.ir.operation_functions.quantized_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.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 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.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 Supported run modes. MLA_MODE : use an implementation that exactly matches execution on the MLA. FAST_MODE : use a fast execution implementation .. py:attribute:: MLA_MODE :value: 1 .. py:attribute:: FAST_MODE :value: 2 .. py:attribute:: FAST_MODE_MODEL_SDK :value: 3 .. py:method:: is_fast_mode() .. 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:: 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:: add(lhs: numpy.ndarray, rhs: numpy.ndarray, axis: Optional[int] = None) -> numpy.ndarray .. 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) .. 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], mode: RunMode = RunMode.MLA_MODE) -> numpy.ndarray .. py:function:: avgpool(attrs: afe.ir.attributes.PoolAttrs, data: numpy.ndarray, pad_value: Union[float, int], quantized: bool = False, rounding_type: Optional[ml_kernels.math_helpers.RoundType] = None, requant: Optional[ml_kernels.requantization.BaseRequantization] = 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:: adaptive_avgpool2d(attrs: afe.ir.attributes.AdaptiveAvgPool2DAttrs, data: numpy.ndarray, quantized: bool = False, rounding_type: Optional[ml_kernels.math_helpers.RoundType] = None, *args, **kwargs) -> numpy.ndarray .. py:function:: broadcast_to(attrs: Union[afe.ir.attributes.BroadcastToAttrs, afe.ir.attributes.BroadcastToQuantAttrs], data: numpy.ndarray) .. py:function:: multiply(lhs: numpy.ndarray, rhs: numpy.ndarray) -> numpy.ndarray Floating-point multiplication. .. py:function:: quantized_multiply(attrs: afe.ir.attributes.MultiplyQuantAttrs, lhs: numpy.ndarray, rhs: numpy.ndarray) -> 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.MeanAttrs, 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: MeanAttrs. 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) -> 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, data_list: List[numpy.ndarray]) -> numpy.ndarray .. 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.SumAttrs, 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) -> numpy.ndarray .. py:function:: swish(data: numpy.ndarray) -> numpy.ndarray .. py:function:: hard_sigmoid(data: numpy.ndarray) -> 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(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:: strided_slice(attrs: afe.ir.attributes.StridedSliceAttrs, data: numpy.ndarray) -> numpy.ndarray .. 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) -> numpy.ndarray .. py:function:: rms_norm(data: numpy.ndarray, attrs: Union[afe.ir.attributes.RMSNormAttrs, afe.ir.attributes.RMSNormQuantAttrs]) -> 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 (D * H * W * 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 4D tensor that will be tessellated on the MLA. If there is the 4D 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: Union[afe.ir.attributes.BatchMatmulAttrs, afe.ir.attributes.BatchMatmulQuantAttrs]) -> numpy.ndarray Execute batch matmul operation. :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. :return: The result of batch matmul operation.