Mugichoko's blog

Mugichoko’s blog

プログラミングを中心としたメモ書き.

ModernGL奮闘記 (11) - Compute Shader -

簡単なCompute Shader(コンピュートシェーダ)を実装する.単に,画面の座標値 (uv) をRG色に割り当てるだけ.これで,大抵のことはできるようになった!

実装方針

今回は特にレンダリングをしたいわけではないので,ウィンドウを作らずにコンテクストだけを作った

f:id:Mugichoko:20211125080326p:plain
PILのImageを経由して表示される最終結

コード (Python)

# glfw_computeshader.py
from pathlib import Path
import moderngl as mgl
import math
from PIL import Image, ImageOps

def source(uri, consts):
    ''' read gl code '''
    with open(uri, 'r') as fp:
        content = fp.read()

    # feed constant values
    for key, value in consts.items():
        content = content.replace(f"%%{key}%%", str(value))
    return content

class App:
    def __init__(self) -> None:
        """
        # Window
        # create a gl window: https://moderngl-window.readthedocs.io/en/latest/reference/settings.conf.settings.html#moderngl_window.conf.Settings.WINDOW
        settings.WINDOW["class"] = "moderngl_window.context.glfw.Window"
        settings.WINDOW["gl_version"] = (4, 3)
        settings.WINDOW["size"] = (width, height)
        settings.WINDOW["aspect_ratio"] = width / height
        self.window = mglw.create_window_from_settings()
        """
        self.ctx = mgl.create_context(require=430, standalone=True)

        # Resources
        resource_dir = Path(__file__).parent.resolve() / "resources"
        shaders_dir = (resource_dir / "shaders").resolve()
        # Compute shader
        self.consts = { "X": 16, "Y": 16, "Z": 1 }
        self.compt_shader = self.ctx.compute_shader(
            source((shaders_dir / "compt_uv.glsl").resolve(), self.consts)
        )

        self.cs_out = self.ctx.texture((512, 512), components=4, dtype="f1")


    def run(self) -> None:
        self.cs_out.bind_to_image(0, read=False, write=True)
        self.compt_shader.run(
            group_x = math.ceil(self.cs_out.width / self.consts["X"]),
            group_y = math.ceil(self.cs_out.height / self.consts["Y"]),
            group_z = math.ceil(1 / self.consts["Z"])
        )

        raw = self.cs_out.read()
        img = Image.frombytes("RGBA", self.cs_out.size, raw, "raw", "RGBA")
        img = ImageOps.flip(img)
        img.show()

if __name__ == "__main__":
    app = App()
    app.run()

シェーダ

// compt_uv.glsl
#version 430

// %%VARIABLE%% will be replaced with consts by python code
// Ref: https://github.com/moderngl/moderngl/blob/master/examples/compute_shader.py
#define X %%X%%
#define Y %%Y%%
#define Z %%Z%%

layout(local_size_x = X, local_size_y = Y, local_size_z = Z) in;

layout(rgba8, binding = 0) writeonly uniform image2D img_out;

void main()
{
    const ivec2 uv = ivec2(gl_GlobalInvocationID.xy);
    const ivec2 size = imageSize(img_out);

    imageStore(img_out, uv, vec4(uv.x / float(size.x), uv.y / float(size.y), 0, 1));
}