diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index cbad1d2..b75bfe9 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -37,6 +37,7 @@ get_target, get_tensor_shape_at_index, get_tensor_type, + get_unnarrowed_output_element_type_from_node, prepare_compute_type_for_norm, process_expanded_indices, process_indices_with_transpose, @@ -2679,7 +2680,7 @@ def replace_sum_dim_intlist( x = _get_operand(values_map, node, 0) args = node.args - target_type = get_output_element_type_from_node(node) + target_type = get_unnarrowed_output_element_type_from_node(node) if x.type.element_type != target_type: x = coreai.cast(x, target_type) @@ -2735,7 +2736,7 @@ def replace_prod_default( values_map: dict[str, Value], node: fx.Node, loc: Location ) -> Value: x = _get_operand(values_map, node, 0) - target_type = get_output_element_type_from_node(node) + target_type = get_unnarrowed_output_element_type_from_node(node) if x.type.element_type != target_type: x = coreai.cast(x, target_type) all_dims = list(range(x.type.rank)) @@ -2746,7 +2747,7 @@ def replace_prod_dim_int( values_map: dict[str, Value], node: fx.Node, loc: Location ) -> Value: x = _get_operand(values_map, node, 0) - target_type = get_output_element_type_from_node(node) + target_type = get_unnarrowed_output_element_type_from_node(node) if x.type.element_type != target_type: x = coreai.cast(x, target_type) dim: int = node.args[1] diff --git a/coreai_torch/_utils.py b/coreai_torch/_utils.py index 7c59846..eb603ce 100644 --- a/coreai_torch/_utils.py +++ b/coreai_torch/_utils.py @@ -473,6 +473,39 @@ def get_output_element_type_from_node(node: fx.Node, index: int | None = None) - return TORCH_TO_COREAI_DTYPE[dtype]() +def get_unnarrowed_output_element_type_from_node( + node: fx.Node, index: int | None = None +) -> Type: + """Return the element type for a node's output, without int64/float64 narrowing. + + Like :func:`get_output_element_type_from_node`, but skips the + ``_NARROW_TORCH_DTYPE`` step. Intended for accumulator-style ops (e.g. + ``sum``/``prod`` reductions) whose torch semantics require accumulating at + the full promoted width (int64) to avoid overflowing early -- narrowing + the accumulator *before* the op runs silently changes its numeric result, + unlike narrowing a plain value cast, which is lossless-by-convention for + values already known to fit (or that the model owner accepts truncating). + coreai's own IR supports int64 (``si64``) as a first-class type, so + producing it here is not a runtime limitation. + """ + val = node.meta["val"] + if index is not None: + val = val[index] + + if isinstance(val, torch.Tensor): + dtype = val.dtype + elif isinstance(val, (float, torch.SymFloat)): + dtype = torch.float32 + elif isinstance(val, (int, torch.SymInt)): + dtype = torch.int32 + elif isinstance(val, (bool, torch.SymBool)): + dtype = torch.bool + else: + dtype = val.dtype # fall back to original behaviour + + return TORCH_TO_COREAI_DTYPE[dtype]() + + @dataclass class _StackedIndexInfo: """Stacked indices and permutations for gather_nd/scatter_nd. diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index 4c24e96..861ab4a 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -6417,6 +6417,104 @@ def forward(self, x: Tensor) -> Tensor: await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) +@pytest.mark.parametrize( + "x", + [ + # torch.sum(int32) promotes to int64: two INT32_MAX values sum to + # 4294967294, which fits in int64 but overflows (wraps) int32. + torch.tensor([2147483647, 2147483647], dtype=torch.int32), + torch.full((3, 4), 2147483647, dtype=torch.int32), + ], +) +async def test_sum_int_promotes_to_int64_no_overflow(x: Tensor) -> None: + """sum's int64 accumulator must not be narrowed to int32, which would + make the reduction silently wrap.""" + + class SumModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.sum(x) + + model = SumModel().eval() + torch_out = model(x) + assert torch_out.dtype == torch.int64 + assert torch_out.item() == x.numel() * 2147483647 + + await validate_numerical_output(model=model, x=x) + + +@pytest.mark.parametrize( + "x,dims,keepdim", + [ + # Explicit dim list + keepdim, exercising the branch of + # replace_sum_dim_intlist that doesn't reduce all dims via the + # empty-list path (see test_sum_int_promotes_to_int64_no_overflow + # above, which only covers torch.sum(x) with no dim/keepdim args). + (torch.full((2, 3), 2147483647, dtype=torch.int32), [0, 1], True), + (torch.full((2, 3), 2147483647, dtype=torch.int32), [0, 1], False), + ], +) +async def test_sum_dim_intlist_int_promotes_to_int64_no_overflow( + x: Tensor, dims: list[int], keepdim: bool +) -> None: + """Explicit dim_IntList + keepdim variant of the int64-promoted sum.""" + + class SumModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.sum(x, dim=dims, keepdim=keepdim) + + model = SumModel().eval() + torch_out = model(x) + assert torch_out.dtype == torch.int64 + + await validate_numerical_output(model=model, x=x) + + +@pytest.mark.parametrize( + "x", + [ + # torch.prod(int32) promotes to int64: 70000 * 70000 = 4900000000, + # which fits in int64 but overflows int32. + torch.tensor([70000, 70000], dtype=torch.int32), + # Negative-value overflow: promotion to int64 is sign-agnostic. + torch.tensor([-70000, 70000], dtype=torch.int32), + ], +) +async def test_prod_int_promotes_to_int64_no_overflow(x: Tensor) -> None: + """prod's int64 accumulator must not be narrowed to int32, which would + make the reduction silently wrap.""" + + class ProdModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.prod(x) + + model = ProdModel().eval() + torch_out = model(x) + assert torch_out.dtype == torch.int64 + assert torch_out.item() == x[0].item() * x[1].item() + + await validate_numerical_output(model=model, x=x) + + +@pytest.mark.parametrize( + "x,dim", + [ + (torch.tensor([[2147483647, 2147483647]], dtype=torch.int32), 1), + ], +) +async def test_prod_dim_int_promotes_to_int64_no_overflow(x: Tensor, dim: int) -> None: + """dim_int variant of the int64-promoted prod.""" + + class ProdModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.prod(x, dim=dim) + + model = ProdModel().eval() + torch_out = model(x) + assert torch_out.dtype == torch.int64 + + await validate_numerical_output(model=model, x=x) + + @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "x,dims",