From 3f4fd89fdeb68a031c03da5da53332e639711a37 Mon Sep 17 00:00:00 2001 From: jonlachmann Date: Tue, 5 Mar 2024 20:03:12 +0100 Subject: [PATCH] Get the symbolic union of two matrices as a separate operation to speed up the addition of them. --- NAMESPACE | 1 + R/extendr-wrappers.R | 6 +- man/beta_draw.Rd | 2 +- src/Makevars | 1 + src/rust/Cargo.lock | 205 ++++++++++++++++++++++++++++++++++++++++++- src/rust/Cargo.toml | 3 +- src/rust/src/lib.rs | 35 +++++++- 7 files changed, 245 insertions(+), 8 deletions(-) diff --git a/NAMESPACE b/NAMESPACE index a5cc999..b0a81cf 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -2,4 +2,5 @@ export(beta_draw) export(beta_symbolic) +export(symbolic_union) useDynLib(sparse.llt, .registration = TRUE) diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R index 3f8a317..c80a7d9 100644 --- a/R/extendr-wrappers.R +++ b/R/extendr-wrappers.R @@ -17,7 +17,11 @@ beta_symbolic <- function(invomega0, x_invsigma_x) .Call(wrap__beta_symbolic, in #' Draw beta from N(omegabar * x_invsigma_y, omegabar) #' @export -beta_draw <- function(invomega0, x_invsigma_x, x_invsigma_y, random_vec, symb) .Call(wrap__beta_draw, invomega0, x_invsigma_x, x_invsigma_y, random_vec, symb) +beta_draw <- function(invomega0, x_invsigma_x, x_invsigma_y, random_vec, symb, symb2) .Call(wrap__beta_draw, invomega0, x_invsigma_x, x_invsigma_y, random_vec, symb, symb2) + +#' Calculate the symbolic union of two sparse matrices to speed up future binary operations between the two +#' @export +symbolic_union <- function(lhs, rhs) .Call(wrap__symbolic_union, lhs, rhs) # nolint end diff --git a/man/beta_draw.Rd b/man/beta_draw.Rd index afb1bda..6e7f97c 100644 --- a/man/beta_draw.Rd +++ b/man/beta_draw.Rd @@ -4,7 +4,7 @@ \alias{beta_draw} \title{Draw beta from N(omegabar * x_invsigma_y, omegabar)} \usage{ -beta_draw(invomega0, x_invsigma_x, x_invsigma_y, random_vec, symb) +beta_draw(invomega0, x_invsigma_x, x_invsigma_y, random_vec, symb, symb2) } \description{ Draw beta from N(omegabar * x_invsigma_y, omegabar) diff --git a/src/Makevars b/src/Makevars index 59f80b8..897cfcf 100644 --- a/src/Makevars +++ b/src/Makevars @@ -17,6 +17,7 @@ $(STATLIB): export CARGO_HOME=$(CARGOTMP); \ fi && \ export PATH="$(PATH):$(HOME)/.cargo/bin" && \ + export RUSTFLAGS='-C target-cpu=native' && \ cargo build --lib --release --manifest-path=./rust/Cargo.toml --target-dir $(TARGET_DIR) if [ "$(NOT_CRAN)" != "true" ]; then \ rm -Rf $(CARGOTMP) && \ diff --git a/src/rust/Cargo.lock b/src/rust/Cargo.lock index 29c88ff..be5d273 100644 --- a/src/rust/Cargo.lock +++ b/src/rust/Cargo.lock @@ -20,6 +20,15 @@ version = "2.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bytemuck" version = "1.14.3" @@ -58,6 +67,15 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e8f1e641542c07631228b1e0dc04b69ae3c1d58ef65d5691a439711d805c698" +[[package]] +name = "cpufeatures" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53fe5e26ff1b7aef8bca9c6080520cfb8d9333c7568e1829cef191a9723e5504" +dependencies = [ + "libc", +] + [[package]] name = "crossbeam-deque" version = "0.8.5" @@ -89,12 +107,32 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "dbgf" version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6ca96b45ca70b8045e0462f191bd209fcb3c3bfe8dbfb1257ada54c4dd59169" +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "dyn-stack" version = "0.10.0" @@ -169,7 +207,6 @@ dependencies = [ [[package]] name = "faer" version = "0.18.0" -source = "git+https://github.com/sarah-ek/faer-rs.git#854e16484ca7f2f8b30c83555472ac90914fec3d" dependencies = [ "bytemuck", "coe-rs", @@ -178,11 +215,15 @@ dependencies = [ "equator", "faer-entity", "gemm", + "libm", "matrixcompare", "matrixcompare-core", + "npyz", "num-complex", "num-traits", "paste", + "rand", + "rand_distr", "rayon", "reborrow", "serde", @@ -190,8 +231,7 @@ dependencies = [ [[package]] name = "faer-entity" -version = "0.17.0" -source = "git+https://github.com/sarah-ek/faer-rs.git#854e16484ca7f2f8b30c83555472ac90914fec3d" +version = "0.18.0" dependencies = [ "bytemuck", "coe-rs", @@ -320,6 +360,16 @@ dependencies = [ "seq-macro", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "half" version = "2.4.0" @@ -372,6 +422,34 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b0bdabb30db18805d5290b3da7ceaccbddba795620b86c02145d688e04900a73" +[[package]] +name = "memchr" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "523dc4f511e55ab87b694dc30d0f820d60906ef06413f93d4d7a1385599cc149" + +[[package]] +name = "npyz" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13f27ea175875c472b3df61ece89a6d6ef4e0627f43704e400c782f174681ebd" +dependencies = [ + "byteorder", + "num-bigint", + "py_literal", +] + +[[package]] +name = "num-bigint" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "608e7659b5c3d7cba262d894801b9ec9d00de989e8a82bd4bef91d08da45cdc0" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.5" @@ -382,6 +460,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.18" @@ -404,6 +491,51 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "de3145af08024dea9fa9914f381a17b8fc6034dfb00f3a84013f7ff43f29ed4c" +[[package]] +name = "pest" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f8023d0fb78c8e03784ea1c7f3fa36e68a723138990b8d5a47d916b651e7a8" +dependencies = [ + "memchr", + "thiserror", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0d24f72393fd16ab6ac5738bc33cdb6a9aa73f8b902e8fe29cf4e67d7dd1026" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdc17e2a6c7d0a492f0158d7a4bd66cc17280308bbaff78d5bef566dca35ab80" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pest_meta" +version = "2.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "934cd7631c050f4674352a6e835d5f6711ffbfb9345c2fc0107155ac495ae293" +dependencies = [ + "once_cell", + "pest", + "sha2", +] + [[package]] name = "proc-macro2" version = "1.0.78" @@ -425,6 +557,19 @@ dependencies = [ "reborrow", ] +[[package]] +name = "py_literal" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1" +dependencies = [ + "num-bigint", + "num-complex", + "num-traits", + "pest", + "pest_derive", +] + [[package]] name = "quote" version = "1.0.35" @@ -434,6 +579,31 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" + +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand", +] + [[package]] name = "raw-cpuid" version = "10.7.0" @@ -504,6 +674,17 @@ dependencies = [ "syn", ] +[[package]] +name = "sha2" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sparse_llt" version = "0.1.0" @@ -558,12 +739,30 @@ dependencies = [ "syn", ] +[[package]] +name = "typenum" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" + +[[package]] +name = "ucd-trie" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9" + [[package]] name = "unicode-ident" version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "version_check" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" + [[package]] name = "walkdir" version = "2.5.0" diff --git a/src/rust/Cargo.toml b/src/rust/Cargo.toml index 6487d22..513fb1a 100644 --- a/src/rust/Cargo.toml +++ b/src/rust/Cargo.toml @@ -9,5 +9,6 @@ name = 'sparse_llt' [dependencies] extendr-api = '*' -faer = { git = "https://github.com/sarah-ek/faer-rs.git" } +#faer = { git = "https://github.com/sarah-ek/faer-rs.git" } +faer = { path = "/Users/jonlachmann/RustroverProjects/faer-rs" } dyn-stack = "0.10.0" diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs index 38b6f6a..91d5c1f 100644 --- a/src/rust/src/lib.rs +++ b/src/rust/src/lib.rs @@ -29,6 +29,22 @@ fn load_dgc_matrix<'a>(mat: S4, x: &'a Robj) -> faer::sparse::SparseColMatRef<'a return a_mat; } +/// Calculate the symbolic union of two sparse matrices to speed up future binary operations between the two +/// @export +#[extendr] +fn symbolic_union(lhs: S4, rhs: S4) -> Robj { + let lhs_x = lhs.get_slot("x").unwrap(); + let lhs_mat = load_dgc_matrix(lhs, &lhs_x); + + let rhs_x = rhs.get_slot("x").unwrap(); + let rhs_mat = load_dgc_matrix(rhs, &rhs_x); + + let symbolic = faer::sparse::ops::union_symbolic(lhs_mat.symbolic(), rhs_mat.symbolic()).unwrap(); + let externalptr_symb = ExternalPtr::new(symbolic); + + return externalptr_symb.into() +} + /// Perform a symbolic decomposition of invomega0 + x_invsigma_x and return a pointer to it /// @export #[extendr] @@ -67,11 +83,15 @@ fn beta_symbolic(invomega0: S4, x_invsigma_x: S4) -> Robj { /// Draw beta from N(omegabar * x_invsigma_y, omegabar) /// @export #[extendr] -unsafe fn beta_draw(invomega0: S4, x_invsigma_x: S4, x_invsigma_y: &[f64], random_vec: &[f64], symb: Robj) -> Vec { +unsafe fn beta_draw(invomega0: S4, x_invsigma_x: S4, x_invsigma_y: &[f64], random_vec: &[f64], symb: Robj, symb2: Robj) -> Vec { // Get the symbolic decomposition from the pointer let symbolic_ptr: ExternalPtr> = symb.try_into().unwrap(); let symbolic = symbolic_ptr.addr(); + // Get the symbolic target matrix from the pointer + let symbolic2_ptr: ExternalPtr> = symb2.try_into().unwrap(); + let symbolic2 = symbolic2_ptr.addr(); + // Map invomega0 and x_invsigma_x let invomega0_x = invomega0.get_slot("x").unwrap(); let invomega0_mat = load_dgc_matrix(invomega0, &invomega0_x); @@ -79,7 +99,17 @@ unsafe fn beta_draw(invomega0: S4, x_invsigma_x: S4, x_invsigma_y: &[f64], rando let x_invsigma_x_x = x_invsigma_x.get_slot("x").unwrap(); let x_invsigma_x_mat = load_dgc_matrix(x_invsigma_x, &x_invsigma_x_x); - let invomegabar = invomega0_mat + x_invsigma_x_mat; + let x = &mut *vec![0f64; symbolic2.row_indices().len()]; + + let mut invomegabar = faer::sparse::SparseColMatMut::::new({ + symbolic2.as_ref() + }, + x + ); + + faer::sparse::ops::add_into(invomegabar.rb_mut(), invomega0_mat, x_invsigma_x_mat); + + //let invomegabar = invomega0_mat + x_invsigma_x_mat; // Copy x_invsigma_y into owned memory and map it to a matrix let mut x_invsigma_y_bind = x_invsigma_y.to_owned(); @@ -137,4 +167,5 @@ extendr_module! { mod sparse_llt; fn beta_symbolic; fn beta_draw; + fn symbolic_union; } -- GitLab