[MPS] Add commitAndContinue to non blocking blits. (#51)

Author: kulinseth

aten/src/ATen/mps/MPSStream.h

@@ -34,6 +34,8 @@ typedef void* MTLDevice_t;
 namespace at {
 namespace mps {
 
+#define USE_MPSCOMMANDBUFFER 1
+
 //-----------------------------------------------------------------
 //  MPSStream
 //-----------------------------------------------------------------
@@ -53,6 +55,7 @@ class TORCH_API MPSStream
   MTLCommandBuffer_t commandBuffer();
   void commit(bool flush);
   void commitAndWait();
+  void commitAndContinue();
   void synchronize();
 
   void flush();

aten/src/ATen/mps/MPSStream.mm

@@ -59,6 +59,11 @@
   _commandBuffer = nil;
 }
 
+void MPSStream::commitAndContinue() {
+  assert(_commandBuffer);
+  [_commandBuffer commitAndContinue];
+}
+
 void MPSStream::flush() {
   if (_commandBuffer) {
     [_commandBuffer commit];
@@ -76,7 +81,6 @@
   [_commandBuffer release];
 }
 
-#define USE_MPSCOMMANDBUFFER 1
 
 void MPSStream::executeMPSGraph(MPSGraph* mpsGraph, NSDictionary* feeds, NSDictionary* results) {
   dispatch_sync(_serialQueue, ^() {

aten/src/ATen/native/mps/operations/Copy.mm

@@ -150,7 +150,11 @@ void copy_cast_mps(at::Tensor& dst, const at::Tensor& src,
         [blitEncoder endEncoding];
 
         if (non_blocking) {
+#if USE_MPSCOMMANDBUFFER
+          stream->commitAndContinue();
+#else
           stream->commit(true);
+#endif
         } else {
           stream->commitAndWait();
         }
@@ -215,7 +219,11 @@ void copy_cast_mps(at::Tensor& dst, const at::Tensor& src,
                                size:(NSUInteger)size];
         [blitEncoder endEncoding];
         if (non_blocking) {
+#if USE_MPSCOMMANDBUFFER
+          stream->commitAndContinue();
+#else
           stream->commit(true);
+#endif
         } else {
           stream->commitAndWait();
         }
@@ -300,7 +308,11 @@ void copy_blit_mps(void* dst, const void* src, size_t size) {
                                size:src_size];
         [blitEncoder endEncoding];
         // GPU to GPU copy needs flushing only, and no synchronization with CPU is necessary
-        stream->commit(true);
+#if USE_MPSCOMMANDBUFFER
+          stream->commitAndContinue();
+#else
+          stream->commit(true);
+#endif
       }
     });
   } else {