From 44e6d7d08cb5e59865c8b8fb69d5dccc34786da0 Mon Sep 17 00:00:00 2001
From: TheJackiMonster <thejackimonster@gmail.com>
Date: Tue, 26 Jul 2022 12:10:22 +0200
Subject: [PATCH] Added voxel rendering

Signed-off-by: TheJackiMonster <thejackimonster@gmail.com>
---
 projects/fire_works/shaders/add.comp          |  33 +++-
 projects/fire_works/shaders/smoke.frag        |   2 +-
 projects/fire_works/shaders/voxel.inc         |   8 +-
 .../fire_works/shaders/voxel_particle.comp    |  18 +-
 projects/fire_works/shaders/voxel_smoke.comp  |  16 +-
 projects/fire_works/shaders/voxel_trail.comp  |  16 +-
 projects/fire_works/src/main.cpp              | 161 +++++++++++++++++-
 src/vkcv/ImageManager.cpp                     |   4 +-
 8 files changed, 233 insertions(+), 25 deletions(-)

diff --git a/projects/fire_works/shaders/add.comp b/projects/fire_works/shaders/add.comp
index 3187e18d..48917d11 100644
--- a/projects/fire_works/shaders/add.comp
+++ b/projects/fire_works/shaders/add.comp
@@ -1,14 +1,24 @@
 #version 440
+#extension GL_GOOGLE_include_directive : enable
 
 layout(set=0, binding=0, rgba16f) readonly uniform image2D inParticles;
 layout(set=0, binding=1, rgba16f) readonly uniform image2D inSmoke;
 layout(set=0, binding=2, rgba16f) readonly uniform image2D inTrails;
 layout(set=0, binding=3, rgba16f) writeonly uniform image2D outImage;
 
