[MPS] Build MPS delegate with Werror=1

backends/apple/mps/operators/node_visitor.py

@@ -374,7 +374,10 @@ def process_placeholder_nodes(
             input_id = placeholder_visitor.define_tensor(node, mps_graph)
             mps_graph.input_ids.append(input_id)
 
-            if placeholder_visitor.convert_model_to_fp16:
+            if (
+                placeholder_visitor.convert_model_to_fp16
+                and node.meta["val"].dtype == torch.float32
+            ):
                 mps_node = MPSNode(
                     mpsnode_union=MPSCast(
                         input1_id=input_id,
@@ -393,7 +396,10 @@ def process_output_node(
     output_id = output_visitor.define_tensor(output_node, mps_graph)
     mps_graph.output_ids.append(output_id)
 
-    if output_visitor.convert_model_to_fp16:
+    if (
+        output_visitor.convert_model_to_fp16
+        and output_node.meta["val"].dtype == torch.float32
+    ):
         mps_node = MPSNode(
             mpsnode_union=MPSCast(
                 input1_id=output_id,

backends/apple/mps/runtime/MPSGraphBuilder.h

@@ -16,9 +16,11 @@
 #include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
 
 // MPS headers
+#include <executorch/backends/apple/mps/runtime/operations/MPSGraphVenturaOps.h>
 #include <executorch/backends/apple/mps/runtime/operations/OperationUtils.h>
 #include <executorch/backends/apple/mps/schema_generated.h>
 
+#include <unordered_map>
 #include <vector>
 
 namespace torch {

backends/apple/mps/runtime/MPSStream.h

@@ -91,9 +91,6 @@ class MPSStream {
   MPSCommandBuffer* _commandBuffer = nil;
   MPSCommandBuffer* _prevCommandBuffer = nil;
   id<MTLComputeCommandEncoder> _commandEncoder = nil;
-  MPSGraphExecutionDescriptor* _executionDescriptor = nil;
-  MPSGraphExecutableExecutionDescriptor* _executableExecutionDescriptor = nil;
-  MPSGraphCompilationDescriptor* _compilationDescriptor = nil;
   dispatch_queue_t _serialQueue = nullptr;
   // CommitAndContinue is disabled by default
   bool _enableCommitAndContinue = false;

backends/apple/mps/runtime/MPSStream.mm

@@ -16,40 +16,18 @@ @interface MPSGraphExecutionDescriptor ()
 namespace mps {
 namespace delegate {
 
-// threshold to perform adaptive commit if the accumulated size
-// of resources encoded on the command buffer exceeds that.
-static const size_t kCmdBufAdaptiveCommitThreshold = MB(64);
-
 //-----------------------------------------------------------------
 //  MPSStream
 //-----------------------------------------------------------------
 
 MPSStream::MPSStream() {
   _commandQueue = [MPSDevice::getInstance()->device() newCommandQueue];
   _serialQueue = dispatch_queue_create("metal gpu stream", nullptr);
-  _executionDescriptor = [MPSGraphExecutionDescriptor new];
-  _executableExecutionDescriptor = [MPSGraphExecutableExecutionDescriptor new];
-  _compilationDescriptor = [MPSGraphCompilationDescriptor new];
-
-  // internal CommitAndContinue heuristic of MPSGraph is disabled, and we
-  // control it via Adaptive Commit in Executorch-side
-  _executionDescriptor.enableCommitAndContinue = false;
-
-  // Choose level which optimizes for GPU
-  _compilationDescriptor.optimizationLevel = MPSGraphOptimizationLevel0;
-  _executionDescriptor.compilationDescriptor =  _compilationDescriptor;
 }
 
 MPSStream::~MPSStream() {
   [_commandQueue release];
   _commandQueue = nil;
-  [_executionDescriptor release];
-  [_compilationDescriptor release];
-  [_executableExecutionDescriptor release];
-
-  _executionDescriptor = nil;
-  _compilationDescriptor = nil;
-  _executableExecutionDescriptor = nil;
 
   assert(_commandBuffer == nil);
 }

backends/apple/mps/runtime/operations/BinaryOps.mm

@@ -118,6 +118,9 @@
     graphNode->input2_id(),                                                        \
     graphNode->output_id()                                                         \
   );                                                                               \
+  ET_CHECK_OR_RETURN_ERROR(                                                        \
+    isMacOS13OrNewer(), NotSupported,                                              \
+    "%s supported by MPS on MacOS13.0+/iOS16.1+", #aot_name);                       \
                                                                                    \
   _idToMPSGraphTensor[graphNode->output_id()] = binaryOpTensor(                    \
     getMPSGraphTensor(graphNode->input1_id()),                                     \
@@ -196,10 +199,7 @@
   MPSGraph* mpsGraph,
   const std::string& op_name) {
   MPSDataType mpsInputDataType = [primaryTensor dataType];
-  MPSDataType mpsOtherDataType = [secondaryTensor dataType];
-
   ScalarType inputDataType = getScalarType(mpsInputDataType);
-  ScalarType otherDataType = getScalarType(mpsOtherDataType);
 
   if(rounding_mode.has_value() && *rounding_mode == "trunc"){
     ET_CHECK_MSG(inputDataType != ScalarType::Half,

backends/apple/mps/runtime/operations/IndexingOps.mm

@@ -22,14 +22,14 @@
   if(castIndexTensor.dataType != MPSDataTypeInt32) {
     castIndexTensor = [mpsGraph castTensor:indexTensor
                                      toType:MPSDataTypeInt32
-                                       name:nil];
+                                       name:@"castTensor"];
   }
 
   return  [mpsGraph gatherWithUpdatesTensor:inputTensor
                               indicesTensor:castIndexTensor
                                        axis:dim
                             batchDimensions:0
-                                       name:nil];
+                                       name:@"indexSelect"];
 }
 
 Error
@@ -48,7 +48,7 @@
   if(castIndexTensor.dataType != MPSDataTypeInt32) {
     castIndexTensor = [_mpsGraph castTensor:indexTensor
                                      toType:MPSDataTypeInt32
-                                       name:nil];
+                                       name:@"castTensor"];
   }
 
   _idToMPSGraphTensor[graphNode->output_id()] =

backends/apple/mps/runtime/operations/MPSGraphVenturaOps.h

@@ -0,0 +1,207 @@
+
+//
+//  Copyright (c) 2023 Apple Inc. All rights reserved.
+//  Provided subject to the LICENSE file in the top level directory.
+//
+
+#pragma once
+
+#include <MetalPerformanceShadersGraph/MetalPerformanceShadersGraph.h>
+
+@interface MPSGraph (VenturaOps)
+
+#if !defined(__MAC_13_0) && (!defined(MAC_OS_X_VERSION_13_0) || (MAC_OS_X_VERSION_MIN_REQUIRED < MAC_OS_X_VERSION_13_0))
+
+typedef NS_ENUM(NSUInteger, MPSGraphResizeNearestRoundingMode) {
+  MPSGraphResizeNearestRoundingModeRoundPreferCeil = 0L,
+  MPSGraphResizeNearestRoundingModeRoundPreferFloor = 1L,
+  MPSGraphResizeNearestRoundingModeCeil = 2L,
+  MPSGraphResizeNearestRoundingModeFloor = 3L,
+  MPSGraphResizeNearestRoundingModeRoundToEven = 4L,
+  MPSGraphResizeNearestRoundingModeRoundToOdd = 5L,
+};
+
+// Define complex enums for MacOS 12
+#define MPSDataTypeComplexBit 0x01000000
+#define MPSDataTypeComplexFloat32 ((MPSDataType)(MPSDataTypeFloatBit | MPSDataTypeComplexBit | 64))
+#define MPSDataTypeComplexFloat16 ((MPSDataType)(MPSDataTypeFloatBit | MPSDataTypeComplexBit | 32))
+#endif
+
+- (MPSGraphTensor *_Nonnull)cumulativeSumWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                               axis:(NSInteger)axis
+                                               name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)sortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                      axis:(NSInteger)axis
+                                      name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)sortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                      axis:(NSInteger)axis
+                                descending:(BOOL)descending
+                                      name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)sortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                axisTensor:(MPSGraphTensor *_Nonnull)axisTensor
+                                descending:(BOOL)descending
+                                      name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)sortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                axisTensor:(MPSGraphTensor *_Nonnull)axisTensor
+                                      name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)argSortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                         axis:(NSInteger)axis
+                                         name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)argSortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                         axis:(NSInteger)axis
+                                   descending:(BOOL)descending
+                                         name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)argSortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                   axisTensor:(MPSGraphTensor *_Nonnull)axisTensor
+                                   descending:(BOOL)descending
+                                         name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)argSortWithTensor:(MPSGraphTensor *_Nonnull)tensor
+                                   axisTensor:(MPSGraphTensor *_Nonnull)axisTensor
+                                         name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)inverseOfTensor:(MPSGraphTensor *_Nonnull)inputTensor name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeNearestWithTensor:(MPSGraphTensor *_Nonnull)imagesTensor
+                                         sizeTensor:(MPSGraphTensor *_Nonnull)size
+                                nearestRoundingMode:(MPSGraphResizeNearestRoundingMode)nearestRoundingMode
+                                       centerResult:(BOOL)centerResult
+                                       alignCorners:(BOOL)alignCorners
+                                             layout:(MPSGraphTensorNamedDataLayout)layout
+                                               name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeNearestWithTensor:(MPSGraphTensor *_Nonnull)imagesTensor
+                                         sizeTensor:(MPSGraphTensor *_Nonnull)size
+                                  scaleOffsetTensor:(MPSGraphTensor *_Nonnull)scaleOffset
+                                nearestRoundingMode:(MPSGraphResizeNearestRoundingMode)nearestRoundingMode
+                                             layout:(MPSGraphTensorNamedDataLayout)layout
+                                               name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeBilinearWithTensor:(MPSGraphTensor *_Nonnull)imagesTensor
+                                          sizeTensor:(MPSGraphTensor *_Nonnull)size
+                                        centerResult:(BOOL)centerResult
+                                        alignCorners:(BOOL)alignCorners
+                                              layout:(MPSGraphTensorNamedDataLayout)layout
+                                                name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeBilinearWithTensor:(MPSGraphTensor *_Nonnull)imagesTensor
+                                          sizeTensor:(MPSGraphTensor *_Nonnull)size
+                                   scaleOffsetTensor:(MPSGraphTensor *_Nonnull)scaleOffset
+                                              layout:(MPSGraphTensorNamedDataLayout)layout
+                                                name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeNearestWithGradientTensor:(MPSGraphTensor *_Nonnull)gradient
+                                                      input:(MPSGraphTensor *_Nonnull)input
+                                        nearestRoundingMode:(MPSGraphResizeNearestRoundingMode)nearestRoundingMode
+                                               centerResult:(BOOL)centerResult
+                                               alignCorners:(BOOL)alignCorners
+                                                     layout:(MPSGraphTensorNamedDataLayout)layout
+                                                       name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeNearestWithGradientTensor:(MPSGraphTensor *_Nonnull)gradient
+                                                      input:(MPSGraphTensor *_Nonnull)input
+                                          scaleOffsetTensor:(MPSGraphTensor *_Nonnull)scaleOffset
+                                        nearestRoundingMode:(MPSGraphResizeNearestRoundingMode)nearestRoundingMode
+                                                     layout:(MPSGraphTensorNamedDataLayout)layout
+                                                       name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeBilinearWithGradientTensor:(MPSGraphTensor *_Nonnull)gradient
+                                                       input:(MPSGraphTensor *_Nonnull)input
+                                                centerResult:(BOOL)centerResult
+                                                alignCorners:(BOOL)alignCorners
+                                                      layout:(MPSGraphTensorNamedDataLayout)layout
+                                                        name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)resizeBilinearWithGradientTensor:(MPSGraphTensor *_Nonnull)gradient
+                                                       input:(MPSGraphTensor *_Nonnull)input
+                                           scaleOffsetTensor:(MPSGraphTensor *_Nonnull)scaleOffset
+                                                      layout:(MPSGraphTensorNamedDataLayout)layout
+                                                        name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)sampleGridWithSourceTensor:(MPSGraphTensor *_Nonnull)source
+                                      coordinateTensor:(MPSGraphTensor *_Nonnull)coordinates
+                                                layout:(MPSGraphTensorNamedDataLayout)layout
+                                  normalizeCoordinates:(BOOL)normalizeCoordinates
+                                   relativeCoordinates:(BOOL)relativeCoordinates
+                                          alignCorners:(BOOL)alignCorners
+                                           paddingMode:(MPSGraphPaddingMode)paddingMode
+                                          samplingMode:(MPSGraphResizeMode)samplingMode
+                                         constantValue:(double)constantValue
+                                                  name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)sampleGridWithSourceTensor:(MPSGraphTensor *_Nonnull)source
+                                      coordinateTensor:(MPSGraphTensor *_Nonnull)coordinates
+                                                layout:(MPSGraphTensorNamedDataLayout)layout
+                                  normalizeCoordinates:(BOOL)normalizeCoordinates
+                                   relativeCoordinates:(BOOL)relativeCoordinates
+                                          alignCorners:(BOOL)alignCorners
+                                           paddingMode:(MPSGraphPaddingMode)paddingMode
+                                   nearestRoundingMode:(MPSGraphResizeNearestRoundingMode)nearestRoundingMode
+                                         constantValue:(double)constantValue
+                                                  name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)truncateWithTensor:(MPSGraphTensor *_Nonnull)tensor name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)transposeTensor:(MPSGraphTensor *_Nonnull)tensor
+                                permutation:(NSArray<NSNumber *> *_Nonnull)permutation
+                                       name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)bitwiseANDWithPrimaryTensor:(MPSGraphTensor *_Nonnull)primaryTensor
+                                        secondaryTensor:(MPSGraphTensor *_Nonnull)secondaryTensor
+                                                   name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)bitwiseORWithPrimaryTensor:(MPSGraphTensor *_Nonnull)primaryTensor
+                                       secondaryTensor:(MPSGraphTensor *_Nonnull)secondaryTensor
+                                                  name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)bitwiseXORWithPrimaryTensor:(MPSGraphTensor *_Nonnull)primaryTensor
+                                        secondaryTensor:(MPSGraphTensor *_Nonnull)secondaryTensor
+                                                   name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nonnull)bitwiseNOTWithTensor:(MPSGraphTensor *_Nonnull)tensor name:(NSString *_Nullable)name;
+
+#if !defined(MAC_OS_X_VERSION_12_2) || (MAC_OS_X_VERSION_MIN_REQUIRED < MAC_OS_X_VERSION_12_2)
+- (MPSGraphTensor *_Nullable)expandDimsOfTensor:(MPSGraphTensor *_Nullable)tensor
+                                           axis:(NSInteger)axis
+                                           name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nullable)expandDimsOfTensor:(MPSGraphTensor *_Nullable)tensor
+                                           axes:(NSArray<NSNumber *> *_Nullable)axes
+                                           name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nullable)squeezeTensor:(MPSGraphTensor *_Nullable)tensor
+                                      axes:(NSArray<NSNumber *> *_Nullable)axes
+                                      name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nullable)squeezeTensor:(MPSGraphTensor *_Nullable)tensor
+                                      axis:(NSInteger)axis
+                                      name:(NSString *_Nullable)name;
+
+- (NSArray<MPSGraphTensor *> *_Nullable)
+    maxPooling2DReturnIndicesWithSourceTensor:(MPSGraphTensor *_Nullable)source
+                                   descriptor:(MPSGraphPooling2DOpDescriptor *_Nullable)descriptor
+                                         name:(NSString *_Nullable)name;
+
+- (MPSGraphTensor *_Nullable)coordinateAlongAxis:(NSInteger)axis
+                                 withShapeTensor:(MPSGraphTensor *_Nullable)shapeTensor
+                                            name:(NSString *_Nullable)name;
+
+- (NSArray<MPSGraphTensor *> *_Nullable)splitTensor:(MPSGraphTensor *_Nullable)tensor
+                                   splitSizesTensor:(MPSGraphTensor *_Nullable)splitSizesTensor
+                                               axis:(NSInteger)axis
+                                               name:(NSString *_Nullable)name;
+
+- (NSArray<MPSGraphTensor *> *_Nullable)splitTensor:(MPSGraphTensor *_Nullable)tensor
+                                         splitSizes:(NSArray<NSNumber *> *_Nullable)splitSizes
+                                               axis:(NSInteger)axis
+                                               name:(NSString *_Nullable)name;
+#endif
+
+@end

