CUDA OpenGL4.6 Interop in Julia Lang: julia set

juliaset.jl

module Draw

import GLFW
using ModernGL
using CUDA
using CUDA: cuGraphicsGLRegisterBuffer,
    CUgraphicsResource,
    cuGraphicsMapResources,
    CUdeviceptr,
    cuGraphicsUnmapResources,
    cuGraphicsResourceGetMappedPointer_v2,
    cuGraphicsUnregisterResource,
    cuGraphicsSubResourceGetMappedArray

const cudaGraphicsMapFlagsNone =  0

mutable struct App
    window::GLFW.Window
    width::Int
    height::Int
    framebufferresized::Bool

    pixelbuffer_ref::Ref{Cuint}
    pboCUDA_ref::Ref{CUgraphicsResource}
    vertexbuffer_ref::Ref{Cuint}
    uvbuffer_ref::Ref{Cuint}
    renderedtexture_ref::Ref{Cuint}

    shaderid::Int
    texID::Int

    App(width, height) = begin
        a = new()
        a.width = width
        a.height = height
        a.framebufferresized = false
        a.pixelbuffer_ref = Ref{Cuint}(0)
        a.vertexbuffer_ref = Ref{Cuint}(0)
        a.uvbuffer_ref = Ref{Cuint}(0)
        a.renderedtexture_ref = Ref{Cuint}(0)

        return a
    end
end

struct RGB{T}
    r::T
    g::T
    b::T
end


Base.length(::RGB) = 3

function resetwindow(app)
    glViewport(0, 0, app.width, app.height)
    app.framebufferresized = false
end

function makeframesizechangedcallback(app)
    function callback(window, width, height)
        app.width = width
        app.height = height
        app.framebufferresized = true
        
        return
    end
end

function initwindow!(app)
    
    window = GLFW.CreateWindow(app.width, app.height, "CUDA GL interop")
    app.window = window
    GLFW.MakeContextCurrent(window)
    # GLFW.SetFramebufferSizeCallback(window, makeframesizechangedcallback(app))
    
    glClearColor(0.0, 0.0, 0.0, 1.0)
    
    return app
end


function initpixelbuffer!(app)
    glGenBuffers(1, app.pixelbuffer_ref)

    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, app.pixelbuffer_ref[])
    datasize = app.width * app.height * 4 * sizeof(UInt8)
    glBufferData(GL_PIXEL_UNPACK_BUFFER, datasize, C_NULL, GL_DYNAMIC_DRAW)
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0)
    pboCUDA_ref = Ref{CUgraphicsResource}()
    @assert (cuGraphicsGLRegisterBuffer(pboCUDA_ref, app.pixelbuffer_ref[], cudaGraphicsMapFlagsNone)) === nothing
    app.pboCUDA_ref = pboCUDA_ref

    return app
end

function initVAO!(app)
    vertexarrayid_ref = Ref{Cuint}()
    glGenVertexArrays(1, vertexarrayid_ref)
    glBindVertexArray(vertexarrayid_ref[])

    g_vertex_buffer_data = Vector{Float32}([
        -1.0,  1.0, 0.0,
        -1.0, -1.0, 0.0,
         1.0, -1.0, 0.0,

        -1.0,  1.0, 0.0,
         1.0,  1.0, 0.0,
         1.0,  -1.0, 0.0,
    ])
    glGenBuffers(1, app.vertexbuffer_ref)
    glBindBuffer(GL_ARRAY_BUFFER, app.vertexbuffer_ref[])
    glBufferData(GL_ARRAY_BUFFER, sizeof(g_vertex_buffer_data), g_vertex_buffer_data, GL_STATIC_DRAW)
    glBindBuffer(GL_ARRAY_BUFFER, C_NULL)
    
    g_uv_buffer_data = Vector{Float32}([
        0.0, 1.0,
        0.0, 0.0,
        1.0, 0.0,

        0.0, 1.0,
        1.0, 1.0,
        1.0, 0.0
    ])
    glGenBuffers(1, app.uvbuffer_ref)
    glBindBuffer(GL_ARRAY_BUFFER, app.uvbuffer_ref[])
    glBufferData(GL_ARRAY_BUFFER, sizeof(g_uv_buffer_data), g_uv_buffer_data, GL_STATIC_DRAW)
    glBindBuffer(GL_ARRAY_BUFFER, C_NULL)

    return app
end

function initframebuffer!(app)
    
    glEnable(GL_TEXTURE_2D)
    
    glGenTextures(1, app.renderedtexture_ref)
    glBindTexture(GL_TEXTURE_2D, app.renderedtexture_ref[])

    glTexStorage2D(GL_TEXTURE_2D, 1, GL_RGB8, app.width, app.height)
    
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST); 
    return app
end

function checkshadererror(shaderid)
    result_ref = Ref{Int32}()
    glGetShaderiv(shaderid, GL_COMPILE_STATUS, result_ref);
    if result_ref != GL_TRUE
        infolength_ref = Ref{Int32}()
        glGetShaderiv(shaderid, GL_INFO_LOG_LENGTH, infolength_ref)
        if infolength_ref[] == 0
            println("shader error but no info")
            return
        end
        errormsg = Vector{UInt8}(undef, infolength_ref[])
        glGetShaderInfoLog(shaderid, infolength_ref[], C_NULL, Ref(errormsg, 1))
        error(unsafe_string(pointer(errormsg)))
    end
end

