fix: fix sum_dim_out and trunc_div type issues (#2105)

Author: physics31415926

src/flag_gems/ops/div.py

@@ -71,17 +71,44 @@ def trunc_div_func(x, y):
 @pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
 @triton.jit
 def trunc_div_func_tensor_scalar(x, y):
-    return trunc(div_rz(x, y))
+    return trunc(div_rz(x, tl.cast(y, x.dtype)))
 
 
 @pointwise_dynamic(is_tensor=[False, True], promotion_methods=[(0, 1, "DEFAULT")])
 @triton.jit
 def trunc_div_func_scalar_tensor(x, y):
-    return trunc(div_rz(x, y))
+    return trunc(div_rz(tl.cast(x, y.dtype), y))
+
+
+# Integer truncation division: Triton's // on integers is C-style (truncates toward zero)
+@pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def trunc_div_int_func(x, y):
+    return x // y
+
+
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def trunc_div_int_func_tensor_scalar(x, y):
+    return x // y
+
+
+@pointwise_dynamic(is_tensor=[False, True], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def trunc_div_int_func_scalar_tensor(x, y):
+    return x // y
 
 
 def trunc_divide(A, B):
     logger.debug("GEMS TRUNC_DIVIDE")
+    # Integer types: use dedicated int kernels (Triton // is C-style truncation)
+    if isinstance(A, torch.Tensor) and not A.is_floating_point():
+        if isinstance(B, torch.Tensor):
+            return trunc_div_int_func(A, B)
+        else:
+            return trunc_div_int_func_tensor_scalar(A, B)
+    if isinstance(B, torch.Tensor) and not B.is_floating_point():
+        return trunc_div_int_func_scalar_tensor(A, B)
     if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
         return trunc_div_func(A, B)
     elif isinstance(A, torch.Tensor):
@@ -95,6 +122,12 @@ def trunc_divide(A, B):
 
 def trunc_divide_(A, B):
     logger.debug("GEMS TRUNC_DIVIDE_")
+    # Integer types: use dedicated int kernels (Triton // is C-style truncation)
+    if not A.is_floating_point():
+        if isinstance(B, torch.Tensor):
+            return trunc_div_int_func(A, B, out0=A)
+        else:
+            return trunc_div_int_func_tensor_scalar(A, B, out0=A)
     if isinstance(B, torch.Tensor):
         return trunc_div_func(A, B, out0=A)
     else:

src/flag_gems/ops/sum.py

@@ -267,7 +267,16 @@ def sum_dim_comm(inp, dim=None, keepdim=False, *, dtype=None, out=None):
         inp = inp.contiguous()
         K = inp.numel() // M // N
         shape[dim] = 1
-        if out is None:
+        _out_provided = out is not None
+        if _out_provided:
+            # Resize out to the expected output shape, matching native PyTorch
+            # sum.out behavior. The caller (e.g. logsumexp) may pass a
+            # zero-size placeholder that needs to be resized before use.
+            if keepdim:
+                out.resize_(shape)
+            else:
+                out.resize_(shape[:dim] + shape[dim + 1 :])
+        else:
             out = torch.empty(shape, dtype=dtype, device=inp.device)
 
         with torch_device_fn.device(inp.device):
@@ -288,7 +297,7 @@ def sum_dim_comm(inp, dim=None, keepdim=False, *, dtype=None, out=None):
                     M,
                     N,
                 )
-        if not keepdim:
+        if not keepdim and not _out_provided:
             out = out.squeeze(dim=dim)
         return out
     else:
@@ -298,14 +307,22 @@ def sum_dim_comm(inp, dim=None, keepdim=False, *, dtype=None, out=None):
             N *= shape[i]
             shape[i] = 1
         M = inp.numel() // N
-        if out is None:
+        _out_provided = out is not None
+        if _out_provided:
+            dim_set = set(dim)
+            if keepdim:
+                out.resize_(shape)
+            else:
+                out.resize_([s for i, s in enumerate(shape) if i not in dim_set])
+        else:
             out = torch.empty(shape, dtype=dtype, device=inp.device)
 
         grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]),)
         with torch_device_fn.device(inp.device):
             sum_dim_kernel[grid](inp, out, M, N)
-        if not keepdim:
-            out = out.squeeze(dim=dim)
+        if not keepdim and not _out_provided:
+            for d in sorted(dim, reverse=True):
+                out = out.squeeze(dim=d)
         return out
 
 