backends/apple/mps/runtime/operations/OperationUtils.mm

@@ -227,7 +227,6 @@
 
 MPSGraphTensor*
 MPSGraphBuilder::getMPSGraphTensor(int32_t id) {
-  static int32_t cacheEntries = _idToMPSGraphTensor.size();
   return _idToMPSGraphTensor[id];
 }
 

backends/apple/mps/runtime/operations/PadOps.mm

@@ -27,7 +27,6 @@
               "invalid padding argument of size %d", padding_size);
 
   auto input_sizes = getMPSShapeVec(input.shape);
-  int64_t nbatch = 1;
   int64_t ndims = input_sizes.size();
 
   ET_CHECK_MSG(
@@ -39,7 +38,6 @@
   int dim_w = padding_dim;
   int dim_h = padding_dim - 1;
   int dim_d = padding_dim - 2;
-  int dim_slices = 0;
 
   if (mode != MPSGraphPaddingModeConstant && ndims > padding_dim) {
     bool valid_dims = input_sizes[1] != 0 && input_sizes[padding_dim] != 0;
@@ -59,8 +57,6 @@
     dim_w += dim_diff;
     dim_h += dim_diff;
     dim_d += dim_diff;
-    dim_slices++;
-    nbatch = input_sizes[0];
   }
 
   int64_t pad_l = padding[0];
@@ -70,13 +66,11 @@
   int64_t pad_front = padding_size > 4 ? padding[4] : 0;
   int64_t pad_back  = padding_size > 4 ? padding[5] : 0;
 
-  int64_t nplane = input_sizes[dim_slices];
   int64_t input_w = input_sizes[dim_w];
   int64_t output_w  = input_w + pad_l + pad_r;
   int64_t input_h = padding_dim > 1 ? input_sizes[dim_h] : 0;
   int64_t output_h = padding_dim > 1 ? input_h + pad_t + pad_b : 0;
   int64_t input_d = padding_dim > 2 ? input_sizes[dim_d] : 0;
-  int64_t output_d = padding_dim > 2 ? input_d + pad_front + pad_back : 0;
 
   ET_CHECK_MSG(
     output_w >= 1 || output_h >= padding_dim - 1,