Start adding compute shader support

examples/histogram_ex.py

@@ -0,0 +1,12 @@
+import yt
+
+import yt_idv
+from yt_idv.scene_annotations.block_histogram import BlockHistogram
+
+ds = yt.load_sample("IsolatedGalaxy")
+
+rc = yt_idv.render_context(height=800, width=800, gui=True)
+sg = rc.add_scene(ds, "density", no_ghost=True)
+bh = BlockHistogram(data=sg.data_objects[0], bins=64)
+sg.annotations = sg.annotations + [bh]
+# rc.run()

yt_idv/opengl_support.py

@@ -93,13 +93,21 @@
         4: (GL.GL_UNSIGNED_BYTE, GL.GL_RGBA8, GL.GL_RGBA),
     },
     "uint32": {
-        1: (GL.GL_UNSIGNED_INT, GL.GL_R32UI, GL.GL_RED),
+        1: (GL.GL_UNSIGNED_INT, GL.GL_R32UI, GL.GL_RED_INTEGER),
         2: (GL.GL_UNSIGNED_INT, GL.GL_RG32UI, GL.GL_RG),
         3: (GL.GL_UNSIGNED_INT, GL.GL_RGB32UI, GL.GL_RGB),
         4: (GL.GL_UNSIGNED_INT, GL.GL_RGBA32UI, GL.GL_RGBA),
     },
+    "int32": {
+        1: (GL.GL_INT, GL.GL_R32I, GL.GL_RED_INTEGER),
+        2: (GL.GL_INT, GL.GL_RG32I, GL.GL_RG),
+        3: (GL.GL_INT, GL.GL_RGB32I, GL.GL_RGB),
+        4: (GL.GL_INT, GL.GL_RGBA32I, GL.GL_RGBA),
+    },
 }
 
+SKIP_MIPMAP = [GL.GL_R32UI, GL.GL_R16UI, GL.GL_R8UI, GL.GL_R32I, GL.GL_R16I, GL.GL_R8I]
+
 
 def coerce_uniform_type(val, gl_type):
     # gl_type here must be in const_types
@@ -168,6 +176,7 @@ class Texture(traitlets.HasTraits):
     channels = GLValue("r32f")
     min_filter = GLValue("linear")
     mag_filter = GLValue("linear")
+    image_mode = GLValue("write only")
 
     @traitlets.default("texture_name")
     def _default_texture_name(self):
@@ -181,6 +190,34 @@ def bind(self, target=0):
         _ = GL.glActiveTexture(TEX_TARGETS[target])
         GL.glBindTexture(self.dim_enum, 0)
 
+    @contextmanager
+    def clear(self):
+        with self.bind_as_image(0, GL.GL_WRITE_ONLY):
+            GL.glClearTexImage(self.texture_name, 0, self.channels, GL.GL_FLOAT, None)
+            yield
+
+    @contextmanager
+    def bind_as_image(self, target=0, override_mode=None):
+        _ = GL.glActiveTexture(TEX_TARGETS[target])
+        mode = override_mode or self.image_mode
+        _ = GL.glBindImageTexture(
+            target, self.texture_name, 0, False, 0, mode, self.channels
+        )
+        yield
+        _ = GL.glActiveTexture(TEX_TARGETS[target])
+        GL.glBindImageTexture(0, 0, 0, False, 0, mode, self.channels)
+
+    def read_as_image(self):
+        if len(self.data.shape) == 2:
+            channels = self.data.shape[-1]
+        else:
+            channels = 1
+        gl_type, _, type2 = TEX_CHANNELS[self.data.dtype.name][channels]
+        with self.bind():
+            data = np.zeros(self.data.shape, dtype=self.data.dtype)
+            GL.glGetTexImage(self.dim_enum, 0, type2, gl_type, data)
+        return data
+
 
 class Texture1D(Texture):
     boundary_x = TextureBoundary()
