1 files changed, 122 insertions, 0 deletions
diff --git a/source/blender/gpu/metal/mtl_query.mm b/source/blender/gpu/metal/mtl_query.mm
new file mode 100644
index 00000000000..33ae8f554c3
--- /dev/null
+++ b/source/blender/gpu/metal/mtl_query.mm
@@ -0,0 +1,122 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#include "mtl_query.hh"
+
+namespace blender::gpu {
+
+static const size_t VISIBILITY_COUNT_PER_BUFFER = 512;
+/* defined in the documentation but not queryable programmatically:
+ * https://developer.apple.com/documentation/metal/mtlvisibilityresultmode/mtlvisibilityresultmodeboolean?language=objc
+ */
+static const size_t VISIBILITY_RESULT_SIZE_IN_BYTES = 8;
+
+MTLQueryPool::MTLQueryPool()
+{
+  allocate_buffer();
+}
+MTLQueryPool::~MTLQueryPool()
+{
+  for (gpu::MTLBuffer *buf : buffer_) {
+    BLI_assert(buf);
+    buf->free();
+  }
+}
+
+void MTLQueryPool::allocate_buffer()
+{
+  /* Allocate Metal buffer for visibility results. */
+  size_t buffer_size_in_bytes = VISIBILITY_COUNT_PER_BUFFER * VISIBILITY_RESULT_SIZE_IN_BYTES;
+  gpu::MTLBuffer *buffer = MTLContext::get_global_memory_manager().allocate_buffer(
+      buffer_size_in_bytes, true);
+  BLI_assert(buffer);
+  buffer_.append(buffer);
+}
+
+static inline MTLVisibilityResultMode to_mtl_type(GPUQueryType type)
+{
+  if (type == GPU_QUERY_OCCLUSION) {
+    return MTLVisibilityResultModeBoolean;
+  }
+  BLI_assert(0);
+  return MTLVisibilityResultModeBoolean;
+}
+
+void MTLQueryPool::init(GPUQueryType type)
+{
+  BLI_assert(initialized_ == false);
+  initialized_ = true;
+  type_ = type;
+  mtl_type_ = to_mtl_type(type);
+  query_issued_ = 0;
+}
+
+void MTLQueryPool::begin_query()
+{
+  MTLContext *ctx = reinterpret_cast<MTLContext *>(GPU_context_active_get());
+
+  /* Ensure our allocated buffer pool has enough space for the current queries. */
+  int query_id = query_issued_;
+  int requested_buffer = query_id / VISIBILITY_COUNT_PER_BUFFER;
+  if (requested_buffer >= buffer_.size()) {
+    allocate_buffer();
+  }
+
+  BLI_assert(requested_buffer < buffer_.size());
+  gpu::MTLBuffer *buffer = buffer_[requested_buffer];
+
+  /* Ensure visibility buffer is set on the context. If visibility buffer changes,
+   * we need to begin a new render pass with an updated reference in the
+   * MTLRenderPassDescriptor. */
+  ctx->set_visibility_buffer(buffer);
+
+  ctx->ensure_begin_render_pass();
+  id<MTLRenderCommandEncoder> rec = ctx->main_command_buffer.get_active_render_command_encoder();
+  [rec setVisibilityResultMode:mtl_type_
+                        offset:(query_id % VISIBILITY_COUNT_PER_BUFFER) *
+                               VISIBILITY_RESULT_SIZE_IN_BYTES];
+  query_issued_ += 1;
+}
+
+void MTLQueryPool::end_query()
+{
+  MTLContext *ctx = reinterpret_cast<MTLContext *>(GPU_context_active_get());
+
+  id<MTLRenderCommandEncoder> rec = ctx->main_command_buffer.get_active_render_command_encoder();
+  [rec setVisibilityResultMode:MTLVisibilityResultModeDisabled offset:0];
+}
+
+void MTLQueryPool::get_occlusion_result(MutableSpan<uint32_t> r_values)
+{
+  MTLContext *ctx = reinterpret_cast<MTLContext *>(GPU_context_active_get());
+
+  /* Create a blit encoder to synchronize the query buffer results between
+   * GPU and CPU when not using shared-memory. */
+  if ([ctx->device hasUnifiedMemory] == false) {
+    id<MTLBlitCommandEncoder> blit_encoder = ctx->main_command_buffer.ensure_begin_blit_encoder();
+    BLI_assert(blit_encoder);
+    for (gpu::MTLBuffer *buf : buffer_) {
+      [blit_encoder synchronizeResource:buf->get_metal_buffer()];
+    }
+    BLI_assert(ctx->get_inside_frame());
+  }
+
+  /* Wait for GPU operatiosn to complete and for query buffer contents
+   * to be synchronised back to host memory. */
+  GPU_finish();
+
+  /* Iterate through all possible visibility buffers and copy results into provided
+   * container. */
+  for (const int i : IndexRange(query_issued_)) {
+    int requested_buffer = i / VISIBILITY_COUNT_PER_BUFFER;
+    const uint64_t *queries = static_cast<const uint64_t *>(
+        buffer_[requested_buffer]->get_host_ptr());
+    r_values[i] = static_cast<uint32_t>(queries[i % VISIBILITY_COUNT_PER_BUFFER]);
+  }
+  ctx->set_visibility_buffer(nullptr);
+}
+
+}  // namespace blender::gpu