Compute fp16 batch norm in fp32 instead of downcasting the params

`nn.BatchNorm2d` keeps gamma/beta/running_mean/running_var in fp32 even
when the activation is fp16. ATen handles that mixed-precision case by
promoting the input to fp32, using the params and eps at full precision,
and narrowing only the final result (`native_batch_norm_helper` in
torch/_decomp/decompositions.py).

The composite body did the opposite: it cast the fp32 params and eps down
to the input element type and ran every op in fp16. That loses mantissa
bits on the params (eps 1e-5 becomes 1.00136e-05), and any channel whose
running_var or running_mean exceeds the fp16 max of 65504 becomes inf --
sqrt(inf) then divides the activation to a silent zero.

Use the existing `prepare_compute_type_for_norm` helper, as layer_norm
and group_norm already do, so the fp16 input is promoted, the arithmetic
runs in fp32, and there is a single narrowing cast before the output.
fp32 inputs are unaffected apart from dropping now-redundant f32->f32
casts, which the optimizer already folded.

Tests: a numerical case with large and tiny running_var (max abs error
4.945 and 0.5 against eager before the fix) and an IR case pinning the
promote/compute/narrow shape of the fp16 composite body.

Author: gokulkrishna98

coreai_torch/_aten_to_core.py

@@ -706,16 +706,27 @@ def replace_batch_norm(
     def batch_norm(
         input: Value, gamma: Value, beta: Value, mean: Value, variance: Value
     ) -> Value:
+        # Match ATen: compute in fp32 and narrow only the result. Downcasting
+        # the fp32 params to fp16 loses precision and overflows to inf.
+        input_compute, compute_type, needs_downcast = prepare_compute_type_for_norm(
+            input, ele_type, loc
+        )
+
         eps_casted = coreai.constant(node.args[6])
 
-        eps_casted = coreai.cast(eps_casted, ele_type)
+        def to_compute_type(v: Value) -> Value:
+            if v.type.element_type == compute_type:
+                return v
+            return coreai.cast(v, compute_type)
+
+        eps_casted = to_compute_type(eps_casted)
 
         expand_dims = [0] + list(
             range(2, x_type.rank)
         )  # expand [C] to [1, C, 1, ...] for broadcasting
 
         def expand_and_cast(v: Value) -> Value:
-            return coreai.cast(coreai.expand_dims(v, expand_dims), ele_type)
+            return to_compute_type(coreai.expand_dims(v, expand_dims))
 
         weight, bias, running_mean, running_var = (
             expand_and_cast(v) for v in [gamma, beta, mean, variance]
@@ -724,12 +735,15 @@ def expand_and_cast(v: Value) -> Value:
         # ((x - mean) / sqrt(var + eps)) * weight + bias
         std = coreai.sqrt(coreai.broadcasting_add(running_var, eps_casted))
         normalized = coreai.broadcasting_divide(
-            coreai.broadcasting_sub(input, running_mean), std
+            coreai.broadcasting_sub(input_compute, running_mean), std
         )
 
-        return coreai.broadcasting_add(
+        result = coreai.broadcasting_add(
             coreai.broadcasting_mul(normalized, weight), bias
         )
+        if needs_downcast:
+            result = coreai.cast(result, ele_type)
+        return result
 
     # In inference mode, outputs[1] and [2] are empty placeholder tensors (shape [0]).
     np_dtype = _get_coreai_to_numpy_dtype()[ele_type]

tests/ops/test_ops.py

@@ -692,17 +692,9 @@ async def test_nan_propagation(self) -> None:
         await validate_numerical_output(model=model, y=y, x=x)
 
 
-@pytest.mark.parametrize(
-    "x",
-    [
-        torch.rand(2, 3, 8, 8),
-        torch.rand(2, 3, 8, 8, dtype=torch.float16),  # fp16
-    ],
-)
-@pytest.mark.parametrize(
-    "dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
-)
-async def test_batchnorm(x: Tensor, dynamic_dims: tuple[int]) -> None:
+class TestBatchNorm:
+    """Tests for native_batch_norm → coreai batch_norm composite."""
+
     class BatchNormModel(nn.Module):
         def __init__(self) -> None:
             super().__init__()
@@ -711,10 +703,47 @@ def __init__(self) -> None:
         def forward(self, x: Tensor) -> Tensor:
             return self.bn(x)
 
-    model = BatchNormModel().eval()
-    dim_names = {0: "batch", 1: "channels", 2: "height", 3: "width"}
-    dynamic_shapes = make_dynamic_shapes(x={d: dim_names[d] for d in dynamic_dims})
-    await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
+    @pytest.mark.parametrize(
+        "x",
+        [
+            torch.rand(2, 3, 8, 8),
+            torch.rand(2, 3, 8, 8, dtype=torch.float16),  # fp16
+        ],
+    )
+    @pytest.mark.parametrize(
+        "dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
+    )
+    async def test_basic(self, x: Tensor, dynamic_dims: tuple[int]) -> None:
+        model = self.BatchNormModel().eval()
+        dim_names = {0: "batch", 1: "channels", 2: "height", 3: "width"}
+        dynamic_shapes = make_dynamic_shapes(x={d: dim_names[d] for d in dynamic_dims})
+        await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
+
+    @pytest.mark.parametrize(
+        "running_var, running_mean, weight",
+        [
+            # running_var past the fp16 max (65504) becomes inf when the params are
+            # downcast before the sqrt, which silently zeroes the whole output.
+            ([1e6, 2e5, 9e4], [0.0, 1.0, -2.0], [2000.0, 1000.0, 500.0]),
+            # Tiny variances make eps and the param mantissas precision-critical.
+            ([1e-6, 4e-5, 2.0], [0.5, -0.25, 100.0], [1.0, 2.0, 0.5]),
+        ],
+    )
+    async def test_fp16_input_fp32_params(
+        self,
+        running_var: list[float],
+        running_mean: list[float],
+        weight: list[float],
+    ) -> None:
+        """An fp16 activation with fp32 params must be computed in fp32."""
+        model = self.BatchNormModel().eval()
+        model.bn.running_var.data = torch.tensor(running_var)
+        model.bn.running_mean.data = torch.tensor(running_mean)
+        model.bn.weight.data = torch.tensor(weight)
+        model.bn.bias.data = torch.tensor([0.0, 1.0, -1.0])
+
+        x = (torch.rand(2, 3, 4, 4) * 2 - 1).half()
+        await validate_numerical_output(model=model, x=x)
 
 
 @pytest.mark.parametrize("dynamic", [False, True])

tests/ops/test_ops_ir.py

@@ -577,16 +577,16 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class TestBatchNormIR:
-    def test_static_composite(self) -> None:
-        class BatchNormModel(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.bn = nn.BatchNorm2d(3)
+    class BatchNormModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.bn = nn.BatchNorm2d(3)
 
-            def forward(self, x: Tensor) -> Tensor:
-                return self.bn(x)
+        def forward(self, x: Tensor) -> Tensor:
+            return self.bn(x)
 
-        ir = get_ir(BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4))
+    def test_static_composite(self) -> None:
+        ir = get_ir(self.BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4))
         filecheck_pattern(
             ir,
             check_file="""
@@ -616,6 +616,43 @@ def forward(self, x: Tensor) -> Tensor:
             """,
         )
 
+    def test_fp16_input_computes_in_fp32(self) -> None:
+        """fp16 input is promoted, the fp32 params are used as-is, and only the
+        result is narrowed back to fp16."""
+        ir = get_ir(
+            self.BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4, dtype=torch.float16)
+        )
+        filecheck_pattern(
+            ir,
+            check_file="""
+                // CHECK-LABEL: module {
+                // CHECK-NEXT:   coreai.graph private noinline @batch_norm_{{.*}}(%[[ARG0:.*]]: tensor<1x3x4x4xf16> {coreai.name = "input"}, %[[ARG1:.*]]: tensor<3xf32> {coreai.name = "gamma"}, %[[ARG2:.*]]: tensor<3xf32> {coreai.name = "beta"}, %[[ARG3:.*]]: tensor<3xf32> {coreai.name = "mean"}, %[[ARG4:.*]]: tensor<3xf32> {coreai.name = "variance"}) -> tensor<1x3x4x4xf16> attributes {__coreai_pure__, composite_decl = #coreai.composite_declaration<"batch_norm" = {input_names = ["input", "gamma", "beta", "mean", "variance"], op_attrs = {eps = 9.99999974E-6 : f32, version = 1 : si64}, output_names = ["output"]}>, template_op = "batch_norm"} {
+                // CHECK-NEXT:     %[[V0:.*]] = coreai.constant dense<[1, 3, 1, 1]> : tensor<4xui32>
+                // CHECK-NEXT:     %[[V1:.*]] = coreai.constant dense<9.99999974E-6> : tensor<f32>
+                // CHECK-NEXT:     %[[V2:.*]] = coreai.cast %[[ARG0]] : tensor<1x3x4x4xf16> to tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V3:.*]] = coreai.reshape %[[ARG1]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V4:.*]] = coreai.reshape %[[ARG2]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V5:.*]] = coreai.reshape %[[ARG3]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V6:.*]] = coreai.decomposable.broadcasting_add %[[ARG4]], %[[V1]] : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+                // CHECK-NEXT:     %[[V7:.*]] = coreai.reshape %[[V6]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V8:.*]] = coreai.sqrt %[[V7]] : tensor<1x3x1x1xf32> -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V9:.*]] = coreai.decomposable.broadcasting_sub %[[V2]], %[[V5]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V10:.*]] = coreai.decomposable.broadcasting_divide %[[V9]], %[[V8]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V11:.*]] = coreai.decomposable.broadcasting_mul %[[V10]], %[[V3]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V12:.*]] = coreai.decomposable.broadcasting_add %[[V11]], %[[V4]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V13:.*]] = coreai.cast %[[V12]] : tensor<1x3x4x4xf32> to tensor<1x3x4x4xf16>
+                // CHECK-NEXT:     coreai.output %[[V13]] : tensor<1x3x4x4xf16>
+                // CHECK-NEXT:   }
+                // CHECK-NEXT:   coreai.graph @main(%[[ARG0]]: tensor<1x3x4x4xf16> {coreai.name = "x"}) -> (tensor<1x3x4x4xf16> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
+                // CHECK-NEXT:     %[[V0]] = coreai.constant dense<1.000000e+00> : tensor<3xf32>
+                // CHECK-NEXT:     %[[V1]] = coreai.constant dense<0.000000e+00> : tensor<3xf32>
+                // CHECK-NEXT:     %[[V2]] = coreai.invoke @batch_norm_{{.*}}(%[[ARG0]], %[[V0]], %[[V1]], %[[V1]], %[[V0]])  : (tensor<1x3x4x4xf16>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>) -> tensor<1x3x4x4xf16>
+                // CHECK-NEXT:     coreai.output %[[V2]] : tensor<1x3x4x4xf16>
+                // CHECK-NEXT:   }
+                // CHECK-NEXT: }
+            """,
+        )
+
 
 class TestToCopyIR:
     def test_cast_f32_to_f16(self) -> None: