Register GlTextureBuffer pool with GpuBuffer

First crack at hooking up pools with the GpuBufferStorage system. Will most likely be superseded later, but for now this works with minimal code impact: just overwrite the factory for a storage type with one that uses the pool.

PiperOrigin-RevId: 488783854

mediapipe/gpu/gpu_buffer_storage.h

@@ -74,13 +74,17 @@ class GpuBufferStorageRegistry {
 
   template <class Storage>
   RegistryToken Register() {
-    return Register(
+    return RegisterFactory<Storage>(
         [](int width, int height,
            GpuBufferFormat format) -> std::shared_ptr<Storage> {
           return CreateStorage<Storage>(overload_priority<10>{}, width, height,
                                         format);
-        },
+        });
-        Storage::GetProviderTypes());
+  }
+
+  template <class Storage, class F>
+  RegistryToken RegisterFactory(F&& factory) {
+    return Register(factory, Storage::GetProviderTypes());
   }
 
   template <class StorageFrom, class StorageTo, class F>
@@ -148,6 +152,13 @@ class GpuBufferStorageImpl : public GpuBufferStorage, public U... {
     return kHashes;
   }
 
+  // Exposing this as a function allows dependent initializers to call this to
+  // ensure proper ordering.
+  static GpuBufferStorageRegistry::RegistryToken RegisterOnce() {
+    static auto registration = GpuBufferStorageRegistry::Get().Register<T>();
+    return registration;
+  }
+
  private:
   virtual const void* down_cast(TypeId to) const override {
     return down_cast_impl(to, types<T, U...>{});
@@ -161,8 +172,7 @@ class GpuBufferStorageImpl : public GpuBufferStorage, public U... {
     return down_cast_impl(to, types<W...>{});
   }
 
-  inline static auto registration =
+  inline static auto registration = RegisterOnce();
-      GpuBufferStorageRegistry::Get().Register<T>();
   using RequireStatics = ForceStaticInstantiation<&registration>;
 };
 

mediapipe/gpu/gpu_shared_data_internal.cc

@@ -200,4 +200,34 @@ GpuSharedData::GpuSharedData() : GpuSharedData(kPlatformGlContextNone) {}
 MPPGraphGPUData* GpuResources::ios_gpu_data() { return ios_gpu_data_; }
 #endif  // __APPLE__
 
+extern const GraphService<GpuResources> kGpuService;
+
+#if !MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+static std::shared_ptr<GlTextureBuffer> GetGlTextureBufferFromPool(
+    int width, int height, GpuBufferFormat format) {
+  std::shared_ptr<GlTextureBuffer> texture_buffer;
+  const auto cc = LegacyCalculatorSupport::Scoped<CalculatorContext>::current();
+
+  if (cc && cc->Service(kGpuService).IsAvailable()) {
+    GpuBufferMultiPool* pool =
+        &cc->Service(kGpuService).GetObject().gpu_buffer_pool();
+    // Note that the "gpu_buffer_pool" serves GlTextureBuffers on non-Apple
+    // platforms. TODO: refactor into storage pools.
+    texture_buffer = pool->GetBuffer(width, height, format)
+                         .internal_storage<GlTextureBuffer>();
+  } else {
+    texture_buffer = GlTextureBuffer::Create(width, height, format);
+  }
+  return texture_buffer;
+}
+
+static auto kGlTextureBufferPoolRegistration = [] {
+  // Ensure that the GlTextureBuffer's own factory is already registered, so we
+  // can override it.
+  GlTextureBuffer::RegisterOnce();
+  return internal::GpuBufferStorageRegistry::Get()
+      .RegisterFactory<GlTextureBuffer>(GetGlTextureBufferFromPool);
+}();
+#endif  // !MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
+
 }  // namespace mediapipe