Preserve command buffer in runMPSGraph() to optimize performance (#45)

* Preserve command buffer in runMPSGraph() to optimize performance

* Add macro to toggle between commandBuffer and commandQueue methods in runMPSGraph

Author: razarmehr

aten/src/ATen/mps/MPSStream.h

@@ -56,6 +56,7 @@ class TORCH_API MPSStream
   void synchronize();
 
   void flush();
+  void executeMPSGraph(MPSGraph* mpsGraph, NSDictionary* feeds, NSDictionary* results);
 
   /// Get the MPS device index that this stream is associated with.
   c10::DeviceIndex device_index() const { return _stream.device_index(); }
@@ -71,6 +72,7 @@ class TORCH_API MPSStream
   Stream _stream;
   MTLCommandQueue_t   _commandQueue = nil;
   MTLCommandBuffer_t  _commandBuffer = nil;
+  MPSGraphExecutionDescriptor *_executionDescriptor = nil;
   void _flush(bool commitAndWait) const;
 
   dispatch_queue_t    _serialQueue = nullptr;

aten/src/ATen/mps/MPSStream.mm

@@ -13,11 +13,16 @@
   _commandQueue = [MPSDevice::getInstance()->device() newCommandQueue];
   TORCH_CHECK(_stream.device_type() == DeviceType::MPS);
   _serialQueue = dispatch_queue_create("metal gpu stream", NULL);
+  _executionDescriptor = [MPSGraphExecutionDescriptor new];
+  _executionDescriptor.completionHandler = ^(NSDictionary<MPSGraphTensor *,
+                                             MPSGraphTensorData *> * resultsDictionary,
+                                             NSError * _Nullable error) { };
 }
 
 MPSStream::~MPSStream() {
-  [_commandQueue autorelease];
+  [_commandQueue release];
   _commandQueue = nil;
+  [_executionDescriptor release];
 
   assert(_commandBuffer == nil);
 }
@@ -71,6 +76,27 @@
   [_commandBuffer release];
 }
 
+#define USE_MPSCOMMANDBUFFER 1
+
+void MPSStream::executeMPSGraph(MPSGraph* mpsGraph, NSDictionary* feeds, NSDictionary* results) {
+  dispatch_sync(_serialQueue, ^() {
+#if USE_MPSCOMMANDBUFFER
+    [mpsGraph encodeToCommandBuffer:commandBuffer()
+                              feeds:feeds
+                   targetOperations:nil
+                  resultsDictionary:results
+                executionDescriptor:_executionDescriptor];
+#else
+    commit(true);
+    [mpsGraph runAsyncWithMTLCommandQueue:_commandQueue
+                                    feeds:feeds
+                         targetOperations:nil
+                        resultsDictionary:results
+                      executionDescriptor:_executionDescriptor];
+#endif
+ });
+}
+
 //-----------------------------------------------------------------
 //  MPSStreamImpl
 //-----------------------------------------------------------------

aten/src/ATen/native/mps/OperationUtils.mm

@@ -60,30 +60,8 @@
   return gen;
 }
 
-void runMPSGraph(
-    MPSStream* mpsStream,
-    MPSGraph* mpsGraph,
-    NSDictionary* feeds,
-    NSDictionary* results) {
-  dispatch_sync(mpsStream->queue(), ^() {
-    @autoreleasepool {
-      mpsStream->commit(true);
-      id<MTLCommandQueue> commandQueue = mpsStream->commandQueue();
-      MPSGraphExecutionDescriptor *executionDescriptor = [[MPSGraphExecutionDescriptor new] autorelease];
-
-      executionDescriptor.completionHandler = ^(NSDictionary<MPSGraphTensor *,
-                                                MPSGraphTensorData *> * resultsDictionary,
-                                                NSError * _Nullable error) {
-      };
-
-      [mpsGraph runAsyncWithMTLCommandQueue:commandQueue
-                                feeds:feeds
-                     targetOperations:nil
-                    resultsDictionary:results
-                  executionDescriptor:executionDescriptor];
-
-    }
-  });
+void runMPSGraph(MPSStream* mpsStream, MPSGraph* mpsGraph, NSDictionary* feeds, NSDictionary* results) {
+  mpsStream->executeMPSGraph(mpsGraph, feeds, results);
 }
 
 MPSDataType getMPSDataType(ScalarType scalar_type) {