Unity でエディタ拡張向けの重い画像生成の改善をしてみた

はじめに

大分ニッチな話だと思うのですが、現在メンテしているプロジェクトで音声の解析結果を示すために画像生成を行い、それをエディタ拡張で表示して見せています。

これは Burst と Job を使って生成しているのですが、エディタ拡張の UI の体系は同期処理となるため、現状では Job を即時実行・回収することで、同期処理系となってしまっています。

tips.hecomi.com

今回はここを何とかしようと試みたお話を共有します。内容はテクスチャ生成ですが、それ以外にも重い処理を行うようなエディタ拡張のヒントになるかと思います。

コードと問題点

具体的に以下のようなコードを想定してみましょう。

AsyncTextureCreator.cs

using UnityEngine;

public class AsyncTextureCreator : MonoBehaviour
{
    public int width = 128;
    public int height = 32;
}

AsyncTextureCreatorEditor.cs

using UnityEngine;
using UnityEditor;
using Unity.Burst;
using Unity.Jobs;
using Unity.Collections;
using Unity.Mathematics;
using UnityEngine.Profiling;

// カラーバッファに適当な float バッファを元に色を書き込むジョブ
[BurstCompile]
internal struct CreateTextureJob : IJob
{
    [WriteOnly] public NativeArray<Color32> texColors;
    [ReadOnly] [DeallocateOnJobCompletion] public NativeArray<float> array;

    Color ToRGB(float hue)
    {
        hue = (1f - math.cos(math.PI * hue)) * 0.5f;
        hue = 1f - hue;
        hue *= 5f;
        var x = 1 - math.abs(hue % 2f - 1f);
        return
            hue < 1f ? new Color(1f, x, 0f) :
            hue < 2f ? new Color(x, 1f, 0f) :
            hue < 3f ? new Color(0f, 1f, x) :
            hue < 4f ? new Color(0f, x, 1f) :
            new Color(x * 0.5f, 0f, 0.5f);
    }

    public void Execute()
    {
        for (int i = 0; i < array.Length; ++i)
        {
            texColors[i] = ToRGB(array[i]);
        }
    }
}

[CustomEditor(typeof(AsyncTextureCreator))]
public class AsyncTextureCreatorEditor : Editor
{
    AsyncTextureCreator creator => target as AsyncTextureCreator;
    Texture2D _texture;
    CustomSampler _sampler = CustomSampler.Create("AsyncTextureCreatorEditor.UpdateTexture");
    
    void UpdateTexture()
    {
        _sampler.Begin();
        
        // テクスチャサイズが変更されたら Texture2D を作成
        int width = creator.width;
        int height = creator.height;
        if (!_texture || _texture.width != width || _texture.height != height)
        {
            _texture = new Texture2D(width, height);
        }

        // Texture2D のカラーバッファのポインタと、ジョブで色塗りする用のバッファを生成
        var texColors = _texture.GetPixelData<Color32>(0);
        var array = new NativeArray<float>(width * height, Allocator.TempJob);

        // バッファには適当な数値を詰めておく
        for (int y = 0; y < height; ++y)
        {
            for (int x = 0; x < width; ++x)
            {
                int i = y * width + x;
                array[i] = i / 256f;
            }
        }

        // ジョブを作成
        var job = new CreateTextureJob()
        {
            texColors = texColors,
            array = array,
        };

        // ジョブを実行・終了待ち
        job.Schedule().Complete();

        // GPU 上のテクスチャへと反映
        _texture.Apply();
        
        _sampler.End();
    }
    
    public override void OnInspectorGUI()
    {
        DrawDefaultInspector();
        EditorGUILayout.Space();
        UpdateTexture();
        var area = EditorGUILayout.GetControlRect(false, creator.height);
        GUI.DrawTexture(area, _texture, ScaleMode.StretchToFill);
    }
}

これを実行すると以下のような UI が表示されます。

UI

こちらのパフォーマンスは以下のような感じです。

Profiler

同期ジョブを実行しているのでジョブが重いときにかなり長い時間ブロッキングしてしまうことが分かります。そして何故か一つの InspectorWindow.Repaint から OnGUI が 3 回呼ばれており謎です。。

改善案

一つは同期ジョブを何とか非同期ジョブにできないかの検討、そしてもう一つは単一描画リクエスト中に複数回呼ばれてしまい、無駄な画像生成が走っているのを止める点です。

複数回呼び出しの抑制

まずは明らかに無駄である複数回呼び出しを何とかします。残念ながらなぜ単一の GUI 描画リクエストから複数回描画が走ってしまうのかは分からなかったですが、単一フレーム内に限らず、基本的には変更があったときのみ新しいテクスチャを生成・更新して、そうでないときはキャッシュしたものを再利用する形のほうが望ましいです。

ここでは横幅や縦幅が変更されたときに更新するようにしてみます。実際のユースケースでは、何かしらのフラグを通じて更新要求を監視するのが良いと思われます。

...
public class AsyncTextureCreatorEditor : Editor
{
    ...
    bool isTextureUpdateRequired => 
        !_texture || 
        _texture.width != creator.width || 
        _texture.height != creator.height;
    
    void UpdateTexture()
    {
        if (!isTextureUpdateRequired) return;

        int width = creator.width;
        int height = creator.height;
        _texture = new Texture2D(width, height);

        ...
    }
}

非同期ジョブ検討

当該フレームに得た情報を描画するのは、先にも述べたように同期 UI の体系なので基本的には難しいと思います。ただし 1 フレーム遅らせることが許容できるケースであれば比較的簡単そうです。戦略としては、更新要求があったら Job をスケジュールし、次の OnInspectorGUI の初めで回収する、というものです。ただ、UI に変化がない大半の時間は OnInspectorGUI は基本的には呼ばれません。そこで、無理やりもう一度 OnInspectorGUI が呼ばれるようにテクスチャを更新した際は、エディタに対して Repaint を要求します。これで次のフレームに再度 OnInspectorGUI が回ってくるのでそこで更新処理を行う、という形になります。

...
public class AsyncTextureCreatorEditor : Editor
{
    ...
    JobHandle _jobHandle;
    bool _isTextureUpdateRequested = false;
    bool isTextureUpdateRequired => ...;
    
    void UpdateTexture()
    {
        if (_isTextureUpdateRequested)
        {
            _jobHandle.Complete();
            _texture.Apply();
            _isTextureUpdateRequested = false;
        }

        if (!isTextureUpdateRequired) return;
        ...

        var job = new CreateTextureJob() { ... };
        _jobHandle = job.Schedule();
        _isTextureUpdateRequested = true;
        
        Repaint();
    }
    ...
}