Metal: Stable argument buffers; GPU rendering crashes; visionOS exports

Supersedes #110683

servers/rendering/rendering_shader_container.h

@@ -34,6 +34,8 @@
 #include "servers/rendering/rendering_device_commons.h"
 
 struct SpvReflectShaderModule;
+struct SpvReflectDescriptorBinding;
+struct SpvReflectSpecializationConstant;
 
 class RenderingShaderContainer : public RefCounted {
 	GDSOFTCLASS(RenderingShaderContainer, RefCounted);
@@ -43,6 +45,8 @@ public:
 	static const uint32_t CONTAINER_VERSION = 2;
 
 protected:
+	using RDC = RenderingDeviceCommons;
+
 	struct ContainerHeader {
 		uint32_t magic_number = 0;
 		uint32_t version = 0;
@@ -96,7 +100,7 @@ protected:
 	Vector<uint32_t> reflection_binding_set_uniforms_count;
 	Vector<ReflectionBindingData> reflection_binding_set_uniforms_data;
 	Vector<ReflectionSpecializationData> reflection_specialization_data;
-	Vector<RenderingDeviceCommons::ShaderStage> reflection_shader_stages;
+	Vector<RDC::ShaderStage> reflection_shader_stages;
 
 	virtual uint32_t _format() const = 0;
 	virtual uint32_t _format_version() const = 0;
@@ -120,30 +124,159 @@ protected:
 	virtual uint32_t _to_bytes_shader_extra_data(uint8_t *p_bytes, uint32_t p_index) const;
 	virtual uint32_t _to_bytes_footer_extra_data(uint8_t *p_bytes) const;
 
-	// This method will be called when set_from_shader_reflection() is finished. Used to update internal structures to match the reflection if necessary.
-	virtual void _set_from_shader_reflection_post(const RenderingDeviceCommons::ShaderReflection &p_reflection);
+	template <class T>
+	struct ReflectSymbol {
+		static constexpr uint32_t STAGE_INDEX[RDC::SHADER_STAGE_MAX] = {
+			0, // SHADER_STAGE_VERTEX
+			1, // SHADER_STAGE_FRAGMENT
+			0, // SHADER_STAGE_TESSELATION_CONTROL
+			1, // SHADER_STAGE_TESSELATION_EVALUATION
+			0, // SHADER_STAGE_COMPUTE
+		};
 
-	class ReflectedShaderStage {
+		BitField<RDC::ShaderStage> stages = {};
+
+	private:
+		const T *_spv_reflect[2] = { nullptr };
+
+	public:
+		_FORCE_INLINE_ constexpr uint32_t get_index_for_stage(RDC::ShaderStage p_stage) const {
+			DEV_ASSERT(stages.has_flag((1 << p_stage)));
+			return STAGE_INDEX[p_stage];
+		}
+
+		const T &get_spv_reflect(RDC::ShaderStage p_stage) const;
+
+		/*! Returns the first valid stage if multiple stages are set.
+		 *
+		 * Crashes if no stages are set.
+		 */
+		const T &get_spv_reflect() const {
+			for (const T *d : _spv_reflect) {
+				if (d != nullptr) {
+					return *d;
+				}
+			}
+			CRASH_NOW_MSG("No stages set in ReflectSymbol");
+		}
+		void set_spv_reflect(RDC::ShaderStage p_stage, const T *p_spv);
+	};
+
+	struct ReflectImageTraits {
+		RDC::DataFormat format = RDC::DATA_FORMAT_MAX;
+	};
+
+	struct ReflectUniform : ReflectSymbol<SpvReflectDescriptorBinding> {
+		RDC::UniformType type = RDC::UniformType::UNIFORM_TYPE_MAX;
+		uint32_t binding = 0;
+
+		ReflectImageTraits image;
+
+		uint32_t length = 0; // Size of arrays (in total elements), or ubos (in bytes * total elements).
+		bool writable = false;
+
+		bool operator<(const ReflectUniform &p_other) const {
+			if (binding != p_other.binding) {
+				return binding < p_other.binding;
+			}
+			if (type != p_other.type) {
+				return type < p_other.type;
+			}
+			if (writable != p_other.writable) {
+				return writable < p_other.writable;
+			}
+			if (stages != p_other.stages) {
+				return stages < p_other.stages;
+			}
+			if (length != p_other.length) {
+				return length < p_other.length;
+			}
+			return false;
+		}
+	};
+
+	struct ReflectSpecializationConstant : ReflectSymbol<SpvReflectSpecializationConstant> {
+		RDC::PipelineSpecializationConstantType type = {};
+		uint32_t constant_id = 0xffffffff;
+		union {
+			uint32_t int_value = 0;
+			float float_value;
+			bool bool_value;
+		};
+
+		bool operator<(const ReflectSpecializationConstant &p_other) const { return constant_id < p_other.constant_id; }
+	};
+
+	class ReflectShaderStage {
 		friend class RenderingShaderContainer;
 
 		Vector<uint8_t> _spirv_data;
 		SpvReflectShaderModule *_module = nullptr;
 
 	public:
-		RenderingDeviceCommons::ShaderStage shader_stage = RenderingDeviceCommons::SHADER_STAGE_MAX;
+		RDC::ShaderStage shader_stage = RDC::SHADER_STAGE_MAX;
 		const SpvReflectShaderModule &module() const;
 		const Span<uint32_t> spirv() const;
 		const Vector<uint8_t> spirv_data() const { return _spirv_data; }
 
-		ReflectedShaderStage();
-		~ReflectedShaderStage();
+		ReflectShaderStage();
+		~ReflectShaderStage();
 	};
 
-	// This method will be called when set_code_from_spirv() is called.
-	virtual bool _set_code_from_spirv(Span<ReflectedShaderStage> p_spirv) = 0;
+	typedef LocalVector<ReflectUniform> ReflectDescriptorSet;
 
-	void set_from_shader_reflection(const RenderingDeviceCommons::ShaderReflection &p_reflection);
-	Error reflect_spirv(const String &p_shader_name, Span<RenderingDeviceCommons::ShaderStageSPIRVData> p_spirv, LocalVector<ReflectedShaderStage> &r_refl);
+	struct ReflectShader {
+		uint64_t vertex_input_mask = 0;
+		uint32_t fragment_output_mask = 0;
+		uint32_t compute_local_size[3] = {};
+		uint32_t push_constant_size = 0;
+		bool has_multiview = false;
+		bool has_dynamic_buffers = false;
+
+		LocalVector<ReflectShaderStage> shader_stages;
+		LocalVector<ReflectDescriptorSet> uniform_sets;
+		LocalVector<ReflectSymbol<SpvReflectDescriptorBinding>> reflect_uniforms;
+		LocalVector<ReflectSpecializationConstant> specialization_constants;
+		LocalVector<ReflectSymbol<SpvReflectSpecializationConstant>> reflect_specialization_constants;
+		LocalVector<RDC::ShaderStage> stages_vector;
+		BitField<RDC::ShaderStage> stages_bits = {};
+		BitField<RDC::ShaderStage> push_constant_stages = {};
+
+		_FORCE_INLINE_ bool is_compute() const {
+			return stages_bits.has_flag(RDC::SHADER_STAGE_COMPUTE_BIT);
+		}
+
+		/*! Returns the uniform at the specified global index.
+		 *
+		 * This is a flattened view of all uniform sets.
+		 */
+		ReflectUniform &uniform_at(uint32_t p_index) {
+			for (LocalVector<ReflectUniform> &set : uniform_sets) {
+				if (p_index < set.size()) {
+					return set[p_index];
+				}
+				p_index -= set.size();
+			}
+			CRASH_NOW_MSG(vformat("Uniform index %d out of range (total %d)", p_index, uniform_count()));
+		}
+
+		uint32_t uniform_count() const {
+			uint32_t count = 0;
+			for (const LocalVector<ReflectUniform> &set : uniform_sets) {
+				count += set.size();
+			}
+			return count;
+		}
+	};
+
+	// This method will be called when set_from_shader_reflection() is finished. Used to update internal structures to match the reflection if necessary.
+	virtual void _set_from_shader_reflection_post(const ReflectShader &p_shader);
+
+	// This method will be called when set_code_from_spirv() is called.
+	virtual bool _set_code_from_spirv(const ReflectShader &p_shader) = 0;
+
+	void set_from_shader_reflection(const ReflectShader &p_reflection);
+	Error reflect_spirv(const String &p_shader_name, Span<RDC::ShaderStageSPIRVData> p_spirv, ReflectShader &r_shader);
 
 public:
 	enum CompressionFlags {
@@ -151,7 +284,7 @@ public:
 	};
 
 	struct Shader {
-		RenderingDeviceCommons::ShaderStage shader_stage = RenderingDeviceCommons::SHADER_STAGE_MAX;
+		RDC::ShaderStage shader_stage = RDC::SHADER_STAGE_MAX;
 		PackedByteArray code_compressed_bytes;
 		uint32_t code_compression_flags = 0;
 		uint32_t code_decompressed_size = 0;
@@ -160,8 +293,8 @@ public:
 	CharString shader_name;
 	Vector<Shader> shaders;
 
-	bool set_code_from_spirv(const String &p_shader_name, Span<RenderingDeviceCommons::ShaderStageSPIRVData> p_spirv);
-	RenderingDeviceCommons::ShaderReflection get_shader_reflection() const;
+	bool set_code_from_spirv(const String &p_shader_name, Span<RDC::ShaderStageSPIRVData> p_spirv);
+	RDC::ShaderReflection get_shader_reflection() const;
 	bool from_bytes(const PackedByteArray &p_bytes);
 	PackedByteArray to_bytes() const;
 	bool compress_code(const uint8_t *p_decompressed_bytes, uint32_t p_decompressed_size, uint8_t *p_compressed_bytes, uint32_t *r_compressed_size, uint32_t *r_compressed_flags) const;