From d93367f219ab1718c36d7afca1d733aaafccd68d Mon Sep 17 00:00:00 2001
From: TheJackiMonster <thejackimonster@gmail.com>
Date: Tue, 26 Jul 2022 14:05:37 +0200
Subject: [PATCH] Separated voxel sampling from blending

Signed-off-by: TheJackiMonster <thejackimonster@gmail.com>
---
 projects/fire_works/shaders/add.comp         |  41 ++----
 projects/fire_works/shaders/clear.comp       |  25 ++++
 projects/fire_works/shaders/sample.comp      |  48 +++++++
 projects/fire_works/shaders/voxel.inc        |   8 +-
 projects/fire_works/shaders/voxel_smoke.comp |  40 ++++--
 projects/fire_works/src/main.cpp             | 144 ++++++++++++++-----
 6 files changed, 227 insertions(+), 79 deletions(-)
 create mode 100644 projects/fire_works/shaders/clear.comp
 create mode 100644 projects/fire_works/shaders/sample.comp

diff --git a/projects/fire_works/shaders/add.comp b/projects/fire_works/shaders/add.comp
index 48917d11..a13dc0ee 100644
--- a/projects/fire_works/shaders/add.comp
+++ b/projects/fire_works/shaders/add.comp
@@ -1,17 +1,13 @@
 #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;
+layout(set=0, binding=0) uniform texture2D voxelTexture;
+layout(set=0, binding=1) uniform sampler voxelSampler;
 
-#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(set=0, binding=2, rgba16f) restrict readonly uniform image2D inParticles;
+layout(set=0, binding=3, rgba16f) restrict readonly uniform image2D inSmoke;
+layout(set=0, binding=4, rgba16f) restrict readonly uniform image2D inTrails;
+layout(set=0, binding=5, rgba16f) restrict writeonly uniform image2D outImage;
 
 layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
 
