Fix full reduction reading outside non-contiguous input views

src/flag_gems/ops/amax.py

@@ -99,6 +99,7 @@ def amax_kernel(
 def amax(inp, dim=None, keepdim=False):
     logger.debug("GEMS AMAX")
     if dim is None or len(dim) == 0:
+        inp = inp.contiguous()
         M = inp.numel()
         block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
         mid_size = triton.cdiv(M, block_size)

src/flag_gems/ops/amin.py

@@ -106,6 +106,7 @@ def amin(inp, dim=None, keepdim=False):
         torch.bfloat16,
     ), "amin only supports float dtypes"
     if dim is None or len(dim) == 0:
+        inp = inp.contiguous()
         M = inp.numel()
         block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
         mid_size = triton.cdiv(M, block_size)
@@ -165,7 +166,9 @@ def amin_(inp, dim=None, keepdim=False):
     if isinstance(dim, int):
         dim = [dim]
     if dim is None or len(dim) == 0:
-        M = inp.numel()
+        # Read from a contiguous copy but write back to the original tensor.
+        inp_flat = inp.contiguous()
+        M = inp_flat.numel()
         block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
         mid_size = triton.cdiv(M, block_size)
         block_mid = triton.next_power_of_2(mid_size)
@@ -180,7 +183,7 @@ def amin_(inp, dim=None, keepdim=False):
             out = torch.empty(shape, dtype=dtype, device=inp.device)
         with torch_device_fn.device(inp.device):
             amin_kernel_1[(mid_size, 1)](
-                inp,
+                inp_flat,
                 mid,
                 M,
                 block_size,

src/flag_gems/ops/argmax.py

@@ -195,6 +195,7 @@ def argmax_kernel_inner(
 def argmax(inp, dim=None, keepdim=False, *, dtype=None):
     logger.debug("GEMS ARGMAX")
     if dim is None:
+        inp = inp.contiguous()
         M = inp.numel()
         if dtype is None:
             dtype = inp.dtype

src/flag_gems/ops/argmin.py

@@ -223,6 +223,7 @@ def argmin_kernel(
 def argmin(inp, dim=None, keepdim=False, *, dtype=None):
     logger.debug("GEMS ARGMIN")
     if dim is None:
+        inp = inp.contiguous()
         M = inp.numel()
         if dtype is None:
             dtype = inp.dtype

src/flag_gems/ops/min.py

@@ -112,6 +112,7 @@ def min_kernel(
 
 def min(inp):
     logger.debug("GEMS MIN")
+    inp = inp.contiguous()
     M = inp.numel()
     block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
     mid_size = triton.cdiv(M, block_size)

src/flag_gems/ops/prod.py

@@ -67,6 +67,7 @@ def prod(inp, *, dtype=None):
     if dtype is None:
         dtype = inp.dtype
 
+    inp = inp.contiguous()
     M = inp.numel()
     block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
     mid_size = triton.cdiv(M, block_size)

src/flag_gems/runtime/backend/_ascend/ops/amax.py

@@ -103,6 +103,7 @@ def amax_kernel(
 def amax(inp, dim=None, keepdim=False):
     logger.debug("GEMS_ASCEND AMAX")
     if dim is None or len(dim) == 0:
+        inp = inp.contiguous()
         M = inp.numel()
         block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
         mid_size = triton.cdiv(M, block_size)

src/flag_gems/runtime/backend/_ascend/ops/argmax.py

@@ -109,16 +109,18 @@ def argmax_kernel(
 def argmax(inp, dim=None, keepdim=False, *, dtype=None):
     logger.debug("GEMS_ASCEND ARGMAX")
     if dim is None:
+        inp = inp.contiguous()
         M = inp.numel()
         if dtype is None:
             dtype = inp.dtype
-        block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
+        block_size = min(1024, triton.next_power_of_2(math.ceil(math.sqrt(M))))
         mid_size = triton.cdiv(M, block_size)
         block_mid = triton.next_power_of_2(mid_size)
 
         mid_value = torch.empty((mid_size,), dtype=dtype, device=inp.device)
         mid_index = torch.empty((mid_size,), dtype=torch.int64, device=inp.device)
-        out = torch.empty([], dtype=torch.int64, device=inp.device)
+        out_shape = [1] * inp.dim() if keepdim else []
+        out = torch.empty(out_shape, dtype=torch.int64, device=inp.device)
 
         with torch_device_fn.device(inp.device):
             argmax_kernel_1[(mid_size, 1, 1)](

src/flag_gems/runtime/backend/_ascend/ops/argmin.py

@@ -125,6 +125,7 @@ def argmin(inp, dim=None, keepdim=False, *, dtype=None):
         result = argmin(inp.to(torch.float32), dim=dim, keepdim=keepdim, dtype=dtype)
         return result
     if dim is None:
+        inp = inp.contiguous()
         M = inp.numel()
         if dtype is None:
             dtype = inp.dtype

src/flag_gems/runtime/backend/_ascend/ops/mean.py

@@ -45,6 +45,7 @@ def mean_kernel_1(
 
 def mean(inp, *, dtype=None):
     logger.debug("GEMS_ASCEND MEAN")
+    inp = inp.contiguous()
     M = inp.numel()
     if dtype is None:
         dtype = inp.dtype

src/flag_gems/runtime/backend/_ascend/ops/min.py

@@ -119,6 +119,7 @@ def min_kernel(
 
 def min(inp):
     logger.debug("GEMS_ASCEND MIN")
+    inp = inp.contiguous()
     M = inp.numel()
     block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
     mid_size = triton.cdiv(M, block_size)

tests/test_amax.py

@@ -44,3 +44,23 @@ def test_amax(shape, dim, keepdim, dtype):
         res_out = torch.amax(inp, dim=dim, keepdim=keepdim)
 
     utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.amax
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_amax_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The values outside the view would change the result.
+    base = torch.full((4, 6), 100.0, dtype=dtype, device=flag_gems.device)
+    base[:, :3] = torch.arange(1, 13, dtype=dtype, device=flag_gems.device).reshape(
+        4, 3
+    )
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.amax(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.amax(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_amin.py

@@ -61,3 +61,23 @@ def test_amin_(shape, dim, keepdim, dtype):
         res_out = res_out.expand_as(inp)
 
     utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.amin
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_amin_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The values outside the view would change the result.
+    base = torch.full((4, 6), -100.0, dtype=dtype, device=flag_gems.device)
+    base[:, :3] = torch.arange(1, 13, dtype=dtype, device=flag_gems.device).reshape(
+        4, 3
+    )
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.amin(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.amin(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_argmax.py

@@ -16,6 +16,7 @@
 import torch
 
 import flag_gems
+from flag_gems.runtime import torch_device_fn
 
 from . import accuracy_utils as utils
 from . import conftest as cfg
@@ -60,10 +61,52 @@ def test_argmax(shape, dim, keepdim, dtype):
     else:
         inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
 
-    ref_inp = utils.to_reference(inp)
+    use_cpu_ref = flag_gems.vendor_name == "ascend" and dim is None and keepdim
+    ref_inp = inp.cpu() if use_cpu_ref else utils.to_reference(inp)
 
     ref_out = torch.argmax(ref_inp, dim=dim, keepdim=keepdim)
+    if use_cpu_ref and not cfg.TO_CPU:
+        ref_out = ref_out.to(inp.device)
     with flag_gems.use_gems():
         res_out = torch.argmax(inp, dim=dim, keepdim=keepdim)
 
     utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.argmax
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_argmax_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The values outside the view would change the result.
+    base = torch.full((4, 6), 100.0, dtype=dtype, device=flag_gems.device)
+    base[:, :3] = torch.arange(1, 13, dtype=dtype, device=flag_gems.device).reshape(
+        4, 3
+    )
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.argmax(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.argmax(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.argmax
+@pytest.mark.skipif(
+    flag_gems.vendor_name != "ascend", reason="Ascend-specific regression test"
+)
+def test_argmax_large_flattened_block_size():
+    # block_size was derived from sqrt(numel) with no upper bound. For an input
+    # this large it went past what the device handles on the flattened path, and
+    # the returned index pointed at an element that was not the maximum.
+    for seed in (0, 2, 3, 9):
+        torch_device_fn.manual_seed_all(seed)
+        inp = torch.randn((200, 2560, 3), dtype=torch.bfloat16, device=flag_gems.device)
+        ref_out = torch.argmax(inp.cpu())
+
+        with flag_gems.use_gems():
+            res_out = torch.argmax(inp)
+
+        torch.testing.assert_close(res_out.cpu(), ref_out, atol=0, rtol=0)

tests/test_argmin.py

@@ -46,3 +46,23 @@ def test_argmin(shape, dim, keepdim, dtype):
         res_out = torch.argmin(inp, dim=dim, keepdim=keepdim)
 
     utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.argmin
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_argmin_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The values outside the view would change the result.
+    base = torch.full((4, 6), -100.0, dtype=dtype, device=flag_gems.device)
+    base[:, :3] = torch.arange(1, 13, dtype=dtype, device=flag_gems.device).reshape(
+        4, 3
+    )
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.argmin(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.argmin(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_mean.py

@@ -113,3 +113,23 @@ def test_mean_dim_large_innerdim(shape, dim, keepdim, dtype):
         res_out = torch.mean(inp, dim, keepdim)
 
     utils.gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.mean
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_mean_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The values outside the view would change the result.
+    base = torch.full((4, 6), 100.0, dtype=dtype, device=flag_gems.device)
+    base[:, :3] = torch.arange(1, 13, dtype=dtype, device=flag_gems.device).reshape(
+        4, 3
+    )
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp, True)
+
+    ref_out = torch.mean(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.mean(inp)
+
+    utils.gems_assert_close(res_out, ref_out, dtype)

tests/test_min.py

@@ -87,3 +87,23 @@ def test_min_dim(shape, dim, keepdim, dtype):
 
     utils.gems_assert_equal(res_out_index, ref_out_index)
     utils.gems_assert_equal(res_out_value, ref_out_value)
+
+
+@pytest.mark.min
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_min_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The values outside the view would change the result.
+    base = torch.full((4, 6), -100.0, dtype=dtype, device=flag_gems.device)
+    base[:, :3] = torch.arange(1, 13, dtype=dtype, device=flag_gems.device).reshape(
+        4, 3
+    )
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp)
+
+    ref_out = torch.min(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.min(inp)
+
+    utils.gems_assert_equal(res_out, ref_out)

tests/test_prod.py

@@ -90,3 +90,22 @@ def test_prod_dim_multi_tile(shape, dim, keepdim):
         res_out = torch.prod(inp, dim=dim, keepdim=keepdim)
 
     utils.gems_assert_close(res_out, ref_out, torch.float32)
+
+
+@pytest.mark.prod
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_prod_full_reduction_noncontiguous(dtype):
+    # The flattened reduction kernel addresses the input linearly, so for a view
+    # whose storage still holds the discarded columns it read those elements
+    # instead. The zeros outside the view would drive the product to 0.
+    base = torch.zeros((4, 6), dtype=dtype, device=flag_gems.device)
+    base[:, :3] = 1.0
+    base[0, 0] = 2.0
+    inp = base[:, :3]
+    ref_inp = utils.to_reference(inp, True)
+
+    ref_out = torch.prod(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.prod(inp)
+
+    utils.gems_assert_close(res_out, ref_out, dtype)