Работа с Gradient через jobs + burst

В Unity есть класс Gradient, который предоставляет удобные средства для управления градиентом в рантайме и редакторе. Но т.к. это класс, а не структура использовать его через Job system и burst нельзя. Это первая проблема. Вторая проблема — это работа с ключами градиента. Получение значений осуществляется через массив, который создаётся в куче. И как следствие напрягает сборщик мусора.

Сейчас я покажу как можно решить эти проблемы. И в качестве бонуса получить увеличение производительности до 8 раз при выполнении метода Evaluate через burst.

Структура

Для начала нужно получить прямой доступ к памяти объекта градиента. Этот адрес является неизменяемым на протяжении всей жизни объекта. Получив его один раз можно работать с прямым доступом к данным градиента, не переживая о том, что адрес может измениться.

Вот простейший метод расширения который поможет прочитать адрес, не прибегая к рефлексии при каждом вызове. Из минусов, объект нужно закрепить. Из плюсов адрес нужно получить только один раз поэтому время, затраченное на закрепление объекта не критично.

public static class GradientExt
{
    private static readonly int m_PtrOffset;
    
    static GradientExt()
    {
        var m_PtrMember = typeof(Gradient).GetField("m_Ptr", BindingFlags.Instance | BindingFlags.NonPublic
        m_PtrOffset = UnsafeUtility.GetFieldOffset(m_PtrMember);
    }
    
    public static unsafe IntPtr Ptr(this Gradient gradient)
    {
        var ptr = (byte*) UnsafeUtility.PinGCObjectAndGetAddress(gradient, out var handle);
        var gradientPtr = *(IntPtr*) (ptr + m_PtrOffset);
        UnsafeUtility.ReleaseGCObject(handle);
        return gradientPtr;
    }
}

UnsafeUtility.GetFieldOffset – возвращает смещение поля относительно структуры или класса, в котором оно содержится.

UnsafeUtility.PinGCObjectAndGetAddress – закрепляет объект. И гарантирует, что объект не будет перемещаться в памяти. Возвращает адрес участка памяти, в котором находится объект.

UnsafeUtility.ReleaseGCObject – освобождает хэндл объекта GC, полученный ранее.

Теперь можно получить адрес на участок памяти, где хранятся данные градиента.

public Gradient gradient;
....
IntPtr gradientPtr = gradient.Ptr();

Дальше нужно немножко поковырять память чтобы понять как именно расположены данные градиента. Для этого я выведу в инспектор Unity этот участок памяти в виде массива. Затем остаётся лишь изменять градиент и смотреть какие именно участки это затрагивает.

[ExecuteAlways]
public class MemoryResearch : MonoBehaviour
{
    public Gradient gradient = new Gradient();
    public float[] gradientMemoryLocation = new float[50];
    private static unsafe void CopyMemory<T>(Gradient gradient, T[] gradientMemoryLocation) where T : unmanaged
    {
        IntPtr gradientPtr = gradient.Ptr();
        fixed (T* gradientMemoryLocationPtr = gradientMemoryLocation)
            UnsafeUtility.MemCpy(gradientMemoryLocationPtr, (void*) gradientPtr, gradientMemoryLocation.Length);
    }
    