@@ -28,6 +24,10 @@ void main() {
     vec4 outSmoke = imageLoad(inSmoke, uv);
     vec4 outTrails = imageLoad(inTrails, uv);
 
+    vec2 pos = (vec2(uv) + vec2(0.5f)) / vec2(res);
+
+    vec4 outSamples = texture(sampler2D(voxelTexture, voxelSampler), pos);
+
     // TODO: add noise to the smoke here!
 
     vec4 result = vec4(
@@ -38,26 +38,7 @@ 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 = outSamples;
 
     result.r = clamp(result.r, 0, 1);
     result.g = clamp(result.g, 0, 1);
diff --git a/projects/fire_works/shaders/clear.comp b/projects/fire_works/shaders/clear.comp
new file mode 100644
index 00000000..4668538c
--- /dev/null
+++ b/projects/fire_works/shaders/clear.comp
@@ -0,0 +1,25 @@
+#version 440
+#extension GL_GOOGLE_include_directive : enable
+
+#include "physics.inc"
+#include "voxel.inc"
+
+layout(set=0, binding=0, r32ui) restrict writeonly uniform uimage3D voxelRed;
+layout(set=0, binding=1, r32ui) restrict writeonly uniform uimage3D voxelGreen;
+layout(set=0, binding=2, r32ui) restrict writeonly uniform uimage3D voxelBlue;
+layout(set=0, binding=3, r32ui) restrict writeonly uniform uimage3D voxelDensity;
+
+layout(local_size_x = 4, local_size_y = 4, local_size_z = 4) in;
+
+void main() {
+    if(any(greaterThanEqual(gl_GlobalInvocationID.xyz, imageSize(voxelDensity)))){
+        return;
+    }
+
+    ivec3 pos = ivec3(gl_GlobalInvocationID.xyz);
+
+    voxel_write(voxelRed, pos, 0.0f);
+    voxel_write(voxelGreen, pos, 0.0f);
+    voxel_write(voxelBlue, pos, 0.0f);
+    voxel_write(voxelDensity, pos, mediumDensity);
+}
\ No newline at end of file
diff --git a/projects/fire_works/shaders/sample.comp b/projects/fire_works/shaders/sample.comp
new file mode 100644
index 00000000..df336992
--- /dev/null
+++ b/projects/fire_works/shaders/sample.comp
@@ -0,0 +1,48 @@
+#version 440
+#extension GL_GOOGLE_include_directive : enable
+
+#include "voxel.inc"
+
+layout(set=0, binding=0, r32ui) readonly uniform uimage3D voxelRed;
+layout(set=0, binding=1, r32ui) readonly uniform uimage3D voxelGreen;
+layout(set=0, binding=2, r32ui) readonly uniform uimage3D voxelBlue;
+layout(set=0, binding=3, r32ui) readonly uniform uimage3D voxelDensity;
+
+layout(set=1, binding=0, rgba16f) restrict writeonly uniform image2D outImage;
+
+layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
+
+void main() {
+    const ivec2 res = imageSize(outImage);
+
+    if(any(greaterThanEqual(gl_GlobalInvocationID.xy, res))){
+        return;
+    }
+
+    ivec2 uv = ivec2(gl_GlobalInvocationID.xy);
+
+    const ivec3 voxelRes = imageSize(voxelDensity);
+
+    vec4 voxel = vec4(0.0f);
+
+    for (int i = 0; i < voxelRes.z; i++) {
+        const ivec3 voxelPos = ivec3(uv, 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)
+        );
+    }
+
+    voxel.r = clamp(voxel.r, 0, 1);
+    voxel.g = clamp(voxel.g, 0, 1);
+    voxel.b = clamp(voxel.b, 0, 1);
+    voxel.a = clamp(voxel.a, 0, 1);
+
+    imageStore(outImage, uv, voxel);
+}
\ No newline at end of file
diff --git a/projects/fire_works/shaders/voxel.inc b/projects/fire_works/shaders/voxel.inc
index bb2e4269..8097da99 100644
--- a/projects/fire_works/shaders/voxel.inc
+++ b/projects/fire_works/shaders/voxel.inc
@@ -1,9 +1,11 @@
 #ifndef VOXEL_INC
 #define VOXEL_INC
 
-#define voxel_add(img, pos, value) imageAtomicAdd(img, ivec3((imageSize(img) - ivec3(1)) * pos), uint(0xFF * value))
+#define VOXEL_NORM_VALUE 0xFF
 
-#define voxel_write(img, pos, value) imageStore(img, pos, uvec4(0xFF * value));
-#define voxel_read(img, pos) imageLoad(img, pos).r / 255.0f;
+#define voxel_add(img, pos, value) imageAtomicAdd(img, ivec3((imageSize(img) - ivec3(1)) * pos), uint(VOXEL_NORM_VALUE * value))
+
+#define voxel_write(img, pos, value) imageStore(img, pos, uvec4(VOXEL_NORM_VALUE * value));
+#define voxel_read(img, pos) imageLoad(img, pos).r / float(VOXEL_NORM_VALUE);
 
 #endif // VOXEL_INC
\ No newline at end of file
diff --git a/projects/fire_works/shaders/voxel_smoke.comp b/projects/fire_works/shaders/voxel_smoke.comp
index 6b1ca497..f399eba7 100644
--- a/projects/fire_works/shaders/voxel_smoke.comp
+++ b/projects/fire_works/shaders/voxel_smoke.comp
@@ -22,6 +22,8 @@ layout( push_constant ) uniform constants{
     mat4 mvp;
 };
 
+#define NUM_SMOKE_SAMPLES 4
+
 void main() {
     uint id = gl_GlobalInvocationID.x;
 
@@ -38,23 +40,33 @@ void main() {
         return;
     }
 
-    vec4 cs_pos = mvp * vec4(position, 1);
+    vec3 offset = vec3(-size);
 
-    if (abs(cs_pos.w) <= 0.0f) {
-        return;
-    }
+    for (;offset.x <= size; offset.x += size / NUM_SMOKE_SAMPLES) {
+        for (;offset.y <= size; offset.y += size / NUM_SMOKE_SAMPLES) {
+            for (;offset.z <= size; offset.z += size / NUM_SMOKE_SAMPLES) {
+                vec4 cs_pos = mvp * vec4(position + offset, 1);
 
-    vec3 ndc_pos = cs_pos.xyz / cs_pos.w;
-    vec3 pos = (ndc_pos + vec3(1, 1, 0)) * vec3(0.5f, 0.5f, 1.0f);
+                if (abs(cs_pos.w) <= 0.0f) {
+                    return;
+                }
 
-    if ((any(greaterThanEqual(pos, vec3(1.5f)))) || (any(lessThanEqual(pos, vec3(-0.5f))))) {
-        return;
-    }
+                vec3 ndc_pos = cs_pos.xyz / cs_pos.w;
+                vec3 pos = (ndc_pos + vec3(1, 1, 0)) * vec3(0.5f, 0.5f, 1.0f);
+
+                if ((any(greaterThanEqual(pos, vec3(1.5f)))) || (any(lessThanEqual(pos, vec3(-0.5f))))) {
+                    return;
+                }
 
-    vec3 color = smokes[id].color;
+                vec3 color = smokes[id].color;
 
-    voxel_add(voxelRed, pos, color.r);
-    voxel_add(voxelGreen, pos, color.g);
-    voxel_add(voxelBlue, pos, color.b);
-    voxel_add(voxelDensity, pos, density);
+                float local_density = density * max(1.0f - length(offset / size), 0.0f);
+
+                voxel_add(voxelRed, pos, color.r);
+                voxel_add(voxelGreen, pos, color.g);
+                voxel_add(voxelBlue, pos, color.b);
+                voxel_add(voxelDensity, pos, local_density);
+            }
+        }
+    }
 }
diff --git a/projects/fire_works/src/main.cpp b/projects/fire_works/src/main.cpp
index 7c2f795c..de2a98b6 100644
--- a/projects/fire_works/src/main.cpp
+++ b/projects/fire_works/src/main.cpp
@@ -707,8 +707,7 @@ int main(int argc, const char **argv) {
 		voxelWidth,
 		voxelHeight,
 		voxelDepth,
-		false,
-		true
+		false, true
 	);
 	
 	vkcv::Image voxelGreen = core.createImage(
@@ -716,8 +715,7 @@ int main(int argc, const char **argv) {
 		voxelWidth,
 		voxelHeight,
 		voxelDepth,
-		false,
-		true
+		false, true
 	);
 	
 	vkcv::Image voxelBlue = core.createImage(
@@ -725,8 +723,7 @@ int main(int argc, const char **argv) {
 		voxelWidth,
 		voxelHeight,
 		voxelDepth,
-		false,
-		true
+		false, true
 	);
 	
 	vkcv::Image voxelDensity = core.createImage(
@@ -734,18 +731,41 @@ int main(int argc, const char **argv) {
 		voxelWidth,
 		voxelHeight,
 		voxelDepth,
-		false,
-		true
+		false, true
+	);
+	
+	vkcv::Image voxelSamples = core.createImage(
+		colorFormat,
+		voxelWidth,
+		voxelHeight,
+		1, false, true
+   	);
+	
+	vkcv::SamplerHandle voxelSampler = core.createSampler(
+		vkcv::SamplerFilterType::LINEAR,
+		vkcv::SamplerFilterType::LINEAR,
+		vkcv::SamplerMipmapMode::LINEAR,
+		vkcv::SamplerAddressMode::CLAMP_TO_EDGE
 	);
 	
+	vkcv::ShaderProgram voxelClearShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/clear.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		voxelClearShader.addShader(shaderStage, path);
+	});
+	
+	const auto& voxelBindings = voxelClearShader.getReflectedDescriptors().at(0);
+	auto voxelDescriptorSetLayout = core.createDescriptorSetLayout(voxelBindings);
+	
+	vkcv::ComputePipelineHandle voxelClearPipeline = core.createComputePipeline({
+		voxelClearShader,
+		{ voxelDescriptorSetLayout }
+	});
+	
 	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 }
@@ -771,6 +791,19 @@ int main(int argc, const char **argv) {
 		{ trailDescriptorLayout, voxelDescriptorSetLayout }
 	});
 	
