😎 Parallel Matrix Multiplication on GPU, using Rust Vulkan Compute API of `vulkano` 🚴🏼

README.md

Using device: Apple M1 Max (type: IntegratedGpu). API version: 1.2.296
Matrix A = (1774 X 2048)
Matrix B = (2048 X 1024)
Matrix C = A X B = (1774 X 1024)
(a) GPU matrix multiply: 47.795959ms
(b) CPU parallel matrix multiply: 1.081648416s
(c) CPU single-threaded matrix multiply: 7.010535s
Speed Up (b/a): 22x
Speed Up (c/a): 146x
Speed Up (c/b): 6x

---

Using device: Intel(R) Graphics (ADL GT2) (type: IntegratedGpu). API version: 1.3.281
Matrix A = (1774 X 2048)
Matrix B = (2048 X 1024)
Matrix C = A X B = (1774 X 1024)
(a) GPU matrix multiply: 280.689909ms
(b) CPU parallel matrix multiply: 1.792873747s
(c) CPU single-threaded matrix multiply: 7.692478067s
Speed Up (b/a): 6x
Speed Up (c/a): 27x
Speed Up (c/b): 4x

Cargo.toml

[package]
name = "gpu-matrix-multiplication"
version = "0.1.0"
edition = "2024"
rust-version = "1.85.0"

[dependencies]
vulkano = "=0.35.1"
vulkano-shaders = "=0.35.0"
rand = "=0.8.4"
rayon = "=1.10.0"

main.rs

use rand::{Rng, SeedableRng, rngs::StdRng};
use rayon::prelude::*;
use std::{ops::Mul, sync::Arc, time::Instant};
use vulkano::{
    VulkanLibrary,
    buffer::{Buffer, BufferCreateInfo, BufferUsage},
    command_buffer::{
        AutoCommandBufferBuilder, CommandBufferUsage, CopyBufferInfo, PrimaryCommandBufferAbstract,
        allocator::StandardCommandBufferAllocator,
    },
    descriptor_set::{
        DescriptorSet, WriteDescriptorSet, allocator::StandardDescriptorSetAllocator,
    },
    device::{
        Device, DeviceCreateInfo, DeviceExtensions, QueueCreateInfo, QueueFlags,
        physical::PhysicalDeviceType,
    },
    instance::{Instance, InstanceCreateFlags, InstanceCreateInfo},
    memory::allocator::{AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator},
    pipeline::{
        ComputePipeline, Pipeline, PipelineBindPoint, PipelineLayout,
        PipelineShaderStageCreateInfo, compute::ComputePipelineCreateInfo,
        layout::PipelineDescriptorSetLayoutCreateInfo,
    },
    sync::{self, GpuFuture},
};

const MATRIX_A_NUM_ROWS: u32 = 1774;
const MATRIX_A_NUM_COLS: u32 = 1u32 << 11;
const MATRIX_B_NUM_ROWS: u32 = 1u32 << 11;
const MATRIX_B_NUM_COLS: u32 = 1u32 << 10;

#[derive(PartialEq, Debug)]
struct Matrix {
    rows: u32,
    cols: u32,
    elems: Vec<u32>,
}

impl Mul for Matrix {
    type Output = Matrix;

    fn mul(self, rhs: Self) -> Self::Output {
        assert_eq!(self.cols, rhs.rows);
        let mut elems = vec![0u32; (self.rows * rhs.cols) as usize];

        for i in 0..self.rows {
            for j in 0..rhs.cols {
                let mut sum = 0u32;
                for k in 0..self.cols {
                    sum = sum.wrapping_add(
                        self.elems[(i * self.cols + k) as usize]
                            .wrapping_mul(rhs.elems[(k * rhs.cols + j) as usize]),
                    );
                }
                elems[(i * rhs.cols + j) as usize] = sum;
            }
        }

        Matrix {
            rows: self.rows,
            cols: rhs.cols,
            elems,
        }
    }
}

