Add burn-hip

.cargo/config.toml

@@ -1,2 +1,3 @@
 [alias]
+bb = "run --release --bin burnbench --"
 xtask = "run --target-dir target/xtask --color always --package xtask --bin xtask --"

.github/workflows/publish.yml

@@ -126,6 +126,19 @@ jobs:
     secrets:
       CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
 
+  publish-burn-hip:
+    uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v1
+    needs:
+      - publish-burn-tensor
+      - publish-burn-autodiff
+      - publish-burn-ndarray
+      - publish-burn-common
+      - publish-burn-jit
+    with:
+      crate: burn-hip
+    secrets:
+      CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
+
   publish-burn-candle:
     uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v1
     needs:

Cargo.toml

@@ -153,8 +153,8 @@ ahash = { version = "0.8.11", default-features = false }
 portable-atomic-util = { version = "0.2.2", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "e3fdc96ec2d68dcdde8135bd011907b4e662388c" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "e3fdc96ec2d68dcdde8135bd011907b4e662388c" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb2a5c87802c9d01801f16d93cada5fe924989f3" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb2a5c87802c9d01801f16d93cada5fe924989f3" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

backend-comparison/Cargo.toml

@@ -17,6 +17,7 @@ candle-cuda = ["burn/candle-cuda"]
 candle-metal = ["burn/candle", "burn/metal"]
 cuda-jit = ["burn/cuda-jit"]
 cuda-jit-fusion = ["cuda-jit", "burn/fusion"]
+hip-jit = ["burn/hip-jit"]
 default = ["burn/std", "burn/autodiff", "burn/wgpu", "burn/autotune"]
 ndarray = ["burn/ndarray"]
 ndarray-blas-accelerate = ["burn/ndarray", "burn/accelerate"]

backend-comparison/README.md

@@ -6,7 +6,7 @@ to complex models.
 ## burnbench CLI
 
 This crate comes with a CLI binary called `burnbench` which can be executed via
-`cargo run  --release --bin burnbench`.
+`cargo run --release --bin burnbench`.
 
 Note that you need to run the `release` target of `burnbench` otherwise you won't
 be able to share your benchmark results.
@@ -15,6 +15,9 @@ The end of options argument `--` is used to pass arguments to the `burnbench`
 application. For instance `cargo run --bin burnbench -- list` passes the `list`
 argument to `burnbench` effectively calling `burnbench list`.
 
+There is also a cargo alias `cargo bb` which simplifies the command line.
+The example command above then becomes: `cargo bb list`.
+
 ### Commands
 
 #### List benches and backends

backend-comparison/src/burnbenchapp/base.rs

@@ -86,6 +86,8 @@ enum BackendValues {
     CudaJit,
     #[strum(to_string = "cuda-jit-fusion")]
     CudaJitFusion,
+    #[strum(to_string = "hip-jit")]
+    HipJit,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, ValueEnum, Display, EnumIter)]

backend-comparison/src/lib.rs

@@ -65,6 +65,8 @@ macro_rules! bench_on_backend {
         let feature_name = "cuda-jit";
         #[cfg(feature = "cuda-jit-fusion")]
         let feature_name = "cuda-jit-fusion";
+        #[cfg(feature = "hip-jit")]
+        let feature_name = "hip-jit";
 
         #[cfg(any(feature = "wgpu"))]
         {
@@ -146,6 +148,13 @@ macro_rules! bench_on_backend {
 
             bench::<Cuda<half::f16>>(&CudaDevice::default(), feature_name, url, token);
         }
+
+        #[cfg(feature = "hip-jit")]
+        {
+            use burn::backend::hip_jit::{Hip, HipDevice};
+
+            bench::<Hip<half::f16>>(&HipDevice::default(), feature_name, url, token);
+        }
     };
 }
 

crates/burn-core/Cargo.toml

