Skip to content

Instantly share code, notes, and snippets.

@flaneur2020
Created December 15, 2023 14:52
Show Gist options
  • Save flaneur2020/eebbb23de1ddac69e8a2183a1b0b4f1c to your computer and use it in GitHub Desktop.
Save flaneur2020/eebbb23de1ddac69e8a2183a1b0b4f1c to your computer and use it in GitHub Desktop.
When performing matrix-vector multiplication on the GPU, you can optimize the operation using parallel reduction to efficiently sum up the elements. In a Vulkan compute shader, this could be done using shared memory (local memory in GLSL) to perform a partial sum for each workgroup before completing the sum for the entire vector.
Here's a GLSL compute shader example that demonstrates a reduction pattern for summing elements. The shader assumes that the size of the matrix row (which is the size of the vector) is a power of two for simplicity. This may not be the case in a real-world scenario, so you would need to modify the shader to handle arbitrary sizes.
```glsl
#version 450
// Define the size of the workgroup
layout(local_size_x = 128) in;
// Binding for the matrix buffer
layout(set = 0, binding = 0) readonly buffer Matrix {
float m[];
} matrix;
// Binding for the vector buffer
layout(set = 0, binding = 1) readonly buffer Vector {
float v[];
} vector;
// Binding for the result buffer
layout(set = 0, binding = 2) buffer Result {
float r[];
} result;
// Shared memory for the reduction
shared float sharedData[gl_WorkGroupSize.x];
void main() {
int localIdx = int(gl_LocalInvocationID.x);
int globalIdx = int(gl_GlobalInvocationID.x);
// Calculate the number of elements in a row (assuming a square matrix)
int N = int(sqrt(float(matrix.m.length())));
// Perform multiplication and store in shared memory
sharedData[localIdx] = matrix.m[globalIdx * N + localIdx] * vector.v[localIdx];
// Synchronize to ensure all writes to shared memory have completed
barrier();
// Perform parallel reduction in shared memory
for (int offset = gl_WorkGroupSize.x / 2; offset > 0; offset >>= 1) {
if (localIdx < offset) {
sharedData[localIdx] += sharedData[localIdx + offset];
}
// Synchronize after each reduction step
barrier();
}
// Write the result from the first thread of each workgroup
if (localIdx == 0) {
result.r[globalIdx] = sharedData[0];
}
}
```
In this shader:
- Each workgroup computes a part of the matrix-vector multiplication and stores it in shared memory.
- A parallel reduction is performed within each workgroup using the shared memory to sum the elements.
- The resulting sum for each row of the matrix is written back by the first thread in each workgroup.
Please note:
- The size of the workgroup (`local_size_x`) is set to 128. The actual optimal size depends on your GPU and the specific problem size.
- The size of the shared memory array is set to the same as the workgroup size. You need to ensure that your GPU's shared memory per workgroup can accommodate this.
- The shader assumes the size of the matrix row is a power of two, which simplifies the reduction logic. If the row size is not a power of two, additional logic is required to handle the remaining elements.
- The shader assumes a square matrix for simplicity. You will need to adjust the code if your matrices have different dimensions.
- The `barrier()` function is used to synchronize threads within a workgroup, ensuring that all threads have completed their writes to shared memory before the reduction starts and between each step of the reduction.
To use this shader in a Vulkan application, you would need to create a compute pipeline, allocate and bind the appropriate buffers for the matrix, vector, and result, and dispatch the compute work specifying the number of workgroups based on the size of your matrix. Additionally, you would need to handle synchronization to ensure the compute shader execution is complete before reading back the result.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment