Skip to content

Commit

Permalink
Use specialization constants in compute shader
Browse files Browse the repository at this point in the history
  • Loading branch information
@lwjglgamedev
lwjglgamedev committed Sep 16, 2023
1 parent 0abb885 commit 4474fe8
Show file tree
Hide file tree
Showing 32 changed files with 270 additions and 67 deletions.
142 changes: 99 additions & 43 deletions bookcontents/chapter-14/chapter-14.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,13 +244,13 @@ Let’s review the `processAnimations` method, which is defined like this:
```java
public class ModelLoader {
...
public static final int MAX_JOINTS = 150;
private static final Matrix4f IDENTITY_MATRIX = new Matrix4f();
...
private static List<ModelData.Animation> processAnimations(AIScene aiScene, List<Bone> boneList,
Node rootNode, Matrix4f globalInverseTransformation) {
List<ModelData.Animation> animations = new ArrayList<>();

int maxJointsMatricesLists = EngineProperties.getInstance().getMaxJointsMatricesLists();
// Process all animations
int numAnimations = aiScene.mNumAnimations();
PointerBuffer aiAnimations = aiScene.mAnimations();
Expand All @@ -263,7 +263,7 @@ public class ModelLoader {
animations.add(animation);

for (int j = 0; j < maxFrames; j++) {
Matrix4f[] jointMatrices = new Matrix4f[MAX_JOINTS];
Matrix4f[] jointMatrices = new Matrix4f[maxJointsMatricesLists];
Arrays.fill(jointMatrices, IDENTITY_MATRIX);
ModelData.AnimatedFrame animatedFrame = new ModelData.AnimatedFrame(jointMatrices);
buildFrameMatrices(aiAnimation, boneList, animatedFrame, j, rootNode,
Expand Down Expand Up @@ -408,7 +408,36 @@ public class ModelLoader {
}
```

All that new information needs to be handled in the `VulkanModel` class so it is loaded into the GPU. This class now defines the following new attribute:
We needed to add new configuration properties in the `EngineProperties` class to properly establish the size of the descriptors set pool in the `AnimationComputeActivity` class, here are the changes:
```java
public class EngineProperties {
...
private static final int DEFAULT_MAX_JOINTS_MATRICES_LISTS = 100;
private static final int DEFAULT_STORAGES_BUFFERS = 100;
...
private int maxJointsMatricesLists;
private int maxStorageBuffers;
...
private EngineProperties() {
...
maxStorageBuffers = Integer.parseInt(props.getOrDefault("maxStorageBuffers", DEFAULT_STORAGES_BUFFERS).toString());
maxJointsMatricesLists = Integer.parseInt(props.getOrDefault("maxJointsMatricesLists", DEFAULT_MAX_JOINTS_MATRICES_LISTS).toString());
...
}
...
public int getMaxJointsMatricesLists() {
return maxJointsMatricesLists;
}

public int getMaxStorageBuffers() {
return maxStorageBuffers;
}
...
}
```


All the information processed by `ModelLoader` class needs to be handled in the `VulkanModel` class so it is loaded into the GPU. This class now defines the following new attribute:
```java
public class VulkanModel {

Expand Down Expand Up @@ -483,6 +512,7 @@ public class VulkanModel {
...
}
```

If the model has animations, we create the buffers that will hold the join transformation matrices for each of the frames. As usual, we use staging buffers with temporary CPU accessible buffers that get copied to a GPU only one, so we need to record those transitions. We also create the wights buffer for each of the meshes and store that in the `VulkanModel.VulkanMesh` instance. The `createJointMatricesBuffers` is defined like this:
```java
public class VulkanModel {
Expand Down Expand Up @@ -603,7 +633,7 @@ public class ComputePipeline {
}
```

