1// Copyright 2020 Google Inc. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at:
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#import "tensorflow/lite/objc/apis/TFLCoreMLDelegate.h"
16
17#ifdef COCOAPODS
18@import TensorFlowLiteCCoreML;
19#else
20#include "tensorflow/lite/delegates/coreml/coreml_delegate.h"
21#endif
22
23NS_ASSUME_NONNULL_BEGIN
24
25@implementation TFLCoreMLDelegateOptions
26
27- (instancetype)init {
28  self = [super init];
29  if (self != nil) {
30    _coreMLVersion = 0;
31    _maxDelegatedPartitions = 0;
32    _minNodesPerPartition = 2;
33    _enabledDevices = TFLCoreMLDelegateEnabledDevicesNeuralEngine;
34  }
35
36  return self;
37}
38
39@end
40
41@implementation TFLCoreMLDelegate
42
43@synthesize cDelegate = _cDelegate;
44
45#pragma mark - NSObject
46
47- (void)dealloc {
48  TfLiteCoreMlDelegateDelete((TfLiteDelegate*)self.cDelegate);
49}
50
51#pragma mark - Public
52
53- (nullable instancetype)init {
54  TFLCoreMLDelegateOptions* options = [[TFLCoreMLDelegateOptions alloc] init];
55  return [self initWithOptions:options];
56}
57
58- (nullable instancetype)initWithOptions:(TFLCoreMLDelegateOptions*)options {
59  self = [super init];
60  if (self != nil) {
61    TfLiteCoreMlDelegateOptions cOptions;
62
63    cOptions.coreml_version = options.coreMLVersion;
64    cOptions.max_delegated_partitions = options.maxDelegatedPartitions;
65    cOptions.min_nodes_per_partition = options.minNodesPerPartition;
66
67    switch (options.enabledDevices) {
68      case TFLCoreMLDelegateEnabledDevicesNeuralEngine:
69        cOptions.enabled_devices = TfLiteCoreMlDelegateDevicesWithNeuralEngine;
70        break;
71
72      case TFLCoreMLDelegateEnabledDevicesAll:
73        cOptions.enabled_devices = TfLiteCoreMlDelegateAllDevices;
74        break;
75    }
76    _cDelegate = TfLiteCoreMlDelegateCreate(&cOptions);
77    if (_cDelegate == nil) {
78      return nil;
79    }
80  }
81  return self;
82}
83
84@end
85
86NS_ASSUME_NONNULL_END
87