💎

Metal-cpp のメモリ管理について

2024/08/30に公開

はじめに

Macやiphone なので GPU を使うために Metal に手を出してみた
メモリ管理方法が通常の c++ と異なっているためメモ代わりにまとめる

概要

Metal-cpp は Objective-C のWrapper でメモリ管理はそれを踏襲している.
ヘッダーを一つ include するだけで使用することができるため非常に便利.

以下のようにマクロを定義してあげないとコンパイル時にリンクエラーになる.
複数定義するとそれもエラーになるので, main.cpp などに一つだけ定義してあげると良い.

#define NS_PRIVATE_IMPLEMENTATION
#define CA_PRIVATE_IMPLEMENTATION
#define MTL_PRIVATE_IMPLEMENTATION
#include <Metal.hpp>

導入までは以下のサイトを参考
https://developer.apple.com/jp/metal/cpp/

Metal でのメモリ管理

基本的には std::shared_ptr のように参照カウンタを使用して参照されるたびにカウントが増え参照されなくなると減るそして0になると開放される. また標準の new, delete 関数は使用することができない.

Metal-cpp ライブラリにおいて以下の接頭辞がついている関数には所有権がある

  • alloc
  • new
  • copy
  • mutableCopy
  • create

もちろん所有権があるということは,
使用後 release 関数を使用してしっかりとメモリ解放してあげる必要がある.
また release したオブジェクトを再度 release すると Segmentation fault となる.
(以下の例が参照カウントの例)

メモリ解放の例
// この pool は次の例で使用する重要なもの
NS::AutoreleasePool *pool
    = NS::AutoreleasePool::alloc()->init();

MTL::Device *device_1
    = MTL::CreateSystemDefaultDevice(); // ここで参照カウントが 1

auto *device_2 = device_1->retain(); // 参照カウントが 2

// 何らかの処理

device_1->release(); // 参照カウントが 1
device_2->release(); // 参照カウントが 0

// device のカウントがすでに 0 なのでこれ以下は使えない

device_2->release(); // セグメンテーションフォルト
pool->release(); 

newなどの接頭辞がついたものについて必ず release する必要があるかについては要調査
release しておくのが安全

release する以外にもう一つ方法がありそれが autorelease

autorelease の使用例
MTL::Device *device_1
    = MTL::CreateSystemDefaultDevice(); // ここで参照カウントが 1

auto *device_2 = device_1->retain(); // 参照カウントが 2

// 何らかの処理

device_1->release(); // 参照カウントが 1

// 参照カウントが 1
// autorelease オブジェクトの所有権は AutoreleasePool に移動する
device_2->autorelease(); 

pool->release(); // カウントを減らしメモリすべて開放する

ちなみに以下の例のように先ほど上げた接頭辞がついていない関数(commandBuffer)の所有権は持っていないのrelease する必要がない. release するとこれもセグメンテーションフォルトになる.

MTL::CommandBuffer *command_buffer = m_command_queue->commandBuffer();

Metal 専用の SharedPtr

new などの接頭辞がついた関数を release する必要があるのがわかったがめんどくさいし忘れることもある.
これを解決するのが NS::SharedPtr .
所有権がないものを SharedPtr にいれるとこれもセグメンテーションフォルトになる.

使い方は以下のような感じ.
newがついた関数を NS::TransferPtr にわたすことで SharedPtr になる.

NS::SharedPtr<MTL::Buffer> buffer = NS::TransferPtr(
      device->newBuffer(
        buffer_length_ * sizeof(T),
        MTL::ResourceStorageModeShared
      )
    );

newなどがついてない関数を渡したいことがある.
例えば MTL::Device である.
retain 関数で 参照カウンタが増えたものの, release を忘れてカウンタが減らないことがあるため.
NS::RetainPtr() を使って所有権がない関数などの所有権を取得しつつ SharedPtr として取得することができる.
これでrelease 忘れがなくなる

MTL::Device *device_1
    = MTL::CreateSystemDefaultDevice();
NS::SharedPtr<MTL::Device> device_2 = NS::RetainPtr(device_1);

MTL::CommandBuffer *command_buffer_1 = m_command_queue->commandBuffer();
NS::SharedPtr<MTL::CommandBuffer> command_buffer_2 = NS::RetainPtr(command_buffer_1);

