Skip to content

Commit

Permalink
Reflection based safe descriptor set updates & set-typed pools
Browse files Browse the repository at this point in the history
  • Loading branch information
@bobcao3
bobcao3 committed Oct 1, 2023
1 parent cfe8e90 commit e112e31
Show file tree
Hide file tree
Showing 9 changed files with 265 additions and 129 deletions.
26 changes: 26 additions & 0 deletions src/main/java/me/cortex/vulkanite/client/Vulkanite.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,16 @@
import me.cortex.vulkanite.client.rendering.VulkanPipeline;
import me.cortex.vulkanite.lib.base.VContext;
import me.cortex.vulkanite.lib.base.initalizer.VInitializer;
import me.cortex.vulkanite.lib.descriptors.VDescriptorPool;
import me.cortex.vulkanite.lib.descriptors.VDescriptorSetLayout;
import me.cortex.vulkanite.lib.descriptors.VTypedDescriptorPool;
import me.jellysquid.mods.sodium.client.render.chunk.RenderSection;
import me.jellysquid.mods.sodium.client.render.chunk.compile.ChunkBuildOutput;
import net.minecraft.util.Util;
import org.lwjgl.vulkan.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import static org.lwjgl.vulkan.EXTDebugUtils.VK_EXT_DEBUG_UTILS_EXTENSION_NAME;
Expand Down Expand Up @@ -50,6 +54,7 @@ public class Vulkanite {
private final ArbitarySyncPointCallback fencedCallback = new ArbitarySyncPointCallback();

private final AccelerationManager accelerationManager;
private final HashMap<VDescriptorSetLayout, VTypedDescriptorPool> descriptorPools = new HashMap<>();

public Vulkanite() {
ctx = createVulkanContext();
Expand All @@ -70,6 +75,24 @@ public void upload(List<ChunkBuildOutput> results) {
accelerationManager.chunkBuilds(results);
}

public VTypedDescriptorPool getPoolByLayout(VDescriptorSetLayout layout) {
synchronized (descriptorPools) {
if (!descriptorPools.containsKey(layout)) {
descriptorPools.put(layout, new VTypedDescriptorPool(ctx, layout, 0));
}
return descriptorPools.get(layout);
}
}

public void removePoolByLayout(VDescriptorSetLayout layout) {
synchronized (descriptorPools) {
if (descriptorPools.containsKey(layout)) {
descriptorPools.get(layout).free();
descriptorPools.remove(layout);
}
}
}

public void sectionRemove(RenderSection section) {
accelerationManager.sectionRemove(section);
}
Expand All @@ -92,6 +115,9 @@ public void addSyncedCallback(Runnable callback) {
}

public void destroy() {
for (var pool : descriptorPools.values()) {
pool.free();
}
accelerationManager.cleanup();
}

Expand Down
186 changes: 76 additions & 110 deletions src/main/java/me/cortex/vulkanite/client/rendering/VulkanPipeline.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,13 +62,6 @@ public class VulkanPipeline {

private record RtPipeline(VRaytracePipeline pipeline, int commonSet, int geomSet, int customTexSet, int ssboSet) {}
private RtPipeline[] raytracePipelines;
private VDescriptorSetLayout commonLayout;
private VDescriptorSetLayout customtexLayout;
private VDescriptorSetLayout storageBufferLayout;

private VDescriptorPool commonDescriptorPool;
private VDescriptorPool customtexDescriptorPool;
private VDescriptorPool storageBufferDescriptorPool;

private final VSampler sampler;
private final VSampler ctexSampler;
Expand Down Expand Up @@ -170,40 +163,6 @@ public VulkanPipeline(VContext ctx, AccelerationManager accelerationManager, Ray
.maxAnisotropy(1.0f));

try {
commonLayout = new DescriptorSetLayoutBuilder()
.binding(0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_ALL)// camera data
.binding(1, VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, VK_SHADER_STAGE_ALL)// funni acceleration buffer
.binding(3, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL)//block texture
.binding(4, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL)//block texture normal
.binding(5, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL)//block texture specular
// Reordered these so output texture is last... this means you can dynamically add more output textures without messing other ids
.binding(6, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, maxIrisRenderTargets, VK_SHADER_STAGE_ALL)// output texture
.build(ctx);

DescriptorSetLayoutBuilder ctexLayoutBuilder = new DescriptorSetLayoutBuilder();
for (int i = 0; i < customTextureViews.length; i++) {
ctexLayoutBuilder.binding(i, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL);
}

customtexLayout = ctexLayoutBuilder.build(ctx);

DescriptorSetLayoutBuilder ssboLayoutBuilder = new DescriptorSetLayoutBuilder();
for (int id : ssboIds) {
ssboLayoutBuilder.binding(id, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_SHADER_STAGE_ALL);
}

storageBufferLayout = ssboLayoutBuilder.build(ctx);

//TODO: use frameahead count instead of just... 10
commonDescriptorPool = new VDescriptorPool(ctx, VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT, 10, commonLayout.types);
commonDescriptorPool.allocateSets(commonLayout);

customtexDescriptorPool = new VDescriptorPool(ctx, VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT, 10, customtexLayout.types);
customtexDescriptorPool.allocateSets(customtexLayout);

storageBufferDescriptorPool = new VDescriptorPool(ctx, VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT, 10, storageBufferLayout.types);
storageBufferDescriptorPool.allocateSets(storageBufferLayout);

var commonSetExpected = new ShaderReflection.Set(new ShaderReflection.Binding[]{
new ShaderReflection.Binding("", 0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 0, false),
new ShaderReflection.Binding("", 1, VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, 0, false),
Expand Down Expand Up @@ -329,46 +288,13 @@ public void renderPostShadows(List<VGImage> outImgs, Camera camera, ShaderStorag
uboBuffer.unmap();
uboBuffer.flush();

long commonSet = commonDescriptorPool.get(fidx);
long ctexSet = customtexDescriptorPool.get(fidx);
long ssboSet = storageBufferDescriptorPool.get(fidx);

var updater = new DescriptorUpdateBuilder(ctx, 7)
.set(commonSet)
.uniform(0, uboBuffer)
.acceleration(1, tlas)
.imageSampler(3, blockAtlasView.getView(), sampler)
.imageSampler(4,
blockAtlasNormalView.getView() != null ? blockAtlasNormalView.getView()
: placeholderNormalsView,
sampler)
.imageSampler(5,
blockAtlasSpecularView.getView() != null ? blockAtlasSpecularView.getView()
: placeholderSpecularView,
sampler);
List<VImageView> outImgViewList = new ArrayList<>(outImgs.size());
for (int i = 0; i < outImgs.size(); i++) {
int index = i;
outImgViewList.add(irisRenderTargetViews[i].getView(() -> outImgs.get(index)));
}
updater.imageStore(6, 0, outImgViewList);
updater.apply();

updater = new DescriptorUpdateBuilder(ctx, customTextureViews.length)
.set(ctexSet);

for (int i = 0; i < customTextureViews.length; i++) {
updater.imageSampler(i, customTextureViews[i].getView(), ctexSampler);
}
updater.apply();

updater = new DescriptorUpdateBuilder(ctx, ssbos.length)
.set(ssboSet);

for (ShaderStorageBuffer ssbo : ssbos) {
updater.buffer(ssbo.getIndex(), ((IVGBuffer) ssbo).getBuffer());
// Call getView() on shared image view trackers to ensure they are created
blockAtlasView.getView();
blockAtlasNormalView.getView();
blockAtlasSpecularView.getView();
for (var v : customTextureViews) {
v.getView();
}
updater.apply();

//TODO: dont use a single use pool for commands like this...
var cmd = singleUsePool.createCommandBuffer();
Expand All @@ -392,23 +318,69 @@ public void renderPostShadows(List<VGImage> outImgs, Camera camera, ShaderStorag
}
}

for (var pipeline : raytracePipelines) {
pipeline.pipeline.bind(cmd);
var sets = new TreeMap<Integer, Long>();
if (pipeline.commonSet != -1) {
sets.put(pipeline.commonSet, commonSet);
for (var record : raytracePipelines) {
var pipeline = record.pipeline;
pipeline.bind(cmd);
var layouts = pipeline.reflection.getLayouts(); // Should be cached already
var sets = new long[layouts.size()];
if (record.commonSet != -1) {
var commonSet = Vulkanite.INSTANCE.getPoolByLayout(layouts.get(record.commonSet)).allocateSet();

var updater = new DescriptorUpdateBuilder(ctx, pipeline.reflection.getSet(record.commonSet))
.set(commonSet.set)
.uniform(0, uboBuffer)
.acceleration(1, tlas)
.imageSampler(3, blockAtlasView.getView(), sampler)
.imageSampler(4,
blockAtlasNormalView.getView() != null ? blockAtlasNormalView.getView()
: placeholderNormalsView,
sampler)
.imageSampler(5,
blockAtlasSpecularView.getView() != null ? blockAtlasSpecularView.getView()
: placeholderSpecularView,
sampler);
List<VImageView> outImgViewList = new ArrayList<>(outImgs.size());
for (int i = 0; i < outImgs.size(); i++) {
int index = i;
outImgViewList.add(irisRenderTargetViews[i].getView(() -> outImgs.get(index)));
}
updater.imageStore(6, 0, outImgViewList);
updater.apply();

sets[record.commonSet] = commonSet.set;
cmd.addTransientResource(commonSet);
}
if (pipeline.geomSet != -1) {
sets.put(pipeline.geomSet, accelerationManager.getGeometrySet());
if (record.geomSet != -1) {
sets[record.geomSet] = accelerationManager.getGeometrySet();
}
if (pipeline.customTexSet != -1) {
sets.put(pipeline.customTexSet, ctexSet);
if (record.customTexSet != -1) {
var ctexSet = Vulkanite.INSTANCE.getPoolByLayout(layouts.get(record.customTexSet)).allocateSet();

var updater = new DescriptorUpdateBuilder(ctx, pipeline.reflection.getSet(record.customTexSet))
.set(ctexSet.set);
for (int i = 0; i < customTextureViews.length; i++) {
updater.imageSampler(i, customTextureViews[i].getView(), ctexSampler);
}
updater.apply();

sets[record.customTexSet] = ctexSet.set;
cmd.addTransientResource(ctexSet);
}
if (pipeline.ssboSet != -1) {
sets.put(pipeline.ssboSet, ssboSet);
if (record.ssboSet != -1) {
var ssboSet = Vulkanite.INSTANCE.getPoolByLayout(layouts.get(record.ssboSet)).allocateSet();

var updater = new DescriptorUpdateBuilder(ctx, pipeline.reflection.getSet(record.ssboSet))
.set(ssboSet.set);
for (ShaderStorageBuffer ssbo : ssbos) {
updater.buffer(ssbo.getIndex(), ((IVGBuffer) ssbo).getBuffer());
}
updater.apply();

sets[record.ssboSet] = ssboSet.set;
cmd.addTransientResource(ssboSet);
}
pipeline.pipeline.bindDSet(cmd, sets.values().stream().mapToLong(a->a).toArray());
pipeline.pipeline.trace(cmd, outImgs.get(0).width, outImgs.get(0).height, 1);
pipeline.bindDSet(cmd, sets);
pipeline.trace(cmd, outImgs.get(0).width, outImgs.get(0).height, 1);

// Barrier on the output images
for (var img : outImgs) {
Expand Down Expand Up @@ -459,23 +431,10 @@ public void renderPostShadows(List<VGImage> outImgs, Camera camera, ShaderStorag
}

public void destory() {
for (var pass : raytracePipelines) {
if (pass != null) {
pass.pipeline.free();
}
}
if (commonLayout != null)
commonLayout.free();
if (customtexLayout != null)
customtexLayout.free();
if (storageBufferLayout != null)
storageBufferLayout.free();
if (commonDescriptorPool != null)
commonDescriptorPool.free();
if (customtexDescriptorPool != null)
customtexDescriptorPool.free();
if (storageBufferDescriptorPool != null)
storageBufferDescriptorPool.free();
vkDeviceWaitIdle(ctx.device);

// Check pending fences first
// Then destroy the cmd pool (which destroys linked transient resources)
ctx.sync.checkFences();
if (singleUsePool != null) {
singleUsePool.doReleases();
Expand All @@ -484,6 +443,13 @@ public void destory() {
if (previousSemaphore != null) {
previousSemaphore.free();
}
// Finally destroy the pipelines
// (Which destroys the descriptor set layouts & releases the VTypedDescriptorPool)
for (var pass : raytracePipelines) {
if (pass != null) {
pass.pipeline.free();
}
}

for (SharedImageViewTracker customTexView : customTextureViews) {
if (customTexView != null)
Expand Down
19 changes: 12 additions & 7 deletions src/main/java/me/cortex/vulkanite/lib/cmd/VCmdBuff.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,19 +25,21 @@
import static org.lwjgl.vulkan.KHRAccelerationStructure.VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR;
import static org.lwjgl.vulkan.KHRAccelerationStructure.VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR;

import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.LinkedList;

//TODO: Track with TrackedResourceObject but need to be careful due to how the freeing works
public class VCmdBuff extends TrackedResourceObject implements Pointer {
private final VCommandPool pool;
public final VkCommandBuffer buffer;

private LinkedList<VBuffer> transientBuffers;
private HashSet<TrackedResourceObject> transientResources;

VCmdBuff(VCommandPool pool, VkCommandBuffer buff) {
this.pool = pool;
this.buffer = buff;
this.transientBuffers = new LinkedList<>();
this.transientResources = new HashSet<>();
}

//Enqueues the pool to be freed by the owning thread
Expand Down Expand Up @@ -69,7 +71,7 @@ public void encodeDataUpload(MemoryManager manager, long src, VBuffer dest, long
vkCmdCopyBuffer(buffer, staging.buffer(), dest.buffer(), copy);
}

transientBuffers.add(staging);
transientResources.add(staging);
}

public void encodeImageUpload(MemoryManager manager, long src, VImage dest, long srcSize, int destLayout) {
Expand All @@ -89,7 +91,7 @@ public void encodeImageUpload(MemoryManager manager, long src, VImage dest, long
vkCmdCopyBufferToImage(buffer, staging.buffer(), dest.image(), destLayout, copy);
}

transientBuffers.add(staging);
transientResources.add(staging);
}

public void encodeMemoryBarrier() {
Expand All @@ -102,6 +104,10 @@ public void encodeMemoryBarrier() {
}
}

public void addTransientResource(TrackedResourceObject resource) {
transientResources.add(resource);
}

public static int dstStageToAccess(int dstStage) {
switch (dstStage) {
case VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT:
Expand Down Expand Up @@ -200,8 +206,7 @@ public void free() {
void freeInternal() {
free0();
vkFreeCommandBuffers(pool.device, pool.pool, buffer);
while (!transientBuffers.isEmpty()) {
transientBuffers.removeFirst().free();
}
transientResources.forEach(TrackedResourceObject::free);
transientResources.clear();
}
}
Loading

0 comments on commit e112e31

Please sign in to comment.