commit 0adc4619c3bb8079d7c0d73bb570d044a5de3f9d
author Lei Zhang <antiagainst@google.com> Sat Mar 11 21:09:54 2023 -0800
committer Lei Zhang <antiagainst@google.com> Tue Jun 13 21:17:31 2023 -0700
tree 2db2ab35ca10bb4320747bd0ea5d3c2ca8cea591
parent 3f64a119bf4d3705e7d51adc29218cc38aa64dbe

[metal] Use MTLFence to synchronize encoders in command buffer

This is needed for `MTLHazardTrackingModeTracked` mode per
Metal documentation:

> An MTLFence synchronizes access to one or more resources across
> different passes within a command buffer. Use fences to specify
> any inter-pass resource dependencies within the same command
> buffer.

experimental/metal/direct_command_buffer.m

diff --git a/experimental/metal/direct_command_buffer.m b/experimental/metal/direct_command_buffer.m
index bfb6081..3a3c202 100644
--- a/experimental/metal/direct_command_buffer.m
+++ b/experimental/metal/direct_command_buffer.m
@@ -106,7 +106,18 @@
 
 static id<MTLComputeCommandEncoder> iree_hal_metal_get_or_begin_compute_encoder(
     iree_hal_metal_command_buffer_t* command_buffer) {
-  if (command_buffer->blit_encoder) iree_hal_metal_end_blit_encoder(command_buffer);
+  id<MTLFence> encoder_fence = nil;
+  if (command_buffer->blit_encoder) {
+    // We would need to use a fence to synchronize "one or more resources across different passes
+    // within a command buffer."
+    // https://developer.apple.com/documentation/metal/resource_synchronization
+    encoder_fence = [command_buffer->command_buffer.device newFence];  // +1
+    [command_buffer->command_buffer addCompletedHandler:^(id<MTLCommandBuffer> cb) {
+      [encoder_fence release];  // -1
+    }];
+    [command_buffer->blit_encoder updateFence:encoder_fence];
+    iree_hal_metal_end_blit_encoder(command_buffer);
+  }
 
   @autoreleasepool {  // Use @autoreleasepool to trigger the autorelease within encoder creation.
     if (!command_buffer->compute_encoder) {
@@ -116,12 +127,27 @@
           computeCommandEncoderWithDispatchType:command_buffer->dispatch_type] retain];  // +1
     }
   }
+
+  if (encoder_fence != nil) {
+    [command_buffer->compute_encoder waitForFence:encoder_fence];
+  }
   return command_buffer->compute_encoder;
 }
 
 static id<MTLBlitCommandEncoder> iree_hal_metal_get_or_begin_blit_encoder(
     iree_hal_metal_command_buffer_t* command_buffer) {
-  if (command_buffer->compute_encoder) iree_hal_metal_end_compute_encoder(command_buffer);
+  id<MTLFence> encoder_fence = nil;
+  if (command_buffer->compute_encoder) {
+    // We would need to use a fence to synchronize "one or more resources across different passes
+    // within a command buffer."
+    // https://developer.apple.com/documentation/metal/resource_synchronization
+    encoder_fence = [command_buffer->command_buffer.device newFence];  // +1
+    [command_buffer->command_buffer addCompletedHandler:^(id<MTLCommandBuffer> cb) {
+      [encoder_fence release];  // -1
+    }];
+    [command_buffer->compute_encoder updateFence:encoder_fence];
+    iree_hal_metal_end_compute_encoder(command_buffer);
+  }
 
   @autoreleasepool {  // Use @autoreleasepool to trigger the autorelease within encoder creation.
     if (!command_buffer->blit_encoder) {
@@ -129,6 +155,10 @@
           [[command_buffer->command_buffer blitCommandEncoder] retain];  // +1
     }
   }
+
+  if (encoder_fence != nil) {
+    [command_buffer->blit_encoder waitForFence:encoder_fence];
+  }
   return command_buffer->blit_encoder;
 }