サンプルコード

ベクトルの要素ごとの和を求めるコード.
とりあえず autorelease しておけばとりあえず安心.

#define NS_PRIVATE_IMPLEMENTATION
#define CA_PRIVATE_IMPLEMENTATION
#define MTL_PRIVATE_IMPLEMENTATION
#include <Metal.hpp>
#include <iostream>


int main()
{
  NS::AutoreleasePool *pool
    = NS::AutoreleasePool::alloc()->init();

  // 先に autorelease をしておけばとりあえずメモリリークは防げる
  MTL::Device *device
    = MTL::CreateSystemDefaultDevice()->autorelease();
  NS::Error *error;

  // カーネル
  std::string kernel = "#include <metal_stdlib>\n"
                       "using namespace metal;\n"
                       "\n"
                       "kernel void add_arrays(device const float* inA,\n"
                       "                       device const float* inB,\n"
                       "                       device float* result,\n"
                       "                       uint index [[thread_position_in_grid]])\n"
                       "{\n"
                       "    result[index] = inA[index] + inB[index];\n"
                       "}";

  // カーネルを読み込みライブラリを作成
  // これも同様に autorelease しておけばメモリリークは防げる
  MTL::Library *library =
    device->newLibrary(NS::String::string(kernel.c_str(),
                                          NS::ASCIIStringEncoding),
                       nullptr,
                       &error)->autorelease();
  if (library == nullptr || error != nullptr)
  {
    throw std::runtime_error("Failed to create Metal library.");
  }


  // これの所有権はない
  auto function_name = NS::String::string("add_arrays", NS::ASCIIStringEncoding);

  // 指定した関数名の関数を取得. これも autorelease しておけばメモリリークは防げる
  auto add_function = library->newFunction(function_name)->autorelease();

  if (!add_function)
  {
    std::cerr << "Failed to find the adder function.";
  }

  auto *add_function_pso = device->newComputePipelineState(add_function, &error)->autorelease();
  auto *command_queue = device->newCommandQueue()->autorelease();

  // buffer を取得
  const int length = 10;
  const int size = length * sizeof(float);
  auto buffer_A = device->newBuffer(size, MTL::ResourceStorageModeShared)->autorelease();
  auto buffer_B = device->newBuffer(size, MTL::ResourceStorageModeShared)->autorelease();
  auto buffer_result = device->newBuffer(size, MTL::ResourceStorageModeShared)->autorelease();

  // buffer にデータを書き込む
  for (int i = 0; i < length; i++)
  {
    auto value = static_cast<float>(i);
    static_cast<float *>(buffer_A->contents())[i] = value;
    static_cast<float *>(buffer_B->contents())[i] = value;
  }

  // 以下の2つは所有権がない
  MTL::CommandBuffer *command_buffer = command_queue->commandBuffer();
  MTL::ComputeCommandEncoder *compute_encoder = command_buffer->computeCommandEncoder();
  
  compute_encoder->setComputePipelineState(add_function_pso);
  compute_encoder->setBuffer(buffer_A, 0, 0);
  compute_encoder->setBuffer(buffer_B, 0, 1);
  compute_encoder->setBuffer(buffer_result, 0, 2);

  MTL::Size grid_size = MTL::Size(array_length, 1, 1);

  NS::UInteger thread_group_size_ = add_function_pso->maxTotalThreadsPerThreadgroup();
  if (thread_group_size_ > array_length)
  {
    thread_group_size_ = array_length;
  }

  MTL::Size thread_group_size = MTL::Size(thread_group_size_, 1, 1);

  compute_encoder->dispatchThreads(grid_size, thread_group_size);
  compute_encoder->endEncoding();
  command_buffer->commit();
  command_buffer->waitUntilCompleted();

  for (int i = 0; i < length; i++)
  {
    std::cout << static_cast<float *>(buffer_result->contents())[i] << std::endl;
  }
  pool->release();

  return 0;

参考文献

https://developer.apple.com/jp/videos/play/wwdc2022/10160/

https://developer.apple.com/metal/cpp/

https://github.com/postmalloc/metal-cpp-example/tree/main

https://developer.apple.com/jp/metal/cpp/

Discussion