VulkanSimulation/shaders/normal.comp

#version 450
#extension GL_EXT_shader_atomic_float : enable



layout (local_size_x = 256) in;

struct Vertex {
    vec3 position;
    vec3 color;
    uvec3 normal;
};

struct Face {
    uint a;
    uint b;
    uint c;
};

layout (std430, set = 0, binding = 0) buffer VertexBuffer {
    Vertex vertices[];
};
layout (std430, set = 0, binding = 1) buffer FaceBuffer {
    Face faces[];
};

layout (std140, set = 0, binding = 5) uniform Sizes {
    uint vertexCount;
    uint faceCount;
};


layout (push_constant) uniform PushConstants {
    uint state;
};

void atomicAddVec3(uint vID, vec3 add){
    for (int i = 0; i < 3; i++) {
        uint expected_memory = vertices[vID].normal[i];
        float floatInput = uintBitsToFloat(vertices[vID].normal[i]) + add[i];
        uint actual_content = atomicCompSwap(vertices[vID].normal[i], expected_memory, floatBitsToUint(floatInput));
        while (actual_content != expected_memory) {
            expected_memory = actual_content;
            floatInput = uintBitsToFloat(expected_memory) + add[i];
            actual_content = atomicCompSwap(vertices[vID].normal[i], expected_memory, floatBitsToUint(floatInput));
        }
    }
}

void reset(uint vID){
    vertices[vID].normal = floatBitsToUint(vec3(0, 0, 0));
}

void accumulate(uint fID){
    Face f = faces[fID];

    Vertex v1 = vertices[f.a];
    Vertex v2 = vertices[f.b];
    Vertex v3 = vertices[f.c];

    vec3 weightedNormal = cross(v3.position - v1.position, v2.position - v1.position);
    atomicAddVec3(f.a, weightedNormal);
    atomicAddVec3(f.b, weightedNormal);
    atomicAddVec3(f.c, weightedNormal);
}

void norm(uint vID){
    vertices[vID].normal = floatBitsToUint(normalize(uintBitsToFloat(vertices[vID].normal)));
}

void main() {
    uint id = gl_GlobalInvocationID.x;
    switch (state){
        case 0:
            if (id < vertexCount){
                reset(id);
            }
            break;
        case 1:
            if (id < faceCount){
                accumulate(id);
            }
            break;
        case 2:
            if (id < vertexCount){
                norm(id);
            }
            break;
    }
}