Working with Gradient via jobs + burst

Unity has a Gradient class that provides convenient means for managing gradients at runtime and in the editor. But because This is a class, not a structure, it cannot be used through the Job system and burst. This is the first problem. The second problem is working with gradient keys. The values are obtained through an array that is created on the heap. And as a result, the garbage collector is annoying.

Now I will show how you can solve these problems. And as a bonus, get a performance increase of up to 8 times when executing the Evaluate method via burst.

Structure

First, you need to directly access the memory of the gradient object. This address remains unchanged throughout the life of the object. Having received it once, you can work with direct access to the gradient data without worrying that the address may change.

m_Ptr is the address that needs to be obtained to access the memory of the C++ part of the engine.

Here is the simplest extension method that will help you read the address without resorting to reflection every time you call it. The downside is that the object needs to be secured. On the plus side, the address needs to be obtained only once, so the time spent on securing the object is not critical.

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 – Returns the offset of a field relative to the structure or class that contains it.

UnsafeUtility.PinGCObjectAndGetAddress – pins the object. And it ensures that the object will not move around in memory. Returns the address of the memory location in which the object is located.

UnsafeUtility.ReleaseGCObject – releases the GC object handle obtained previously.

Now you can get the address of the memory location where the gradient data is stored.

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

Next, you need to poke around your memory a little to understand exactly how the gradient data is located. To do this, I will display this piece of memory as an array in the Unity inspector. Then all that remains is to change the gradient and see which areas it affects.

[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 – copies the specified number of bytes from one memory area to another.

Demonstration of how values change in memory when color changes.

By simple manipulations and changing the memory type float/ushort/byte, etc. I found the complete location of each gradient parameter. In this article I will give examples for Unity 22.3, but there are slight differences for different versions. The full version of the code can be found at the end of the article.

//Позиции ключей хранятся как 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; //цветовое пространство
}

I also add an extension method to get a pointer to the GradientStruct structure:

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

Gradient.colorKeys via NativeArray

Knowing the gradient memory structure, you can write methods to work with Gradient.colorKeys and Gradient.alphaKeys via 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;
}

As a result

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

can be replaced by

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

and forget about the garbage collector when reading values. And also use these methods inside the Job system. The result of gradient.DirectAccess() can be cached and used throughout the life of the object.

Final preparation for Job system

You need to make your own implementation of the Evaluate method, because the native method remained, along with the class, beyond the reach of the new structure. I will not go into details of the algorithm. It is too trivial and not relevant to the topic of the article.

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);
}

Multithreading

The structure obtained above can both read values and write them. If you try to use it simultaneously in different recording streams, you will get Race Conditions. Never use it for multi-threaded jobs. For this I will prepare a readonly version.

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);
  }
}

And the extension method:

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

This reading structure also just needs to be created once and can be passed to any multi-threaded job or used somewhere else throughout the entire life of the object.

Usage example:

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();

Performance test

A processor with AVX2 support was used for tests. With this test, I did not set out to show the most objective results. But the trend should be clear. The essence of the test: one hundred thousand iterations are done in one thread and the color of the gradient is calculated using the Evaluate method. In all interpolation modes, custom implementation leads by a wide margin. Which came as a big surprise to me. I was sure that the C++ version would be faster.

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();
    }
}

Bottom line. As a result of the simplest manipulations, I got direct access to the gradient memory of the C++ part of the engine. I sent a pointer to this memory to the Job system and was able to perform calculations inside the job, taking advantage of all the advantages of the burst compiler.

Compatibility

Performance has been tested in all versions of Unity from 2020.3 to 2023.2. 0a19. Most likely there will be no changes until Unity decides to add new features for the gradient. In recent years, this has happened only once in version 2022.2. But I strongly recommend that before using this code in untested versions, you make sure it works.