DSL description of the layernorm operator's mathematical calculation process for non_aligned scene
(input_x, input_gamma, input_beta,
output_y, output_mean, output_variance,
begin_norm_axis, begin_params_axis,
ori_shape, epsilon, kernel_name="layer_norm",
impl_mode="high_performance")
| 352 | |
| 353 | # 'pylint: disable=too-many-locals,too-many-statements,too-many-branches |
| 354 | def nz_non_aligned(input_x, input_gamma, input_beta, |
| 355 | output_y, output_mean, output_variance, |
| 356 | begin_norm_axis, begin_params_axis, |
| 357 | ori_shape, epsilon, kernel_name="layer_norm", |
| 358 | impl_mode="high_performance"): |
| 359 | """ |
| 360 | DSL description of the layernorm operator's mathematical calculation process for non_aligned scene |
| 361 | """ |
| 362 | shape_x = shape_util.shape_to_list(input_x.shape) |
| 363 | dtype = input_x.dtype.lower() |
| 364 | cast_dtype = "float16" |
| 365 | if dtype == "float16" and \ |
| 366 | ((tbe_platform.cce_conf.api_check_support |
| 367 | ("te.lang.cce.vexp", "float32") and |
| 368 | impl_mode == "high_performance") or |
| 369 | impl_mode == "high_precision"): |
| 370 | cast_dtype = "float32" |
| 371 | input_x = tbe.cast_to(input_x, "float32") |
| 372 | input_gamma = tbe.cast_to(input_gamma, "float32") |
| 373 | input_beta = tbe.cast_to(input_beta, "float32") |
| 374 | else: |
| 375 | input_x = tbe.vadds(input_x, 0) |
| 376 | |
| 377 | # Calculate the scaling ratio of the average |
| 378 | reduce_elts = 1.0 |
| 379 | index_list = tuple(index for index, _ in enumerate(ori_shape)) |
| 380 | reduce_axis = index_list[begin_norm_axis:] |
| 381 | for i in reduce_axis: |
| 382 | reduce_elts *= ori_shape[i] |
| 383 | reduce_axis = to_frac_z_axis(ori_shape, reduce_axis) |
| 384 | mean_cof = reduce_elts ** (-1) |
| 385 | |
| 386 | # DSL description of the mean calculation process |
| 387 | with tvm.tag_scope("tail_block_pretreatment"): |
| 388 | lambda_func = lambda *indice: tvm.const(0, input_x.dtype) |
| 389 | temp = tvm.compute(input_x.shape, lambda_func, name="tail_block_pretreatment") |
| 390 | |
| 391 | input_x = tbe.vadd(input_x, temp) |
| 392 | mean_muls = tbe.vmuls(input_x, mean_cof) |
| 393 | mean = tbe.sum(mean_muls, axis=reduce_axis, keepdims=True) |
| 394 | |
| 395 | mean_square = tbe.vmul(mean, mean) |
| 396 | x_square = tbe.vmul(input_x, input_x) |
| 397 | x_square = tbe.vmuls(x_square, mean_cof) |
| 398 | x_square_mean = tbe.sum(x_square, axis=reduce_axis, keepdims=True) |
| 399 | variance = tbe.vsub(x_square_mean, mean_square) |
| 400 | |
| 401 | # DSL description of the normalize calculation process |
| 402 | mean_normalize_broadcast = _broadcast_nz(mean, shape_x) |
| 403 | normalize_sub = tbe.vsub(input_x, mean_normalize_broadcast) |
| 404 | epsilon = tvm.const(epsilon, dtype=cast_dtype) |
| 405 | |
| 406 | normalize_add = tbe.vadds(variance, epsilon) |
| 407 | normalize_log = tbe.vlog(normalize_add) |
| 408 | normalize_log_mul = \ |
| 409 | tbe.vmuls(normalize_log, tvm.const(-0.5, dtype=cast_dtype)) |
| 410 | normalize_exp = tbe.vexp(normalize_log_mul) |
| 411 | variance_normalize_broadcast = _broadcast_nz(normalize_exp, shape_x) |
no test coverage detected