Start adding compute shader support

yt_idv/opengl_support.py

@@ -523,6 +523,66 @@ def __iter__(self, target=0):
         GL.glBindTexture(GL.GL_TEXTURE_3D, 0)
 
 
+class ImageTexture(traitlets.HasTraits):
+    image_texture_name = traitlets.CInt(-1)
+    # This is, generally speaking, read-only.  We don't want to write it from
+    # these objects.  Instead, we will read it from a texture object that is
+    # similarly bound.
+    data = traittypes.Array(None, allow_none=True)
+    channels = GLValue("r32f")
+    min_filter = GLValue("linear")
+    mag_filter = GLValue("linear")
+    image_mode = GLValue("write only")
+
+    @traitlets.default("image_texture_name")
+    def _default_image_texture_name(self):
+        return GL.glGenTextures(1)
+
+    @contextmanager
+    def clear(self, target=0):
+        with self.bind(0, GL.GL_WRITE_ONLY):
+            GL.glClearTexImage(
+                self.image_texture_name, 0, self.channels, GL.GL_FLOAT, None
+            )
+            yield
+
+    @contextmanager
+    def bind(self, target=0, override_mode=None):
+        _ = GL.glActiveTexture(TEX_TARGETS[target])
+        mode = override_mode or self.image_mode
+        _ = GL.glBindImageTexture(
+            0, self.image_texture_name, 0, False, 0, mode, self.channels
+        )
+        yield
+        _ = GL.glActiveTexture(TEX_TARGETS[target])
+        GL.glBindImageTexture(0, 0, 0, False, 0, 0, 0)
+
+
+# These require different semantics for creating the textures, as we still need
+# to allocate storage for them on the texture unit.
+
+
+class ImageTexture1D(ImageTexture):
+    boundary_x = TextureBoundary()
+    dims = 1
+    dims_enum = GLValue("texture 1d")
+
+
+class ImageTexture2D(ImageTexture):
+    boundary_x = TextureBoundary()
+    boundary_y = TextureBoundary()
+    dims = 2
+    dims_enum = GLValue("texture 2d")
+
+
+class ImageTexture3D(ImageTexture):
+    boundary_x = TextureBoundary()
+    boundary_y = TextureBoundary()
+    boundary_z = TextureBoundary()
+    dims = 3
+    dims_enum = GLValue("texture 3d")
+
+
 def compute_box_geometry(left_edge, right_edge):
     move = get_translate_matrix(*left_edge)
     width = right_edge - left_edge

yt_idv/scene_annotations/block_histogram.py

@@ -0,0 +1,39 @@
+import numpy as np
+import traitlets
+from OpenGL import GL
+
+from yt_idv.scene_annotations.base_annotation import SceneAnnotation
+from yt_idv.scene_data.block_collection import BlockCollection
+
+
+class BlockHistogram(SceneAnnotation):
+    """
+    A class that computes and displays a histogram of block data.
+    """
+
+    name = "block_histogram"
+    data = traitlets.Instance(BlockCollection)
+    bins = traitlets.CInt(64)
+    min_val = traitlets.CFloat(0.0)
+    max_val = traitlets.CFloat(1.0)
+
+    def _set_compute_uniforms(self, scene, shader_program):
+        shader_program._set_uniform("min_val", self.min_val)
+        shader_program._set_uniform("max_val", self.max_val)
+        shader_program._set_uniform("bins", self.bins)
+
+    def compute(self, scene, program):
+        for _tex_ind, tex, bitmap_tex in self.data.viewpoint_iter(scene.camera):
+            # We now need to bind our textures.  We don't care about positions.
+            with tex.bind(target=0):
+                with bitmap_tex.bind(target=1):
+                    # This will need to be carefully chosen based on our
+                    # architecture, I guess.  That aspect of running compute
+                    # shaders, CUDA, etc, is one of my absolute least favorite
+                    # parts.
+                    GL.glDispatchCompute(self.bins, 1, 1)
+
+    def draw(self, scene, program):
+        # This will probably need to have somewhere to draw the darn thing.  So
+        # we'll need display coordinates, size, etc.
+        pass

yt_idv/scene_components/base_component.py