tests/test_binary_pointwise_ops.py

@@ -722,6 +722,58 @@ def test_accuracy_trunc_divide_scalar_scalar(dtype):
         gems_assert_close(res_out, ref_out, dtype)
 
 
+@pytest.mark.div
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", INT_DTYPES + [torch.int64])
+def test_trunc_div_int(shape, dtype):
+    # Regression test: integer types must be dispatched at Python layer to avoid
+    # passing int tensors to div_rz which only supports floating point.
+    inp1 = torch.randint(1, 100, shape, dtype=dtype, device=flag_gems.device)
+    inp2 = torch.randint(1, 100, shape, dtype=dtype, device=flag_gems.device)
+    ref_inp1 = to_reference(inp1, False)
+    ref_inp2 = to_reference(inp2, False)
+
+    ref_out = torch.div(ref_inp1, ref_inp2, rounding_mode="trunc")
+    with flag_gems.use_gems():
+        res_out = torch.div(inp1, inp2, rounding_mode="trunc")
+
+    gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.div
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", INT_DTYPES + [torch.int64])
+def test_trunc_div_tensor_scalar_int(shape, dtype):
+    # Regression test: integer types must be dispatched at Python layer to avoid
+    # passing int tensors to div_rz which only supports floating point.
+    inp1 = torch.randint(1, 100, shape, dtype=dtype, device=flag_gems.device)
+    scalar = random.randint(1, 10)
+    ref_inp1 = to_reference(inp1, False)
+
+    ref_out = torch.div(ref_inp1, scalar, rounding_mode="trunc")
+    with flag_gems.use_gems():
+        res_out = torch.div(inp1, scalar, rounding_mode="trunc")
+
+    gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.div
+@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
+@pytest.mark.parametrize("dtype", INT_DTYPES + [torch.int64])
+def test_trunc_div_scalar_tensor_int(shape, dtype):
+    # Regression test: integer types must be dispatched at Python layer to avoid
+    # passing int tensors to div_rz which only supports floating point.
+    inp2 = torch.randint(1, 100, shape, dtype=dtype, device=flag_gems.device)
+    scalar = random.randint(1, 100)
+    ref_inp2 = to_reference(inp2, False)
+
+    ref_out = torch.div(scalar, ref_inp2, rounding_mode="trunc")
+    with flag_gems.use_gems():
+        res_out = torch.div(scalar, inp2, rounding_mode="trunc")
+
+    gems_assert_equal(res_out, ref_out)
+
+
 # TODO: failed at large size, eg. (65536 * 2048,)
 @pytest.mark.floor_divide
 @pytest.mark.parametrize("shape", POINTWISE_SHAPES)

tests/test_general_reduction_ops.py

@@ -475,6 +475,34 @@ def test_accuracy_sum_dim(shape, dim, keepdim, dtype):
     gems_assert_close(res_out, ref_out, dtype, reduce_dim=_dim)
 
 
+@pytest.mark.sum
+@pytest.mark.parametrize("shape", REDUCTION_SHAPES)
+@pytest.mark.parametrize("keepdim, dim", KEEPDIM_DIM)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_sum_dim_out(shape, dim, keepdim, dtype):
+    # Regression test: sum_dim_out must resize external out tensor and skip squeeze.
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp = to_reference(inp, True)
+
+    ref_result = torch.sum(ref_inp, dim=dim, keepdim=keepdim)
+
+    # Pre-allocate out tensor with wrong shape to test resize logic
+    out = torch.empty((1,), dtype=dtype, device=flag_gems.device)
+    with flag_gems.use_gems():
+        res_result = torch.sum(inp, dim=dim, keepdim=keepdim, out=out)
+
+    if isinstance(dim, int):
+        dim = [dim]
+    dim = [d % inp.ndim for d in dim]
+    _dim = 1
+    for d in dim:
+        _dim *= shape[d]
+    if dim == []:
+        _dim = inp.numel()
+    gems_assert_close(res_result, ref_result, dtype, reduce_dim=_dim)
+    gems_assert_close(out, ref_result, dtype, reduce_dim=_dim)
+
+
 QUANTILE_SHAPES = REDUCTION_SMALL_SHAPES + [(10, 64, 196), (65535, 1)]
 QUANTILE_FLOAT_DTYPES = [torch.float32]
 QUANTILE_Q = (