WebGPU writeBuffer overWrites with zeros in special cases

Bugs
,
8

Here’s my untested attempt. I hope it’s useful for discussion.

  • setSubDataRaw - the lowest level. implemening chunking and calling GPUBuffer.writeBuffer
  • setSubDataAlign - wrapper around setSubDataRaw() that aligns dstOffset and bytesLength and copies from src.buffer outside the specified range. Suitable if src is a full copy of dst.
  • setSubDataPadLength - wrapper around setSubDataRaw() duplicating original setSubData()

My preference is to never use setSubDataPadLength(), but I include it because I don’t know if I’ve persuaded you. To insert this code as a non-breaking change, rename setSubDataPadLength to setSubData (and remove the original setSubData).

To insert this code as a breaking change, rename setSubDataRaw as setSubData and change the calls in setSubDataAlign and setSubDataPadLength to use the new setSubData. My preference would be to replace the existing setSubData with (renamed) setSubDataRaw and udpdate setSubDataAlign to call (the new) setSubData and don’t bother adding setSubDataPadLength.

// This calls GPUBuffer.writeBuffer() with no alignment corrections
// dstByteOffset and byteLength must both be aligned to 4 bytes
// and bytes moved must be within src and dst arrays
// note that 4th and 5th arguments are in BYTES, not elements, unlike GPUBuffer.writeBuffer()
// set maxChunk = Infinity for no chunking
// chunking ON by default
public setSubDataRaw(dataBuffer: WebGPUDataBuffer, dstByteOffset: number, src: ArrayBufferView, srcByteOffset = 0, byteLength = 0, maxChunk = 1024 * 1024 * 15): void {
    const buffer = dataBuffer.underlyingResource as GPUBuffer;
    
    // this I think fixes an error in original setSubData() by including srcByteOffset
    byteLength = byteLength || src.byteLength-srcByteOffset;;
    
    // if src is ArrayBuffer, then src.byteLength is available, but byteOffset is not.
    // src.byteOffset will interpret null or undefined as NaN so use src.byteOffset??0
    // src.buffer??src will be src because .buffer is undefined in ArrayBuffer
    
    // make srcByteOffset relative to src.buffer (allowing for ArrayBuffer with ??...)
    srcByteOffset += src.byteOffset??0;
    // since we convert to ArrayBuffer, data specified beyond src TypedArray WILL be copied from src.buffer 
    // we can check for this with "if (srcByteOffset+src.byteLength>src.byteLength) throw..."
    src = src.buffer??src; // here, make sure we're using an ArrayBuffer so args to writeBuffer are in bytes
    
    // or use (this is a little cleaner)
    //if (TypedArray.isView(src)) srcByteOffset += src.byteOffset;
    //else src = src.buffer;
    
    while (byteLength > maxChunk) {
        this._device.queue.writeBuffer(buffer, dstByteOffset, src, srcByteOffset, maxChunk);
        dstByteOffset += maxChunk;
        srcByteOffset += maxChunk;
        byteLength -= maxChunk;
    }
    
    this._device.queue.writeBuffer(buffer, dstByteOffset, src, srcByteOffset, byteLength);
    

}

// assume src is a full copy of dataBuffer, so handle alignment by copying 
// data outside srcByteOffset through byteLength. Use data from src.buffer if needed.
// chunking OFF by default
// if dstByteOffset is not a multiple of 4, copy additional bytes from src.buffer prior to src[srcByteOffset]
// and writing to buffer prior to buffer[dstByteOffset].
// Pad byteLength (adjusted by dstOffset alignment) the same way:
// by copying data after src[srcByteOffset+byteOffset]
public setSubDataAlign(dataBuffer: WebGPUDataBuffer, dstByteOffset: number, src: ArrayBufferView, srcByteOffset = 0, byteLength = 0, maxChunk = Infinity): void {
    
    // we might copy more than requested to make sure the write is aligned
    startPre = dstOffset & 3 // use 3 for 4-byte alignment. Use 7 for 8-byte alignment.
    srcByteOffset -= startPre
    dstByteOffset -= startPre
    byteLength = (byteLength+startPre+3) & ~3
    
    setSubDataRaw(dataBuffer, dstByteOffset, src, srcByteOffset, byteLength, maxChunk);
}
// pad byteLength to 4-byte boundary with 0. copy src to tmp with byteLength aligned to 4 bytes.
// this matches original setSubData
public setSubDataPadLength(dataBuffer: WebGPUDataBuffer, dstByteOffset: number, src: ArrayBufferView, srcByteOffset = 0, byteLength = 0, maxChunk = 1024 * 1024 * 15): void {
    
    // if alignment is needed, pad with 0 by copying to a new larger TypedArray
    
    const newbyteLength = (byteLength+startPre+3) & ~3
    if (newbyteLength!=byteLength) {
        const tempView = new Uint8Array(src.buffer??src,src.byteOffset+srcByteOffset,byteLength);
        src = new Uint8Array(newbyteLength)
        src.set(tempView)
        srcByteOffset = 0
    }
    setSubDataRaw(dataBuffer, dstByteOffset, src, srcByteOffset, byteLength, maxChunk);
}