wgpu_helper.py

import wgpu
import numpy as np
from rendercanvas.auto import RenderCanvas, loop

class WGPUHelper:
    """
    Class that is responsible to abstract away the gpu dirty mess...
    """

    def __init__(self, workgroup_size = 16, shader_source='kernel.wgsl'):
        self.adapter = None
        self.device = None
        self.workgroup_size = workgroup_size
        self.shader_source = shader_source
        self.bind_group_1 = None
        self.bind_group_2 = None
        self.source = ''

        # FD-TD variables
        self.c1x = None
        self.c1y = None
        self.c2x = None
        self.c2y = None
        self.c3 = None
        self.c4 = None
        self.dx = None
        self.dy = None
        self.dt = None
        self.t = None
        self.wire = None

        self.Ezix = None
        self.Eziy = None
        self.Hxi = None
        self.Hyi = None

        self.sources = None
        
        # GPU versions of FD-TD variables
        self.wgpu_c1x = None
        self.wgpu_c1y = None
        self.wgpu_c2x = None
        self.wgpu_c2y = None
        self.wgpu_c3 = None
        self.wgpu_c4 = None
        self.wgpu_dx = None
        self.wgpu_dy = None
        self.wgpu_dt = None
        self.wgpu_t = None
        self.wgpu_wire = None
        
        # double buffer to avoid concurrence and data race in the spatial difference equation update
        self.wgpu_Ezix_1 = None
        self.wgpu_Eziy_1 = None
        self.wgpu_Hxi_1 = None
        self.wgpu_Hyi_1 = None
        
        self.wgpu_Ezix_2 = None
        self.wgpu_Eziy_2 = None
        self.wgpu_Hxi_2 = None
        self.wgpu_Hyi_2 = None

        self.wgpu_sources = None
    
    # Frame Loop
    def compute(self, n):
        N = self.Ezix.shape[0]
        M = self.Eziy.shape[1]
        self.call_kernel(1, 1, 0, 2) # 2 = update_source
        self.call_kernel(N, M, n, 0) # 0 = leap_frog1 (Compute Ezi)
        self.call_kernel(N, M, n, 1) # 1 = leap_frog2 (compute Hxi, Hyi with the UPDATED version of Ezi)



    def setup_compute(self, c1x, c1y, c2x, c2y, c3, c4, dx, dy, dt, t, wire, Ezix, Eziy, Hxi, Hyi, sources):

        """
        sources format: [(position, type, parameters), (position, type, parameters), ...]
        """

        print('preparing the GPU')

        N = Ezix.shape[0]
        M = Ezix.shape[1]
        
        # 0. store data (caution, these are copies)
        self.c1x = c1x.astype(np.float32)
        self.c1y = c1y.astype(np.float32)
        self.c2x = c2x.astype(np.float32)
        self.c2y = c2y.astype(np.float32)
        self.c3 = c3.astype(np.float32)
        self.c4 = c4.astype(np.float32)
        self.dx = dx.astype(np.float32)
        self.dy = dy.astype(np.float32)
        self.dt = dt.astype(np.float32)
        self.t = t.astype(np.float32)
        self.wire = wire.astype(np.float32)

        # Ezix and Eziy will both add up to form Ezi (in 2D we have only Ez, but for PML, we need to use Berenger strategy
        # of spliting Ez into a component what will propagate in x and y (so that the vertical PML will absorb Ex and horizontal will
        # absorb only Ey
        self.Ezix = Ezix.astype(np.float32)
        self.Eziy = Eziy.astype(np.float32)
        self.Hxi = Hxi.astype(np.float32)
        self.Hyi = Hyi.astype(np.float32)

        self.sources = sources

        # 1. Setup: Create a device to communicate with the GPU
        self.adapter = wgpu.gpu.request_adapter_sync(
		power_preference="high-performance",
		force_fallback_adapter=False)
        self.device = self.adapter.request_device_sync()


    	# Print out exactly what hardware Python bound to
        print(f"Active Adapter Summary: {self.adapter.summary}")
        print(f"Active Adapter Info: {self.adapter.info}")

        # 2. Shader: Define the WGSL code
        with open(self.shader_source,'rt') as fid:
            shader_source = fid.read()
        shader_source = shader_source.replace('[[workgoup_size]]', str(self.workgroup_size))
        shader_source = shader_source.replace('[[N]]', str(N))
        shader_source = shader_source.replace('[[M]]', str(M))
        self.source = shader_source
        # print(shader_source)
        self.shader_module = self.device.create_shader_module(code=shader_source)



        # 3. Buffers and uniforms: Move data to the GPU
        # We need a buffer that is "Storage" (for compute) and "CopySrc" (so we can read it back) and some variables 
        # We also need some uniforms (things that never change during the lifetime of the shader)
        # to use as parameters to the shader (uniforms)
        usage_storage = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST |  wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.VERTEX
        usage_uniform = wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST
        self.wgpu_c1x = self.device.create_buffer_with_data(
            data=self.c1x, 
            usage=usage_storage
        )
        self.wgpu_c1y = self.device.create_buffer_with_data(
            data=self.c1y, 
            usage=usage_storage
        )
        self.wgpu_c2x = self.device.create_buffer_with_data(
            data=self.c2x, 
            usage=usage_storage
        )
        self.wgpu_c2y = self.device.create_buffer_with_data(
            data=self.c2y, 
            usage=usage_storage
        )
        self.wgpu_c3 = self.device.create_buffer_with_data(
            data=self.c3, 
            usage=usage_storage
        )
        self.wgpu_c4 = self.device.create_buffer_with_data(
            data=self.c4, 
            usage=usage_storage
        )
        self.wgpu_wire = self.device.create_buffer_with_data(
            data=self.wire, 
            usage=usage_storage
        )
        self.wgpu_dx = self.device.create_buffer_with_data(
            data=self.dx, 
            usage=usage_uniform
        )
        self.wgpu_dy = self.device.create_buffer_with_data(
            data=self.dy, 
            usage=usage_uniform
        )
        self.wgpu_dt = self.device.create_buffer_with_data(
            data=self.dt, 
            usage=usage_uniform
        )
        self.wgpu_t = self.device.create_buffer_with_data(
            data=self.t, 
            usage=usage_storage
        )
        self.wgpu_Ezix_1 = self.device.create_buffer_with_data(
            data=self.Ezix, 
            usage=usage_storage
        )
        self.wgpu_Eziy_1 = self.device.create_buffer_with_data(
            data=self.Eziy, 
            usage=usage_storage
        )
        self.wgpu_Hxi_1 = self.device.create_buffer_with_data(
            data=self.Hxi, 
            usage=usage_storage
        )
        self.wgpu_Hyi_1 = self.device.create_buffer_with_data(
            data=self.Hyi, 
            usage=usage_storage
        )
        self.wgpu_Ezix_2 = self.device.create_buffer_with_data(
            data=self.Ezix, 
            usage=usage_storage
        )
        self.wgpu_Eziy_2 = self.device.create_buffer_with_data(
            data=self.Eziy, 
            usage=usage_storage
        )
        self.wgpu_Hxi_2 = self.device.create_buffer_with_data(
            data=self.Hxi, 
            usage=usage_storage
        )
        self.wgpu_Hyi_2 = self.device.create_buffer_with_data(
            data=self.Hyi, 
            usage=usage_storage
        )
        self.wgpu_sources = self.device.create_buffer_with_data(
            data=self.sources.tobytes(),
            usage=usage_storage
        )
        bind_group_layout_entries = [
                {"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # c1x
                {"binding": 1, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # c1y
                {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # c2x
                {"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # c2y
                {"binding": 4, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # c3
                {"binding": 5, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # c4
                {"binding": 6, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # wire
                {"binding": 7, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.uniform}}, # dx
                {"binding": 8, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.uniform}}, # dy
                {"binding": 9, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.uniform}}, # dt
                {"binding": 10, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # t
                {"binding": 11, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Ezix_1(or 2)
                {"binding": 12, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Eziy_1(or 2)
                {"binding": 13, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Hxi_1(or 2)
                {"binding": 14, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Hyi_1(or 2)
                {"binding": 15, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Ezix_2(or 1)
                {"binding": 16, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Eziy_2(or 1)
                {"binding": 17, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Hxi_2(or 1)
                {"binding": 18, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # Hyi_2(or 1)
                {"binding": 19, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.storage}}, # sources
        ]


        # 4. Pipeline: Compile the shader and set up the layout
        #              It will organize the runing of the code. For instance will tell which binding is which (storage, uniform, etc) 
        compute_pipelines = [
            self.device.create_compute_pipeline(
                layout=self.device.create_pipeline_layout(bind_group_layouts=[
                    self.device.create_bind_group_layout(entries=bind_group_layout_entries)
                ]),
                compute={"module": self.shader_module, "entry_point": 'leap_frog1'}
            ),
            self.device.create_compute_pipeline(
                layout=self.device.create_pipeline_layout(bind_group_layouts=[
                    self.device.create_bind_group_layout(entries=bind_group_layout_entries)
                ]),
                compute={"module": self.shader_module, "entry_point": 'leap_frog2'}
            ),
            self.device.create_compute_pipeline(
                layout=self.device.create_pipeline_layout(bind_group_layouts=[
                    self.device.create_bind_group_layout(entries=bind_group_layout_entries)
                ]),
                compute={"module": self.shader_module, "entry_point": 'update_source'}
            ),
        ]

        # 6. Bind Group: Link the buffers and uniforms to the pipeline
        #                Pipeline just told what the date will be... Here we bind them together with the created gpu_xxx variables
        # We have two for the double buffer pointer swap trick
        bind_group_entries_1 = [
            {"binding":  0, "resource": {"buffer": self.wgpu_c1x, "offset": 0, "size": self.wgpu_c1x.size}},
            {"binding":  1, "resource": {"buffer": self.wgpu_c1y, "offset": 0, "size": self.wgpu_c1y.size}},
            {"binding":  2, "resource": {"buffer": self.wgpu_c2x, "offset": 0, "size": self.wgpu_c2x.size}},
            {"binding":  3, "resource": {"buffer": self.wgpu_c2y, "offset": 0, "size": self.wgpu_c2y.size}},
            {"binding":  4, "resource": {"buffer": self.wgpu_c3, "offset": 0, "size": self.wgpu_c3.size}},
            {"binding":  5, "resource": {"buffer": self.wgpu_c4, "offset": 0, "size": self.wgpu_c4.size}},
            {"binding":  6, "resource": {"buffer": self.wgpu_wire, "offset": 0, "size": self.wgpu_wire.size}},
            {"binding":  7, "resource": {"buffer": self.wgpu_dx, "offset": 0, "size": self.wgpu_dx.size}},
            {"binding":  8, "resource": {"buffer": self.wgpu_dy, "offset": 0, "size": self.wgpu_dy.size}},
            {"binding":  9, "resource": {"buffer": self.wgpu_dt, "offset": 0, "size": self.wgpu_dt.size}},
            {"binding": 10, "resource": {"buffer": self.wgpu_t, "offset": 0, "size": self.wgpu_t.size}},
            {"binding": 11, "resource": {"buffer": self.wgpu_Ezix_1, "offset": 0, "size": self.wgpu_Ezix_1.size}},
            {"binding": 12, "resource": {"buffer": self.wgpu_Eziy_1, "offset": 0, "size": self.wgpu_Eziy_1.size}},
            {"binding": 13, "resource": {"buffer": self.wgpu_Hxi_1, "offset": 0, "size": self.wgpu_Hxi_1.size}},
            {"binding": 14, "resource": {"buffer": self.wgpu_Hyi_1, "offset": 0, "size": self.wgpu_Hyi_1.size}},
            {"binding": 15, "resource": {"buffer": self.wgpu_Ezix_2, "offset": 0, "size": self.wgpu_Ezix_2.size}},
            {"binding": 16, "resource": {"buffer": self.wgpu_Eziy_2, "offset": 0, "size": self.wgpu_Eziy_2.size}},
            {"binding": 17, "resource": {"buffer": self.wgpu_Hxi_2, "offset": 0, "size": self.wgpu_Hxi_2.size}},
            {"binding": 18, "resource": {"buffer": self.wgpu_Hyi_2, "offset": 0, "size": self.wgpu_Hyi_2.size}},
            {"binding": 19, "resource": {"buffer": self.wgpu_sources, "offset": 0, "size": self.wgpu_sources.size}},
        ]
        bind_group_entries_2 = [
            {"binding":  0, "resource": {"buffer": self.wgpu_c1x, "offset": 0, "size": self.wgpu_c1x.size}},
            {"binding":  1, "resource": {"buffer": self.wgpu_c1y, "offset": 0, "size": self.wgpu_c1y.size}},
            {"binding":  2, "resource": {"buffer": self.wgpu_c2x, "offset": 0, "size": self.wgpu_c2x.size}},
            {"binding":  3, "resource": {"buffer": self.wgpu_c2y, "offset": 0, "size": self.wgpu_c2y.size}},
            {"binding":  4, "resource": {"buffer": self.wgpu_c3, "offset": 0, "size": self.wgpu_c3.size}},
            {"binding":  5, "resource": {"buffer": self.wgpu_c4, "offset": 0, "size": self.wgpu_c4.size}},
            {"binding":  6, "resource": {"buffer": self.wgpu_wire, "offset": 0, "size": self.wgpu_wire.size}},
            {"binding":  7, "resource": {"buffer": self.wgpu_dx, "offset": 0, "size": self.wgpu_dx.size}},
            {"binding":  8, "resource": {"buffer": self.wgpu_dy, "offset": 0, "size": self.wgpu_dy.size}},
            {"binding":  9, "resource": {"buffer": self.wgpu_dt, "offset": 0, "size": self.wgpu_dt.size}},
            {"binding": 10, "resource": {"buffer": self.wgpu_t, "offset": 0, "size": self.wgpu_t.size}},
            {"binding": 11, "resource": {"buffer": self.wgpu_Ezix_2, "offset": 0, "size": self.wgpu_Ezix_2.size}},
            {"binding": 12, "resource": {"buffer": self.wgpu_Eziy_2, "offset": 0, "size": self.wgpu_Eziy_2.size}},
            {"binding": 13, "resource": {"buffer": self.wgpu_Hxi_2, "offset": 0, "size": self.wgpu_Hxi_2.size}},
            {"binding": 14, "resource": {"buffer": self.wgpu_Hyi_2, "offset": 0, "size": self.wgpu_Hyi_2.size}},
            {"binding": 15, "resource": {"buffer": self.wgpu_Ezix_1, "offset": 0, "size": self.wgpu_Ezix_1.size}},
            {"binding": 16, "resource": {"buffer": self.wgpu_Eziy_1, "offset": 0, "size": self.wgpu_Eziy_1.size}},
            {"binding": 17, "resource": {"buffer": self.wgpu_Hxi_1, "offset": 0, "size": self.wgpu_Hxi_1.size}},
            {"binding": 18, "resource": {"buffer": self.wgpu_Hyi_1, "offset": 0, "size": self.wgpu_Hyi_1.size}},
            {"binding": 19, "resource": {"buffer": self.wgpu_sources, "offset": 0, "size": self.wgpu_sources.size}},
        ]

        # All compute_pilelines use the same bindgroup, so anyone can be used here
        bind_group_1 = self.device.create_bind_group(
            layout=compute_pipelines[0].get_bind_group_layout(0),
            entries=bind_group_entries_1
        )
        bind_group_2 = self.device.create_bind_group(
            layout=compute_pipelines[0].get_bind_group_layout(0),
            entries=bind_group_entries_2
        )

        # store for later use...
        self.compute_pipelines = compute_pipelines
        self.bind_group_1 = bind_group_1
        self.bind_group_2 = bind_group_2



    def call_kernel(self, N_workers, M_workers, n, entry_point_idx):
        # Execution: Dispatch the work
        encoder = self.device.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self.compute_pipelines[entry_point_idx])
        if n%2 == 0:
            compute_pass.set_bind_group(0, self.bind_group_1)
        else:
            compute_pass.set_bind_group(0, self.bind_group_2)
        compute_pass.dispatch_workgroups(1+N_workers//self.workgroup_size, 1+M_workers//self.workgroup_size)  # Run once for each row and coll. In the kernel, global_invocation_id (generally the first parameter) will have .x and .y since we are passing N_workers and M_workers
        compute_pass.end()
        self.device.queue.submit([encoder.finish()])


    def sync_gpu_host_data(self):
        
        Ezix_ = self.device.queue.read_buffer(self.wgpu_Ezix_1).cast('f')
        Eziy_ = self.device.queue.read_buffer(self.wgpu_Eziy_1).cast('f')
        # reshape to familiar shape
        new_Ezix = np.array(Ezix_).reshape(self.Ezix.shape)
        new_Eziy = np.array(Eziy_).reshape(self.Eziy.shape)
        
        return new_Ezix+new_Eziy


    def sync_host_gpu_data(self):

        self.device.queue.write_buffer(self.wgpu_c1x, 0, self.c1x)
        self.device.queue.write_buffer(self.wgpu_c1y, 0, self.c1y)
        self.device.queue.write_buffer(self.wgpu_c2x, 0, self.c2x)
        self.device.queue.write_buffer(self.wgpu_c2y, 0, self.c2y)
        self.device.queue.write_buffer(self.wgpu_c3, 0, self.c3)
        self.device.queue.write_buffer(self.wgpu_c4, 0, self.c4)
        self.device.queue.write_buffer(self.wgpu_wire, 0, self.wire)
        self.device.queue.write_buffer(self.wgpu_dx, 0, self.dx)
        self.device.queue.write_buffer(self.wgpu_dy, 0, self.dy)
        self.device.queue.write_buffer(self.wgpu_dt, 0, self.dt)
        self.device.queue.write_buffer(self.wgpu_t, 0, self.t)
        self.device.queue.write_buffer(self.wgpu_Ezix_1, 0, self.Ezix)
        self.device.queue.write_buffer(self.wgpu_Eziy_1, 0, self.Eziy)
        self.device.queue.write_buffer(self.wgpu_Hxi_1, 0, self.Hxi)
        self.device.queue.write_buffer(self.wgpu_Hyi_1, 0, self.Hyi)
        self.device.queue.write_buffer(self.wgpu_Ezix_2, 0, self.Ezix)
        self.device.queue.write_buffer(self.wgpu_Eziy_2, 0, self.Eziy)
        self.device.queue.write_buffer(self.wgpu_Hxi_2, 0, self.Hxi)
        self.device.queue.write_buffer(self.wgpu_Hyi_2, 0, self.Hyi)