.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "gallery/compute_noop.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_gallery_compute_noop.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_gallery_compute_noop.py:


Compute Noop
------------

Example compute shader that does ... nothing but copy a value from one
buffer into another.

.. GENERATED FROM PYTHON SOURCE LINES 8-13

.. code-block:: Python


    import wgpu
    from wgpu.utils.compute import compute_with_buffers  # Convenience function



.. GENERATED FROM PYTHON SOURCE LINES 14-38

.. code-block:: Python


    shader_source = """

    @group(0) @binding(0)
    var<storage,read> data1: array<i32>;

    @group(0) @binding(1)
    var<storage,read_write> data2: array<i32>;

    @compute
    @workgroup_size(1)
    fn main(@builtin(global_invocation_id) index: vec3<u32>) {
        let i: u32 = index.x;
        data2[i] = data1[i];
    }
    """

    # Create input data as a memoryview
    n = 20
    data = memoryview(bytearray(n * 4)).cast("i")
    for i in range(n):
        data[i] = i



.. GENERATED FROM PYTHON SOURCE LINES 39-51

.. code-block:: Python


    # The first arg is the input data, per binding
    # The second arg is the output types, per binding
    out = compute_with_buffers({0: data}, {1: (n, "i")}, shader_source, n=n)

    # The result is a dict matching the output types
    # Select data from buffer at binding 1
    result = out[1].tolist()
    print(result)
    assert result == list(range(20))



.. GENERATED FROM PYTHON SOURCE LINES 52-61

.. code-block:: Python


    # import numpy as np
    #
    # numpy_data = np.frombuffer(data, np.int32)
    # out = compute_with_buffers({0: numpy_data}, {1: numpy_data.nbytes}, shader_source, n=n)
    # result = np.frombuffer(out[1], dtype=np.int32)
    # print(result.tolist())



.. GENERATED FROM PYTHON SOURCE LINES 64-65

Create device and shader object

.. GENERATED FROM PYTHON SOURCE LINES 65-81

.. code-block:: Python

    device = wgpu.utils.get_default_device()

    # Show all available adapters
    adapters = wgpu.gpu.enumerate_adapters_sync()
    for a in adapters:
        print(a.summary)

    # You can select specific GPU from the available adapters
    # adapter = None
    # for a in adapters:
    #     if "NVIDIA" in a.summary:
    #         adapter = a
    #         break
    # assert adapter is not None
    # device = adapter.request_device_sync()


.. GENERATED FROM PYTHON SOURCE LINES 82-142

.. code-block:: Python

    cshader = device.create_shader_module(code=shader_source)

    # Create buffer objects, input buffer is mapped.
    buffer1 = device.create_buffer_with_data(data=data, usage=wgpu.BufferUsage.STORAGE)
    buffer2 = device.create_buffer(
        size=data.nbytes, usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC
    )

    # Setup layout and bindings
    binding_layouts = [
        {
            "binding": 0,
            "visibility": wgpu.ShaderStage.COMPUTE,
            "buffer": {
                "type": wgpu.BufferBindingType.read_only_storage,
            },
        },
        {
            "binding": 1,
            "visibility": wgpu.ShaderStage.COMPUTE,
            "buffer": {
                "type": wgpu.BufferBindingType.storage,
            },
        },
    ]
    bindings = [
        {
            "binding": 0,
            "resource": {"buffer": buffer1, "offset": 0, "size": buffer1.size},
        },
        {
            "binding": 1,
            "resource": {"buffer": buffer2, "offset": 0, "size": buffer2.size},
        },
    ]

    # Put everything together
    bind_group_layout = device.create_bind_group_layout(entries=binding_layouts)
    pipeline_layout = device.create_pipeline_layout(bind_group_layouts=[bind_group_layout])
    bind_group = device.create_bind_group(layout=bind_group_layout, entries=bindings)

    # Create and run the pipeline
    compute_pipeline = device.create_compute_pipeline(
        layout=pipeline_layout,
        compute={"module": cshader, "entry_point": "main"},
    )
    command_encoder = device.create_command_encoder()
    compute_pass = command_encoder.begin_compute_pass()
    compute_pass.set_pipeline(compute_pipeline)
    compute_pass.set_bind_group(0, bind_group)
    compute_pass.dispatch_workgroups(n, 1, 1)  # x y z
    compute_pass.end()
    device.queue.submit([command_encoder.finish()])

    # Read result
    # result = buffer2.read_data().cast("i")
    out = device.queue.read_buffer(buffer2).cast("i")
    result = out.tolist()
    print(result)
    assert result == list(range(20))


.. _sphx_glr_download_gallery_compute_noop.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: compute_noop.ipynb <compute_noop.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: compute_noop.py <compute_noop.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: compute_noop.zip <compute_noop.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_

.. only:: html

    Interactive example
    ===================

    This uses Pyodide. If this does not work, your browser may not have sufficient support for wasm/pyodide/wgpu (check your browser dev console).
    Stdout (print statements) will also appear in the browser console.

    .. raw:: html

        <iframe src="pyodide.html#compute_noop.py"></iframe>