Skip to content

Add compute_workgroups example#801

Open
mrinalchaturvedi27 wants to merge 2 commits intopygfx:mainfrom
mrinalchaturvedi27:workgroup-exp
Open

Add compute_workgroups example#801
mrinalchaturvedi27 wants to merge 2 commits intopygfx:mainfrom
mrinalchaturvedi27:workgroup-exp

Conversation

@mrinalchaturvedi27
Copy link

@mrinalchaturvedi27 mrinalchaturvedi27 commented Mar 2, 2026

This PR adds a compute_workgroups.py example demonstrating GPU workgroups and invocation IDs.
While reviewing the WGPU documentation and the Rust hello_workgroups example, I noticed that wgpu-py does not include a focused example for workgroups.
This example provides a minimal compute shader that writes global_invocation_id, local_invocation_id, and workgroup_id into a storage buffer so the relationship between them can be inspected.

Copilot AI review requested due to automatic review settings March 2, 2026 19:48
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds a new compute_workgroups.py example to the examples/ suite to help users understand how compute workgroups relate to invocation IDs in wgpu-py, by writing invocation/workgroup identifiers into a storage buffer and printing/verifying the mapping on the CPU.

Changes:

  • Add a minimal WGSL compute shader that records global_invocation_id, local_invocation_id, and workgroup_id (currently x-components) into a storage buffer.
  • Add a Python driver that dispatches workgroups, reads back the buffer, prints a table, and validates the expected ID relationships.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +4 to +5
Each thread writes its global, local, and workgroup IDs into a storage buffer
so the relationship between them can be inspected.
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docstring says each thread writes its global/local/workgroup IDs, but the shader only writes the .x component of each builtin (y/z are always 0 with the current 1D dispatch). Consider clarifying in the docstring that this example records the x-components (or expand the buffer to store all 3 components per ID).

Suggested change
Each thread writes its global, local, and workgroup IDs into a storage buffer
so the relationship between them can be inspected.
Each thread in this 1D dispatch writes the x-components of its global, local,
and workgroup IDs into a storage buffer so the relationship between them can
be inspected.

Copilot uses AI. Check for mistakes.
f"Dispatched {workgroups} workgroup(s) of {workgroup_size} thread(s) each {total_threads} threads total.\n"
)

print(f"{'Thread':>6} {'global_id':>9} {'local_id':>8} {'workgroup_id':>12} ")
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There’s an extra trailing space at the end of the header format string, which results in an extra space in the printed table header. Consider removing it to keep output tidy.

Suggested change
print(f"{'Thread':>6} {'global_id':>9} {'local_id':>8} {'workgroup_id':>12} ")
print(f"{'Thread':>6} {'global_id':>9} {'local_id':>8} {'workgroup_id':>12}")

Copilot uses AI. Check for mistakes.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
@mrinalchaturvedi27
Copy link
Author

Hi @Korijn @almarklein, just checking in on this PR. Whenever you get a chance, could you please take a look? Happy to make any changes if needed. Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants