Simplify Buffer.fill() API (NVIDIA#1366)

Author: Andy-Jost

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from libc.stdint cimport uintptr_t, int64_t, uint64_t
+from libc.stdint cimport uintptr_t
 
 from cuda.bindings cimport cydriver
 from cuda.core.experimental._memory._device_memory_resource cimport DeviceMemoryResource
@@ -14,8 +14,14 @@ from cuda.core.experimental._stream cimport Stream_accept, Stream
 from cuda.core.experimental._utils.cuda_utils cimport HANDLE_RETURN
 
 import abc
+import sys
 from typing import TypeVar, Union
 
+if sys.version_info >= (3, 12):
+    from collections.abc import Buffer as BufferProtocol
+else:
+    BufferProtocol = object
+
 from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
 from cuda.core.experimental._utils.cuda_utils import driver
 
@@ -190,58 +196,85 @@ cdef class Buffer:
                 s
             ))
 
-    def fill(self, value: int, width: int, *, stream: Stream | GraphBuilder):
-        """Fill this buffer with a value pattern asynchronously on the given stream.
+    def fill(self, value: int | BufferProtocol, *, stream: Stream | GraphBuilder):
+        """Fill this buffer with a repeating byte pattern.
 
         Parameters
         ----------
-        value : int
-            Integer value to fill the buffer with
-        width : int
-            Width in bytes for each element (must be 1, 2, or 4)
+        value : int | :obj:`collections.abc.Buffer`
+            - int: Must be in range [0, 256). Converted to 1 byte.
+            - :obj:`collections.abc.Buffer`: Must be 1, 2, or 4 bytes.
         stream : :obj:`~_stream.Stream` | :obj:`~_graph.GraphBuilder`
-            Keyword argument specifying the stream for the asynchronous fill
+            Stream for the asynchronous fill operation.
 
         Raises
         ------
+        TypeError
+            If value is not an int and does not support the buffer protocol.
         ValueError
-            If width is not 1, 2, or 4, if value is out of range for the width,
-            or if buffer size is not divisible by width
+            If value byte length is not 1, 2, or 4.
+            If buffer size is not divisible by value byte length.
+        OverflowError
+            If int value is outside [0, 256).
 
         """
         cdef Stream s_stream = Stream_accept(stream)
         cdef unsigned char c_value8
         cdef unsigned short c_value16
         cdef unsigned int c_value32
         cdef size_t N
-
-        # Validate width
-        if width not in (1, 2, 4):
-            raise ValueError(f"width must be 1, 2, or 4, got {width}")
+        cdef size_t width
+        cdef unsigned int int_value
+
+        # Get fill pattern from value
+        if isinstance(value, int):
+            # We define the int input to mean a 1-byte pattern.
+            # Match int.to_bytes(1, "little") behavior: raise OverflowError if not in [0, 256).
+            if value < 0 or value >= 256:
+                raise OverflowError("int value must be in range [0, 256)")
+            width = 1
+            int_value = <unsigned int>value
+        else:
+            try:
+                mv = memoryview(value)
+            except TypeError:
+                raise TypeError(
+                    f"value must be an int or support the buffer protocol, got {type(value).__name__}"
+                ) from None
+            width = mv.nbytes
+
+            # Validate width early to avoid copying/processing large invalid inputs.
+            if width not in (1, 2, 4):
+                raise ValueError(f"value must be 1, 2, or 4 bytes, got {width}")
+
+            # Convert to a 1-D view of bytes.
+            #
+            # Note: NumPy scalar memoryviews are 0-D, and int.from_bytes(mv, ...) errors with
+            # "0-dim memory has no length". Casting to 'B' gives us a byte-addressable view.
+            try:
+                int_value = int.from_bytes(mv.cast("B"), "little")
+            except TypeError:
+                int_value = int.from_bytes(mv.tobytes(), "little")
 
         # Validate buffer size modulus.
         cdef size_t buffer_size = self._size
         if buffer_size % width != 0:
-            raise ValueError(f"buffer size ({buffer_size}) must be divisible by width ({width})")
-
-        # Map width (bytes) to bitwidth and validate value
-        cdef int bitwidth = width * 8
-        _validate_value_against_bitwidth(bitwidth, value, is_signed=False)
+            raise ValueError(f"buffer size ({buffer_size}) must be divisible by {width}")
 
-        # Validate value fits in width and perform fill
+        # Perform fill based on width
         cdef cydriver.CUstream s = s_stream._handle
         if width == 1:
-            c_value8 = <unsigned char>value
+            c_value8 = <unsigned char>int_value
             N = buffer_size
             with nogil:
                 HANDLE_RETURN(cydriver.cuMemsetD8Async(<cydriver.CUdeviceptr>self._ptr, c_value8, N, s))
         elif width == 2:
-            c_value16 = <unsigned short>value
+            c_value16 = <unsigned short>int_value
             N = buffer_size // 2
             with nogil:
                 HANDLE_RETURN(cydriver.cuMemsetD16Async(<cydriver.CUdeviceptr>self._ptr, c_value16, N, s))
         else:  # width == 4
-            c_value32 = <unsigned int>value
+            c_value32 = <unsigned int>int_value
             N = buffer_size // 4
             with nogil:
                 HANDLE_RETURN(cydriver.cuMemsetD32Async(<cydriver.CUdeviceptr>self._ptr, c_value32, N, s))
@@ -409,43 +442,3 @@ cdef class MemoryResource:
             and document the behavior.
         """
         ...
