Using Unity’s New Job System for Diamond-Square Generation

As I’ve mentioned in previous posts, I’ve begun work on a new project that likely will be some kind of follow-up to Starcom: Nexus. For this week’s update I thought I’d do a technical discussion of one of my efforts.

Unity has been making some big changes to the engine to support a more performant Entity-Component-Systems model, which they call DOTS (Data-Oriented Technology Stack). Based on my investigations to date, I don’t think the overall system is close to ready for production. It does work, but the documentation is currently extremely sparse and bugs are common. Hopefully that will improve later this year as the technology matures and Unity can devote more resources to documenting it.

Still, this week I was able to implement one feature from Starcom: Nexus using Unity’s new Job System. In the original game, most of the planets have procedurally generated textures produced by combining 3 different base textures using a color map which I call a “thermal map” because it sort of looks like heat vision. It starts with an image like this:

“Thermal” color map

That image is passed to a shader which uses the red, green and blue components to selectively blend three different textures that represent the low, medium and high elevations respectively. The shader does some other stuff as well, but the bulk of the procedural aspect comes from the generation of the above image.

The end result looks something like this:

Or this…

The thermal map is generated using the diamond-square algorithm.  This relatively simple (as far as procedural generation goes) algorithm produces a pseudo-random grid of values with some very convenient properties:

  1. First, and most importantly, it is a nice approximation of an elevation map.
  2. It can be made to wrap-around.
  3. Using the same seed, it will always generate the same output. This means that we can end up with an arbitrarily high resolution image but only store a single integer.
  4. We can adjust the resolution without changing the overall output: if we start with a 128×128 texture map, then later switch to 512×512, it will look like a higher resolution version of the same image rather than something wildly different.

A downside is that it is not super fast. Generating a 512×512 texture map involves iterating over a quarter million coordinates, each of which also involves averaging four other coordinates and a separate R, G, and B evaluation. On my machine I think it took maybe a dozen milliseconds. Not a lot, but if you suddenly need to load a dozen planets, it would produce a noticeable stutter in the frame rate. In the original game I implemented the algorithm as an IEnumerable so that the task could be spread out over several dozen frames. But it still happened in the Unity Main thread, which meant it was potentially competing with other operations.

Enter the new Job System. Jobs have the advantage that they allow game logic to run in separate threads, like Unity’s internal rendering system and physics already do, but without the developer needing to worry (as much) about the nightmare of race conditions.

I decided that the diamond-square algorithm was a good candidate to try out this system.

Here is the component in its entirety, warts and all:

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

namespace Wx3.StarcomGameplay
{
  /// <summary>
  /// A diamond-square texture generator using Unity jobs
  /// </summary>
  public class DiamondSquare : MonoBehaviour, IProcedural
  {
    public Gradient elevationGradient;
    public int scale = 8;
    public int seed;
    public float roughness = 0.5f;

    private JobHandle jobHandle;
    private NativeArray<Color> pixels;
    [ReadOnly]
    private NativeArray<GradientColorKey> gradientColors;
    private bool started = false;

    public struct DiamondSquareJob : IJob
    {
      public NativeArray<Color> pixels;
      public int seed;
      public int size;
      public float roughness;
      public NativeArray<GradientColorKey> gradientColors;

      public void Execute()
      {
        System.Random random = new System.Random(seed);
        // The basic algorithm is based on this SO answer:
        // https://stackoverflow.com/questions/2755750/diamond-square-algorithm
        int gridDim = size + 1;
        float[,] grid = new float[gridDim, gridDim];
        float offset = roughness;
        grid[0, 0] = 0.5f;
        grid[0, gridDim - 1] = 0.5f;
        grid[gridDim - 1, 0] = 0.5f;
        grid[gridDim - 1, gridDim - 1] = 0.5f;
        for (int side = gridDim - 1; side >= 2; side /= 2)
        {
          int halfSide = side / 2;
          // Square values:
          for (int x = 0; x < gridDim - 1; x += side)
          {
            for (int y = 0; y < gridDim - 1; y += side)
            {
              float avg = grid[x, y] + grid[x + side, y] + grid[x, y + side] + grid[x + side, y + side];
              if (avg < 0) avg = 0;
              avg = avg / 4.0f;
              grid[x + halfSide, y + halfSide] = avg + Range(-offset, offset, random);
            }
          }
          // Diamond values:
          for (int x = 0; x < gridDim - 1; x += halfSide)
          {
            for (int y = (x + halfSide) % side; y < gridDim - 1; y += side)
            {
              float avg =
                grid[(x - halfSide + gridDim - 1) % (gridDim - 1), y] +
                grid[(x + halfSide) % (gridDim - 1), y] +
                grid[x, (y + halfSide) % (gridDim - 1)] +
                grid[x, (y - halfSide + gridDim - 1) % (gridDim - 1)];
              if (avg < 0) avg = 0;
              avg = avg / 4.0f;
              avg = Mathf.Clamp(avg + Range(-offset, offset, random), 0, 1);
              grid[x, y] = avg;
              if (x == 0) grid[gridDim - 1, y] = avg;
              if (y == 0) grid[x, gridDim - 1] = avg;
            }
          }
          offset *= 0.5f;
        }
        
        for(int y = 1; y <= size; y++)
        {
          for(int x = 1; x <= size; x++)
          {
            // Our diamond-square is 1 bigger than our texture so we average the four points 
            // around each pixel:
            float avg = grid[x - 1, y - 1] + grid[x, y - 1] + grid[x - 1, y] + grid[x, y];
            avg /= 4.0f;
            int i = ((y - 1) * size) + (x - 1);
            pixels[i] = Evaluate(avg); 
          }
        }
      }