@@ -199,7 +236,7 @@ def _set_data(self, change):
             gl_type, type1, type2 = TEX_CHANNELS[data.dtype.name][channels]
             GL.glPixelStorei(GL.GL_UNPACK_ALIGNMENT, 1)
             if not isinstance(change["old"], np.ndarray):
-                GL.glTexStorage1D(GL.GL_TEXTURE_1D, 6, type1, dx)
+                GL.glTexStorage1D(GL.GL_TEXTURE_1D, 1, type1, dx)
             GL.glTexSubImage1D(GL.GL_TEXTURE_1D, 0, 0, dx, type2, gl_type, data)
             GL.glTexParameterf(GL.GL_TEXTURE_1D, GL.GL_TEXTURE_WRAP_S, self.boundary_x)
             GL.glTexParameteri(
@@ -208,7 +245,8 @@ def _set_data(self, change):
             GL.glTexParameteri(
                 GL.GL_TEXTURE_1D, GL.GL_TEXTURE_MAG_FILTER, self.mag_filter
             )
-            GL.glGenerateMipmap(GL.GL_TEXTURE_1D)
+            if type1 not in SKIP_MIPMAP:
+                GL.glGenerateMipmap(GL.GL_TEXTURE_1D)
 
 
 class ColormapTexture(Texture1D):
@@ -268,7 +306,8 @@ def _set_data(self, change):
             GL.glTexParameteri(
                 GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MAG_FILTER, self.mag_filter
             )
-            GL.glGenerateMipmap(GL.GL_TEXTURE_2D)
+            if type1 not in SKIP_MIPMAP:
+                GL.glGenerateMipmap(GL.GL_TEXTURE_2D)
 
 
 class TransferFunctionTexture(Texture2D):
@@ -329,7 +368,8 @@ def _set_data(self, change):
             GL.glTexParameteri(
                 GL.GL_TEXTURE_3D, GL.GL_TEXTURE_MAG_FILTER, self.mag_filter
             )
-            GL.glGenerateMipmap(GL.GL_TEXTURE_3D)
+            if type1 not in SKIP_MIPMAP:
+                GL.glGenerateMipmap(GL.GL_TEXTURE_3D)
 
 
 class VertexAttribute(traitlets.HasTraits):

yt_idv/scene_annotations/block_histogram.py

@@ -0,0 +1,64 @@
+import numpy as np
+import traitlets
+from OpenGL import GL
+
+from yt_idv.opengl_support import Texture1D
+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)
+    output_data = traitlets.Instance(Texture1D)
+
+    bins = traitlets.CInt(64)
+    min_val = traitlets.CFloat(0.0)
+    max_val = traitlets.CFloat(1.0)
+
+    _recompute = True
+
+    @traitlets.default("output_data")
+    def _default_output_data(self):
+        return Texture1D(data=np.zeros(self.bins, dtype="u4"), channels="r32ui")
+
+    def _set_uniforms(self, scene, shader_program):
+        pass
+
+    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):
+        # We need a place to dump our stuff.
+        self.visible = False
+        self.output_data.clear()
+        total = 0
+        with self.output_data.bind_as_image(2):
+            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.
+                total += bitmap_tex.data.size
+                with tex.bind(target=0):
+                    with bitmap_tex.bind(target=1):
+                        GL.glUseProgram(program.program)
+                        self._set_compute_uniforms(scene, program)
+                        # 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(
+                            bitmap_tex.data.shape[0],
+                            bitmap_tex.data.shape[1],
+                            bitmap_tex.data.shape[2],
+                        )
+                        GL.glMemoryBarrier(GL.GL_SHADER_IMAGE_ACCESS_BARRIER_BIT)
+
+    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

@@ -13,6 +13,7 @@
 )
 from yt_idv.scene_data.base_data import SceneData
 from yt_idv.shader_objects import (
+    ComputeShaderProgram,
     ShaderProgram,
     ShaderTrait,
     component_shaders,
@@ -49,18 +50,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)
 
@@ -166,6 +171,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
@@ -201,6 +210,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"]
@@ -234,6 +247,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:
@@ -270,9 +293,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: