feat: Separate API for writing from copying. (#615)

Author: iwoplaza

packages/typegpu/src/core/buffer/buffer.ts

@@ -40,22 +40,21 @@ export interface TgpuBuffer<TData extends AnyData> extends TgpuNamable {
   readonly label: string | undefined;
 
   readonly buffer: GPUBuffer;
-  readonly device: GPUDevice;
   readonly destroyed: boolean;
 
   $usage<T extends ('uniform' | 'storage' | 'vertex')[]>(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>>;
   $addFlags(flags: GPUBufferUsageFlags): this;
-  $device(device: GPUDevice): this;
 
-  write(data: Infer<TData> | TgpuBuffer<TData>): void;
+  write(data: Infer<TData>): void;
+  copyFrom(srcBuffer: TgpuBuffer<TData>): void;
   read(): Promise<Infer<TData>>;
   destroy(): void;
 }
 
-export function createBufferImpl<TData extends AnyData>(
-  group: ExperimentalTgpuRoot | undefined,
+export function INTERNAL_createBuffer<TData extends AnyData>(
+  group: ExperimentalTgpuRoot,
   typeSchema: TData,
   initialOrBuffer?: Infer<TData> | GPUBuffer,
 ): TgpuBuffer<TData> {
@@ -88,8 +87,8 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   public readonly resourceType = 'buffer';
   public flags: GPUBufferUsageFlags =
     GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC;
-  private _device: GPUDevice | null = null;
   private _buffer: GPUBuffer | null = null;
+  private _ownBuffer: boolean;
   private _destroyed = false;
 
   private _label: string | undefined;
@@ -100,13 +99,15 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   public usableAsVertex = false;
 
   constructor(
-    private readonly _group: ExperimentalTgpuRoot | undefined,
+    private readonly _group: ExperimentalTgpuRoot,
     public readonly dataType: TData,
     public readonly initialOrBuffer?: Infer<TData> | GPUBuffer | undefined,
   ) {
     if (isGPUBuffer(initialOrBuffer)) {
+      this._ownBuffer = false;
       this._buffer = initialOrBuffer;
     } else {
+      this._ownBuffer = true;
       this.initial = initialOrBuffer;
     }
   }
@@ -116,18 +117,14 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
   }
 
   get buffer() {
-    if (!this._device) {
-      throw new Error(
-        'Create this buffer using `root.createBuffer` instead of `tgpu.createBuffer`.',
-      );
-    }
+    const device = this._group.device;
 
     if (this._destroyed) {
       throw new Error('This buffer has been destroyed');
     }
 
     if (!this._buffer) {
-      this._buffer = this._device.createBuffer({
+      this._buffer = device.createBuffer({
         size: sizeOf(this.dataType),
         usage: this.flags,
         mappedAtCreation: !!this.initial,
@@ -144,15 +141,6 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     return this._buffer;
   }
 
-  get device() {
-    if (!this._device) {
-      throw new Error(
-        'This buffer has not been assigned a device. Use .$device(device) to assign a device',
-      );
-    }
-    return this._device;
-  }
-
   get destroyed() {
     return this._destroyed;
   }
@@ -185,56 +173,42 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     return this;
   }
 
-  $device(device: GPUDevice) {
-    this._device = device;
-    return this;
-  }
-
-  write(dataOrBuffer: Infer<TData> | TgpuBuffer<TData>): void {
+  write(data: Infer<TData>): void {
     const gpuBuffer = this.buffer;
-    const device = this.device;
+    const device = this._group.device;
 
     if (gpuBuffer.mapState === 'mapped') {
       const mapped = gpuBuffer.getMappedRange();
-      if (isBuffer(dataOrBuffer)) {
-        throw new Error('Cannot copy to a mapped buffer.');
-      }
-      writeData(new BufferWriter(mapped), this.dataType, dataOrBuffer);
+      writeData(new BufferWriter(mapped), this.dataType, data);
       return;
     }
 
     const size = sizeOf(this.dataType);
-    if (isBuffer(dataOrBuffer)) {
-      const sourceBuffer = dataOrBuffer.buffer;
-
-      if (this._group) {
-        const encoder = this._group.commandEncoder;
-        encoder.copyBufferToBuffer(sourceBuffer, 0, gpuBuffer, 0, size);
-      } else {
-        const commandEncoder = device.createCommandEncoder();
-        commandEncoder.copyBufferToBuffer(sourceBuffer, 0, gpuBuffer, 0, size);
-        device.queue.submit([commandEncoder.finish()]);
-      }
-    } else {
-      if (this._group) {
-        // Flushing any commands yet to be encoded.
-        this._group.flush();
-      }
 
-      const hostBuffer = new ArrayBuffer(size);
-      writeData(new BufferWriter(hostBuffer), this.dataType, dataOrBuffer);
-      device.queue.writeBuffer(gpuBuffer, 0, hostBuffer, 0, size);
+    // Flushing any commands yet to be encoded.
+    this._group.flush();
+
+    const hostBuffer = new ArrayBuffer(size);
+    writeData(new BufferWriter(hostBuffer), this.dataType, data);
+    device.queue.writeBuffer(gpuBuffer, 0, hostBuffer, 0, size);
+  }
+
+  copyFrom(srcBuffer: TgpuBuffer<TData>): void {
+    if (this.buffer.mapState === 'mapped') {
+      throw new Error('Cannot copy to a mapped buffer.');
     }
+
+    const size = sizeOf(this.dataType);
+    const encoder = this._group.commandEncoder;
+    encoder.copyBufferToBuffer(srcBuffer.buffer, 0, this.buffer, 0, size);
   }
 
   async read(): Promise<Infer<TData>> {
-    if (this._group) {
-      // Flushing any commands yet to be encoded.
-      this._group.flush();
-    }
+    // Flushing any commands yet to be encoded.
+    this._group.flush();
 
     const gpuBuffer = this.buffer;
-    const device = this.device;
+    const device = this._group.device;
 
     if (gpuBuffer.mapState === 'mapped') {
       const mapped = gpuBuffer.getMappedRange();
@@ -283,7 +257,9 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
       return;
     }
     this._destroyed = true;
-    this._buffer?.destroy();
+    if (this._ownBuffer) {
+      this._buffer?.destroy();
+    }
   }
 
   toString(): string {

packages/typegpu/src/core/root/init.ts

@@ -22,7 +22,11 @@ import {
   isBindGroupLayout,
 } from '../../tgpuBindGroupLayout';
 import type { TgpuSlot } from '../../types';
-import { type TgpuBuffer, createBufferImpl, isBuffer } from '../buffer/buffer';
+import {
+  INTERNAL_createBuffer,
+  type TgpuBuffer,
+  isBuffer,
+} from '../buffer/buffer';
 import type { IOLayout } from '../function/fnTypes';
 import type { TgpuComputeFn } from '../function/tgpuComputeFn';
 import type { TgpuFragmentFn } from '../function/tgpuFragmentFn';
@@ -189,12 +193,8 @@ class TgpuRootImpl extends WithBindingImpl implements ExperimentalTgpuRoot {
     typeSchema: TData,
     initialOrBuffer?: Infer<TData> | GPUBuffer,
   ): TgpuBuffer<TData> {
-    const buffer = createBufferImpl(this, typeSchema, initialOrBuffer).$device(
-      this.device,
-    );
-
+    const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer);
     this._disposables.push(buffer);
-
     return buffer;
   }
 

packages/typegpu/src/index.ts

@@ -3,9 +3,7 @@
  */
 
 import { init, initFromDevice } from './core/root/init';
-import { createBuffer } from './legacyBufferApi';
 import { bindGroupLayout } from './tgpuBindGroupLayout';
-import { read, write } from './tgpuBufferUtils';
 
 export const tgpu = {
   /** @hidden @deprecated Use `'uniform'` string literal instead. */
@@ -19,13 +17,6 @@ export const tgpu = {
 
   init,
   initFromDevice,
-
-  /** @hidden */
-  createBuffer,
-  /** @hidden */
-  read,
-  /** @hidden */
-  write,
 };
 export default tgpu;