fix(@typegpu/gl): support array texture loads

packages/typegpu-gl/src/glslGenerator.ts

@@ -399,6 +399,22 @@ function isF32VecfSchema(
   );
 }
 
+const UNSIGNED_TO_SIGNED = {
+  u32: d.i32,
+  vec2u: d.vec2i,
+  vec3u: d.vec3i,
+  vec4u: d.vec4i,
+} as const;
+
+function sampledTextureGlslType(data: d.BaseData): string | undefined {
+  if (data.type !== 'texture_2d' && data.type !== 'texture_2d_array') {
+    return undefined;
+  }
+  const sampleType = (data as d.WgslTexture2d | d.WgslTexture2dArray).sampleType.type;
+  const prefix = sampleType === 'u32' ? 'u' : sampleType === 'i32' ? 'i' : '';
+  return data.type === 'texture_2d' ? `${prefix}sampler2D` : `${prefix}sampler2DArray`;
+}
+
 const HELPERS = {
   // TODO(#2821): Make signature more accurate when std.sign and std.abs
   // accept a wider union
@@ -411,6 +427,11 @@ const HELPERS = {
     'use gpu';
     return std.select(coords, d.vec2f(coords.x, 1 - coords.y), flip);
   },
+  flipTexelYConditionally: (coords: d.v2i, height: number, flip: boolean): d.v2i => {
+    'use gpu';
+    // Integer mix() requires GLSL ES 3.10, so selecting on the scalar component instead
+    return d.vec2i(coords.x, std.select(coords.y, height - 1 - coords.y, flip));
+  },
   // WGSL specifies floor(x + 0.5); GLSL round leaves half ties implementation-defined.
   pack4x8unorm: (value: d.v4f): number => {
     'use gpu';
@@ -579,10 +600,8 @@ export class GlslGenerator extends WgslGenerator {
         return snip(options.id, options.dataType, 'handle');
       }
 
-      if (options.dataType.type === 'texture_2d') {
-        const sampleType = (options.dataType as d.WgslTexture2d).sampleType.type;
-        const glslType =
-          sampleType === 'u32' ? 'usampler2D' : sampleType === 'i32' ? 'isampler2D' : 'sampler2D';
+      const glslType = sampledTextureGlslType(options.dataType);
+      if (glslType) {
         let flipId = this.#crossShaderStageState.textureFlipIdentifiers.get(options.id);
         if (!flipId) {
           flipId = this.ctx.makeUniqueIdentifier(`${options.id}_flipY`, 'global');
@@ -614,13 +633,9 @@ export class GlslGenerator extends WgslGenerator {
   }
 
   override emitTypeAnnotation(data: d.BaseData): string {
-    if (data.type === 'texture_2d') {
-      const sampleType = (data as d.WgslTexture2d).sampleType.type;
-      return sampleType === 'u32'
-        ? 'usampler2D'
-        : sampleType === 'i32'
-          ? 'isampler2D'
-          : 'sampler2D';
+    const textureGlslType = sampledTextureGlslType(data);
+    if (textureGlslType) {
+      return textureGlslType;
     }
 
     if (data.type === 'sampler' || data.type === 'sampler_comparison') {
@@ -775,18 +790,48 @@ export class GlslGenerator extends WgslGenerator {
     }
 
     if (name === 'textureLoad') {
-      const [texture, coords, level] = args;
-      if (!texture || !coords || !level) {
+      const [texture, coords, ...rest] = args;
+      if (!texture || !coords) {
+        throw new Error(`Invalid number of arguments for '${name}'`);
+      }
+
+      const textureType = texture.dataType as d.WgslTexture | d.WgslStorageTexture;
+      if (textureType.type.startsWith('texture_storage_')) {
+        throw new Error(`Storage textures are not supported when generating GLSL ('${name}').`);
+      }
+
+      const isTextureArray = textureType.dimension === '2d-array';
+      // Array textures receive (coords, arrayIndex, level), the rest receive (coords, level | sampleIndex)
+      const [arrayIndex, level] = isTextureArray ? rest : [undefined, ...rest];
+      if (!level || (isTextureArray && !arrayIndex)) {
         throw new Error(`Invalid number of arguments for '${name}'`);
       }
+
       const textureName = this.ctx.resolveSnippet(texture).value;
-      const coordsValue = this.ctx.resolveSnippet(coords).value;
-      const levelValue = this.ctx.resolveSnippet(level).value;
+      const levelValue = this.ctx.resolveSnippet(this.#toSigned(level)).value;
       const flipId = this.#crossShaderStageState.textureFlipIdentifiers.get(textureName);
-      const orientedCoords = flipId
-        ? `ivec2(${coordsValue}.x, ${flipId} ? textureSize(${textureName}, ${levelValue}).y - 1 - int(${coordsValue}.y) : int(${coordsValue}.y))`
-        : coordsValue;
-      return `texelFetch(${textureName}, ${orientedCoords}, ${levelValue})`;
+
+      let texelCoords: Snippet = this.#toSigned(coords);
+      if (flipId && (textureType.dimension === '2d' || isTextureArray)) {
+        // Multisampled textures have no mip levels, so the size query takes no level
+        const size = (textureType as d.WgslTexture).multisampled
+          ? `textureSize(${textureName})`
+          : `textureSize(${textureName}, ${levelValue})`;
+        texelCoords =
+          this._callShellless(HELPERS.flipTexelYConditionally, [
+            texelCoords,
+            snip(`${size}.y`, d.i32, 'runtime'),
+            snip(flipId, d.bool, 'uniform'),
+          ]) ?? texelCoords;
+      }
+      if (isTextureArray && arrayIndex) {
+        // GLSL addresses array layers through the last coordinate component
+        texelCoords = this.typeInstantiation(d.vec3i, [texelCoords, arrayIndex]);
+      }
+
+      const texel = `texelFetch(${textureName}, ${this.ctx.resolveSnippet(texelCoords).value}, ${levelValue})`;
+      // Depth textures are fetched through regular samplers, with depth in the first channel
+      return textureType.type.startsWith('texture_depth_') ? `${texel}.x` : texel;
     }
 
     if (name === 'textureDimensions') {
@@ -872,6 +917,17 @@ export class GlslGenerator extends WgslGenerator {
     return super.emitCall(name, templateParams, args);
   }
 
+  /**
+   * GLSL's texelFetch only accepts signed integer coordinates, levels and sample indices.
+   */
+  #toSigned(snippet: Snippet): Snippet {
+    const signedSchema =
+      snippet.dataType === UnknownData
+        ? undefined
+        : UNSIGNED_TO_SIGNED[snippet.dataType.type as keyof typeof UNSIGNED_TO_SIGNED];
+    return signedSchema ? this.typeInstantiation(signedSchema, [snippet]) : snippet;
+  }
+
   #orientedTextureOffset(offset: Snippet, flipId: string | undefined): string {
     const value = this.ctx.resolveSnippet(offset).value;
     return flipId ? `(${flipId} ? ivec2(${value}.x, -${value}.y) : ${value})` : value;

packages/typegpu-gl/tests/glslGenerator.test.ts

@@ -1,7 +1,11 @@
 import { describe, expect, vi } from 'vitest';
 import { tgpu, d, std } from 'typegpu';
 import { dualGlOptions, glOptions } from '@typegpu/gl';
-import { translateWgslTypeToGlsl } from '../src/glslGenerator.ts';
+import {
+  CrossShaderStageState,
+  GlslGenerator,
+  translateWgslTypeToGlsl,
+} from '../src/glslGenerator.ts';
 import { it } from './utils/extendedTest.ts';
 
 describe('translateWgslTypeToGlsl', () => {
@@ -409,6 +413,172 @@ describe('GlslGenerator - standard function calls', () => {
     `);
   });
 
+  it('combines coordinates and array index when loading a 2d-array texture', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet('palette', d.texture2dArray(), 'handle');
+
+    function loadArrayTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec2u(2, 3), d.u32(4), 1);
+    }
+
+    expect(tgpu.resolve([loadArrayTexel], glOptions())).toMatchInlineSnapshot(`
+      "vec4 loadArrayTexel() {
+        return texelFetch(palette, ivec3(ivec2(uvec2(2, 3)), 4u), 1);
+      }"
+    `);
+  });
+
+  it('preserves the flip adaptation when loading a 2d-array texture', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet('palette', d.texture2dArray(), 'handle');
+    const state = new CrossShaderStageState();
+    state.textureFlipIdentifiers.set('palette', 'palette_flipY');
+
+    function loadArrayTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec2i(2, 3), 4, 1);
+    }
+
+    expect(
+      tgpu.resolve([loadArrayTexel], {
+        unstable_shaderGenerator: new GlslGenerator('neutral', state),
+      }),
+    ).toMatchInlineSnapshot(`
+      "ivec2 flipTexelYConditionally(ivec2 coords, int height, bool flip) {
+        return ivec2(coords.x, (flip ? ((height - 1) - coords.y) : coords.y));
+      }
+
+      vec4 loadArrayTexel() {
+        return texelFetch(palette, ivec3(flipTexelYConditionally(ivec2(2, 3), textureSize(palette, 1).y, palette_flipY), 4), 1);
+      }"
+    `);
+  });
+
+  it('converts unsigned 3d load coordinates to signed coordinates', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet('volume', d.texture3d(), 'handle');
+
+    function loadTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec3u(2, 3, 4), 1);
+    }
+
+    expect(tgpu.resolve([loadTexel], glOptions())).toMatchInlineSnapshot(`
+      "vec4 loadTexel() {
+        return texelFetch(volume, ivec3(uvec3(2, 3, 4)), 1);
+      }"
+    `);
+  });
+
+  it('loads depth 2d-array textures through the first channel', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet(
+      'depthPalette',
+      d.textureDepth2dArray(),
+      'handle',
+    );
+
+    function loadArrayTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec2i(2, 3), 4, 1);
+    }
+
+    expect(tgpu.resolve([loadArrayTexel], glOptions())).toMatchInlineSnapshot(`
+      "float loadArrayTexel() {
+        return texelFetch(depthPalette, ivec3(ivec2(2, 3), 4), 1).x;
+      }"
+    `);
+  });
+
+  it('converts unsigned load levels to signed levels', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet('palette', d.texture2d(), 'handle');
+    const state = new CrossShaderStageState();
+    state.textureFlipIdentifiers.set('palette', 'palette_flipY');
+
+    function loadTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec2u(1, 2), d.u32(3));
+    }
+
+    expect(
+      tgpu.resolve([loadTexel], {
+        unstable_shaderGenerator: new GlslGenerator('neutral', state),
+      }),
+    ).toMatchInlineSnapshot(`
+      "ivec2 flipTexelYConditionally(ivec2 coords, int height, bool flip) {
+        return ivec2(coords.x, (flip ? ((height - 1) - coords.y) : coords.y));
+      }
+
+      vec4 loadTexel() {
+        return texelFetch(palette, flipTexelYConditionally(ivec2(uvec2(1, 2)), textureSize(palette, int(3u)).y, palette_flipY), int(3u));
+      }"
+    `);
+  });
+
+  it('does not flip 3d load coordinates', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet('volume', d.texture3d(), 'handle');
+    const state = new CrossShaderStageState();
+    state.textureFlipIdentifiers.set('volume', 'volume_flipY');
+
+    function loadTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec3i(2, 3, 4), 1);
+    }
+
+    expect(
+      tgpu.resolve([loadTexel], {
+        unstable_shaderGenerator: new GlslGenerator('neutral', state),
+      }),
+    ).toMatchInlineSnapshot(`
+      "vec4 loadTexel() {
+        return texelFetch(volume, ivec3(2, 3, 4), 1);
+      }"
+    `);
+  });
+
+  it('annotates 2d-array texture parameters with array sampler types', () => {
+    const paletteSchema = d.texture2dArray(d.u32);
+    const texture = tgpu['~unstable'].rawCodeSnippet('palette', paletteSchema, 'handle');
+
+    const loadFrom = (tex: d.Infer<typeof paletteSchema>, coords: d.v2i) => {
+      'use gpu';
+      return std.textureLoad(tex, coords, 4, 1);
+    };
+
+    function loadArrayTexel() {
+      'use gpu';
+      return loadFrom(texture.$, d.vec2i(2, 3));
+    }
+
+    expect(tgpu.resolve([loadArrayTexel], glOptions())).toMatchInlineSnapshot(`
+      "uvec4 loadFrom(usampler2DArray tex, ivec2 coords) {
+        return texelFetch(tex, ivec3(coords, 4), 1);
+      }
+
+      uvec4 loadArrayTexel() {
+        return loadFrom(palette, ivec2(2, 3));
+      }"
+    `);
+  });
+
+  it('throws when loading from storage textures', () => {
+    const texture = tgpu['~unstable'].rawCodeSnippet(
+      'storagePalette',
+      d.textureStorage2dArray('rgba8unorm', 'read-only'),
+      'handle',
+    );
+
+    function loadTexel() {
+      'use gpu';
+      return std.textureLoad(texture.$, d.vec2i(2, 3), 4);
+    }
+
+    expect(() => tgpu.resolve([loadTexel], glOptions())).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:loadTexel
+      - fn*:loadTexel()
+      - fn:textureLoad: Storage textures are not supported when generating GLSL ('textureLoad').]
+    `);
+  });
+
   it('translates textureSample() to texture() with the combined sampler', () => {
     const texture = tgpu['~unstable'].rawCodeSnippet('palette', d.texture2d(), 'handle');
     const sampler = tgpu['~unstable'].rawCodeSnippet('paletteSampler', d.sampler(), 'handle');