    private void Update()
    {
        CopyMemory(gradient, gradientMemoryLocation);
    }
}

UnsafeUtility.MemCpy – копирует указанное количество байт из одной области памяти в другую.

Путём нехитрых манипуляций и смены типа памяти float/ushort/byte и т.д. я нашёл полное расположение каждого параметра градиента. В статье буду приводить примеры для Unity 22.3, но есть небольшие различия для разных версий. С полной версией кода можно ознакомится в конце статьи.

//Позиции ключей хранятся как ushort где 0 = 0%, а 65535 = 100%.
public unsafe struct GradientStruct
{
    private fixed byte colors[sizeof(float) * 4 * 8]; //8 rgba цветовых значений (128 байт)
    private fixed byte colorTimes[sizeof(ushort) * 8]; //время для каждого цветового ключа (16 байт)
    private fixed byte alphaTimes[sizeof(ushort) * 8]; //время для каждого альфа ключа (16 байт)
    private byte colorCount; //количество цветовых ключей
    private byte alphaCount; //количество альфа ключей
    private byte mode; //режим смешивания цветов
    private byte colorSpace; //цветовое пространство
}

Также добавляю метод расширения для получения указателя на структуру GradientStruct:

public static unsafe GradientStruct* DirectAccess(this Gradient gradient)
{
    return (GradientStruct*) gradient.Ptr();
}

Gradient.colorKeys через NativeArray

Зная структуру памяти градиента, можно написать методы для работы с Gradient.colorKeys и Gradient.alphaKeys через NativeArray.

private float4* Colors(int index)
{
    fixed(byte* colorsPtr = colors) return (float4*) colorsPtr + index;
}

private ushort* ColorsTimes(int index)
{
    fixed(byte* colorTimesPtr = colorTimes) return (ushort*) colorTimesPtr + index;
}

private ushort* AlphaTimes(int index)
{
    fixed(byte* alphaTimesPtr = alphaTimes) return (ushort*) alphaTimesPtr + index;
}

public void SetColorKey(int index, GradientColorKeyBurst value)
{
    #if ENABLE_UNITY_COLLECTIONS_CHECKS
      if (index < 0 || index > 7) IncorrectIndex();
    #endif
    Colors(index)->xyz = value.color.xyz;
    *ColorsTimes(index) = (ushort) (65535 * value.time);
}

public GradientColorKeyBurst GetColorKey(int index)
{
    #if ENABLE_UNITY_COLLECTIONS_CHECKS
        if (index < 0 || index > 7) IncorrectIndex();
    #endif
    return new GradientColorKeyBurst(*Colors(index), *ColorsTimes(index) / 65535f);
}

public void SetColorKeys(NativeArray<GradientColorKeyBurst> colorKeys)
{
    #if ENABLE_UNITY_COLLECTIONS_CHECKS
        if (colorKeys.Length < 2 || colorKeys.Length > 8) IncorrectLength();
    #endif
    
    for (var i = 0; i < colorCount; i++)
    {
        SetColorKey(i, colorKeys[i]);
    }
}

public NativeArray<GradientColorKeyBurst> GetColorKeys(Allocator allocator)
{
    var colorKeys = new NativeArray<GradientColorKeyBurst>(colorCount, allocator);
    
    for (var i = 0; i < colorCount; i++)
    {
        colorKeys[i] = GetColorKey(i);
    }
    return colorKeys;
}

public void SetAlphaKey(int index, GradientAlphaKey value)
{
    #if ENABLE_UNITY_COLLECTIONS_CHECKS
      if (index < 0 || index > 7) IncorrectIndex();
    #endif
    Colors(index)->w = value.alpha;
    *AlphaTimes(index) = (ushort) (65535 * value.time);
}

public GradientAlphaKey GetAlphaKey(int index)
{
    #if ENABLE_UNITY_COLLECTIONS_CHECKS
        if (index < 0 || index > 7) IncorrectIndex();
    #endif
    return new GradientAlphaKey(Colors(index)->w, *AlphaTimes(index) / 65535f);
}

public void SetAlphaKeys(NativeArray<GradientAlphaKey> alphaKeys)
{
    #if ENABLE_UNITY_COLLECTIONS_CHECKS
        if (alphaKeys.Length < 2 || alphaKeys.Length > 8) IncorrectLength();
    #endif
    
    for (var i = 0; i < colorCount; i++)
    {
        SetAlphaKey(i, alphaKeys[i]);
    }
}

public NativeArray<GradientAlphaKey> GetAlphaKeys(Allocator allocator)
{
    var alphaKeys = new NativeArray<GradientAlphaKey>(alphaCount, allocator);
    
    for (var i = 0; i < alphaCount; i++)
    {
        alphaKeys[i] = GetAlphaKey(i);
    }
    return alphaKeys;
}

В результате

var colorKeys = gradient.colorKeys;
var alphaKeys = gradient.alphaKeys;

можно заменить на

var gradientPtr = gradient.DirectAccess();
var colorKeys = gradientPtr->GetColorKeys(Allocator.Temp);
var alphaKeys = gradientPtr->GetAlphaKeys(Allocator.Temp);

и забыть о сборщике мусора при чтении значений. А также использовать эти методы внутри Job system. Результат gradient.DirectAccess() можно закешировать и использовать на протяжении всей жизни объекта.

Финальная подготовка для Job system

Нужно сделать свою реализацию метода Evaluate, потому что нативный метод остался вместе с классом за пределами досягаемости из новой структуры. Вдаваться в подробности алгоритма не буду. Он слишком тривиален и не имеет отношения к теме статьи.

public float4 EvaluateBurst(float time)
{
    float3 color = default;
    var colorCalculated = false;
    var colorKey = GetColorKeyBurst(0);
    if (time <= colorKey.time)
    {
        color = colorKey.color.xyz;
        colorCalculated = true;
    }
    
    if (!colorCalculated)
        for (var i = 0; i < colorCount - 1; i++)
        {
            var colorKeyNext = GetColorKeyBurst(i + 1);
                
            if (time <= colorKeyNext.time)
            {
                if (Mode == GradientMode.Blend)
                {
                    var localTime = (time - colorKey.time) / (colorKeyNext.time - colorKey.time);
                    color = math.lerp(colorKey.color.xyz, colorKeyNext.color.xyz, localTime);
                }
                else if (Mode == GradientMode.PerceptualBlend)
                {
                    var localTime = (time - colorKey.time) / (colorKeyNext.time - colorKey.time);
                    color = OklabToLinear(math.lerp(LinearToOklab(colorKey.color.xyz), LinearToOklab(colorKeyNext.color.xyz), localTime));
                }
                else
                {
                    color = colorKeyNext.color.xyz;
                }
                colorCalculated = true;
                break;
            }
            
            colorKey = colorKeyNext;
        }
  
    if (!colorCalculated) color = colorKey.color.xyz;
    
    
    
    float alpha = default;
    var alphaCalculated = false;
    
    var alphaKey = GetAlphaKey(0);
    if (time <= alphaKey.time)
    {
        alpha = alphaKey.alpha;
        alphaCalculated = true;
    }
    
    if (!alphaCalculated)
        for (var i = 0; i < alphaCount - 1; i++)
        {
            var alphaKeyNext = GetAlphaKey(i + 1);
                
            if (time <= alphaKeyNext.time)
            {
                if (Mode == GradientMode.Blend || Mode == GradientMode.PerceptualBlend)
                {
                    var localTime = (time - alphaKey.time) / (alphaKeyNext.time - alphaKey.time);
                    alpha = math.lerp(alphaKey.alpha, alphaKeyNext.alpha, localTime);
                }
                else
                {
                    alpha = alphaKeyNext.alpha;
                }
                alphaCalculated = true;
                break;
            }
            
            alphaKey = alphaKeyNext;
        }
    
    if (!alphaCalculated) alpha = alphaKey.alpha;
        
    return new float4(color, alpha);
}

Многопоточность

Полученная выше структура умеет как читать значения, так и писать их. Если попытаться её использовать одновременно в разных потоках для записи, то будет Race Conditions. Никогда не используйте её для многопоточных заданий. Для этого я подготовлю readonly версию.

internal unsafe struct GradientStruct
{

