Entry point called by bench.py. Must match reference.reduce_sum_ref signature.
(x: torch.Tensor, dim: int = -1)
| 109 | |
| 110 | |
| 111 | def kernel_fn(x: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 112 | """Entry point called by bench.py. Must match reference.reduce_sum_ref signature.""" |
| 113 | assert x.is_cuda |
| 114 | |
| 115 | # Normalize dim |
| 116 | if dim < 0: |
| 117 | dim = x.ndim + dim |
| 118 | assert 0 <= dim < x.ndim |
| 119 | |
| 120 | orig_dtype = x.dtype |
| 121 | |
| 122 | # For the common case of reducing the last dimension on a 2D tensor |
| 123 | if dim == x.ndim - 1: |
| 124 | if x.dtype != torch.float16: |
| 125 | x = x.to(torch.float16) |
| 126 | |
| 127 | # Ensure 2D |
| 128 | orig_shape = list(x.shape) |
| 129 | if x.dim() == 1: |
| 130 | x = x.unsqueeze(0) |
| 131 | elif x.dim() > 2: |
| 132 | x = x.reshape(-1, x.shape[-1]) |
| 133 | |
| 134 | mod = _get_module() |
| 135 | out = mod.reduce_sum_cuda(x) |
| 136 | |
| 137 | if orig_dtype != torch.float16: |
| 138 | out = out.to(orig_dtype) |
| 139 | |
| 140 | # Restore output shape (input shape minus last dim) |
| 141 | out_shape = orig_shape[:-1] |
| 142 | if not out_shape: |
| 143 | out_shape = [1] |
| 144 | return out.view(out_shape) |
| 145 | else: |
| 146 | # General case: move reduction dim to last, then reduce |
| 147 | perm = list(range(x.ndim)) |
| 148 | perm.pop(dim) |
| 149 | perm.append(dim) |
| 150 | x = x.permute(*perm).contiguous() |
| 151 | return kernel_fn(x, dim=-1) |
nothing calls this directly
no test coverage detected