julius/wgpu_compute_shader
1
0
Fork 1
Code Issues Pull Requests Packages Projects Releases Wiki Activity

🚀 Builds and runs on Linux and MacOS, result still wrong

Browse Source 
Signed-off-by: Julius Koskela <julius.koskela@unikie.com>
This commit is contained in:
Julius Koskela 2023-12-22 10:51:26 +02:00
parent e1e9b8dd18
commit bf16edd3aa
Signed by: julius
GPG Key ID: 5A7B7F4897C2914B
10 changed files with 254 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@ -1 +1,2 @@
/target
target
result

17
Cargo.lock generated
View File

@ -247,6 +247,12 @@ version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
[[package]]
name = "fixedbitset"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
[[package]]
name = "flume"
version = "0.11.0"
@ -598,6 +604,7 @@ dependencies = [
"indexmap",
"log",
"num-traits",
"petgraph",
"rustc-hash",
"spirv",
"termcolor",
@ -711,6 +718,16 @@ version = "1.0.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "de3145af08024dea9fa9914f381a17b8fc6034dfb00f3a84013f7ff43f29ed4c"
[[package]]
name = "petgraph"
version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e1d3afd2628e69da2be385eb6f2fd57c8ac7977ceeff6dc166ff1657b0e386a9"
dependencies = [
"fixedbitset",
"indexmap",
]
[[package]]
name = "pkg-config"
version = "0.3.28"

2
Cargo.toml
View File

@ -6,7 +6,7 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
wgpu = { version = "0.18", features = ["vulkan-portability"] }
wgpu = { version = "0.18", features = ["vulkan-portability", "spirv" ] }
env_logger = "0.9.1"
pollster = "0.2.5"
futures-intrusive = "0.4"

9
build.rs
View File

@ -1,9 +0,0 @@
fn main() {
if cfg!(target_os = "linux") {
// println!("cargo:rustc-link-lib=X11");
// println!("cargo:rustc-link-lib=Xcursor");
// println!("cargo:rustc-link-lib=Xrandr");
// println!("cargo:rustc-link-lib=Xi");
println!("cargo:rustc-link-lib=vulkan");
}
}

2
flake.nix
View File

@ -7,7 +7,7 @@
nixpkgs,
}: {
packages.x86_64-linux.default = nixpkgs.legacyPackages.x86_64-linux.callPackage ./default.nix {};
devShells.x86_64-linux.default = nixpkgs.legacyPackages.x86_64-linux.callPackage ./shell.nix {};
devShells.x86_64-linux = nixpkgs.legacyPackages.x86_64-linux.callPackage ./shell.nix {};
formatter.x86_64-linux = nixpkgs.legacyPackages.x86_64-linux.alejandra;
};
}

65
shell.nix
View File

