.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "gallery/compute_int64.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_int64.py>`
        to download the full example code.

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

.. _sphx_glr_gallery_compute_int64.py:


Compute INT64
-------------

simple example of using the int64 shader feature

.. GENERATED FROM PYTHON SOURCE LINES 7-74

.. code-block:: Python


    import wgpu

    adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance")
    device = adapter.request_device_sync(required_features=["shader-int64"])

    add_shader = """
    @group(0) @binding(0)
    var<storage, read_write> data: array<i64>;

    @compute
    @workgroup_size(1)
    fn main(@builtin(global_invocation_id) index: vec3<u32>) {
        let i: u32 = index.x;
        let a_idx = i * 3u;
        let b_idx = a_idx + 1u;
        let c_idx = a_idx + 2u;
        data[c_idx] = data[a_idx] + data[b_idx];
    }
    """
    a = 0x1234567890
    b = 0x9876543210
    data = memoryview(bytearray(6 * 8)).cast("q")  # signed long long >= 64 bits
    data[0] = a
    data[1] = b
    data[3] = -a
    data[4] = -b

    buffer = device.create_buffer_with_data(
        data=data,
        usage=wgpu.BufferUsage.STORAGE
        | wgpu.BufferUsage.COPY_SRC
        | wgpu.BufferUsage.COPY_DST,
    )
    pipeline = device.create_compute_pipeline(
        layout="auto",
        compute={
            "module": device.create_shader_module(code=add_shader),
            "entry_point": "main",
        },
    )

    bind_group = device.create_bind_group(
        layout=pipeline.get_bind_group_layout(0),
        entries=[
            {
                "binding": 0,
                "resource": {"buffer": buffer, "offset": 0, "size": buffer.size},
            }
        ],
    )

    command_encoder = device.create_command_encoder()
    compute_pass = command_encoder.begin_compute_pass()
    compute_pass.set_pipeline(pipeline)
    compute_pass.set_bind_group(0, bind_group)
    compute_pass.dispatch_workgroups(2)  # we do two calculations
    compute_pass.end()
    device.queue.submit([command_encoder.finish()])
    result = device.queue.read_buffer(buffer)
    res = result.cast("q").tolist()
    c = a + b
    assert res[2] == c, f"expected {c}, got {res[2]}"
    print(f"{a:#x} + {b:#x} = {c:#x}")
    d = -a + -b
    assert res[5] == d, f"expected {d}, got {res[5]}"
    print(f"{-a:#x} + {-b:#x} = {d:#x}")


.. _sphx_glr_download_gallery_compute_int64.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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