1import { assert, unreachable } from '../../../common/framework/util/util.js';
2import { kSizedTextureFormatInfo, SizedTextureFormat } from '../../capability_info.js';
3import { align, isAligned } from '../math.js';
4
5export const kBytesPerRowAlignment = 256;
6export const kBufferCopyAlignment = 4;
7
8export interface LayoutOptions {
9  mipLevel: number;
10  bytesPerRow?: number;
11  rowsPerImage?: number;
12}
13
14const kDefaultLayoutOptions = { mipLevel: 0, bytesPerRow: undefined, rowsPerImage: undefined };
15
16export function getMipSizePassthroughLayers(
17  dimension: GPUTextureDimension,
18  size: [number, number, number],
19  mipLevel: number
20): [number, number, number] {
21  const shiftMinOne = (n: number) => Math.max(1, n >> mipLevel);
22  switch (dimension) {
23    case '1d':
24      assert(size[2] === 1);
25      return [shiftMinOne(size[0]), size[1], size[2]];
26    case '2d':
27      return [shiftMinOne(size[0]), shiftMinOne(size[1]), size[2]];
28    case '3d':
29      return [shiftMinOne(size[0]), shiftMinOne(size[1]), shiftMinOne(size[2])];
30    default:
31      unreachable();
32  }
33}
34
35export interface TextureCopyLayout {
36  bytesPerBlock: number;
37  byteLength: number;
38  minBytesPerRow: number;
39  bytesPerRow: number;
40  rowsPerImage: number;
41  mipSize: [number, number, number];
42}
43
44export function getTextureCopyLayout(
45  format: SizedTextureFormat,
46  dimension: GPUTextureDimension,
47  size: [number, number, number],
48  options: LayoutOptions = kDefaultLayoutOptions
49): TextureCopyLayout {
50  const { mipLevel } = options;
51  let { bytesPerRow, rowsPerImage } = options;
52
53  const mipSize = getMipSizePassthroughLayers(dimension, size, mipLevel);
54
55  const { blockWidth, blockHeight, bytesPerBlock } = kSizedTextureFormatInfo[format];
56
57  // We align mipSize to be the physical size of the texture subresource.
58  mipSize[0] = align(mipSize[0], blockWidth);
59  mipSize[1] = align(mipSize[1], blockHeight);
60
61  const minBytesPerRow = (mipSize[0] / blockWidth) * bytesPerBlock;
62  const alignedMinBytesPerRow = align(minBytesPerRow, kBytesPerRowAlignment);
63  if (bytesPerRow !== undefined) {
64    assert(bytesPerRow >= alignedMinBytesPerRow);
65    assert(isAligned(bytesPerRow, kBytesPerRowAlignment));
66  } else {
67    bytesPerRow = alignedMinBytesPerRow;
68  }
69
70  if (rowsPerImage !== undefined) {
71    assert(rowsPerImage >= mipSize[1]);
72  } else {
73    rowsPerImage = mipSize[1];
74  }
75
76  assert(isAligned(rowsPerImage, blockHeight));
77  const bytesPerSlice = bytesPerRow * (rowsPerImage / blockHeight);
78  const sliceSize =
79    bytesPerRow * (mipSize[1] / blockHeight - 1) + bytesPerBlock * (mipSize[0] / blockWidth);
80  const byteLength = bytesPerSlice * (mipSize[2] - 1) + sliceSize;
81
82  return {
83    bytesPerBlock,
84    byteLength: align(byteLength, kBufferCopyAlignment),
85    minBytesPerRow,
86    bytesPerRow,
87    rowsPerImage,
88    mipSize,
89  };
90}
91
92export function fillTextureDataWithTexelValue(
93  texelValue: ArrayBuffer,
94  format: SizedTextureFormat,
95  dimension: GPUTextureDimension,
96  outputBuffer: ArrayBuffer,
97  size: [number, number, number],
98  options: LayoutOptions = kDefaultLayoutOptions
99): void {
100  const { blockWidth, blockHeight, bytesPerBlock } = kSizedTextureFormatInfo[format];
101  assert(bytesPerBlock === texelValue.byteLength);
102
103  const { byteLength, rowsPerImage, bytesPerRow } = getTextureCopyLayout(
104    format,
105    dimension,
106    size,
107    options
108  );
109
110  assert(byteLength <= outputBuffer.byteLength);
111
112  const mipSize = getMipSizePassthroughLayers(dimension, size, options.mipLevel);
113
114  const texelValueBytes = new Uint8Array(texelValue);
115  const outputTexelValueBytes = new Uint8Array(outputBuffer);
116  for (let slice = 0; slice < mipSize[2]; ++slice) {
117    for (let row = 0; row < mipSize[1]; row += blockHeight) {
118      for (let col = 0; col < mipSize[0]; col += blockWidth) {
119        const byteOffset =
120          slice * rowsPerImage * bytesPerRow + row * bytesPerRow + col * texelValue.byteLength;
121        outputTexelValueBytes.set(texelValueBytes, byteOffset);
122      }
123    }
124  }
125}
126
127export function createTextureUploadBuffer(
128  texelValue: ArrayBuffer,
129  device: GPUDevice,
130  format: SizedTextureFormat,
131  dimension: GPUTextureDimension,
132  size: [number, number, number],
133  options: LayoutOptions = kDefaultLayoutOptions
134): {
135  buffer: GPUBuffer;
136  bytesPerRow: number;
137  rowsPerImage: number;
138} {
139  const { byteLength, bytesPerRow, rowsPerImage, bytesPerBlock } = getTextureCopyLayout(
140    format,
141    dimension,
142    size,
143    options
144  );
145
146  const buffer = device.createBuffer({
147    mappedAtCreation: true,
148    size: byteLength,
149    usage: GPUBufferUsage.COPY_SRC,
150  });
151  const mapping = buffer.getMappedRange();
152
153  assert(texelValue.byteLength === bytesPerBlock);
154  fillTextureDataWithTexelValue(texelValue, format, dimension, mapping, size, options);
155  buffer.unmap();
156
157  return {
158    buffer,
159    bytesPerRow,
160    rowsPerImage,
161  };
162}
163