feat!: Make ug dependency optional (#3268)

* feat!: Make `ug` dep optional

* fix(example/mnist-training): Run all epochs

* doc(`candle-ug`): Crate documentation

* fix: feature-gate the `ComputePipeline` import

---------

Co-authored-by: ivarflakstad <[email protected]>

Author: DanikVitek

Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "candle-nn",
     "candle-pyo3",
     "candle-transformers",
+    "candle-ug",
     "candle-wasm-examples/*",
     "candle-wasm-tests",
     "tensor-tools",
@@ -43,6 +44,7 @@ candle-metal-kernels = { path = "./candle-metal-kernels", version = "0.9.2-alpha
 candle-nn = { path = "./candle-nn", version = "0.9.2-alpha.2" }
 candle-onnx = { path = "./candle-onnx", version = "0.9.2-alpha.2" }
 candle-transformers = { path = "./candle-transformers", version = "0.9.2-alpha.2" }
+candle-ug = { path = "./candle-ug", version = "0.9.2-alpha.2" }
 clap = { version = "4.2.4", features = ["derive"] }
 criterion = { version = "0.7.0", default-features = false }
 cudarc = { version = "0.18.1", features = [
@@ -65,10 +67,7 @@ half = { version = "2.5.0", features = [
     "use-intrinsics",
     "rand_distr",
 ] }
-float8 = { version = "0.5.0", features = [
-    "num-traits",
-    "rand_distr",
-] }
+float8 = { version = "0.5.0", features = ["num-traits", "rand_distr"] }
 hound = "3.5.1"
 image = { version = "0.25.2", default-features = false, features = [
     "jpeg",

candle-core/Cargo.toml

@@ -35,9 +35,7 @@ yoke = { workspace = true }
 zip = { workspace = true }
 
 [target.'cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))'.dependencies]
-ug = { workspace = true }
-ug-cuda = { workspace = true, optional = true }
-ug-metal = { workspace = true, optional = true }
+candle-ug = { workspace = true, optional = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }
@@ -46,7 +44,7 @@ criterion = { workspace = true }
 
 [features]
 default = []
-cuda = ["cudarc", "dep:candle-kernels", "dep:ug-cuda", "float8/cuda"]
+cuda = ["cudarc", "dep:candle-kernels", "candle-ug?/cuda", "float8/cuda"]
 cudnn = ["cuda", "cudarc/cudnn"]
 nccl = ["cuda", "cudarc/nccl"]
 mkl = ["dep:libc", "dep:intel-mkl-src"]
@@ -55,8 +53,9 @@ metal = [
     "dep:objc2-metal",
     "dep:objc2-foundation",
     "dep:candle-metal-kernels",
-    "dep:ug-metal",
+    "candle-ug?/metal",
 ]
+ug = ["dep:candle-ug"]
 
 [[bench]]
 name = "bench_main"

candle-core/src/cuda_backend/device.rs

@@ -173,14 +173,14 @@ impl CudaDevice {
         self.context.is_event_tracking()
     }
 
-    #[cfg(not(target_arch = "wasm32"))]
+    #[cfg(all(feature = "ug", not(target_arch = "wasm32")))]
     pub fn compile(
         &self,
         func_name: &'static str,
-        kernel: ug::lang::ssa::Kernel,
+        kernel: candle_ug::lang::ssa::Kernel,
     ) -> Result<CudaFunc> {
         let mut buf = vec![];
-        ug_cuda::code_gen::gen(&mut buf, func_name, &kernel)?;
+        candle_ug::cuda::code_gen::gen(&mut buf, func_name, &kernel)?;
         let cuda_code = String::from_utf8(buf)?;
         let opts = cudarc::nvrtc::CompileOptions {
             use_fast_math: Some(true),

candle-core/src/custom_op.rs

@@ -376,6 +376,7 @@ impl Tensor {
     }
 }
 
+#[cfg(feature = "ug")]
 pub struct UgIOp1 {
     name: &'static str,
     #[cfg(feature = "cuda")]
@@ -384,12 +385,13 @@ pub struct UgIOp1 {
     func: candle_metal_kernels::metal::ComputePipeline,
 }
 
+#[cfg(feature = "ug")]
 impl UgIOp1 {
     #[allow(unused)]
     #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
     pub fn new(
         name: &'static str,
-        kernel: ug::lang::ssa::Kernel,
+        kernel: candle_ug::lang::ssa::Kernel,
         device: &crate::Device,
     ) -> Result<Self> {
         #[cfg(feature = "cuda")]
@@ -414,6 +416,7 @@ impl UgIOp1 {
     }
 }
 
+#[cfg(feature = "ug")]
 impl InplaceOp1 for UgIOp1 {
     fn name(&self) -> &'static str {
         self.name

candle-core/src/error.rs

@@ -174,9 +174,9 @@ pub enum Error {
     #[error("Metal error {0}")]
     Metal(#[from] MetalError),
 
-    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
+    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios"), feature = "ug"))]
     #[error(transparent)]
-    Ug(#[from] ug::Error),
+    Ug(#[from] candle_ug::Error),
 
     #[error(transparent)]
     TryFromIntError(#[from] core::num::TryFromIntError),

candle-core/src/lib.rs

@@ -92,7 +92,9 @@ mod variable;
 pub use cuda_backend::cudnn;
 
 pub use cpu_backend::{CpuStorage, CpuStorageRef};
-pub use custom_op::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3, UgIOp1};
+#[cfg(feature = "ug")]
+pub use custom_op::UgIOp1;
+pub use custom_op::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3};
 pub use device::{Device, DeviceLocation, NdArray};
 pub use dtype::{DType, DTypeParseError, FloatDType, IntDType, WithDType};
 pub use dummy_dtype::{F4, F6E2M3, F6E3M2, F8E8M0};

candle-core/src/metal_backend/device.rs

@@ -1,13 +1,17 @@
 use crate::{DType, Result};
+
+#[cfg(feature = "ug")]
+use candle_metal_kernels::metal::ComputePipeline;
 use candle_metal_kernels::{
     metal::{
-        BlitCommandEncoder, Buffer, BufferMap, Commands, ComputeCommandEncoder, ComputePipeline,
-        Device, MTLResourceOptions,
+        BlitCommandEncoder, Buffer, BufferMap, Commands, ComputeCommandEncoder, Device,
+        MTLResourceOptions,
     },
     Kernels,
 };
 use objc2_foundation::NSURL;
 use objc2_metal::{MTLCaptureDescriptor, MTLCaptureDestination, MTLCaptureManager};
+
 use std::path::Path;
 use std::sync::{Arc, Mutex, RwLock};
 
@@ -88,14 +92,14 @@ impl std::ops::Deref for MetalDevice {
 }
 
 impl MetalDevice {
-    #[cfg(all(not(target_arch = "wasm32"), not(target_os = "ios")))]
+    #[cfg(all(feature = "ug", not(target_arch = "wasm32"), not(target_os = "ios")))]
     pub fn compile(
         &self,
         func_name: &'static str,
-        kernel: ug::lang::ssa::Kernel,
+        kernel: candle_ug::lang::ssa::Kernel,
     ) -> Result<ComputePipeline> {
         let mut buf = vec![];
-        ug_metal::code_gen::gen(&mut buf, func_name, &kernel)?;
+        candle_ug::metal::code_gen::gen(&mut buf, func_name, &kernel)?;
         let metal_code = String::from_utf8(buf)?;
         let lib = self
             .device

candle-core/tests/custom_op_tests.rs

@@ -145,20 +145,20 @@ fn inplace_op1() -> Result<()> {
     Ok(())
 }
 
-#[cfg(any(feature = "cuda", feature = "metal"))]
+#[cfg(all(feature = "ug", any(feature = "cuda", feature = "metal")))]
 #[allow(clippy::approx_constant)]
 #[test]
 fn ug_op() -> Result<()> {
     let kernel = {
-        use ug::lang::op;
+        use candle_ug::lang::op;
 
-        let layout = ug::Layout::from_shape(&[12]);
-        let ptr = op::Arg::ptr(ug::DType::F32);
-        let src = op::load(ptr.id(), layout.clone(), ug::DType::F32)?;
+        let layout = candle_ug::Layout::from_shape(&[12]);
+        let ptr = op::Arg::ptr(candle_ug::DType::F32);
+        let src = op::load(ptr.id(), layout.clone(), candle_ug::DType::F32)?;
         let src = op::unary(op::UnaryOp::Exp, src)?;
         let st = op::store(ptr.id(), layout, src)?;
         let kernel = op::Kernel::new("exp".to_string(), vec![ptr], vec![st]);
-        let opts: ug::lower_op::Opts = Default::default();
+        let opts: candle_ug::lower_op::Opts = Default::default();
         kernel.lower(&opts)?
     };
     let device = if candle_core::utils::cuda_is_available() {

candle-examples/examples/mnist-training/main.rs

@@ -137,7 +137,7 @@ fn training_loop_cnn(
     let test_labels = m.test_labels.to_dtype(DType::U32)?.to_device(&dev)?;
     let n_batches = train_images.dim(0)? / BSIZE;
     let mut batch_idxs = (0..n_batches).collect::<Vec<usize>>();
-    for epoch in 1..args.epochs {
+    for epoch in 1..=args.epochs {
         let mut sum_loss = 0f32;
         batch_idxs.shuffle(&mut rng());
         for batch_idx in batch_idxs.iter() {
@@ -194,7 +194,7 @@ fn training_loop<M: Model>(
     let mut sgd = candle_nn::SGD::new(varmap.all_vars(), args.learning_rate)?;
     let test_images = m.test_images.to_device(&dev)?;
     let test_labels = m.test_labels.to_dtype(DType::U32)?.to_device(&dev)?;
-    for epoch in 1..args.epochs {
+    for epoch in 1..=args.epochs {
         let logits = model.forward(&train_images)?;
         let log_sm = ops::log_softmax(&logits, D::Minus1)?;
         let loss = loss::nll(&log_sm, &train_labels)?;

candle-ug/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "candle-ug"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+repository.workspace = true
+keywords.workspace = true
+categories.workspace = true
+license.workspace = true
+
+[dependencies]
+ug = { workspace = true }
+ug-cuda = { workspace = true, optional = true }
+ug-metal = { workspace = true, optional = true }
+
+[features]
+default = []
+cuda = ["dep:ug-cuda"]
+metal = ["dep:ug-metal"]