Address review

- plumb `use_aten` through the driver
- use bitmask-style values for the mode enum

Signed-off-by: Alex Zinenko <git@ozinenko.com>

Author: ftynse

iree/turbine/kernel/boo/op_exports/layer_norm.py

@@ -17,11 +17,11 @@
 class Mode(ModeBase, IntEnum):
     """Mode selector for layer normalization, with each gradient being its own mode."""
 
-    FORWARD = 0
-    INPUT_BACKWARD = 1
-    WEIGHT_BACKWARD = 2
-    BIAS_BACKWARD = 3
-    FULL_BACKWARD = 4
+    FORWARD = 1
+    INPUT_BACKWARD = 2
+    WEIGHT_BACKWARD = 4
+    BIAS_BACKWARD = 8
+    FULL_BACKWARD = INPUT_BACKWARD | WEIGHT_BACKWARD | BIAS_BACKWARD
 
 
 class LayerNormSignature(OpSignature):
@@ -46,6 +46,7 @@ def __init__(
         dtype=torch.bfloat16,
         mode: str | Mode = Mode.FORWARD,
         forwarded_args_dtype: torch.dtype | None = None,
+        use_aten: bool = True,
     ):
         if (
             len(normalized_shape) > len(input_shape)
@@ -63,6 +64,7 @@ def __init__(
         self.dtype = dtype
         self.mode = Mode.parse(mode)
         self.forwarded_args_dtype = forwarded_args_dtype or dtype
+        self.use_aten = use_aten
 
     @property
     def output_shape(self) -> list[int]:
@@ -124,6 +126,7 @@ def func_name(self) -> str:
             "x".join(str(i) for i in self.input_shape),
             "w" if self.elementwise_affine is not None else "",
             "b" if self.bias is not None else "",
+            "aten" if self.use_aten else "",
         ]
         return "_".join(name_items)
 
@@ -183,6 +186,7 @@ def as_init_kwargs(self) -> dict[str, Any]:
             "dtype": self.dtype,
             "mode": self.Mode,
             "forwarded_args_dtype": self.forwarded_args_dtype,
+            "use_aten": self.use_aten,
         }
 
     def get_output_size(self) -> int:
@@ -395,9 +399,9 @@ class LayerNormBackwardFull(torch.nn.Module):
     weights, and bias of the layer normalization given the gradient of its
     output."""
 
-    def __init__(self, signature: LayerNormSignature, *, use_aten=True):
+    def __init__(self, signature: LayerNormSignature):
         super().__init__()
-        self.use_aten = use_aten
+        self.use_aten = signature.use_aten
         self.normalized_shape = signature.normalized_shape
         self.need_bias = signature.bias
         self.need_weight = signature.elementwise_affine
@@ -438,12 +442,14 @@ def forward(
         # Recompute norm instead of saving it. Judging by the signature, this is the same
         # decision as ATen.
         norm = (input - mean) * rstd
+        # norm = norm.to(dtype=input.dtype)
         dnorm = grad_output * weight if weight is not None else grad_output
         dx = (
             dnorm
             - dnorm.mean(dim=self.normalized_dim, keepdim=True)
             - norm * (dnorm * norm).mean(dim=self.normalized_dim, keepdim=True)
         ) * rstd
+        # dx = dx.to(dtype=input.dtype)
         dw = None
         if self.need_weight:
             dw = (grad_output * norm).sum(self.keep_dim)
@@ -489,19 +495,10 @@ def get_signature(args: argparse.Namespace) -> LayerNormSignature:
         ), "Can only normalize one trailing dimension for now (MIOpen limitation)."
         normalized_shape = shape[args.normalized_dim :]
 
-        match args.forw:
-            case 1:
-                mode = Mode.FORWARD
-            case 2:
-                mode = Mode.INPUT_BACKWARD
-            case 3:
-                mode = Mode.WEIGHT_BACKWARD
-            case 4:
-                mode = Mode.BIAS_BACKWARD
-            case 5:
-                mode = Mode.FULL_BACKWARD
-            case _:
-                raise ValueError(f"Unsupported mode {args.forw}.")
+        try:
+            mode = Mode(args.forw)
+        except Exception as e:
+            raise ValueError(f"Unsupported mode {args.forw}.") from e
 
         return LayerNormSignature(
             input_shape=shape,
@@ -511,6 +508,7 @@ def get_signature(args: argparse.Namespace) -> LayerNormSignature:
             bias=True,
             dtype=_DTypeCommandDispatcher.get_dtype(args.command),
             mode=mode,
+            use_aten=args.use_aten,
         )
 
     def get_miopen_parser() -> argparse.ArgumentParser:
@@ -519,7 +517,11 @@ def get_miopen_parser() -> argparse.ArgumentParser:
             "command", default="layernorm", choices=_DTypeCommandDispatcher.choices()
         )
         parser.add_argument(
-            "--forw", "-F", type=int, default=1, help="Run only forward LayerNorm"
+            "--forw",
+            "-F",
+            type=int,
+            default=1,
+            help="Kind of kernel to run, not compatible with MIOpen (1 forward, 2 backward input, 4 backward weight, 8 backward bias, 14 full backward)",
         )
         parser.add_argument(
             "--input",
@@ -539,6 +541,12 @@ def get_miopen_parser() -> argparse.ArgumentParser:
         parser.add_argument(
             "--normalized_dim", "-o", type=int, default=3, help="Normalized dim"
         )
+        parser.add_argument(
+            "--use-aten",
+            type=bool,
+            default=True,
+            help="Use core ATen op instead of a manual implementation",
+        )
         return parser
 
     @classmethod

tests/kernel/boo/op_exports/layer_norm_backward_impl_test.py

@@ -114,14 +114,16 @@ def test_layer_norm_impl(
 @pytest.mark.parametrize(
     "elementwise_affine_bias", [(False, False), (True, False), (True, True)]
 )
+@pytest.mark.parametrize("use_aten", [True, False])
 def test_layer_norm_combined_impl(
     input_shape: tuple[int, ...],
     device: str,
     dtype: torch.dtype,
     elementwise_affine_bias: tuple[bool, bool],
+    use_aten: bool,
 ):
     # Account for ATen weirdness on GPU.
-    if device == "cuda" and dtype == torch.bfloat16:
+    if device == "cuda" and dtype == torch.bfloat16 and use_aten:
         forwarded_args_dtype = torch.float32
     else:
         forwarded_args_dtype = dtype
@@ -134,6 +136,7 @@ def test_layer_norm_combined_impl(
         "bias": bias,
         "dtype": dtype,
         "forwarded_args_dtype": forwarded_args_dtype,
+        "use_aten": use_aten,
     }
     fwd_sig = LayerNormSignature(**kwargs)
     args = fwd_sig.get_sample_args(seed=1)
@@ -150,15 +153,14 @@ def test_layer_norm_combined_impl(
 
     main_result = fwd_results[fwd_sig.main_result_index]
     main_result.retain_grad()
-    # TODO: this is not a good loss function (#1021).
-    loss = main_result.sum()
+    loss = main_result.mean() / main_result.numel()
     loss.backward(retain_graph=True)
 
     bwd_input_args = bwd_sig.arrange_backward_launch_args(args, fwd_results)
     grads = tuple(x for x in bwd(main_result.grad, *bwd_input_args) if x is not None)
 
-    rtol = 1e-4
-    atol = 1e-4
+    rtol = 1e-6
+    atol = 1e-6
     assert len(grads) == len(args)
     results = [
         torch.allclose(arg.grad, grad, rtol=rtol, atol=atol)