function initShader!(app)
    VertexShaderID = glCreateShader(GL_VERTEX_SHADER)
    FragmentShaderID = glCreateShader(GL_FRAGMENT_SHADER)
    vertcode = """
    #version 460 core
    layout(location = 0) in vec3 vertexPosition_modelspace;
    layout(location = 1) in vec2 vertexUV;
    out vec2 UV;

    void main(){

        gl_Position.xyz = vertexPosition_modelspace;
        gl_Position.w = 1.0;
        UV = vertexUV;
    }
    """
    frgcode = """
    #version 460 core
    in vec2 UV;
    layout(location = 0) out vec3 color;
    uniform sampler2D tex;
    void main(){
        //color = vec3(1,0,0);
        color = texture(tex, UV).rgb;
    }
    """
    VertexSourcePointer = Ref{Ptr{UInt8}}(pointer(vertcode))
    FragmentSourcePointer  = Ref{Ptr{UInt8}}(pointer(frgcode))
    GC.@preserve vertcode frgcode begin
        glShaderSource(VertexShaderID, 1, VertexSourcePointer, C_NULL);
        glCompileShader(VertexShaderID)
        checkshadererror(VertexShaderID)
        
        glShaderSource(FragmentShaderID, 1, FragmentSourcePointer, C_NULL)
        glCompileShader(FragmentShaderID)
        checkshadererror(FragmentShaderID)
    end
    programid = glCreateProgram()
    glAttachShader(programid, VertexShaderID)
    glAttachShader(programid, FragmentShaderID)
    app.texID = glGetUniformLocation(programid, "tex")
    glLinkProgram(programid)
    glDeleteShader(VertexShaderID)
    glDeleteShader(FragmentShaderID)
    app.shaderid = programid
    return app
end

function init!(app::App)
    
    app |> 
        initwindow! |>
        initpixelbuffer! |>
        initframebuffer! |>
        initVAO! |>
        initShader!
end

function mainloop(app::App)
    
    while (!GLFW.WindowShouldClose(app.window))
        GLFW.PollEvents()
        drawframe(app)
        
        sleep(1 / 30)
    end
    return app
end

function julia_set_count(c, z, N=255)
    for k in 1:N 
        z = z^2 + c
        abs2(z) > 4 && return k
    end
    return N
end


function fragmentfunc(x, y, width, height, t)
    nx = (x * 2 - width) / width
    ny = (y * 2 - height) / height
    r = julia_set_count(complex(-0.38, 0.598 + rem(t * 0.001, 0.08)), complex(nx, ny))
    return RGB{UInt8}(r, 0x00, 0x00)
end


function kernel(texture, width, height, t)    
    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x

    for i in index:stride:length(texture)
        @inbounds texture[i] = fragmentfunc(rem(i, width), div(i, width), width, height, t)
    end 
    return
end

function cudainteropdraw(app)
    @assert cuGraphicsMapResources(1, app.pboCUDA_ref, C_NULL) === nothing
    num_bytes_ref = Ref{Csize_t}()
    pDevPtr_ref = Ref{CUdeviceptr}()
    @assert cuGraphicsResourceGetMappedPointer_v2(pDevPtr_ref, num_bytes_ref, app.pboCUDA_ref[]) === nothing
    p = Base.unsafe_convert(CuPtr{RGB{UInt8}}, pDevPtr_ref[])
    texture = Base.unsafe_wrap(CuArray, p, Int(floor(num_bytes_ref[] / 3)))
    
    threads = (32, 32)
    numblocks = (app.width ÷ 32, app.height ÷ 32)
    t = time()
    CUDA.@sync begin
        @cuda threads = threads blocks = numblocks kernel(texture, app.width, app.height, t)
    end
    
    @assert (cuGraphicsUnmapResources(1, app.pboCUDA_ref, C_NULL)) === nothing
    
end

function drawframe(app)
    
    glClear(GL_COLOR_BUFFER_BIT)

    glUseProgram(app.shaderid)
    
    println(glGetError())

    glActiveTexture(GL_TEXTURE0)
    glBindTexture(GL_TEXTURE_2D, app.renderedtexture_ref[])
    
    
    glUniform1i(app.texID, 0) 
    
    cudainteropdraw(app)

    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, app.pixelbuffer_ref[])
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, app.width, app.height, GL_RGB, GL_UNSIGNED_BYTE, C_NULL)

    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0)
    
    

    glEnableVertexAttribArray(0)
    glBindBuffer(GL_ARRAY_BUFFER, app.vertexbuffer_ref[])
    glVertexAttribPointer(
        0,3,GL_FLOAT,GL_FALSE, 0, C_NULL
    )
    glBindBuffer(GL_ARRAY_BUFFER, C_NULL)
    

    glEnableVertexAttribArray(1)
    glBindBuffer(GL_ARRAY_BUFFER, app.uvbuffer_ref[])
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 0, C_NULL)
    glBindBuffer(GL_ARRAY_BUFFER, C_NULL)

    
    glDrawArrays(GL_TRIANGLES, 0, 6)

    glDisableVertexAttribArray(0)
    glDisableVertexAttribArray(1)

    GLFW.SwapBuffers(app.window)

    return
end
    
function cleanup!(app::App)
    glDeleteBuffers(1, app.vertexbuffer_ref)
    glDeleteBuffers(1, app.uvbuffer_ref)
    glDeleteTextures(1, app.renderedtexture_ref)
    cuGraphicsUnregisterResource(app.pboCUDA_ref[])
    glDeleteBuffers(1, app.pixelbuffer_ref)
    GLFW.DestroyWindow(app.window)
    
    return app
end

function main()
    app = App(512, 512) |>
        init! |>
        mainloop |>
        cleanup!
end


end

Draw.main()