[Reduce] Move guard for zero-size inside custom_op

Author: tridao

quack/rmsnorm.py

@@ -351,6 +351,8 @@ def _rmsnorm_fwd(
         assert weight.dtype in supported_types, "Weight must be float32, float16 or bfloat16"
     if residual is not None:
         assert residual.dtype in supported_types, "Residual must be float16, bfloat16, or float32"
+    if x.numel() == 0:
+        return
 
     N = x.size(-1)
     per_head = (weight is not None and weight.dim() == 2) or (bias is not None and bias.dim() == 2)
@@ -494,8 +496,7 @@ def rmsnorm_fwd(
         )
     else:
         residual_out = None
-    if x.numel() > 0:
-        _rmsnorm_fwd(x, weight, out, bias, rstd, None, residual, residual_out, eps, False)
+    _rmsnorm_fwd(x, weight, out, bias, rstd, None, residual, residual_out, eps, False)
     # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0
     if residual_out is None:
         residual_out = x
@@ -964,6 +965,8 @@ def _rmsnorm_bwd(
     if dresidual is not None:
         assert dresidual.shape == x.shape
         assert dresidual.dtype in supported_types, "Residual must be float16, bfloat16, or float32"
+    if x.numel() == 0:
+        return
 
     N = x.size(-1)
     if dw_partial is None and db_partial is None:

quack/softmax.py

@@ -200,6 +200,8 @@ def _softmax_fwd(x: torch.Tensor, out: torch.Tensor) -> None:
     """
     assert x.dim() == 2, "Input must be 2D"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype"
+    if x.numel() == 0:
+        return
     N = x.size(1)
     dtype, out_dtype = [torch2cute_dtype_map[t.dtype] for t in [x, out]]
     _compile_softmax_fwd(dtype, out_dtype, N)(x, out)
@@ -224,8 +226,7 @@ def _softmax_fwd_fake(x: torch.Tensor, out: torch.Tensor) -> None:
 
 def softmax_fwd(x: torch.Tensor) -> torch.Tensor:
     out = torch.empty_like(x)
-    if x.numel() > 0:
-        _softmax_fwd(x, out)
+    _softmax_fwd(x, out)
     return out
 
 
@@ -401,6 +402,8 @@ def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor) -> No
     assert dy.shape == y.shape, "dy and y must have same shape"
     assert dy.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype"
     assert y.dtype == dy.dtype, "dy and y must have same dtype"
+    if dy.numel() == 0:
+        return
     N = dy.size(1)
     dtype, y_dtype, dx_dtype = [torch2cute_dtype_map[t.dtype] for t in [dy, y, dx]]
     _compile_softmax_backward(dtype, y_dtype, dx_dtype, N)(dy, y, dx)
@@ -419,8 +422,7 @@ def _softmax_backward_fake(dy: torch.Tensor, y: torch.Tensor, dx: torch.Tensor)
 
 def softmax_bwd(dy: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
     dx = torch.empty_like(dy)
-    if dy.numel() > 0:
-        _softmax_backward(dy, y, dx)
+    _softmax_backward(dy, y, dx)
     return dx
 
 

quack/topk.py

@@ -230,6 +230,8 @@ def _topk_fwd(
     assert x.dim() == 2, "Input must be 2D"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype"
     assert k > 0 and k <= x.shape[1], "k must be positive and <= N"
+    if x.numel() == 0:
+        return
 
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
@@ -284,8 +286,7 @@ def topk_fwd(x: torch.Tensor, k: int, softmax: bool = False):
     M = x.size(0)
     values = torch.empty((M, k), dtype=x.dtype, device=x.device)
     indices = torch.empty((M, k), dtype=torch.int32, device=x.device)
-    if x.numel() > 0:
-        _topk_fwd(x, k, softmax, values, indices)
+    _topk_fwd(x, k, softmax, values, indices)
     return values, indices
 
 
@@ -479,6 +480,8 @@ def _topk_bwd(
         assert values.dim() == 2, "values must be 2D"
     assert indices.dim() == 2, "indices must be 2D"
     assert dvalues.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported dtype"
+    if dvalues.numel() == 0:
+        return
 
     N = dx.size(1)
     dtype = torch2cute_dtype_map[dvalues.dtype]
@@ -548,8 +551,7 @@ def topk_bwd(
     """
     M, k = dvalues.shape
     dx = torch.zeros((M, N), dtype=dvalues.dtype, device=dvalues.device)
-    if dvalues.numel() > 0:
-        _topk_bwd(dvalues, values, indices, k, softmax, dx)
+    _topk_bwd(dvalues, values, indices, k, softmax, dx)
     return dx