@ -1,51 +1,30 @@
{pkgs}:
with pkgs; let
build = pkgs.callPackage ./default.nix {};
let
nixpkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/archive/nixos-unstable.tar.gz") {};
in
with nixpkgs; let
buildInputs = [
udev
alsa-lib
vulkan-loader
vulkan-headers
vulkan-tools
vulkan-validation-layers
xorg.libX11
xorg.libXcursor
xorg.libXi
xorg.libXrandr # To use the x11 feature
libxkbcommon
wayland # To use the wayland feature
];
in
mkShell {
inherit build;
packages = [
libX11
libXcursor
libXrandr
libXi
vulkan-headers
vulkan-loader
vulkan-validation-layers
vulkan-tools
inherit buildInputs;
nativeBuildInputs = [
pkg-config
git
gcc
gnumake
cmake
glibc
python3
shaderc
];
inputsFrom = [
cmake
shaderc
];
RUST_BACKTRACE = "1";
LD_LIBRARY_PATH="${pkgs.libX11}/lib:${pkgs.libXcursor}/lib:${pkgs.libXrandr}/lib:${pkgs.libXi}/lib:${pkgs.vulkan-loader}/lib:${pkgs.vulkan-validation-layers}/lib:${pkgs.vulkan-tools}/lib:${pkgs.vulkan-headers}/lib:${pkgs.stdenv.cc.cc.lib}:${pkgs.stdenv.cc.cc.lib64}:$LD_LIBRARY_PATH";
cargoBuildFlags = ["--release --features build-from-source"];
LD_LIBRARY_PATH = lib.makeLibraryPath buildInputs;
shellHook = ''
export VK_ICD_FILENAMES=${vulkan-loader}/share/vulkan/icd.d/radeon_icd64.json
export VK_LAYER_PATH=${vulkan-validation-layers}/share/vulkan/explicit_layer.d
export VK_INSTANCE_LAYERS=VK_LAYER_KHRONOS_validation
export VK_DEVICE_LAYERS=VK_LAYER_KHRONOS_validation
export VK_LOADER_DEBUG=all
export VK_LOADER_DEBUG_FILE=/tmp/vulkan.log
export VK_INSTANCE_EXTENSIONS=VK_EXT_debug_utils
export VK_DEVICE_EXTENSIONS=VK_EXT_debug_utils
export VK_LAYER_ENABLES=VK_LAYER_KHRONOS_validation
export VK_LAYER_DISABLES=VK_LAYER_LUNARG_api_dump
export VK_LAYER_PATH=${vulkan-validation-layers}/share/vulkan/explicit_layer.d
pkg-config --cflags fontconfig fontconfig >= 2.11.1 --libs vulkan
export LD_LIBRARY_PATH=${lib.makeLibraryPath buildInputs}
'';
}

124
src/main.rs
View File

@ -1,18 +1,23 @@
// ... existing imports ...
mod ppu;
use std::time::Instant;
use wgpu::util::DeviceExt;
use bytemuck;
async fn run() {
let instance = wgpu::Instance::default();
let instance_desc = wgpu::InstanceDescriptor {
backends: wgpu::Backends::VULKAN,
dx12_shader_compiler: wgpu::Dx12Compiler::Fxc,
gles_minor_version: wgpu::Gles3MinorVersion::default(),
flags: wgpu::InstanceFlags::empty(),
};
let instance = wgpu::Instance::new(instance_desc);
println!("instance {:?}", instance);
let adapter = instance
.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
force_fallback_adapter: false,
force_fallback_adapter: true,
compatible_surface: None,
})
.await;
@ -48,14 +53,24 @@ async fn run() {
let start_instant = Instant::now();
let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: None,
//source: wgpu::ShaderSource::SpirV(bytes_to_u32(include_bytes!("alu.spv")).into()),
// source: wgpu::ShaderSource::SpirV(bytes_to_u32(include_bytes!("shader.spv")).into()),
source: wgpu::ShaderSource::Wgsl(include_str!("shader.wgsl").into()),
});
println!("shader compilation {:?}", start_instant.elapsed());
let matrix_size = std::mem::size_of::<[[f32; 4]; 4]>();
let matrix_a_data = [[1.0; 4]; 4]; // Example data
let matrix_b_data = [[2.0; 4]; 4]; // Example data
let matrix_a_data = [
[1.0, 2.0, 3.0, 4.0], // Example data
[5.0, 6.0, 7.0, 8.0],
[9.0, 10.0, 11.0, 12.0],
[13.0, 14.0, 15.0, 16.0],
];
let matrix_b_data = [
[1.0, 2.0, 3.0, 4.0], // Example data
[5.0, 6.0, 7.0, 8.0],
[9.0, 10.0, 11.0, 12.0],
[13.0, 14.0, 15.0, 16.0],
];
let matrix_a_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("Matrix A Buffer"),
@ -67,18 +82,38 @@ async fn run() {
contents: bytemuck::bytes_of(&matrix_b_data),
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
});
// Create Matrix C as a storage buffer
let matrix_c_buf = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Matrix C Buffer"),
label: Some("Matrix C Storage Buffer"),
size: matrix_size as u64,
usage: wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::MAP_READ
| wgpu::BufferUsages::COPY_DST,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, // Changed to COPY_SRC
mapped_at_creation: false,
});
let query_buf = if features.contains(wgpu::Features::TIMESTAMP_QUERY) {
// Create a separate staging buffer for reading back data to CPU
let matrix_c_staging_buf = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Matrix C Staging Buffer"),
size: matrix_size as u64,
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, // Staging buffer
mapped_at_creation: false,
});
// Query Result Buffer - For GPU to write query results
let query_result_buf = if features.contains(wgpu::Features::TIMESTAMP_QUERY) {
Some(device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Query Buffer"),
label: Some("Query Result Buffer"),
size: 16, // Enough for two 64-bit timestamps
usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
}))
} else {
None
};
// Readback Buffer - For CPU to read query results
let readback_buf = if features.contains(wgpu::Features::TIMESTAMP_QUERY) {
Some(device.create_buffer(&wgpu::BufferDescriptor {
label: Some("Readback Buffer"),
size: 16, // Enough for two 64-bit timestamps
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
@ -89,17 +124,41 @@ async fn run() {
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: None,
entries: &[wgpu::BindGroupLayoutEntry {
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true }, // Adjust based on shader needs
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true }, // Adjust based on shader needs
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
}],
},
// ... any other entries needed ...
],
});
let compute_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: None,
bind_group_layouts: &[&bind_group_layout],
@ -145,19 +204,31 @@ async fn run() {
if let Some(query_set) = &query_set {
encoder.write_timestamp(query_set, 1);
}
encoder.copy_buffer_to_buffer(&matrix_c_buf, 0, &matrix_c_buf, 0, matrix_size as u64);
encoder.copy_buffer_to_buffer(
&matrix_c_buf,
0,
&matrix_c_staging_buf,
0,
matrix_size as u64,
);
if let Some(query_set) = &query_set {
if let Some(query_buf) = &query_buf {
if let Some(query_result_buf) = &query_result_buf {
encoder.write_timestamp(query_set, 1);
encoder.resolve_query_set(query_set, 0..2, query_buf, 0);
encoder.resolve_query_set(query_set, 0..2, query_result_buf, 0);
// Ensure to copy after the query set is resolved
if let Some(readback_buf) = &readback_buf {
encoder.copy_buffer_to_buffer(query_result_buf, 0, readback_buf, 0, 16);
}
}
}
queue.submit(Some(encoder.finish()));
// Assuming query_buf has been properly initialized earlier
// Assuming readback_buf has been properly initialized earlier
let buf_slice = matrix_c_buf.slice(..);
let query_slice = query_buf.as_ref().map(|buf| buf.slice(..)); // Adjust this line
let buf_slice = matrix_c_staging_buf.slice(..);
let query_slice = readback_buf.as_ref().map(|buf| buf.slice(..)); // Use readback_buf instead of query_buf
let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
@ -175,7 +246,14 @@ async fn run() {
if let Some(Ok(())) = receiver.receive().await {
let data_raw = &*buf_slice.get_mapped_range();
println!("compute shader result: {:?}", data_raw);
// map the data to a slice of f64s
let parsed = data_raw
.chunks_exact(std::mem::size_of::<f32>())
.map(|b| {
let a = bytemuck::from_bytes::<f32>(b);
a.clone()
}).collect::<Vec<f32>>();
println!("matrix_c_staging_buf {:?}", parsed);
}
if features.contains(wgpu::Features::TIMESTAMP_QUERY) {

73
src/ppu.rs Normal file
View File

@ -0,0 +1,73 @@
use std::time::Instant;
use wgpu::util::DeviceExt;
use bytemuck;
pub struct PPU {
instance: wgpu::Instance,
adapter: wgpu::Adapter,
device: wgpu::Device,
queue: wgpu::Queue,
query_set: Option<wgpu::QuerySet>,
// shader: wgpu::ShaderModule,
}
impl PPU {
pub async fn init() -> Self {
let instance_desc = wgpu::InstanceDescriptor {
backends: wgpu::Backends::VULKAN,
dx12_shader_compiler: wgpu::Dx12Compiler::Fxc,
gles_minor_version: wgpu::Gles3MinorVersion::default(),
flags: wgpu::InstanceFlags::empty(),
};
let instance = wgpu::Instance::new(instance_desc);
let adapter = instance
.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
force_fallback_adapter: true,
compatible_surface: None,
}).await.expect("Failed to find an appropriate adapter");
let features = adapter.features();
let (device, queue) = adapter
.request_device(
&wgpu::DeviceDescriptor {
label: None,
features: features & wgpu::Features::TIMESTAMP_QUERY,
limits: Default::default(),
},
None,
)
.await
.unwrap();
let query_set = if features.contains(wgpu::Features::TIMESTAMP_QUERY) {
Some(device.create_query_set(&wgpu::QuerySetDescriptor {
count: 2,
ty: wgpu::QueryType::Timestamp,
label: None,
}))
} else {
None
};
println!("PPU initialized");
Self {
instance,
adapter,
device,
queue,
query_set,
}
}
}
impl std::fmt::Display for PPU {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "PPU")?;
write!(f, "instance {:?}", self.instance)?;
write!(f, "adapter {:?}", self.adapter)?;
write!(f, "device {:?}", self.device)?;
write!(f, "queue {:?}", self.queue)?;
write!(f, "query_set {:?}", self.query_set)?;
Ok(())
}
}