+	vkcv::ShaderProgram voxelSampleShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/sample.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		voxelSampleShader.addShader(shaderStage, path);
+	});
+	
+	const auto& sampleBindings = voxelSampleShader.getReflectedDescriptors().at(1);
+	auto samplesDescriptorSetLayout = core.createDescriptorSetLayout(sampleBindings);
+	
+	vkcv::ComputePipelineHandle voxelSamplePipeline = core.createComputePipeline({
+		voxelSampleShader,
+		{ voxelDescriptorSetLayout, samplesDescriptorSetLayout }
+	});
+	
 	auto voxelDescriptorSet = core.createDescriptorSet(voxelDescriptorSetLayout);
 	
 	{
@@ -782,6 +815,14 @@ int main(int argc, const char **argv) {
 		core.writeDescriptorSet(voxelDescriptorSet, writes);
 	}
 	
+	auto samplesDescriptorSet = core.createDescriptorSet(samplesDescriptorSetLayout);
+	
+	{
+		vkcv::DescriptorWrites writes;
+		writes.writeStorageImage(0, voxelSamples.getHandle());
+		core.writeDescriptorSet(samplesDescriptorSet, writes);
+	}
+	
 	vkcv::effects::BloomAndFlaresEffect bloomAndFlares (core);
 	bloomAndFlares.setUpsamplingLimit(3);
 	
@@ -792,9 +833,10 @@ int main(int argc, const char **argv) {
 	
 	vkcv::DescriptorSetLayoutHandle addDescriptorLayout = core.createDescriptorSetLayout(addShader.getReflectedDescriptors().at(0));
 	vkcv::DescriptorSetHandle addDescriptor = core.createDescriptorSet(addDescriptorLayout);
+	
 	vkcv::ComputePipelineHandle addPipe = core.createComputePipeline({
 		addShader,
-		{ addDescriptorLayout, voxelDescriptorSetLayout }
+		{ addDescriptorLayout }
 	});
 	
 	vkcv::ShaderProgram tonemappingShader;
@@ -826,11 +868,11 @@ int main(int argc, const char **argv) {
 			if ((core.getImageWidth(colorBuffers[i]) != swapchainWidth) ||
 				(core.getImageHeight(colorBuffers[i]) != swapchainHeight)) {
 				colorBuffers[i] = core.createImage(
-										   colorFormat,
-										   swapchainWidth,
-										   swapchainHeight,
-										   1, false, true, true
-										   ).getHandle();
+					colorFormat,
+					swapchainWidth,
+					swapchainHeight,
+					1, false, true, true
+				).getHandle();
 			}
 		}
 		
@@ -847,13 +889,28 @@ 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);
 		
