feat: Minify shaders in production

[iwoplaza]

Note

This issue assumes that #2019 and #2414 are implemented, and builds upon them in certain sections

Let's consider the following code:

const Boid = d.struct({
  pos: d.vec3f,
  dir: d.vec3f,
});

const boids = root.createReadonly(d.arrayOf(Boid, 100));

const getAlignment = () => {
  'use gpu';
  const firstBoid = boids.$[0]!;
  const alignment = std.dot(firstBoid.dir, std.normalize(d.vec3f(0.1, 0.4, 1)));
  return alignment;
};

When calling tgpu.resolve([getAlignment]), the following gets generates:

struct Boid {
  pos: vec3f,
  dir: vec3f,
}

@group(0) @binding(0) var<storage, read> boids: array<Boid, 100>;

fn getAlignment() -> f32 {
  let firstBoid = (&boids[0i]);
  let alignment = dot((*firstBoid).dir, vec3f(0.09245003759860992, 0.3698001503944397, 0.9245003461837769));
  return alignment;
}

We would like to have the option to generate the following instead (new lines added for clarity):

struct z{pos:vec3f,dir:vec3f}
@group(0)@binding(0)var<storage,read>y:array<z,100>;
fn x()->f32{let e=(&y[0i]);let f=dot((*e).dir,vec3f(0.09245003759860992,0.3698001503944397,0.9245003461837769));return f;}

Proposed API

We could have the plugin and runtime library detect whether it's being run in production or in development mode, and choose to minify in production by default. In other scenarios, devs can use the following APIs to control minification:

// vite.config.mjs
import { defineConfig } from 'vite';
import typegpu from 'unplugin-typegpu/vite';

export default defineConfig({
  // letting the plugin know that it can optimize the code in a way that doesn't have
  // to preserve the original names.
  plugins: [typegpu({ minify: true })],
});

// Letting the shader generator know that it's supposed to minify things like whitespace

const root = await tgpu.init({ minify: true });
const code = tgpu.resolve([foo], { minify: true });

We could also inject some sort of config so that the root and tgpu.resolve APIs automatically conform to the setting used by the plugin, but then we wouldn't have the ability to minify some snippets, and not minify others.

Reasoning and implementation notes

During bundling, unplugin-typegpu currently emits the following (I highlighted the parts that have an effect on the names we use in the shaders):

  const Boid = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(d.struct({
    pos: d.vec3f,
    dir: d.vec3f
+ }), "Boid"));

  const boids = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(
    root.createReadonly(d.arrayOf(Boid, 100)),
+   "boids"
  ));

  const getAlignment = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
    "use gpu";
    const firstBoid = boids.$[0];
    const alignment = std.dot(firstBoid.dir, std.normalize(d.vec3f(.1, .4, 1)));
    return alignment;
  }), {
    v: 2,
+   name: "getAlignment",
+   ast: {"params":[],"body":[0,[[13,"firstBoid",[8,"boids.$",[5,"0"]]],[13,"alignment",[6,"std.dot",[[7,"firstBoid","dir"],[6,"std.normalize",[[6,"d.vec3f",[[5,"0.1"],[5,"0.4"],[5,"1"]]]]]]]],[10,"alignment"]]]},
+   ref: {
+     'boids.$': () => boids.$,
+     'std.dot': () => std.dot,
+     'std.normalize': () => std.normalize,
+     'd.vec3f': () => d.vec3f,
+   },
  }) && $.f)({}));

After the bundler minifies the code for production, those name strings will stay the same, meaning they will still contribute to generating the same shader code as they did during development, with their proper names. This is great for debugging, but if minifying is used a form of obfuscation, then the original names would still be retrievable.

Renaming definitions

With minification turned on, each time we auto-name a definition (constant, variable, struct, function, ...), we can generate a random 1 character name for them. We don't have to worry about clashes, because our naming system will ensure unique names across the program.

  const Boid = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(d.struct({
    pos: d.vec3f,
    dir: d.vec3f
- }), "Boid"));
+ }), "z"));

  const boids = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(
    root.createReadonly(d.arrayOf(Boid, 100)),
-   "boids"
+   "y"
  ));

  const getAlignment = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
    "use gpu";
    const firstBoid = boids.$[0];
    const alignment = std.dot(firstBoid.dir, std.normalize(d.vec3f(.1, .4, 1)));
    return alignment;
  }), {
    v: 2,
-    name: "getAlignment",
+    name: "x",
   ast: {"params":[],"body":[0,[[13,"firstBoid",[8,"boids.$",[5,"0"]]],[13,"alignment",[6,"std.dot",[[7,"firstBoid","dir"],[6,"std.normalize",[[6,"d.vec3f",[[5,"0.1"],[5,"0.4"],[5,"1"]]]]]]]],[10,"alignment"]]]},
   ref: {
     'boids.$': () => boids.$,
     'std.dot': () => std.dot,
     'std.normalize': () => std.normalize,
     'd.vec3f': () => d.vec3f,
   },
  }) && $.f)({}));

