[Advanced Recipe] Mechanism for tracking resource use (introspection)

[iwoplaza]

I received a feature request, to be able to inspect whether a function itself (or transitively) used a slot/accessor/struct, etc.. The motivation is e.g. being able to determine if a shader uses a provided resource, and if not, not bothering to populate it with expensive to compute data, as an optimization trick.

It seems like it's a utility that can be constructed out of the public API surface today. Let's look at how the API of such utility could look like:

import { tgpu, d, std } from 'typegpu';
import { track, getTrackedResources } from './track.ts';

const textureAccess = tgpu.accessor(d.texture2d(d.f32));

function sampleTexture(uv: d.v2f): d.v4f {
  'use gpu';
  // Calling the `track` function at comptime, which adds
  // the `textureAccess` accessor to the set of tracked
  // resources when the `sampleTexture` function is used
  // in the shader.
  track(textureAccess);

  return std.textureSample(textureAccess.$, uv);
}

function main() {
  'use gpu';
  return sampleTexture(d.vec2f());
}

const resources = getTrackedResources([main]);
// The function `main` does use `sampleTexture`, so the value of 
// `resources` will be `Set(1) { textureAccess }`

If we make track return it's input as well, then we can even write:

function sampleTexture(uv: d.v2f): d.v4f {
  'use gpu';
  return std.textureSample(track(textureAccess).$, uv);
}

By hiding the accessor, and only exposing it through functions or getters that call track, we know when a shader used a particular resource.

Here is how track and getTrackedResources can be implemented:

import { tgpu, std } from 'typegpu';

// Can be provided with a Set where tracked resources should be added to.
const resourceUsageSlot = tgpu.slot<{ usedResources: Set<unknown> } | null>(null);

const track = tgpu.comptime(<T>(resource: T): T => {
  if (std.getTargetShaderLanguage() === undefined) {
    // If we're not resolving a shader, we don't need to track anything
    return resource;
  }

  // We track the resource if the Set has been provided
  resourceUsageSlot.$?.usedResources.add(resource);
  return resource;
});

function getTrackedResources(resolvables: ResolvableObject[]): Set<unknown> {
  const usedResources = new Set();

  // Resolving the objects, providing the `resourceUsageSlot` with a Set where
  // each tracked resource can be added to
  void tgpu.resolve(resolvables, {
    config: (cfg) => cfg.with(resourceUsageSlot, { usedResources }),
  });

  return usedResources;
}

Generalizing

This mechanism can be used for determining whether any specific branch was taken during shader generation, not just to collect a set of objects.

[aleksanderkatan]

I think that I'd rather have these resources bound to resolution, not to function.

import { tgpu, d, std } from 'typegpu';
import { track, getTrackedResources } from './track.ts';

const textureAccess = tgpu.accessor(d.texture2d(d.f32));
let arbitraryCounter = 0;

function sampleTexture(uv: d.v2f): d.v4f {
  'use gpu';
  if (arbitraryCounter === 0) {
    track(textureAccess);
    return std.textureSample(textureAccess.$, uv);
  }
  return d.vec4f();
}

function main() {
  'use gpu';
  return sampleTexture(d.vec2f());
}

// using main in a shader
arbitraryCounter++;
const resources = getTrackedResources([main]); // texture not included!!

In this case, this is debugable by the user, but there could be worse cases, like using a comptime that changes its value with each call, where user will have no idea what did they do wrong.
Not to mention the extra resolve's overhead.

I don't recall, why didn't we want to just return the used resources from resolve?
If the reason was that the pipeline couldn't do it, I think we could easily give the pipeline an extra prop.
I think that the approach of "was this used in this pipeline/resolve?" is cleaner than "was this used in either overload of this particular function, with the current context?"