Fix aten.min.dim argmin indices at dtype-extremal minima

replace_min_dim computed argmin as argmax(x * -1), which is wrong at integer
dtype extremes: uint8 min 0 negates to 0 (not the range top), and int8 min
-128 overflows on negation. Reverse order with the bitwise complement
~x = x ^ -1 (broadcasting_bitwise_xor) for integer dtypes -- a strictly
decreasing, overflow-free bijection that preserves first-index tie-breaking;
keep x * -1 for float. Adds dtype-extremal and tie-break coverage.

Author: arames

coreai_torch/_aten_to_core.py

@@ -2303,19 +2303,35 @@ def replace_min_dim(
 ) -> OpResultList:
     """Computes min along a specific dimension, returning (values, indices).
 
-    Argmin is computed as argmax(-x), since Core AI has argmax but not argmin.
+    Core AI has argmax but not argmin, so argmin is computed as argmax of an
+    order-reversed copy of ``x``. For integer (non-bool) types the reversal is
+    the bitwise complement ``~x`` (``x ^ -1``), which is a strictly decreasing
+    bijection over the full range of every signed/unsigned integer dtype and
+    never overflows -- unlike the ``-x`` negation trick, which is wrong at
+    dtype-extremal minima (e.g. uint8 ``0`` or int8 ``-128``). For float (and
+    bool) types the ``-x`` negation is exact, so it is kept.
     """
     x = _get_operand(values_map, node, 0)
     dim = node.args[1]
     keepdim = len(node.args) >= 3 and bool(node.args[2])
     dim = dim + x.type.rank if dim < 0 else dim
 
+    element_type = x.type.element_type
+    is_integer = isinstance(element_type, IntegerType)
+    is_bool = element_type == IntegerType.get_signless(1)
+    if is_integer and not is_bool:
+        # ~x == x ^ ALL_ONES; -1 has all bits set in two's complement.
+        reversed_x = coreai.broadcasting_bitwise_xor(
+            x, coreai.constant(-1, dtype=element_type)
+        )
+    else:
+        reversed_x = coreai.broadcasting_mul(
+            x, coreai.constant(-1, dtype=element_type)
+        )
+
     min_values = coreai.reduce_min(x, [dim])
     argmin_indices = coreai.cast(
-        coreai.argmax(
-            coreai.broadcasting_mul(x, coreai.constant(-1, dtype=x.type.element_type)),
-            dim,
-        ),
+        coreai.argmax(reversed_x, dim),
         np.int32,
     )
 

tests/ops/test_ops.py

@@ -3003,6 +3003,106 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
     await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
 
 
