def matrix_diag_part(input_diagonal, input_help): """ Calculate the batched diagonal part of a batched tensor. Note: input_help is a tensor with a diagonal element of 1 and other positions of 0, the last two dimensions can be unequal. Args: input_diagonal (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8. The last two dimensions can be unequal. input_help (tvm.tensor.Tensor): Tensor of float32, float16, int32, int8, uint8, and with a diagonal element of 1 and other positions of 0. Returns: tvm.tensor.Tensor, has the same type as input_diagonal, the shape dims is equal to dims(input_diagonal) - 1. """ dtype_input_diagonal = input_diagonal.dtype dtype_input_help = input_help.dtype utils.elemwise_shape_check(input_help.shape, input_diagonal.shape) if len(input_help.shape) < 2: raise ValueError("Input tensors of rank>=2 are supported!") utils.ops_dtype_check([dtype_input_diagonal, dtype_input_help], [ utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT8, utils.DtypeForDavinci.INT32, utils.DtypeForDavinci.UINT8 ]) res = matrix_diag_part_compute(input_diagonal, input_help) return res
def Sin(x, target=utils.CCE): """ Computes sine value of a tensor with Taylor's theorem. .. math:: \\begin{array}{ll} \\\\ sin(x) = x - \\frac{x^3}{3!} + \\frac{x^5}{5!} + ... + (-1)^k \\cdot \\frac{x^{2(k+1)}}{(2(k+1))!} \\end{array} Args: x (tvm.tensor.Tensor): Tensor of type float16, float32. Rerurns: tvm.tensor.Tensor of same type and shape as in_data. Supported Platforms: 'Ascend' """ utils.ops_dtype_check(x.dtype, utils.DtypeForDavinci.ALL_FLOAT) utils.check_shape(x.shape) use_call = True if use_call: return sin_call(x) return sin_compute(x)
def reciprocal(data, high_precision=True, target=utils.CCE): """ Computes the reciprocal of data element-wise. Args: data (list[tvm.tensor.Tensor]): a list of tvm.tensor.Tensor of type float16, float32. high_precision (bool): a bool value, whether to use high-precision version. Returns: tvm.tensor.Tensor of same type and shape as data. Supported Platforms: 'Ascend', 'GPU' """ utils.ops_dtype_check(data.dtype, utils.DtypeForDavinci.ALL_FLOAT) shape = [x.value for x in data.shape] utils.check_shape(shape) res = akg.tvm.compute( shape, lambda *indice: akg.tvm.const(1, data.dtype) / data(*indice), name="res") # When product is mini, using Newtom iteration method to achieve higher precision. if product_is_mini() and high_precision: steps = 1 for _ in range(steps): temp1 = data * res temp2 = temp1 * akg.tvm.const(-1, data.dtype) temp3 = temp2 + akg.tvm.const(2, data.dtype) res = temp3 * res return res
def xdivy(data_x1, data_x2, target=utils.CCE): """ Calculate data_x1 divided by data_x2. .. math:: y = \\left\\{ \\begin{aligned} 0, && if \\quad x1 == 0 \\\\ \\dfrac{x1}{x2}, && otherwise \\end{aligned} \\right. Args: data_x1 (tvm.tensor.Tensor): Tensor of dtype "float16" or "float32" data_x2 (tvm.tensor.Tensor): Tensor of dtype "float16" or "float32" Returns: tvm.tensor.Tensor """ shape_x1 = get_shape(data_x1) shape_x2 = get_shape(data_x2) utils.check_shape(shape_x1) utils.check_shape(shape_x2) utils.elemwise_dtype_check(data_x1.dtype, data_x2.dtype) dtype = data_x1.dtype utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) return xdivy_compute(data_x1, data_x2)
def pad(data, paddings, padtype, target="cce"): """add paddings to the tensor :shape: The shape of the tensor, now only support two dimension Tensor :paddings: The shape of the paddings, shape [N,2], N is the dimension of the tensor, For each dimension D of input, paddings[D, 0] indicates how many values to add before the contents of tensor in that dimension, and paddings[D, 1] indicates how many values to add after the contents of tensor in that dimension. :dtype: The type of the input, float16, float32 :padtype: One of "CONSTANT", "REFLECT", or "SYMMETRIC". """ # check shape utils.check_shape(data.shape) # check types utils.ops_dtype_check(data.dtype, utils.DtypeForDavinci.ALL_TYPES) # check padding types ptype_checklist = ['constant'] if not (padtype in ptype_checklist): raise RuntimeError("pad_cce only support %s while padtype is %s" % (",".join(ptype_checklist), padtype)) dtype = data.dtype if dtype == 'int8' or dtype == 'uint8': data = Cast(data, "float16", target=target) rank = len(data.shape) pad_before = [] pad_after = [] for i in range(rank): pad_before.append(paddings[i][0]) pad_after.append(paddings[i][1]) B = tvm_pad(data, pad_before, pad_after=pad_after, name='B') if dtype == 'int8' or dtype == 'uint8': B = Cast(B, dtype, target=target) return B
def flatten(x): """ reshape into (batch, c*h*w). Args: x (akg.tvm.tensor.Tensor): the first dimension is batch Returns: akg.tvm.tensor.Tensor """ # check shape utils.check_shape(x) shape = get_shape(x) # check input tensor data_type utils.ops_dtype_check(x.dtype, [ utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT8, utils.DtypeForDavinci.INT16, utils.DtypeForDavinci.INT32, utils.DtypeForDavinci.INT64, utils.DtypeForDavinci.UINT8, utils.DtypeForDavinci.UINT16, utils.DtypeForDavinci.UINT32, utils.DtypeForDavinci.UINT64 ]) size = 1 for i in range(1, len(shape)): size = size * shape[i] new_shape = [shape[0], size] res = akg.topi.reshape(x, new_shape) return res
def sum_v2(inputs, axis=None, keepdims=True, target=utils.CCE): """ another implementation of sum with topi api. Supported Platforms: 'Ascend' """ if target != utils.CCE: raise RuntimeError('operator not supported on %s' % utils.get_backend(target)) dtype = inputs.dtype utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) axis = ft_util.refine_reduce_axis(inputs, axis) utils.check_shape(inputs.shape) if not axis: output = akg.topi.identity(inputs) else: if dtype == "float16": step_sum = Cast(inputs, "float32", target) else: step_sum = inputs step_sum = akg.topi.sum(step_sum, axis=axis, keepdims=keepdims) if dtype == "float16": output = Cast(step_sum, "float16", target) else: output = step_sum return output
def resize_nearest(input, output_shape): """ Resize images using Nearest-neighbor interpolation. Args: input (tvm.tensor.Tensor): 4-D tensor of type float16 or float32 `("NHWC")`. output_shape (Union[tuple, list]): New size of image 4 integers `("NHWC")`. Note: The batch_num("N") of input and output must be equal, channel_num("C") is also. Returns: tvm.tensor.Tensor, has the same type as `input`. """ input_shape = get_shape(input) utils.check_shape(input, 4, "input") utils.check_shape(output_shape, 4, "output_shape") utils.ops_dtype_check(input.dtype, utils.DtypeForDavinci.ALL_FLOAT) utils.check_equal("input batchsize", "output batchsize", input_shape[0], output_shape[0]) utils.check_equal("input channel num", "output channel num", input_shape[3], output_shape[3]) res = process_integer_scale(input, output_shape) if res == None: res = process_non_integer_scale(input, output_shape) return res
def leaky_relu(data, negative_slop=0): """ leaky_relu op for input tensor (N,C,H,W) OR (N,C1,H,W,C0). ..math:`max(x,negative_slop*x)` Args: data (tvm.tensor.Tensor): tensor with type float16 or float32. negative_slop (float): 0<=negative_slop<1 Returns: tvm.tensor.Tensor. """ dtype = data.dtype utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) utils.check_shape(data.shape) if negative_slop >= 1 or negative_slop < 0: raise RuntimeError( "leaky_relu only support negative_slop between [0,1)") slop_tmp = akg.tvm.const(negative_slop, dtype=dtype) tmp = akg.lang.ascend.vmuls(data, slop_tmp) res = akg.lang.ascend.vmax(tmp, data) return res
def logsoftmax_grad(Y, dY, axis): """ Computes the back propagation gradients by chain rule. Args: Y: Tensor, holds the logsoftmax activation output. dY: Tensor, holds the initial gradients. axis: Integer, on which dimension the softmax is applied. Returns: Tensor, the overall gradients. """ shape = [x.value for x in Y.shape] utils.check_shape(shape) dtype = Y.dtype utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) if axis == -1: axis = len(shape) + axis if axis >= len(shape): raise RuntimeError("axis should be less than dimension") if axis < -1: raise RuntimeError( "negative axis only support -1, please specify the axis in positive value" ) softmax = akg.topi.exp(Y) dy_sum = akg.lang.ascend.sum(dY, axis=axis) dy_sum_broadcast = akg.lang.ascend.broadcast(dy_sum, shape) mul_result = akg.lang.ascend.vmul(softmax, dy_sum_broadcast) res = akg.lang.ascend.vsub(dY, mul_result) attrs = {"pragma_modshift": 1} return res, attrs
def reduce_any_d(x, axis=None, keepdims=False): """ Reduce a tensor on a certain axis based on max. Args: x (tvm.tensor.Tensor): The input tensor to reduce. Should be of type int8. axis (Union[list, tuple, int, None]): The dimensions to reduce. If None, all dimensions will be reduced. each dim must be in the range [-len(data.shape), len(data.shape) - 1]. keepdims (Union[bool, None]): If True, retains reduced dimensions with length 1, defaults to False. Returns: tvm.tensor.Tensor of same type as input tensor x. """ # check type utils.ops_dtype_check(x.dtype, utils.DtypeForDavinci.INT8) utils.check_shape(x.shape) # check axis utils.reduce_axis_check(x.shape, axis) refined_axis = refine_reduce_axis(x, axis) if len(set(refined_axis)) == len(x.shape) and not keepdims: keepdims = True res = _reduce_any_d_compute(x, refined_axis, keepdims) if len(set(refined_axis)) == len(x.shape): res = topi.reshape(res, (1, )) return res
def Concat(data, axis, target=utils.CCE): """ Concatenates data along the dimension set by axis. Args: data (Union[list, tuple]): list or tuple of tvm.tensor.Tensor of type float16, float32, int32, int8, uint8 axis (int): Specifies the axis along which to concatenate. Must be in the range [-rank(data), rank(data)) Returns: tvm.tensor.Tensor of same type as data. """ data_size = len(data) if data_size < min_size: raise RuntimeError("The size of data must be greater equal 1") dtype = data[0].dtype utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_TYPES) shape_0 = data[0].shape utils.check_shape(shape_0) if axis < 0: axis += len(shape_0) for i in range(1, data_size): shape_i = data[i].shape utils.check_shape(shape_i) if len(shape_i) != len(shape_0): raise ValueError("Input tensors must have same dimensions.") res = akg.lang.ascend.concat(data, axis) return res
def sum(inputs, axis=None, keepdims=False, target=utils.CCE): """ Compute the sum of elements across dimensions of a tensor. Args: inputs (tvm.tensor.Tensor): Tensor. axis (Union[list, tuple, int, None]): If the list or tuple is empty, the axis equal to None. keepdims (bool): If keepdims equal to True, the result shape length is same to input shape length. Returns: tvm.tensor.Tensor, has same type as input. If keepdims is True, all reduced dimensions are retained with length 1, else these reduced axis will be eliminate. Supported Platforms: 'Ascend', 'GPU', 'CPU' """ # Check types if target == utils.CCE: dtype = inputs.dtype utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) axis = ft_util.refine_reduce_axis(inputs, axis) utils.check_shape(inputs.shape) if not axis: output = akg.topi.identity(inputs) else: output = akg.topi.sum(inputs, axis=axis, keepdims=keepdims) return output
def bitwise_xor(x1, x2, target=utils.CCE): """ Computes the bitwise or of `x1` and `x2`. Args: x1 (tvm.tensor.Tensor): Tensor of type int16, uint16. x2 (tvm.tensor.Tensor): Tensor of type int16, uint16. Returns: tvm.tensor.Tensor, has the same type as x1. """ # check shape utils.check_shape(x1) utils.check_shape(x2) _, _, output_shape = produce_shapes(get_shape(x1), get_shape(x2)) # check input tensor data_type utils.ops_dtype_check( [x1.dtype, x2.dtype], [utils.DtypeForDavinci.INT16, utils.DtypeForDavinci.UINT16]) dtype = x1.dtype if dtype != x2.dtype: raise RuntimeError("input type must be same, but got %s vs %s", dtype, x2.dtype) x1 = akg.topi.broadcast_to(x1, output_shape) x2 = akg.topi.broadcast_to(x2, output_shape) # x1 ^ x2 = (x1 | x2) & (~ (x1 & x2)) res = akg.tvm.compute( output_shape, lambda *indice: (x1(*indice) | x2(*indice)) & (~(x1(*indice) & x2(*indice)))) return res
def dropout_do_mask(data_tensor, data_mask, keep_prob): dtype = data_tensor.dtype shape_tensor = [x.value for x in data_tensor.shape] utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) utils.check_shape(shape_tensor) strides = [1] for x in reversed(shape_tensor): strides.append(strides[-1] * x) if keep_prob < 0 or keep_prob > 1: raise RuntimeError("keep_prob must in [0,1]") keep_prob_const = akg.tvm.const(1.0 / keep_prob, dtype=dtype) data_scale_ub = akg.tvm.compute( shape_tensor, lambda *indices: data_tensor(*indices) * keep_prob_const, name='data_scale_ub') def get_index(indices): idx = 0 for i in range(len(indices)): idx += indices[len(indices) - i - 1] * strides[i] return idx // 8 if dtype == "float32": data_scale_ub_16 = akg.topi.cast(data_scale_ub, "float16") res_ub_16 = akg.tvm.compute(shape_tensor, lambda *indice: dav.dropout(data_mask[get_index(indice)], data_scale_ub_16(*indice))) res = akg.topi.cast(res_ub_16, "float32") else: res = akg.tvm.compute(shape_tensor, lambda *indice: dav.dropout(data_mask[get_index(indice)], data_scale_ub(*indice))) return res
def Addn(data, target=utils.CCE): """ Compute sum of all elements in tensor. Args: data (tvm.tensor.Tensor): Tensor of of type float16, float32. Returns: tvm.tensor.Tensor, compute result, get all elements' sum. Supported Platforms: 'Ascend', 'GPU', 'CPU' """ utils.check_supported_target(target) # check types dtype = data[0].dtype if target == utils.CCE: utils.ops_dtype_check(dtype, utils.DtypeForDavinci.ALL_FLOAT) res = data[0] for i in range(1, len(data)): utils.elemwise_dtype_check(res.dtype, data[i].dtype) utils.elemwise_shape_check(res.shape, data[i].shape) res = akg.topi.elemwise_sum(data) return res
def truncate_div(input_x1, input_x2): """ Calculating data's truncate_div, res = floor(x1/x2) if x1/x2>0 else ceil(x1/x2). Args: input_x1 (tvm.tensor.Tensor): Input tensor, support float16, float32 on mini device, while support int32, int8, uint8, float16, float32 on cloud ones. input_x2 (tvm.tensor.Tensor): Input tensor, with same dtype as input_x1. Returns: A tvm.tensor.Tensor as result of truncate_div. """ utils.check_shape(get_shape(input_x1)) utils.check_shape(get_shape(input_x2)) utils.elemwise_dtype_check(input_x1.dtype, input_x2.dtype) utils.ops_dtype_check( input_x1.dtype, (utils.DtypeForDavinci.ALL_FLOAT) if product_is_mini() \ else (utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT32, utils.DtypeForDavinci.INT8, utils.DtypeForDavinci.UINT8)) return truncate_div_compute(input_x1, input_x2)
def case_1(data_shape, dtype, kernel_name, attrs): """elemwise chain case 1""" utils.ops_dtype_check(dtype, utils.DtypeForDavinci.FLOAT16) utils.check_shape_length_equal("data", data_shape, 2) m, k = data_shape A = akg.tvm.placeholder((m, k), name='A', dtype=dtype) B = akg.tvm.placeholder((k, ), name='B', dtype=dtype) C = akg.tvm.placeholder((m, k), name='C', dtype=dtype) E = akg.tvm.compute((m, k), lambda i, j: A[i, j] * (B[j] + C[i, j]), name="E") forward_s = akg.tvm.create_schedule(E.op) op_vars = [A, B, C, E] akg.lower(forward_s, op_vars, simple_mode=True, polyhedral=True) kernel_name = gen_name_kernel(kernel_name, dtype, data_shape) with akg.build_config(add_lower_pass=debug_mode(0), dump_pass_ir=True): mod = akg.build(forward_s, op_vars, "cce", name="test", attrs=attrs, polyhedral=True) return mod
def mean_v2(data, axis=None, keepdims=False, target=utils.CCE): """ Simple implementation of mean. Supported Platforms: 'Ascend' """ # Check types utils.ops_dtype_check(data.dtype, utils.DtypeForDavinci.ALL_FLOAT) # Check shape shape = [x.value for x in data.shape] utils.reduce_axis_check(shape, axis) axis = ft_util.refine_reduce_axis(data, axis) dtype = data.dtype count = 1 for i in axis: count *= shape[i] count_rec = 1 / count output = sum_v2(data, axis, keepdims, target=target) res = output * akg.tvm.const(count_rec, dtype) attrs = get_attrs(data) if shape_is_dynamic(data): attrs["custom_tiling"] = mean_dynamic_tiling_strategy(data, axis) return res, attrs
def apply_rms_prop_mixed_precision(var, ms, mom, grad, lr, momentum, rho, epsilon): """ Mixed precision version for apply_rms_prop. Args: var (tvm.tensor.Tensor): The tensor to be updated. Should be float32. ms (tvm.tensor.Tensor): Mean square, a tensor of same shape and type as var. mom (tvm.tensor.Tensor): A tensor of same shape and type as var. grad (tvm.tensor.Tensor): A tensor of same shape and type as var. lr (tvm.tensor.Tensor): Learning rate, a scalar tensor of same type as var. momentum (float): Coefficient for calculate new mom, 0.0 <= momentum <= 1.0. rho (float): Coefficient for calculate new ms, 0.0 <= rho <= 1.0. epsilon (float): A small value to prevent division by 0. Returns: tvm.tensor.Tensor, Updated var of type float32. tvm.tensor.Tensor, Updated var of type float16. tvm.tensor.Tensor, Updated ms. tvm.tensor.Tensor, Updated mom. """ utils.ops_dtype_check(var.dtype, utils.DtypeForDavinci.FLOAT32) _apply_rms_prop_check(var, ms, mom, grad, lr, momentum, rho, epsilon) out_var, out_var_fp16, out_ms, out_mom = _apply_rms_prop_mixed_precision_compute( var, ms, mom, grad, lr, momentum, rho, epsilon) out_var, binds_info = TensorUtils.inplace_set(var, out_var, "var_buf") out_ms, binds_info2 = TensorUtils.inplace_set(ms, out_ms, "ms_buf") out_mom, binds_info3 = TensorUtils.inplace_set(mom, out_mom, "mom_buf") binds_info.update(binds_info2) binds_info.update(binds_info3) attrs = {utils.BINDS: binds_info} return out_var, out_var_fp16, out_ms, out_mom, attrs
def discontinous_mov(data, out_shape, target=utils.CCE): """ Extract the element with the odd index from the original data and copy it into a tensor with a dimension of 2 * original dimension/2. Args: data (tvm.tensor.Tensor): Tensor of type float16, float32. out_shape (list): a list of output's shape. Returns: tvm.tensor.Tensor, has the same type as data, but it's shape changes to out_shape not data's shape. Example: if data = [1,2,3,4,5,6,7,8,9,10] then the output = [[1,3,5,7,9],[1,3,5,7,9]]. """ # check types utils.ops_dtype_check(data.dtype, utils.DtypeForDavinci.ALL_FLOAT) shape = [x.value for x in data.shape] utils.check_shape(shape) output = akg.tvm.compute(out_shape, lambda j, i: data[i * 2], name="output") return output
def matrix_diag(data, out_shape): """ Generate a batched tensor whose value in diagonal lines are defined in `data`. Args: data (tvm.tensor.Tensor): A tensor of type float16, float32 or int32. Rank is L. out_shape (Union[list, tuple]): Output shape of length L + 1. The value of `out_shape[0, ..., L-1]` should be equal to `data.shape[0, ..., L-1]`. Returns: tvm.tensor.Tensor, has same type as "data", shape is "out_shape". """ dtype = data.dtype utils.ops_dtype_check(dtype, [utils.DtypeForDavinci.ALL_FLOAT, utils.DtypeForDavinci.INT32]) shape = get_shape(data) utils.check_shape(data) utils.check_shape(out_shape, length=len(shape) + 1) if tuple(shape[:-1]) != tuple(out_shape[:-2]): raise RuntimeError("The value of out_shape[:-2] should be equal to data.shape[:-1]") res = akg.tvm.compute(out_shape, lambda *i: akg.tvm.if_then_else(akg.tvm.all(i[-1] == i[-2], i[-1] < shape[-1]), data(*i[:-1]), zero_const(dtype)), name="diag") return res
def sparse_softmax_cross_entropy_with_logits(labels, logits, reduction='mean'): """ Computes sparse softmax cross entropy between `logits` and `labels`. Note: Softmax calculation of Logits is done inside the op. Args: labels (tvm.tensor.Tensor): int32 tensor of shape [batch_size]. Each entry in it must be an index in `[0, num_classes)`. logits (tvm.tensor.Tensor): float32 or float16 tensor of shape [batch_size, num_class]. reduction (str): Specifies the reduction to apply to the output: 'none' or 'mean' or 'sum'. Default: 'mean'. 'none': no reduction for the output 'sum': the sum for the output 'mean': the mean for the output. Returns: tvm.tensor.Tensor, has the same dtype as logits. If reduction is 'none', shape of the tensor is the same as logits, otherwise shape of the tensor is the same as labels. Supported Platforms: 'Ascend' """ utils.ops_dtype_check(logits.dtype, utils.DtypeForDavinci.ALL_FLOAT) strategy, cost, _ = sparse_softmax_cross_entropy_with_logits_impl( labels, logits, reduction) attr_map = {"custom_tiling": strategy} return cost, attr_map
def asinh(x, target=utils.CCE): r""" Compute asinh function. .. math:: asinh(x) = log(x+\sqrt{x*x+1}) Args: x (tvm.tensor.Tensor): Tensor of type float16, float32. Returns: tvm.tensor.Tensor, has the same type and shape as x. Supported Platforms: 'Ascend' """ # check shape utils.check_shape(x) # check input tensor data_type utils.ops_dtype_check(x.dtype, utils.DtypeForDavinci.ALL_FLOAT) dtype = x.dtype # Known that, asinh(x) = log(x + sqrt(x*x+1)), and, asinh(-x) = -asinh(x) # If x is a large negative number, (x + sqrt(x*x+1)) will be close to zero. # So, asinh(x) = sign(x) * log(|x| + sqrt(|x|*|x| + 1)) compute_dtype = dtype if dtype == "float16": # To avoid overflow and higher accuracy, x is casted to float32 compute_dtype = "float32" x = topi.cast(x, compute_dtype) x_abs = topi.abs(x) if product_is_mini(): # sqrt(|x|*|x| + 1) = |x| * sqrt(1 + 1/(|x|*|x|)) vsquare_add_one = topi.add(1, topi.divide(1, topi.multiply(x_abs, x_abs))) sqrt_compute_value = sqrt_mini_newton_iter_impl(vsquare_add_one) sqrt_value = topi.multiply(x_abs, sqrt_compute_value) else: x_abs_square_add_one = topi.add(topi.multiply(x_abs, x_abs), 1) sqrt_value = topi.sqrt(x_abs_square_add_one) x_add_sqrt = topi.add(x_abs, sqrt_value) if product_is_mini(): log_value = log_compute_mini_impl(x_add_sqrt, target) else: log_value = topi.log(x_add_sqrt) res = topi.multiply(Sign(x, target), log_value) if res.dtype != dtype: res = topi.cast(res, dtype) if product_is_mini(): attrs = {"enable_auto_inline": False} return res, attrs return res
def quantize_chk_cfg_and_gen_outdtype(quant_algo, scale_mode, scale_sqrt, qdrtensors): """check all the params is valid, and general output dtype""" # check quantize algorithm and quantize scale type if quant_algo is None: # quantize switch off if scale_mode is not None or scale_sqrt is not None \ or qdrtensors is not None: raise RuntimeError("Invalid Quantize Config.") out_dtype = "float16" return out_dtype # quantize switch on, all quantize params should not be None if scale_mode is None or scale_sqrt is None \ or qdrtensors is None: raise RuntimeError("Invalid Quantize Config!") if len(quant_algo) != 2 or any([i not in [0, 1] for i in quant_algo]): raise RuntimeError("Invalid Quantize Config!!!") # check quantize algorithm if quant_algo[0] not in (0, 1): raise RuntimeError("Quantize algorithm just support 0 for non " "offset and 1 for half offset, but get {}." "".format(quant_algo[0])) # check quantize scale type if quant_algo[1] != 0: raise RuntimeError("Quantize scale only support SCALAR now.") # non offset get int8, half offset get uint8 out_dtype = "int8" if quant_algo[0] == 0 else "uint8" if scale_mode not in (0, 1, 2): raise ValueError("Invalid scale mode, just support '0,1,2' but get " "{}!".format(scale_mode)) # now scale mode limit if scale_mode in (0, 1): raise RuntimeError("quantized_avg_pool just support requantize now!") # check scale method if scale_sqrt not in (0, 1): raise RuntimeError("Invalid scale moethod!") # scalar scale type for now if len(qdrtensors) != 2: raise RuntimeError("qdrtensors should contain two tensors for scale " "and offset!") if get_shape(qdrtensors[0]) != [1] or get_shape(qdrtensors[1]) != [1]: raise RuntimeError("Scale for dequantize or requantize only " "support scalar tensor.") utils.ops_dtype_check(qdrtensors[0].dtype, utils.DtypeForDavinci.FLOAT16) utils.ops_dtype_check(qdrtensors[1].dtype, utils.DtypeForDavinci.FLOAT16) # utils.ops_dtype_check(qdrtensors[0].dtype, # utils.DtypeForDavinci.ALL_FLOAT) # utils.ops_dtype_check(qdrtensors[1].dtype, # utils.DtypeForDavinci.ALL_FLOAT) return out_dtype
def fused_bn2(mean, var_part, running_mean, running_var, momentum=0.8): """ Calculating mean, variance and update running variables. Read fused_bn1 docs for details. Note: Apply reduction of 'N' axis to calculating mean and variance. Args: mean (tvm.tensor.Tensor): Tensor of type float32 as mean. var_part (tvm.tensor.Tensor): Tensor of type float32, intermediate variables for variance. running_mean (tvm.tensor.Tensor): Tensor of type float32 as trained mean used in inference stage. running_var (tvm.tensor.Tensor): Tensor of type float32 as trained variance used in inference stage. momentum (float): A float number used for updating running values, must meet condition '0.0 < momentum < 1.0'. Returns: variance (tvm.tensor.Tensor): A float32 tensor as data's variance. running_mean_updated (tvm.tensor.Tensor): A float32 tensor as updated running_mean (updated inplace). running_var_updated (tvm.tensor.Tensor): A float32 tensor, updated running_var (updated inplace). """ utils.ops_dtype_check([mean.dtype, var_part.dtype], utils.DtypeForDavinci.FLOAT32) dim_info, _ = bn2_set_dim_func(mean, var_part, running_mean, running_var, momentum) attrs = {**ATTR_MAP_BN2} in_tensors = (var_part, mean, running_mean, running_var) sub_mean_square = akg.tvm.compute( mean.shape, lambda *i: akg.tvm.const(-1.0, dtype=mean.dtype) * mean(*i) * mean(*i), name="sub_mean_square") variance = akg.tvm.compute(mean.shape, lambda *i: var_part(*i) + sub_mean_square(*i), name="variance") # update running mean and variance running_mean_updated = \ update_by_moving_average(running_mean, mean, momentum) running_var_updated = \ update_by_moving_average(running_var, variance, momentum) out_tensors = (variance, running_mean_updated, running_var_updated) tensors_and_binds = inplace_operate_bind(in_tensors, out_tensors, ((2, 1), (3, 2))) out_tensors = tensors_and_binds[0] attrs[kernel_exec.BINDS] = tensors_and_binds[1] if dim_info != "": attrs["dim"] = dim_info return (*out_tensors, attrs)
def _div_ascend(data1, data2): """ Calculates x/y, and returns an integer when inputs are all integers. When both arguments are integers, use integer division (also known as "floor division"). When arguments are float numbers, use normal floating point division Note: div supports broadcasting. Args: data1 (tvm.tensor.Tensor): Tensor of type float16, float32, int32, int8 and uint8. data2 (tvm.tensor.Tensor): Tensor of type float16, float32, int32, int8 and uint8. Returns: tvm.tensor.Tensor, has the same type as data1 and data2. """ utils.ops_dtype_check([data1.dtype, data2.dtype], utils.DtypeForDavinci.ALL_TYPES) utils.elemwise_dtype_check(data1.dtype, data2.dtype) dtype = data1.dtype shape1 = [x.value for x in data1.shape] shape2 = [x.value for x in data2.shape] utils.check_shape(shape1) utils.check_shape(shape2) utils.auto_broadcast_check(shape1, shape2) n_shape1, n_shape2, out_shape = produce_shapes(shape1, shape2) if n_shape1 != out_shape: input1_cast = akg.topi.broadcast_to(data1, out_shape) else: input1_cast = data1 if n_shape2 != out_shape: input2_cast = akg.topi.broadcast_to(data2, out_shape) else: input2_cast = data2 if dtype in ("int32", "int8", "uint8"): input1p = Case(input1_cast, "float16", utils.CCE) input2p = Cast(input2_cast, "float16", utils.CCE) else: input1p = input1_cast input2p = input2_cast if product_is_mini(): input2p_rec = reciprocal(input2p, target=utils.CCE) res = akg.topi.multiply(input1p, input2p_rec) else: res = akg.topi.divide(input1p, input2p) if dtype in ("int8", "uint8"): res = floor(res, utils.CCE) res = Cast(res, "float16", utils.CCE) if dtype in ("int32", "int8", "uint8"): res = Cast(res, dtype, utils.CCE) return res
def check_op1(*args): dy, data, mean = args utils.ops_dtype_check([dy.dtype, data.dtype], utils.DtypeForDavinci.ALL_FLOAT) utils.ops_dtype_check(mean.dtype, utils.DtypeForDavinci.FLOAT32) shape_nc1hwc0 = get_shape(dy) check_shape("NC1HWC0", data, shape_nc1hwc0, "data") check_shape("C1C0", mean, shape_nc1hwc0, "mean") return shape_nc1hwc0
def fill(shape, value, dtype, target="cce"): utils.ops_dtype_check( dtype, [utils.DtypeForDavinci.FLOAT16, utils.DtypeForDavinci.INT32]) utils.check_shape(shape) A = akg.tvm.const(value, dtype) res = akg.tvm.compute(shape, lambda *i: A, name="fill") return res
def matmul(x, y, b, out_dtype, left_format="zZ", right_format="nZ", out_format="zN", transpose_x=False, transpose_y=False, attrs=None, target=utils.CCE): """ Computes matrix multiplication x * y + b. Args: x: akg.tvm.Tensor of type int8, uint8, float16, float32, int32. Left matrix. y: akg.tvm.Tensor of same type as x. Right matrix. b: akg.tvm.Tensor of same type as x. Bias tensor. out_dtype: str. Data type of output tensor. left_format: str. Data format of left matrix. Supported data format list ["zZ", "nZ", "zN"]. right_format: str. Data format of right matrix. Supported data format list ["zZ", "nZ", "zN"]. out_format: str. Data format of output tensor. Supported data format list ["zZ", "nZ", "zN"]. transpose_x: Boolean. Specifies whether x is transposed or not. transpose_y: Boolean. Specifies whether y is transposed or not. attrs: Dict. Used in matmul computation. Note: before call matmul, 2d to Fractal is needed. Returns: akg.tvm.Tensor with type out_dtype. Supported Platforms: 'Ascend' """ utils.ops_dtype_check([x.dtype, y.dtype], utils.DtypeForDavinci.ALL_FLOAT) shape_x = [shape_element.value for shape_element in x.shape] utils.check_shape(shape_x) shape_y = [shape_element.value for shape_element in y.shape] utils.check_shape(shape_y) if left_format not in ["zZ", "zN"]: raise ValueError("unsupport left_format now: %s" % left_format) if right_format not in ["nZ", "zZ", "zN"]: raise ValueError("unsupport right_format now: %s" % right_format) if out_format not in ["zN", "zZ"]: raise ValueError("unsupport out_format now: %s" % out_format) out = matmul4d_compute(x, y, b, out_dtype, left_format, right_format, out_format, transpose_x, transpose_y, attrs) attr_map = {"pragma_rmselfdep": False} dims_info, _ = matmul_set_dim(x, y, b, out_dtype, left_format, right_format, out_format, transpose_x, transpose_y) attr_map["dim"] = dims_info return out, attr_map