+#include "voxel.inc"
+
+layout(set=1, binding=0, r32ui) uniform uimage3D voxelRed;
+layout(set=1, binding=1, r32ui) uniform uimage3D voxelGreen;
+layout(set=1, binding=2, r32ui) uniform uimage3D voxelBlue;
+layout(set=1, binding=3, r32ui) uniform uimage3D voxelDensity;
+
 layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
 
 void main() {
-    if(any(greaterThanEqual(gl_GlobalInvocationID.xy, imageSize(outImage)))){
+    const ivec2 res = imageSize(outImage);
+
+    if(any(greaterThanEqual(gl_GlobalInvocationID.xy, res))){
         return;
     }
 
@@ -28,6 +38,27 @@ void main() {
         outParticles.a + outSmoke.a + outTrails.a
     );
 
+    const ivec3 voxelRes = imageSize(voxelDensity);
+
+    ivec2 voxelUV = uv * voxelRes.xy / res;
+    vec4 voxel = vec4(0.0f);
+
+    for (int i = 0; i < voxelRes.z; i++) {
+        const ivec3 voxelPos = ivec3(voxelUV, i);
+
+        const float red = voxel_read(voxelRed, voxelPos);
+        const float green = voxel_read(voxelGreen, voxelPos);
+        const float blue = voxel_read(voxelBlue, voxelPos);
+        const float density = voxel_read(voxelDensity, voxelPos);
+
+        voxel = vec4(
+            (voxel.rgb + vec3(red, green, blue) * density) * (1.0f - voxel.a),
+            voxel.a + (density) * (1.0f - voxel.a)
+        );
+    }
+
+    result = voxel;
+
     result.r = clamp(result.r, 0, 1);
     result.g = clamp(result.g, 0, 1);
     result.b = clamp(result.b, 0, 1);
diff --git a/projects/fire_works/shaders/smoke.frag b/projects/fire_works/shaders/smoke.frag
index ac5740d8..e5549df1 100644
--- a/projects/fire_works/shaders/smoke.frag
+++ b/projects/fire_works/shaders/smoke.frag
@@ -49,7 +49,7 @@ void main()	{
     result.r = clamp(result.r, 0, 1);
     result.g = clamp(result.g, 0, 1);
     result.b = clamp(result.b, 0, 1);
-    result.a = clamp(result.a, 0, 1);;
+    result.a = clamp(result.a, 0, 1);
 
     if (result.a < 1.0f) {
         outColor = result;
diff --git a/projects/fire_works/shaders/voxel.inc b/projects/fire_works/shaders/voxel.inc
index bbeb3446..bb2e4269 100644
--- a/projects/fire_works/shaders/voxel.inc
+++ b/projects/fire_works/shaders/voxel.inc
@@ -1,9 +1,9 @@
 #ifndef VOXEL_INC
 #define VOXEL_INC
 
-struct voxel_t {
-    vec3 color;
-    float density;
-};
+#define voxel_add(img, pos, value) imageAtomicAdd(img, ivec3((imageSize(img) - ivec3(1)) * pos), uint(0xFF * value))
+
+#define voxel_write(img, pos, value) imageStore(img, pos, uvec4(0xFF * value));
+#define voxel_read(img, pos) imageLoad(img, pos).r / 255.0f;
 
 #endif // VOXEL_INC
\ No newline at end of file
diff --git a/projects/fire_works/shaders/voxel_particle.comp b/projects/fire_works/shaders/voxel_particle.comp
index 8489e960..301c4c9d 100644
--- a/projects/fire_works/shaders/voxel_particle.comp
+++ b/projects/fire_works/shaders/voxel_particle.comp
@@ -1,4 +1,4 @@
-#version 450 core
+#version 450
 #extension GL_GOOGLE_include_directive : enable
 #extension GL_ARB_separate_shader_objects : enable
 
@@ -13,9 +13,10 @@ layout(set=0, binding=0, std430) readonly buffer particleBuffer {
 
 #include "voxel.inc"
 
-layout(set=1, binding=0, std430) buffer voxelBuffer {
-    voxel_t voxel [];
-};
+layout(set=1, binding=0, r32ui) uniform uimage3D voxelRed;
+layout(set=1, binding=1, r32ui) uniform uimage3D voxelGreen;
+layout(set=1, binding=2, r32ui) uniform uimage3D voxelBlue;
+layout(set=1, binding=3, r32ui) uniform uimage3D voxelDensity;
 
 layout( push_constant ) uniform constants{
     mat4 mvp;
@@ -44,10 +45,15 @@ void main() {
     vec3 ndc_pos = cs_pos.xyz / cs_pos.w;
     vec3 pos = (ndc_pos + vec3(1, 1, 0)) * vec3(0.5f, 0.5f, 1.0f);
 
-    // clipping!
+    if ((any(greaterThanEqual(pos, vec3(1.5f)))) || (any(lessThanEqual(pos, vec3(-0.5f))))) {
+        return;
+    }
 
     float size = particles[id].size;
     vec3 color = particles[id].color;
 
-    // write color into voxel at `pos * (resolution-1)` atomically!
+    voxel_add(voxelRed, pos, color.r);
+    voxel_add(voxelGreen, pos, color.g);
+    voxel_add(voxelBlue, pos, color.b);
+    voxel_add(voxelDensity, pos, 1.0f);
 }
diff --git a/projects/fire_works/shaders/voxel_smoke.comp b/projects/fire_works/shaders/voxel_smoke.comp
index da33a8ea..6b1ca497 100644
--- a/projects/fire_works/shaders/voxel_smoke.comp
+++ b/projects/fire_works/shaders/voxel_smoke.comp
@@ -13,9 +13,10 @@ layout(set=0, binding=0, std430) readonly buffer smokeBuffer {
 
 #include "voxel.inc"
 
-layout(set=1, binding=0, std430) buffer voxelBuffer {
-    voxel_t voxel [];
-};
+layout(set=1, binding=0, r32ui) uniform uimage3D voxelRed;
+layout(set=1, binding=1, r32ui) uniform uimage3D voxelGreen;
+layout(set=1, binding=2, r32ui) uniform uimage3D voxelBlue;
+layout(set=1, binding=3, r32ui) uniform uimage3D voxelDensity;
 
 layout( push_constant ) uniform constants{
     mat4 mvp;
@@ -46,9 +47,14 @@ void main() {
     vec3 ndc_pos = cs_pos.xyz / cs_pos.w;
     vec3 pos = (ndc_pos + vec3(1, 1, 0)) * vec3(0.5f, 0.5f, 1.0f);
 
-    // clipping!
+    if ((any(greaterThanEqual(pos, vec3(1.5f)))) || (any(lessThanEqual(pos, vec3(-0.5f))))) {
+        return;
+    }
 
     vec3 color = smokes[id].color;
 
-    // write (color, density) into voxel at `pos * (resolution-1)` atomically!
+    voxel_add(voxelRed, pos, color.r);
+    voxel_add(voxelGreen, pos, color.g);
+    voxel_add(voxelBlue, pos, color.b);
+    voxel_add(voxelDensity, pos, density);
 }
diff --git a/projects/fire_works/shaders/voxel_trail.comp b/projects/fire_works/shaders/voxel_trail.comp
index 6419cf72..667e402c 100644
--- a/projects/fire_works/shaders/voxel_trail.comp
+++ b/projects/fire_works/shaders/voxel_trail.comp
@@ -20,9 +20,10 @@ layout(set=0, binding=1, std430) buffer pointBuffer {
 
 #include "voxel.inc"
 
-layout(set=1, binding=0, std430) buffer voxelBuffer {
-    voxel_t voxel [];
-};
+layout(set=1, binding=0, r32ui) uniform uimage3D voxelRed;
+layout(set=1, binding=1, r32ui) uniform uimage3D voxelGreen;
+layout(set=1, binding=2, r32ui) uniform uimage3D voxelBlue;
+layout(set=1, binding=3, r32ui) uniform uimage3D voxelDensity;
 
 #include "smoke.inc"
 
@@ -74,8 +75,13 @@ void main() {
         vec3 ndc_pos = cs_pos.xyz / cs_pos.w;
         vec3 pos = (ndc_pos + vec3(1, 1, 0)) * vec3(0.5f, 0.5f, 1.0f);
 
-        // clipping!
+        if ((any(greaterThanEqual(pos, vec3(1.5f)))) || (any(lessThanEqual(pos, vec3(-0.5f))))) {
+            continue;
+        }
 
-        // write (color, density) into voxel at `pos * (resolution-1)` atomically!
+        voxel_add(voxelRed, pos, color.r);
+        voxel_add(voxelGreen, pos, color.g);
+        voxel_add(voxelBlue, pos, color.b);
+        voxel_add(voxelDensity, pos, density);
     }
 }
diff --git a/projects/fire_works/src/main.cpp b/projects/fire_works/src/main.cpp
index 692fbc14..7c2f795c 100644
--- a/projects/fire_works/src/main.cpp
+++ b/projects/fire_works/src/main.cpp
@@ -695,6 +695,93 @@ int main(int argc, const char **argv) {
 		{ descriptorSetLayout }
 	});
 	
+	const uint32_t voxelWidth = 160;
+	const uint32_t voxelHeight = 90;
+	const uint32_t voxelDepth = 64;
+	
+	std::vector<uint32_t> zeroVoxel;
+	zeroVoxel.resize(voxelWidth * voxelHeight * voxelDepth, 0);
+	
+	vkcv::Image voxelRed = core.createImage(
+		vk::Format::eR32Uint,
+		voxelWidth,
+		voxelHeight,
+		voxelDepth,
+		false,
+		true
+	);
+	
+	vkcv::Image voxelGreen = core.createImage(
+		vk::Format::eR32Uint,
+		voxelWidth,
+		voxelHeight,
+		voxelDepth,
+		false,
+		true
+	);
+	
+	vkcv::Image voxelBlue = core.createImage(
+		vk::Format::eR32Uint,
+		voxelWidth,
+		voxelHeight,
+		voxelDepth,
+		false,
+		true
+	);
+	
+	vkcv::Image voxelDensity = core.createImage(
+		vk::Format::eR32Uint,
+		voxelWidth,
+		voxelHeight,
+		voxelDepth,
+		false,
+		true
+	);
+	
+	vkcv::ShaderProgram voxelParticleShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/voxel_particle.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		voxelParticleShader.addShader(shaderStage, path);
+	});
+	
+	const auto& voxelBindings = voxelParticleShader.getReflectedDescriptors().at(1);
+	auto voxelDescriptorSetLayout = core.createDescriptorSetLayout(voxelBindings);
+	
+	vkcv::ComputePipelineHandle voxelParticlePipeline = core.createComputePipeline({
+		voxelParticleShader,
+		{ descriptorSetLayout, voxelDescriptorSetLayout }
+	});
+	
+	vkcv::ShaderProgram voxelSmokeShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/voxel_smoke.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		voxelSmokeShader.addShader(shaderStage, path);
+	});
+	
+	vkcv::ComputePipelineHandle voxelSmokePipeline = core.createComputePipeline({
+		voxelSmokeShader,
+		{ smokeDescriptorLayout, voxelDescriptorSetLayout }
+	});
+	
+	vkcv::ShaderProgram voxelTrailShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/voxel_trail.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		voxelTrailShader.addShader(shaderStage, path);
+	});
+	
+	vkcv::ComputePipelineHandle voxelTrailPipeline = core.createComputePipeline({
+		voxelTrailShader,
+		{ trailDescriptorLayout, voxelDescriptorSetLayout }
+	});
+	
+	auto voxelDescriptorSet = core.createDescriptorSet(voxelDescriptorSetLayout);
+	
+	{
+		vkcv::DescriptorWrites writes;
+		writes.writeStorageImage(0, voxelRed.getHandle());
+		writes.writeStorageImage(1, voxelGreen.getHandle());
+		writes.writeStorageImage(2, voxelBlue.getHandle());
+		writes.writeStorageImage(3, voxelDensity.getHandle());
+		core.writeDescriptorSet(voxelDescriptorSet, writes);
+	}
+	
 	vkcv::effects::BloomAndFlaresEffect bloomAndFlares (core);
 	bloomAndFlares.setUpsamplingLimit(3);
 	
@@ -707,7 +794,7 @@ int main(int argc, const char **argv) {
 	vkcv::DescriptorSetHandle addDescriptor = core.createDescriptorSet(addDescriptorLayout);
 	vkcv::ComputePipelineHandle addPipe = core.createComputePipeline({
 		addShader,
-		{ addDescriptorLayout }
+		{ addDescriptorLayout, voxelDescriptorSetLayout }
 	});
 	
 	vkcv::ShaderProgram tonemappingShader;
@@ -760,6 +847,11 @@ int main(int argc, const char **argv) {
 		
 		std::cout << time_values[0] << " " << time_values[1] << std::endl;
 		
+		voxelRed.fill(zeroVoxel.data());
+		voxelGreen.fill(zeroVoxel.data());
+		voxelBlue.fill(zeroVoxel.data());
+		voxelDensity.fill(zeroVoxel.data());
+		
 		auto cmdStream = core.createCommandStream(vkcv::QueueType::Graphics);
 		
 		core.recordBufferMemoryBarrier(cmdStream, eventBuffer.getHandle());
@@ -870,6 +962,27 @@ int main(int argc, const char **argv) {
 		);
 		core.recordEndDebugLabel(cmdStream);
 		
+		vkcv::PushConstants pushConstantsVoxel (sizeof(glm::mat4));
+		pushConstantsVoxel.appendDrawcall(camera.getMVP());
+		
+		core.recordBeginDebugLabel(cmdStream, "Particle voxel update", { 1.0f, 0.5f, 0.8f, 1.0f });
+		core.prepareImageForStorage(cmdStream, voxelRed.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelGreen.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelBlue.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelDensity.getHandle());
+		
+		core.recordComputeDispatchToCmdStream(
+			cmdStream,
+			voxelParticlePipeline,
+			particleDispatchCount,
+			{
+				vkcv::DescriptorSetUsage(0, descriptorSet),
+				vkcv::DescriptorSetUsage(1, voxelDescriptorSet)
+			},
+			pushConstantsVoxel
+		);
+		core.recordEndDebugLabel(cmdStream);
+		
 		core.recordBufferMemoryBarrier(cmdStream, smokeBuffer.getHandle());
 		
 		draw_smoke_t draw_smoke {
@@ -892,6 +1005,24 @@ int main(int argc, const char **argv) {
 		);
 		core.recordEndDebugLabel(cmdStream);
 		
+		core.recordBeginDebugLabel(cmdStream, "Smoke voxel update", { 1.0f, 0.7f, 0.8f, 1.0f });
+		core.prepareImageForStorage(cmdStream, voxelRed.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelGreen.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelBlue.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelDensity.getHandle());
+		
+		core.recordComputeDispatchToCmdStream(
+			cmdStream,
+			voxelSmokePipeline,
+			smokeDispatchCount,
+			{
+				vkcv::DescriptorSetUsage(0, smokeDescriptorSet),
+				vkcv::DescriptorSetUsage(1, voxelDescriptorSet)
+			},
+			pushConstantsVoxel
+		);
+		core.recordEndDebugLabel(cmdStream);
+		
 		core.recordBufferMemoryBarrier(cmdStream, trailBuffer.getHandle());
 		core.recordBufferMemoryBarrier(cmdStream, pointBuffer.getHandle());
 		
@@ -907,6 +1038,24 @@ int main(int argc, const char **argv) {
 		);
 		core.recordEndDebugLabel(cmdStream);
 		
+		core.recordBeginDebugLabel(cmdStream, "Trail voxel update", { 1.0f, 0.9f, 0.8f, 1.0f });
+		core.prepareImageForStorage(cmdStream, voxelRed.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelGreen.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelBlue.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelDensity.getHandle());
+		
+		core.recordComputeDispatchToCmdStream(
+			cmdStream,
+			voxelTrailPipeline,
+			trailDispatchCount,
+			{
+				vkcv::DescriptorSetUsage(0, trailDescriptorSet),
+				vkcv::DescriptorSetUsage(1, voxelDescriptorSet)
+			},
+			pushConstantsVoxel
+		);
+		core.recordEndDebugLabel(cmdStream);
+		
 		core.recordBeginDebugLabel(cmdStream, "Add rendered images", { 0.5f, 0.5f, 1.0f, 1.0f });
 		
 		vkcv::DescriptorWrites addDescriptorWrites;
@@ -917,6 +1066,11 @@ int main(int argc, const char **argv) {
 		
 		core.writeDescriptorSet(addDescriptor, addDescriptorWrites);
 		
+		core.prepareImageForStorage(cmdStream, voxelRed.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelGreen.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelBlue.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelDensity.getHandle());
+		
 		uint32_t colorDispatchCount[3];
 		colorDispatchCount[0] = std::ceil(swapchainWidth / 8.f);
 		colorDispatchCount[1] = std::ceil(swapchainHeight / 8.f);
@@ -926,7 +1080,10 @@ int main(int argc, const char **argv) {
 			cmdStream,
 			addPipe,
 			colorDispatchCount,
-			{vkcv::DescriptorSetUsage(0, addDescriptor) },
+			{
+				vkcv::DescriptorSetUsage(0, addDescriptor),
+				vkcv::DescriptorSetUsage(1, voxelDescriptorSet)
+			},
 			vkcv::PushConstants(0)
 		);
 		
diff --git a/src/vkcv/ImageManager.cpp b/src/vkcv/ImageManager.cpp
index 585c1faf..87837528 100644
--- a/src/vkcv/ImageManager.cpp
+++ b/src/vkcv/ImageManager.cpp
@@ -470,8 +470,10 @@ namespace vkcv {
 		switch (format) {
 			case vk::Format::eR8Unorm:
 				return 1;
+			case vk::Format::eR16Unorm:
+				return 2;
+			case vk::Format::eR32Uint:
 			case vk::Format::eR8G8B8A8Srgb:
-				return 4;
 			case vk::Format::eR8G8B8A8Unorm:
 				return 4;
 			case vk::Format::eR16G16B16A16Sfloat:
-- 
GitLab