From b0b38a0013c819a6db4156330cbbe2e0dab11bd8 Mon Sep 17 00:00:00 2001
From: Camillo Lugaresi <camillol@google.com>
Date: Tue, 6 Dec 2022 08:29:35 -0800
Subject: [PATCH] Internal change

PiperOrigin-RevId: 493313240
---
 mediapipe/gpu/metal_shared_resources.h       | 40 +++++++++++
 mediapipe/gpu/metal_shared_resources.mm      | 73 ++++++++++++++++++++
 mediapipe/gpu/metal_shared_resources_test.mm | 49 +++++++++++++
 3 files changed, 162 insertions(+)
 create mode 100644 mediapipe/gpu/metal_shared_resources.h
 create mode 100644 mediapipe/gpu/metal_shared_resources.mm
 create mode 100644 mediapipe/gpu/metal_shared_resources_test.mm

diff --git a/mediapipe/gpu/metal_shared_resources.h b/mediapipe/gpu/metal_shared_resources.h
new file mode 100644
index 000000000..341860a2d
--- /dev/null
+++ b/mediapipe/gpu/metal_shared_resources.h
@@ -0,0 +1,40 @@
+#ifndef MEDIAPIPE_GPU_METAL_SHARED_RESOURCES_H_
+#define MEDIAPIPE_GPU_METAL_SHARED_RESOURCES_H_
+
+#import <CoreVideo/CVMetalTextureCache.h>
+#import <CoreVideo/CoreVideo.h>
+#import <Foundation/NSObject.h>
+#import <Metal/Metal.h>
+
+#ifndef __OBJC__
+#error This class must be built as Objective-C++.
+#endif  // !__OBJC__
+
+@interface MPPMetalSharedResources : NSObject {
+}
+
+- (instancetype)init NS_DESIGNATED_INITIALIZER;
+
+@property(readonly) id<MTLDevice> mtlDevice;
+@property(readonly) id<MTLCommandQueue> mtlCommandQueue;
+#if COREVIDEO_SUPPORTS_METAL
+@property(readonly) CVMetalTextureCacheRef mtlTextureCache;
+#endif
+
+@end
+
+namespace mediapipe {
+
+class MetalSharedResources {
+ public:
+  MetalSharedResources();
+  ~MetalSharedResources();
+  MPPMetalSharedResources* resources() { return resources_; }
+
+ private:
+  MPPMetalSharedResources* resources_;
+};
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_GPU_METAL_SHARED_RESOURCES_H_
diff --git a/mediapipe/gpu/metal_shared_resources.mm b/mediapipe/gpu/metal_shared_resources.mm
new file mode 100644
index 000000000..80d755a01
--- /dev/null
+++ b/mediapipe/gpu/metal_shared_resources.mm
@@ -0,0 +1,73 @@
+#import "mediapipe/gpu/metal_shared_resources.h"
+
+@interface MPPMetalSharedResources ()
+@end
+
+@implementation MPPMetalSharedResources {
+}
+
+@synthesize mtlDevice = _mtlDevice;
+@synthesize mtlCommandQueue = _mtlCommandQueue;
+#if COREVIDEO_SUPPORTS_METAL
+@synthesize mtlTextureCache = _mtlTextureCache;
+#endif
+
+- (instancetype)init {
+  self = [super init];
+  if (self) {
+  }
+  return self;
+}
+
+- (void)dealloc {
+#if COREVIDEO_SUPPORTS_METAL
+  if (_mtlTextureCache) {
+    CFRelease(_mtlTextureCache);
+    _mtlTextureCache = NULL;
+  }
+#endif
+}
+
+- (id<MTLDevice>)mtlDevice {
+  @synchronized(self) {
+    if (!_mtlDevice) {
+      _mtlDevice = MTLCreateSystemDefaultDevice();
+    }
+  }
+  return _mtlDevice;
+}
+
+- (id<MTLCommandQueue>)mtlCommandQueue {
+  @synchronized(self) {
+    if (!_mtlCommandQueue) {
+      _mtlCommandQueue = [self.mtlDevice newCommandQueue];
+    }
+  }
+  return _mtlCommandQueue;
+}
+
+#if COREVIDEO_SUPPORTS_METAL
+- (CVMetalTextureCacheRef)mtlTextureCache {
+  @synchronized(self) {
+    if (!_mtlTextureCache) {
+      CVReturn __unused err =
+          CVMetalTextureCacheCreate(NULL, NULL, self.mtlDevice, NULL, &_mtlTextureCache);
+      NSAssert(err == kCVReturnSuccess, @"Error at CVMetalTextureCacheCreate %d ; device %@", err,
+               self.mtlDevice);
+      // TODO: register and flush metal caches too.
+    }
+  }
+  return _mtlTextureCache;
+}
+#endif
+
+@end
+
+namespace mediapipe {
+
+MetalSharedResources::MetalSharedResources() {
+  resources_ = [[MPPMetalSharedResources alloc] init];
+}
+MetalSharedResources::~MetalSharedResources() {}
+
+}  // namespace mediapipe
diff --git a/mediapipe/gpu/metal_shared_resources_test.mm b/mediapipe/gpu/metal_shared_resources_test.mm
new file mode 100644
index 000000000..9eb53a9b7
--- /dev/null
+++ b/mediapipe/gpu/metal_shared_resources_test.mm
@@ -0,0 +1,49 @@
+// Copyright 2019 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#import <UIKit/UIKit.h>
+#import <XCTest/XCTest.h>
+
+#include <memory>
+
+#include "absl/memory/memory.h"
+#include "mediapipe/framework/port/threadpool.h"
+
+#import "mediapipe/gpu/gpu_shared_data_internal.h"
+#import "mediapipe/gpu/metal_shared_resources.h"
+
+@interface MPPMetalSharedResourcesTests : XCTestCase {
+}
+@end
+
+@implementation MPPMetalSharedResourcesTests
+
+// This test verifies that the internal Objective-C object is correctly
+// released when the C++ wrapper is released.
+- (void)testCorrectlyReleased {
+  __weak id metalRes = nil;
+  std::weak_ptr<mediapipe::GpuResources> weakGpuRes;
+  @autoreleasepool {
+    auto maybeGpuRes = mediapipe::GpuResources::Create();
+    XCTAssertTrue(maybeGpuRes.ok());
+    weakGpuRes = *maybeGpuRes;
+    metalRes = (**maybeGpuRes).metal_shared().resources();
+    XCTAssertNotEqual(weakGpuRes.lock(), nullptr);
+    XCTAssertNotNil(metalRes);
+  }
+  XCTAssertEqual(weakGpuRes.lock(), nullptr);
+  XCTAssertNil(metalRes);
+}
+
+@end