Skip to content

Commit feffc76

Browse files
gokulkrishna98gokulkrishna98
authored
Merge branch 'main' into dev/gokul/layer_norm_test_fix
2 parents 5ef1bd1 + 46d15b7 commit feffc76

3 files changed

Lines changed: 107 additions & 27 deletions

File tree

‎coreai_torch/_aten_to_core.py‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -706,16 +706,27 @@ def replace_batch_norm(
706706
def batch_norm(
707707
input: Value, gamma: Value, beta: Value, mean: Value, variance: Value
708708
) -> Value:
709+
# Match ATen: compute in fp32 and narrow only the result. Downcasting
710+
# the fp32 params to fp16 loses precision and overflows to inf.
711+
input_compute, compute_type, needs_downcast = prepare_compute_type_for_norm(
712+
input, ele_type, loc
713+
)
714+
709715
eps_casted = coreai.constant(node.args[6])
710716

711-
eps_casted = coreai.cast(eps_casted, ele_type)
717+
def to_compute_type(v: Value) -> Value:
718+
if v.type.element_type == compute_type:
719+
return v
720+
return coreai.cast(v, compute_type)
721+
722+
eps_casted = to_compute_type(eps_casted)
712723

713724
expand_dims = [0] + list(
714725
range(2, x_type.rank)
715726
) # expand [C] to [1, C, 1, ...] for broadcasting
716727

717728
def expand_and_cast(v: Value) -> Value:
718-
return coreai.cast(coreai.expand_dims(v, expand_dims), ele_type)
729+
return to_compute_type(coreai.expand_dims(v, expand_dims))
719730

720731
weight, bias, running_mean, running_var = (
721732
expand_and_cast(v) for v in [gamma, beta, mean, variance]
@@ -724,12 +735,15 @@ def expand_and_cast(v: Value) -> Value:
724735
# ((x - mean) / sqrt(var + eps)) * weight + bias
725736
std = coreai.sqrt(coreai.broadcasting_add(running_var, eps_casted))
726737
normalized = coreai.broadcasting_divide(
727-
coreai.broadcasting_sub(input, running_mean), std
738+
coreai.broadcasting_sub(input_compute, running_mean), std
728739
)
729740

730-
return coreai.broadcasting_add(
741+
result = coreai.broadcasting_add(
731742
coreai.broadcasting_mul(normalized, weight), bias
732743
)
744+
if needs_downcast:
745+
result = coreai.cast(result, ele_type)
746+
return result
733747

734748
# In inference mode, outputs[1] and [2] are empty placeholder tensors (shape [0]).
735749
np_dtype = _get_coreai_to_numpy_dtype()[ele_type]

‎tests/ops/test_ops.py‎

Lines changed: 44 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -692,17 +692,9 @@ async def test_nan_propagation(self) -> None:
692692
await validate_numerical_output(model=model, y=y, x=x)
693693

694694

695-
@pytest.mark.parametrize(
696-
"x",
697-
[
698-
torch.rand(2, 3, 8, 8),
699-
torch.rand(2, 3, 8, 8, dtype=torch.float16), # fp16
700-
],
701-
)
702-
@pytest.mark.parametrize(
703-
"dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
704-
)
705-
async def test_batchnorm(x: Tensor, dynamic_dims: tuple[int]) -> None:
695+
class TestBatchNorm:
696+
"""Tests for native_batch_norm → coreai batch_norm composite."""
697+
706698
class BatchNormModel(nn.Module):
707699
def __init__(self) -> None:
708700
super().__init__()
@@ -711,10 +703,47 @@ def __init__(self) -> None:
711703
def forward(self, x: Tensor) -> Tensor:
712704
return self.bn(x)
713705

714-
model = BatchNormModel().eval()
715-
dim_names = {0: "batch", 1: "channels", 2: "height", 3: "width"}
716-
dynamic_shapes = make_dynamic_shapes(x={d: dim_names[d] for d in dynamic_dims})
717-
await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
706+
@pytest.mark.parametrize(
707+
"x",
708+
[
709+
torch.rand(2, 3, 8, 8),
710+
torch.rand(2, 3, 8, 8, dtype=torch.float16), # fp16
711+
],
712+
)
713+
@pytest.mark.parametrize(
714+
"dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
715+
)
716+
async def test_basic(self, x: Tensor, dynamic_dims: tuple[int]) -> None:
717+
model = self.BatchNormModel().eval()
718+
dim_names = {0: "batch", 1: "channels", 2: "height", 3: "width"}
719+
dynamic_shapes = make_dynamic_shapes(x={d: dim_names[d] for d in dynamic_dims})
720+
await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
721+
722+
@pytest.mark.parametrize(
723+
"running_var, running_mean, weight",
724+
[
725+
# running_var past the fp16 max (65504) becomes inf when the params are
726+
# downcast before the sqrt, which silently zeroes the whole output.
727+
([1e6, 2e5, 9e4], [0.0, 1.0, -2.0], [2000.0, 1000.0, 500.0]),
728+
# Tiny variances make eps and the param mantissas precision-critical.
729+
([1e-6, 4e-5, 2.0], [0.5, -0.25, 100.0], [1.0, 2.0, 0.5]),
730+
],
731+
)
732+
async def test_fp16_input_fp32_params(
733+
self,
734+
running_var: list[float],
735+
running_mean: list[float],
736+
weight: list[float],
737+
) -> None:
738+
"""An fp16 activation with fp32 params must be computed in fp32."""
739+
model = self.BatchNormModel().eval()
740+
model.bn.running_var.data = torch.tensor(running_var)
741+
model.bn.running_mean.data = torch.tensor(running_mean)
742+
model.bn.weight.data = torch.tensor(weight)
743+
model.bn.bias.data = torch.tensor([0.0, 1.0, -1.0])
744+
745+
x = (torch.rand(2, 3, 4, 4) * 2 - 1).half()
746+
await validate_numerical_output(model=model, x=x)
718747

719748

720749
@pytest.mark.parametrize("dynamic", [False, True])

‎tests/ops/test_ops_ir.py‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -577,16 +577,16 @@ def forward(self, x: Tensor) -> Tensor:
577577

578578

579579
class TestBatchNormIR:
580-
def test_static_composite(self) -> None:
581-
class BatchNormModel(nn.Module):
582-
def __init__(self):
583-
super().__init__()
584-
self.bn = nn.BatchNorm2d(3)
580+
class BatchNormModel(nn.Module):
581+
def __init__(self):
582+
super().__init__()
583+
self.bn = nn.BatchNorm2d(3)
585584

586-
def forward(self, x: Tensor) -> Tensor:
587-
return self.bn(x)
585+
def forward(self, x: Tensor) -> Tensor:
586+
return self.bn(x)
588587

589-
ir = get_ir(BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4))
588+
def test_static_composite(self) -> None:
589+
ir = get_ir(self.BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4))
590590
filecheck_pattern(
591591
ir,
592592
check_file="""
@@ -616,6 +616,43 @@ def forward(self, x: Tensor) -> Tensor:
616616
""",
617617
)
618618

619+
def test_fp16_input_computes_in_fp32(self) -> None:
620+
"""fp16 input is promoted, the fp32 params are used as-is, and only the
621+
result is narrowed back to fp16."""
622+
ir = get_ir(
623+
self.BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4, dtype=torch.float16)
624+
)
625+
filecheck_pattern(
626+
ir,
627+
check_file="""
628+
// CHECK-LABEL: module {
629+
// 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"} {
630+
// CHECK-NEXT: %[[V0:.*]] = coreai.constant dense<[1, 3, 1, 1]> : tensor<4xui32>
631+
// CHECK-NEXT: %[[V1:.*]] = coreai.constant dense<9.99999974E-6> : tensor<f32>
632+
// CHECK-NEXT: %[[V2:.*]] = coreai.cast %[[ARG0]] : tensor<1x3x4x4xf16> to tensor<1x3x4x4xf32>
633+
// CHECK-NEXT: %[[V3:.*]] = coreai.reshape %[[ARG1]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
634+
// CHECK-NEXT: %[[V4:.*]] = coreai.reshape %[[ARG2]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
635+
// CHECK-NEXT: %[[V5:.*]] = coreai.reshape %[[ARG3]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
636+
// CHECK-NEXT: %[[V6:.*]] = coreai.decomposable.broadcasting_add %[[ARG4]], %[[V1]] : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
637+
// CHECK-NEXT: %[[V7:.*]] = coreai.reshape %[[V6]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
638+
// CHECK-NEXT: %[[V8:.*]] = coreai.sqrt %[[V7]] : tensor<1x3x1x1xf32> -> tensor<1x3x1x1xf32>
639+
// CHECK-NEXT: %[[V9:.*]] = coreai.decomposable.broadcasting_sub %[[V2]], %[[V5]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
640+
// CHECK-NEXT: %[[V10:.*]] = coreai.decomposable.broadcasting_divide %[[V9]], %[[V8]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
641+
// CHECK-NEXT: %[[V11:.*]] = coreai.decomposable.broadcasting_mul %[[V10]], %[[V3]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
642+
// CHECK-NEXT: %[[V12:.*]] = coreai.decomposable.broadcasting_add %[[V11]], %[[V4]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
643+
// CHECK-NEXT: %[[V13:.*]] = coreai.cast %[[V12]] : tensor<1x3x4x4xf32> to tensor<1x3x4x4xf16>
644+
// CHECK-NEXT: coreai.output %[[V13]] : tensor<1x3x4x4xf16>
645+
// CHECK-NEXT: }
646+
// CHECK-NEXT: coreai.graph @main(%[[ARG0]]: tensor<1x3x4x4xf16> {coreai.name = "x"}) -> (tensor<1x3x4x4xf16> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
647+
// CHECK-NEXT: %[[V0]] = coreai.constant dense<1.000000e+00> : tensor<3xf32>
648+
// CHECK-NEXT: %[[V1]] = coreai.constant dense<0.000000e+00> : tensor<3xf32>
649+
// CHECK-NEXT: %[[V2]] = coreai.invoke @batch_norm_{{.*}}(%[[ARG0]], %[[V0]], %[[V1]], %[[V1]], %[[V0]]) : (tensor<1x3x4x4xf16>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>) -> tensor<1x3x4x4xf16>
650+
// CHECK-NEXT: coreai.output %[[V2]] : tensor<1x3x4x4xf16>
651+
// CHECK-NEXT: }
652+
// CHECK-NEXT: }
653+
""",
654+
)
655+
619656

620657
class TestToCopyIR:
621658
def test_cast_f32_to_f16(self) -> None:

0 commit comments

Comments
 (0)