+class TestMinDimArgminDtypeExtremes:
+    """min.dim argmin indices must be correct even at dtype-extremal minima.
+
+    Core AI has argmax but not argmin, so ``replace_min_dim`` derives the
+    argmin index by reversing the order of ``x`` and taking ``argmax``. The
+    reversal must be exact at the extremes of the input dtype's range:
+
+      * uint8 minimum ``0`` — a plain ``-x`` (i.e. ``x * -1``) negation maps
+        ``0`` to ``0`` rather than the top of the unsigned range, so the true
+        minimum stops being the argmax.
+      * int8 minimum ``-128`` — ``-128 * -1`` overflows the signed range and
+        (under wrapping arithmetic) lands back on ``-128``.
+
+    Both were previously wrong on interpreter/cpu/gpu. The fix reverses integer
+    inputs with the bitwise complement ``~x`` (a strictly decreasing bijection
+    over the whole integer range that never overflows), leaving the exact float
+    negation path untouched. Ground truth is ``torch.min(x, dim=...)``.
+    """
+
+    class MinDimModel(nn.Module):
+        def __init__(self, dim: int, keepdim: bool) -> None:
+            super().__init__()
+            self.dim = dim
+            self.keepdim = keepdim
+
+        def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
+            return torch.min(x, dim=self.dim, keepdim=self.keepdim)
+
+    @pytest.mark.parametrize("keepdim", [False, True])
+    @pytest.mark.parametrize(
+        "values,dtype",
+        [
+            # uint8: true minimum is 0 (bottom of the unsigned range).
+            ([3, 1, 0, 5], torch.uint8),
+            # Signed dtypes at their exact minimum (INT_MIN), which overflows
+            # under the old ``x * -1`` negation trick.
+            ([3, 1, -128, 5], torch.int8),
+            ([3, 1, -(2**15), 5], torch.int16),
+            ([3, 1, -(2**31), 5], torch.int32),
+            # NOTE: coreai-torch narrows int64 inputs to int32 at the graph
+            # boundary, so the exercisable extreme for the int64 dtype *path*
+            # is int32's minimum (a true int64 min of -2**63 would be
+            # truncated at that boundary, unrelated to this argmin fix).
+            ([3, 1, -(2**31), 5], torch.int64),
+        ],
+    )
+    async def test_dtype_min_value(
+        self, values: list[int], dtype: torch.dtype, keepdim: bool
+    ) -> None:
+        x = torch.tensor(values, dtype=dtype)
+        model = self.MinDimModel(dim=0, keepdim=keepdim).eval()
+        await validate_numerical_output(model=model, x=x)
+
+    @pytest.mark.parametrize("keepdim", [False, True])
+    @pytest.mark.parametrize(
+        "values,dtype",
+        [
+            # Duplicate minima at the dtype extreme: torch returns the FIRST
+            # index (index 0 here), which the reversal must preserve.
+            ([-128, -128, 5], torch.int8),
+            ([0, 0, 5], torch.uint8),
+        ],
+    )
+    async def test_duplicate_minima_first_index(
+        self, values: list[int], dtype: torch.dtype, keepdim: bool
+    ) -> None:
+        x = torch.tensor(values, dtype=dtype)
+        # Sanity-check the ground-truth tie-break we are asserting against.
+        assert int(torch.min(x, dim=0).indices) == 0
+        model = self.MinDimModel(dim=0, keepdim=keepdim).eval()
+        await validate_numerical_output(model=model, x=x)
+
+    @pytest.mark.parametrize("keepdim", [False, True])
+    async def test_float_path_untouched(self, keepdim: bool) -> None:
+        # Float negation is exact/no-overflow; confirm it still validates.
+        x = torch.tensor([3.0, 1.0, -5.0, 2.0], dtype=torch.float32)
+        model = self.MinDimModel(dim=0, keepdim=keepdim).eval()
+        await validate_numerical_output(model=model, x=x)
+
+    @pytest.mark.parametrize("keepdim", [False, True])
+    async def test_bool_input_reduce_min_unsupported(self, keepdim: bool) -> None:
+        # torch.min(dim) on a bool tensor is legal in torch, and the argmin fix
+        # routes bool down the (mul-based) reversal branch — NOT the integer
+        # bitwise-complement branch — since ~x on i1 is not what we want. That
+        # branch is fine, but bool min.dim is independently unsupported at the
+        # *values* path: coreai.reduce_min rejects an i1 operand (unrelated to
+        # this argmin fix, and true on origin/main too). Pin that limitation so
+        # a future reduce_min-on-bool change is noticed here; if reduce_min
+        # gains i1 support, this test should flip to validate_numerical_output.
+        x = torch.tensor([True, False, True, False])
+        assert x.dtype == torch.bool
+        # First False is the minimum; torch returns its first index.
+        assert int(torch.min(x, dim=0).indices) == 1
+        model = self.MinDimModel(dim=0, keepdim=keepdim).eval()
+        # coreai.reduce_min's MLIR verifier rejects the i1 operand; save_asset
+        # surfaces it as a RuntimeError wrapping the underlying MLIRError.
+        with pytest.raises(RuntimeError):
+            await validate_numerical_output(model=model, x=x)
+
+
 @pytest.mark.parametrize("dynamic", [False, True])
 @pytest.mark.parametrize(
     "shape,dim,keepdim,dtype",