【Unity】Burstコードからマネージドコードを呼ぶ

Unity Loggingのソースを見ていたら、Burst2ManagedCallというクラスを見かけました。

https://github.com/needle-mirror/com.unity.logging/blob/6adb6404216ee7b4f87668c45984b11fc63de525/Runtime/Burst2ManagedCall.cs#L117

Burstコンパイル対象のC#コードは事前にネイティブコードにコンパイルされるため、その中では様々な制限がかかります。特に参照型が使用できないのは難しいです。また、Burstコンパイルされたコードからそのまま別のメソッドを呼び出すと、その内容もBurstコンパイルされて実行されるため、同様の制限が適用されます。

https://docs.unity3d.com/ja/Packages/com.unity.burst@1.8/manual/csharp-language-support.html

Burst2ManagedCallは名前の通り、Burstコンパイルされたコードからマネージドコードを呼び出すためのユーティリティです。以下はそのままUnity Loggingからの抜粋です。

using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
using Unity.Burst;

internal static class Burst2ManagedCall<T, Key>
{
    private static T s_Delegate;
    private static readonly SharedStatic<FunctionPointer<T>> s_SharedStatic = SharedStatic<FunctionPointer<T>>.GetOrCreate<FunctionPointer<T>, Key>(16);
    public static bool IsCreated => s_SharedStatic.Data.IsCreated;

    public static void Init(T @delegate)
    {
        CheckIsNotCreated();
        s_Delegate = @delegate;
        s_SharedStatic.Data = new FunctionPointer<T>(Marshal.GetFunctionPointerForDelegate(s_Delegate));
    }

    public static ref FunctionPointer<T> Ptr()
    {
        CheckIsCreated();
        return ref s_SharedStatic.Data;
    }

    [Conditional("ENABLE_UNITY_COLLECTIONS_CHECKS"), Conditional("UNITY_DOTS_DEBUG")] // ENABLE_UNITY_COLLECTIONS_CHECKS or UNITY_DOTS_DEBUG
    private static void CheckIsCreated()
    {
        if (IsCreated == false)
            throw new InvalidOperationException("Burst2ManagedCall was NOT created!");
    }

    [Conditional("ENABLE_UNITY_COLLECTIONS_CHECKS"), Conditional("UNITY_DOTS_DEBUG")] // ENABLE_UNITY_COLLECTIONS_CHECKS or UNITY_DOTS_DEBUG
    private static void CheckIsNotCreated()
    {
        if (IsCreated)
            throw new InvalidOperationException("Burst2ManagedCall was already created!");
    }
}

使い方

あらかじめ、呼び出す対象のメソッドをデリゲートとしてBurst2ManagedCallにセットしておきます。

public static class Burst2ManagedCallTest
{
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
    private delegate void MyManagedDelegate(int param);

    private struct MyManagedDelegateKey
    {
    }

    [AOT.MonoPInvokeCallback(typeof(MyManagedDelegate))]
    private static void MyManagedFunc(int param)
    {
        // 参照型もバリバリ使える
        // JobSystem経由だとワーカースレッド上で実行されるので注意
        new GameObject().transform.position = new Vector3(param, 0, 0);
    }

    private static void RegisterMyManagedFunc()
    {
        Burst2ManagedCall<MyManagedDelegate, MyManagedDelegateKey>.Init(MyManagedFunc);
    }
}

次に、Burstコードから呼び出します。

[BurstCompile]
private static void BurstFunc()
{
    Burst2ManagedCall<MyManagedDelegate, MyManagedDelegateKey>.Ptr().Invoke(100);
}

しくみを理解する

Burst2ManagedCall<T, TKey>にはふたつの型引数があり、Tには呼び出す対象のデリゲート型を、TKeyには渡したメソッドを一意に定めるためのキーとなる型を指定します。

TKeyのように、型引数を静的なデータのキーとして利用するテクニックはstatic type cachingと呼ばれており、System.TypeをキーとするDictionaryを利用するよりもパフォーマンスがよいことが知られています。

Burst2ManagedCall<T, TKey>.Init()

private static T s_Delegate;
private static readonly SharedStatic<FunctionPointer<T>> s_SharedStatic = SharedStatic<FunctionPointer<T>>.GetOrCreate<FunctionPointer<T>, Key>(16);

