Added fluid simulation for the smoke

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

projects/fire_works/shaders/fluid.comp

+#version 450 core
+#extension GL_GOOGLE_include_directive : enable
+#extension GL_ARB_separate_shader_objects : enable
+#extension GL_EXT_nonuniform_qualifier : enable
+
+layout(local_size_x = 4, local_size_y = 4, local_size_z = 4) in;
+
+#include "physics.inc"
+
+layout(set=0, binding=0) uniform sampler smokeSampler;
+layout(set=0, binding=1) uniform texture3D smokeTextures [];
+
+layout(set=1, binding=0, rgba16) uniform image3D fluidImage;
+
+layout( push_constant ) uniform constants{
+    uint smokeTextureIndex;
+    float dt;
+};
+
+vec4 getDataFrom(vec3 position, vec3 offset) {
+    return texture(
+        sampler3D(
+            smokeTextures[nonuniformEXT(smokeTextureIndex)],
+            smokeSampler
+        ),
+        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();
+
+    const float flowSpeed = 1.0f;
+    const float flowChange = flowSpeed * dt;
+
+    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), flowChange);
+
+    imageStore(fluidImage, ivec3(id), data);
+}

projects/fire_works/src/main.cpp

@@ -49,7 +49,7 @@ struct draw_particles_t {
 };
 
 #define PARTICLE_COUNT 1024
-#define SMOKE_COUNT 1024
+#define SMOKE_COUNT 256
 #define SMOKE_RESOLUTION 4
 #define RANDOM_DATA_LENGTH 1024
 
@@ -130,7 +130,10 @@ int main(int argc, const char **argv) {
 		generationShader.addShader(shaderStage, path);
 	});
 	
-	vkcv::DescriptorSetLayoutHandle generationDescriptorLayout = core.createDescriptorSetLayout(generationShader.getReflectedDescriptors().at(1));
+	vkcv::DescriptorSetLayoutHandle generationDescriptorLayout = core.createDescriptorSetLayout(
+		generationShader.getReflectedDescriptors().at(1)
+	);
+	
 	vkcv::DescriptorSetHandle generationDescriptorSet = core.createDescriptorSet(generationDescriptorLayout);
 	
 	vkcv::DescriptorBindings descriptorBindings1;
