[rhi] MetalSurface functions

taichi/rhi/metal/CMakeLists.txt

@@ -18,4 +18,4 @@ target_include_directories(${METAL_RHI}
     ${PROJECT_SOURCE_DIR}/external/glad/include
     ${PROJECT_SOURCE_DIR}/external/glfw/include
   )
-target_link_libraries(${METAL_RHI} PRIVATE spirv-cross-msl spirv-cross-core)
+target_link_libraries(${METAL_RHI} PRIVATE spirv-cross-msl spirv-cross-core "-framework QuartzCore")

taichi/rhi/metal/metal_device.h

@@ -1,18 +1,21 @@
 #pragma once
-#include <memory>
 #include "taichi/common/logging.h"
 #include "taichi/rhi/device.h"
-#include "taichi/rhi/metal/metal_api.h"
 #include "taichi/rhi/impl_support.h"
+#include "taichi/rhi/metal/metal_api.h"
+#include <memory>
 
 #if defined(__APPLE__) && defined(__OBJC__)
+#import <CoreGraphics/CoreGraphics.h>
 #import <Foundation/Foundation.h>
 #import <Metal/Metal.h>
 #import <MetalKit/MetalKit.h>
-#import <CoreGraphics/CoreGraphics.h>
+#import <QuartzCore/QuartzCore.h>
 #define DEFINE_METAL_ID_TYPE(x) typedef id<x> x##_id;
+#define DEFINE_OBJC_TYPE(x) @class x;
 #else
 #define DEFINE_METAL_ID_TYPE(x) typedef struct x##_t *x##_id;
+#define DEFINE_OBJC_TYPE(x) typedef void x;
 #endif
 
 DEFINE_METAL_ID_TYPE(MTLDevice);
@@ -26,8 +29,11 @@ DEFINE_METAL_ID_TYPE(MTLCommandQueue);
 DEFINE_METAL_ID_TYPE(MTLCommandBuffer);
 DEFINE_METAL_ID_TYPE(MTLBlitCommandEncoder);
 DEFINE_METAL_ID_TYPE(MTLComputeCommandEncoder);
+DEFINE_METAL_ID_TYPE(CAMetalDrawable);
+DEFINE_OBJC_TYPE(CAMetalLayer);
 
 #undef DEFINE_METAL_ID_TYPE
