Caffe2 - C++ API
A deep learning, cross platform ML framework
mpscnn_context.h
1 
18 #pragma once
19 
20 #import <Metal/MTLBuffer.h>
21 #import <Metal/MTLDevice.h>
22 #import <Metal/MTLLibrary.h>
23 
24 #include <array>
25 #include <mutex>
26 #include <string>
27 #include <thread>
28 #include <unordered_map>
29 
30 namespace caffe2 {
31 
32 struct MPSCNNContext {
33  public:
34  id<MTLDevice> device;
35  id<MTLCommandQueue> commandQueue;
36  id<MTLLibrary> library;
37 
38  id<MTLComputePipelineState> getPipelineState(NSString* kernel);
39  id<MTLComputePipelineState> getSpecializedPipelineState(NSString* kernel,
40  const std::vector<ushort>& constants);
41 
42  private:
43  std::mutex pipelineCacheMutex_;
44  std::unordered_map<std::string, id<MTLComputePipelineState>> pipelineCache_;
45 };
46 
47 // get the singleton instance.
48 MPSCNNContext& getMPSCNNContext();
49 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.