+		uint32_t voxelDispatchCount[3];
+		voxelDispatchCount[0] = std::ceil(voxelWidth / 4.f);
+		voxelDispatchCount[1] = std::ceil(voxelHeight / 4.f);
+		voxelDispatchCount[2] = std::ceil(voxelDepth / 4.f);
+		
+		core.recordBeginDebugLabel(cmdStream, "Voxel clear", { 0.5f, 0.25f, 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,
+			voxelClearPipeline,
+			voxelDispatchCount,
+			{ vkcv::DescriptorSetUsage(0, voxelDescriptorSet) },
+			vkcv::PushConstants(0)
+		);
+		core.recordEndDebugLabel(cmdStream);
+		
 		core.recordBufferMemoryBarrier(cmdStream, eventBuffer.getHandle());
 		core.recordBufferMemoryBarrier(cmdStream, smokeBuffer.getHandle());
 		core.recordBufferMemoryBarrier(cmdStream, smokeIndexBuffer.getHandle());
@@ -1056,20 +1113,46 @@ int main(int argc, const char **argv) {
 		);
 		core.recordEndDebugLabel(cmdStream);
 		
+		core.recordBeginDebugLabel(cmdStream, "Sample voxels", { 0.5f, 0.5f, 1.0f, 1.0f });
+		
+		core.prepareImageForStorage(cmdStream, voxelRed.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelGreen.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelBlue.getHandle());
+		core.prepareImageForStorage(cmdStream, voxelDensity.getHandle());
+		
+		core.prepareImageForStorage(cmdStream, voxelSamples.getHandle());
+		
+		uint32_t sampleDispatchCount[3];
+		sampleDispatchCount[0] = std::ceil(voxelWidth / 8.f);
+		sampleDispatchCount[1] = std::ceil(voxelHeight / 8.f);
+		sampleDispatchCount[2] = 1;
+		
+		core.recordComputeDispatchToCmdStream(
+			cmdStream,
+			voxelSamplePipeline,
+			sampleDispatchCount,
+			{
+				vkcv::DescriptorSetUsage(0, voxelDescriptorSet),
+				vkcv::DescriptorSetUsage(1, samplesDescriptorSet)
+			},
+			vkcv::PushConstants(0)
+		);
+		
+		core.recordEndDebugLabel(cmdStream);
+		
 		core.recordBeginDebugLabel(cmdStream, "Add rendered images", { 0.5f, 0.5f, 1.0f, 1.0f });
 		
 		vkcv::DescriptorWrites addDescriptorWrites;
+		addDescriptorWrites.writeSampledImage(0, voxelSamples.getHandle());
+		addDescriptorWrites.writeSampler(1, voxelSampler);
+		
 		for (size_t i = 0; i < colorBuffers.size(); i++) {
-			addDescriptorWrites.writeStorageImage(i, colorBuffers[i]);
+			addDescriptorWrites.writeStorageImage(2 + i, colorBuffers[i]);
 			core.prepareImageForStorage(cmdStream, colorBuffers[i]);
 		}
 		
 		core.writeDescriptorSet(addDescriptor, addDescriptorWrites);
-		
-		core.prepareImageForStorage(cmdStream, voxelRed.getHandle());
-		core.prepareImageForStorage(cmdStream, voxelGreen.getHandle());
-		core.prepareImageForStorage(cmdStream, voxelBlue.getHandle());
-		core.prepareImageForStorage(cmdStream, voxelDensity.getHandle());
+		core.prepareImageForSampling(cmdStream, voxelSamples.getHandle());
 		
 		uint32_t colorDispatchCount[3];
 		colorDispatchCount[0] = std::ceil(swapchainWidth / 8.f);
@@ -1080,10 +1163,7 @@ int main(int argc, const char **argv) {
 			cmdStream,
 			addPipe,
 			colorDispatchCount,
-			{
-				vkcv::DescriptorSetUsage(0, addDescriptor),
-				vkcv::DescriptorSetUsage(1, voxelDescriptorSet)
-			},
+			{ vkcv::DescriptorSetUsage(0, addDescriptor) },
 			vkcv::PushConstants(0)
 		);
 		
-- 
GitLab