refactor: seperate crabml-vulkan and crabml-wgpu (#192)

* extract wgpu-backend

* extract wgpu crate

* seperate crabml-vulkan

* seperate test cases

* fix workflow

* fix workflow

* tune workflow

* restrict threads

* retry test

.github/workflows/rust.yml

@@ -42,4 +42,4 @@ jobs:
           sudo apt install vulkan-validationlayers
 
       - name: Run tests
-        run: cargo test --workspace
+        run: RUST_TEST_THREADS=1 cargo test --workspace

Cargo.toml

@@ -2,6 +2,8 @@
 resolver = "2"
 members = [
     "crabml-core",
+    "crabml-wgpu",
+    "crabml-vulkan",
     "crabml-llama2",
     "crabml-cli",
 ]
@@ -18,3 +20,5 @@ repository = "https://github.com/crabml/crabml"
 [workspace.dependencies]
 crabml = { version = "0.1.0", path = "crabml-core" }
 crabml-llama2 = { version = "0.1.0", path = "crabml-llama2" }
+crabml-wgpu = { version = "0.1.0", path = "crabml-wgpu" }
+crabml-vulkan = { version = "0.1.0", path = "crabml-vulkan" }

crabml-cli/Cargo.toml

@@ -11,6 +11,7 @@ rayon = "1"
 num_cpus = "1.16.0"
 clap = { version = "4.0", features = ["derive"] }
 crabml-llama2 = { workspace = true }
+crabml-wgpu = { workspace = true }
 crabml = { workspace = true }
 jemallocator = "0.3"
 rustyline = "9.0.0"

crabml-cli/src/main.rs

@@ -5,9 +5,6 @@ use std::time::Instant;
 
 use clap::Parser;
 use clap::ValueEnum;
-use crabml::backends::wgpu::WgpuTensor;
-use crabml::backends::wgpu::WgpuTensorDevice;
-use crabml::backends::wgpu::WgpuTensorDeviceOptions;
 use crabml::error::Result;
 use crabml::gguf::GGUFFile;
 use crabml::gguf::GGUFFileLoader;
@@ -18,6 +15,9 @@ use crabml_llama2::llama2::Llama2Runner;
 use crabml_llama2::model::CpuLlamaModelLoader;
 use crabml_llama2::GpuLlamaModel;
 use crabml_llama2::Llama2Chat;
+use crabml_wgpu::WgpuTensor;
+use crabml_wgpu::WgpuTensorDevice;
+use crabml_wgpu::WgpuTensorDeviceOptions;
 use rustyline::error::ReadlineError;
 use rustyline::Editor;
 

crabml-core/Cargo.toml

@@ -10,7 +10,6 @@ description = "crabml core package"
 int-enum = "0.5.0"
 memmap2 = "0.7.1"
 half = { version = "2.3.1" }
-wgpu = "0.19.1"
 env_logger = "0.10"
 pollster = "0.2.4"
 bytemuck = { version = "1.14.0", features = ["derive"] }

crabml-core/src/lib.rs

@@ -1,15 +1,11 @@
 #![feature(portable_simd)]
 #![feature(slice_as_chunks)]
 #![cfg_attr(target_arch = "aarch64", feature(stdarch_neon_dotprod))]
-#![feature(thread_local)]
-#![feature(lazy_cell)]
 #![feature(iter_array_chunks)]
 #![feature(lint_reasons)]
 #![allow(clippy::map_entry)]
 #![allow(clippy::comparison_chain)]
 
-#[allow(unreachable_patterns)]
-pub mod backends;
 pub mod cpu;
 pub mod error;
 pub mod gguf;

crabml-llama2/Cargo.toml

@@ -11,6 +11,7 @@ memmap2 = "0.7.1"
 rand = "0.8.5"
 num_cpus = "1.16.0"
 crabml = { workspace = true }
+crabml-wgpu = { workspace = true }
 half = { version = "2.3.1" }
 
 [dev-dependencies]

crabml-llama2/src/llama2.rs

@@ -528,11 +528,11 @@ impl<'a, T: Tensor> Llama2Runner<T> {
 #[cfg(test)]
 mod tests {
     use approx::assert_relative_eq;
-    use crabml::backends::wgpu::WgpuTensor;
-    use crabml::backends::wgpu::WgpuTensorDevice;
-    use crabml::backends::wgpu::WgpuTensorDeviceOptions;
     use crabml::cpu::CpuTensorDeviceOptions;
     use crabml::gguf::GGUFFileLoader;
+    use crabml_wgpu::WgpuTensor;
+    use crabml_wgpu::WgpuTensorDevice;
+    use crabml_wgpu::WgpuTensorDeviceOptions;
 
     use super::*;
     use crate::model::CpuLlamaModelLoader;

crabml-vulkan/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "crabml-vulkan"
+version = { workspace = true }
+edition = { workspace = true }
+license = { workspace = true }
+repository = { workspace = true }
+description = "crabml vulkan support"
+
+[dependencies]
+int-enum = "0.5.0"
+half = { version = "2.3.1" }
+env_logger = "0.10"
+pollster = "0.2.4"
+bytemuck = { version = "1.14.0", features = ["derive"] }
+vulkano = "0.34.1"
+vulkano-shaders = "0.34.0"
+crabml = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = "1.2.1"
+bencher = "0.1.5"
+approx = "0.5.1"

crabml-core/src/backends/vulkan/vulkan_device.rs → crabml-vulkan/src/vulkan_device.rs

@@ -95,7 +95,7 @@ impl VulkanTensorDevice {
 
     fn load_shaders(&mut self) {
         mod add {
-            vulkano_shaders::shader! { ty: "compute", path: "./src/backends/vulkan/shaders/add.comp" }
+            vulkano_shaders::shader! { ty: "compute", path: "./src/shaders/add.comp" }
         }
 
         let device = self.inner.device.clone();

crabml-core/src/backends/vulkan/vulkan_tensor.rs → crabml-vulkan/src/vulkan_tensor.rs

@@ -1,12 +1,13 @@
 #![allow(dead_code, unused_variables)]
 
+use crabml::error::Error;
+use crabml::error::ErrorKind;
+use crabml::error::Result;
+use crabml::gguf::GGMLType;
+use crabml::tensor::Tensor;
+use crabml::tensor::TensorStrider;
+
 use super::vulkan_device::VulkanTensorDeviceRef;
-use crate::error::Error;
-use crate::error::ErrorKind;
-use crate::error::Result;
-use crate::gguf::GGMLType;
-use crate::tensor::Tensor;
-use crate::tensor::TensorStrider;
 
 #[derive(Clone)]
 pub struct VulkanTensor {
@@ -129,7 +130,7 @@ impl Tensor for VulkanTensor {
 
     fn rope_inplace(
         self,
-        mode: crate::tensor::RopeMode,
+        mode: crabml::tensor::RopeMode,
         pos: usize,
         rope_dims: usize,
     ) -> Result<Self> {
@@ -187,14 +188,14 @@ impl Tensor for VulkanTensor {
 
 #[cfg(test)]
 mod tests {
+    use crabml::error::Result;
+    use crabml::tensor::Tensor;
+
     use super::VulkanTensor;
-    use crate::backends::vulkan::vulkan_device::VulkanTensorDevice;
-    use crate::backends::vulkan::vulkan_device::VulkanTensorDeviceOptions;
-    use crate::error::Result;
-    use crate::tensor::Tensor;
+    use crate::vulkan_device::VulkanTensorDevice;
+    use crate::vulkan_device::VulkanTensorDeviceOptions;
 
     #[test]
-    #[ignore]
     fn test_add() -> Result<()> {
         let d = VulkanTensorDevice::new(VulkanTensorDeviceOptions::default());
 

crabml-wgpu/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "crabml-wgpu"
+version = { workspace = true }
+edition = { workspace = true }
+license = { workspace = true }
+repository = { workspace = true }
+description = "crabml wgpu support"
+
+[dependencies]
+int-enum = "0.5.0"
+half = { version = "2.3.1" }
+wgpu = "0.19.1"
+env_logger = "0.10"
+pollster = "0.2.4"
+bytemuck = { version = "1.14.0", features = ["derive"] }
+byteorder = "1.5.0"
+crabml = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = "1.2.1"
+bencher = "0.1.5"
+approx = "0.5.1"

crabml-core/src/backends/wgpu/mod.rs → crabml-wgpu/src/lib.rs

@@ -1,3 +1,6 @@
+#![feature(thread_local)]
+#![feature(lazy_cell)]
+
 mod meta;
 mod wgpu_device;
 mod wgpu_tensor;

crabml-core/src/backends/wgpu/wgpu_device.rs → crabml-wgpu/src/wgpu_device.rs

@@ -3,9 +3,9 @@ use std::cell::RefCell;
 use std::collections::HashMap;
 use std::rc::Rc;
 
+use crabml::tensor::Tensor;
 use wgpu::util::DeviceExt;
 
-use crate::tensor::Tensor;
 pub struct WgpuTensorDeviceOptions {
     pub staging_buf_bytes: usize,
 

crabml-core/src/backends/wgpu/wgpu_tensor.rs → crabml-wgpu/src/wgpu_tensor.rs

@@ -1,20 +1,20 @@
 use std::rc::Rc;
 
+use crabml::error::ErrorKind;
+use crabml::error::Result;
+use crabml::gguf::GGMLType;
+use crabml::tensor::RopeMode;
+use crabml::tensor::Tensor;
+use crabml::tensor::TensorStrider;
 use wgpu::util::DeviceExt;
 
 use super::meta::ConcatenateMeta;
 use super::meta::MatmulMeta;
 use super::meta::RmsNormMeta;
 use super::WgpuTensorDeviceRef;
-use crate::backends::wgpu::meta::BatchMatmulMeta;
-use crate::backends::wgpu::meta::ContiguousMeta;
-use crate::backends::wgpu::meta::RopeMeta;
-use crate::error::ErrorKind;
-use crate::error::Result;
-use crate::gguf::GGMLType;
-use crate::tensor::RopeMode;
-use crate::tensor::Tensor;
-use crate::tensor::TensorStrider;
+use crate::meta::BatchMatmulMeta;
+use crate::meta::ContiguousMeta;
+use crate::meta::RopeMeta;
 
 #[derive(Clone)]
 pub struct WgpuTensor {
@@ -782,15 +782,15 @@ mod tests {
     use std::sync::LazyLock;
 
     use approx::assert_relative_eq;
+    use crabml::error::Result;
+    use crabml::gguf::GGMLType;
+    use crabml::tensor::RopeMode;
+    use crabml::tensor::Tensor;
 
     use super::WgpuTensor;
-    use crate::backends::wgpu::WgpuTensorDevice;
-    use crate::backends::wgpu::WgpuTensorDeviceOptions;
-    use crate::backends::wgpu::WgpuTensorDeviceRef;
-    use crate::error::Result;
-    use crate::gguf::GGMLType;
-    use crate::tensor::RopeMode;
-    use crate::tensor::Tensor;
+    use crate::WgpuTensorDevice;
+    use crate::WgpuTensorDeviceOptions;
+    use crate::WgpuTensorDeviceRef;
 
     #[thread_local]
     static DEVICE: LazyLock<WgpuTensorDeviceRef> = LazyLock::new(|| {