[FRONTEND] Support negative axis for reduce and scan (triton-lang#2849)

Author: Jokeren

python/test/unit/language/test_core.py

@@ -1688,10 +1688,11 @@ def kernel(X, Z, BLOCK: tl.constexpr):
                    for shape in reduce3d_shapes
                    for axis in [0, 1, 2]]
 invalid_config = [('sum', 'float32', (32, 32), axis) for axis in [2, 3]]
+negative_config = [('sum', 'float32', (32, 32), -1)]
 
 
 @pytest.mark.parametrize("op, dtype_str, shape, axis",
-                         reduce_configs1 + reduce_configs2 + reduce_configs3 + invalid_config)
+                         reduce_configs1 + reduce_configs2 + reduce_configs3 + invalid_config + negative_config)
 @pytest.mark.parametrize("num_ctas", num_ctas_list)
 def test_reduce(op, dtype_str, shape, axis, num_ctas, device):
     check_type_supported(dtype_str, device)  # bfloat16 on cc < 80 will not be tested
@@ -1739,8 +1740,8 @@ def kernel(X, Z, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.const
     z_dtype_str = get_reduced_dtype(dtype_str, op)
     z_tri_dtype_str = z_dtype_str
     # triton result
-    ret_numel = 1 if axis is None else shape[1 - axis]
-    z_shape = (1, ) if axis is None else tuple(shape_i for i, shape_i in enumerate(shape) if i != axis)
+    non_negative_axis = axis if axis is None or axis >= 0 else len(shape) + axis
+    z_shape = (1, ) if axis is None else tuple(shape_i for i, shape_i in enumerate(shape) if i != non_negative_axis)
     z_tri = to_triton(numpy_random(z_shape, dtype_str=z_dtype_str, rs=rs), device=device, dst_type=z_tri_dtype_str)
     BLOCK_K = 1 if len(shape) == 2 else shape[2]
     IS_3D = bool(len(shape) == 3)
@@ -1787,6 +1788,7 @@ def kernel(X, Z, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.const
                 for axis in [1, 0]
                 for shape in scan2d_shapes
                 for op in ['cumsum', 'cumprod', 'get_first_element']]
+negative_config = [('cumsum', 'float32', (32, 32), -1, 4)]
 
 
 @triton.jit
@@ -1795,7 +1797,7 @@ def get_first_element(a, b):
     return a
 
 
-@pytest.mark.parametrize("op, dtype_str, shape, axis, num_warps", scan_configs)
+@pytest.mark.parametrize("op, dtype_str, shape, axis, num_warps", scan_configs + negative_config)
 def test_scan2d(op, dtype_str, shape, axis, num_warps, device):
     if is_hip():
         pytest.skip("test_scan2d is not supported in HIP")

python/triton/language/core.py

@@ -1477,8 +1477,9 @@ def make_combine_region(reduce_op):
                 handles = [r.handle for r in results]
             _builder.create_reduce_ret(*handles)
 
+    axis = _constexpr_to_value(axis)
     if axis is not None:
-        axis = _constexpr_to_value(axis)
+        axis = _wrap_axis(axis, len(input[0].shape))
     return semantic.reduction(input, axis, make_combine_region, _builder)
 
 
@@ -1557,6 +1558,8 @@ def make_combine_region(scan_op):
             _builder.create_scan_ret(*handles)
 
     axis = _constexpr_to_value(axis)
+    if axis is not None:
+        axis = _wrap_axis(axis, len(input[0].shape))
     return semantic.associative_scan(input, axis, make_combine_region, _builder)