@@ -12,6 +12,7 @@
 )
 from yt_idv.scene_data.base_data import SceneData
 from yt_idv.shader_objects import (
+    ComputeShaderProgram,
     ShaderProgram,
     ShaderTrait,
     component_shaders,
@@ -44,18 +45,22 @@ class SceneComponent(traitlets.HasTraits):
     clear_region = traitlets.Bool(False)
 
     render_method = traitlets.Unicode(allow_none=True)
+    compute_shader = ShaderTrait(allow_none=True).tag(shader_type="compute")
     fragment_shader = ShaderTrait(allow_none=True).tag(shader_type="fragment")
     geometry_shader = ShaderTrait(allow_none=True).tag(shader_type="geometry")
     vertex_shader = ShaderTrait(allow_none=True).tag(shader_type="vertex")
     fb = traitlets.Instance(Framebuffer)
     colormap_fragment = ShaderTrait(allow_none=True).tag(shader_type="fragment")
     colormap_vertex = ShaderTrait(allow_none=True).tag(shader_type="vertex")
     colormap = traitlets.Instance(ColormapTexture)
+    _compute_program = traitlets.Instance(ShaderProgram, allow_none=True)
     _program1 = traitlets.Instance(ShaderProgram, allow_none=True)
     _program2 = traitlets.Instance(ShaderProgram, allow_none=True)
     _program1_invalid = True
     _program2_invalid = True
     _cmap_bounds_invalid = True
+    _recompute = False
+    _compute_program_invalid = True
 
     display_name = traitlets.Unicode(allow_none=True)
 
@@ -161,6 +166,10 @@ def _add_initial_isolayer(self, change):
     def _fb_default(self):
         return Framebuffer()
 
+    @traitlets.observe("compute_shader")
+    def _change_compute(self, change):
+        self._compute_program_invalid = True
+
     @traitlets.observe("fragment_shader")
     def _change_fragment(self, change):
         # Even if old/new are the same
@@ -196,6 +205,10 @@ def _default_colormap(self):
         cm.colormap_name = "arbre"
         return cm
 
+    @traitlets.default("compute_shader")
+    def _compute_shader_default(self):
+        return component_shaders[self.name][self.render_method]["compute"]
+
     @traitlets.default("vertex_shader")
     def _vertex_shader_default(self):
         return component_shaders[self.name][self.render_method]["first_vertex"]
@@ -229,6 +242,16 @@ def _default_base_quad(self):
         )
         return bq
 
+    @property
+    def compute_program(self):
+        if self._compute_program_invalid:
+            if self._compute_program is not None:
+                self._compute_program.delete_program()
+            self._compute_program = ComputeShaderProgram(self.compute_shader)
+            self._compute_program_invalid = False
+            self._recompute = True
+        return self._compute_program
+
     @property
     def program1(self):
         if self._program1_invalid:
@@ -265,9 +288,20 @@ def _set_iso_uniforms(self, p):
         p._set_uniform("iso_min", float(self.data.min_val))
         p._set_uniform("iso_max", float(self.data.max_val))
 
+    def compute(self, scene, program):
+        pass
+
+    def _set_compute_uniforms(self, scene, program):
+        pass
+
     def run_program(self, scene):
         # Store this info, because we need to render into a framebuffer that is the
         # right size.
+        if self._recompute:
+            with self.compute_program.enable() as p:
+                self._set_compute_uniforms(scene, p)
+                self.compute(scene, p)
+            self._recompute = False
         x0, y0, w, h = GL.glGetIntegerv(GL.GL_VIEWPORT)
         GL.glViewport(0, 0, w, h)
         if not self.visible:

yt_idv/shader_objects.py

@@ -228,6 +228,48 @@ def enable(self):
         GL.glUseProgram(0)
 
 
+class ComputeShaderProgram(ShaderProgram):
+    """
+    Wrapper class that compiles and links compute shaders.
+
+    This has very little shared code with a ShaderProgram, and is designed to be
+    used for compute shaders only.
+
+    Parameters
+    ----------
+
+    compute_shader : string
+        or :class:`yt.visualization.volume_rendering.shader_objects.ComputeShader`
+        The vertex shader used in the Interactive Data Visualization pipeline.
+    """
+
+    def __init__(self, compute_shader):
+        self.link(compute_shader)
+        self._uniform_funcs = OrderedDict()
+
+    def link(self, compute_shader):
+        """
+        Links the compute shader to the program.
+
+        Parameters
+        ----------
+        compute_shader : string
+            or :class:`yt.visualization.volume_rendering.shader_objects.ComputeShader`
+            The compute shader used in the Interactive Data Visualization pipeline.
+        """
+        self.program = GL.glCreateProgram()
+        if not isinstance(compute_shader, Shader):
+            compute_shader = Shader(source=compute_shader)
+        self.compute_shader = compute_shader
+        GL.glAttachShader(self.program, self.compute_shader.shader)
+        GL.glLinkProgram(self.program)
+        result = GL.glGetProgramiv(self.program, GL.GL_LINK_STATUS)
+        if not result:
+            raise RuntimeError(GL.glGetProgramInfoLog(self.program))
+        self.compute_shader.delete_shader()
+        self.introspect()
+
+
 class Shader(traitlets.HasTraits):
     """
     Creates a shader from source
@@ -252,7 +294,9 @@ class Shader(traitlets.HasTraits):
     source = traitlets.Any()
     shader_name = traitlets.CUnicode()
     info = traitlets.CUnicode()
-    shader_type = traitlets.CaselessStrEnum(("vertex", "fragment", "geometry"))
+    shader_type = traitlets.CaselessStrEnum(
+        ("vertex", "fragment", "geometry", "compute")
+    )
     blend_func = traitlets.Tuple(
         GLValue(), GLValue(), default_value=("src alpha", "dst alpha")
     )
@@ -382,7 +426,7 @@ def _validate_shader(shader_type, value, allow_null=True):
 
 class ShaderTrait(traitlets.TraitType):
     default_value = None
-    info_text = "A shader (vertex, fragment or geometry)"
+    info_text = "A shader (vertex, fragment, geometry or compute)"
 
     def validate(self, obj, value):
         if isinstance(value, str):