public static void Init(T @delegate)
{
    CheckIsNotCreated();
    s_Delegate = @delegate;
    s_SharedStatic.Data = new FunctionPointer<T>(Marshal.GetFunctionPointerForDelegate(s_Delegate));
}

Burst2ManagedCall<T, TKey>.Init()にデリゲートを渡すと、そのデリゲートはMarshal.GetFunctionPointerForDelegate()に渡されます。

アンマネージ コードから呼び出すことができる関数ポインターにデリゲートを変換します。

マネージド コードからガベージ コレクターによってデリゲートが収集されないように手動で保持する必要があります。 ガベージ コレクターは、アンマネージ コードへの参照を追跡しません。

https://learn.microsoft.com/ja-jp/dotnet/api/system.runtime.interopservices.marshal.getfunctionpointerfordelegate?view=net-8.0

上記の説明の通り、デリゲートがGCに回収されないようs_Delegateフィールドにデリゲート自体を保持しています。Burst2ManagedCallはこれらの操作をラップして安全性を担保する役割も果たしているようです。

また、取得された関数ポインタはFunctionPointer<T>構造体にラップされます。これはBurstで関数ポインタを呼び出すためのラッパーのようです。

https://docs.unity3d.com/ja/Packages/com.unity.burst@1.8/manual/csharp-function-pointers.html

private static readonly SharedStatic<FunctionPointer<T>> s_SharedStatic = SharedStatic<FunctionPointer<T>>.GetOrCreate<FunctionPointer<T>, Key>(16);

s_SharedStatic.Data = new FunctionPointer<T>(Marshal.GetFunctionPointerForDelegate(s_Delegate));

FunctionPointer<T>はSharedStatic<FunctionPointer<T>>型のフィールドに保持されます。SharedStatic<T>はマネージドとBurst間でミュータブル (コンパイル時に決定できない) データをやり取りするためのユーティリティです。詳しくはBurstのドキュメントに記載されています。

https://docs.unity3d.com/ja/Packages/com.unity.burst@1.8/manual/csharp-shared-static.html

Burst2ManagedCall<T, TKey>.Ptr()

public static ref FunctionPointer<T> Ptr()
{
    CheckIsCreated();
    return ref s_SharedStatic.Data;
}

登録したFunctionPointer<T>をBurstから取得します。

FunctionPointer<T>.Invoke()

public T Invoke
{
    get
    {
        CheckIsCreated();
        return Marshal.GetDelegateForFunctionPointer<T>(_ptr);
    }
}

Marshal.GetDelegateForFunctionPointer()で関数ポインタをデリゲートに戻す処理が書かれています。ただ、おそらくこれは非Burstから呼んだ場合の処理で、Burstでは直接ポインタを呼ぶコードに変換されているのではないかと思います。

以下はBurst Inspectorの当該呼び出し部分（x64）です。ちょっと自信ないですが、s_SharedStaticの内容と思しきポインタに直接jmpしてる（よね……？）

mov        rax, qword ptr [rip + "Unity.Burst.SharedStatic`1<Unity.Burst.FunctionPointer`1<Burst2ManagedCallTest/MyManagedDelegate>> Burst2ManagedCall`2<Burst2ManagedCallTest/MyManagedDelegate,Burst2ManagedCallTest/MyManagedDelegateKey>::s_SharedStatic"]
.Ltmp1:
#DEBUG_VALUE: Unity.Burst.SharedStatic`1<Unity.Burst.FunctionPointer`1<Burst2ManagedCallTest.MyManagedDelegate>>.get_Data:this <- undef
.cv_loc        2 1 72 0                        # Burst2ManagedCallTest.cs:72:0
mov        ecx, 100
pop        rbp
rex64 jmp        qword ptr [rax]         # TAILCALL

おわり

Unity Loggingはログ出力処理（Sink）のカスタマイズを行えるのですが、Unity Loggingのログ出力は全体をBurst + Job Systemでワーカースレッド上で実行することでメインスレッドへの負荷を抑える設計になっているため、Sinkの実装もBurstコンパイルが必須になっています。Burst2ManagedCallみたいなユーティリティが必要とされるのも納得です。

マーシャリングの扱いとかがどうなっているかとか、マネージドコード - Burstコード間で行ったり来たりしすぎるとパフォーマンスに悪影響があるんじゃないかなどなど疑問はいろいろありますが、とりあえずマネージドコードを呼べたのでよしとします。

参考

https://adarapata.hatenablog.com/entry/2024/01/23/010658