vello/piet-gpu-hal/examples/shader/gen/collatz.msl

#pragma clang diagnostic ignored "-Wmissing-prototypes"

#include <metal_stdlib>
#include <simd/simd.h>

using namespace metal;

// Implementation of the GLSL mod() function, which is slightly different than Metal fmod()
template<typename Tx, typename Ty>
inline Tx mod(Tx x, Ty y)
{
    return x - y * floor(x / y);
}

struct PrimeIndices
{
    uint indices[1];
};

constant uint3 gl_WorkGroupSize [[maybe_unused]] = uint3(1u);

static inline __attribute__((always_inline))
uint collatz_iterations(thread uint& n)
{
    uint i = 0u;
    while (n != 1u)
    {
        if (mod(float(n), 2.0) == 0.0)
        {
            n /= 2u;
        }
        else
        {
            n = (3u * n) + 1u;
        }
        i++;
    }
    return i;
}

kernel void main0(device PrimeIndices& _57 [[buffer(0)]], uint3 gl_GlobalInvocationID [[thread_position_in_grid]])
{
    uint index = gl_GlobalInvocationID.x;
    uint param = _57.indices[index];
    uint _65 = collatz_iterations(param);
    _57.indices[index] = _65;
}