[Impeller] track the sizes of all outstanding MTLTexture allocations and report per frame in MB, matching Vulkan implementation. (#53618)

Does what it says!

Fixes flutter/flutter#150936

Author: jonahwilliams

ci/licenses_golden/excluded_files

@@ -176,6 +176,7 @@
 ../../../flutter/impeller/golden_tests/README.md
 ../../../flutter/impeller/playground
 ../../../flutter/impeller/renderer/backend/gles/test
+../../../flutter/impeller/renderer/backend/metal/allocator_mtl_unittests.mm
 ../../../flutter/impeller/renderer/backend/metal/texture_mtl_unittests.mm
 ../../../flutter/impeller/renderer/backend/vulkan/allocator_vk_unittests.cc
 ../../../flutter/impeller/renderer/backend/vulkan/command_encoder_vk_unittests.cc

impeller/core/allocator.h

@@ -47,9 +47,12 @@ class Allocator {
   /// @brief Write debug memory usage information to the dart timeline in debug
   ///        and profile modes.
   ///
-  ///        This is only supported on the Vulkan backend.
+  ///        This is supported on both the Metal and Vulkan backends.
   virtual void DebugTraceMemoryStatistics() const {};
 
+  // Visible for testing.
+  virtual size_t DebugGetHeapUsage() const { return 0; }
+
  protected:
   Allocator();
 

impeller/core/texture_descriptor.h

@@ -5,6 +5,7 @@
 #ifndef FLUTTER_IMPELLER_CORE_TEXTURE_DESCRIPTOR_H_
 #define FLUTTER_IMPELLER_CORE_TEXTURE_DESCRIPTOR_H_
 
+#include <cstdint>
 #include "impeller/core/formats.h"
 #include "impeller/geometry/size.h"
 
@@ -51,6 +52,22 @@ struct TextureDescriptor {
     return size.Area() * BytesPerPixelForPixelFormat(format);
   }
 
+  constexpr size_t GetByteSizeOfAllMipLevels() const {
+    if (!IsValid()) {
+      return 0u;
+    }
+    size_t result = 0u;
+    int64_t width = size.width;
+    int64_t height = size.height;
+    for (auto i = 0u; i < mip_count; i++) {
+      result +=
+          ISize(width, height).Area() * BytesPerPixelForPixelFormat(format);
+      width /= 2;
+      height /= 2;
+    }
+    return result;
+  }
+
   constexpr size_t GetBytesPerRow() const {
     if (!IsValid()) {
       return 0u;

impeller/playground/playground_test.h

@@ -42,7 +42,7 @@ class PlaygroundTest : public Playground,
 
  private:
   // |Playground|
-  bool ShouldKeepRendering() const;
+  bool ShouldKeepRendering() const override;
 
 #if FML_OS_MACOSX
   fml::ScopedNSAutoreleasePool autorelease_pool_;

impeller/renderer/backend/metal/BUILD.gn

@@ -65,10 +65,14 @@ impeller_component("metal") {
 impeller_component("metal_unittests") {
   testonly = true
 
-  sources = [ "texture_mtl_unittests.mm" ]
+  sources = [
+    "allocator_mtl_unittests.mm",
+    "texture_mtl_unittests.mm",
+  ]
 
   deps = [
     ":metal",
+    "//flutter/impeller/playground:playground_test",
     "//flutter/testing:testing_lib",
   ]
 

impeller/renderer/backend/metal/allocator_mtl.h

@@ -6,18 +6,42 @@
 #define FLUTTER_IMPELLER_RENDERER_BACKEND_METAL_ALLOCATOR_MTL_H_
 
 #include <Metal/Metal.h>
+#include <atomic>
 
+#include "impeller/base/thread.h"
 #include "impeller/core/allocator.h"
 
 namespace impeller {
 
+class DebugAllocatorStats {
+ public:
+  DebugAllocatorStats() {}
+
+  ~DebugAllocatorStats() {}
+
+  /// Increment the tracked allocation size in bytes.
+  void Increment(size_t size);
+
+  /// Decrement the tracked allocation size in bytes.
+  void Decrement(size_t size);
+
+  /// Get the current tracked allocation size in MB.
+  size_t GetAllocationSizeMB();
+
+ private:
+  std::atomic<size_t> size_ = 0;
+};
+
 class AllocatorMTL final : public Allocator {
  public:
   AllocatorMTL();
 
   // |Allocator|
   ~AllocatorMTL() override;
 
+  // |Allocator|
+  size_t DebugGetHeapUsage() const override;
+
  private:
   friend class ContextMTL;
 
@@ -26,6 +50,12 @@ class AllocatorMTL final : public Allocator {
   bool supports_memoryless_targets_ = false;
   bool supports_uma_ = false;
   bool is_valid_ = false;
+
+#ifdef IMPELLER_DEBUG
+  std::shared_ptr<DebugAllocatorStats> debug_allocater_ =
+      std::make_shared<DebugAllocatorStats>();
+#endif  // IMPELLER_DEBUG
+
   ISize max_texture_supported_;
 
   AllocatorMTL(id<MTLDevice> device, std::string label);
@@ -47,6 +77,9 @@ class AllocatorMTL final : public Allocator {
   // |Allocator|
   ISize GetMaxTextureSizeSupported() const override;
 
+  // |Allocator|
+  void DebugTraceMemoryStatistics() const override;
+
   AllocatorMTL(const AllocatorMTL&) = delete;
 
   AllocatorMTL& operator=(const AllocatorMTL&) = delete;

impeller/renderer/backend/metal/allocator_mtl.mm

@@ -6,7 +6,9 @@
 
 #include "flutter/fml/build_config.h"
 #include "flutter/fml/logging.h"
+#include "fml/trace_event.h"
 #include "impeller/base/validation.h"
+#include "impeller/core/formats.h"
 #include "impeller/renderer/backend/metal/device_buffer_mtl.h"
 #include "impeller/renderer/backend/metal/formats_mtl.h"
 #include "impeller/renderer/backend/metal/texture_mtl.h"
@@ -88,6 +90,19 @@ static bool SupportsLossyTextureCompression(id<MTLDevice> device) {
 #endif
 }
 
+void DebugAllocatorStats::Increment(size_t size) {
+  size_.fetch_add(size, std::memory_order_relaxed);
+}
+
+void DebugAllocatorStats::Decrement(size_t size) {
+  size_.fetch_sub(size, std::memory_order_relaxed);
+}
+
+size_t DebugAllocatorStats::GetAllocationSizeMB() {
+  size_t new_value = size_ / 1000000;
+  return new_value;
+}
+
 AllocatorMTL::AllocatorMTL(id<MTLDevice> device, std::string label)
     : device_(device), allocator_label_(std::move(label)) {
   if (!device_) {
@@ -212,11 +227,23 @@ static MTLStorageMode ToMTLStorageMode(StorageMode mode,
     }
   }
 
+#ifdef IMPELLER_DEBUG
+  if (desc.storage_mode != StorageMode::kDeviceTransient) {
+    debug_allocater_->Increment(desc.GetByteSizeOfAllMipLevels());
+  }
+#endif  // IMPELLER_DEBUG
+
   auto texture = [device_ newTextureWithDescriptor:mtl_texture_desc];
   if (!texture) {
     return nullptr;
   }
-  return TextureMTL::Create(desc, texture);
+  std::shared_ptr<TextureMTL> result_texture =
+      TextureMTL::Create(desc, texture);
+#ifdef IMPELLER_DEBUG
+  result_texture->SetDebugAllocator(debug_allocater_);
+#endif  // IMPELLER_DEBUG
+
+  return result_texture;
 }
 
 uint16_t AllocatorMTL::MinimumBytesPerRow(PixelFormat format) const {
@@ -228,4 +255,21 @@ static MTLStorageMode ToMTLStorageMode(StorageMode mode,
   return max_texture_supported_;
 }
 
+size_t AllocatorMTL::DebugGetHeapUsage() const {
+#ifdef IMPELLER_DEBUG
+  return debug_allocater_->GetAllocationSizeMB();
+#else
+  return 0u;
+#endif  // IMPELLER_DEBUG
+}
+
+void AllocatorMTL::DebugTraceMemoryStatistics() const {
+#ifdef IMPELLER_DEBUG
+  size_t allocated_size = DebugGetHeapUsage();
+  FML_TRACE_COUNTER("flutter", "AllocatorMTL",
+                    reinterpret_cast<int64_t>(this),  // Trace Counter ID
+                    "MemoryBudgetUsageMB", allocated_size);
+#endif  // IMPELLER_DEBUG
+}
+
 }  // namespace impeller

impeller/renderer/backend/metal/allocator_mtl_unittests.mm

@@ -0,0 +1,70 @@
+// Copyright 2013 The Flutter Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "flutter/testing/testing.h"
+#include "impeller/core/formats.h"
+#include "impeller/core/texture_descriptor.h"
+#include "impeller/playground/playground_test.h"
+#include "impeller/renderer/backend/metal/allocator_mtl.h"
+#include "impeller/renderer/backend/metal/context_mtl.h"
+#include "impeller/renderer/backend/metal/formats_mtl.h"
+#include "impeller/renderer/backend/metal/lazy_drawable_holder.h"
+#include "impeller/renderer/backend/metal/texture_mtl.h"
+#include "impeller/renderer/capabilities.h"
+
+#include <QuartzCore/CAMetalLayer.h>
+#include <memory>
+#include <thread>
+
+#include "gtest/gtest.h"
+
+namespace impeller {
+namespace testing {
+
+using AllocatorMTLTest = PlaygroundTest;
+INSTANTIATE_METAL_PLAYGROUND_SUITE(AllocatorMTLTest);
+
+TEST_P(AllocatorMTLTest, DebugTraceMemoryStatistics) {
+  auto& context_mtl = ContextMTL::Cast(*GetContext());
+  const auto& allocator = context_mtl.GetResourceAllocator();
+
+  EXPECT_EQ(allocator->DebugGetHeapUsage(), 0u);
+
+  // Memoryless texture does not increase allocated size.
+  {
+    TextureDescriptor desc;
+    desc.format = PixelFormat::kR8G8B8A8UNormInt;
+    desc.storage_mode = StorageMode::kDeviceTransient;
+    desc.size = {1000, 1000};
+    auto texture_1 = allocator->CreateTexture(desc);
+
+    EXPECT_EQ(allocator->DebugGetHeapUsage(), 0u);
+
+    // Private storage texture increases allocated size.
+    desc.storage_mode = StorageMode::kDevicePrivate;
+    auto texture_2 = allocator->CreateTexture(desc);
+
+#ifdef IMPELLER_DEBUG
+    EXPECT_EQ(allocator->DebugGetHeapUsage(), 4u);
+#else
+    EXPECT_EQ(allocator->DebugGetHeapUsage(), 0u);
+#endif  // IMPELLER_DEBUG
+
+    // Host storage texture increases allocated size.
+    desc.storage_mode = StorageMode::kHostVisible;
+    auto texture_3 = allocator->CreateTexture(desc);
+
+#ifdef IMPELLER_DEBUG
+    EXPECT_EQ(allocator->DebugGetHeapUsage(), 8u);
+#else
+    EXPECT_EQ(allocator->DebugGetHeapUsage(), 0u);
+#endif  // IMPELLER_DEBUG
+  }
+
+  // After all textures are out of scope, memory has been decremented.
+  EXPECT_EQ(allocator->DebugGetHeapUsage(), 0u);
+}
+
+}  // namespace testing
+}  // namespace impeller

impeller/renderer/backend/metal/surface_mtl.mm

@@ -229,6 +229,10 @@ - (void)flutterPrepareForPresent:(nonnull id<MTLCommandBuffer>)commandBuffer;
     return false;
   }
 
+#ifdef IMPELLER_DEBUG
+  context->GetResourceAllocator()->DebugTraceMemoryStatistics();
+#endif  // IMPELLER_DEBUG
+
   if (requires_blit_) {
     if (!(source_texture_ && destination_texture_)) {
       return false;

impeller/renderer/backend/metal/texture_mtl.h

@@ -9,6 +9,7 @@
 
 #include "impeller/base/backend_cast.h"
 #include "impeller/core/texture.h"
+#include "impeller/renderer/backend/metal/allocator_mtl.h"
 
 namespace impeller {
 
@@ -47,7 +48,16 @@ class TextureMTL final : public Texture,
 
   bool GenerateMipmap(id<MTLBlitCommandEncoder> encoder);
 
+#ifdef IMPELLER_DEBUG
+  void SetDebugAllocator(
+      const std::shared_ptr<DebugAllocatorStats>& debug_allocator);
+#endif  // IMPELLER_DEBUG
+
  private:
+#ifdef IMPELLER_DEBUG
+  std::shared_ptr<DebugAllocatorStats> debug_allocator_ = nullptr;
+#endif  // IMPELLER_DEBUG
+
   AcquireTextureProc aquire_proc_ = {};
   bool is_valid_ = false;
   bool is_wrapped_ = false;

impeller/renderer/backend/metal/texture_mtl.mm

@@ -6,6 +6,7 @@
 #include <memory>
 
 #include "impeller/base/validation.h"
+#include "impeller/core/formats.h"
 #include "impeller/core/texture_descriptor.h"
 
 namespace impeller {
@@ -59,7 +60,17 @@ new TextureMTL(
   return std::make_shared<TextureMTL>(desc, [texture]() { return texture; });
 }
 
-TextureMTL::~TextureMTL() = default;
+TextureMTL::~TextureMTL() {
+#ifdef IMPELLER_DEBUG
+  if (debug_allocator_) {
+    auto desc = GetTextureDescriptor();
+    if (desc.storage_mode == StorageMode::kDeviceTransient) {
+      return;
+    }
+    debug_allocator_->Decrement(desc.GetByteSizeOfBaseMipLevel());
+  }
+#endif  // IMPELLER_DEBUG
+}
 
 void TextureMTL::SetLabel(std::string_view label) {
   if (is_drawable_) {
@@ -76,6 +87,13 @@ new TextureMTL(
   return OnSetContents(mapping->GetMapping(), mapping->GetSize(), slice);
 }
 
+#ifdef IMPELLER_DEBUG
+void TextureMTL::SetDebugAllocator(
+    const std::shared_ptr<DebugAllocatorStats>& debug_allocator) {
+  debug_allocator_ = debug_allocator;
+}
+#endif  // IMPELLER_DEBUG
+
 // |Texture|
 bool TextureMTL::OnSetContents(const uint8_t* contents,
                                size_t length,

impeller/renderer/backend/vulkan/allocator_vk.h

@@ -21,8 +21,8 @@ class AllocatorVK final : public Allocator {
   // |Allocator|
   ~AllocatorVK() override;
 
-  // Visible for testing
-  size_t DebugGetHeapUsage() const;
+  // |Allocator|
+  size_t DebugGetHeapUsage() const override;
 
   /// @brief Select a matching memory type for the given
   ///        [memory_type_bits_requirement], or -1 if none is found.