@@ -23,6 +23,7 @@ default = [
     "burn-tensor/default",
     "burn-wgpu?/default",
     "burn-cuda?/default",
+    "burn-hip?/default",
     "burn-autodiff?/default",
 ]
 doc = [
@@ -35,6 +36,7 @@ doc = [
     "tch",
     "wgpu",
     "cuda-jit",
+    "hip-jit",
     "vision",
     "autodiff",
     # Doc features
@@ -46,6 +48,7 @@ doc = [
     "burn-tensor/doc",
     "burn-wgpu/doc",
     "burn-cuda/doc",
+    "burn-hip/doc",
 ]
 network = ["burn-common/network"]
 sqlite = ["burn-dataset?/sqlite"]
@@ -59,6 +62,7 @@ std = [
     "burn-tensor/std",
     "burn-wgpu?/std",
     "burn-cuda?/std",
+    "burn-hip?/std",
     "flate2",
     "half/std",
     "log",
@@ -86,6 +90,7 @@ template = ["burn-wgpu?/template"]
 candle = ["burn-candle"]
 candle-cuda = ["candle", "burn-candle/cuda"]
 cuda-jit = ["burn-cuda"]
+hip-jit = ["burn-hip"]
 ndarray = ["burn-ndarray"]
 tch = ["burn-tch"]
 wgpu = ["burn-wgpu"]
@@ -101,6 +106,7 @@ experimental-named-tensor = ["burn-tensor/experimental-named-tensor"]
 record-backward-compat = []
 
 test-cuda = ["cuda-jit"] # To use cuda during testing, default uses ndarray.
+test-hip = ["hip-jit"] # To use hip during testing, default uses ndarray.
 test-tch = ["tch"] # To use tch during testing, default uses ndarray.
 test-wgpu = ["wgpu"] # To use wgpu during testing, default uses ndarray.
 test-wgpu-spirv = [
@@ -121,6 +127,7 @@ burn-tensor = { path = "../burn-tensor", version = "0.15.0", default-features =
 burn-autodiff = { path = "../burn-autodiff", version = "0.15.0", optional = true }
 burn-candle = { path = "../burn-candle", version = "0.15.0", optional = true }
 burn-cuda = { path = "../burn-cuda", version = "0.15.0", optional = true, default-features = false }
+burn-hip = { path = "../burn-hip", version = "0.15.0", optional = true, default-features = false }
 burn-ndarray = { path = "../burn-ndarray", version = "0.15.0", optional = true, default-features = false }
 burn-tch = { path = "../burn-tch", version = "0.15.0", optional = true }
 burn-wgpu = { path = "../burn-wgpu", version = "0.15.0", optional = true, default-features = false }

crates/burn-core/src/backend.rs

@@ -28,6 +28,9 @@ pub use burn_candle as candle;
 #[cfg(feature = "candle")]
 pub use burn_candle::Candle;
 
+#[cfg(feature = "hip-jit")]
+pub use burn_hip as hip_jit;
+
 #[cfg(feature = "tch")]
 pub use burn_tch as libtorch;
 

crates/burn-hip/Cargo.toml

@@ -0,0 +1,40 @@
+[package]
+authors = ["nathanielsimard <nathaniel.simard.42@gmail.com>"]
+categories = ["science"]
+description = "ROCm HIP backend for the Burn framework"
+edition.workspace = true
+keywords = ["deep-learning", "machine-learning", "gpu", "rocm", "hip"]
+license.workspace = true
+name = "burn-hip"
+readme.workspace = true
+repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-hip"
+documentation = "https://docs.rs/burn-hip"
+version.workspace = true
+
+[features]
+default = ["fusion", "burn-jit/default", "cubecl/default"]
+fusion = ["burn-fusion", "burn-jit/fusion"]
+autotune = ["burn-jit/autotune"]
+doc = ["burn-jit/doc"]
+std = ["burn-jit/std", "cubecl/std"]
+
+[dependencies]
+cubecl = { workspace = true, features = ["hip"] }
+burn-jit = { path = "../burn-jit", version = "0.15.0", default-features = false }
+burn-tensor = { path = "../burn-tensor", version = "0.15.0", features = ["cubecl-hip"] }
+burn-fusion = { path = "../burn-fusion", version = "0.15.0", optional = true }
+
+half = { workspace = true }
+bytemuck = { workspace = true }
+
+log = { workspace = true }
+derive-new = { workspace = true }
+
+[dev-dependencies]
+burn-jit = { path = "../burn-jit", version = "0.15.0", default-features = false, features = [
+  "export_tests",
+] }
+
+[package.metadata.docs.rs]
+features = ["doc"]
+rustdoc-args = ["--cfg", "docsrs"]

crates/burn-hip/README.md

@@ -0,0 +1,7 @@
+# Burn-hip
+
+Backend using ROCm HIP runtime.
+
+To execute the tests for this backend set an environment variable called `ROCM_PATH` or `CUBECL_ROCM_PATH` to the installation path of ROCm. It is often `/opt/rocm`.
+
+For now this backend requires the version `6.2.2` of ROCm or a compatible version.

crates/burn-hip/src/lib.rs

@@ -0,0 +1,22 @@
+#![cfg_attr(docsrs, feature(doc_auto_cfg))]
+
+extern crate alloc;
+
+use burn_jit::JitBackend;
+pub use cubecl::hip::HipDevice;
+use cubecl::hip::HipRuntime;
+
+#[cfg(not(feature = "fusion"))]
+pub type Hip<F = f32, I = i32> = JitBackend<HipRuntime, F, I>;
+
+#[cfg(feature = "fusion")]
+pub type Hip<F = f32, I = i32> = burn_fusion::Fusion<JitBackend<HipRuntime, F, I>>;
+
+#[cfg(test)]
+mod tests {
+    use burn_jit::JitBackend;
+
+    pub type TestRuntime = cubecl::hip::HipRuntime;
+
+    burn_jit::testgen_all!();
+}

crates/burn-tensor/Cargo.toml

@@ -28,6 +28,7 @@ repr = []
 cubecl = ["dep:cubecl"]
 cubecl-wgpu = ["cubecl", "cubecl/wgpu"]
 cubecl-cuda = ["cubecl", "cubecl/cuda"]
+cubecl-hip = ["cubecl", "cubecl/hip"]
 
 [dependencies]
 burn-common = { path = "../burn-common", version = "0.15.0", default-features = false }

crates/burn-tensor/src/lib.rs

@@ -89,3 +89,15 @@ mod cube_cuda {
         }
     }
 }
+
+#[cfg(feature = "cubecl-hip")]
+mod cube_hip {
+    use crate::backend::{DeviceId, DeviceOps};
+    use cubecl::hip::HipDevice;
+
+    impl DeviceOps for HipDevice {
+        fn id(&self) -> DeviceId {
+            DeviceId::new(0, self.index as u32)
+        }
+    }
+}

crates/burn/Cargo.toml

@@ -50,6 +50,7 @@ template = ["burn-core/template"]
 
 candle = ["burn-core/candle"]
 cuda-jit = ["burn-core/cuda-jit"]
+hip-jit = ["burn-core/hip-jit"]
 ndarray = ["burn-core/ndarray"]
 tch = ["burn-core/tch"]
 wgpu = ["burn-core/wgpu"]