impl Matrix {
    fn random(seed: u64, rows: u32, cols: u32) -> Matrix {
        let mut rng = StdRng::seed_from_u64(seed);
        Matrix {
            rows,
            cols,
            elems: (0..rows * cols).map(|_| rng.r#gen::<u32>()).collect(),
        }
    }

    fn par_mul(&self, rhs: &Matrix) -> Matrix {
        assert_eq!(self.cols, rhs.rows);
        let mut elems = vec![0u32; (self.rows * rhs.cols) as usize];

        elems.par_iter_mut().enumerate().for_each(|(lin_idx, v)| {
            let r_idx = lin_idx / rhs.cols as usize;
            let c_idx = lin_idx - r_idx * rhs.cols as usize;

            *v = (0..self.cols as usize).fold(0u32, |acc, k| {
                acc.wrapping_add(
                    self.elems[r_idx * self.cols as usize + k]
                        .wrapping_mul(rhs.elems[k * rhs.cols as usize + c_idx]),
                )
            });
        });

        Matrix {
            rows: self.rows,
            cols: rhs.cols,
            elems,
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        let offset0 = 0;
        let offset1 = offset0 + std::mem::size_of_val(&self.rows);
        let offset2 = offset1 + std::mem::size_of_val(&self.cols);
        let encoded_elems_byte_len = std::mem::size_of::<u32>() * (self.rows * self.cols) as usize;
        let encoded_mat_byte_len = offset2 + encoded_elems_byte_len;

        let elems_as_bytes = unsafe {
            let ptr_elems = self.elems.as_ptr();
            let ptr_elem_bytes: *const u8 = ptr_elems.cast();

            core::slice::from_raw_parts(ptr_elem_bytes, encoded_elems_byte_len)
        };

        let mut bytes = vec![0u8; encoded_mat_byte_len];

        bytes[offset0..offset1].copy_from_slice(&self.rows.to_le_bytes());
        bytes[offset1..offset2].copy_from_slice(&self.cols.to_le_bytes());
        bytes[offset2..].copy_from_slice(elems_as_bytes);

        bytes
    }

    fn from_bytes(bytes: &[u8]) -> Matrix {
        let offset0 = 0;
        let offset1 = offset0 + std::mem::size_of::<u32>();
        let offset2 = offset1 + std::mem::size_of::<u32>();

        let rows = u32::from_le_bytes(bytes[offset0..offset1].try_into().unwrap());
        let cols = u32::from_le_bytes(bytes[offset1..offset2].try_into().unwrap());
        let num_elems = (rows * cols) as usize;

        let encoded_elems_byte_len = std::mem::size_of::<u32>() * num_elems;
        let remaining_num_bytes = bytes.len() - offset2;

        assert_eq!(encoded_elems_byte_len, remaining_num_bytes);

        let elems = unsafe {
            let ptr_elems_bytes = bytes[offset2..].as_ptr();
            let ptr_elems: *const u32 = ptr_elems_bytes.cast();

            core::slice::from_raw_parts(ptr_elems, num_elems)
        }
        .to_vec();

        Matrix { rows, cols, elems }
    }
}

fn main() {
    let library = VulkanLibrary::new().unwrap();
    let instance = Instance::new(
        library,
        InstanceCreateInfo {
            flags: InstanceCreateFlags::ENUMERATE_PORTABILITY,
            ..Default::default()
        },
    )
    .unwrap();

    let device_extensions = DeviceExtensions {
        khr_storage_buffer_storage_class: true,
        ..DeviceExtensions::empty()
    };
    let (physical_device, queue_family_index) = instance
        .enumerate_physical_devices()
        .unwrap()
        .filter(|p| p.supported_extensions().contains(&device_extensions))
        .filter_map(|p| {
            p.queue_family_properties()
                .iter()
                .position(|q| {
                    q.queue_flags
                        .intersects(QueueFlags::COMPUTE | QueueFlags::TRANSFER)
                })
                .map(|i| (p, i as u32))
        })
        .min_by_key(|(p, _)| match p.properties().device_type {
            PhysicalDeviceType::DiscreteGpu => 0,
            PhysicalDeviceType::IntegratedGpu => 1,
            PhysicalDeviceType::VirtualGpu => 2,
            PhysicalDeviceType::Cpu => 3,
            PhysicalDeviceType::Other => 4,
            _ => 5,
        })
        .unwrap();

    println!(
        "Using device: {} (type: {:?}). API version: {}",
        physical_device.properties().device_name,
        physical_device.properties().device_type,
        physical_device.api_version(),
    );

    let (device, queue) = {
        let (device, mut queues) = Device::new(
            physical_device,
            DeviceCreateInfo {
                enabled_extensions: device_extensions,
                queue_create_infos: vec![QueueCreateInfo {
                    queue_family_index,
                    ..Default::default()
                }],
                ..Default::default()
            },
        )
        .unwrap();

        (device, queues.next().unwrap())
    };

    let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));

    let matrix_a = Matrix::random(13, MATRIX_A_NUM_ROWS, MATRIX_A_NUM_COLS);
    let matrix_b = Matrix::random(17, MATRIX_B_NUM_ROWS, MATRIX_B_NUM_COLS);

