"""GPU FFT ocean compute chain (design D11, RM-E7gpu).
Moves the per-frame ocean hot path off the CPU: instead of a ``numpy`` inverse
FFT + a per-frame texture upload (:class:`OceanFFT.evaluate`), three compute
shaders run entirely on the GPU each frame, writing straight into the sampled
displacement / slope texture-2D-arrays the ocean shader reads (so ``ocean.vert``
/ ``ocean.frag`` are unchanged):
1. ``ocean_spectrum.comp`` -- time-evolve the frozen ``h0`` spectrum (uploaded
once per :meth:`OceanFFT.configure`) into four packed complex IFFT inputs per
cascade (two real fields per complex transform, design D11).
2. ``ocean_fft.comp`` -- radix-2 Stockham autosort inverse FFT: ``log2(N)`` row
stages then ``log2(N)`` column stages, ping-ponging two RGBA32F storage
image-2D-arrays (one dispatch per stage covers every cascade + packed field).
3. ``ocean_assemble.comp`` -- unpack + scale the transforms, fold in choppiness,
accumulate the Jacobian folding foam, and write the ``b2`` displacement +
``b3`` slope arrays the forward ocean pass samples.
Owns the FFT ping-pong images, the static spectrum SSBO and the two output
arrays; :class:`OceanPass` binds the output views into its render descriptor set.
Reuses the project compute idiom (glslc ``.comp`` via
:func:`create_compute_pipeline`, storage images, push constants, ping-pong) from
the SSR / SSGI / particle passes. Everything here only runs when a visible
``OceanSurface3D`` is in the scene, so with no ocean the frame is byte-identical.
"""
from __future__ import annotations
import logging
from typing import Any
import numpy as np
import vulkan as vk
from ..gpu.descriptors import (
DescriptorWriteBatch,
allocate_descriptor_set,
create_descriptor_set_layout,
create_pool_for_types,
)
from ..gpu.memory import _find_memory_type, create_buffer, upload_numpy
from ..gpu.pipeline_compute import create_compute_pipeline
from .ocean_fft import OceanFFT
__all__ = ["OceanCompute"]
log = logging.getLogger(__name__)
_SPEC_FMT = vk.VK_FORMAT_R32G32B32A32_SFLOAT # complex ping-pong (RG used); storage-guaranteed
_OUT_FMT = vk.VK_FORMAT_R16G16B16A16_SFLOAT # displacement / slope arrays (sampled by ocean shader)
_NUM_FIELDS = 4 # packed complex transforms per cascade (8 real fields, 2 per complex FFT)
[docs]
class OceanCompute:
"""GPU compute FFT: time-evolve + inverse-FFT + assemble the ocean fields."""
def __init__(self, engine: Any) -> None:
self._engine = engine
self._ready = False
# Shared pipelines (created once in setup).
self._spectrum_pipe: Any = None
self._spectrum_layout: Any = None
self._spectrum_module: Any = None
self._spectrum_set_layout: Any = None
self._fft_pipe: Any = None
self._fft_layout: Any = None
self._fft_module: Any = None
self._fft_set_layout: Any = None
self._assemble_pipe: Any = None
self._assemble_layout: Any = None
self._assemble_module: Any = None
self._assemble_set_layout: Any = None
# Per-size resources (recreated on quality change).
self._size = 0
self._cascades = 0
self._spec_buf: Any = None
self._spec_mem: Any = None
self._ping: list[dict[str, Any]] = [] # two {image, mem, view} complex buffers
self._disp: dict[str, Any] = {}
self._grad: dict[str, Any] = {}
self._pool: Any = None
self._spectrum_set: Any = None
self._fft_set_ab: Any = None # src=ping0 -> dst=ping1
self._fft_set_ba: Any = None # src=ping1 -> dst=ping0
self._assemble_set: Any = None
self._spectrum_generation = -1
self._reset_foam = True
self._disp_layout = vk.VK_IMAGE_LAYOUT_UNDEFINED
self._grad_layout = vk.VK_IMAGE_LAYOUT_UNDEFINED
# --------------------------------------------------------------- lifecycle
[docs]
def setup(self) -> None:
"""Create the three shared compute pipelines (size-independent)."""
e = self._engine
device = e.ctx.device
cs = vk.VK_SHADER_STAGE_COMPUTE_BIT
sb = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
si = vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE
self._spectrum_set_layout = create_descriptor_set_layout(device, [(0, sb, cs, 1), (1, si, cs, 1)])
self._fft_set_layout = create_descriptor_set_layout(device, [(0, si, cs, 1), (1, si, cs, 1)])
self._assemble_set_layout = create_descriptor_set_layout(
device, [(0, si, cs, 1), (1, si, cs, 1), (2, si, cs, 1)]
)
sd = e.shader_dir
self._spectrum_pipe, self._spectrum_layout, self._spectrum_module = create_compute_pipeline(
device, sd / "ocean_spectrum.comp", [self._spectrum_set_layout], 8
)
self._fft_pipe, self._fft_layout, self._fft_module = create_compute_pipeline(
device, sd / "ocean_fft.comp", [self._fft_set_layout], 12
)
self._assemble_pipe, self._assemble_layout, self._assemble_module = create_compute_pipeline(
device, sd / "ocean_assemble.comp", [self._assemble_set_layout], 32
)
self._ready = True
log.debug("Ocean compute initialised (spectrum + FFT + assemble pipelines)")
[docs]
@property
def disp_view(self) -> Any:
return self._disp.get("view")
[docs]
@property
def grad_view(self) -> Any:
return self._grad.get("view")
[docs]
@property
def size(self) -> int:
return self._size
[docs]
@property
def cascades(self) -> int:
return self._cascades
# ------------------------------------------------------------- resources
def _create_array_image(self, size: int, layers: int, fmt: int, usage: int) -> dict[str, Any]:
device = self._engine.ctx.device
phys = self._engine.ctx.physical_device
info = vk.VkImageCreateInfo(
imageType=vk.VK_IMAGE_TYPE_2D,
format=fmt,
extent=vk.VkExtent3D(width=size, height=size, depth=1),
mipLevels=1,
arrayLayers=layers,
samples=vk.VK_SAMPLE_COUNT_1_BIT,
tiling=vk.VK_IMAGE_TILING_OPTIMAL,
usage=usage,
sharingMode=vk.VK_SHARING_MODE_EXCLUSIVE,
initialLayout=vk.VK_IMAGE_LAYOUT_UNDEFINED,
)
image = vk.vkCreateImage(device, info, None)
reqs = vk.vkGetImageMemoryRequirements(device, image)
mem = vk.vkAllocateMemory(
device,
vk.VkMemoryAllocateInfo(
allocationSize=reqs.size,
memoryTypeIndex=_find_memory_type(phys, reqs.memoryTypeBits, vk.VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT),
),
None,
)
vk.vkBindImageMemory(device, image, mem, 0)
view = vk.vkCreateImageView(
device,
vk.VkImageViewCreateInfo(
image=image,
viewType=vk.VK_IMAGE_VIEW_TYPE_2D_ARRAY,
format=fmt,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=layers,
),
),
None,
)
return {"image": image, "mem": mem, "view": view, "layers": layers}
def _destroy_array_image(self, res: dict[str, Any]) -> None:
device = self._engine.ctx.device
if res.get("view"):
vk.vkDestroyImageView(device, res["view"], None)
if res.get("image"):
vk.vkDestroyImage(device, res["image"], None)
if res.get("mem"):
vk.vkFreeMemory(device, res["mem"], None)
res.clear()
[docs]
def resize(self, fft: OceanFFT, size: int, cascades: int) -> None:
"""(Re)allocate the size-dependent resources for an ``N x N`` / ``cascades`` ocean.
Recreates the two complex ping-pong arrays, the two output arrays, the
static spectrum SSBO and the descriptor sets, then transitions the images
to their steady-state layouts. Foam is reset on the next dispatch (the new
output image starts undefined).
"""
self._destroy_size_resources()
device = self._engine.ctx.device
phys = self._engine.ctx.physical_device
self._size = int(size)
self._cascades = int(cascades)
layers = self._cascades * _NUM_FIELDS
st = vk.VK_IMAGE_USAGE_STORAGE_BIT
self._ping = [self._create_array_image(size, layers, _SPEC_FMT, st) for _ in range(2)]
self._disp = self._create_array_image(size, self._cascades, _OUT_FMT, st | vk.VK_IMAGE_USAGE_SAMPLED_BIT)
self._grad = self._create_array_image(size, self._cascades, _OUT_FMT, st | vk.VK_IMAGE_USAGE_SAMPLED_BIT)
# Static spectrum SSBO (h0 + k-vectors), uploaded once per configure().
floats = self._cascades * size * size * 8
self._spec_buf, self._spec_mem = create_buffer(
device,
phys,
floats * 4,
vk.VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
vk.VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | vk.VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
)
self._spectrum_generation = -1 # force re-upload
# The ping-pong complex arrays live permanently in GENERAL (written +
# read as storage images every stage); one-time transition off UNDEFINED.
for res in self._ping:
self._one_time_transition(res["image"], layers, vk.VK_IMAGE_LAYOUT_UNDEFINED, vk.VK_IMAGE_LAYOUT_GENERAL)
self._disp_layout = vk.VK_IMAGE_LAYOUT_UNDEFINED
self._grad_layout = vk.VK_IMAGE_LAYOUT_UNDEFINED
self._reset_foam = True
self._create_descriptor_sets()
def _create_descriptor_sets(self) -> None:
device = self._engine.ctx.device
si = vk.VK_DESCRIPTOR_TYPE_STORAGE_IMAGE
sb = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
self._pool = create_pool_for_types(device, {si: 8, sb: 1}, max_sets=4)
self._spectrum_set = allocate_descriptor_set(device, self._pool, self._spectrum_set_layout)
self._fft_set_ab = allocate_descriptor_set(device, self._pool, self._fft_set_layout)
self._fft_set_ba = allocate_descriptor_set(device, self._pool, self._fft_set_layout)
self._assemble_set = allocate_descriptor_set(device, self._pool, self._assemble_set_layout)
gen = vk.VK_IMAGE_LAYOUT_GENERAL
a, b = self._ping[0]["view"], self._ping[1]["view"]
spec_size = self._cascades * self._size * self._size * 8 * 4
with DescriptorWriteBatch(device) as w:
w.ssbo(self._spectrum_set, 0, self._spec_buf, spec_size)
w.storage_image(self._spectrum_set, 1, a, gen) # spectrum writes ping0
w.storage_image(self._fft_set_ab, 0, a, gen)
w.storage_image(self._fft_set_ab, 1, b, gen)
w.storage_image(self._fft_set_ba, 0, b, gen)
w.storage_image(self._fft_set_ba, 1, a, gen)
w.storage_image(self._assemble_set, 0, a, gen) # final FFT result lands in ping0
w.storage_image(self._assemble_set, 1, self._disp["view"], gen)
w.storage_image(self._assemble_set, 2, self._grad["view"], gen)
def _one_time_transition(self, image: Any, layers: int, old: int, new: int) -> None:
e = self._engine
device = e.ctx.device
cmd = vk.vkAllocateCommandBuffers(
device,
vk.VkCommandBufferAllocateInfo(
commandPool=e.ctx.command_pool, level=vk.VK_COMMAND_BUFFER_LEVEL_PRIMARY, commandBufferCount=1
),
)[0]
vk.vkBeginCommandBuffer(cmd, vk.VkCommandBufferBeginInfo(flags=vk.VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT))
self._image_barrier(
cmd,
image,
layers,
old,
new,
0,
vk.VK_ACCESS_SHADER_WRITE_BIT,
vk.VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
)
vk.vkEndCommandBuffer(cmd)
vk.vkQueueSubmit(e.ctx.graphics_queue, 1, [vk.VkSubmitInfo(commandBufferCount=1, pCommandBuffers=[cmd])], None)
vk.vkQueueWaitIdle(e.ctx.graphics_queue)
vk.vkFreeCommandBuffers(device, e.ctx.command_pool, 1, [cmd])
# ------------------------------------------------------------- per-frame
[docs]
def upload_spectrum_if_dirty(self, fft: OceanFFT) -> None:
"""Re-upload the frozen static spectrum when :meth:`OceanFFT.configure` rebuilt it."""
if fft.spectrum_generation == self._spectrum_generation:
return
upload_numpy(self._engine.ctx.device, self._spec_mem, fft.gpu_static_spectrum())
self._spectrum_generation = fft.spectrum_generation
[docs]
def dispatch(self, cmd: Any, time: float, choppiness: float, foam_threshold: float, foam_decay: float) -> None:
"""Record the full GPU FFT chain into ``cmd`` (outside any render pass).
Evolves the spectrum, runs the Stockham inverse FFT and assembles the
displacement / slope arrays, leaving both in ``SHADER_READ_ONLY_OPTIMAL``
ready for the forward ocean pass to sample. Must be called after
:meth:`upload_spectrum_if_dirty` and only when a visible ocean exists.
"""
if not self._ready or self._size <= 0:
return
n = self._size
ffi = vk.ffi
groups = (n + 7) // 8
# 1. Time-evolve the frozen spectrum into ping0.
pc = ffi.new("char[]", np.array([time], dtype=np.float32).tobytes() + np.array([n], dtype=np.int32).tobytes())
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._spectrum_pipe)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._spectrum_layout, 0, 1, [self._spectrum_set], 0, None
)
vk._vulkan.lib.vkCmdPushConstants(cmd, self._spectrum_layout, vk.VK_SHADER_STAGE_COMPUTE_BIT, 0, 8, pc)
vk.vkCmdDispatch(cmd, groups, groups, self._cascades)
self._memory_barrier(cmd)
# 2. Stockham inverse FFT: log2(N) row stages then log2(N) column stages.
stages = int(np.log2(n))
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._fft_pipe)
stage_idx = 0
for axis in (0, 1):
for stage in range(stages):
src_is_a = (stage_idx % 2) == 0
fft_set = self._fft_set_ab if src_is_a else self._fft_set_ba
pcf = ffi.new("char[]", np.array([n, stage, axis], dtype=np.int32).tobytes())
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._fft_layout, 0, 1, [fft_set], 0, None
)
vk._vulkan.lib.vkCmdPushConstants(cmd, self._fft_layout, vk.VK_SHADER_STAGE_COMPUTE_BIT, 0, 12, pcf)
if axis == 0:
vk.vkCmdDispatch(cmd, (n // 2 + 7) // 8, groups, self._cascades * _NUM_FIELDS)
else:
vk.vkCmdDispatch(cmd, groups, (n // 2 + 7) // 8, self._cascades * _NUM_FIELDS)
self._memory_barrier(cmd)
stage_idx += 1
# Even total stage count (2*log2N) -> final result is back in ping0.
# 3. Assemble: move disp/grad to GENERAL, run, then back to SHADER_READ.
self._transition_output(cmd, self._disp, self._disp_layout, vk.VK_IMAGE_LAYOUT_GENERAL)
self._transition_output(cmd, self._grad, self._grad_layout, vk.VK_IMAGE_LAYOUT_GENERAL)
inv_n2 = 1.0 / float(n * n)
p0 = np.array([choppiness, foam_threshold, foam_decay, inv_n2], dtype=np.float32)
p1 = np.array([n, 1 if self._reset_foam else 0, 0, 0], dtype=np.int32)
pca = ffi.new("char[]", p0.tobytes() + p1.tobytes())
vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._assemble_pipe)
vk.vkCmdBindDescriptorSets(
cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, self._assemble_layout, 0, 1, [self._assemble_set], 0, None
)
vk._vulkan.lib.vkCmdPushConstants(cmd, self._assemble_layout, vk.VK_SHADER_STAGE_COMPUTE_BIT, 0, 32, pca)
vk.vkCmdDispatch(cmd, groups, groups, self._cascades)
self._reset_foam = False
# 4. Compute writes -> vertex/fragment sampling of the displacement/slope arrays.
self._transition_output(
cmd, self._disp, vk.VK_IMAGE_LAYOUT_GENERAL, vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
)
self._transition_output(
cmd, self._grad, vk.VK_IMAGE_LAYOUT_GENERAL, vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
)
self._disp_layout = vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
self._grad_layout = vk.VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
# ------------------------------------------------------------- barriers
def _memory_barrier(self, cmd: Any) -> None:
"""Compute-write -> compute-read global memory barrier between FFT dispatches."""
b = vk.VkMemoryBarrier(
srcAccessMask=vk.VK_ACCESS_SHADER_WRITE_BIT,
dstAccessMask=vk.VK_ACCESS_SHADER_READ_BIT | vk.VK_ACCESS_SHADER_WRITE_BIT,
)
vk.vkCmdPipelineBarrier(
cmd,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
0,
1,
[b],
0,
None,
0,
None,
)
def _transition_output(self, cmd: Any, res: dict[str, Any], old: int, new: int) -> None:
if new == vk.VK_IMAGE_LAYOUT_GENERAL:
src_access = vk.VK_ACCESS_SHADER_READ_BIT
dst_access = vk.VK_ACCESS_SHADER_READ_BIT | vk.VK_ACCESS_SHADER_WRITE_BIT
src_stage = vk.VK_PIPELINE_STAGE_VERTEX_SHADER_BIT | vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT
dst_stage = vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT
if old == vk.VK_IMAGE_LAYOUT_UNDEFINED:
src_access = 0
src_stage = vk.VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT
else:
src_access = vk.VK_ACCESS_SHADER_WRITE_BIT
dst_access = vk.VK_ACCESS_SHADER_READ_BIT
src_stage = vk.VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT
dst_stage = vk.VK_PIPELINE_STAGE_VERTEX_SHADER_BIT | vk.VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT
self._image_barrier(cmd, res["image"], res["layers"], old, new, src_access, dst_access, src_stage, dst_stage)
def _image_barrier(self, cmd, image, layers, old, new, src_access, dst_access, src_stage, dst_stage) -> None:
b = vk.VkImageMemoryBarrier(
srcAccessMask=src_access,
dstAccessMask=dst_access,
oldLayout=old,
newLayout=new,
srcQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
dstQueueFamilyIndex=vk.VK_QUEUE_FAMILY_IGNORED,
image=image,
subresourceRange=vk.VkImageSubresourceRange(
aspectMask=vk.VK_IMAGE_ASPECT_COLOR_BIT,
baseMipLevel=0,
levelCount=1,
baseArrayLayer=0,
layerCount=layers,
),
)
vk.vkCmdPipelineBarrier(cmd, src_stage, dst_stage, 0, 0, None, 0, None, 1, [b])
# ------------------------------------------------------------- cleanup
def _destroy_size_resources(self) -> None:
device = self._engine.ctx.device
if self._pool:
vk.vkDestroyDescriptorPool(device, self._pool, None)
self._pool = None
for res in self._ping:
self._destroy_array_image(res)
self._ping = []
if self._disp:
self._destroy_array_image(self._disp)
if self._grad:
self._destroy_array_image(self._grad)
self._disp, self._grad = {}, {}
if self._spec_buf:
vk.vkDestroyBuffer(device, self._spec_buf, None)
vk.vkFreeMemory(device, self._spec_mem, None)
self._spec_buf = self._spec_mem = None
[docs]
def cleanup(self) -> None:
if not self._ready:
return
device = self._engine.ctx.device
self._destroy_size_resources()
for obj, fn in [
(self._spectrum_pipe, vk.vkDestroyPipeline),
(self._spectrum_layout, vk.vkDestroyPipelineLayout),
(self._spectrum_module, vk.vkDestroyShaderModule),
(self._spectrum_set_layout, vk.vkDestroyDescriptorSetLayout),
(self._fft_pipe, vk.vkDestroyPipeline),
(self._fft_layout, vk.vkDestroyPipelineLayout),
(self._fft_module, vk.vkDestroyShaderModule),
(self._fft_set_layout, vk.vkDestroyDescriptorSetLayout),
(self._assemble_pipe, vk.vkDestroyPipeline),
(self._assemble_layout, vk.vkDestroyPipelineLayout),
(self._assemble_module, vk.vkDestroyShaderModule),
(self._assemble_set_layout, vk.vkDestroyDescriptorSetLayout),
]:
if obj:
fn(device, obj, None)
self._ready = False