DSL description of the layernorm operator's mathematical calculation process Parameters ---------- input_x: TVM tensor the placeholder of x input data input_gamma: TVM tensor the placeholder of gamma input data input_beta: TVM tensor the placeholder
(input_x, input_gamma, input_beta,
output_y, output_mean, output_variance,
begin_norm_axis, begin_params_axis,
epsilon, kernel_name="layer_norm",
impl_mode="high_performance")
| 569 | # 'pylint: disable=too-many-statements,too-many-branches |
| 570 | @tbe_platform.fusion_manager.fusion_manager.register("layer_norm") |
| 571 | def layer_norm_compute(input_x, input_gamma, input_beta, |
| 572 | output_y, output_mean, output_variance, |
| 573 | begin_norm_axis, begin_params_axis, |
| 574 | epsilon, kernel_name="layer_norm", |
| 575 | impl_mode="high_performance"): |
| 576 | """ |
| 577 | DSL description of the layernorm operator's mathematical calculation process |
| 578 | |
| 579 | Parameters |
| 580 | ---------- |
| 581 | input_x: TVM tensor |
| 582 | the placeholder of x input data |
| 583 | input_gamma: TVM tensor |
| 584 | the placeholder of gamma input data |
| 585 | input_beta: TVM tensor |
| 586 | the placeholder of beta input data |
| 587 | output_data: dict |
| 588 | shape and dtype of output |
| 589 | begin_norm_axis: int |
| 590 | The first normalization dimension: normalization will be |
| 591 | performed along dimensions `begin_norm_axis : rank(inputs)` |
| 592 | begin_params_axis: int |
| 593 | The first parameter (beta, gamma) dimension: scale |
| 594 | and centering parameters will have dimensions |
| 595 | `begin_params_axis : rank(inputs)` and will be broadcast with the |
| 596 | normalized inputs accordingly. |
| 597 | epsilon: float, |
| 598 | Minimum positive number greater than 0 |
| 599 | kernel_name: str |
| 600 | cce kernel name, default value is "cce_layernorm" |
| 601 | |
| 602 | Returns |
| 603 | ------- |
| 604 | res_tuple: tuple |
| 605 | (mean, variance, result) |
| 606 | """ |
| 607 | shape_x = shape_util.shape_to_list(input_x.shape) |
| 608 | dtype = input_x.dtype.lower() |
| 609 | cast_dtype, cast_fp16_dtype = "float16", "float16" |
| 610 | cast_dtype_precision = dtype |
| 611 | if dtype == "float16" and \ |
| 612 | ((tbe_platform.cce_conf.api_check_support |
| 613 | ("te.lang.cce.vexp", "float32") and |
| 614 | impl_mode == "high_performance") or |
| 615 | impl_mode == "high_precision"): |
| 616 | cast_dtype = "float32" |
| 617 | cast_dtype_precision = "float32" |
| 618 | input_x = tbe.cast_to(input_x, "float32") |
| 619 | input_gamma = tbe.cast_to(input_gamma, "float32") |
| 620 | input_beta = tbe.cast_to(input_beta, "float32") |
| 621 | |
| 622 | # Calculate the scaling ratio of the average |
| 623 | index_list = tuple(index for index, _ in enumerate(shape_x)) |
| 624 | reduce_axis = index_list[begin_norm_axis:] |
| 625 | |
| 626 | reduce_elts = 1.0 |
| 627 | for i in reduce_axis: |
| 628 | reduce_elts *= shape_x[i] |