e.blog

主にUnity/UE周りのことについてまとめていきます

C# Job System + Burst compilerを使ってPNG画像の展開を最適化してみる

C# Unity Job System Burst

概要

前回の記事、前々回の記事でPNG画像について書きました。

前回の記事で、最適化について書けたら書きますと書いていたのですが、C# Job SystemとBurst compilerを利用して最適化してみたのでそれをメモがてら書いておきたいと思います。

ちなみに、UnityのAPIである Texture2D.LoadImage(byte[] data); に比べると解凍部分の時間以外は大体同じくらいの時間で展開できているみたいです。（ただ、解凍処理が重くて合計すると倍くらいの時間かかってしまっていますが・・）

PNG画像展開の実装についてはGitHubにあげてあるので、実際に動くものを見たい方はそちらをご覧ください。

github.com

実装したものをAndroidの実機で動かしたデモです↓

C# Job Systemとは

C# Job System自体の解説記事ではないのでここでは簡単な実装方法の解説にとどめます。詳細については以下の記事がとても分かりやすくまとめてくれているのでそちらをご覧ください。 tsubakit1.hateblo.jp

概要についてドキュメントから引用すると、

Unity C# Job System を利用すると、Unity とうまく相互作用する マルチスレッドコード を書くことができ、正しいコードを書くことを容易にします。

マルチスレッドでコードを書くと、高いパフォーマンスを得ることができます。これらには、フレームレートの大幅な向上が含まれます。C#ジョブで Burst コンパイラーを使用すると、改良された コード生成 (英語)の品質が提供され、モバイルデバイスのバッテリー消費を大幅に削減します。

C# Job System の本質的な性質は Unity が内部で使用するもの (Unity のネイティブジョブシステム) との統合性です。ユーザーが作成したコードと Unity は ワーカースレッド を共有します。この連携により、CPU コア より多くのスレッドを作成すること (CPU リソースの競合の原因となります) を回避できます。

マルチスレッドプログラミングでは、競合の問題やスレッドを立ち上げることのコスト、スレッドの立ち上げすぎに寄るコンテキストスイッチのコストなどが問題になることがあります。しかしC# Job Systemを利用すると、Unityが起動しているスレッドの空いている時間を使って効率よく処理を行うことができるため低コストかつ安全にコードを書くことを可能にしてくれます。

加えて、後述するBurst Compilerを併用することでさらに高速化を望むことができます。

ただ、メインスレッド外で動作するためUnity APIを使うことができないという制約はそのままです。ただし Transform に限定して、それを操作する方法が提供されています。（今回は主題ではないので詳細は割愛します）

実装方法

C# Job Systemを利用するためには必要なインターフェースを実装し、それをスレッドで実行されるようにスケジュールする必要があります。

簡単な利用方法について以下の記事を参考にさせていただきました。 gametukurikata.com

IJobを実装する

まず1番シンプルな方法は IJob インターフェースを実装することです。

using System.Collections;
using System.Collections.Generic;
using Unity.Collections;
using Unity.Jobs;
using UnityEngine;

// 必ず struct にする必要がある
public struct AnyJob : IJob
{
    // 計算に利用する値
    public float value;

    // 結果を返すNativeArray
    public NativeArray<float> result;

    public void Execute()
    {
        for (int i = 0; i < 100; ++i)
        {
            result[0] += value;
        }
    }
}

※ Burst compilerで利用する場合、 class は使えないため必ず struct で定義する必要があります。

ここで、なぜ NativeArray を、しかも1要素のものを利用しているかというと、これはジョブシステムの安全性に起因するものです。

参考にした記事から引用すると、

C# Job Systemの安全機能によってジョブの結果が共有出来ない為、Nativeコンテナと呼ばれる共有メモリを使って結果を保存します。

と書かれています。そしてこのNativeコンテナはいくつか種類があり、以下のようなものが標準で用意されています。

  • NativeArray
  • NativeList
  • NativeHashMap
  • NativeMultiHashMap
  • NativeQueue

また特殊ケースに対応するため、コンテナは自身で定義、作成することもできるようになっています。

並列処理用のJob（IJobParallelFor）を実装する

上記のジョブはひとつのタスクをワーカースレッドで実行するものでした。スレッドで処理したいものの中には並列実行したいものも多数存在します。（それこそ今回はまさにこちらを利用しました）

