【Unity】ComponentDataArray<T>からの効率的なコピーについて

執筆者環境

• Unity 2018.2.2f1 Personal
• ECS version0.0.12-preview8

問題提起

私が以前書いた記事ではComponentDataArrayからNativeArrayにデータをコピーし、ComputeBufferにNativeArrayをSetDataするという処理がありました。
この処理はOnUpdateメソッド内部で呼ばれるため、毎フレームGC と私の腸から血が流れる程のストレスが発生し続けました。

using(var nativeArray = new NativeArray<float2>(length, Allocator.Temp, NativeArrayOptions.UninitializedMemory))
{
    componentDataArray.CopyTo(nativeArray);
    computeBuffer.SetData(nativeArray);
}

解決策

ComponentDataArrayにはGetChunkArrayというメソッドが定義されていました。

public NativeArray<T> GetChunkArray(int startIndex, int maxCount)

公式のリファレンスが2018年8月現在未整備もいいところなので前記事を書いた時点では使い方がわからず、当然不使用でした。
このGetChunkArrayの使用法は次のメソッドで明瞭に示されます。

Unity.Entities/Iterators/ComponentDataArray.cs(抜粋)

public void CopyTo(NativeSlice<T> dst, int startIndex = 0)
{
    var copiedCount = 0;
    while (copiedCount < dst.Length)
    {
        var chunkArray = GetChunkArray(startIndex + copiedCount, dst.Length - copiedCount);
        dst.Slice(copiedCount, chunkArray.Length).CopyFrom(chunkArray);

        copiedCount += chunkArray.Length;
    }
}

FileStreamのint Read(Span)と同様の使い方ですね！
そうかそうか、while文で回しながら引っ張ってくるのが正しいやり方でしたか。

GetChunkArrayの内部も見てみましょう。

Unity.Entities/Iterators/ComponentDataArray.cs(抜粋)

public NativeArray<T> GetChunkArray(int startIndex, int maxCount)
{
    int count;
    //@TODO: How should we declare read / write here?
    var ptr = GetUnsafeChunkPtr(startIndex, maxCount, out count, true);

    var arr = NativeArrayUnsafeUtility.ConvertExistingDataToNativeArray<T>(ptr, count, Allocator.Invalid);

#if ENABLE_UNITY_COLLECTIONS_CHECKS
    NativeArrayUnsafeUtility.SetAtomicSafetyHandle(ref arr, m_Safety);
#endif

    return arr;
}

internal void* GetUnsafeChunkPtr(int startIndex, int maxCount, out int actualCount, bool isWriting)
{
#if ENABLE_UNITY_COLLECTIONS_CHECKS
    GetUnsafeChunkPtrCheck(startIndex, maxCount);
#endif

    m_Iterator.UpdateCache(startIndex, out m_Cache, isWriting);

    void* ptr = (byte*) m_Cache.CachedPtr + startIndex * m_Cache.CachedSizeOf;
    actualCount = Math.Min(maxCount, m_Cache.CachedEndIndex - startIndex);

    return ptr;
}

NativeArrayUnsafeUtility.ConvertExistingDataToNativeArrayはアロケーションフリーですからGetChunkArrayで戻されるNativeArrayもまたアロケーションフリーです。
適切にGetChunkArrayすればアロケーションフリーでComponentDataArrayからコピペできますね。

実際の例

public static class ComputeBufferHelper
{
    public static void CopyFrom<T>(this ComputeBuffer dst, ref ComponentDataArray<T> src, int length) where T : struct, IComponentData
    {
        NativeArray<T> chunkArray;
        for (int copiedCount = 0; copiedCount < length; copiedCount += chunkArray.Length)
        {
            // Allocator.InvalidであるからDisposeしてはならぬ。
            chunkArray = src.GetChunkArray(copiedCount, length - copiedCount);
            dst.SetData(chunkArray, 0, copiedCount, chunkArray.Length);
        }
    }
    public static unsafe void CopyFromUnsafe<T>(this ComputeBuffer dst, ref ComponentDataArray<T> src, int length) where T : struct, IComponentData
    {
        var stride = dst.stride;
        var ptr = dst.GetNativeBufferPtr();
        NativeArray<T> chunkArray;
        for (int copiedCount = 0, chunkLength = 0; copiedCount < length; copiedCount += chunkLength)
        {
            chunkArray = src.GetChunkArray(copiedCount, length - copiedCount);
            chunkLength = chunkArray.Length;
            UnsafeUtility.MemCpy(ptr.ToPointer(), NativeArrayUnsafeUtility.GetUnsafeBufferPointerWithoutChecks(chunkArray), chunkLength * stride);
            ptr += chunkLength * stride;
        }
    }
}

unsafe使う気があるなら下の拡張メソッドを使うのもいいかもしれませんね。
応用としてGetChunkArrayで得たNativeArrayをIJobParallelForBatchを実装したJobに渡すというのもいいでしょう。
JobHandleをNativeListに追加していってJobHandle.CombineDependenciesで一纏めにするという使い方も想定できます。