Going back to the `ComputePipeline` constructor, we first initialize the `VkPipelineShaderStageCreateInfo` structure with the compute shader information. In this specific case, we receive a `ShaderProgram` instance through the `ComputePipeline.PipeLineCreationInfo` record. The `ShaderProgram` class is also used for the graphics pipeline creation holding shader modules that refer to vertex and fragment shaders. In our case, we just need a compute shader, so we just need a shader module. We will assume that in the first position of the shader modules we will receive that reference.
Going back to the `ComputePipeline` constructor, we first initialize the `VkPipelineShaderStageCreateInfo` structure with the compute shader information. In this specific case, we receive a `ShaderProgram` instance through the `ComputePipeline.PipeLineCreationInfo` record. The `ShaderProgram` class is also used for the graphics pipeline creation holding shader modules that refer to vertex and fragment shaders. In our case, we just need a compute shader, so we just need a shader module. We will assume that in the first position of the shader modules we will receive that reference. As in other types of shaders, compute shaders can use specialization constants, so we need to add support for them if `ShaderModule` instance defines them.
```java
public class ComputePipeline {
...
Expand All @@ -625,6 +655,9 @@ public class ComputePipeline {
.stage(shaderModule.shaderStage())
.module(shaderModule.handle())
.pName(main);
if (shaderModule.specInfo() != null) {
shaderStage.pSpecializationInfo(shaderModule.specInfo());
}
...
}
...
Expand Down Expand Up @@ -811,19 +844,21 @@ public class AnimationComputeActivity {
private static final String ANIM_COMPUTE_SHADER_FILE_SPV = ANIM_COMPUTE_SHADER_FILE_GLSL + ".spv";
private static final int LOCAL_SIZE_X = 32;

private MemoryBarrier memoryBarrier;
private final Device device;
private final MemoryBarrier memoryBarrier;
private final Queue.ComputeQueue computeQueue;
// Key is the entity id
private final Map<String, List<EntityAnimationBuffer>> entityAnimationsBuffers;
// Key is the model id
private final Map<String, ModelDescriptorSets> modelDescriptorSetsMap;
private final Scene scene;
private final AnimationSpecConstants animationSpecConstants;

private CommandBuffer commandBuffer;
private ComputePipeline computePipeline;
private Queue.ComputeQueue computeQueue;
private DescriptorPool descriptorPool;
private DescriptorSetLayout[] descriptorSetLayouts;
private Device device;
// Key is the entity id
private Map<String, List<EntityAnimationBuffer>> entityAnimationsBuffers;
private Fence fence;
// Key is the model id
private Map<String, ModelDescriptorSets> modelDescriptorSetsMap;
private Scene scene;
private ShaderProgram shaderProgram;
private DescriptorSetLayout.StorageDescriptorSetLayout storageDescriptorSetLayout;
private DescriptorSetLayout.UniformDescriptorSetLayout uniformDescriptorSetLayout;
Expand All @@ -832,6 +867,7 @@ public class AnimationComputeActivity {
this.scene = scene;
device = pipelineCache.getDevice();
computeQueue = new Queue.ComputeQueue(device, 0);
animationSpecConstants = new AnimationSpecConstants();
createDescriptorPool();
createDescriptorSets();
createShaders();
Expand All @@ -845,12 +881,60 @@ public class AnimationComputeActivity {
}
```

In the `cleanup` method, as usual, we just free the resources:
We create an instance of `AnimationSpecConstants` which will be used to set specialization constants. Specifically, it will hold the maximum number of joint matrices lists which will be used in the compute shader. It is defined liked this:

```java
package org.vulkanb.eng.graph.animation;

import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.*;
import org.vulkanb.eng.EngineProperties;
import org.vulkanb.eng.graph.vk.GraphConstants;

import java.nio.ByteBuffer;

public class AnimationSpecConstants {

private final ByteBuffer data;
private final VkSpecializationMapEntry.Buffer specEntryMap;
private final VkSpecializationInfo specInfo;

public AnimationSpecConstants() {
EngineProperties engineProperties = EngineProperties.getInstance();
data = MemoryUtil.memAlloc(GraphConstants.INT_LENGTH);
data.putInt(engineProperties.getMaxJointsMatricesLists());
data.flip();

specEntryMap = VkSpecializationMapEntry.calloc(1);
specEntryMap.get(0)
.constantID(0)
.size(GraphConstants.INT_LENGTH)
.offset(0);

specInfo = VkSpecializationInfo.calloc();
specInfo.pData(data)
.pMapEntries(specEntryMap);
}

public void cleanup() {
MemoryUtil.memFree(specEntryMap);
specInfo.free();
MemoryUtil.memFree(data);
}

public VkSpecializationInfo getSpecInfo() {
return specInfo;
}
}
```