-
-
-# Helper Functions
-# ----------------
-cdef void _validate_value_against_bitwidth(int bitwidth, int64_t value, bint is_signed=False) except *:
-    """Validate that a value fits within the representable range for a given bitwidth.
-
-    Parameters
-    ----------
-    bitwidth : int
-        Number of bits (e.g., 8, 16, 32)
-    value : int64_t
-        Value to validate
-    is_signed : bool, optional
-        Whether the value is signed (default: False)
-
-    Raises
-    ------
-    ValueError
-        If value is outside the representable range for the bitwidth
-    """
-    cdef int max_bits = bitwidth
-    assert max_bits < 64, f"bitwidth ({max_bits}) must be less than 64"
-
-    cdef int64_t min_value
-    cdef uint64_t max_value_unsigned
-    cdef int64_t max_value
-
-    if is_signed:
-        min_value = -(<int64_t>1 << (max_bits - 1))
-        max_value = (<int64_t>1 << (max_bits - 1)) - 1
-    else:
-        min_value = 0
-        max_value_unsigned = (<uint64_t>1 << max_bits) - 1
-        max_value = <int64_t>max_value_unsigned
-
-    if not min_value <= value <= max_value:
-        raise ValueError(
-            f"value must be in range [{min_value}, {max_value}]"
-        )

cuda_core/tests/test_graph_mem.py

@@ -112,7 +112,7 @@ def apply_kernels(mr, stream, out):
         # Fills out with 3
         def apply_kernels(mr, stream, out):
             buffer = mr.allocate(NBYTES, stream=stream)
-            buffer.fill(3, width=1, stream=stream)
+            buffer.fill(3, stream=stream)
             out.copy_from(buffer, stream=stream)
             buffer.close()
 

cuda_core/tests/test_memory.py

@@ -219,86 +219,105 @@ def test_buffer_copy_from():
     buffer_copy_from(DummyPinnedMemoryResource(device), device, check=True)
 
 