その場合は IJobParallelFor インターフェースを実装します。基本的な使い方は IJob と変わりありませんが、実装すべきメソッドが少しだけ異なります。

// IJobParallelForのメソッド
void Execute(int index);

違いは引数に int 型の index を受け取るところです。並列実行されるため、そのジョブが「今何番目の位置の処理をすべきか」を index で知ることができるわけですね。

なのでデータを並列化可能な状態で用意したあと、この index を頼りに必要なデータを取り出し、計算を行って結果を返す処理を書いていくことになります。

Jobを実行する

Jobが実装できたら、次はこれを実行する方法を見てみましょう。

using System.Collections;
using System.Collections.Generic;
using Unity.Collections;
using Unity.Jobs;
using UnityEngine;

public class JobExecutor : MonoBehaviour
{
    private void Start()
    {
        NativeArray<float> resultArray = new NativeArray<float>(1, Allocator.TempJob);

        AnyJob job = new AnyJob
        {
            value = 0.5f,
            result = resultArray,
        };

        // ジョブをスケジュールすると
        // ジョブシステムによって空いているワーカースレッドで実行される
        JobHandle handle = job.Schedule();

        // ジョブの完了を待つ
        handle.Complete();

        // 結果を取り出す
        float result = resultArray[0];

        Debug.Log($"Result: {reuslt}");

        resultArray.Dispose();
    }
}

大事な点としては、結果を受け取るための NativeArray<T> を定義し、メモリをアンマネージド領域に確保します。そしてそれをジョブに設定しスケジュールします。

ジョブはそのまま実行するのではなく、空いているワーカースレッドで実行されるようにスケジュールする必要があります。

なお、ジョブに対して入力が必要な場合は同様に NativeArray<T> を利用してデータを渡す必要があります。

Jobを直列につなげる

場合によっては「このジョブが終わったあとに次のジョブを実行する」という形で、前のジョブに依存するような処理もあります。（今回の実装でも、並列化可能な部分は並列化したジョブで処理を行い、その後の処理はこの並列化した処理に依存した形で進んでいきます）

直列化させるのは簡単で、スケジュールする際に 前に終わっていたほしいジョブのハンドル を引数に渡します。

PreviousJob previousJob = new PreviousJob { ... };
NextJob nextJob = new NextJob { ... };

JobHandle previousHandle = preivousJob.Schedule();
JobHandle nextHandle = nextJob.Schedule(previousHandle);

Burst Compilerとは

Unity Blogから引用すると以下のように書かれています。

Burst は、新しいデータ指向技術スタック（DOTS）とUnity Job System を使って作成された、Unity プロジェクトのパフォーマンスを向上させるために使用できる事前コンパイラー技術です。Burst は、高性能 C#（HPC#）として知られる C# 言語のサブセットをコンパイルすることで動作し、LLVM コンパイラーフレームワークの上に構築された高度な最適化を展開することで、デバイスのパワーを効率的に利用します。

また、ドキュメントから引用すると、

Burst is a compiler, it translates from IL/.NET bytecode to highly optimized native code using LLVM. It is released as a unity package and integrated into Unity using the Unity Package Manager.

と書かれています。以下の動画でも言及がありますが、通常はC#はIL（Intermediate Language）に変換されます。しかし、Burst CompilerではIR（Intermediate Representation）に変換し、LLVMでさらに機械語に変換される、という手順を踏みます。 （動画ではLLVMを「人類の英知の結晶」と呼んでいましたｗ）

ちょっとまだしっかりと理解したわけではないですが、大きく最適化を施したコードが生成される、と思っておけばいいでしょう。ただ、無償でそうした最適化が手に入るわけではなく、それなりの制約を課したコードを書く必要がある点に注意が必要です。

www.youtube.com

Burst Compilerを適用するには BurstCompile アトリビュートを付与するだけです。（後述のコードを参照ください）

なお、2022.02.25時点で Package Managerから検索してインストール ができないようです。インストールについてはこちらの記事を参考にさせていただきました。

ざっくりとした解説は以上です。以下から、実際のコードを元に解説をしていきます。

並列化して高速化

前回の記事でPNGデータの展開処理について書きました。しかし前回の実装ではすべてのピクセルを順次処理していました。しかし、展開の仕組みを理解すると、Filter Type 1の場合、つまり左のピクセルにのみ依存している場合はその行だけの処理で完結します。言い換えるとすべての行は並列に処理できるということです。

