SIMD Mandelbrot calculation

Question

I was messing around with GPU compute shaders the other day and created a Mandelbrot shader. Unfortunately, Metal doesn't support double-precision in compute shaders, so beyond a certain zoom level, I need to switch back to the CPU. In doing so, I decided to try writing SIMD code for the calculations to make it faster.

In the code below I'm using AVX512 instructions, and I do get a speedup over the scalar code. I break the image into 64x64 pixel tiles and farm them out to available cores. For the scalar code on one particular test image, the average time to calculate a tile is 0.757288 seconds. For the SIMD version below it's 0.466437. That's about a 33% increase, which is OK. Given that I'm calculating 8 times as many pixels at once, I was hoping for more.

These are some useful types I use in the code.

#include <immintrin.h>
typedef struct RGBA8Pixel {
    uint8_t red;
    uint8_t green;
    uint8_t blue;
    uint8_t alpha;
} RGBA8Pixel;

typedef union intVec8 {
    __m512i ivec;
    int64_t vec[8];
} intVec8;

typedef union doubleVec8 {
    __m512d dvec;
    double vec[8];
} doubleVec8;

And here's my function for calculating 1 64x64 tile:

- (void)calculateSIMDFromRow:(int)startPixelRow
                       toRow:(int)endPixelRow
                     fromCol:(int)startPixelCol
                       toCol:(int)endPixelCol;
{
    if (!_keepRendering)
    {
        return;
    }

    const doubleVec8 k0s = {
        .vec[0] = 0.0,
        .vec[1] = 0.0,
        .vec[2] = 0.0,
        .vec[3] = 0.0,
        .vec[4] = 0.0,
        .vec[5] = 0.0,
        .vec[6] = 0.0,
        .vec[7] = 0.0,
    };

    const intVec8   k1s = {
        .vec[0] = 1,
        .vec[1] = 1,
        .vec[2] = 1,
        .vec[3] = 1,
        .vec[4] = 1,
        .vec[5] = 1,
        .vec[6] = 1,
        .vec[7] = 1,
    };

    const doubleVec8    k2s = {
        .vec[0] = 2.0,
        .vec[1] = 2.0,
        .vec[2] = 2.0,
        .vec[3] = 2.0,
        .vec[4] = 2.0,
        .vec[5] = 2.0,
        .vec[6] = 2.0,
        .vec[7] = 2.0,
    };

    const doubleVec8    k4s = {
        .vec[0] = 4.0,
        .vec[1] = 4.0,
        .vec[2] = 4.0,
        .vec[3] = 4.0,
        .vec[4] = 4.0,
        .vec[5] = 4.0,
        .vec[6] = 4.0,
        .vec[7] = 4.0,
    };

    UInt64      maxIterations   = [self maxIterations];
    NSSize      viewportSize    = [self viewportSize];
    for (int row = startPixelRow; (row < endPixelRow) && (_keepRendering); ++row)
    {
        RGBA8Pixel* nextPixel   = _outputBitmap + (row * (int)viewportSize.width) + startPixelCol;
        double      yCoord      = _yCoords [ row ];
        doubleVec8  yCoords;
        for (int i = 0; i < 8; i++)
        {
            yCoords.vec [ i ] = yCoord;
        }
        double*     nextXCoord  = &_xCoords [ startPixelCol ];
        for (int col = startPixelCol; (col < endPixelCol) && (_keepRendering); col += 8)
        {
            __m512d as = _mm512_load_pd(nextXCoord);
            nextXCoord += 8;
            __m512d bs = yCoords.dvec;
            __m512d cs = as;
            __m512d ds = bs;

            UInt64 scalarIters = 1;
            __m512i iterations  = k1s.ivec;
            __m512d dists       = k0s.dvec;
            __mmask8 allDone    = 0;
            while ((allDone != 0xFF) && (_keepRendering))
            {
                // newA = a * a - b * b + c
                __m512d newA;
                __m512d newB;
                newA = _mm512_mul_pd(as, as);
                newA = _mm512_sub_pd(newA, _mm512_mul_pd(bs, bs));
                newA = _mm512_add_pd(newA, cs);

                //double    newB    = 2 * a * b + d;
                newB = _mm512_mul_pd(_mm512_mul_pd(k2s.dvec, as), bs);
                newB = _mm512_add_pd(newB, ds);

                as = newA;
                bs = newB;

                dists = _mm512_mul_pd(newB, newB);
                dists = _mm512_add_pd(_mm512_mul_pd(newA, newA), dists);
                __mmask8 escaped = _mm512_cmplt_pd_mask(dists, k4s.dvec);

                iterations = _mm512_mask_add_epi64(iterations, escaped, iterations, k1s.ivec);
                scalarIters++;
                __mmask8 hitMaxIterations = (scalarIters == maxIterations) ? 0xFF : 0;

                allDone = ~escaped | hitMaxIterations;
            }

            intVec8 iters = { .ivec = iterations };
            for (int i = 0; i < 8; i++)
            {
                UInt64  nextIteration = iters.vec [ i ];
                if (nextIteration == maxIterations)
                {
                    *nextPixel = kBlack;
                }
                else
                {
                    *nextPixel = kPalette [ nextIteration % kPaletteSize ];
                }
                nextPixel++;
            }
        }
    }
}