-def buffer_fill(dummy_mr: MemoryResource, device: Device, check=False):
-    stream = device.create_stream()
-
-    # Test width=1 (byte fill)
-    buffer1 = dummy_mr.allocate(size=1024)
-    buffer1.fill(0x42, width=1, stream=stream)
-    device.sync()
-
-    if check:
-        ptr = ctypes.cast(buffer1.handle, ctypes.POINTER(ctypes.c_byte))
-        for i in range(10):
-            assert ptr[i] == 0x42
+def _bytes_repeat(pattern: bytes, size: int) -> bytes:
+    assert len(pattern) > 0
+    assert size % len(pattern) == 0
+    return pattern * (size // len(pattern))
 
-    # Test error: invalid width
-    for bad_width in [w for w in range(-10, 10) if w not in (1, 2, 4)]:
-        with pytest.raises(ValueError, match="width must be 1, 2, or 4"):
-            buffer1.fill(0x42, width=bad_width, stream=stream)
 
-    # Test error: value out of range for width=1
-    for bad_value in [-42, -1, 256]:
-        with pytest.raises(ValueError, match="value must be in range \\[0, 255\\]"):
-            buffer1.fill(bad_value, width=1, stream=stream)
+def _pattern_bytes(value) -> bytes:
+    if isinstance(value, int):
+        return bytes([value])
+    return bytes(memoryview(value).cast("B"))
 
-    # Test error: buffer size not divisible by width
-    for bad_size in [1025, 1027, 1029, 1031]:  # Not divisible by 2
-        buffer_err = dummy_mr.allocate(size=1025)
-        with pytest.raises(ValueError, match="must be divisible"):
-            buffer_err.fill(0x1234, width=2, stream=stream)
-        buffer_err.close()
-
-    buffer1.close()
-
-    # Test width=2 (16-bit fill)
-    buffer2 = dummy_mr.allocate(size=1024)  # Divisible by 2
-    buffer2.fill(0x1234, width=2, stream=stream)
-    device.sync()
-
-    if check:
-        ptr = ctypes.cast(buffer2.handle, ctypes.POINTER(ctypes.c_uint16))
-        for i in range(5):
-            assert ptr[i] == 0x1234
-
-    # Test error: value out of range for width=2
-    for bad_value in [-42, -1, 65536, 65537, 100000]:
-        with pytest.raises(ValueError, match="value must be in range \\[0, 65535\\]"):
-            buffer2.fill(bad_value, width=2, stream=stream)
-
-    buffer2.close()
 
-    # Test width=4 (32-bit fill)
-    buffer4 = dummy_mr.allocate(size=1024)  # Divisible by 4
-    buffer4.fill(0xDEADBEEF, width=4, stream=stream)
-    device.sync()
-
-    if check:
-        ptr = ctypes.cast(buffer4.handle, ctypes.POINTER(ctypes.c_uint32))
-        for i in range(5):
-            assert ptr[i] == 0xDEADBEEF
-
-    # Test error: value out of range for width=4
-    for bad_value in [-42, -1, 4294967296, 4294967297, 5000000000]:
-        with pytest.raises(ValueError, match="value must be in range \\[0, 4294967295\\]"):
-            buffer4.fill(bad_value, width=4, stream=stream)
-
-    # Test error: buffer size not divisible by width
-    for bad_size in [1025, 1026, 1027, 1029, 1030, 1031]:  # Not divisible by 4
-        buffer_err2 = dummy_mr.allocate(size=bad_size)
-        with pytest.raises(ValueError, match="must be divisible"):
-            buffer_err2.fill(0xDEADBEEF, width=4, stream=stream)
-        buffer_err2.close()
-
-    buffer4.close()
-
-
-def test_buffer_fill():
+@pytest.fixture(params=["device", "unified", "pinned"])
+def fill_env(request):
     device = Device()
     device.set_current()
-    buffer_fill(DummyDeviceMemoryResource(device), device)
-    buffer_fill(DummyUnifiedMemoryResource(device), device)
-    buffer_fill(DummyPinnedMemoryResource(device), device, check=True)
+    if request.param == "device":
+        mr = DummyDeviceMemoryResource(device)
+    elif request.param == "unified":
+        mr = DummyUnifiedMemoryResource(device)
+    else:
+        mr = DummyPinnedMemoryResource(device)
+    return device, mr
+
+
+_FILL_SIZE = 64  # Keep small; divisible by 1/2/4.
+
+_FILL_CASES = [
+    # int -> 1-byte pattern
+    pytest.param(0x42, _FILL_SIZE, None, id="int-0x42"),
+    pytest.param(-1, _FILL_SIZE, OverflowError, id="int-neg"),
+    pytest.param(256, _FILL_SIZE, OverflowError, id="int-256"),
+    pytest.param(1000, _FILL_SIZE, OverflowError, id="int-1000"),
+    # bad type
+    pytest.param("invalid", _FILL_SIZE, TypeError, id="bad-type-str"),
+    # bytes-like patterns
+    pytest.param(b"\x7f", _FILL_SIZE, None, id="bytes-1"),
+    pytest.param(b"\x34\x12", _FILL_SIZE, None, id="bytes-2"),
+    pytest.param(b"\xef\xbe\xad\xde", _FILL_SIZE, None, id="bytes-4"),
+    pytest.param(b"\x34\x12", _FILL_SIZE + 1, ValueError, id="bytes-2-bad-size"),
+    pytest.param(b"\xef\xbe\xad\xde", _FILL_SIZE + 2, ValueError, id="bytes-4-bad-size"),
+    pytest.param(b"", _FILL_SIZE, ValueError, id="bytes-0"),
+    pytest.param(b"\x01\x02\x03", _FILL_SIZE, ValueError, id="bytes-3"),
+]
+
+if np is not None:
+    _FILL_CASES.extend(
+        [
+            # 8-bit patterns
+            pytest.param(np.uint8(0), _FILL_SIZE, None, id="np-uint8-0"),
+            pytest.param(np.uint8(255), _FILL_SIZE, None, id="np-uint8-255"),
+            pytest.param(np.int8(-1), _FILL_SIZE, None, id="np-int8--1"),
+            pytest.param(np.int8(127), _FILL_SIZE, None, id="np-int8-127"),
+            pytest.param(np.int8(-128), _FILL_SIZE, None, id="np-int8--128"),
+            # 16-bit patterns
+            pytest.param(np.uint16(0x1234), _FILL_SIZE, None, id="np-uint16-0x1234"),
+            pytest.param(np.uint16(0xFFFF), _FILL_SIZE, None, id="np-uint16-0xFFFF"),
+            pytest.param(np.int16(-1), _FILL_SIZE, None, id="np-int16--1"),
+            pytest.param(np.int16(32767), _FILL_SIZE, None, id="np-int16-max"),
+            pytest.param(np.int16(-32768), _FILL_SIZE, None, id="np-int16-min"),
+            pytest.param(np.uint16(0x1234), _FILL_SIZE + 1, ValueError, id="np-uint16-bad-size"),
+            # 32-bit patterns
+            pytest.param(np.uint32(0xDEADBEEF), _FILL_SIZE, None, id="np-uint32-0xDEADBEEF"),
+            pytest.param(np.uint32(0xFFFFFFFF), _FILL_SIZE, None, id="np-uint32-0xFFFFFFFF"),
+            pytest.param(np.int32(-1), _FILL_SIZE, None, id="np-int32--1"),
+            pytest.param(np.int32(2147483647), _FILL_SIZE, None, id="np-int32-max"),
+            pytest.param(np.int32(-2147483648), _FILL_SIZE, None, id="np-int32-min"),
+            pytest.param(np.uint32(0xDEADBEEF), _FILL_SIZE + 2, ValueError, id="np-uint32-bad-size"),
+            # float32 (bit-pattern fill)
+            pytest.param(np.float32(1.0), _FILL_SIZE, None, id="np-float32-1.0"),
+            # 64-bit patterns should error (8-byte pattern)
+            pytest.param(np.uint64(0), _FILL_SIZE, ValueError, id="np-uint64-err"),
+            pytest.param(np.int64(0), _FILL_SIZE, ValueError, id="np-int64-err"),
+            pytest.param(np.float64(0), _FILL_SIZE, ValueError, id="np-float64-err"),
+        ]
+    )
+
+
+@pytest.mark.parametrize("value,size,exc", _FILL_CASES)
+def test_buffer_fill(fill_env, value, size, exc):
+    device, mr = fill_env
+    stream = device.create_stream()
+    buffer = mr.allocate(size=size)
+    try:
+        if exc is not None:
+            with pytest.raises(exc):
+                buffer.fill(value, stream=stream)
+            return
+
+        buffer.fill(value, stream=stream)
+        device.sync()
+
+        # Verify contents only for host-accessible buffers.
+        if buffer.is_host_accessible:
+            pat = _pattern_bytes(value)
+            got = ctypes.string_at(int(buffer.handle), size)
+            assert got == _bytes_repeat(pat, size)
+    finally:
+        buffer.close()
 
 
 def buffer_close(dummy_mr: MemoryResource):