ということで、Filter Type 1の行を抜き出してそこを並列化してみます。並列化には前述したC# Job Systemを採用しました。（Burst compilerによる高速化も見込めるかなと思ったので）

Filter Type 1の行を抜き出す

まずは該当の行を抜き出します。

// Filter Type 1の行とそれ以外の行を抜き出す
LineInfo info = PngParser.ExtractLines(data, _metaData);

// -----------------

// 抜き出し処理
// シンプルに行の頭のFilter Typeを判定しているだけ
public static LineInfo ExtractLines(byte[] data, PngMetaData metaData)
{
    List<int> type1 = new List<int>();
    List<int> other = new List<int>();

    for (int h = 0; h < metaData.height; ++h)
    {
        int idx = metaData.rowSize * h;
        byte filterType = data[idx];

        if (filterType == 1)
        {
            type1.Add(h);
        }
        else
        {
            other.Add(h);
        }
    }

    return new LineInfo
    {
        filterType1 = type1.ToArray(),
        otherType = other.ToArray(),
    };
}

Filter Type 1の行とそれ以外を抜き出したら、それぞれを分けてジョブを設定しスケジュールします。

LineInfo info = PngParser.ExtractLines(data, _metaData);

_type1Indices = new NativeArray<int>(info.filterType1, Allocator.Persistent);
_otherIndices = new NativeArray<int>(info.otherType, Allocator.Persistent);
_dataArray = new NativeArray<byte>(data, Allocator.Persistent);
_pixelArray = new NativeArray<Pixel32>(_metaData.width * _metaData.height, Allocator.Persistent);

ExpandType1Job type1Job = new ExpandType1Job
{
    indices = _type1Indices,
    data = _dataArray,
    pixels = _pixelArray,
    metaData = _metaData,
};

ExpandJob job = new ExpandJob
{
    indices = _otherIndices,
    data = _dataArray,
    pixels = _pixelArray,
    metaData = _metaData,
};

JobHandle type1JobHandle = type1Job.Schedule(info.filterType1.Length, 32);
_jobHandle = job.Schedule(type1JobHandle);

実際に、左のピクセルだけに依存する行を処理するジョブの実装を以下に示します。

using System;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;

[BurstCompile]
public struct ExpandType1Job : IJobParallelFor
{
    // Filter type 1の行のIndexを格納しているNativeArray。
    [ReadOnly] public NativeArray<int> indices;
    
    // PNGの生データ
    public NativeArray<byte> data;

    // 計算結果を格納するピクセル配列
    public NativeArray<Pixel32> pixels;

    // PNG画像のメタデータ
    public PngMetaData metaData;

    // 各行ごとに並列処理する
    public void Execute(int index)
    {
        int y = indices[index];

        int idx = metaData.rowSize * y;
        int startIndex = idx + 1;

        if (data.Length < startIndex + (metaData.width * metaData.stride))
        {
            throw new IndexOutOfRangeException("Index out of range.");
        }

        Expand(startIndex, y);
    }

    private unsafe void Expand(int startIndex, int y)
    {
        Pixel32* pixelPtr = (Pixel32*)pixels.GetUnsafePtr();
        pixelPtr += (metaData.width * (metaData.height - 1 - y));

        byte* ptr = (byte*)data.GetUnsafePtr();
        ptr += startIndex;

        Pixel32 left = Pixel32.Zero;
        Pixel32 current = default;

        for (int x = 0; x < metaData.width; ++x)
        {
            *(uint*)&current = *(uint*)ptr;

            left = Pixel32.CalculateFloor(current, left);

            ptr += metaData.stride;

            *pixelPtr = left;
            ++pixelPtr;
        }
    }
}

前回の実装のFilter Type 1の場合の処理を並列処理するジョブで実装した例です。

基本的な処理は前回のものとほとんど変わりありませんが、データの取得とその格納部分で、行のIndexを利用して処理をしている点が異なります。

Filter Type 1の行は他の行に依存しないため並列に処理しても競合は起きません。なので今回はここを IJobParallelFor の並列ジョブで実装しました。データを見てみると、だいたいの場合において半分くらいはこのタイプのようなので（自分のチェックした観測範囲内では）、半分のデータを並列処理して高速化することが見込めます。