I'm new to Intel SIMD instructions and frankly find them quite confusing. If there are better ways to do any of the above, please let me know. I tried using the fused multiply-add and multiply-add-negate instructions, and they made the code significantly slower than using 2 or 3 separate instructions, unfortunately.

I'm working on macOS using Xcode 10.2.1 using the Intel data types and intrinsics found in <immintrin.h>.

Answer 1

A couple of thoughts:

1. You say

   Given that I'm calculating 8 times as many pixels at once, I was hoping for more.

   Yes, simd delivers some pretty spectacular performance improvements when doing vector/matrix operations, but for the Mandelbrot, all you’re doing is elementwise addition and multiplication, so you should see improvements, but nothing approaching 8× for these simple elementwise calculations. In my tests, the simd achieved just about twice the performance of the scalar rendition on both i9 Mac and iPhone Xs Max.

2. Two algorithmic observations:

   • I notice that you’re squaring as and bs twice, once during the algorithm and again in the escaping test. I’d suggest refactoring this so you use the results of this squaring for both the algorithm and for the escaping test.

   • I notice that you’re doing vector multiplication for 2 in the 2×a×b portion. I used the vector × scalar product rather than the vector elementwise product. That might be a tad faster.

3. If you’d like to eliminate those unintuitive _mm512_xxx calls, you might consider the simd library, which is part of the Accelerate framework. This is a higher level of abstraction and, especially in Swift, you end up with very natural looking code which is, IMHO, easier to read.

4. Conceptually, it should be noted that the simd performance gains are going to be offset by those boundary cases where some pixels have escaped and others haven’t, as you’re going to be calculating iterations for all eight pixels in the vector, including those that have already escaped.

   This probably doesn’t have a material impact on the performance, but it’s worth noting, especially if dealing with cases where dealing with intricate portions of the Mandelbrot set (which are the most interesting parts).

---

For what it’s worth, this is what a Swift simd rendition might look like:

import simd

func calculate(real: simd_double8, imaginary: simd_double8) -> simd_double8 {
    var zReal = real // simd_double8.zero
    var zImaginary = imaginary // simd_double8.zero

    let thresholds = simd_double8(repeating: 4)
    let maxIterations = 10_000.0

    var notEscaped = SIMDMask<SIMD8<Double.SIMDMaskScalar>>(repeating: true)
    let isDone = SIMDMask<SIMD8<Double.SIMDMaskScalar>>(repeating: false)

    var currentIterations = 0.0
    var iterations = simd_double8.zero

    repeat {                                                    // z = z^2 + c
        currentIterations += 1.0
        iterations.replace(with: currentIterations, where: notEscaped)

        let zRealSquared = zReal * zReal
        let zImaginarySquared = zImaginary * zImaginary

        zImaginary = 2.0 * zReal * zImaginary + imaginary       // 2 × zr × zi + ci
        zReal = zRealSquared - zImaginarySquared + real         // zr^2 - zi^2 + cr

        notEscaped = zRealSquared + zImaginarySquared .< thresholds
    } while notEscaped != isDone && currentIterations < maxIterations

    iterations.replace(with: 0, where: notEscaped)

    return iterations
}

What’s nice about that, is that it’s very similar to the scalar rendition, free of cryptic method references. For example, here is the scalar version:

func calculate(real: Double, imaginary: Double) -> Int {
    var zReal = real
    var zImaginary = imaginary

    let thresholds = 4.0
    let maxIterations = 10_000

    var notEscaped = false

    var currentIterations = 0

    repeat {                                                    // z = z^2 + c
        currentIterations += 1

        let zRealSquared = zReal * zReal
        let zImaginarySquared = zImaginary * zImaginary

        zImaginary = 2.0 * zReal * zImaginary + imaginary       // 2 × zr × zi + ci
        zReal = zRealSquared - zImaginarySquared + real         // zr^2 - zi^2 + cr

        notEscaped = zRealSquared + zImaginarySquared < thresholds
    } while notEscaped && currentIterations < maxIterations

    return currentIterations >= maxIterations ? 0 : currentIterations
}

Answer 2

I tried using the fused multiply-add and multiply-add-negate instructions, and they made the code significantly slower than using 2 or 3 separate instructions, unfortunately.

This is usually indicative of a latency bottleneck,

That's about a 33% increase, which is OK. Given that I'm calculating 8 times as many pixels at once, I was hoping for more.

And so is this.

Actually this is expected for Mandelbrot, because it based on iterated function application, so it inherently has a non-trivial loop-carried dependency. Floating point operations on Intel have a high throughput, but they are still slow operations in the sense of having a high latency compared to their throughput. On Skylake X (which I guess you are using, from your use of AVX512), an FMA takes 4 cycles but the processor can start two of them every cycle. So if they are too "tied up" (and with FMA things get even more tied up, because every FMA is waiting for 3 instead of 2 inputs to be ready), it might be that there is always some floating point operation being executed, but actually on Skylake X we would want 8 operations to be "busy" at any time. On Haswell it was even worse, taking 10 "overlapping" FMAs to saturate the floating point units.

That situation can be improved by interleaving the calculation of several (4? 8?) independent rows of 8 pixels, though an unfortunate side-effect of this is that it would also "round up" the loop count to the max count among all pixels in the block. That already happens at a smaller scale now but it will get worse, and suppress the potential gain from doing this.