  ...
  
  public static ReadOnly AsReadOnly(GradientStruct* data) => new ReadOnly(data);

  public readonly struct ReadOnly
  {
      private readonly GradientStruct* ptr;
      public ReadOnly(GradientStruct* ptr)
      {
          this.ptr = ptr;
      }
      
      public int ColorCount => ptr->ColorCount;
      
      public int AlphaCount => ptr->AlphaCount;
      
      public GradientMode Mode => ptr->Mode;
      
      #if UNITY_2022_2_OR_NEWER
          public ColorSpace ColorSpace => ptr->ColorSpace;
      #endif
      
      public GradientColorKeyBurst GetColorKey(int index) => ptr->GetColorKey(index);
      
      public NativeArray<GradientColorKeyBurst> GetColorKeys(Allocator allocator) => ptr->GetColorKeys(allocator);
      
      public GradientAlphaKey GetAlphaKey(int index) => ptr->GetAlphaKey(index);
      
      public NativeArray<GradientAlphaKey> GetAlphaKeys(Allocator allocator) => ptr->GetAlphaKeys(allocator);
      
      public float4 Evaluate(float time)=> ptr->Evaluate(time);
  }
}

И метод расширения:

public static unsafe GradientStruct.ReadOnly DirectAccessReadOnly(this Gradient gradient)
{
    return GradientStruct.AsReadOnly(gradient.DirectAccess());
}

Эту структуру для чтения точно также достаточно создать один раз и можно передать в любое многопоточное задание или использовать где-то ещё на протяжении всей жизни объекта.

Пример использования:

var gradientReadOnly = gradient.DirectAccessReadOnly();
var colorKeys = gradientReadOnly.GetColorKeys(Allocator.Temp);
var alphaKeys = gradientReadOnly.GetAlphaKeys(Allocator.Temp);
var color = gradientReadOnly.Evaluate(0.6f);
colorKeys.Dispose();
alphaKeys.Dispose();

Тест производительности

Для тестов использован процессор с поддержкой AVX2. Данным тестом я не ставил цель показать максимально объективные результаты. Но тенденция должны быть понятна. Суть теста: в одном потоке делается сто тысяч итераций и вычисляется цвет градиента с помощью метода Evaluate. Во всех режимах интерполяции с большим отрывом лидирует кастомная реализация. Что стало для меня большим сюрпризом. Был уверен, что c++ версия будет быстрее.

public class PerformanceTest : MonoBehaviour
{
    public Gradient gradient = new Gradient();
    
