Fix depth in SDSM shadows (jMonkeyEngine#2630)

* Fix jMonkeyEngine#2620

* Address feedback

* Address more feedback

* Use GLSL compat

* Fix shader

---------

Co-authored-by: 3gg <3gg@shellblade.net>

Author: jeannekamikaze

jme3-core/src/main/java/com/jme3/renderer/opengl/ComputeShader.java

@@ -50,25 +50,55 @@
 public class ComputeShader extends NativeObject {
 
     private final GL4 gl;
+    // The source need not be stored, but it helps with debugging.
     private final String source;
+
     /**
      * Creates a new compute shader from GLSL source code.
      */
     public ComputeShader(GL4 gl, String source) {
         super();
         this.gl = gl;
         this.source = source;
-        //Load this upfront to surface any problems at init time
+        // Load this up front to surface any problems at init time.
+        createComputeShader();
+    }
+
+    /**
+     * Creates a new compute shader from GLSL source code and a set of defines.
+     *
+     * @param defines An array of string pairs. The first element of the pair
+     *                is the macro name; the second element, the definition.
+     */
+    public ComputeShader(GL4 gl, String source, String[][] defines) {
+        super();
+        this.gl = gl;
+        this.source = addDefines(source, defines);
+        // Load this up front to surface any problems at init time.
         createComputeShader();
     }
-    private ComputeShader(ComputeShader source){
+
+    private ComputeShader(ComputeShader source) {
         super();
         this.gl = source.gl;
         this.id = source.id;
         this.source = null;
     }
 
-    private void createComputeShader(){
+    private String addDefines(String source, String[][] defines) {
+        // The #version pragma must appear before anything else. Insert the
+        // defines after it.
+        String[] sourceLines = (String[])source.split("\\r?\\n", 2);
+        StringBuilder builder = new StringBuilder();
+        builder.append(sourceLines[0] + "\n");
+        for (String[] pair : defines) {
+            builder.append("#define " + pair[0] + " " + pair[1] + "\n");
+        }
+        builder.append(sourceLines[1] + "\n");
+        return builder.toString();
+    }
+
+    private void createComputeShader() {
         // Create and compile the shader
         int shaderId = gl.glCreateShader(GL4.GL_COMPUTE_SHADER);
         if (shaderId <= 0) {

jme3-core/src/main/java/com/jme3/shadow/SdsmDirectionalLightShadowRenderer.java

@@ -147,7 +147,6 @@ private void initGL() {
 
         sdsmFitter = new SdsmFitter(gl4, renderer, assetManager);
         glInitialized = true;
-
     }
 
     /**

jme3-core/src/main/java/com/jme3/shadow/SdsmFitter.java

@@ -58,10 +58,16 @@ public class SdsmFitter {
 
     private final GL4 gl4;
     private final Renderer renderer;
+    private final AssetManager assetManager;
     private int maxFrameLag = 3;
 
-    private final ComputeShader depthReduceShader;
-    private final ComputeShader fitFrustumsShader;
+    // Because the multisampling status of the depth texture may change anytime,
+    // we store both ms and no-ms versions of the shaders, both initialized
+    // lazily. This avoids recompiling the same set of shaders twice and also
+    // compiling unnecessary shaders if the multisampling status does not
+    // change.
+    private InternalShaders shadersNoMultisampling;
+    private InternalShaders shadersMultisampling;
 
     private final LinkedList<SdsmResultHolder> resultHoldersInFlight = new LinkedList<>();
     private final LinkedList<SdsmResultHolder> resultHoldersReady = new LinkedList<>();
@@ -310,18 +316,58 @@ void cleanup() {
         }
     }
 
+    /**
+     * Manages the pair of depth-reduce and fit-frustums shaders.
+     */
+    private class InternalShaders {
+        public final ComputeShader depthReduceShader;
+        public final ComputeShader fitFrustumsShader;
+
+        InternalShaders(AssetManager assetManager, boolean multisampling) {
+            String reduceSource = (String)assetManager.loadAsset(REDUCE_DEPTH_SHADER);
+            String fitSource = (String)assetManager.loadAsset(FIT_FRUSTUMS_SHADER);
+
+            depthReduceShader = buildShader(reduceSource, multisampling);
+            fitFrustumsShader = buildShader(fitSource, multisampling);
+        }
+
+        private ComputeShader buildShader(String source, boolean multisampling) {
+            ComputeShader shader =
+                    multisampling
+                    ? new ComputeShader(gl4, source, new String[][]{{"RESOLVE_DEPTH_MS", "1"}})
+                    : new ComputeShader(gl4, source);
+            renderer.registerNativeObject(shader);
+            return shader;
+        }
+        
+        /**
+         * Cleans up GPU resources.
+         */
+        public void cleanup(Renderer renderer) {
+            depthReduceShader.deleteObject(renderer);
+            fitFrustumsShader.deleteObject(renderer);
+        }
+    }
+
     public SdsmFitter(GL4 gl, Renderer renderer, AssetManager assetManager) {
         this.gl4 = gl;
         this.renderer = renderer;
-
-        // Load compute shaders
-        String reduceSource = (String)assetManager.loadAsset(REDUCE_DEPTH_SHADER);
-        String fitSource = (String)assetManager.loadAsset(FIT_FRUSTUMS_SHADER);
-
-        depthReduceShader = new ComputeShader(gl, reduceSource);
-        renderer.registerNativeObject(depthReduceShader);
-        fitFrustumsShader = new ComputeShader(gl, fitSource);
-        renderer.registerNativeObject(fitFrustumsShader);
+        this.assetManager = assetManager;
+    }
+    
+    private InternalShaders initShaders(Texture depthTexture) {
+        boolean multisampling = depthTexture.getImage().getMultiSamples() > 1;
+        if (multisampling) {
+            if (shadersMultisampling == null) {
+                shadersMultisampling = new InternalShaders(assetManager, true);
+            }
+            return shadersMultisampling;
+        } else {
+            if (shadersNoMultisampling == null) {
+                shadersNoMultisampling = new InternalShaders(assetManager, false);
+            }
+            return shadersNoMultisampling;
+        }
     }
 
     /**
@@ -336,6 +382,12 @@ public SdsmFitter(GL4 gl, Renderer renderer, AssetManager assetManager) {
     public void fit(Texture depthTexture, int splitCount, Matrix4f cameraToLight,
                     float cameraNear, float cameraFar) {
 
+        int depthMultiSamples = depthTexture.getImage().getMultiSamples();
+
+        InternalShaders shaders = initShaders(depthTexture);
+        ComputeShader depthReduceShader = shaders.depthReduceShader;
+        ComputeShader fitFrustumsShader = shaders.fitFrustumsShader;
+
         SdsmResultHolder holder = getResultHolderForUse();
         holder.parameters = new FitParameters(cameraToLight, splitCount, cameraNear, cameraFar);
 
@@ -360,6 +412,8 @@ public void fit(Texture depthTexture, int splitCount, Matrix4f cameraToLight,
         } catch (TextureUnitException e) {
             throw new RendererException(e);
         }
+        int loc = depthReduceShader.getUniformLocation("m_NumSamplesDepth");
+        depthReduceShader.setUniform(loc, depthMultiSamples);
         depthReduceShader.bindShaderStorageBuffer(1, holder.minMaxDepthSsbo);
         depthReduceShader.dispatch(xGroups, yGroups, 1);
         gl4.glMemoryBarrier(GL4.GL_SHADER_STORAGE_BARRIER_BIT);
@@ -373,6 +427,8 @@ public void fit(Texture depthTexture, int splitCount, Matrix4f cameraToLight,
         } catch (TextureUnitException e) {
             throw new RendererException(e);
         }
+        loc = fitFrustumsShader.getUniformLocation("m_NumSamplesDepth");
+        fitFrustumsShader.setUniform(loc, depthMultiSamples);
         fitFrustumsShader.bindShaderStorageBuffer(1, holder.minMaxDepthSsbo);
         fitFrustumsShader.bindShaderStorageBuffer(2, holder.fitFrustumSsbo);
 
@@ -437,11 +493,11 @@ public void cleanup() {
         }
         resultHoldersReady.clear();
 
-        if (depthReduceShader != null) {
-            depthReduceShader.deleteObject(renderer);
+        if (shadersMultisampling != null) {
+            shadersMultisampling.cleanup(renderer);
         }
-        if (fitFrustumsShader != null) {
-            fitFrustumsShader.deleteObject(renderer);
+        if (shadersNoMultisampling != null) {
+            shadersNoMultisampling.cleanup(renderer);
         }
     }
 

jme3-core/src/main/resources/Common/MatDefs/Shadow/Sdsm/FitLightFrustums.comp

@@ -1,12 +1,15 @@
 #version 430
 
+#import "Common/ShaderLib/GLSLCompat.glsllib"
+#import "Common/ShaderLib/MultiSample.glsllib"
+
 /**
  * Computes tight bounding boxes for each shadow cascade via min/max on lightspace locations of depth samples that fall within each cascade.
  */
 
 layout(local_size_x = 16, local_size_y = 16) in;
 
-layout(binding = 0) uniform sampler2D inputDepth;
+layout(binding = 0) uniform DEPTHTEXTURE inputDepth;
 
 layout(std430, binding = 1) readonly buffer MinMaxBuffer {
     uint gMin;
@@ -103,7 +106,7 @@ void main() {
     ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
     ivec2 lid = ivec2(gl_LocalInvocationID.xy);
     uint tid = gl_LocalInvocationIndex;
-    ivec2 inputSize = textureSize(inputDepth, 0);
+    ivec2 inputSize = getTextureSize(inputDepth);
     ivec2 baseCoord = gid * 2;
 
     // Initialize local bounds to infinity
@@ -126,11 +129,11 @@ void main() {
         for (int x = 0; x < 2; x++) {
             ivec2 coord = baseCoord + ivec2(x, y);
             if (coord.x < inputSize.x && coord.y < inputSize.y) {
-                float depth = texelFetch(inputDepth, coord, 0).r;
+                float depth = getDepthMax(inputDepth, coord, m_NumSamplesDepth).r;
                 // Skip background (depth == 1.0)
                 if (depth != 1.0) {
                     // Reconstruct clip-space position from depth
-                    vec2 uv = (vec2(coord) + 0.5) / vec2(textureSize(inputDepth, 0));
+                    vec2 uv = (vec2(coord) + 0.5) / vec2(getTextureSize(inputDepth));
                     vec4 clipPos = vec4(uv * 2.0 - 1.0, depth * 2.0 - 1.0, 1.0);
 
                     // Transform to light view space

jme3-core/src/main/resources/Common/MatDefs/Shadow/Sdsm/ReduceDepth.comp

@@ -1,12 +1,15 @@
 #version 430
 
+#import "Common/ShaderLib/GLSLCompat.glsllib"
+#import "Common/ShaderLib/MultiSample.glsllib"
+
 /**
  * Finds the global minimum/maximum values of a depth texture.
  */
 
 layout(local_size_x = 16, local_size_y = 16) in;
 
-layout(binding = 0) uniform sampler2D inputDepth;
+layout(binding = 0) uniform DEPTHTEXTURE inputDepth;
 
 layout(std430, binding = 1) buffer MinMaxBuffer {
     uint gMin;
@@ -33,7 +36,7 @@ void main() {
     ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
     ivec2 lid = ivec2(gl_LocalInvocationID.xy);
     uint tid = gl_LocalInvocationIndex;
-    ivec2 inputSize = textureSize(inputDepth, 0);
+    ivec2 inputSize = getTextureSize(inputDepth);
 
     // Each thread samples a 2x2 block
     ivec2 baseCoord = gid * 2;
@@ -43,7 +46,7 @@ void main() {
         for (int x = 0; x < 2; x++) {
             ivec2 coord = baseCoord + ivec2(x, y);
             if (coord.x < inputSize.x && coord.y < inputSize.y) {
-                float depth = texelFetch(inputDepth, coord, 0).r;
+                float depth = getDepthMax(inputDepth, coord, m_NumSamplesDepth).r;
                 // Discard depth == 1.0 (background/sky)
                 if (depth != 1.0) {
                     minMax.x = min(minMax.x, depth);

jme3-core/src/main/resources/Common/MatDefs/Shadow/Sdsm/SdsmPostShadow.frag

@@ -1,9 +1,10 @@
 #import "Common/ShaderLib/GLSLCompat.glsllib"
+#import "Common/ShaderLib/MultiSample.glsllib"
 #import "Common/ShaderLib/Shadows.glsllib"
 
 //Stripped version of the usual shadow fragment shader for SDSM; it intentionally leaves out some features.
-uniform sampler2D m_Texture;
-uniform sampler2D m_DepthTexture;
+uniform COLORTEXTURE m_Texture;
+uniform DEPTHTEXTURE m_DepthTexture;
 uniform mat4 m_ViewProjectionMatrixInverse;
 uniform vec4 m_ViewProjectionMatrixRow2;
 
@@ -43,8 +44,8 @@ float determineShadow(int index, vec4 worldPos){
 }
 
 void main() {
-    float depth = texture2D(m_DepthTexture,texCoord).r;
-    vec4 color = texture2D(m_Texture,texCoord);
+    float depth = getDepth(m_DepthTexture,texCoord).r;
+    vec4 color  = getColor(m_Texture,texCoord);
 
     //Discard shadow computation on the sky
     if(depth == 1.0){

jme3-core/src/main/resources/Common/MatDefs/Shadow/Sdsm/SdsmPostShadow.j3md

@@ -31,7 +31,8 @@ MaterialDef Post Shadow {
         Texture2D DepthTexture
 
         Boolean BackfaceShadows //Not used.
-        Int NumSamples //Not used.
+        Int NumSamples
+        Int NumSamplesDepth
     }
 
     Technique {
@@ -47,6 +48,8 @@ MaterialDef Post Shadow {
             FILTER_MODE : FilterMode
             PCFEDGE : PCFEdge
             SHADOWMAP_SIZE : ShadowMapSize
+            RESOLVE_MS : NumSamples
+            RESOLVE_DEPTH_MS : NumSamplesDepth
         }
     }
 }

jme3-core/src/main/resources/Common/ShaderLib/MultiSample.glsllib

@@ -18,12 +18,12 @@ uniform int m_NumSamplesDepth;
 // NOTE: Only define multisample functions if multisample is available
 #if defined(GL_ARB_texture_multisample) || (defined GL_ES && __VERSION__>=310)
 vec4 textureFetch(in sampler2DMS tex,in vec2 texC, in int numSamples){
-      ivec2 iTexC = ivec2(texC * vec2(textureSize(tex)));
-      vec4 color = vec4(0.0);
-      for (int i = 0; i < numSamples; i++){
-         color += texelFetch(tex, iTexC, i);
-      }
-      return color / float(numSamples);
+    ivec2 iTexC = ivec2(texC * vec2(textureSize(tex)));
+    vec4 color = vec4(0.0);
+    for (int i = 0; i < numSamples; i++){
+       color += texelFetch(tex, iTexC, i);
+    }
+    return color / float(numSamples);
 }
 
 vec4 fetchTextureSample(in sampler2DMS tex,in vec2 texC,in int sampleId){
@@ -40,11 +40,23 @@ vec4 getColorSingle(in sampler2DMS tex, in vec2 texC){
     return texelFetch(tex, iTexC, 0);
 }
 
-vec4 getDepth(in sampler2DMS tex,in vec2 texC){
-      return textureFetch(tex,texC,m_NumSamplesDepth);
+vec4 getDepth(in sampler2DMS tex, in vec2 texC){
+    return textureFetch(tex,texC,m_NumSamplesDepth);
+
 }
 
-#endif
+vec4 getDepthMax(in sampler2DMS tex, in ivec2 coord, in int numSamples){
+    vec4 result = vec4(0.0);
+    for (int i = 0; i < numSamples; i++){
+       result = max(result, texelFetch(tex, coord, i));
+    }
+    return result;
+}
+
+ivec2 getTextureSize(in sampler2DMS tex) {
+    return textureSize(tex);
+}
+#endif // Multisampling
 
 vec4 fetchTextureSample(in sampler2D tex,in vec2 texC,in int sampleId){
     return texture2D(tex,texC);
@@ -55,10 +67,17 @@ vec4 getColor(in sampler2D tex, in vec2 texC){
 }
 
 vec4 getColorSingle(in sampler2D tex, in vec2 texC){
-    return texture2D(tex, texC);
+    return texture2D(tex,texC);
 }
 
-vec4 getDepth(in sampler2D tex,in vec2 texC){
+vec4 getDepth(in sampler2D tex, in vec2 texC){
     return texture2D(tex,texC);
 }
 
+vec4 getDepthMax(in sampler2D tex, in ivec2 coord, in int numSamples){
+    return texelFetch(tex, coord, 0);
+}
+
+ivec2 getTextureSize(in sampler2D tex){
+    return textureSize(tex, 0);
+}

jme3-core/src/plugins/java/com/jme3/shader/plugins/GLSLLoader.java

@@ -108,6 +108,9 @@ private ShaderDependencyNode loadNode(Reader reader, String nodeName) {
                     }
                 } else if (tln.startsWith("#extension ")) {
                     sbExt.append(ln).append('\n');
+                } else if (tln.startsWith("#version ")) {
+                    // #version must appear before the extensions, so treat it like one.
+                    sbExt.append(ln).append('\n');
                 } else {
                     sb.append(ln).append('\n');
                 }