layer_norm converter

Layer norm linting correction

ops file correction

fixing lint

Acc_ops layer_norm correction

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -30,6 +30,7 @@
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
 from torch_tensorrt.fx.converters.impl.elementwise import fmod
 from torch_tensorrt.fx.converters.impl.normalization import batch_norm
+from torch_tensorrt.fx.converters.impl.normalization import layer_norm
 from torch_tensorrt.fx.converters.impl.unary import sign
 from torch_tensorrt.fx.converters.impl.elementwise.base import (
     convert_binary_elementwise,
@@ -649,171 +650,16 @@ def acc_ops_batch_norm(
 
 @tensorrt_converter(acc_ops.layer_norm)
 def acc_ops_layer_norm(network, target, args, kwargs, name):
-    input_val = kwargs["input"]
-
-    if not isinstance(input_val, trt.tensorrt.ITensor):
-        raise RuntimeError(
-            f"LayerNorm received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    gamma = kwargs["weight"].detach().cpu().float().numpy()
-    gamma_field = trt.PluginField("gamma", gamma, trt.PluginFieldType.FLOAT32)
-    beta = kwargs["bias"].detach().cpu().float().numpy()
-    beta_field = trt.PluginField("beta", beta, trt.PluginFieldType.FLOAT32)
-    eps_field = trt.PluginField(
-        "eps", np.array([kwargs["eps"]], dtype=np.float32), trt.PluginFieldType.FLOAT32
-    )
-    normalized_shape = kwargs["normalized_shape"]
-    try:
-        normalized_shape = np.array(normalized_shape, dtype=np.int32)
-    except TypeError:
-        _LOGGER.error(
-            f"Unable to convert normalized_shape with value {normalized_shape} to a field, fall back to []"
-        )
-        normalized_shape = np.array([], dtype=np.int32)
-
-    normalized_shape_filed = trt.PluginField(
-        "normalized_shape", normalized_shape, trt.PluginFieldType.INT32
-    )
-    field_collection = trt.PluginFieldCollection(
-        [gamma_field, beta_field, eps_field, normalized_shape_filed]
-    )
-
-    try:
-        if network.has_implicit_batch_dimension:
-            plugin = get_trt_plugin("layer_norm", field_collection, "1", "fx2trt")
-        else:
-            plugin = get_trt_plugin("LayerNormDynamic", field_collection, "1", "fx2trt")
-    except AssertionError:
-        _LOGGER.error(
-            "Unable to find layer norm plugin, fall back to TensorRT implementation."
-        )
-        return layer_norm(network, target, args, kwargs, name)
-    layer = network.add_plugin_v2([input_val], plugin)
-    layer.name = name
-    return layer.get_output(0)
-
-
-def layer_norm(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"LayerNorm received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    shape = kwargs["weight"].shape  # type: ignore[union-attr]
-    broadcasted_shape = (1,) * (len(input_val.shape) - len(shape)) + shape
-    gamma = to_numpy(kwargs["weight"].reshape(*shape))  # type: ignore[union-attr]
-    beta = to_numpy(kwargs["bias"].reshape(*shape))  # type: ignore[union-attr]
-    eps = kwargs["eps"]
-
-    axes = 0
-    for d in range(len(shape)):
-        axes |= 1 << (len(input_val.shape) - d - 1)
-
-    # E[x]
-    mean_expected_layer = network.add_reduce(
-        input_val, trt.ReduceOperation.AVG, axes, keep_dims=True
-    )
-    set_layer_name(mean_expected_layer, target, f"{name}_mean_expected")
-
-    # X-E[x]
-    sub_trt = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        f"{name}_sub",
-        trt.ElementWiseOperation.SUB,
-        input_val,
-        mean_expected_layer.get_output(0),
-    )
-    # Variance = mean(pow(x_sub_mean,2))
-    pow_tensor = network.add_constant(
-        (1,) * len(input_val.shape),
-        trt.Weights(np.ascontiguousarray([2.0], dtype=np.float32)),
-    )
-    pow_tensor.name = f"{name}_power"
-    pow_var = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        f"{name}_pow_var",
-        trt.ElementWiseOperation.POW,
-        sub_trt,
-        pow_tensor.get_output(0),
-    )
-    mean_trt_layer = network.add_reduce(
-        pow_var, trt.ReduceOperation.AVG, axes, keep_dims=True
-    )
-    set_layer_name(mean_trt_layer, target, f"{name}_mean")
-    # Variance + eps
-    eps_tensor = network.add_constant(
-        (1,) * len(input_val.shape),
-        trt.Weights(np.ascontiguousarray([eps], dtype=np.float32)),
-    )
-    eps_tensor.name = f"{name}_eps"
-    add_trt = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        f"{name}_add",
-        trt.ElementWiseOperation.SUM,
-        mean_trt_layer.get_output(0),
-        eps_tensor.get_output(0),
-    )
-    # SQRT((Var + eps))
-    sqrt_trt = convert_unary(
-        network,
-        target,
-        SourceIR.ACC,
-        f"{name}_sqrt",
-        trt.UnaryOperation.SQRT,
-        add_trt,
-    )
-    # (x - E[x]) / sqrt((var + eps))
-    div_trt = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        f"{name}_div_trt",
-        trt.ElementWiseOperation.DIV,
-        sub_trt,
-        sqrt_trt,
-    )
-
-    assert gamma is not None
-    gamma_tensor = network.add_constant(gamma.shape, trt.Weights(np.ascontiguousarray(gamma)))  # type: ignore[attr-defined]
-    gamma_tensor.name = f"{name}_gamma"
-    assert beta is not None
-    beta_tensor = network.add_constant(gamma.shape, trt.Weights(np.ascontiguousarray(beta)))  # type: ignore[attr-defined]
-    beta_tensor.name = f"{name}_beta"
-    # y * gamma + beta
-    scale_layer = convert_binary_elementwise(
-        network,
-        target,
-        SourceIR.ACC,
-        f"{name}_scale",
-        trt.ElementWiseOperation.PROD,
-        div_trt,
-        gamma_tensor.get_output(0),
-    )
-    return convert_binary_elementwise(
+    return layer_norm(
         network,
         target,
         SourceIR.ACC,
         name,
-        trt.ElementWiseOperation.SUM,
-        scale_layer,
-        beta_tensor.get_output(0),
+        kwargs["input"],
+        kwargs["normalized_shape"],
+        kwargs["weight"],
+        kwargs["bias"],
+        kwargs["eps"],
     )
 
 

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -26,6 +26,7 @@
 from torch_tensorrt.fx.converters.impl.elementwise import fmod
 from torch_tensorrt.fx.converters.impl.elementwise import rsub
 from torch_tensorrt.fx.converters.impl.normalization import batch_norm
+from torch_tensorrt.fx.converters.impl.normalization import layer_norm
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -263,10 +264,30 @@ def aten_ops_leaky_relu(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-
     return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
 
 
+@tensorrt_converter(torch.ops.aten.layer_norm.default)
+def aten_ops_layernorm(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return layer_norm(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args[2],
+        args[3],
+        args[4],
+    )
+
+
 @tensorrt_converter(torch.ops.aten.linear)
 def aten_ops_linear(
     network: TRTNetwork,