Speed up Reaction-Diffusion simulation calculations in Unity

Question

I'm trying to use Unity Job system to speed up some calculation related to Reaction-Diffusion simulation. The calculations requires 2 matrixes, current and next and the current will be used to calculated next, then next will become current. I wanted each job to handle a single row of the matrix.

The problem lies in the Laplace() function, which when inputed a coordination, return a value based on that cell and 8 other cells around it:

private float laplaceA(int x, int y)
{
    float sum = 0f; 
    for (int idx = -1; idx <= 1; idx++)
    {
        for (int j = -1; j <= 1; j++)
        {
            int newX = x + idx, newY = y + j;
            if (withinBorder(newX, newY))
            {
                float val = cur[newX*height + newY].x;

                if (idx == 0 && j == 0) { sum += val * -1f;} //center
                else if (Absy(idx) == 1 && Absy(j) == 1) {sum += val * 0.05f;} //diagonal
                else { sum += val * 0.2f;} //adjacent
            }
        }
    }
    return sum;
    //NOTE: Absy() is just Mathf.Abs() rewrite for using jobs
}

I tried some ways to make it works:

1. IJobParallel seems promising at first, but then I realize Laplace() can't access what it need, as I need the previous and next row, not just the current row (hence this post).
2. Tried NativeArray<NativeArray<float2>> to pass in each array one of a time, but failed as Unity doesn't support this yet. List<NativeArray<float2>> to try and sneak in failed as well.
3. Tried using just IJob and creating a NativeArray<> each row and a job with it, but I haven't figured out a way to store each array to dispose of it later.

Any idea to fix this problem, or other methods for speeding up calculations?

Answer 1

After some experiements I figured out what you need to do. The key appears to be to use an IJobParallelFor and to mark the array you only want to read from with the [ReadOnly] attribute. This IJobParallelFor interface is specifically designed for operating on arrays. The [ReadOnly] attribute allows you to access cells at other indices than the one currently being iterated.

1. Flatten the 2d array into a 1d array. This is important because IJobParallelFor isn't designed to work with 2d arrays.
2. Keep two copies of the whole matrix as persistent NativeArrays within the MonoBehaviour / System from which you intend to spawn the jobs. The reason why we want two persistent copies is because we don't want to constantly allocate and deallocate memory. What we are going to do is implement double buffering. One copy is being created while the other one serves as a template. And then for the next iteration we simply switch them, so that the new buffer becomes the template and the old buffer gets overwritten by the next version.
3. Create your job as an IJobParallelFor. You do not want to create copies of the persistent arrays, so the job should take its data as NativeSlices. One for the current array which should be declared with the [ReadOnly] attribute. This is important, because otherwise you get "ReadWriteBuffers are restricted to only read & write the element at the job index." errors as soon as you try to access another cell than the one being iterated. The other is for the next array which it is going to overwrite. We are going to tag this with the [WriteOnly] attribute. It does not appear strictly necessary to do that and it doesn't really seem to affect performance in my particular test-case. But who knows, perhaps it allows some under-the-hood optimizations.
4. Schedule the IJobParallelFor from Update (or where else you want to schedule it)
5. When the job is completed, switch which NativeArray is "current" and which one is the "next" array.

Here is an example:

public class LaplaceVisualisation : MonoBehaviour {
    // some constants to avoid magic numbers
    private static int WIDTH = 256;
    private static int HEIGHT = 256;
    private static int NUM_CELLS = WIDTH * HEIGHT;
    
    // the two persistent data buffers
    private NativeArray<float2> bufferA;
    private NativeArray<float2> bufferB; 

    // the double buffering implementation
    private bool buffersFlipped;
    private NativeSlice<float2> frontBuffer{ get => buffersFlipped ? bufferB : bufferA; }
    private NativeSlice<float2> backBuffer { get => buffersFlipped ? bufferA : bufferB; }


    void Start() {
        // create the two persistent buffers
        bufferA = new NativeArray<float2>(NUM_CELLS, Allocator.Persistent);
        bufferB = new NativeArray<float2>(NUM_CELLS, Allocator.Persistent);

        // TODO: Initialize bufferA with the initial configuration
        // You don't need to initialize bufferB, because it will get overwritten anyway
    }

    // the parallel job
    struct LaplaceJob : IJobParallelFor {
        [ReadOnly] public NativeSlice<vector2> current;
        [WriteOnly] public NativeSlice<vector2> next;

        public void Execute(int i) {

            var x = i % WIDTH;
            var y = i / WIDTH;
            
            // TODO: your logic here 
            // read ONLY from current
            // write ONLY to next[i]
        }
    }
    
    void Update() {
        // create the job.
        // note that frontBuffer and backBuffer are implemented as 
        // properties which return a slice of whatever buffer
        // has the respective role at the moment
        var laplaceJob = new LaplaceJob {
            current = frontBuffer,
            next = backBuffer
        };

        // run it
        var laplaceHandle = laplaceJob.Schedule(NUM_CELLS, 256);
        laplaceHandle.Complete();

        // flip frontBuffer and backBuffer
        buffersFlipped = !buffersFlipped;

        // visualise the data
        var data = frontBuffer;
        // TODO: Your visualisation code here
        // Do that by reading from data
    }

    private void OnDestroy() {
        // and don't forget to dispose the two NativeArray's when
        // Unity quits, so we don't get an ugly error message
        bufferA.Dispose();
        bufferB.Dispose();
    }
}