@@ -202,9 +205,29 @@ int main(int argc, const char **argv) {
 	auto smokeImagesBindings = smokeShaderProgram.getReflectedDescriptors().at(1);
 	smokeImagesBindings[1].descriptorCount = SMOKE_COUNT;
 	
+	smokeImagesBindings[0].shaderStages |= vkcv::ShaderStage::COMPUTE;
+	smokeImagesBindings[1].shaderStages |= vkcv::ShaderStage::COMPUTE;
+	
 	vkcv::DescriptorSetLayoutHandle smokeImagesDescriptorLayout = core.createDescriptorSetLayout(smokeImagesBindings);
 	vkcv::DescriptorSetHandle smokeImagesDescriptorSet = core.createDescriptorSet(smokeImagesDescriptorLayout);
 	
+	vkcv::ShaderProgram fluidShader;
+	compiler.compile(vkcv::ShaderStage::COMPUTE, "shaders/fluid.comp", [&](vkcv::ShaderStage shaderStage, const std::filesystem::path& path) {
+		fluidShader.addShader(shaderStage, path);
+	});
+	
+	vkcv::DescriptorSetLayoutHandle fluidImageDescriptorSetLayout = core.createDescriptorSetLayout(
+		fluidShader.getReflectedDescriptors().at(1)
+	);
+	
+	vkcv::ComputePipelineHandle fluidPipeline = core.createComputePipeline({
+		fluidShader,
+		{
+			smokeImagesDescriptorLayout,
+			fluidImageDescriptorSetLayout
+		}
+	});
+	
 	std::vector<particle_t> particles;
 	particles.reserve(PARTICLE_COUNT);
 	
@@ -342,7 +365,7 @@ int main(int argc, const char **argv) {
 	}
 	
 	std::vector<vkcv::ImageHandle> smokeImages;
-	smokes.reserve(smokes.size());
+	smokeImages.reserve(smokes.size());
 	
 	std::vector<uint64_t> imageData;
 	imageData.resize(SMOKE_RESOLUTION * SMOKE_RESOLUTION * SMOKE_RESOLUTION, 0);
@@ -362,6 +385,22 @@ int main(int argc, const char **argv) {
 		smokeImages.push_back(image.getHandle());
 	}
 	
+	std::vector<vkcv::ImageHandle> fluidImages;
+	fluidImages.reserve(smokeImages.size());
+	
+	for (size_t i = 0; i < smokeImages.size(); i++) {
+		vkcv::Image image = core.createImage(
+			vk::Format::eR16G16B16A16Sfloat,
+			SMOKE_RESOLUTION,
+			SMOKE_RESOLUTION,
+			SMOKE_RESOLUTION,
+			false,
+			true
+		);
+		
+		fluidImages.push_back(image.getHandle());
+	}
+	
 	vkcv::SamplerHandle smokeSampler = core.createSampler(
 		vkcv::SamplerFilterType::LINEAR,
 		vkcv::SamplerFilterType::LINEAR,
@@ -372,14 +411,18 @@ int main(int argc, const char **argv) {
 	{
 		vkcv::DescriptorWrites writes;
 		writes.writeSampler(0, smokeSampler);
-		
-		for (size_t i = 0; i < smokeImages.size(); i++) {
-			writes.writeSampledImage(1, smokeImages[i], 0, false, i);
-		}
-		
 		core.writeDescriptorSet(smokeImagesDescriptorSet, writes);
 	}
 	
+	std::vector<vkcv::DescriptorSetHandle> fluidImagesDescriptorSets;
+	fluidImagesDescriptorSets.reserve(smokeImages.size());
+	
+	for (size_t i = 0; i < smokeImages.size(); i++) {
+		fluidImagesDescriptorSets.push_back(core.createDescriptorSet(
+			fluidImageDescriptorSetLayout
+		));
+	}
+	
 	vkcv::Buffer<glm::vec3> cubePositions = core.createBuffer<glm::vec3>(vkcv::BufferType::VERTEX, 8);
 	cubePositions.fill({
 		glm::vec3(-1.0f, -1.0f, -1.0f),
@@ -595,6 +638,22 @@ int main(int argc, const char **argv) {
 		
 		current = next;
 		
+		{
+			vkcv::DescriptorWrites writes;
+		
+			for (size_t i = 0; i < smokeImages.size(); i++) {
+				writes.writeSampledImage(1, smokeImages[i], 0, false, i);
+			}
+		
+			core.writeDescriptorSet(smokeImagesDescriptorSet, writes);
+		}
+		
+		for (size_t i = 0; i < smokeImages.size(); i++) {
+			vkcv::DescriptorWrites writes;
+			writes.writeStorageImage(0, fluidImages[i]);
+			core.writeDescriptorSet(fluidImagesDescriptorSets[i], writes);
+		}
+		
 		float time_values [2];
 		time_values[0] = static_cast<float>(0.000001 * static_cast<double>(time.count()));
 		time_values[1] = static_cast<float>(0.000001 * static_cast<double>(deltatime.count()));
@@ -648,6 +707,36 @@ int main(int argc, const char **argv) {
 		);
 		core.recordEndDebugLabel(cmdStream);
 		
+		uint32_t fluidDispatchCount[3];
+		fluidDispatchCount[0] = std::ceil(SMOKE_RESOLUTION / 4.f);
+		fluidDispatchCount[1] = std::ceil(SMOKE_RESOLUTION / 4.f);
+		fluidDispatchCount[2] = std::ceil(SMOKE_RESOLUTION / 4.f);
+		
+		core.recordBeginDebugLabel(cmdStream, "Flow simulation", { 0.1f, 0.5f, 1.0f, 1.0f });
+		for (size_t i = 0; i < smokeImages.size(); i++) {
+			float fluid_values [2];
+			reinterpret_cast<uint*>(fluid_values)[0] = static_cast<uint>(i);
+			fluid_values[1] = time_values[1];
+			
+			vkcv::PushConstants pushConstantsFluid (2 * sizeof(float));
+			pushConstantsFluid.appendDrawcall(fluid_values);
+			
+			core.prepareImageForSampling(cmdStream, smokeImages[i]);
+			core.prepareImageForStorage(cmdStream, fluidImages[i]);
+			
+			core.recordComputeDispatchToCmdStream(
+				cmdStream,
+				fluidPipeline,
+				fluidDispatchCount,
+				{
+					vkcv::DescriptorSetUsage(0, smokeImagesDescriptorSet),
+					vkcv::DescriptorSetUsage(1, fluidImagesDescriptorSets[i])
+				},
+				pushConstantsFluid
+			);
+		}
+		core.recordEndDebugLabel(cmdStream);
+		
 		cameraManager.update(time_values[1]);
 		
 		const auto& camera = cameraManager.getActiveCamera();
@@ -736,6 +825,10 @@ int main(int argc, const char **argv) {
 		gui.endGUI();
 		
 		core.endFrame(windowHandle);
+		
+		for (size_t i = 0; i < smokeImages.size(); i++) {
+			std::swap(smokeImages[i], fluidImages[i]);
+		}
 	}
 	
 	return 0;