実際、前回の実装との負荷を比較してみると2～3倍くらいには速くなっていました。

後続の処理もIJobで実装

前述のように、左ピクセルのみに依存する行については並列に処理することができます。一方、それ以外のFilter Typeの場合は上の行にも依存するため逐一計算していく必要があります。（実際には、依存する上の行だけを抜き出してそれをループ処理することで並列化は可能だと思いますが、今回は Filter Type 1 とそれ以外という形で実装しました）

以下にその実装のコードを示します。

using System;
using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;

[BurstCompile]
public struct ExpandJob : IJob
{
    [ReadOnly] public NativeArray<int> indices;
    
    public NativeArray<byte> data;
    public NativeArray<Pixel32> pixels;

    public PngMetaData metaData;

    public void Execute()
    {
        for (int i = 0; i < indices.Length; ++i)
        {
            int y = indices[i];

            int idx = metaData.rowSize * y;
            int startIndex = idx + 1;

            if (data.Length < startIndex + (metaData.width * metaData.stride))
            {
                throw new IndexOutOfRangeException("Index out of range.");
            }

            byte filterType = data[idx];

            switch (filterType)
            {
                case 0:
                    ExpandType0(startIndex, y);
                    break;

                // case 1:
                //     ExpandType1(data, startIndex, stride, h, pixels, metaData);
                //     break;

                case 2:
                    ExpandType2(startIndex, y);
                    break;

                case 3:
                    ExpandType3(startIndex, y);
                    break;

                case 4:
                    ExpandType4(startIndex, y);
                    break;
            }
        }
    }
    
    private unsafe void ExpandType0(int startIndex, int y)
    {
        Pixel32* pixelPtr = (Pixel32*)pixels.GetUnsafePtr();
        pixelPtr += (metaData.width * (metaData.height - 1 - y));

        byte* ptr = (byte*)data.GetUnsafePtr();
        ptr += startIndex;

        Pixel32 current = default;

        for (int x = 0; x < metaData.width; ++x)
        {
            *(uint*)&current = *(uint*)ptr;

            ptr += metaData.stride;

            *pixelPtr = current;
            ++pixelPtr;
        }
    }

    private unsafe void ExpandType2(int startIndex, int y)
    {
        Pixel32* pixelPtr = (Pixel32*)pixels.GetUnsafePtr();
        pixelPtr += (metaData.width * (metaData.height - 1 - y));

        byte* ptr = (byte*)data.GetUnsafePtr();
        ptr += startIndex;

        Pixel32 up = Pixel32.Zero;
        Pixel32 current = default;

        int upStride = metaData.width;

        for (int x = 0; x < metaData.width; ++x)
        {
            *(uint*)&current = *(uint*)ptr;

            if (y == 0)
            {
                up = Pixel32.Zero;
            }
            else
            {
                *(uint*)&up = *(uint*)(pixelPtr + upStride);
            }

            up = Pixel32.CalculateFloor(current, up);

            ptr += metaData.stride;

            *pixelPtr = up;
            ++pixelPtr;
        }
    }

    private unsafe void ExpandType3(int startIndex, int y)
    {
        Pixel32* pixelPtr = (Pixel32*)pixels.GetUnsafePtr();
        pixelPtr += (metaData.width * (metaData.height - 1 - y));

        byte* ptr = (byte*)data.GetUnsafePtr();
        ptr += startIndex;

        Pixel32 up = Pixel32.Zero;
        Pixel32 left = Pixel32.Zero;
        Pixel32 current = default;

        int upStride = metaData.width;

        for (int x = 0; x < metaData.width; ++x)
        {
            *(uint*)&current = *(uint*)ptr;

            if (y == 0)
            {
                up = Pixel32.Zero;
            }
            else
            {
                *(uint*)&up = *(uint*)(pixelPtr + upStride);
            }

            left = Pixel32.CalculateAverage(current, left, up);

            ptr += metaData.stride;

            *pixelPtr = left;
            ++pixelPtr;
        }
    }