Renaming external references

External reference names do not directly contribute to what shader code gets generated, but they could be used to easily reverse-engineer original names of values from the outer scope. We rename these names with the same mechanism as we do for definitions.

  const getAlignment = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
    "use gpu";
    const firstBoid = boids.$[0];
    const alignment = std.dot(firstBoid.dir, std.normalize(d.vec3f(.1, .4, 1)));
    return alignment;
  }), {
    v: 2,
    name: "x",
-   ast: {"params":[],"body":[0,[[13,"firstBoid",[8,"boids.$",[5,"0"]]],[13,"alignment",[6,"std.dot",[[7,"firstBoid","dir"],[6,"std.normalize",[[6,"d.vec3f",[[5,"0.1"],[5,"0.4"],[5,"1"]]]]]]]],[10,"alignment"]]]},
+   ast: {"params":[],"body":[0,[[13,"firstBoid",[8,"a"      ,[5,"0"]]],[13,"alignment",[6,"b"      ,[[7,"firstBoid","dir"],[6,"c"            ,[[6,"d"      ,[[5,"0.1"],[5,"0.4"],[5,"1"]]]]]]]],[10,"alignment"]]]},
   ref: {
-     'boids.$': () => boids.$,
-     'std.dot': () => std.dot,
-     'std.normalize': () => std.normalize,
-     'd.vec3f': () => d.vec3f,
+     'a': () => boids.$,
+     'b': () => std.dot,
+     'c': () => std.normalize,
+     'd': () => d.vec3f,
   },
  }) && $.f)({}));

Renaming local definitions

Local definitions can be renamed with the same mechanism as all of the above.

  const getAlignment = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
    "use gpu";
    const firstBoid = boids.$[0];
    const alignment = std.dot(firstBoid.dir, std.normalize(d.vec3f(.1, .4, 1)));
    return alignment;
  }), {
    v: 2,
    name: "x",
-   ast: {"params":[],"body":[0,[[13,"firstBoid",[8,"a",[5,"0"]]],[13,"alignment",[6,"b",[[7,"firstBoid","dir"],[6,"c",[[6,"d",[[5,"0.1"],[5,"0.4"],[5,"1"]]]]]]]],[10,"alignment"]]]},
+   ast: {"params":[],"body":[0,[[13,"e"        ,[8,"a",[5,"0"]]],[13,"f"        ,[6,"b",[[7,"e"        ,"dir"],[6,"c",[[6,"d",[[5,"0.1"],[5,"0.4"],[5,"1"]]]]]]]],[10,"f"        ]]]},
   ref: {
     'a': () => boids.$,
     'b': () => std.dot,
     'c': () => std.normalize,
     'd': () => d.vec3f,
   },
  }) && $.f)({}));

Note how we don't rename the locals in the JS functions itself, only in the tinyest representation, as that's what going to be used to generate the shader. The JS function will be minified by the bundler, outside of our minification process.

Resulting shader

struct z {
  pos: vec3f,
  dir: vec3f,
}

@group(0) @binding(0) var<storage, read> y: array<z, 100>;

fn x() -> f32 {
  let e = (&y[0i]);
  let f = dot((*e).dir, vec3f(0.09245003759860992, 0.3698001503944397, 0.9245003461837769));
  return f;
}

Reducing whitespace

The tinyest representation of source code doesn't encode any whitespace, so the whitespace seen in the resulting WGSL shader above is recreated by the shader generator (WgslGenerator by default). This means the simplest solution would be to omit whitespace based on a minify: boolean parameter accepted by tgpu.init() and tgpu.resolve() APIs.

const code = tgpu.resolve([getAlignment], { minify: true });
// struct z{pos:vec3f,dir:vec3f}@group(0)@binding(0)var<storage,read>y:array<z,100>;fn x()->f32{let e=(&y[0i]);let f=dot((*e).dir,vec3f(0.09245003759860992,0.3698001503944397,0.9245003461837769));return f;}

Granularly opting out of renaming

With this setup, granularly opting out of renaming for specific definitions is very straight-forward and doesn't require any new APIs. We simply respect the names given explicitly by devs with .$name(...).

const Boid = d.struct({
  pos: d.vec3f,
  dir: d.vec3f,
}).$name("Boid"); // Explicit name!

// ...

const code = tgpu.resolve([getAlignment], { minify: true });
// struct Boid{pos:vec3f,dir:vec3f}@group(0)@binding(0)var<storage,read>y:array<Boid,100>;fn x()->f32{let e=(&y[0i]);let f=dot((*e).dir,vec3f(0.09245003759860992,0.3698001503944397,0.9245003461837769));return f;}

This is helpful when using TypeGPU with a framework that expects functions or other definitions to be named very specifically.