Allow shader containers to be loaded in parallel.

Author: blueskythlikesclouds

core/io/compression.cpp

@@ -40,11 +40,34 @@
 #include <brotli/decode.h>
 #endif
 
-// Caches for zstd.
-static BinaryMutex mutex;
-static ZSTD_DCtx *current_zstd_d_ctx = nullptr;
-static bool current_zstd_long_distance_matching;
-static int current_zstd_window_log_size;
+namespace {
+struct ZstdDecompressorContext {
+	ZSTD_DCtx *zstd_d_ctx = nullptr;
+	bool zstd_long_distance_matching = false;
+	int zstd_window_log_size = 0;
+
+	~ZstdDecompressorContext() {
+		if (zstd_d_ctx) {
+			ZSTD_freeDCtx(zstd_d_ctx);
+		}
+	}
+
+	void invalidate(bool p_zstd_long_distance_matching, int p_zstd_window_log_size) {
+		if (!zstd_d_ctx || zstd_long_distance_matching != p_zstd_long_distance_matching || zstd_window_log_size != p_zstd_window_log_size) {
+			if (zstd_d_ctx) {
+				ZSTD_freeDCtx(zstd_d_ctx);
+			}
+
+			zstd_d_ctx = ZSTD_createDCtx();
+			if (p_zstd_long_distance_matching) {
+				ZSTD_DCtx_setParameter(zstd_d_ctx, ZSTD_d_windowLogMax, p_zstd_window_log_size);
+			}
+			zstd_long_distance_matching = p_zstd_long_distance_matching;
+			zstd_window_log_size = p_zstd_window_log_size;
+		}
+	}
+};
+} //namespace
 
 int64_t Compression::compress(uint8_t *p_dst, const uint8_t *p_src, int64_t p_src_size, Mode p_mode) {
 	switch (p_mode) {
@@ -197,22 +220,10 @@ int64_t Compression::decompress(uint8_t *p_dst, int64_t p_dst_max_size, const ui
 			return total;
 		} break;
 		case MODE_ZSTD: {
-			MutexLock lock(mutex);
-
-			if (!current_zstd_d_ctx || current_zstd_long_distance_matching != zstd_long_distance_matching || current_zstd_window_log_size != zstd_window_log_size) {
-				if (current_zstd_d_ctx) {
-					ZSTD_freeDCtx(current_zstd_d_ctx);
-				}
-
-				current_zstd_d_ctx = ZSTD_createDCtx();
-				if (zstd_long_distance_matching) {
-					ZSTD_DCtx_setParameter(current_zstd_d_ctx, ZSTD_d_windowLogMax, zstd_window_log_size);
-				}
-				current_zstd_long_distance_matching = zstd_long_distance_matching;
-				current_zstd_window_log_size = zstd_window_log_size;
-			}
+			thread_local ZstdDecompressorContext decompressor_ctx;
+			decompressor_ctx.invalidate(zstd_long_distance_matching, zstd_window_log_size);
 
-			size_t ret = ZSTD_decompressDCtx(current_zstd_d_ctx, p_dst, p_dst_max_size, p_src, p_src_size);
+			size_t ret = ZSTD_decompressDCtx(decompressor_ctx.zstd_d_ctx, p_dst, p_dst_max_size, p_src, p_src_size);
 			return (int64_t)ret;
 		} break;
 	}

servers/rendering/renderer_rd/shader_rd.cpp

@@ -258,6 +258,7 @@ void ShaderRD::_initialize_version(Version *p_version) {
 	p_version->variants.resize_initialized(variant_defines.size());
 	p_version->variant_data.resize(variant_defines.size());
 	p_version->group_compilation_tasks.resize_initialized(group_enabled.size());
+	p_version->group_loaded_from_cache.resize_initialized(group_enabled.size());
 }
 
 void ShaderRD::_clear_version(Version *p_version) {
@@ -611,6 +612,15 @@ String ShaderRD::_get_cache_file_path(Version *p_version, int p_group, const Str
 	return shader_cache_dir.path_join(relative_path);
 }
 
+void ShaderRD::_load_variant_from_cache(uint32_t p_variant, CompileData p_data) {
+	uint32_t variant = group_to_variant_map[p_data.group][p_variant];
+	if (!variants_enabled[variant]) {
+		return; // Variant is disabled, return.
+	}
+
+	p_data.version->variants.write[variant] = RD::get_singleton()->shader_create_from_bytecode_with_samplers(p_data.version->variant_data[variant], p_data.version->variants[variant], immutable_samplers);
+}
+
 bool ShaderRD::_load_from_cache(Version *p_version, int p_group) {
 	String api_safe_name = String(RD::get_singleton()->get_device_api_name()).validate_filename().to_lower();
 	Ref<FileAccess> f;
@@ -662,28 +672,14 @@ bool ShaderRD::_load_from_cache(Version *p_version, int p_group) {
 		p_version->variant_data.write[variant_id] = variant_bytes;
 	}
 
-	for (uint32_t i = 0; i < variant_count; i++) {
-		int variant_id = group_to_variant_map[p_group][i];
-		if (!variants_enabled[variant_id]) {
-			p_version->variants.write[variant_id] = RID();
-			continue;
-		}
-		print_verbose(vformat("Loading cache for shader %s, variant %d", name, i));
-		{
-			RID shader = RD::get_singleton()->shader_create_from_bytecode_with_samplers(p_version->variant_data[variant_id], p_version->variants[variant_id], immutable_samplers);
-			if (shader.is_null()) {
-				for (uint32_t j = 0; j < i; j++) {
-					int variant_free_id = group_to_variant_map[p_group][j];
-					RD::get_singleton()->free_rid(p_version->variants[variant_free_id]);
-				}
-				ERR_FAIL_COND_V(shader.is_null(), false);
-			}
+	CompileData compile_data;
+	compile_data.version = p_version;
+	compile_data.group = p_group;
 
-			p_version->variants.write[variant_id] = shader;
-		}
-	}
+	WorkerThreadPool::GroupID group_task = WorkerThreadPool::get_singleton()->add_template_group_task(this, &ShaderRD::_load_variant_from_cache, compile_data, variant_count, -1, true, SNAME("LoadVariantFromCache"));
+	p_version->group_compilation_tasks.write[p_group] = group_task;
+	p_version->group_loaded_from_cache.write[p_group] = true;
 
-	p_version->valid = true;
 	return true;
 }
 
@@ -733,6 +729,7 @@ void ShaderRD::_compile_version_start(Version *p_version, int p_group) {
 
 	WorkerThreadPool::GroupID group_task = WorkerThreadPool::get_singleton()->add_template_group_task(this, &ShaderRD::_compile_variant, compile_data, group_to_variant_map[p_group].size(), -1, true, SNAME("ShaderCompilation"));
 	p_version->group_compilation_tasks.write[p_group] = group_task;
+	p_version->group_loaded_from_cache.write[p_group] = false;
 }
 
 void ShaderRD::_compile_version_end(Version *p_version, int p_group) {
@@ -772,7 +769,7 @@ void ShaderRD::_compile_version_end(Version *p_version, int p_group) {
 		return;
 	}
 #if ENABLE_SHADER_CACHE
-	else if (shader_cache_user_dir_valid) {
+	else if (shader_cache_user_dir_valid && !p_version->group_loaded_from_cache[p_group]) {
 		_save_to_cache(p_version, p_group);
 	}
 #endif

servers/rendering/renderer_rd/shader_rd.h

@@ -82,6 +82,7 @@ class ShaderRD {
 		HashMap<StringName, CharString> code_sections;
 		Vector<CharString> custom_defines;
 		Vector<WorkerThreadPool::GroupID> group_compilation_tasks;
+		Vector<bool> group_loaded_from_cache;
 
 		Vector<Vector<uint8_t>> variant_data;
 		Vector<RID> variants;
@@ -178,6 +179,7 @@ class ShaderRD {
 	String _version_get_sha1(Version *p_version) const;
 	String _get_cache_file_relative_path(Version *p_version, int p_group, const String &p_api_name);
 	String _get_cache_file_path(Version *p_version, int p_group, const String &p_api_name, bool p_user_dir);
+	void _load_variant_from_cache(uint32_t p_variant, CompileData p_data);
 	bool _load_from_cache(Version *p_version, int p_group);
 	void _save_to_cache(Version *p_version, int p_group);
 	void _initialize_cache();

servers/rendering/rendering_device.cpp

@@ -4170,8 +4170,6 @@ RID RenderingDevice::shader_create_from_bytecode(const Vector<uint8_t> &p_shader
 }
 
 RID RenderingDevice::shader_create_from_bytecode_with_samplers(const Vector<uint8_t> &p_shader_binary, RID p_placeholder, const Vector<PipelineImmutableSampler> &p_immutable_samplers) {
-	_THREAD_SAFE_METHOD_
-
 	Ref<RenderingShaderContainer> shader_container = driver->get_shader_container_format().create_container();
 	ERR_FAIL_COND_V(shader_container.is_null(), RID());
 
@@ -4195,6 +4193,8 @@ RID RenderingDevice::shader_create_from_bytecode_with_samplers(const Vector<uint
 	RDD::ShaderID shader_id = driver->shader_create_from_container(shader_container, driver_immutable_samplers);
 	ERR_FAIL_COND_V(!shader_id, RID());
 
+	_THREAD_SAFE_METHOD_
+
 	// All good, let's create modules.
 
 	RID id;