Metal: Performance improvements and bug fixes

1 files changed, 38 insertions, 15 deletions

diff --git a/drivers/metal/metal_objects.h b/drivers/metal/metal_objects.h
index 030b353ee8..38d5b53ffa 100644
--- a/drivers/metal/metal_objects.h
+++ b/drivers/metal/metal_objects.h
@@ -96,6 +96,22 @@ _FORCE_INLINE_ ShaderStageUsage &operator|=(ShaderStageUsage &p_a, int p_b) {
 	return p_a;
 }
 
+enum StageResourceUsage : uint32_t {
+	VertexRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_VERTEX * 2),
+	VertexWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_VERTEX * 2),
+	FragmentRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_FRAGMENT * 2),
+	FragmentWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_FRAGMENT * 2),
+	TesselationControlRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_TESSELATION_CONTROL * 2),
+	TesselationControlWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_TESSELATION_CONTROL * 2),
+	TesselationEvaluationRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_TESSELATION_EVALUATION * 2),
+	TesselationEvaluationWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_TESSELATION_EVALUATION * 2),
+	ComputeRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_COMPUTE * 2),
+	ComputeWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_COMPUTE * 2),
+};
+
+typedef LocalVector<__unsafe_unretained id<MTLResource>> ResourceVector;
+typedef HashMap<StageResourceUsage, ResourceVector> ResourceUsageMap;
+
 enum class MDCommandBufferStateType {
 	None,
 	Render,
@@ -230,6 +246,7 @@ public:
 		uint32_t index_offset = 0;
 		LocalVector<id<MTLBuffer> __unsafe_unretained> vertex_buffers;
 		LocalVector<NSUInteger> vertex_offsets;
+		ResourceUsageMap resource_usage;
 		// clang-format off
 		enum DirtyFlag: uint8_t {
 			DIRTY_NONE     = 0b0000'0000,
@@ -271,8 +288,14 @@ public:
 			blend_constants.reset();
 			vertex_buffers.clear();
 			vertex_offsets.clear();
+			// Keep the keys, as they are likely to be used again.
+			for (KeyValue<StageResourceUsage, LocalVector<__unsafe_unretained id<MTLResource>>> &kv : resource_usage) {
+				kv.value.clear();
+			}
 		}
 
+		void end_encoding();
+
 		_FORCE_INLINE_ void mark_viewport_dirty() {
 			if (viewports.is_empty()) {
 				return;
@@ -356,13 +379,20 @@ public:
 	} render;
 
 	// State specific for a compute pass.
-	struct {
+	struct ComputeState {
 		MDComputePipeline *pipeline = nullptr;
 		id<MTLComputeCommandEncoder> encoder = nil;
+		ResourceUsageMap resource_usage;
 		_FORCE_INLINE_ void reset() {
 			pipeline = nil;
 			encoder = nil;
+			// Keep the keys, as they are likely to be used again.
+			for (KeyValue<StageResourceUsage, LocalVector<__unsafe_unretained id<MTLResource>>> &kv : resource_usage) {
+				kv.value.clear();
+			}
 		}
+
+		void end_encoding();
 	} compute;
 
 	// State specific to a blit pass.
@@ -632,19 +662,6 @@ public:
 	MDRenderShader(CharString p_name, Vector<UniformSet> p_sets, MDLibrary *p_vert, MDLibrary *p_frag);
 };
 
-enum StageResourceUsage : uint32_t {
-	VertexRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_VERTEX * 2),
-	VertexWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_VERTEX * 2),
-	FragmentRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_FRAGMENT * 2),
-	FragmentWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_FRAGMENT * 2),
-	TesselationControlRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_TESSELATION_CONTROL * 2),
-	TesselationControlWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_TESSELATION_CONTROL * 2),
-	TesselationEvaluationRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_TESSELATION_EVALUATION * 2),
-	TesselationEvaluationWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_TESSELATION_EVALUATION * 2),
-	ComputeRead = (MTLResourceUsageRead << RDD::SHADER_STAGE_COMPUTE * 2),
-	ComputeWrite = (MTLResourceUsageWrite << RDD::SHADER_STAGE_COMPUTE * 2),
-};
-
 _FORCE_INLINE_ StageResourceUsage &operator|=(StageResourceUsage &p_a, uint32_t p_b) {
 	p_a = StageResourceUsage(uint32_t(p_a) | p_b);
 	return p_a;
@@ -667,7 +684,13 @@ struct HashMapComparatorDefault<RDD::ShaderID> {
 
 struct BoundUniformSet {
 	id<MTLBuffer> buffer;
-	HashMap<id<MTLResource>, StageResourceUsage> bound_resources;
+	ResourceUsageMap usage_to_resources;
+
+	/// Perform a 2-way merge each key of `ResourceVector` resources from this set into the
+	/// destination set.
+	///
+	/// Assumes the vectors of resources are sorted.
+	void merge_into(ResourceUsageMap &p_dst) const;
 };
 
 class API_AVAILABLE(macos(11.0), ios(14.0)) MDUniformSet {