From a8d7b885c52e9acd25775bbd50de238da5a5624d Mon Sep 17 00:00:00 2001 From: Tao Xu Date: Fri, 17 Sep 2021 09:16:39 -0700 Subject: [PATCH] [CoreML][iOS/MacOS] Add the CoreML executor (#64522) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64522 The `PTMCoreMLExecutor` serves as a bridge between the delegate APIs and Core ML runtime. ghstack-source-id: 138324217 allow-large-files Test Plan: iOS: Run the CoreML tests in the playground app MacOS: ``` buck test pp-macos PASS 633ms 1 Passed 0 Skipped 0 Failed CoreMLTests ``` {F657776101} Reviewed By: raziel, iseeyuan Differential Revision: D30594042 fbshipit-source-id: a42a5307a24c2f364333829f3a84f7b9a51e1b3e --- .../jit/backends/coreml/objc/PTMCoreMLBackend.mm | 261 +++++++++++++++++++++ .../jit/backends/coreml/objc/PTMCoreMLExecutor.h | 34 +++ .../jit/backends/coreml/objc/PTMCoreMLExecutor.mm | 172 ++++++++++++++ 3 files changed, 467 insertions(+) create mode 100644 torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm create mode 100644 torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h create mode 100644 torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.mm diff --git a/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm new file mode 100644 index 0000000..720b137 --- /dev/null +++ b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLBackend.mm @@ -0,0 +1,261 @@ +#include +#include +#include + +#import + +namespace torch { +namespace jit { +namespace mobile { +namespace coreml { + +static constexpr int SUPPORTED_COREML_VER = 4; + +enum TensorType { + Float, + Double, + Int, + Long, + Undefined, +}; + +static inline c10::ScalarType scalarType(TensorType type) { + switch (type) { + case TensorType::Float: + return c10::ScalarType::Float; + case TensorType::Double: + return c10::ScalarType::Double; + case TensorType::Int: + return c10::ScalarType::Int; + case TensorType::Long: + return c10::ScalarType::Long; + case TensorType::Undefined: + return c10::ScalarType::Undefined; + default: + return c10::ScalarType::Undefined; + } +} + +static id parse(NSString* jsonStr) { + NSData* data = [jsonStr dataUsingEncoding:NSUTF8StringEncoding]; + NSError* error; + id result = [NSJSONSerialization JSONObjectWithData:data + options:0 + error:&error]; + if (error || !result) { + TORCH_CHECK( + false, + "parsing JSON string failed!", + error.localizedDescription.UTF8String); + } + + return result; +} + +struct TensorSpec { + public: + TensorSpec() = delete; + TensorSpec(NSArray* spec) { + TORCH_CHECK(spec.count == 3); + name_ = spec[0]; + dtype_ = (TensorType)spec[1].intValue; + NSArray* sizes = parse(spec[2]); + for (NSString* dim in sizes) { + sizes_.emplace_back(dim.integerValue); + } + } + int64_t numel() const { + return std::accumulate( + begin(sizes_), end(sizes_), 1, std::multiplies()); + } + NSString* name() { + return name_; + } + std::vector sizes() { + return sizes_; + } + TensorType dtype() { + return dtype_; + } + + private: + NSString* name_ = @""; + TensorType dtype_ = TensorType::Float; + std::vector sizes_{}; +}; + +struct CoreMLConfig { + public: + CoreMLConfig() = delete; + CoreMLConfig(NSDictionary* dict) + : coreMLVersion_([dict[@"spec_ver"] intValue]), + backend_([dict[@"backend"] lowercaseString]), + allow_low_precision_([dict[@"allow_low_precision"] boolValue]) { + TORCH_CHECK( + coreMLVersion_ >= SUPPORTED_COREML_VER, + "Only Core ML version 4 and above are supported"); + } + int64_t coreMLVersion() const { + return coreMLVersion_; + } + NSString* backend() const { + return backend_; + } + bool allowLowPrecision() const { + return allow_low_precision_; + } + + private: + int64_t coreMLVersion_ = SUPPORTED_COREML_VER; + NSString* backend_ = @"CPU"; + bool allow_low_precision_ = true; +}; + +struct MetaData { + public: + MetaData(NSDictionary* dict) + : torchVer_(dict[@"torch_ver"]), + coremltoolVer_(dict[@"coremltool_ver"]) {} + NSString* torchVer() const { + return torchVer_; + } + NSString* coremltoolVer() const { + return coremltoolVer_; + } + + private: + NSString* torchVer_ = @""; + NSString* coremltoolVer_ = @""; +}; + +// Wrap the Objective-C executor into a C++ to be able to pack into IValue +struct API_AVAILABLE(ios(11.0), macos(10.13)) CoreMLExecutorWrapper + : public CustomClassHolder { + public: + CoreMLExecutorWrapper( + PTMCoreMLExecutor* executor, + std::vector& inputs, + std::vector& outputs, + CoreMLConfig config) + : executor_(executor), + inputs_(inputs), + outputs_(outputs), + config_(config) {} + c10::List execute(c10::impl::GenericList inputs) { + std::vector inputSpecs; + std::vector outputSpecs; + int inputSpecIndex = 0; + // pack the inputs + for (int i = 0; i < inputs.size(); ++i) { + auto val = inputs.get(i); + if (val.isTuple()) { + auto tuples = val.toTuple()->elements(); + for (auto& ival : tuples) { + TORCH_CHECK(ival.isTensor()); + auto tensor = ival.toTensor(); + PTMCoreMLFeatureSpecs spec{ + .name = inputs_[inputSpecIndex].name(), + .tensor = tensor, + }; + inputSpecs.emplace_back(spec); + ++inputSpecIndex; + } + } else { + TORCH_CHECK(val.isTensor()); + auto tensor = val.toTensor(); + PTMCoreMLFeatureSpecs spec{ + .name = inputs_[inputSpecIndex].name(), + .tensor = tensor, + }; + inputSpecs.emplace_back(spec); + ++inputSpecIndex; + } + } + // pack the outputs + c10::List outputs; + id results = [executor_ forwardWithInputs:inputSpecs]; + for (auto& spec : outputs_) { + MLFeatureValue* val = [results featureValueForName:spec.name()]; + TORCH_CHECK(val.multiArrayValue); + // Currently, only Float type is supported + TORCH_CHECK(val.multiArrayValue.dataType == MLMultiArrayDataTypeFloat32); + auto tensor = at::empty(spec.sizes(), scalarType(spec.dtype())); + int64_t count = val.multiArrayValue.count; + memcpy( + tensor.data_ptr(), + (float*)val.multiArrayValue.dataPointer, + count * sizeof(float)); + outputs.push_back(tensor); + } + return outputs; + } + + private: + PTMCoreMLExecutor* executor_ = nullptr; + std::vector inputs_; + std::vector outputs_; + CoreMLConfig config_; +}; + +class API_AVAILABLE(ios(11.0), macos(10.13)) CoreMLBackend + : public torch::jit::PyTorchBackendInterface { + public: + c10::impl::GenericDict compile( + c10::IValue processed, + c10::impl::GenericDict method_compile_spec) override { + auto modelDict = processed.toGenericDict(); + NSString* specs = [[NSString alloc] + initWithCString:modelDict.at("extra").toStringRef().c_str() + encoding:NSUTF8StringEncoding]; + NSDictionary* dict = parse(specs); + NSArray* inputs = dict[@"inputs"]; + NSArray* outputs = dict[@"outputs"]; + std::vector inputSpecs, outputSpecs; + for (NSArray* input in inputs) { + inputSpecs.emplace_back(TensorSpec(input)); + } + for (NSArray* output in outputs) { + outputSpecs.emplace_back(TensorSpec(output)); + } + auto config = CoreMLConfig(dict[@"config"]); + const std::string& model = modelDict.at("model").toStringRef(); + const std::string& sha256 = modelDict.at("hash").toStringRef(); + PTMCoreMLExecutor* executor = [PTMCoreMLExecutor new]; + bool result = [executor compileMLModel:model identifier:sha256]; + TORCH_CHECK(result, "Compiling MLModel failed!"); + auto executorWrapper = c10::make_intrusive( + executor, inputSpecs, outputSpecs, config); + auto handle = IValue::make_capsule(executorWrapper); + c10::Dict ret(StringType::get(), c10::AnyType::get()); + ret.insert("forward", handle); + return c10::impl::toGenericDict(ret); + } + + c10::impl::GenericList execute( + c10::IValue handle, + c10::impl::GenericList inputs) override { + auto executor = c10::static_intrusive_pointer_cast( + handle.toCapsule()); + auto outputs = executor->execute(inputs); + return c10::impl::toList(outputs); + } + bool is_available() override { +#if !defined(__APPLE__) + return false; +#else + if (@available(iOS 14, macOS 10.13, *)) { + return true; + } else { + return false; + } +#endif + } +}; + +API_AVAILABLE(ios(11.0), macos(10.13)) +static auto cls = torch::jit::backend("coreml"); + +} // namespace +} +} +} diff --git a/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h new file mode 100644 index 0000000..7cac4b8 --- /dev/null +++ b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h @@ -0,0 +1,34 @@ +#import +#include + +#include +#include + +struct PTMCoreMLFeatureSpecs { + NSString* name; + at::Tensor tensor; +}; + +API_AVAILABLE(ios(11.0), macos(10.13)) +@interface PTMCoreMLFeatureProvider : NSObject +- (instancetype)initWithFeatureSpecs: + (const std::vector&)specs + CoreMLVersion:(NSUInteger)ver; +@end + +API_AVAILABLE(ios(11.0), macos(10.13)) +@interface PTMCoreMLExecutor : NSObject + +@property(nonatomic, readonly, copy) NSString* modelPath; +@property(nonatomic, readonly, copy) NSString* compiledModelPath; +@property(nonatomic, copy) NSString* backend; +@property(nonatomic, assign) BOOL allowLowPrecision; +@property(nonatomic, assign) NSUInteger coreMLVersion; + +- (BOOL)compileMLModel:(const std::string&)modelSpecs + identifier:(const std::string&)identifier; +- (id)forwardWithInputs: + (const std::vector&)inputs; +- (BOOL)cleanup; + +@end diff --git a/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.mm b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.mm new file mode 100644 index 0000000..7647e16 --- /dev/null +++ b/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.mm @@ -0,0 +1,172 @@ +#include +#include + +#import + +#include +#include +#include + +@implementation PTMCoreMLFeatureProvider { + NSUInteger _coremlVersion; + std::vector _specs; +} + +@synthesize featureNames = _featureNames; + +- (instancetype)initWithFeatureSpecs: + (const std::vector&)specs + CoreMLVersion:(NSUInteger)ver { + self = [super init]; + if (self) { + _coremlVersion = ver; + _specs = specs; + NSMutableArray* names = [NSMutableArray new]; + for (auto& spec : _specs) { + [names addObject:spec.name]; + } + _featureNames = [[NSSet alloc] initWithArray:names]; + } + return self; +} + +- (nullable MLFeatureValue*)featureValueForName:(NSString*)featureName { + for (auto& spec : _specs) { + if ([spec.name isEqualToString:featureName]) { + NSMutableArray* shape = [NSMutableArray new]; + for (auto& dim : spec.tensor.sizes().vec()) { + [shape addObject:@(dim)]; + } + NSMutableArray* strides = [NSMutableArray new]; + for (auto& step : spec.tensor.strides().vec()) { + [strides addObject:@(step)]; + } + NSError* error = nil; + TORCH_CHECK(spec.tensor.dtype() == c10::kFloat); + MLMultiArray* mlArray = [[MLMultiArray alloc] + initWithDataPointer:spec.tensor.data_ptr() + shape:shape + dataType:MLMultiArrayDataTypeFloat32 + strides:strides + deallocator:(^(void* bytes){ + })error:&error]; + return [MLFeatureValue featureValueWithMultiArray:mlArray]; + } + } + return nil; +} + +@end + +@implementation PTMCoreMLExecutor { + MLModel* _mlModel; +} + +- (BOOL)compileMLModel:(const std::string&)modelSpecs + identifier:(const std::string&)identifier + API_AVAILABLE(ios(11.0), macos(10.13)) { + _modelPath = [self _save:modelSpecs + identifier:[NSString stringWithCString:identifier.c_str() + encoding:NSUTF8StringEncoding]]; + NSError* error; + NSURL* compiledModelPath = nil; + if (@available(iOS 11.0, macOS 10.13, *)) { + compiledModelPath = + [MLModel compileModelAtURL:[NSURL URLWithString:_modelPath] + error:&error]; + } else { + TORCH_CHECK(false, "CoreML is not available on your deivce"); + } + if (error || !compiledModelPath) { + // remove cached models if compalition failed. + [self cleanup]; + TORCH_CHECK( + false, + "Error compiling model", + [error localizedDescription].UTF8String); + return NO; + } + if (@available(iOS 12.0, macOS 10.14, *)) { + MLModelConfiguration* config = [MLModelConfiguration alloc]; + MLComputeUnits backend = MLComputeUnitsCPUOnly; + if ([self.backend isEqualToString:@"cpuandgpu"]) { + backend = MLComputeUnitsCPUAndGPU; + } else if ([self.backend isEqualToString:@"all"]) { + backend = MLComputeUnitsAll; + } + config.computeUnits = backend; + config.allowLowPrecisionAccumulationOnGPU = self.allowLowPrecision; + _mlModel = [MLModel modelWithContentsOfURL:compiledModelPath + configuration:config + error:&error]; + } else { + _mlModel = [MLModel modelWithContentsOfURL:compiledModelPath error:&error]; + } + if (error || !_mlModel) { + TORCH_CHECK( + false, "Error loading MLModel", error.localizedDescription.UTF8String); + } + + _compiledModelPath = compiledModelPath.path; + return YES; +} + +- (id)forwardWithInputs: + (const std::vector&)inputs { + NSError* error; + PTMCoreMLFeatureProvider* inputFeature = [[PTMCoreMLFeatureProvider alloc] + initWithFeatureSpecs:inputs + CoreMLVersion:self.coreMLVersion]; + if (inputFeature == nil) { + NSLog(@"inputFeature is not initialized."); + return nil; + } + if (@available(iOS 11.0, macOS 10.13, *)) { + MLPredictionOptions* options = [[MLPredictionOptions alloc] init]; + id outputFeature = + [_mlModel predictionFromFeatures:inputFeature + options:options + error:&error]; + if (error || !outputFeature) { + TORCH_CHECK( + false, + "Error running the prediction", + error.localizedDescription.UTF8String); + } + + return outputFeature; + } else { + TORCH_CHECK("Core ML is available on iOS 11.0 and above"); + return nil; + } +} + +- (BOOL)cleanup { + NSFileManager* fileManager = [NSFileManager defaultManager]; + NSError* error; + if (![fileManager fileExistsAtPath:_modelPath]) { + [fileManager removeItemAtPath:_modelPath error:&error]; + } + if (![fileManager fileExistsAtPath:_compiledModelPath]) { + [fileManager removeItemAtPath:_compiledModelPath error:&error]; + } + return !error; +} + +- (NSString*)_save:(const std::string&)spec identifier:(NSString*)identifier { + NSURL* temporaryDirectoryURL = [NSURL fileURLWithPath:NSTemporaryDirectory() + isDirectory:YES]; + NSString* modelPath = [NSString + stringWithFormat:@"%@/%@", temporaryDirectoryURL.path, identifier]; + if (![[NSFileManager defaultManager] fileExistsAtPath:modelPath]) { + // Note that the serialized protobuf binary contains bytes, not text; + // see + // https://developers.google.com/protocol-buffers/docs/pythontutorial#parsing-and-serialization + NSData* data = [NSData dataWithBytes:spec.c_str() length:spec.length()]; + BOOL ret = [data writeToFile:modelPath atomically:YES]; + TORCH_CHECK(ret, "Save MLModel failed!", modelPath.UTF8String); + } + return modelPath; +} + +@end -- 2.7.4