Fix fill operators crash on non-contiguous tensors (stride != 1) in hopper backend (flagos-ai#2560)

Author: wuhulalala

src/flag_gems/runtime/backend/_nvidia/hopper/ops/fill.py

@@ -43,6 +43,17 @@ def fill_tensor_kernel(
     tl.store(out_ptr + offsets, val, mask=mask)
 
 
+def _as_contiguous(tensor):
+    """Return tensor.contiguous() view for use with flat-offset kernels.
+
+    For non-contiguous tensors this allocates a new buffer; callers that
+    need in-place semantics must copy back afterwards.
+    """
+    if tensor.is_contiguous():
+        return tensor, False
+    return tensor.contiguous(), True
+
+
 def fill_scalar(input, value):
     logger.debug("GEMS_HOPPER FILL_SCALAR")
     out = torch.empty_like(input)
@@ -57,10 +68,13 @@ def fill_scalar_out(input, value, *, out=None):
     logger.debug("GEMS_HOPPER FILL_SCALAR_OUT")
     if out is None:
         return fill_scalar(input, value)
-    n_elements = out.numel()
+    out_contig, need_copy = _as_contiguous(out)
+    n_elements = out_contig.numel()
     grid = (triton.cdiv(n_elements, 1024),)
     with torch_device_fn.device(input.device):
-        fill_scalar_kernel[grid](out, value, n_elements, BLOCK_SIZE=1024)
+        fill_scalar_kernel[grid](out_contig, value, n_elements, BLOCK_SIZE=1024)
+    if need_copy:
+        out.copy_(out_contig)
     return out
 
 
@@ -90,10 +104,13 @@ def fill_tensor_out(input, value, *, out=None):
         raise RuntimeError(
             f"fill only supports 0-dimension value tensor but got tensor with {value.ndim} dimensions."
         )
-    n_elements = out.numel()
+    out_contig, need_copy = _as_contiguous(out)
+    n_elements = out_contig.numel()
     grid = (triton.cdiv(n_elements, 1024),)
     with torch_device_fn.device(input.device):
-        fill_tensor_kernel[grid](out, value, n_elements, BLOCK_SIZE=1024)
+        fill_tensor_kernel[grid](out_contig, value, n_elements, BLOCK_SIZE=1024)
+    if need_copy:
+        out.copy_(out_contig)
     return out
 
 
@@ -105,17 +122,33 @@ def fill_tensor_(self, value):
         raise RuntimeError(
             f"fill only supports 0-dimension value tensor but got tensor with {value.ndim} dimensions."
         )
-    n_elements = self.numel()
-    grid = (triton.cdiv(n_elements, 1024),)
-    with torch_device_fn.device(self.device):
-        fill_tensor_kernel[grid](self, value, n_elements, BLOCK_SIZE=1024)
+    if self.is_contiguous():
+        n_elements = self.numel()
+        grid = (triton.cdiv(n_elements, 1024),)
+        with torch_device_fn.device(self.device):
+            fill_tensor_kernel[grid](self, value, n_elements, BLOCK_SIZE=1024)
+    else:
+        tmp = self.contiguous()
+        n_elements = tmp.numel()
+        grid = (triton.cdiv(n_elements, 1024),)
+        with torch_device_fn.device(self.device):
+            fill_tensor_kernel[grid](tmp, value, n_elements, BLOCK_SIZE=1024)
+        self.copy_(tmp)
     return self
 
 
 def fill_scalar_(self, value):
     logger.debug("GEMS_HOPPER FILL_SCALAR_")
-    n_elements = self.numel()
-    grid = (triton.cdiv(n_elements, 1024),)
-    with torch_device_fn.device(self.device):
-        fill_scalar_kernel[grid](self, value, n_elements, BLOCK_SIZE=1024)
+    if self.is_contiguous():
+        n_elements = self.numel()
+        grid = (triton.cdiv(n_elements, 1024),)
+        with torch_device_fn.device(self.device):
+            fill_scalar_kernel[grid](self, value, n_elements, BLOCK_SIZE=1024)
+    else:
+        tmp = self.contiguous()
+        n_elements = tmp.numel()
+        grid = (triton.cdiv(n_elements, 1024),)
+        with torch_device_fn.device(self.device):
+            fill_scalar_kernel[grid](tmp, value, n_elements, BLOCK_SIZE=1024)
+        self.copy_(tmp)
     return self

tests/test_fill.py

@@ -97,6 +97,32 @@ def test_fill_scalar_(value, shape, dtype):
     with flag_gems.use_gems():
         x.fill_(value)
 
+
+FILL_SLICE_CASES = [
+    # (shape, slice)
+    ((4, 128), (slice(None), slice(64, None))),
+    ((2, 1, 1, 512), (slice(None), slice(None), slice(None), slice(358, None))),
+    ((8, 32, 64), (slice(None), slice(16, None))),
+]
+
+
+@pytest.mark.fill_scalar_
+@pytest.mark.parametrize("shape, slc", FILL_SLICE_CASES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.BOOL_TYPES)
+@pytest.mark.parametrize(
+    "value", [0, 1, True, float("-inf")], ids=["zero", "one", "true", "neginf"]
+)
+def test_fill_sliced_view_scalar(shape, slc, dtype, value):
+    if dtype == torch.bool and value == float("-inf"):
+        pytest.skip("bool tensor does not support -inf")
+
+    x = torch.randn(shape, device=flag_gems.device).to(dtype)
+    ref_x = utils.to_reference(x, False)
+
+    ref_x[slc] = value
+    with flag_gems.use_gems():
+        x[slc] = value
+
     utils.gems_assert_equal(x, ref_x)
 
 
@@ -120,3 +146,20 @@ def test_fill_(value, shape, dtype):
         x.fill_(value_tensor)
 
     utils.gems_assert_equal(x, ref_x)
+
+
+@pytest.mark.fill_tensor_
+@pytest.mark.parametrize("shape, slc", FILL_SLICE_CASES)
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.BOOL_TYPES)
+@pytest.mark.parametrize("value", [0, 1, True], ids=["zero", "one", "true"])
+def test_fill_sliced_view_tensor(shape, slc, dtype, value):
+    x = torch.randn(shape, device=flag_gems.device).to(dtype)
+    ref_x = utils.to_reference(x, False)
+
+    value_tensor = torch.tensor(value, device=flag_gems.device, dtype=dtype)
+    ref_value_tensor = utils.to_reference(value_tensor, False)
+    ref_x[slc] = ref_value_tensor
+    with flag_gems.use_gems():
+        x[slc] = value_tensor
+
+    utils.gems_assert_equal(x, ref_x)