25
src/shader.spv Normal file
View File

@ -0,0 +1,25 @@
#version 450
layout(set = 0, binding = 0) buffer MatrixA {
float data[4][4];
} matrixA;
layout(set = 0, binding = 1) buffer MatrixB {
float data[4][4];
} matrixB;
layout(set = 0, binding = 2) buffer MatrixC {
float data[4][4];
} matrixC;
void main() {
uint row = gl_GlobalInvocationID.y;
uint col = gl_GlobalInvocationID.x;
float sum = 0.0;
for (uint k = 0; k < 4; ++k) {
sum += matrixA.data[row][k] * matrixB.data[k][col];
}
matrixC.data[row][col] = sum;
}

4
src/shader.wgsl
View File

@ -1,5 +1,5 @@
struct Matrix {
data: array<array<f32, 4>, 4>; // Assuming 4x4 matrices
data: array<array<f32, 4>, 4>, // Corrected: Use a comma instead of a semicolon
};
@group(0) @binding(0)
@ -17,7 +17,7 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
let col: u32 = global_id.x;
var sum: f32 = 0.0;
for (var k: u32 = 0; k < 4u; k = k + 1u) {
for (var k: u32 = 0u; k < 4u; k = k + 1u) {
sum = sum + matrixA.data[row][k] * matrixB.data[k][col];
}
matrixC.data[row][col] = sum;