+#undef DEFINE_OBJC_TYPE
 
 namespace taichi::lang {
 
@@ -41,7 +47,7 @@ class MetalStream;
 class MetalDevice;
 
 struct MetalMemory {
- public:
+public:
   // `mtl_buffer` should be already retained.
   explicit MetalMemory(MTLBuffer_id mtl_buffer, bool host_access);
   ~MetalMemory();
@@ -52,14 +58,14 @@ struct MetalMemory {
   size_t size() const;
   RhiResult mapped_ptr(void **mapped_ptr) const;
 
- private:
+private:
   MTLBuffer_id mtl_buffer_;
   bool can_map_{false};
   bool dont_destroy_{false};
 };
 
 struct MetalImage {
- public:
+public:
   // `mtl_texture` should be already retained.
   explicit MetalImage(MTLTexture_id mtl_texture);
   ~MetalImage();
@@ -68,20 +74,20 @@ struct MetalImage {
 
   MTLTexture_id mtl_texture() const;
 
- private:
+private:
   MTLTexture_id mtl_texture_;
   bool dont_destroy_{false};
 };
 
 struct MetalSampler {
- public:
+public:
   // `mtl_texture` should be already retained.
   explicit MetalSampler(MTLSamplerState_id mtl_sampler_state);
   ~MetalSampler();
 
   MTLSamplerState_id mtl_sampler_state() const;
 
- private:
+private:
   MTLSamplerState_id mtl_sampler_state_;
 };
 
@@ -91,19 +97,17 @@ struct MetalWorkgroupSize {
   uint32_t z{0};
 };
 class MetalPipeline final : public Pipeline {
- public:
+public:
   // `mtl_library`, `mtl_function`, `mtl_compute_pipeline_state` should be
   // already retained.
-  explicit MetalPipeline(const MetalDevice &device,
-                         MTLLibrary_id mtl_library,
+  explicit MetalPipeline(const MetalDevice &device, MTLLibrary_id mtl_library,
                          MTLFunction_id mtl_function,
                          MTLComputePipelineState_id mtl_compute_pipeline_state,
                          MetalWorkgroupSize workgroup_size);
   ~MetalPipeline() final;
 
   static MetalPipeline *create(const MetalDevice &device,
-                               const uint32_t *spv_data,
-                               size_t spv_size,
+                               const uint32_t *spv_data, size_t spv_size,
                                const std::string &name);
   void destroy();
 
@@ -114,7 +118,7 @@ class MetalPipeline final : public Pipeline {
     return workgroup_size_;
   }
 
- private:
+private:
   const MetalDevice *device_;
   MTLLibrary_id mtl_library_;
   MTLFunction_id mtl_function_;
@@ -145,37 +149,34 @@ struct MetalShaderResource {
   };
 };
 class MetalShaderResourceSet final : public ShaderResourceSet {
- public:
+public:
   explicit MetalShaderResourceSet(const MetalDevice &device);
   ~MetalShaderResourceSet() final;
 
-  ShaderResourceSet &rw_buffer(uint32_t binding,
-                               DevicePtr ptr,
+  ShaderResourceSet &rw_buffer(uint32_t binding, DevicePtr ptr,
                                size_t size) final;
   ShaderResourceSet &rw_buffer(uint32_t binding, DeviceAllocation alloc) final;
 
   ShaderResourceSet &buffer(uint32_t binding, DevicePtr ptr, size_t size) final;
   ShaderResourceSet &buffer(uint32_t binding, DeviceAllocation alloc) final;
 
-  ShaderResourceSet &image(uint32_t binding,
-                           DeviceAllocation alloc,
+  ShaderResourceSet &image(uint32_t binding, DeviceAllocation alloc,
                            ImageSamplerConfig sampler_config) override;
 
-  ShaderResourceSet &rw_image(uint32_t binding,
-                              DeviceAllocation alloc,
+  ShaderResourceSet &rw_image(uint32_t binding, DeviceAllocation alloc,
                               int lod) override;
 
   inline const std::vector<MetalShaderResource> &resources() const {
     return resources_;
   }
 
- private:
+private:
   const MetalDevice *device_;
-  std::vector<MetalShaderResource> resources_;
+  std::vector<MetalShaderResource> resources_; // TODO: need raster resources
 };
 
 class MetalCommandList final : public CommandList {
- public:
+public:
   explicit MetalCommandList(const MetalDevice &device,
                             MTLCommandQueue_id cmd_queue);
   ~MetalCommandList() final;
@@ -192,13 +193,12 @@ class MetalCommandList final : public CommandList {
   void buffer_fill(DevicePtr ptr, size_t size, uint32_t data) noexcept final;
   RhiResult dispatch(uint32_t x, uint32_t y = 1, uint32_t z = 1) noexcept final;
 
-  void image_transition(DeviceAllocation img,
-                        ImageLayout old_layout,
+  void image_transition(DeviceAllocation img, ImageLayout old_layout,
                         ImageLayout new_layout) final;
 
   MTLCommandBuffer_id finalize();
 
- private:
+private:
   friend class MetalStream;
 
   const MetalDevice *device_;
@@ -210,7 +210,7 @@ class MetalCommandList final : public CommandList {
 };
 
 class MetalStream final : public Stream {
- public:
+public:
   // `mtl_command_queue` should be already retained.
   explicit MetalStream(const MetalDevice &device,
                        MTLCommandQueue_id mtl_command_queue);
@@ -219,47 +219,78 @@ class MetalStream final : public Stream {
   static MetalStream *create(const MetalDevice &device);
   void destroy();
 
-  MTLCommandQueue_id mtl_command_queue() const {
-    return mtl_command_queue_;
-  }
+  MTLCommandQueue_id mtl_command_queue() const { return mtl_command_queue_; }
 
   RhiResult new_command_list(CommandList **out_cmdlist) noexcept final;
-  StreamSemaphore submit(
-      CommandList *cmdlist,
-      const std::vector<StreamSemaphore> &wait_semaphores = {}) final;
-  StreamSemaphore submit_synced(
-      CommandList *cmdlist,
-      const std::vector<StreamSemaphore> &wait_semaphores = {}) final;
+  StreamSemaphore
+  submit(CommandList *cmdlist,
+         const std::vector<StreamSemaphore> &wait_semaphores = {}) final;
+  StreamSemaphore
+  submit_synced(CommandList *cmdlist,
+                const std::vector<StreamSemaphore> &wait_semaphores = {}) final;
 
   void command_sync() override;
 
- private:
+private:
   const MetalDevice *device_;
   MTLCommandQueue_id mtl_command_queue_;
   std::vector<MTLCommandBuffer_id> pending_cmdbufs_;
   bool is_destroyed_{false};
 };
 
+class MetalSurface final : public Surface {
+public:
+  MetalSurface(MetalDevice *device, const SurfaceConfig &config);
+  ~MetalSurface() override;
+
+  CAMetalLayer *mtl_layer() { return layer_; }
+
+  StreamSemaphore acquire_next_image() override;
+  DeviceAllocation get_target_image() override;
+
+  void present_image(
+      const std::vector<StreamSemaphore> &wait_semaphores = {}) override;
+  std::pair<uint32_t, uint32_t> get_size() override;
+  int get_image_count() override;
+  BufferFormat image_format() override;
+  void resize(uint32_t width, uint32_t height) override;
+
+  DeviceAllocation get_depth_data(DeviceAllocation &depth_alloc) override {
+    TI_NOT_IMPLEMENTED;
+  }
+  DeviceAllocation get_image_data() override { TI_NOT_IMPLEMENTED; }
+
+private:
+  void destroy_swap_chain();
+
+  SurfaceConfig config_;
+
+  BufferFormat image_format_{BufferFormat::unknown};
+
+  uint32_t width_{0};
+  uint32_t height_{0};
+
+  MTLTexture_id current_swap_chain_texture_;
+  std::unordered_map<MTLTexture_id, DeviceAllocation> swapchain_images_;
+  CAMetalDrawable_id current_drawable_;
+
+  MetalDevice *device_{nullptr};
+  CAMetalLayer *layer_;
+};
+
 class MetalDevice final : public GraphicsDevice {
- public:
+public:
   // `mtl_device` should be already retained.
   explicit MetalDevice(MTLDevice_id mtl_device);
   ~MetalDevice() override;
 
-  Arch arch() const override {
-    return Arch::metal;
-  }
-  MTLDevice_id mtl_device() const {
-    return mtl_device_;
-  }
+  Arch arch() const override { return Arch::metal; }
+  MTLDevice_id mtl_device() const { return mtl_device_; }
 
   static MetalDevice *create();
   void destroy();
 
-  std::unique_ptr<Surface> create_surface(
-      const SurfaceConfig &config) override {
-    TI_NOT_IMPLEMENTED;
-  }
+  std::unique_ptr<Surface> create_surface(const SurfaceConfig &config) override;
 
   RhiResult allocate_memory(const AllocParams &params,
                             DeviceAllocation *out_devalloc) override;
@@ -282,34 +313,29 @@ class MetalDevice final : public GraphicsDevice {
   void unmap(DeviceAllocation ptr) override;
 
   RhiResult create_pipeline(Pipeline **out_pipeline,
-                            const PipelineSourceDesc &src,
-                            std::string name,
+                            const PipelineSourceDesc &src, std::string name,
                             PipelineCache *cache) noexcept final;
   ShaderResourceSet *create_resource_set() override;
 
-  std::unique_ptr<Pipeline> create_raster_pipeline(
-      const std::vector<PipelineSourceDesc> &src,
-      const RasterParams &raster_params,
-      const std::vector<VertexInputBinding> &vertex_inputs,
-      const std::vector<VertexInputAttribute> &vertex_attrs,
-      std::string name = "Pipeline") override {
-    TI_NOT_IMPLEMENTED;
-  }
-  RasterResources *create_raster_resources() override {
+  std::unique_ptr<Pipeline>
+  create_raster_pipeline(const std::vector<PipelineSourceDesc> &src,
+                         const RasterParams &raster_params,
+                         const std::vector<VertexInputBinding> &vertex_inputs,
+                         const std::vector<VertexInputAttribute> &vertex_attrs,
+                         std::string name = "Pipeline") override {
     TI_NOT_IMPLEMENTED;
   }
+  RasterResources *create_raster_resources() override { TI_NOT_IMPLEMENTED; }
 
   Stream *get_compute_stream() override;
   Stream *get_graphics_stream() override;
   void wait_idle() override;
 
   void memcpy_internal(DevicePtr dst, DevicePtr src, uint64_t size) override;
 
-  const MetalSampler &get_default_sampler() const {
-    return *default_sampler_;
-  }
+  const MetalSampler &get_default_sampler() const { return *default_sampler_; }
 
- private:
+private:
   MTLDevice_id mtl_device_;
   rhi_impl::SyncedPtrStableObjectList<MetalMemory> memory_allocs_;
   rhi_impl::SyncedPtrStableObjectList<MetalImage> image_allocs_;
@@ -320,5 +346,5 @@ class MetalDevice final : public GraphicsDevice {
   bool is_destroyed_{false};
 };
 
-}  // namespace metal
-}  // namespace taichi::lang
+} // namespace metal
+} // namespace taichi::lang