    println!("Matrix A = ({} X {})", matrix_a.rows, matrix_a.cols);
    println!("Matrix B = ({} X {})", matrix_b.rows, matrix_b.cols);
    println!("Matrix C = A X B = ({} X {})", matrix_a.rows, matrix_b.cols);

    let matrix_a_as_bytes = matrix_a.to_bytes();
    let matrix_b_as_bytes = matrix_b.to_bytes();
    let matrix_a_byte_len = matrix_a_as_bytes.len() as u64;
    let matrix_b_byte_len = matrix_b_as_bytes.len() as u64;
    let matrix_c_byte_len = (std::mem::size_of_val(&matrix_a.rows)
        + std::mem::size_of_val(&matrix_b.cols)
        + std::mem::size_of::<u32>() * (matrix_a.rows * matrix_b.cols) as usize)
        as u64;

    let matrix_a_host_local = Buffer::from_iter(
        memory_allocator.clone(),
        BufferCreateInfo {
            usage: BufferUsage::TRANSFER_SRC,
            ..Default::default()
        },
        AllocationCreateInfo {
            memory_type_filter: MemoryTypeFilter::HOST_SEQUENTIAL_WRITE
                | MemoryTypeFilter::PREFER_DEVICE,
            ..Default::default()
        },
        matrix_a_as_bytes,
    )
    .unwrap();

    let matrix_b_host_local = Buffer::from_iter(
        memory_allocator.clone(),
        BufferCreateInfo {
            usage: BufferUsage::TRANSFER_SRC,
            ..Default::default()
        },
        AllocationCreateInfo {
            memory_type_filter: MemoryTypeFilter::HOST_SEQUENTIAL_WRITE
                | MemoryTypeFilter::PREFER_DEVICE,
            ..Default::default()
        },
        matrix_b_as_bytes,
    )
    .unwrap();

    let matrix_a_device_local = Buffer::new_slice::<u8>(
        memory_allocator.clone(),
        BufferCreateInfo {
            usage: BufferUsage::STORAGE_BUFFER | BufferUsage::TRANSFER_DST,
            ..Default::default()
        },
        AllocationCreateInfo {
            memory_type_filter: MemoryTypeFilter::PREFER_DEVICE,
            ..Default::default()
        },
        matrix_a_byte_len,
    )
    .unwrap();

    let matrix_b_device_local = Buffer::new_slice::<u8>(
        memory_allocator.clone(),
        BufferCreateInfo {
            usage: BufferUsage::STORAGE_BUFFER | BufferUsage::TRANSFER_DST,
            ..Default::default()
        },
        AllocationCreateInfo {
            memory_type_filter: MemoryTypeFilter::PREFER_DEVICE,
            ..Default::default()
        },
        matrix_b_byte_len,
    )
    .unwrap();

    let matrix_c_buf = Buffer::new_slice::<u8>(
        memory_allocator.clone(),
        BufferCreateInfo {
            usage: BufferUsage::STORAGE_BUFFER,
            ..Default::default()
        },
        AllocationCreateInfo {
            memory_type_filter: MemoryTypeFilter::HOST_RANDOM_ACCESS
                | MemoryTypeFilter::PREFER_DEVICE,
            ..Default::default()
        },
        matrix_c_byte_len,
    )
    .unwrap();

    let command_buffer_allocator = Arc::new(StandardCommandBufferAllocator::new(
        device.clone(),
        Default::default(),
    ));

    // Transfer data to the GPU
    let transfer_command_buffer = {
        let mut builder = AutoCommandBufferBuilder::primary(
            command_buffer_allocator.clone(),
            queue.queue_family_index(),
            CommandBufferUsage::OneTimeSubmit,
        )
        .unwrap();

        builder
            .copy_buffer(CopyBufferInfo::buffers(
                matrix_a_host_local,
                matrix_a_device_local.clone(),
            ))
            .unwrap()
            .copy_buffer(CopyBufferInfo::buffers(
                matrix_b_host_local,
                matrix_b_device_local.clone(),
            ))
            .unwrap();

        builder.build().unwrap()
    };

    transfer_command_buffer
        .execute(queue.clone())
        .unwrap()
        .then_signal_fence_and_flush()
        .unwrap()
        .wait(None)
        .unwrap();