    [BurstCompile(OptimizeFor = OptimizeFor.Performance)]
    private unsafe struct GradientBurstJob : IJob
    {
        public NativeArray<float4> result;
        [NativeDisableUnsafePtrRestriction] public GradientStruct* gradient;
        
        public void Execute()
        {
            var time = 1f;
            var color = float4.zero;
            
            for (var i = 0; i < 100000; i++)
            {
                time *= 0.9999f;
                color += gradient->EvaluateBurst(time);
            }
            result[0] = color;
        }
    }
    
    private unsafe void Update()
    {
        var nativeArrayResult = new NativeArray<float4>(1, Allocator.TempJob);
        var job = new GradientBurstJob
        {
            result = nativeArrayResult,
            gradient = gradient.DirectAccess()
        };
        var jobHandle = job.ScheduleByRef();
        JobHandle.ScheduleBatchedJobs();
        
        Profiler.BeginSample("NativeGradient");
        var time = 1f;
        var result = new Color(0, 0, 0, 0);
        
        for (var i = 0; i < 100000; i++)
        {
            time *= 0.9999f;
            result += gradient.Evaluate(time);
        }
        Profiler.EndSample();
        
        jobHandle.Complete();
        nativeArrayResult.Dispose();
    }
}

Итог. В результате простейших манипуляций я получил прямой доступ к памяти градиента предназначенной для c++ части движка. Отправил указатель на эту память в Job system и смог произвести вычисления внутри задания воспользовавшись всеми преимуществами компилятора burst.

Совместимость

Работоспособность проверена во всех версиях Unity начиная с 2020.3 и заканчивая 2023.2. 0a19.Скорее всего каких-то изменений не будет до тех пор, пока в Unity не решат добавить новые фичи для градиента. За последние годы такое случилось лишь единожды в версии 2022.2. Но я настоятельно рекомендую, прежде чем воспользоваться этим кодом в непроверенных версиях убедиться в его работоспособности.

Ссылка на полную версию

Как и обещал вот ссылка на полный исходный код https://gist.github.com/viruseg/791789d63775d26a79ca32c0f5d31114