Added voxel smoke rendering

Signed-off-by: TheJackiMonster <thejackimonster@gmail.com>

projects/fire_works/shaders/add.comp

@@ -9,8 +9,16 @@ 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(set=1, binding=0, std430) readonly buffer randomBuffer {
+    float randomData [];
+};
+
 layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
 
+#include "physics.inc"
+
+#define NUM_VOXEL_SAMPLES 32
+
 void main() {
     const ivec2 res = imageSize(outImage);
 
@@ -25,20 +33,33 @@ void main() {
     vec4 outTrails = imageLoad(inTrails, uv);
 
     vec2 pos = (vec2(uv) + vec2(0.5f)) / vec2(res);
+    vec2 size = vec2(textureSize(sampler2D(voxelTexture, voxelSampler), 0));
 
-    vec4 outSamples = texture(sampler2D(voxelTexture, voxelSampler), pos);
+    vec4 outSamples = vec4(0.0f);
 
-    // TODO: add noise to the smoke here!
+    const uint globalID = uv.y + uv.x * res.y;
 
-    vec4 result = vec4(
-        outParticles.rgb * outParticles.a +
-        outSmoke.rgb * outSmoke.a +
-        outTrails.rgb * outTrails.a,
+    for (uint i = 0; i < NUM_VOXEL_SAMPLES; i++) {
+        vec2 noise = vec2(
+            randomData[(globalID * NUM_VOXEL_SAMPLES * 2 + i * 2 + 0) % randomData.length()],
+            randomData[(globalID * NUM_VOXEL_SAMPLES * 2 + i * 2 + 1) % randomData.length()]
+        );
 
-        outParticles.a + outSmoke.a + outTrails.a
-    );
+        outSamples += texture(
+            sampler2D(voxelTexture, voxelSampler),
+            pos + noise * (NUM_VOXEL_SAMPLES - 1.0f) / size
+        ) * max(0.0f, 1.0f - length(noise));
+    }
+
+    outSamples /= NUM_VOXEL_SAMPLES;
+
+    // TODO: add noise to the smoke here!
 
-    result = outSamples;
+    vec4 result = vec4(0.0f);
+    result += vec4(outParticles.rgb * outParticles.a, outParticles.a);
+    result += vec4(outSmoke.rgb * outSmoke.a, outSmoke.a);
+    result += vec4(outTrails.rgb * outTrails.a, outTrails.a);
+    result += vec4(outSamples.rgb * outSamples.a, outSamples.a);
 
     result.r = clamp(result.r, 0, 1);
     result.g = clamp(result.g, 0, 1);

projects/fire_works/shaders/fluid.comp

+#version 450 core
+#extension GL_GOOGLE_include_directive : enable
+#extension GL_ARB_separate_shader_objects : enable
+
+layout(local_size_x = 4, local_size_y = 4, local_size_z = 4) in;
+
+#include "physics.inc"
+#include "voxel.inc"
+
+layout(set=0, binding=0) uniform texture3D voxelTexture;
+layout(set=0, binding=1) uniform sampler voxelSampler;
+layout(set=0, binding=2, rgba16) restrict writeonly uniform image3D fluidImage;
+
+vec4 getDataFrom(vec3 position, vec3 offset) {
+    return texture(
+        sampler3D(
+            voxelTexture,
+            voxelSampler
+        ),
+        position + offset
+    );
+}
+
+shared vec4 cachedData [4][4][4];
+
+void storeCachedData(vec3 position) {
+    uvec3 localId = gl_LocalInvocationID;
+    cachedData[localId.x][localId.y][localId.z] = getDataFrom(position, vec3(0));
+}
+
+vec4 getCachedData() {
+    uvec3 localId = gl_LocalInvocationID;
+    return cachedData[localId.x][localId.y][localId.z];
+}
+
+vec4 loadCachedData(vec3 position, ivec3 offset, ivec3 size) {
+    uvec3 localId = gl_LocalInvocationID;
+    ivec3 index = ivec3(localId) + offset;
+
+    if ((any(lessThan(index, ivec3(0)))) || (any(greaterThan(index, ivec3(gl_WorkGroupSize))))) {
+        return getDataFrom(position, vec3(offset) / vec3(size));
+    } else {
+        return cachedData[index.x][index.y][index.z];
+    }
+}
+
+void main() {
+    uvec3 id = gl_GlobalInvocationID;
+    ivec3 size = imageSize(fluidImage);
+
+    if (any(greaterThanEqual(id, size))) {
+        return;
+    }
+
+    vec3 position = (vec3(id) + vec3(0.5f)) / vec3(size);
+
+    storeCachedData(position);
+    memoryBarrierShared();
+    barrier();
+
+    vec4 extData [6];
+
+    extData[0] = loadCachedData(position, ivec3(+1, 0, 0), size);
+    extData[1] = loadCachedData(position, ivec3(-1, 0, 0), size);
+    extData[2] = loadCachedData(position, ivec3(0, +1, 0), size);
+    extData[3] = loadCachedData(position, ivec3(0, -1, 0), size);
+    extData[4] = loadCachedData(position, ivec3(0, 0, +1), size);
+    extData[5] = loadCachedData(position, ivec3(0, 0, -1), size);
+
+    vec4 data = vec4(0);
+
+    for (uint i = 0; i < 6; i++) {
+        data += extData[i];
+    }
+
+    data = mix(getCachedData(), (data / 6), flowRate);
+
+    imageStore(fluidImage, ivec3(id), data);
+}

projects/fire_works/shaders/physics.inc

@@ -6,7 +6,7 @@ const float pi = 3.14159f;
 const float g = 9.81f;
 const float friction = 0.001f;
 
-const float flowSpeed = 1.0f;
+const float flowRate = 0.75f;
 
 const float mediumDensity = 0.0001f;
 

projects/fire_works/shaders/sample.comp

@@ -3,11 +3,7 @@
 
 #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=0, binding=0, rgba16) restrict readonly uniform image3D voxelImage;
 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;
@@ -21,21 +17,18 @@ void main() {
 
     ivec2 uv = ivec2(gl_GlobalInvocationID.xy);
 
-    const ivec3 voxelRes = imageSize(voxelDensity);
+    const ivec3 voxelRes = imageSize(voxelImage);
 
     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);
+        vec4 data = imageLoad(voxelImage, voxelPos);
 
         voxel = vec4(
-            (voxel.rgb + vec3(red, green, blue) * density) * (1.0f - voxel.a),
-            voxel.a + (density) * (1.0f - voxel.a)
+            (voxel.rgb + data.rgb * data.a) * (1.0f - voxel.a),
+            voxel.a + (data.a) * (1.0f - voxel.a)
         );
     }
 

projects/fire_works/shaders/voxel.comp

+#version 450 core
+#extension GL_GOOGLE_include_directive : enable
+#extension GL_ARB_separate_shader_objects : enable
+
+layout(local_size_x = 4, local_size_y = 4, local_size_z = 4) in;
+
+#include "physics.inc"
+#include "voxel.inc"
+
+layout(set=0, binding=0, r32ui) restrict readonly uniform uimage3D voxelRed;
+layout(set=0, binding=1, r32ui) restrict readonly uniform uimage3D voxelGreen;
+layout(set=0, binding=2, r32ui) restrict readonly uniform uimage3D voxelBlue;
+layout(set=0, binding=3, r32ui) restrict readonly uniform uimage3D voxelDensity;
+
+layout(set=1, binding=0, rgba16) restrict writeonly uniform image3D voxelImage;
+
+void main() {
+    ivec3 pos = ivec3(gl_GlobalInvocationID);
+    ivec3 size = imageSize(voxelImage);
+
+    if (any(greaterThanEqual(pos, size))) {
+        return;
+    }
+
+    const float red = voxel_read(voxelRed, pos);
+    const float green = voxel_read(voxelGreen, pos);
+    const float blue = voxel_read(voxelBlue, pos);
+    const float density = voxel_read(voxelDensity, pos);
+
+    imageStore(voxelImage, pos, vec4(
+        red,
+        green,
+        blue,
+        density
+    ));
+}

projects/fire_works/src/main.cpp

@@ -76,7 +76,7 @@ struct draw_smoke_t {
 #define PARTICLE_COUNT (1024)
 #define SMOKE_COUNT (512)
 #define TRAIL_COUNT (2048)
-#define RANDOM_DATA_LENGTH (1024)
+#define RANDOM_DATA_LENGTH (4096)
 #define POINT_COUNT (2048 * 256)
 
 void InitializeParticles(std::vector<particle_t> &particles) {
@@ -734,6 +734,23 @@ int main(int argc, const char **argv) {
 		false, true
 	);
 	
+	std::array<vkcv::ImageHandle, 2> voxelData {
+		core.createImage(
+			vk::Format::eR16G16B16A16Sfloat,
+			voxelWidth,
+			voxelHeight,
+			voxelDepth,
+			false, true
+		).getHandle(),
+		core.createImage(
+			vk::Format::eR16G16B16A16Sfloat,
+			voxelWidth,
+			voxelHeight,
+			voxelDepth,
+			false, true
+		).getHandle()
+	};
+	
 	vkcv::Image voxelSamples = core.createImage(
 		colorFormat,
 		voxelWidth,
@@ -791,6 +808,32 @@ int main(int argc, const char **argv) {
 		{ trailDescriptorLayout, voxelDescriptorSetLayout }
 	});
 	
+	vkcv::ShaderProgram voxelShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/voxel.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		voxelShader.addShader(shaderStage, path);
+	});
+	
+	const auto& voxelOutBindings = voxelShader.getReflectedDescriptors().at(1);
+	auto voxelOutDescriptorSetLayout = core.createDescriptorSetLayout(voxelOutBindings);
+	
+	vkcv::ComputePipelineHandle voxelPipeline = core.createComputePipeline({
+		voxelShader,
+		{ voxelDescriptorSetLayout, voxelOutDescriptorSetLayout }
+	});
+	
+	vkcv::ShaderProgram fluidShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/fluid.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		fluidShader.addShader(shaderStage, path);
+	});
+	
+	const auto& fluidBindings = fluidShader.getReflectedDescriptors().at(0);
+	auto fluidDescriptorSetLayout = core.createDescriptorSetLayout(fluidBindings);
+	
+	vkcv::ComputePipelineHandle fluidPipeline = core.createComputePipeline({
+		fluidShader,
+		{ fluidDescriptorSetLayout }
+	});
+	
 	vkcv::ShaderProgram voxelSampleShader;
 	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/sample.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
 		voxelSampleShader.addShader(shaderStage, path);
