int-reduce-int64-narrowing-overflow: fix sum/prod int64 accumulator narrowing

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]

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.

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",