feat: add "runtime" and "comptime" features to select when slang shaders need to be compiled (#2)

* feat: make stensor compatible with slang-hal’s new comptime feature

* Release v0.3.0

Author: sebcrozet

.github/workflows/ci.yaml

@@ -96,7 +96,7 @@ jobs:
           sweep-cache: true
 
       - name: Run clippy lints
-        run: SLANG_DIR=$SLANG_DIR cargo clippy --locked --workspace --all-targets -- --deny warnings
+        run: SLANG_DIR=$SLANG_DIR cargo clippy --locked --workspace --all-targets --features runtime -- --deny warnings
 
   # Check documentation.
   doc:
@@ -128,7 +128,7 @@ jobs:
           sweep-cache: true
 
       - name: Check documentation
-        run: SLANG_DIR=$SLANG_DIR cargo doc --locked --workspace --document-private-items --no-deps
+        run: SLANG_DIR=$SLANG_DIR cargo doc --locked --workspace --document-private-items --features runtime --no-deps
     # Testing.
   test:
     needs: setup-slang # Depends on setup-slang
@@ -163,4 +163,4 @@ jobs:
           sweep-cache: true
       - name: Run Cargo Tests
         run: |
-          SLANG_DIR=$SLANG_DIR LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose
+          SLANG_DIR=$SLANG_DIR LIBGL_ALWAYS_SOFTWARE=1 cargo test --verbose --features runtime

.run/Check.run.xml .run/Check (runtime).run.xml.run/Check.run.xml renamed to .run/Check (runtime).run.xml

@@ -1,7 +1,7 @@
 <component name="ProjectRunConfigurationManager">
-  <configuration default="false" name="Check" type="CargoCommandRunConfiguration" factoryName="Cargo Command" nameIsGenerated="true">
+  <configuration default="false" name="Check (runtime)" type="CargoCommandRunConfiguration" factoryName="Cargo Command">
     <option name="buildProfileId" value="dev" />
-    <option name="command" value="check" />
+    <option name="command" value="check --features runtime" />
     <option name="workingDirectory" value="file://$PROJECT_DIR$" />
     <envs />
     <option name="emulateTerminal" value="true" />

Cargo.toml

@@ -3,27 +3,33 @@ name = "stensor"
 authors = ["Sébastien Crozet <sebcrozet@dimforge.com>"]
 description = "Cross-platform GPU tensor library with Slang and Rust."
 repository = "https://github.com/dimforge/stensor"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2024"
 license = "Apache-2.0"
 
 [features]
+comptime = [ "slang-hal/comptime" ]
+runtime = [ "slang-hal/runtime" ]
+
+webgpu = [ "slang-hal/webgpu" ]
+vulkan = [ "slang-hal/vulkan" ]
+metal = [ "slang-hal/metal" ]
+cpu = [ "slang-hal/cpu" ]
 cuda = [ "cudarc", "slang-hal/cuda" ]
 cublas = [ "slang-hal/cublas" ]
 
 [dependencies]
-wgpu = "27"
 encase = "0.12"
 bytemuck = "1"
 nalgebra = { version = "0.34", features = ["encase"] }
 
 cudarc = {  version = "0.16", optional = true }
 
-minislang = "0.2"
-slang-hal = { version = "0.2", features = ["derive"] }
+slang-hal = { version = "0.3", features = ["derive"] }
 include_dir = "0.7"
 
 [dev-dependencies]
+minislang = "0.3"
 nalgebra = { version = "0.34", features = ["rand"] }
 futures-test = "0.3"
 serial_test = "3"
@@ -32,12 +38,16 @@ async-std = { version = "1", features = ["attributes"] }
 plotly = "0.12.1"
 indexmap = "2"
 anyhow = "1"
+wgpu = "27"
 
 [build-dependencies]
-minislang = "0.1"
+minislang = "0.3"
+slang-hal-build = "0.3"
+include_dir = "0.7"
 
 [patch.crates-io]
 #shader-slang = { path = "../slang-rs" }
 #minislang = { path = "../slang-hal/crates/minislang" }
+#slang-hal-build = { path = "../slang-hal/crates/slang-hal-build" }
 #slang-hal-derive = { path = "../slang-hal/crates/slang-hal-derive" }
 #slang-hal = { path = "../slang-hal/crates/slang-hal" }

build.rs

@@ -1,17 +1,24 @@
-use minislang::{SlangCompiler, shader_slang::CompileTarget};
-use std::path::PathBuf;
-use std::str::FromStr;
+#[cfg(feature = "comptime")]
+use std::env;
 
+#[cfg(not(feature = "comptime"))]
+pub fn main() {}
+
+#[cfg(feature = "comptime")]
 pub fn main() {
-    let slang = SlangCompiler::new(vec![PathBuf::from_str("./shaders").unwrap()]);
+    use slang_hal_build::ShaderCompiler;
+
+    const SLANG_SRC_DIR: include_dir::Dir<'_> =
+        include_dir::include_dir!("$CARGO_MANIFEST_DIR/shaders");
 
-    let targets = [
-        CompileTarget::Wgsl,
-        #[cfg(feature = "cuda")]
-        CompileTarget::CudaSource,
-    ];
+    let out_dir = env::var("OUT_DIR").expect("Couldn't determine output directory.");
+    let mut compiler = ShaderCompiler::new(vec![], &out_dir);
+    compiler.add_dir(SLANG_SRC_DIR);
 
-    for target in targets {
-        slang.compile_all(target, "../shaders", "./src/autogen", &[]);
-    }
+    // Compile all shaders from examples/shaders directory.
+    // Note: slang-hal-build will automatically detect which backends to compile for
+    // based on the cargo features enabled during the build.
+    compiler
+        .compile_shaders_dir("shaders", &[])
+        .expect("Failed to compile shaders");
 }

examples/gemm_bench.rs

@@ -1,13 +1,13 @@
 use indexmap::IndexMap;
 use minislang::SlangCompiler;
 use nalgebra::DMatrix;
-use slang_hal::Shader;
 use slang_hal::backend::WebGpu;
 use slang_hal::backend::{Backend, Encoder};
+use slang_hal::{BufferUsages, Shader};
 use stensor::linalg::{Gemm, GemmVariant};
 use stensor::shapes::ViewShapeBuffers;
 use stensor::tensor::GpuTensor;
-use wgpu::{BufferUsages, Features, Limits};
+use wgpu::{Features, Limits};
 
 #[async_std::main]
 async fn main() -> anyhow::Result<()> {

src/lib.rs

@@ -5,8 +5,7 @@
 
 pub use geometry::*;
 pub use linalg::*;
-
-use minislang::SlangCompiler;
+use slang_hal::SlangCompiler;
 
 pub mod geometry;
 pub mod linalg;
@@ -15,7 +14,8 @@ pub mod tensor;
 
 // pub mod utils;
 
-const SLANG_SRC_DIR: include_dir::Dir<'_> =
+/// Directory of slang shaders from `stensor`.
+pub const SLANG_SRC_DIR: include_dir::Dir<'_> =
     include_dir::include_dir!("$CARGO_MANIFEST_DIR/shaders");
 
 /// Register all the shaders from this crate (and its dependencies) as modules accessible to the

src/linalg/contiguous.rs

@@ -62,10 +62,10 @@ mod test {
     use crate::tensor::GpuTensor;
     use minislang::SlangCompiler;
     use nalgebra::DMatrix;
-    use slang_hal::Shader;
     use slang_hal::backend::WebGpu;
     use slang_hal::backend::{Backend, Encoder};
-    use wgpu::{BufferUsages, Features, Limits};
+    use slang_hal::{BufferUsages, Shader};
+    use wgpu::{Features, Limits};
 
     #[futures_test::test]
     #[serial_test::serial]

src/linalg/gemm.rs

@@ -202,9 +202,9 @@ mod test {
     use approx::relative_eq;
     use minislang::SlangCompiler;
     use nalgebra::DMatrix;
-    use slang_hal::Shader;
     use slang_hal::backend::{Backend, Encoder, WebGpu};
-    use wgpu::{BufferUsages, Features, Limits};
+    use slang_hal::{BufferUsages, Shader};
+    use wgpu::{Features, Limits};
 
     #[futures_test::test]
     #[serial_test::serial]

src/linalg/gemv.rs

@@ -313,10 +313,10 @@ mod test {
     use approx::assert_relative_eq;
     use minislang::SlangCompiler;
     use nalgebra::{DMatrix, DVector};
-    use slang_hal::Shader;
     use slang_hal::backend::WebGpu;
     use slang_hal::backend::{Backend, Encoder};
-    use wgpu::{BufferUsages, Features, Limits};
+    use slang_hal::{BufferUsages, Shader};
+    use wgpu::{Features, Limits};
 
     #[futures_test::test]
     #[serial_test::serial]

src/linalg/op_assign.rs

@@ -159,10 +159,10 @@ mod test {
     use crate::tensor::GpuTensor;
     use minislang::SlangCompiler;
     use nalgebra::DVector;
+    use slang_hal::BufferUsages;
     use slang_hal::backend::WebGpu;
     use slang_hal::backend::{Backend, Buffer, Encoder};
     use slang_hal::shader::Shader;
-    use wgpu::BufferUsages;
 
     #[futures_test::test]
     #[serial_test::serial]