Back to the `AnimationComputeActivity` class, in the `cleanup` method, as usual, we just free the resources:
```java
public class AnimationComputeActivity {
...
public void cleanup() {
computePipeline.cleanup();
animationSpecConstants.cleanup();
shaderProgram.cleanup();
commandBuffer.cleanup();
descriptorPool.cleanup();
Expand All @@ -866,7 +950,7 @@ public class AnimationComputeActivity {
}
```

Now we will review the `createXX` methods called in the constructor. The first one is the `createCommandBuffers`, which creates a command buffer that will be used to record the compute dispatch commands. It also creates a fence to prevent reusing the command buffer while in use. The `createDescriptorPool` method creates a descriptor pool, defining the maximum number of each descriptor type that we are going to use. In our case, we will be using storage descriptor sets for the data that we will use for the animation and uniform buffers that we will use to pass the joint transformation matrices. The `createDescriptorSets` method just creates the layouts of the descriptor sets that we will use in the compute shader. The `createPipeline` method just creates our compute pipeline with the descriptor sets layouts information created previously. Finally, the `createShaders` just creates a shader program, which will contain a shader module which holds the compute shader code. As you can see, they are similar as the ones used in geometry, shadow and lighting phases in previous chapters.
Now we will review the `createXX` methods called in the `AnimationComputeActivity` class constructor. The first one is the `createCommandBuffers`, which creates a command buffer that will be used to record the compute dispatch commands. It also creates a fence to prevent reusing the command buffer while in use. The `createDescriptorPool` method creates a descriptor pool, defining the maximum number of each descriptor type that we are going to use. In our case, we will be using storage descriptor sets for the data that we will use for the animation and uniform buffers that we will use to pass the joint transformation matrices. The `createDescriptorSets` method just creates the layouts of the descriptor sets that we will use in the compute shader. The `createPipeline` method just creates our compute pipeline with the descriptor sets layouts information created previously. Finally, the `createShaders` just creates a shader program, which will contain a shader module which holds the compute shader code. As you can see, they are similar as the ones used in geometry, shadow and lighting phases in previous chapters.
```java
public class AnimationComputeActivity {
...
Expand Down Expand Up @@ -916,34 +1000,6 @@ public class AnimationComputeActivity {
}
```

We needed to add new configuration properties in the `EngineProperties` class to properly establish the size of the descriptors set pool in the `AnimationComputeActivity` class, here are the changes:
```java
public class EngineProperties {
...
private static final int DEFAULT_MAX_JOINTS_MATRICES_LISTS = 100;
private static final int DEFAULT_STORAGES_BUFFERS = 100;
...
private int maxJointsMatricesLists;
private int maxStorageBuffers;
...
private EngineProperties() {
...
maxStorageBuffers = Integer.parseInt(props.getOrDefault("maxStorageBuffers", DEFAULT_STORAGES_BUFFERS).toString());
maxJointsMatricesLists = Integer.parseInt(props.getOrDefault("maxJointsMatricesLists", DEFAULT_MAX_JOINTS_MATRICES_LISTS).toString());
...
}
...
public int getMaxJointsMatricesLists() {
return maxJointsMatricesLists;
}

public int getMaxStorageBuffers() {
return maxStorageBuffers;
}
...
}
```

