Make shaders tree-shakable (inspired by threeify's approach)

[bhouston]

Is your feature request related to a problem? Please describe.

Currently the shader chunks are taking up a significant amount of space in the Three.js final bundle. 126KB of 666KB - so roughly 20% of the space.

This is clear from this image shared on Twitter: https://twitter.com/threejs_org/status/1382639620119216131

This is the first of two feature requests on how to reduce further Three.j's build sizes based on my learnings from threeify. The other is here: https://github.com/mrdoob/three.js/issues/21667

Describe the solution you'd like

To make the shader chunks tree shakable, one can take inspiration from how threeify did this. Threeify can have build sizes as small as 12KB compressed JS. https://threeify.org/examples/brdf_clear_coat_specular

The way it did this is:

• The glsl files in the project are transpiled to *.js files where the "#include " statements are replaced with JavaScript import directives. This happens automatically behind the scenes.
• When using .glsl files you can import then directly as you would any other module. They are actually .js files but because many include systems can just auto append .js, you can include them as .glsl. Thus your code looks like this: https://github.com/threeify/threeify/blob/master/src//examples/brdfs/clearcoat/index.ts https://github.com/threeify/threeify/tree/master/src/examples/brdfs/clearcoat
• I wrote this transpiler to do this live .glsl to .js translation: https://github.com/threeify/threeify-glsl-transpiler
• Because everything is based on JavaScript imports, everything is tree shaken just to the code you are using.
• There is no need for a run-time include system for the GLSL files because they just piggy back on the JS include system.
• The glsl-to-JS transpiler I wrote even has options to minimize the produced GLSL files so that you can remove debug information. This can further reduce the resulting bundle size when you just want production code.

Describe alternatives you've considered

I would suggest the above.

Additional context

• There is no dependent upon typescript for this solution.
• It is likely the glsl-transpiler can be directly used in three.js for this purpose. Maybe you could do it inline so that the .glsl.js files are directly beside the source .glsl files?
• If the build/server system doesn't support .glsl completion to .glsl.js, you can just decide to import *.glsl.js files directly if needed. A little less clean but all the benefits of this system work the same.

Example of transpile

An original rgbe.glsl file:

#pragma once
#pragma include "../../math/math.glsl"

vec4 rgbeToLinear( in vec4 value ) {
    return vec4( value.rgb * exp2( value.a * 255.0 - 128.0 ), 1.0 );
}

vec4 linearToRGBE( in vec4 value ) {
    float maxComponent = max( max( value.r, value.g ), value.b );
    float fExp = clamp( ceil( log2( maxComponent ) ), -128.0, 127.0 );
    return vec4( value.rgb / exp2( fExp ), ( fExp + 128.0 ) / 255.0 );
}

Will be transformed into a rgbe.glsl.js JavaScript module:

import _math_glsl from "../../math/math.glsl.js";

export default /* glsl */ `
#ifndef _rgbe_glsl // start of include guard
#define _rgbe_glsl

${_math_glsl}

vec4 rgbeToLinear( in vec4 value ) {
    return vec4( value.rgb * exp2( value.a * 255.0 - 128.0 ), 1.0 );
}

vec4 linearToRGBE( in vec4 value ) {
    float maxComponent = max( max( value.r, value.g ), value.b );
    float fExp = clamp( ceil( log2( maxComponent ) ), -128.0, 127.0 );
    return vec4( value.rgb / exp2( fExp ), ( fExp + 128.0 ) / 255.0 );
}

#endif // end of include guard
`;

Advanced

• It is possible to do this type of include system on a function level rather than just on a file level. That is more advanced by it is similar to how you can import just a few functions from a JS module rather than everything. This could be extended to glsl with more advanced transpilation options. But that can be done in the future.

[marcofugaro]

Really interesting! So with your plugin this next bit of code can be moved to run at compile time, right? In rollup where currently the glsl minification happens. Is that what you're suggesting?

https://github.com/mrdoob/three.js/blob/a80cccc98780666857b2239f1c9e10aedf86dfd9/src/renderers/webgl/WebGLProgram.js#L196-L254

2nd question, does #pragma include "../../math/math.glsl" work in the main() function as well?

[bhouston]

So with your plugin this next bit of code can be moved to run at compile time, right? In rollup where currently the glsl minification happens. Is that what you're suggesting?

Yes, that could happen at compile time. Or it could happen at run-time. That is a choice.

2nd question, does #pragma include "../../math/math.glsl" work in the main() function as well?

Basically any "#pragma include" would be replaced with the equivalent "import" statement. It is a one to one mapping using the same paths. Because every glsl files gets converted to *.glsl.js file.

[bhouston]

To be honest, I think unrolling of loops should happen at run-time because it is a form of compression. Doing so before run-time will increase the bundle size.

[marcofugaro]

I see, that makes sense.

For this to work effectively we would need to make ShaderLib tree-shakeable as well, do you have any tips on how to do that?

I tested it by converting it to esmodules like this:

const basic = {

  uniforms: mergeUniforms( [
    UniformsLib.common,
    UniformsLib.specularmap,
    UniformsLib.envmap,
    UniformsLib.aomap,
    UniformsLib.lightmap,
    UniformsLib.fog
  ] ),

  vertexShader: ShaderChunk.meshbasic_vert,
  fragmentShader: ShaderChunk.meshbasic_frag

}

const lambert = {

  uniforms: mergeUniforms( [
    UniformsLib.common,
    UniformsLib.specularmap,
    UniformsLib.envmap,
    UniformsLib.aomap,
    UniformsLib.lightmap,
    UniformsLib.emissivemap,
    UniformsLib.fog,
    UniformsLib.lights,
    {
      emissive: { value: new Color( 0x000000 ) }
    }
  ] ),

  vertexShader: ShaderChunk.meshlambert_vert,
  fragmentShader: ShaderChunk.meshlambert_frag

}

// ...

export { basic, lambert, ... }

However tree-shaking doesn't work for 2 reasons:

1. the call to mergeUniforms() in the object literal
2. this next line which requires every material in ShaderLib to be included: https://github.com/mrdoob/three.js/blob/a80cccc98780666857b2239f1c9e10aedf86dfd9/src/renderers/webgl/WebGLPrograms.js#L140

Do you have any ideas on how to rearrange the ShaderLib code?

[bhouston]

This seems tree-shakable no?

import meshbasic_vert from 'meshbasic_vert.glsl';
import meshbasic_frag from 'meshbasic_vert.glsl';

const basic = {

  uniforms: mergeUniforms( [
    UniformsLib.common,
    UniformsLib.specularmap,
    UniformsLib.envmap,
    UniformsLib.aomap,
    UniformsLib.lightmap,
    UniformsLib.fog
  ] ),

  vertexShader: meshbasic_vert,
  fragmentShader: meshbasic_frag

}

To deal with this:

const shader = ShaderLib[ shaderID ]; 

We need to remove ShaderLib. Basically if you import PhysicalMaterial or BasicMaterial, it should then bring in the necessary glsl as well via an import. And then that glsl goes along with material somehow. So that it is only imported if used.

This concept of a global ShaderLib or ShaderChunks is fundamentally incompatible with tree shaking.

[raphaelameaume]

I agree, we could still have an index.js file at the root of the shader folder exporting all the shaders so we can do

import * as ShaderLib from "path/to/all/shaders/index.js"

I think we should be able to do so in the library and it would still be tree-shakable:

import { meshbasic_vert, meshphysical_vert, ... } from "path/to/all/shaders/index.js";

[bhouston]

I think that we could modify BasicMaterial to be something like:

import meshbasic_vert from 'shaders/meshbascic_vert.glsl';
import meshbasic_frag from 'shaders/meshbasic_frag.glsl';

export class BasicMaterial {

  vertexShader = meshbasic_vert;
  fragmentShader = meshbasic_frag;

...
};

This way when you use BasicMaterial, it brings along the shaders it needs.

[marcofugaro]

This seems tree-shakable no?

It's not, I have tested this, there is the mergeUniforms() which modifies the object literal, so it's not pure anymore.

---

I think that we could modify BasicMaterial to be something like:

import meshbasic_vert from 'shaders/meshbascic_vert.glsl';
import meshbasic_frag from 'shaders/meshbasic_frag.glsl';

export class BasicMaterial {

  vertexShader = meshbasic_vert;
  fragmentShader = meshbasic_frag;

   uniforms = mergeUniforms( [
    UniformsLib.common,
    UniformsLib.specularmap,
    UniformsLib.envmap,
    UniformsLib.aomap,
    UniformsLib.lightmap,
    UniformsLib.fog
  ] );

  // ...

};

I agree, this is the way to go to enable tree-shaking. ShaderLib needs to go, users will be able to access the shaders/uniforms from the class. Maybe we could even make them static.

However we can still leave ShaderChunk (converting it to modules), as long as it's not used directly. It can be still exported like MathUtils without breaking the API.

https://github.com/mrdoob/three.js/blob/52ec6876d693670629f9618ba7d37f05be5de227/src/Three.js#L109

[bhouston]

In Threeify I never created a "uniform library." Instead I just make sure that introspection on the shaders worked really well. So that the uniforms were introspected automatically with their types and names.

Then you could just pass in a uniform set and if the names matched they would set their values:

  const uniforms = {
    localToWorld: new Matrix4(),
    worldToView: makeMatrix4Translation(new Vector3(0, 0, -3.0)),
    viewToScreen: makeMatrix4PerspectiveFov(25, 0.1, 4.0, 1.0, canvasFramebuffer.aspectRatio),
    cubeMap: lambertianCubeMap,
  };

  renderBufferGeometry(canvasFramebuffer, program, uniforms, bufferGeometry, depthTestState);

Basically I have a class called ProgramUniform that is created for each uniform once you compile a shader into in Program. It introspects itself on creation and then you can just set it with a value and it sets properly. Having a shader full introspectable is really nice: https://github.com/threeify/threeify/blob/master/src/lib/renderers/webgl/programs/ProgramUniform.ts

In my opinion UniformLib is mostly useless. It sets defaults but most of them are 0 or null, thus not actually needed. Other than that it is required to mirror the shader you'll use but that seems to be unnecessary book keeping. There are only a few defaults that are actually useful and I am not sure that justifies its complexity of bookkeeping:

https://github.com/mrdoob/three.js/blob/dev/src/renderers/shaders/UniformsLib.js

[marcofugaro]

@bhouston so how do you think it would look like in three? where would the simple a simple diffuse uniform (starting from color) be created for example?

[bhouston]

How about we separate it out into two steps? We just replace the shader chunks first and leave UniformLib and ShaderLib alone. And then if that works, we replace UniformLib and ShaderLib separately. My concern is that if we try to refactor everything all at once, it won't get started?

(I think that the best case scenario is to minimize the differences between the Material classes and the underlying program uniforms. I think that Three.js's WebGLProgram(s) classes have way to much specialized knowledge of shaders in them, also so does a lot of WebGLRenderer. But this can be fixed incrementally.)

[marcofugaro]

Sure I agree, let's handle shaders first.

Regarding ShaderChunks, another breaking change that came to me is the onBeforeCompile method.

People are using it like this, to attach custom shader code on certain points of the shader:

material.onBeforeCompile = (shader) => {
  shader.vertexShader = shader.vertexShader.replace('#include <begin_vertex>', `
    #include <begin_vertex>
    // custom code here...
    transformed += ...
  `)
}

I guess, if we move the include logic to compile time, users won't be able to do this anymore? Shaders will appear unwrapped as well, is that right?

[donmccurdy]

Tree-shaking shaders sounds like an unambiguous improvement, +1 for the direction.

To be honest, I think unrolling of loops should happen at run-time because it is a form of compression. Doing so before run-time will increase the bundle size.

I'm fine with either one, but it's worth noting that the numbers quoted above are minified but not gzipped. Unrolled loops are exactly the sort of thing that gzip will compress dramatically, and (in my opinion) "minified + gzipped" is the better metric.

---

One "catch" here — many users rely on a loader. GLTFLoader, for example, may output MeshBasicMaterial, MeshStandardMaterial, MeshPhysicalMaterial, or a custom ShaderMaterial depending on the model it loads, so none of those can be tree-shaken when the loader is included.

We are trying to move GLTFLoader's extensions into implementations a "plugin system", but currently all extensions are enabled by default. Perhaps we could export two versions of GLTFLoader from the same file (e.g. GLTFLoader and LiteGLTFLoader?) only one of which includes those exports?

Or I've taken an opt-in approach in the glTF-Transform project, shown here:

https://gltf-transform.donmccurdy.com/extensions.html#installation

[bhouston]

We are trying to move GLTFLoader's extensions into implementations a "plugin system", but currently all extensions are enabled by default. Perhaps we could export two versions of GLTFLoader from the same file (e.g. GLTFLoader and LiteGLTFLoader?) only one of which includes those exports?

If your glTF load properly imports those materials, then they will be included in a build if you wanted to build with extreme tree-shaking. That would be how it works.

If we took the approach I am proposing here, we could merge in all of the examples into the main three.js /src tree and it wouldn't matter. Now we do not have to do that, but maybe some of those post effects should be considered to be part of the three.js /src tree?

When one does tree-shaking properly and design the library around it, there is less distinctions require between core and auxillary code. In some ways, that may be freeing.

[marcofugaro]

I think the main blocker of this feature is that shaders inside onBeforeCompile will be fully expanded (no #includes). I don't know how the userbase will react to it. The other parts are feasible.

[mrdoob]

I'm not sure this change is worth the effort and ramifications at this point in time...

WebGPURenderer will use @sunag's NodeMaterial which is already modularized.

[mrdoob]

https://github.com/mrdoob/three.js/issues/21667#issuecomment-850920437