Merge pull request #10726 from CesiumGS/simpler-cache-key

Improve performance of `ShaderProgram.fromCache()`

Source/Renderer/ShaderCache.js

@@ -56,6 +56,11 @@ ShaderCache.prototype.replaceShaderProgram = function (options) {
   return this.getShaderProgram(options);
 };
 
+function toSortedJson(dictionary) {
+  const sortedKeys = Object.keys(dictionary).sort();
+  return JSON.stringify(dictionary, sortedKeys);
+}
+
 /**
  * Returns a shader program from the cache, or creates and caches a new shader program,
  * given the GLSL vertex and fragment shader source and attribute locations.
@@ -87,24 +92,33 @@ ShaderCache.prototype.getShaderProgram = function (options) {
     });
   }
 
-  const vertexShaderText = vertexShaderSource.createCombinedVertexShader(
-    this._context
-  );
-  const fragmentShaderText = fragmentShaderSource.createCombinedFragmentShader(
-    this._context
-  );
+  // Since ShaderSource.createCombinedXxxShader() can be expensive, use a
+  // simpler key for caching. This way, the function does not have to be called
+  // for each cache lookup.
+  const vertexShaderKey = vertexShaderSource.getCacheKey();
+  const fragmentShaderKey = fragmentShaderSource.getCacheKey();
+  // Sort the keys in the JSON to ensure a consistent order
+  const attributeLocationKey = defined(attributeLocations)
+    ? toSortedJson(attributeLocations)
+    : "";
+  const keyword = `${vertexShaderKey}:${fragmentShaderKey}:${attributeLocationKey}`;
 
-  const keyword =
-    vertexShaderText + fragmentShaderText + JSON.stringify(attributeLocations);
   let cachedShader;
-
   if (defined(this._shaders[keyword])) {
     cachedShader = this._shaders[keyword];
 
     // No longer want to release this if it was previously released.
     delete this._shadersToRelease[keyword];
   } else {
     const context = this._context;
+
+    const vertexShaderText = vertexShaderSource.createCombinedVertexShader(
+      context
+    );
+    const fragmentShaderText = fragmentShaderSource.createCombinedFragmentShader(
+      context
+    );
+
     const shaderProgram = new ShaderProgram({
       gl: context._gl,
       logShaderCompilation: context.logShaderCompilation,

Source/Renderer/ShaderSource.js

@@ -358,6 +358,26 @@ ShaderSource.replaceMain = function (source, renamedMain) {
   return source.replace(/void\s+main\s*\(\s*(?:void)?\s*\)/g, renamedMain);
 };
 
+/**
+ * Since {@link ShaderSource#createCombinedVertexShader} and
+ * {@link ShaderSource#createCombinedFragmentShader} are both expensive to
+ * compute, create a simpler string key for lookups in the {@link ShaderCache}.
+ *
+ * @returns {String} A key for identifying this shader
+ *
+ * @private
+ */
+ShaderSource.prototype.getCacheKey = function () {
+  // Sort defines to make the key comparison deterministic
+  const sortedDefines = this.defines.slice().sort();
+  const definesKey = sortedDefines.join(",");
+  const pickKey = this.pickColorQualifier;
+  const builtinsKey = this.includeBuiltIns;
+  const sourcesKey = this.sources.join("\n");
+
+  return `${definesKey}:${pickKey}:${builtinsKey}:${sourcesKey}`;
+};
+
 /**
  * Create a single string containing the full, combined vertex shader with all dependencies and defines.
  *

Specs/Renderer/ShaderCacheSpec.js

@@ -1,4 +1,4 @@
-import { ShaderCache } from "../../Source/Cesium.js";
+import { ShaderCache, ShaderSource } from "../../Source/Cesium.js";
 import createContext from "../createContext.js";
 
 describe(
@@ -72,6 +72,17 @@ describe(
         "attribute vec4 position; void main() { gl_Position = position; }";
       const fs = "void main() { gl_FragColor = vec4(1.0); }";
 
+      // These functions can be expensive for large shaders, so they should
+      // only be called the first time a shader is created.
+      spyOn(
+        ShaderSource.prototype,
+        "createCombinedVertexShader"
+      ).and.callThrough();
+      spyOn(
+        ShaderSource.prototype,
+        "createCombinedFragmentShader"
+      ).and.callThrough();
+
       const cache = new ShaderCache(context);
       const sp = cache.getShaderProgram({
         vertexShaderSource: vs,
@@ -92,6 +103,75 @@ describe(
       expect(sp._cachedShader.count).toEqual(2);
       expect(cache.numberOfShaders).toEqual(1);
 
+      expect(
+        ShaderSource.prototype.createCombinedVertexShader
+      ).toHaveBeenCalledTimes(1);
+      expect(
+        ShaderSource.prototype.createCombinedFragmentShader
+      ).toHaveBeenCalledTimes(1);
+
+      sp.destroy();
+      sp2.destroy();
+      cache.destroyReleasedShaderPrograms();
+
+      expect(sp.isDestroyed()).toEqual(true);
+      expect(cache.numberOfShaders).toEqual(0);
+
+      cache.destroy();
+    });
+
+    it("cache handles unordered attributeLocations dictionary", function () {
+      const vs =
+        "attribute vec4 position; void main() { gl_Position = position; }";
+      const fs = "void main() { gl_FragColor = vec4(1.0); }";
+
+      // Create a case where JSON.stringify(x) may produce two different results
+      // without sorting the keys
+      const attributeLocations = {
+        position: 0,
+        normal: 1,
+        color: 2,
+      };
+      const attributeLocationsReordered = {
+        color: 2,
+        position: 0,
+        normal: 1,
+      };
+
+      // These functions can be expensive for large shaders, so they should
+      // only be called the first time a shader is created.
+      spyOn(
+        ShaderSource.prototype,
+        "createCombinedVertexShader"
+      ).and.callThrough();
+      spyOn(
+        ShaderSource.prototype,
+        "createCombinedFragmentShader"
+      ).and.callThrough();
+
+      const cache = new ShaderCache(context);
+      const sp = cache.getShaderProgram({
+        vertexShaderSource: vs,
+        fragmentShaderSource: fs,
+        attributeLocations: attributeLocations,
+      });
+      const sp2 = cache.getShaderProgram({
+        vertexShaderSource: vs,
+        fragmentShaderSource: fs,
+        attributeLocations: attributeLocationsReordered,
+      });
+
+      expect(sp).toBe(sp2);
+      expect(sp._cachedShader.count).toEqual(2);
+      expect(cache.numberOfShaders).toEqual(1);
+
+      expect(
+        ShaderSource.prototype.createCombinedVertexShader
+      ).toHaveBeenCalledTimes(1);
+      expect(
+        ShaderSource.prototype.createCombinedFragmentShader
+      ).toHaveBeenCalledTimes(1);
+
       sp.destroy();
       sp2.destroy();
       cache.destroyReleasedShaderPrograms();

Specs/Renderer/ShaderSourceSpec.js

@@ -91,4 +91,46 @@ describe("Renderer/ShaderSource", function () {
     expect(clone.pickColorQualifier).toEqual(source.pickColorQualifier);
     expect(clone.includeBuiltIns).toEqual(source.includeBuiltIns);
   });
+
+  it("creates cache key for empty shader", function () {
+    const source = new ShaderSource();
+    expect(source.getCacheKey()).toBe(":undefined:true:");
+  });
+
+  it("creates cache key", function () {
+    const source = new ShaderSource({
+      defines: ["A", "B", "C"],
+      sources: ["void main() { gl_FragColor = vec4(1.0); }"],
+      pickColorQualifier: "varying",
+      includeBuiltIns: false,
+    });
+
+    expect(source.getCacheKey()).toBe(
+      "A,B,C:varying:false:void main() { gl_FragColor = vec4(1.0); }"
+    );
+  });
+
+  it("uses sorted list of defines in cache key", function () {
+    const defines1 = ["A", "B", "C"];
+    const defines2 = ["B", "C", "A"];
+
+    const source1 = new ShaderSource({ defines: defines1 });
+    const source2 = new ShaderSource({ defines: defines2 });
+    const key1 = source1.getCacheKey();
+    expect(key1).toBe(source2.getCacheKey());
+    expect(key1).toBe("A,B,C:undefined:true:");
+  });
+
+  it("cache key includes all sources", function () {
+    const source = new ShaderSource({
+      sources: [
+        "vec4 getColor() { return vec4(1.0, 0.0, 0.0, 1.0); }",
+        "void main() { gl_fragColor = getColor(); }",
+      ],
+    });
+
+    expect(source.getCacheKey()).toBe(
+      ":undefined:true:vec4 getColor() { return vec4(1.0, 0.0, 0.0, 1.0); }\nvoid main() { gl_fragColor = getColor(); }"
+    );
+  });
 });