We now will create a new method method that should be called when new models that contain animations are created. That method, named `registerModels`, is defined like this:
```java
public class AnimationComputeActivity {
Expand Down Expand Up @@ -1121,7 +1177,7 @@ The next step is to write the compute shader which performs the calculations. Th
```glsl
#version 450
const int MAX_JOINTS = 150;
layout (constant_id = 0) const int MAX_JOINTS = 150;
layout (std430, set=0, binding=0) readonly buffer srcBuf {
float data[];
Expand Down
2 changes: 1 addition & 1 deletion bookcontents/chapter-16/chapter-16.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1907,7 +1907,7 @@ public class AnimationComputeActivity {
...
}
```
Instead of having multiple descriptor sets and buffers per model, we will have all that data in combined storage buffers. Therefore, we need to create the new descriptor sets for them, as storage descriptor sets. Those descriptor sets are created when the animated entities are loaded. This will be done in the `onAnimatedEntitiesLoaded` which will be called from the `Render` instance.
Instead of having multiple descriptor sets and buffers per model, we will have all that data in combined storage buffers. Therefore, we need to create the new descriptor sets for them, as storage descriptor sets. We will not need also specialization constants since everything will by handled by storage buffers (no more uniforms). Those descriptor sets are created when the animated entities are loaded. This will be done in the `onAnimatedEntitiesLoaded` which will be called from the `Render` instance.
```java
public class AnimationComputeActivity {
...
Expand Down
8 changes: 7 additions & 1 deletion booksamples/appendix-01/src/main/java/org/vulkanb/eng/EngineProperties.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ public class EngineProperties {
private static final int DEFAULT_JOINT_MATRICES_BUF = 2000000;
private static final int DEFAULT_MAX_ANIM_WEIGHTS_BUF = 100000;
private static final int DEFAULT_MAX_INDICES_BUF = 5000000;
private static final int DEFAULT_MAX_JOINTS_MATRICES_LISTS = 100;
private static final int DEFAULT_MAX_MATERIALS = 500;
private static final int DEFAULT_MAX_VERTICES_BUF = 20000000;
private static final int DEFAULT_REQUESTED_IMAGES = 3;
Expand All @@ -26,6 +27,7 @@ public class EngineProperties {
private int maxAnimWeightsBuffer;
private int maxIndicesBuffer;
private int maxJointMatricesBuffer;
private int maxJointsMatricesLists;
private int maxMaterials;
private int maxTextures;
private int maxVerticesBuffer;
Expand Down Expand Up @@ -68,7 +70,7 @@ private EngineProperties() {
maxIndicesBuffer = Integer.parseInt(props.getOrDefault("maxIndicesBuffer", DEFAULT_MAX_INDICES_BUF).toString());
maxAnimWeightsBuffer = Integer.parseInt(props.getOrDefault("maxAnimWeightsBuffer", DEFAULT_MAX_ANIM_WEIGHTS_BUF).toString());
maxJointMatricesBuffer = Integer.parseInt(props.getOrDefault("maxJointMatricesBuffer", DEFAULT_JOINT_MATRICES_BUF).toString());
enableCheckPoints = Boolean.parseBoolean(props.getOrDefault("enableCheckPoints", false).toString());
maxJointsMatricesLists = Integer.parseInt(props.getOrDefault("maxJointsMatricesLists", DEFAULT_MAX_JOINTS_MATRICES_LISTS).toString());
} catch (IOException excp) {
Logger.error("Could not read [{}] properties file", FILENAME, excp);
}
Expand Down Expand Up @@ -101,6 +103,10 @@ public int getMaxJointMatricesBuffer() {
return maxJointMatricesBuffer;
}

public int getMaxJointsMatricesLists() {
return maxJointsMatricesLists;
}

public int getMaxMaterials() {
return maxMaterials;
}
Expand Down
3 changes: 3 additions & 0 deletions booksamples/appendix-01/src/main/java/org/vulkanb/eng/graph/vk/ComputePipeline.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,9 @@ public ComputePipeline(PipelineCache pipelineCache, ComputePipeline.PipeLineCrea
.stage(shaderModule.shaderStage())
.module(shaderModule.handle())
.pName(main);
if (shaderModule.specInfo() != null) {
shaderStage.pSpecializationInfo(shaderModule.specInfo());
}

VkPushConstantRange.Buffer vpcr = null;
if (pipeLineCreationInfo.pushConstantsSize() > 0) {
Expand Down
5 changes: 3 additions & 2 deletions booksamples/appendix-01/src/main/java/org/vulkanb/eng/scene/ModelLoader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import org.lwjgl.assimp.*;
import org.lwjgl.system.MemoryStack;
import org.tinylog.Logger;
import org.vulkanb.eng.EngineProperties;

import java.io.File;
import java.lang.Math;
Expand All @@ -16,7 +17,6 @@

public class ModelLoader {

public static final int MAX_JOINTS = 150;
public static final int MAX_WEIGHTS = 4;
private static final Matrix4f IDENTITY_MATRIX = new Matrix4f();

Expand Down Expand Up @@ -187,6 +187,7 @@ private static List<ModelData.Animation> processAnimations(AIScene aiScene, List
Node rootNode, Matrix4f globalInverseTransformation) {
List<ModelData.Animation> animations = new ArrayList<>();

int maxJointsMatricesLists = EngineProperties.getInstance().getMaxJointsMatricesLists();
// Process all animations
int numAnimations = aiScene.mNumAnimations();
PointerBuffer aiAnimations = aiScene.mAnimations();
Expand All @@ -199,7 +200,7 @@ private static List<ModelData.Animation> processAnimations(AIScene aiScene, List
animations.add(animation);

for (int j = 0; j < maxFrames; j++) {
Matrix4f[] jointMatrices = new Matrix4f[MAX_JOINTS];
Matrix4f[] jointMatrices = new Matrix4f[maxJointsMatricesLists];
Arrays.fill(jointMatrices, IDENTITY_MATRIX);
ModelData.AnimatedFrame animatedFrame = new ModelData.AnimatedFrame(jointMatrices);
buildFrameMatrices(aiAnimation, boneList, animatedFrame, j, rootNode,
Expand Down
1 change: 1 addition & 0 deletions booksamples/appendix-01/src/main/resources/eng.properties
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ defaultTexturePath=resources/models/default/default.png
maxAnimWeightsBuffer=1000000
maxIndicesBuffer=5000000
maxJointMatricesBuffer=20000000
maxJointsMatricesLists=150
maxMaterials=500
maxVerticesBuffer=20000000
shadowPcf=true
Expand Down
8 changes: 7 additions & 1 deletion booksamples/appendix-02/src/main/java/org/vulkanb/eng/EngineProperties.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ public class EngineProperties {
private static final int DEFAULT_JOINT_MATRICES_BUF = 2000000;
private static final int DEFAULT_MAX_ANIM_WEIGHTS_BUF = 100000;
private static final int DEFAULT_MAX_INDICES_BUF = 5000000;
private static final int DEFAULT_MAX_JOINTS_MATRICES_LISTS = 100;
private static final int DEFAULT_MAX_MATERIALS = 500;
private static final int DEFAULT_MAX_VERTICES_BUF = 20000000;
private static final int DEFAULT_REQUESTED_IMAGES = 3;
Expand All @@ -26,6 +27,7 @@ public class EngineProperties {
private int maxAnimWeightsBuffer;
private int maxIndicesBuffer;
private int maxJointMatricesBuffer;
private int maxJointsMatricesLists;
private int maxMaterials;
private int maxTextures;
private int maxVerticesBuffer;
Expand Down Expand Up @@ -68,7 +70,7 @@ private EngineProperties() {
maxIndicesBuffer = Integer.parseInt(props.getOrDefault("maxIndicesBuffer", DEFAULT_MAX_INDICES_BUF).toString());
maxAnimWeightsBuffer = Integer.parseInt(props.getOrDefault("maxAnimWeightsBuffer", DEFAULT_MAX_ANIM_WEIGHTS_BUF).toString());
maxJointMatricesBuffer = Integer.parseInt(props.getOrDefault("maxJointMatricesBuffer", DEFAULT_JOINT_MATRICES_BUF).toString());
enableCheckPoints = Boolean.parseBoolean(props.getOrDefault("enableCheckPoints", false).toString());
maxJointsMatricesLists = Integer.parseInt(props.getOrDefault("maxJointsMatricesLists", DEFAULT_MAX_JOINTS_MATRICES_LISTS).toString());
} catch (IOException excp) {
Logger.error("Could not read [{}] properties file", FILENAME, excp);
}
Expand Down Expand Up @@ -101,6 +103,10 @@ public int getMaxJointMatricesBuffer() {
return maxJointMatricesBuffer;
}

public int getMaxJointsMatricesLists() {
return maxJointsMatricesLists;
}

public int getMaxMaterials() {
return maxMaterials;
}
Expand Down
Loading

0 comments on commit 4474fe8

Please sign in to comment.