    private unsafe void ExpandType4(int startIndex, int y)
    {
        Pixel32* pixelPtr = (Pixel32*)pixels.GetUnsafePtr();
        pixelPtr += (metaData.width * (metaData.height - 1 - y));

        byte* ptr = (byte*)data.GetUnsafePtr();
        ptr += startIndex;

        Pixel32 up = Pixel32.Zero;
        Pixel32 left = Pixel32.Zero;
        Pixel32 leftUp = Pixel32.Zero;
        Pixel32 current = default;

        int upStride = metaData.width;

        for (int x = 0; x < metaData.width; ++x)
        {
            *(uint*)&current = *(uint*)ptr;

            if (y == 0)
            {
                up = Pixel32.Zero;
            }
            else
            {
                *(uint*)&up = *(uint*)(pixelPtr + upStride);
            }

            if (y == 0 || x == 0)
            {
                leftUp = Pixel32.Zero;
            }
            else
            {
                *(uint*)&leftUp = *(uint*)(pixelPtr + upStride - 1);
            }

            left = Pixel32.CalculatePaeth(left, up, leftUp, current);

            ptr += metaData.stride;

            *pixelPtr = left;
            ++pixelPtr;
        }
    }
}

実装したジョブを実行する

上記のジョブを実際に実行しているコードは以下になります。

private async void StartJob()
{
    string filePath = PngImageManager.GetSavePath(_urlField.text);

    (PngMetaData metaData, byte[] data) = await Task.Run(() =>
    {
        byte[] rawData = File.ReadAllBytes(filePath);
        return PngParser.Decompress(rawData);
    });
    
    _metaData = metaData;
    LineInfo info = PngParser.ExtractLines(data, _metaData);

    _type1Indices = new NativeArray<int>(info.filterType1, Allocator.Persistent);
    _otherIndices = new NativeArray<int>(info.otherType, Allocator.Persistent);
    _dataArray = new NativeArray<byte>(data, Allocator.Persistent);
    _pixelArray = new NativeArray<Pixel32>(_metaData.width * _metaData.height, Allocator.Persistent);
    
    _stopwatch.Restart();

    ExpandType1Job type1Job = new ExpandType1Job
    {
        indices = _type1Indices,
        data = _dataArray,
        pixels = _pixelArray,
        metaData = _metaData,
    };

    ExpandJob job = new ExpandJob
    {
        indices = _otherIndices,
        data = _dataArray,
        pixels = _pixelArray,
        metaData = _metaData,
    };

    JobHandle type1JobHandle = type1Job.Schedule(info.filterType1.Length, 32);
    _jobHandle = job.Schedule(type1JobHandle);

    _started = true;
    
    ShowTexture();
}

private unsafe void ShowTexture()
{
    // Needs to complete even if it checked `IsCompleted`.
    // This just avoids an error.
    _jobHandle.Complete();

    IntPtr pointer = (IntPtr)_pixelArray.GetUnsafePtr();

    Texture2D texture = new Texture2D(_metaData.width, _metaData.height, TextureFormat.RGBA32, false);

    texture.LoadRawTextureData(pointer, _metaData.width * _metaData.height * 4);
    texture.Apply();

    _preview.texture = texture;

    Dispose();

    _started = false;

    _stopwatch.Stop();

    Debug.Log($"Elapsed time: {_stopwatch.ElapsedMilliseconds.ToString()}ms");
}

以下の部分が、Filter Type 1 の行を抜き出し、それぞれのインデックスバッファを生成している箇所です。

LineInfo info = PngParser.ExtractLines(data, _metaData);

_type1Indices = new NativeArray<int>(info.filterType1, Allocator.Persistent);
_otherIndices = new NativeArray<int>(info.otherType, Allocator.Persistent);

しれっと「インデックスバッファ」と書きましたが、CGのレンダリングでインデックスバッファと頂点バッファを生成しているくだりに似ているのでそう表現しました。ここで得られた行番号（ _type1Indices ）のリストが、並列化可能な行となります。

そして _otherIndices が残りの行となります。あとはこれらのデータと結果を受け取る配列の NativeArray<Pixel32> をアロケートし、それをジョブに渡すことでJob SystemによるPNGの展開処理が完了します。

最後に

今回、初めてしっかりとC# Job SystemとBurst Compilerを使ってみましたが下準備が少し面倒なだけで、安全にマルチスレッドでの処理を書けるのはとてもいいと思いました。うまくすればBurst Compilerによって高速化も望めます。並列化できそうな処理があったら積極的に使っていきたいと思います。

ランキング参加中
Unity