    let pipeline = {
        let compute_shader = cs::load(device.clone()).unwrap();
        let compute_shader_entry_point = compute_shader.entry_point("main").unwrap();
        let compute_stage = PipelineShaderStageCreateInfo::new(compute_shader_entry_point);

        let layout = PipelineLayout::new(
            device.clone(),
            PipelineDescriptorSetLayoutCreateInfo::from_stages([&compute_stage])
                .into_pipeline_layout_create_info(device.clone())
                .unwrap(),
        )
        .unwrap();

        ComputePipeline::new(
            device.clone(),
            None,
            ComputePipelineCreateInfo::stage_layout(compute_stage, layout.clone()),
        )
        .unwrap()
    };

    let descriptor_set_allocator = Arc::new(StandardDescriptorSetAllocator::new(
        device.clone(),
        Default::default(),
    ));
    let descriptor_set_layout = pipeline.layout().set_layouts()[0].clone();
    let descriptor_set = DescriptorSet::new(
        descriptor_set_allocator,
        descriptor_set_layout,
        [
            WriteDescriptorSet::buffer(0, matrix_a_device_local.clone()),
            WriteDescriptorSet::buffer(1, matrix_b_device_local.clone()),
            WriteDescriptorSet::buffer(2, matrix_c_buf.clone()),
        ],
        [],
    )
    .unwrap();

    let command_buffer = {
        let mut command_buffer_builder = AutoCommandBufferBuilder::primary(
            command_buffer_allocator,
            queue.queue_family_index(),
            CommandBufferUsage::MultipleSubmit,
        )
        .unwrap();

        command_buffer_builder
            .bind_pipeline_compute(pipeline.clone())
            .unwrap()
            .bind_descriptor_sets(
                PipelineBindPoint::Compute,
                pipeline.layout().clone(),
                0,
                descriptor_set,
            )
            .unwrap();
        unsafe {
            command_buffer_builder
                .dispatch([
                    MATRIX_A_NUM_ROWS.div_ceil(8),
                    MATRIX_B_NUM_COLS.div_ceil(8),
                    1,
                ])
                .unwrap();
        }

        command_buffer_builder.build().unwrap()
    };

    // Computing Matrix Multiplication on GPU
    let start = Instant::now();
    sync::now(device.clone())
        .then_execute(queue.clone(), command_buffer.clone())
        .unwrap()
        .then_signal_fence_and_flush()
        .unwrap()
        .wait(None)
        .unwrap();
    let gpu_tm = start.elapsed();
    println!("(a) GPU matrix multiply: {:?}", gpu_tm);

    // reading GPU-computed matrix multiplication result
    let gpu_result_as_bytes = matrix_c_buf.read().unwrap();
    let gpu_result = Matrix::from_bytes(&gpu_result_as_bytes);

    // Computing Parallel Matrix Multiplication on CPU, using Rayon
    let start = Instant::now();
    let cpu_result_rayon = matrix_a.par_mul(&matrix_b);
    let cpu_tm_rayon = start.elapsed();
    println!("(b) CPU parallel matrix multiply: {:?}", cpu_tm_rayon);

    let start = Instant::now();
    let cpu_result_single_threaded = matrix_a * matrix_b;
    let cpu_tm_single_threaded = start.elapsed();
    println!(
        "(c) CPU single-threaded matrix multiply: {:?}",
        cpu_tm_single_threaded
    );

    assert_eq!(gpu_result, cpu_result_single_threaded);
    assert_eq!(cpu_result_single_threaded, cpu_result_rayon);

    println!(
        "Speed Up (b/a): {}x\nSpeed Up (c/a): {}x\nSpeed Up (c/b): {}x",
        cpu_tm_rayon.as_nanos() / gpu_tm.as_nanos(),
        cpu_tm_single_threaded.as_nanos() / gpu_tm.as_nanos(),
        cpu_tm_single_threaded.as_nanos() / cpu_tm_rayon.as_nanos(),
    );
}

mod cs {
    // Compiles shader
    vulkano_shaders::shader! {
        ty: "compute",
        path: "./matrix_multiply.glsl",
        vulkan_version: "1.2",
    }
}

matrix_multiply.glsl

#version 460
#pragma shader_stage(compute)

layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

layout(set = 0, binding = 0) buffer readonly MatrixA
{
  uint rows;
  uint cols;
  uint[] elems;
} matrix_a;

layout(set = 0, binding = 1) buffer readonly MatrixB
{
  uint rows;
  uint cols;
  uint[] elems;
} matrix_b;

layout(set = 0, binding = 2) buffer writeonly MatrixC
{
  uint rows;
  uint cols;
  uint[] elems;
} matrix_c;

