Skip to content

Commit

Permalink
Merge pull request #16 from bobcao3/master
Browse files Browse the repository at this point in the history 
BINDLESS GEOMETRY BUFFERS!
  • Loading branch information
@MCRcortex
MCRcortex authored Sep 22, 2023
2 parents 3995adc + 301ee58 commit 230fee3
Show file tree
Hide file tree
Showing 12 changed files with 372 additions and 154 deletions.
2 changes: 1 addition & 1 deletion src/main/java/me/cortex/vulkanite/acceleration/AccelerationBlasBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ private void run() {

private VBuffer uploadTerrainGeometry(BuiltSectionMeshParts meshParts, VCmdBuff cmd) {
var buff = context.memory.createBuffer(meshParts.getVertexData().getLength(),
VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

cmd.encodeDataUpload(context.memory, MemoryUtil.memAddress(meshParts.getVertexData().getDirectBuffer()), buff, 0, meshParts.getVertexData().getLength());
Expand Down
9 changes: 7 additions & 2 deletions src/main/java/me/cortex/vulkanite/acceleration/AccelerationManager.java
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package me.cortex.vulkanite.acceleration;

import me.cortex.vulkanite.lib.base.VContext;
import me.cortex.vulkanite.lib.descriptors.VDescriptorSetLayout;
import me.cortex.vulkanite.lib.memory.VAccelerationStructure;
import me.cortex.vulkanite.lib.memory.VBuffer;
import me.cortex.vulkanite.lib.other.sync.VSemaphore;
Expand Down Expand Up @@ -72,7 +73,11 @@ public void cleanup() {
tlasManager.cleanupTick();
}

public VBuffer getReferenceBuffer() {
return tlasManager.getReferenceBuffer();
public long getGeometrySet() {
return tlasManager.getGeometrySet();
}

public VDescriptorSetLayout getGeometryLayout() {
return tlasManager.getGeometryLayout();
}
}
109 changes: 89 additions & 20 deletions src/main/java/me/cortex/vulkanite/acceleration/AccelerationTLASManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,25 +7,29 @@
import me.cortex.vulkanite.lib.base.VContext;
import me.cortex.vulkanite.lib.cmd.VCmdBuff;
import me.cortex.vulkanite.lib.cmd.VCommandPool;
import me.cortex.vulkanite.lib.descriptors.DescriptorSetLayoutBuilder;
import me.cortex.vulkanite.lib.descriptors.DescriptorUpdateBuilder;
import me.cortex.vulkanite.lib.descriptors.VDescriptorPool;
import me.cortex.vulkanite.lib.descriptors.VDescriptorSetLayout;
import me.cortex.vulkanite.lib.memory.VAccelerationStructure;
import me.cortex.vulkanite.lib.memory.VBuffer;
import me.cortex.vulkanite.lib.other.sync.VFence;
import me.cortex.vulkanite.lib.other.sync.VSemaphore;
import me.jellysquid.mods.sodium.client.render.chunk.RenderSection;
import net.minecraft.util.math.ChunkSectionPos;
import net.minecraft.world.chunk.ChunkSection;
import org.joml.Matrix4x3f;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.*;

import java.util.*;
import java.util.concurrent.LinkedBlockingDeque;

import static org.lwjgl.system.MemoryStack.stackPush;
import static org.lwjgl.util.vma.Vma.VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT;
import static org.lwjgl.vulkan.KHRAccelerationStructure.*;
import static org.lwjgl.vulkan.KHRBufferDeviceAddress.VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT_KHR;
import static org.lwjgl.vulkan.VK10.*;
import static org.lwjgl.vulkan.VK12.VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
import static org.lwjgl.vulkan.VK12.*;

public class AccelerationTLASManager {
private final TLASSectionManager buildDataManager = new TLASSectionManager();
Expand All @@ -41,6 +45,7 @@ public AccelerationTLASManager(VContext context, int queue) {
this.context = context;
this.queue = queue;
this.singleUsePool = context.cmd.createSingleUsePool();
this.buildDataManager.resizeBindlessSet(0, null);
}

//Returns a sync semaphore to chain in the next command submit
Expand Down Expand Up @@ -309,25 +314,88 @@ protected void update(int id, VkAccelerationStructureInstanceKHR data) {
}
}

private static int roundUpPow2(int v) {
v--;
v |= v >> 1;
v |= v >> 2;
v |= v >> 4;
v |= v >> 8;
v |= v >> 16;
v++;
return v;
}

private final class TLASSectionManager extends TLASGeometryManager {
private final TlasPointerArena arena = new TlasPointerArena(30000);
private final long arrayRef = MemoryUtil.nmemCalloc(30000 * 3, 8);
public VBuffer geometryReferenceBuffer;

private VDescriptorSetLayout geometryBufferSetLayout;
private VDescriptorPool geometryBufferDescPool;
private long geometryBufferDescSet = 0;

private int setCapacity = 0;

private record DescUpdateJob(int binding, int dstArrayElement, List<VBuffer> buffers) {}
private final LinkedBlockingDeque<DescUpdateJob> descUpdateJobs = new LinkedBlockingDeque<>();

public void resizeBindlessSet(int newSize, VFence fence) {
if (geometryBufferSetLayout == null) {
var layoutBuilder = new DescriptorSetLayoutBuilder(VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT);
layoutBuilder.binding(0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 65536, VK_SHADER_STAGE_ALL);
layoutBuilder.setBindingFlags(0, VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT | VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT | VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT);
geometryBufferSetLayout = layoutBuilder.build(context);
}

if (newSize > setCapacity) {
int newCapacity = roundUpPow2(Math.max(newSize, 32));
var newGeometryBufferDescPool = new VDescriptorPool(context, VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT, 1, newCapacity, geometryBufferSetLayout.types);
newGeometryBufferDescPool.allocateSets(geometryBufferSetLayout, new int[]{newCapacity});
long newGeometryBufferDescSet = newGeometryBufferDescPool.get(0);

System.out.println("New geometry desc set: " + Long.toHexString(newGeometryBufferDescSet) + " with capacity " + newCapacity);

if (geometryBufferDescSet != 0) {
try (var stack = stackPush()) {
var setCopy = VkCopyDescriptorSet.calloc(1, stack);
setCopy.get(0)
.sType$Default()
.srcSet(geometryBufferDescSet)
.dstSet(newGeometryBufferDescSet)
.descriptorCount(setCapacity);
vkUpdateDescriptorSets(context.device, null, setCopy);
}

// This breaks the shit out of it
// context.sync.addCallback(fence, () -> {
// geometryBufferDescPool.free();
// });

vkDeviceWaitIdle(context.device);
geometryBufferDescPool.free();
}

geometryBufferDescPool = newGeometryBufferDescPool;
geometryBufferDescSet = newGeometryBufferDescSet;
setCapacity = newCapacity;
}

}

@Override
public void setGeometryUpdateMemory(VCmdBuff cmd, VFence fence, VkAccelerationStructureGeometryKHR struct) {
super.setGeometryUpdateMemory(cmd, fence, struct);
var ref = geometryReferenceBuffer;
if (ref != null) {
context.sync.addCallback(fence, ref::free);
resizeBindlessSet(arena.maxIndex, fence);

if (descUpdateJobs.isEmpty()) {
return;
}
geometryReferenceBuffer = context.memory.createBuffer(8L * arena.maxIndex,
VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
VK_MEMORY_HEAP_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
0, VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT);
long ptr = geometryReferenceBuffer.map();
MemoryUtil.memCopy(arrayRef, ptr, 8L * arena.maxIndex);
geometryReferenceBuffer.unmap();

var dub = new DescriptorUpdateBuilder(context, descUpdateJobs.size());
dub.set(geometryBufferDescSet);
while (!descUpdateJobs.isEmpty()) {
var job = descUpdateJobs.poll();
dub.buffer(job.binding, job.dstArrayElement, job.buffers);
}
dub.apply();
}

//TODO: mixinto RenderSection and add a reference to a holder for us, its much faster than a hashmap
Expand Down Expand Up @@ -362,10 +430,7 @@ public void update(AccelerationBlasBuilder.BLASBuildResult result) {
holder.geometryBuffers = data.geometryBuffers();
holder.geometryIndex = arena.allocate(holder.geometryBuffers.size());

for (int i = 0; i < holder.geometryBuffers.size(); i++) {
MemoryUtil.memPutAddress(arrayRef + 8L*(holder.geometryIndex+i), holder.geometryBuffers.get(i).deviceAddress());
}

descUpdateJobs.add(new DescUpdateJob(0, holder.geometryIndex, holder.geometryBuffers));

try (var stack = stackPush()) {
var asi = VkAccelerationStructureInstanceKHR.calloc(stack)
Expand Down Expand Up @@ -432,8 +497,12 @@ public void free(int pos, int count) {
}
}

public VBuffer getReferenceBuffer() {
return buildDataManager.geometryReferenceBuffer;
public long getGeometrySet() {
return buildDataManager.geometryBufferDescSet;
}

public VDescriptorSetLayout getGeometryLayout() {
return buildDataManager.geometryBufferSetLayout;
}

//Called for cleaning up any remaining loose resources
Expand Down
49 changes: 27 additions & 22 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 @@ -61,7 +61,10 @@ public class VulkanPipeline {
private VDescriptorSetLayout commonLayout;
private VDescriptorSetLayout customtexLayout;
private VDescriptorSetLayout storageBufferLayout;
private VDescriptorPool descriptors;

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

private final VSampler sampler;
private final VSampler ctexSampler;
Expand Down Expand Up @@ -104,7 +107,7 @@ public VulkanPipeline(VContext ctx, AccelerationManager accelerationManager, Ray
PBRTextureHolder holder = PBRTextureManager.INSTANCE.getOrLoadHolder(blockAtlas.getGlId());//((TextureAtlasExtension)blockAtlas).getPBRHolder()
return ((IVGImage)holder.getSpecularTexture()).getVGImage();
});
this.placeholderImage = ctx.memory.createImage2D(1, 1, 1, VK_FORMAT_R8G8B8A8_UNORM, VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
this.placeholderImage = ctx.memory.createImage2D(4, 4, 1, VK_FORMAT_R8G8B8A8_UNORM, VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
this.placeholderImageView = new VImageView(ctx, placeholderImage);

try (var stack = stackPush()) {
Expand Down Expand Up @@ -146,14 +149,13 @@ public VulkanPipeline(VContext ctx, AccelerationManager accelerationManager, Ray

try {
commonLayout = new DescriptorSetLayoutBuilder()
.binding(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_ALL)// camera data
.binding(VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, VK_SHADER_STAGE_ALL)// funni acceleration buffer
.binding(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_SHADER_STAGE_ALL)// funni buffer buffer
.binding(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL)//block texture
.binding(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL)//block texture normal
.binding(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_ALL)//block texture specular
.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(VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, VK_SHADER_STAGE_ALL)// output texture
.binding(6, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, VK_SHADER_STAGE_ALL)// output texture
.build(ctx);

DescriptorSetLayoutBuilder ctexLayoutBuilder = new DescriptorSetLayoutBuilder();
Expand All @@ -170,19 +172,21 @@ public VulkanPipeline(VContext ctx, AccelerationManager accelerationManager, Ray

storageBufferLayout = ssboLayoutBuilder.build(ctx);

// Using commonLayout.types is good enough because both VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER and VK_DESCRIPTOR_TYPE_STORAGE_BUFFER are used there already...
//TODO: use frameahead count instead of just... 10
descriptors = new VDescriptorPool(ctx, VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT, 10, commonLayout.types);
descriptors.allocateSets(new VDescriptorSetLayout[]{
commonLayout,
customtexLayout,
storageBufferLayout
});
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);

raytracePipelines = new VRaytracePipeline[passes.length];
for (int i = 0; i < passes.length; i++) {
var builder = new RaytracePipelineBuilder()
.addLayout(commonLayout)
.addLayout(accelerationManager.getGeometryLayout())
.addLayout(customtexLayout)
.addLayout(storageBufferLayout);
passes[i].apply(builder);
Expand Down Expand Up @@ -265,15 +269,14 @@ public void renderPostShadows(VGImage outImg, Camera camera, ShaderStorageBuffer
uboBuffer.unmap();
uboBuffer.flush();

long commonSet = descriptors.get(0);
long ctexSet = descriptors.get(1);
long ssboSet = descriptors.get(2);
long commonSet = commonDescriptorPool.get(fidx);
long ctexSet = customtexDescriptorPool.get(fidx);
long ssboSet = storageBufferDescriptorPool.get(fidx);

var updater = new DescriptorUpdateBuilder(ctx, 7, placeholderImageView)
.set(commonSet)
.uniform(0, uboBuffer)
.acceleration(1, tlas)
.buffer(2, accelerationManager.getReferenceBuffer())
.imageSampler(3, blockAtlasView.getView(), sampler)
.imageSampler(4, blockAtlasNormalView.getView(), sampler)
.imageSampler(5, blockAtlasSpecularView.getView(), sampler)
Expand Down Expand Up @@ -320,7 +323,7 @@ public void renderPostShadows(VGImage outImg, Camera camera, ShaderStorageBuffer

for (var pipeline : raytracePipelines) {
pipeline.bind(cmd);
pipeline.bindDSet(cmd, new long[]{commonSet, ctexSet, ssboSet});
pipeline.bindDSet(cmd, commonSet, accelerationManager.getGeometrySet(), ctexSet, ssboSet);
pipeline.trace(cmd, outImg.width, outImg.height, 1);
}

Expand Down Expand Up @@ -359,7 +362,9 @@ public void destory() {
commonLayout.free();
customtexLayout.free();
storageBufferLayout.free();
descriptors.free();
commonDescriptorPool.free();
customtexDescriptorPool.free();
storageBufferDescriptorPool.free();
ctx.sync.checkFences();
singleUsePool.doReleases();
singleUsePool.free();
Expand Down
19 changes: 19 additions & 0 deletions src/main/java/me/cortex/vulkanite/lib/descriptors/DescriptorSetLayoutBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,18 @@
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.VkDescriptorSetLayoutBinding;
import org.lwjgl.vulkan.VkDescriptorSetLayoutCreateInfo;
import org.lwjgl.vulkan.VkDescriptorSetLayoutBindingFlagsCreateInfo;

import java.nio.LongBuffer;
import java.util.HashMap;

import static me.cortex.vulkanite.lib.other.VUtil._CHECK_;
import static org.lwjgl.system.MemoryStack.stackPush;
import static org.lwjgl.vulkan.VK10.vkCreateDescriptorSetLayout;

public class DescriptorSetLayoutBuilder {
private IntArrayList types = new IntArrayList();
private HashMap<Integer, Integer> bindingFlagsMap = new HashMap<>();
private VkDescriptorSetLayoutBinding.Buffer bindings = VkDescriptorSetLayoutBinding.calloc(0);
public DescriptorSetLayoutBuilder binding(int binding, int type, int count, int stages) {
bindings = VkDescriptorSetLayoutBinding.create(MemoryUtil.nmemRealloc(bindings.address(), (bindings.capacity() + 1L) * VkDescriptorSetLayoutBinding.SIZEOF), bindings.capacity() + 1);
Expand All @@ -30,6 +33,10 @@ public DescriptorSetLayoutBuilder binding(int type, int stages) {
return binding(bindings.capacity(), type, stages);
}

public void setBindingFlags(int binding, int flag) {
bindingFlagsMap.put(binding, flag);
}

int flags;
public DescriptorSetLayoutBuilder() {
this(0);
Expand All @@ -45,6 +52,18 @@ public VDescriptorSetLayout build(VContext ctx) {
.pBindings(bindings)
.flags(flags);

if (!bindingFlagsMap.isEmpty()) {
var bindingFlags = new int[bindings.remaining()];
for (var i = 0; i < bindings.remaining(); i++) {
bindingFlags[i] = bindingFlagsMap.getOrDefault(bindings.get(i).binding(), 0);
}

var bindingInfo = VkDescriptorSetLayoutBindingFlagsCreateInfo.calloc(stack)
.sType$Default()
.pBindingFlags(stack.ints(bindingFlags));
info.pNext(bindingInfo);
}

LongBuffer pBuffer = stack.mallocLong(1);
_CHECK_(vkCreateDescriptorSetLayout(ctx.device, info, null, pBuffer));
return new VDescriptorSetLayout(ctx, pBuffer.get(0), types.toIntArray());
Expand Down
22 changes: 22 additions & 0 deletions src/main/java/me/cortex/vulkanite/lib/descriptors/DescriptorUpdateBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import static org.lwjgl.vulkan.KHRAccelerationStructure.VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR;
import static org.lwjgl.vulkan.VK10.*;

import java.util.List;

public class DescriptorUpdateBuilder {
private final VContext ctx;
private final MemoryStack stack;
Expand Down Expand Up @@ -60,6 +62,26 @@ public DescriptorUpdateBuilder buffer(int binding, VBuffer buffer, long offset,
return this;
}

public DescriptorUpdateBuilder buffer(int binding, int dstArrayElement, List<VBuffer> buffers) {
var bufInfo = VkDescriptorBufferInfo.calloc(buffers.size(), stack);
for (int i = 0; i < buffers.size(); i++) {
bufInfo.get(i)
.buffer(buffers.get(i).buffer())
.offset(0)
.range(VK_WHOLE_SIZE);
}
updates.get()
.sType$Default()
.dstBinding(binding)
.dstSet(set)
.dstArrayElement(dstArrayElement)
.descriptorType(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER)
.descriptorCount(buffers.size())
.pBufferInfo(bufInfo);

return this;
}


public DescriptorUpdateBuilder uniform(int binding, VBuffer buffer) {
return uniform(binding, buffer, 0, VK_WHOLE_SIZE);
Expand Down
Loading

0 comments on commit 230fee3

Please sign in to comment.