@@ -801,7 +844,7 @@ int main(int argc, const char **argv) {
 	
 	vkcv::ComputePipelineHandle voxelSamplePipeline = core.createComputePipeline({
 		voxelSampleShader,
-		{ voxelDescriptorSetLayout, samplesDescriptorSetLayout }
+		{ voxelOutDescriptorSetLayout, samplesDescriptorSetLayout }
 	});
 	
 	auto voxelDescriptorSet = core.createDescriptorSet(voxelDescriptorSetLayout);
@@ -815,6 +858,35 @@ int main(int argc, const char **argv) {
 		core.writeDescriptorSet(voxelDescriptorSet, writes);
 	}
 	
+	auto voxelOutDescriptorSet = core.createDescriptorSet(voxelOutDescriptorSetLayout);
+	
+	{
+		vkcv::DescriptorWrites writes;
+		writes.writeStorageImage(0, voxelData[0]);
+		core.writeDescriptorSet(voxelOutDescriptorSet, writes);
+	}
+	
+	std::array<vkcv::DescriptorSetHandle, 2> fluidDescriptorSet {
+		core.createDescriptorSet(fluidDescriptorSetLayout),
+		core.createDescriptorSet(fluidDescriptorSetLayout)
+	};
+	
+	{
+		vkcv::DescriptorWrites writes;
+		writes.writeSampledImage(0, voxelData[0]);
+		writes.writeSampler(1, voxelSampler);
+		writes.writeStorageImage(2, voxelData[1]);
+		core.writeDescriptorSet(fluidDescriptorSet[0], writes);
+	}
+	
+	{
+		vkcv::DescriptorWrites writes;
+		writes.writeSampledImage(0, voxelData[1]);
+		writes.writeSampler(1, voxelSampler);
+		writes.writeStorageImage(2, voxelData[0]);
+		core.writeDescriptorSet(fluidDescriptorSet[1], writes);
+	}
+	
 	auto samplesDescriptorSet = core.createDescriptorSet(samplesDescriptorSetLayout);
 	
 	{
@@ -836,7 +908,7 @@ int main(int argc, const char **argv) {
 	
 	vkcv::ComputePipelineHandle addPipe = core.createComputePipeline({
 		addShader,
-		{ addDescriptorLayout }
+		{ addDescriptorLayout, generationDescriptorLayout }
 	});
 	
 	vkcv::ShaderProgram tonemappingShader;
@@ -1113,13 +1185,48 @@ int main(int argc, const char **argv) {
 		);
 		core.recordEndDebugLabel(cmdStream);
 		
-		core.recordBeginDebugLabel(cmdStream, "Sample voxels", { 0.5f, 0.5f, 1.0f, 1.0f });
+		core.recordBeginDebugLabel(cmdStream, "Combine voxel data", { 0.5f, 0.5f, 0.5f, 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, voxelData[0]);
+		
+		core.recordComputeDispatchToCmdStream(
+			cmdStream,
+			voxelPipeline,
+			voxelDispatchCount,
+			{
+				vkcv::DescriptorSetUsage(0, voxelDescriptorSet),
+				vkcv::DescriptorSetUsage(1, voxelOutDescriptorSet)
+			},
+			vkcv::PushConstants(0)
+		);
+		
+		core.recordEndDebugLabel(cmdStream);
+		
+		core.recordBeginDebugLabel(cmdStream, "Fluid voxel data", { 0.2f, 0.2f, 0.9f, 1.0f });
+		
+		for (size_t i = 0; i < 8; i++) {
+			core.prepareImageForSampling(cmdStream, voxelData[i % 2]);
+			core.prepareImageForStorage(cmdStream, voxelData[(i + 1) % 2]);
+		
+			core.recordComputeDispatchToCmdStream(
+				cmdStream,
+				fluidPipeline,
+				voxelDispatchCount,
+				{ vkcv::DescriptorSetUsage(0, fluidDescriptorSet[i % 2]) },
+				vkcv::PushConstants(0)
+			);
+		}
+		
+		core.recordEndDebugLabel(cmdStream);
+		
+		core.recordBeginDebugLabel(cmdStream, "Sample voxels", { 0.5f, 0.5f, 1.0f, 1.0f });
+		
+		core.prepareImageForStorage(cmdStream, voxelData[0]);
 		core.prepareImageForStorage(cmdStream, voxelSamples.getHandle());
 		
 		uint32_t sampleDispatchCount[3];
@@ -1132,7 +1239,7 @@ int main(int argc, const char **argv) {
 			voxelSamplePipeline,
 			sampleDispatchCount,
 			{
-				vkcv::DescriptorSetUsage(0, voxelDescriptorSet),
+				vkcv::DescriptorSetUsage(0, voxelOutDescriptorSet),
 				vkcv::DescriptorSetUsage(1, samplesDescriptorSet)
 			},
 			vkcv::PushConstants(0)
@@ -1163,7 +1270,10 @@ int main(int argc, const char **argv) {
 			cmdStream,
 			addPipe,
 			colorDispatchCount,
-			{ vkcv::DescriptorSetUsage(0, addDescriptor) },
+			{
+				vkcv::DescriptorSetUsage(0, addDescriptor),
+				vkcv::DescriptorSetUsage(1, generationDescriptorSet)
+			},
 			vkcv::PushConstants(0)
 		);