void
main()
{
  const uint row_idx = gl_GlobalInvocationID.x;
  const uint col_idx = gl_GlobalInvocationID.y;

  if (row_idx >= matrix_a.rows || col_idx >= matrix_b.cols) {
    return;
  }

  if ((row_idx == 0) && (col_idx == 0)) {
    matrix_c.rows = matrix_a.rows;
    matrix_c.cols = matrix_b.cols;
  }

  uint sum = 0;
  for (uint i = 0; i < matrix_a.cols; i++) {
    sum += matrix_a.elems[row_idx * matrix_a.cols + i] * matrix_b.elems[i * matrix_b.cols + col_idx];
  }

  matrix_c.elems[row_idx * matrix_b.cols + col_idx] = sum;
}

Comparation is not fair, CPU is doing calculation in single thread. Dot product those two arrays (that is more than shader does), using ndarray takes 275 micro seconds on 16 vCore CPU.

Using device: Radeon RX 590 Series (type: DiscreteGpu). API version: 1.3.209
GPU matrix multiply: 28.4929ms
CPU matrix multiply: 548.6µs <-- dot product using ndarray
Speed Up: 0x

Anyhow, example is nice :) TNX!

Tested by RX6700XT and Ryzen 7 5800X3D

Device: AMD Radeon RX 6700 XT
Vulkan API: 1.2.0
Queue Count: 1  Compute: true   Graphics: true
Queue Count: 2  Compute: true   Graphics: false
Queue Count: 2  Compute: false  Graphics: false
GPU matrix multiply: 10.1086ms
CPU matrix multiply: 2.7680681s
Speed Up: 273

Thanks for reporting this @macchky.

Totally agreed with you @mjaric.

I was getting a really poor performance on a rtx 2070 super (~650 ms). Does running this in WSL affect the performance?

I was getting a really poor performance on a rtx 2070 super (~650 ms). Does running this in WSL affect the performance?

There could be many possible reasons for why it ran slow on some platform. As I've never tried running it on WSL, can't really say anything.

I was getting a really poor performance on a rtx 2070 super (~650 ms). Does running this in WSL affect the performance?

There could be many possible reasons for why it ran slow on some platform. As I've never tried running it on WSL, can't really say anything.

Seems like WSL is indeed the cause of low performance. The same code takes 10 ms to complete in Windows but takes over 600 ms in WSL

I was getting a really poor performance on a rtx 2070 super (~650 ms). Does running this in WSL affect the performance?

There could be many possible reasons for why it ran slow on some platform. As I've never tried running it on WSL, can't really say anything.

Seems like WSL is indeed the cause of low performance. The same code takes 10 ms to complete in Windows but takes over 600 ms in WSL

Totally possible.

random comment, you should rework this to use VK_KHR_cooperative_matrix as a good example to compare it vs native shader

random comment, you should rework this to use VK_KHR_cooperative_matrix as a good example to compare it vs native shader

Thanks for the suggestion though I'm not actively maintaining it. You may send me a patch and I'll update the gist.

I just updated this to work with the latest version of vulkano, numbers look better. See https://gist.github.com/itzmeanjan/84613bc7595372c5e6b6c22481d42f9a?permalink_comment_id=3883582#gistcomment-3883582.

Addressed @mjaric 's concern of comparison not being fair, which I totally agree with, by adding rayon -based CPU parallel matrix multiplication implementation.

Updated the gist again, now it supports parallel matrix multiplication on GPU with any two matrices of compatible dimensions.

Tested by RX6700XT and Ryzen 7 5800X3D

Device: AMD Radeon RX 6700 XT
Vulkan API: 1.2.0
Queue Count: 1  Compute: true   Graphics: true
Queue Count: 2  Compute: true   Graphics: false
Queue Count: 2  Compute: false  Graphics: false
GPU matrix multiply: 10.1086ms
CPU matrix multiply: 2.7680681s
Speed Up: 273

Tested by RX6700XT on Ryzen 7 9700X

Using device: AMD Radeon RX 6700 XT (type: DiscreteGpu). API version: 1.3.281
Matrix A = (1774 X 2048)
Matrix B = (2048 X 1024)
Matrix C = A X B = (1774 X 1024)
(a) GPU matrix multiply: 46.6836ms
(b) CPU parallel matrix multiply: 786.9838ms
(c) CPU single-threaded matrix multiply: 13.3806369s
Speed Up (b/a): 16x
Speed Up (c/a): 286x
Speed Up (c/b): 17x

GPU time is slower than before.

I don't think so. Before $A_{1024x1024}$ X $B_{1024x1024}$ was being benchmarked. Now matrix dimensions have changed to ensure that operand matrices don't need to be square.