      /// <summary>
      /// Struct-friendly version of Gradient evaluation:
      /// </summary>
      /// <param name="val">Time equivalent parameter</param>
      /// <returns>Color at time</returns>
      Color Evaluate(float val)
      {
        if (val < gradientColors[0].time) return gradientColors[0].color;
        if (val > 1) val = 1;
        for(int i = 0; i < gradientColors.Length - 1; i++)
        {
          float start = gradientColors[i].time;
          float stop = gradientColors[i + 1].time;
          if (val > start && val < stop)
          {
            float t = (val - start) / (stop - start);
            return Color.Lerp(gradientColors[i].color, gradientColors[i + 1].color, t);
          }
        }
        return gradientColors[gradientColors.Length - 1].color;
      }

      float Range(float min, float max, System.Random random)
      {
        float diff = max - min;
        return min + (float)(diff * random.NextDouble());
      }
    }

    public int Size
    {
      get
      {
        return (int)Mathf.Pow(2, scale);
      }
    }

    public bool IsComplete
    {
      get
      {
        return started && jobHandle.IsCompleted;
      }
    }

    public void Generate(int seed)
    {
      DiamondSquareJob job = new DiamondSquareJob();
      pixels = new NativeArray<Color>(Size * Size, Allocator.Persistent);
      job.pixels = pixels;
      job.size = Size;
      job.seed = seed;
      job.roughness = Util.ExponentialRange(0.5f, roughness, new System.Random(seed));
      gradientColors = new NativeArray<GradientColorKey>(elevationGradient.colorKeys.Length, Allocator.Persistent);
      started = true;
      for (int i = 0; i < elevationGradient.colorKeys.Length; i++)
      {
        gradientColors[i] = elevationGradient.colorKeys[i];
      }
      job.gradientColors = gradientColors;
      jobHandle = job.Schedule();
      JobCallbackManager.RegisterJob(jobHandle, GenerationComplete);
    }

    void GenerationComplete()
    {
      if (this == null || this.gameObject == null)
      {
        return;
      }
      Texture2D thermal = new Texture2D(Size, Size);
      Color[] colors = new Color[Size * Size];
      pixels.CopyTo(colors);
      thermal.SetPixels(colors, 0);
      thermal.Apply();
      GetComponent<Renderer>().material.mainTexture = thermal;
      pixels.Dispose();
    }

    private void OnDestroy()
    {
      if (started)
      {
        jobHandle.Complete();
        if (pixels.IsCreated) pixels.Dispose();
        if(gradientColors.IsCreated) gradientColors.Dispose();
      }

    }
  }
}

How it works:

This is a normal Monobehaviour component that implements IProcedural, which is simply an interface I’ve defined for elements that can be recreated from a seed:

using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace Wx3.StarcomGameplay
{
  /// <summary>
  /// An IProcedural object can be reproducably generated from a seed. 
  /// </summary>
  public interface IProcedural 
  {
    void Generate(int seed);
    bool IsComplete { get; }
  }
}

The “elevation” Gradient defines how numbers from 0-1 get turned into colors and is set from the inspector like so:

Here’s where we run into our first challenge. A Gradient is an object, so our job can’t use it. So I had to write a struct-friendly version of the gradient’s evaluate algorithm and pass in a NativeArray (which is a special “safe” struct for accessing memory directly).

So when our Generate(seed) method is called, we pass the relevant data to the job and call “Schedule()” on it. Behind the scenes a worker thread will run the actual code while we’re free to do other stuff.

As an aside that confused me for a while, the IJob interface doesn’t actually have a Schedule() method. It’s an extension method which is a feature of C# that I just learned about.

Now our diamond square algorithm is off running and not getting in the way of more pressing tasks. Some frame soon it will hopefully finish and… then what? The Job system does not have any callbacks or events. It just has a handle which you can check “IsCompleted”. I assume this is because Jobs were designed with ECS in mind and your System would be watching individual handles. But in this case I’m using regular Monobehaviours.

One could check to see if it’s done in every Update(), but the logic for that is a bit ugly because you’ll have a block at the top that says “if(jobStarted && jobDone && !jobCleanUpComplete)” and you need to remember to add that logic to every component that uses jobs. As a general rule, I try to avoid having “run once” logic inside Update() methods.

Instead, I created a JobCallbackManager singleton that provides a single point to manage jobhandles and associate them with callbacks actions. I’ll probably revisit this approach as I progress further with ECS. Anyway, that’s what JobCallbackManager.RegisterJob() is doing.

Earlier I mentioned that you don’t need to worry as much about race conditions. Which is to say you still need to worry about them a little.

NativeArrays are not managed automatically by garbage collection. Unity won’t clean these up for you when you’re not using them, but it will at least yell at you if it notices you aren’t disposing of them.

In the GenerateComplete() callback method which copies the texture to the shader material, we also call Dispose on the pixels NativeArray.

That’s all fine, but what if our GameObject gets destroyed before generation is finished? In that case our NativeArrays won’t get disposed. So we need to add an OnDestroy method to handle that case. This is a classic example of a race condition: if we hadn’t thought about this scenario now, our code would work just fine. Then months later, maybe in production, there’d be a scenario where planet objects are deleted before the diamond square algorithm completed and we’d have a memory leak. So something to keep in mind when working with jobs: consider edge cases where the Job consumer goes away before the Job is complete.

Here is a gif of the generation in action. Note that this not the actual planet texture, but test objects showing the raw “thermal map” colors. Also, the generation really occurs at a fairly steady rate, the pauses toward the end are due to the video recording software competing with the additional threads:

Thanks for reading! I hope other devs find this post helpful.

Posted in