diff --git a/Cargo.lock b/Cargo.lock index 03f9dd9..9f0e999 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1148,6 +1148,7 @@ dependencies = [ "anyhow", "approx", "arcstr", + "bindgen 0.72.1", "cfgrammar", "const_format", "derive-where", @@ -1157,9 +1158,13 @@ dependencies = [ "indexmap 2.12.1", "itertools 0.14.0", "klayout-lyp", + "libc", "lrlex", "lrpar", "nalgebra", + "nalgebra-sparse", + "rand 0.8.5", + "rayon", "regex", "rgb", "serde", @@ -2403,9 +2408,9 @@ dependencies = [ [[package]] name = "generic-array" -version = "0.14.7" +version = "0.14.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" dependencies = [ "typenum", "version_check", @@ -4011,6 +4016,7 @@ dependencies = [ "num-complex", "num-rational", "num-traits", + "rayon", "simba", "typenum", ] diff --git a/README.md b/README.md index 193fff6..fa11231 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,18 @@ Future versions of Argon will hopefully support: To use Argon, you will need: - [Rust (tested on 1.90.0)](https://www.rust-lang.org/tools/install) - One of [Neovim (version 0.11.0 or above)](https://github.com/neovim/neovim/blob/master/INSTALL.md) or [VS Code (version 1.100.0 or above)](https://code.visualstudio.com/download) +- [SuiteSparse](https://github.com/DrTimothyAldenDavis/SuiteSparse) +### SuiteSparse Installation +#### macOS +If you use Homebrew, all you need is +```bash +brew install suitesparse +``` +#### Windows/Linux +In this case, you will need to download SuiteSparse manually; you can use something like ``` conda ``` or ```vcpkg```. Then, locate the folder that contains the ```.h``` files, such as ```cholmod.h```, and the folder containing the ```.lib``` files (```.so``` or ```.a``` on Linux). Set the environment variables ```SUITESPARSE_INCLUDE_DIR``` and ```SUITESPARSE_LIB_DIR``` to the two paths you obtained from before, respectively. + +### Argon Installation Begin by cloning and compiling the Argon source code: ```bash @@ -42,7 +53,6 @@ git clone https://github.com/ucb-substrate/argon.git cd argon cargo b --release ``` - ### Neovim Add the following to your Neovim Lua configuration: diff --git a/core/compiler/.gitignore b/core/compiler/.gitignore index d163863..dc84959 100644 --- a/core/compiler/.gitignore +++ b/core/compiler/.gitignore @@ -1 +1,2 @@ -build/ \ No newline at end of file +build/ + diff --git a/core/compiler/Cargo.toml b/core/compiler/Cargo.toml index 1b6574b..5722230 100644 --- a/core/compiler/Cargo.toml +++ b/core/compiler/Cargo.toml @@ -5,8 +5,12 @@ edition = "2024" [dependencies] derive-where = { version = "1", features = ["serde"] } -nalgebra = "0.34" +nalgebra = { version = "0.34", features = ["rayon"] } +nalgebra-sparse = "0.11" klayout-lyp = "0.1.1" +rayon = "1.10" +libc = "0.2" +rand = "0.8.5" gds21 = "0.2" anyhow = { workspace = true } @@ -32,6 +36,7 @@ const_format = "0.2" [build-dependencies] lrpar = { version = "0.13", features = ["serde"] } +bindgen = "0.72.1" cfgrammar = { workspace = true } lrlex = { workspace = true } diff --git a/core/compiler/build.rs b/core/compiler/build.rs index 8a4bf93..3da3b53 100644 --- a/core/compiler/build.rs +++ b/core/compiler/build.rs @@ -1,3 +1,4 @@ +use std::env; use std::path::PathBuf; use cfgrammar::yacc::YaccKind; @@ -30,4 +31,51 @@ fn main() { .mod_name("cell_l") .build() .unwrap(); + + println!("cargo:rerun-if-changed=wrapper.h"); + println!("cargo:rerun-if-changed=build.rs"); + + println!("cargo:rerun-if-env-changed=SUITESPARSE_INCLUDE_DIR"); + println!("cargo:rerun-if-env-changed=SUITESPARSE_LIB_DIR"); + + let include_dir = + env::var("SUITESPARSE_INCLUDE_DIR").unwrap_or_else(|_| "/opt/homebrew/include".to_string()); + let lib_dir = + env::var("SUITESPARSE_LIB_DIR").unwrap_or_else(|_| "/opt/homebrew/lib".to_string()); + + println!("cargo:rustc-link-search=native={}", lib_dir); + println!("cargo:rustc-link-lib=spqr"); + println!("cargo:rustc-link-lib=cholmod"); + println!("cargo:rustc-link-lib=suitesparseconfig"); + println!("cargo:rustc-link-lib=amd"); + println!("cargo:rustc-link-lib=colamd"); + println!("cargo:rustc-link-lib=ccolamd"); + println!("cargo:rustc-link-lib=camd"); + + let bindings = bindgen::Builder::default() + .header("wrapper.h") + .clang_arg(format!("-I{}", include_dir)) + .clang_arg(format!("-I{}/suitesparse", include_dir)) + .wrap_unsafe_ops(true) + .parse_callbacks(Box::new(bindgen::CargoCallbacks::new())) + .allowlist_function("SuiteSparseQR_C_QR") + .allowlist_function("cholmod_l_start") + .allowlist_function("cholmod_l_finish") + .allowlist_function("cholmod_l_free") + .allowlist_function("cholmod_l_free_sparse") + .allowlist_function("cholmod_l_free_dense") + .allowlist_function("cholmod_l_sparse_to_dense") + .allowlist_function("cholmod_l_allocate_triplet") + .allowlist_function("cholmod_l_free_triplet") + .allowlist_function("cholmod_l_triplet_to_sparse") + .allowlist_var("SPQR_ORDERING_DEFAULT") + .allowlist_var("SPQR_DEFAULT_TOL") + .allowlist_var("CHOLMOD_REAL") + .generate() + .expect("Unable to generate bindings"); + + let out_path = PathBuf::from(env::var("OUT_DIR").unwrap()); + bindings + .write_to_file(out_path.join("spqr_bindings.rs")) + .expect("Couldn't write bindings"); } diff --git a/core/compiler/src/SPQR.rs b/core/compiler/src/SPQR.rs new file mode 100644 index 0000000..3a5bdfd --- /dev/null +++ b/core/compiler/src/SPQR.rs @@ -0,0 +1,536 @@ +#![allow(non_upper_case_globals)] +#![allow(non_camel_case_types)] +#![allow(non_snake_case)] +#![allow(dead_code)] +#![allow(improper_ctypes)] + +mod ffi { + include!(concat!(env!("OUT_DIR"), "/spqr_bindings.rs")); +} + +use nalgebra::DMatrix; +use nalgebra::DVector; +use nalgebra_sparse::{CooMatrix, CsrMatrix}; +use rayon::prelude::*; +use std::ptr; +use std::ptr::NonNull; + +///SpqrFactorization struct +///Used for SuiteSparse, +pub struct SpqrFactorization { + q_a: *mut ffi::cholmod_sparse, //Q matrix for AP = QR + q_at: *mut ffi::cholmod_sparse, //Q' matrix for AT P' = Q'R' + r_a: *mut ffi::cholmod_sparse, //R matrix for AP = QR + r_at: *mut ffi::cholmod_sparse, //R' matrix for AT P' = Q'R' + e_a: *mut i64, //permutation vector for AP = QR + e_at: *mut i64, //permutation vector for AT P' = Q'R' + rank: usize, //rank of A (and hence AT) + cc_a: *mut ffi::cholmod_common, //cholmod struct for AP = QR + cc_at: *mut ffi::cholmod_common, //cholmod struct for AT P' = Q'R' + m: usize, //number of rows of A (or columns of AT) + n: usize, //number of columns of A (or rows of AT) +} + +unsafe impl Send for SpqrFactorization {} +unsafe impl Sync for SpqrFactorization {} + +//use a struct as a wrapper for a pointer that's shared across threads +pub struct pointer_wrapper { + pointer: NonNull, +} + +impl pointer_wrapper { + fn as_ptr(&self) -> *mut T { + self.pointer.as_ptr() + } +} + +unsafe impl Send for pointer_wrapper {} +unsafe impl Sync for pointer_wrapper {} + +impl SpqrFactorization { + ///Creates new SpqrFactorization struct + /// takes in as input + /// triplet : a vector of triplets (row_idx, col_idx, value) + /// m : number of rows of the matrix A + /// n : number of columns of the matrix A + pub fn from_triplets( + triplet: &Vec<(usize, usize, f64)>, + m: usize, + n: usize, + ) -> Result { + unsafe { + let mut cc_a = Box::new(std::mem::zeroed::()); + let mut cc_at = Box::new(std::mem::zeroed::()); + + ffi::cholmod_l_start(cc_a.as_mut()); + ffi::cholmod_l_start(cc_at.as_mut()); + + cc_a.nthreads_max = 0; + cc_at.nthreads_max = 0; + + let (A, AT) = + Self::triplet_to_cholmod_sparse(triplet, m, n, cc_a.as_mut(), cc_at.as_mut()) + .unwrap(); + + let mut q_a: *mut ffi::cholmod_sparse = ptr::null_mut(); + let mut q_at: *mut ffi::cholmod_sparse = ptr::null_mut(); + + let mut r_a: *mut ffi::cholmod_sparse = ptr::null_mut(); + let mut r_at: *mut ffi::cholmod_sparse = ptr::null_mut(); + + let mut e_a: *mut i64 = ptr::null_mut(); + let mut e_at: *mut i64 = ptr::null_mut(); + + let a_econ: i64 = 0; + + let rank_a = ffi::SuiteSparseQR_C_QR( + ffi::SPQR_ORDERING_DEFAULT as i32, + ffi::SPQR_DEFAULT_TOL as f64, + a_econ, + A, + &mut q_a, + &mut r_a, + &mut e_a, + cc_a.as_mut(), + ); + + //rank_at = rank_a + + let at_econ: i64 = 1 + (n as i64); + + let _rank_at = ffi::SuiteSparseQR_C_QR( + ffi::SPQR_ORDERING_DEFAULT as i32, + ffi::SPQR_DEFAULT_TOL as f64, + at_econ, //want full version of QR for AT + AT, + &mut q_at, + &mut r_at, + &mut e_at, + cc_at.as_mut(), + ); + + ffi::cholmod_l_free_sparse(&mut (A as *mut _), cc_a.as_mut()); + ffi::cholmod_l_free_sparse(&mut (AT as *mut _), cc_at.as_mut()); + + if rank_a == -1 { + //failed + ffi::cholmod_l_finish(cc_a.as_mut()); + ffi::cholmod_l_finish(cc_at.as_mut()); + return Err("failed".to_string()); + } + + Ok(SpqrFactorization { + q_a: q_a, + q_at: q_at, + r_a: r_a, + r_at: r_at, + e_a: e_a, + e_at: e_at, + rank: rank_a as usize, + cc_a: Box::into_raw(cc_a), + cc_at: Box::into_raw(cc_at), + m, + n, + }) + } + } + + ///returns the Dense R matrix from the QR of A + pub fn ra_matrix(&self) -> Result, String> { + unsafe { self.cholmod_sparse_to_dense(self.r_a, self.cc_a) } + } + + ///Returns the Dense Q matrix from the QR of A + pub fn qa_matrix(&self) -> Result, String> { + unsafe { self.cholmod_sparse_to_dense(self.q_a, self.cc_a) } + } + + ///Returns the Dense R (FULL) matrix from the QR of AT + pub fn rat_matrix(&self) -> Result, String> { + unsafe { self.cholmod_sparse_to_dense(self.r_at, self.cc_at) } + } + + ///Returns the Dense Q (FULL) matrix from the QR of AT + pub fn qat_matrix(&self) -> Result, String> { + unsafe { self.cholmod_sparse_to_dense(self.q_at, self.cc_at) } + } + + ///Returns the csr Q (FULL) matrix from the QR of AT + pub fn qat_csr_matrix(&self) -> Result, String> { + unsafe { self.cholmod_sparse_to_csr(self.q_at) } + } + + ///Converts the cholmod_sparse matrix to a csr matrix + pub unsafe fn cholmod_sparse_to_csr( + &self, + mat: *mut ffi::cholmod_sparse, + ) -> Result, String> { + unsafe { + let sparse = &*mat; + let sparse_m = sparse.nrow; + let sparse_n = sparse.ncol; + + let col_pointer = sparse.p as *mut i64; + let row_pointer = sparse.i as *mut i64; + let val = sparse.x as *mut f64; + + let col_pointer_wrapper = pointer_wrapper { + pointer: NonNull::new(col_pointer).unwrap(), + }; + let row_pointer_wrapper = pointer_wrapper { + pointer: NonNull::new(row_pointer).unwrap(), + }; + let val_pointer_wrapper = pointer_wrapper { + pointer: NonNull::new(val).unwrap(), + }; + + let triplets: Vec<(usize, usize, f64)> = (0..sparse_n) + .into_par_iter() + .flat_map(|j| { + let start = *col_pointer_wrapper.as_ptr().add(j); + let end = *col_pointer_wrapper.as_ptr().add(j + 1); + + let mut curr_column_triplets = + Vec::with_capacity(end as usize - start as usize); + + for index in start..end { + let i = *row_pointer_wrapper.as_ptr().add(index as usize); + let value = *val_pointer_wrapper.as_ptr().add(index as usize); + curr_column_triplets.push((i as usize, j as usize, value)); + } + curr_column_triplets + }) + .collect(); + + let coo = CooMatrix::try_from_triplets_iter(sparse_m, sparse_n, triplets).unwrap(); + + Ok(CsrMatrix::from(&coo)) + } + } + + pub fn get_nspace_sparse(&self) -> Result, String> { + unsafe { + let qt = &*self.q_at; + + let col_pointer = qt.p as *mut i64; + let row_pointer = qt.i as *mut i64; + let vals = qt.x as *mut f64; + + let col_pointer_wrapper = pointer_wrapper { + pointer: NonNull::new(col_pointer).unwrap(), + }; + let row_pointer_wrapper = pointer_wrapper { + pointer: NonNull::new(row_pointer).unwrap(), + }; + let vals_pointer_wrapper = pointer_wrapper { + pointer: NonNull::new(vals).unwrap(), + }; + + if self.rank >= self.n { + return Ok(CsrMatrix::zeros(self.m, 0)); + } + let triplets: Vec<(usize, usize, f64)> = (0..self.n - self.rank) + .into_par_iter() + .flat_map(|j| { + let col_index = self.rank + j; + let start = *col_pointer_wrapper.as_ptr().add(col_index) as usize; + let end = *col_pointer_wrapper.as_ptr().add(col_index + 1) as usize; + let mut local_triplets = Vec::with_capacity(end - start); + for index in start..end { + let i = *row_pointer_wrapper.as_ptr().add(index); + let val = *vals_pointer_wrapper.as_ptr().add(index); + local_triplets.push((i as usize, j as usize, val.abs())); + } + local_triplets + }) + .collect(); + let coo = + CooMatrix::try_from_triplets_iter(self.n, self.n - self.rank, triplets).unwrap(); + Ok(CsrMatrix::from(&coo)) + } + } + + ///Returns the permutation vector from QR of A + pub fn permutation_a(&self) -> Result, String> { + unsafe { + // if e is null, permutation is I + if self.e_a.is_null() { + return Ok((0..self.n).collect()); + } + + let perm_pointer = self.e_a as *const i64; + + let mut perm = Vec::with_capacity(self.n); + for i in 0..self.n { + perm.push(*perm_pointer.add(i) as usize); + } + Ok(perm) + } + } + ///Returns the permutation vector from QR of AT + pub fn permutation_at(&self) -> Result, String> { + unsafe { + // if e is null, permutation is I + if self.e_at.is_null() { + return Ok((0..self.m).collect()); + } + + let perm_pointer = self.e_at as *const i64; + + let mut perm = Vec::with_capacity(self.m); + for i in 0..self.m { + perm.push(*perm_pointer.add(i) as usize); + } + Ok(perm) + } + } + ///Returns the rank of the matrix A/AT obtained from QR + ///not always the actual rank, see Kahan matrices + pub fn rank(&self) -> usize { + self.rank + } + + ///Converts a vector of triplets to two cholmod sparse matrices A and AT. + ///takes as input + ///triplet : vector of (row_idx, col_idx, value) + ///m : number of rows in A + ///n : number of columns in A + pub unsafe fn triplet_to_cholmod_sparse( + triplet: &Vec<(usize, usize, f64)>, + m: usize, + n: usize, + cc_a: *mut ffi::cholmod_common, + cc_at: *mut ffi::cholmod_common, + ) -> Result<(*mut ffi::cholmod_sparse, *mut ffi::cholmod_sparse), String> { + unsafe { + let nnz = triplet.len(); + + let cholmod_triplet_a = + ffi::cholmod_l_allocate_triplet(m, n, nnz, 0, ffi::CHOLMOD_REAL as i32, cc_a); + let cholmod_triplet_at = + ffi::cholmod_l_allocate_triplet(n, m, nnz, 0, ffi::CHOLMOD_REAL as i32, cc_at); + + let cholmod_triplet_a_ref = &mut *cholmod_triplet_a; + let cholmod_triplet_at_ref = &mut *cholmod_triplet_at; + + let j_a = cholmod_triplet_a_ref.j as *mut i64; + let i_a = cholmod_triplet_a_ref.i as *mut i64; + let x_a = cholmod_triplet_a_ref.x as *mut f64; + + let j_at = cholmod_triplet_at_ref.j as *mut i64; + let i_at = cholmod_triplet_at_ref.i as *mut i64; + let x_at = cholmod_triplet_at_ref.x as *mut f64; + + let j_a_wrapper = pointer_wrapper { + pointer: NonNull::new(j_a).unwrap(), + }; + let i_a_wrapper = pointer_wrapper { + pointer: NonNull::new(i_a).unwrap(), + }; + let x_a_wrapper = pointer_wrapper { + pointer: NonNull::new(x_a).unwrap(), + }; + + let j_at_wrapper = pointer_wrapper { + pointer: NonNull::new(j_at).unwrap(), + }; + let i_at_wrapper = pointer_wrapper { + pointer: NonNull::new(i_at).unwrap(), + }; + let x_at_wrapper = pointer_wrapper { + pointer: NonNull::new(x_at).unwrap(), + }; + + triplet + .par_iter() + .enumerate() + .for_each(|(idx, (i, j, val))| { + let i_a_pointer = i_a_wrapper.as_ptr(); + let j_a_pointer = j_a_wrapper.as_ptr(); + let x_a_pointer = x_a_wrapper.as_ptr(); + + let i_at_pointer = i_at_wrapper.as_ptr(); + let j_at_pointer = j_at_wrapper.as_ptr(); + let x_at_pointer = x_at_wrapper.as_ptr(); + + *i_a_pointer.add(idx) = *i as i64; + *j_a_pointer.add(idx) = *j as i64; + *x_a_pointer.add(idx) = *val; + + //for at, swap (i, j) -> (j, i) + *i_at_pointer.add(idx) = *j as i64; + *j_at_pointer.add(idx) = *i as i64; + *x_at_pointer.add(idx) = *val; + }); + + cholmod_triplet_a_ref.nnz = nnz; + cholmod_triplet_at_ref.nnz = nnz; + + let a_sparse = ffi::cholmod_l_triplet_to_sparse(cholmod_triplet_a, nnz, cc_a); + let at_sparse = ffi::cholmod_l_triplet_to_sparse(cholmod_triplet_at, nnz, cc_at); + + ffi::cholmod_l_free_triplet(&mut (cholmod_triplet_a as *mut _), cc_a); + ffi::cholmod_l_free_triplet(&mut (cholmod_triplet_at as *mut _), cc_at); + + Ok((a_sparse, at_sparse)) + } + } + + ///Takes in as input sparse matrix and cholmod struct, returns the dense DMatrix version of the data + unsafe fn cholmod_sparse_to_dense( + &self, + sparse: *const ffi::cholmod_sparse, + cc: *mut ffi::cholmod_common, + ) -> Result, String> { + unsafe { + let dense = ffi::cholmod_l_sparse_to_dense(sparse as *mut _, &mut *cc); + + let result = self.cholmod_dense_to_nalgebra(dense).unwrap(); + ffi::cholmod_l_free_dense(&mut (dense as *mut _), &mut *cc); + + Ok(result) + } + } + + ///Takes in as input a cholmod dense and converts it into a dense DMatrix + unsafe fn cholmod_dense_to_nalgebra( + &self, + dense: *const ffi::cholmod_dense, + ) -> Result, String> { + unsafe { + let dense_ref = &*dense; + let m = dense_ref.nrow; + let n = dense_ref.ncol; + let data_pointer = dense_ref.x as *mut f64; + let acc_data_pointer = pointer_wrapper { + pointer: NonNull::new(data_pointer).unwrap(), + }; + + let mut matrix = DMatrix::zeros(m, n); + + matrix + .par_column_iter_mut() + .enumerate() + .for_each(|(j, mut col_slice)| { + let col_pointer = acc_data_pointer.as_ptr().add(j * m); + for i in 0..m { + col_slice[i] = *col_pointer.add(i); + } + }); + + Ok(matrix) + } + } + + ///Complete solve function; determines what path of action to take depending on dimensions of matrix A + pub fn solve(&self, b: &DVector) -> Result, String> { + if self.m >= self.n { + return self.solve_regular(b); + } else { + return self.solve_underconstrained(b); + } + } + + ///Solves the system and returns the least squares solution in the case where the matrix has m >= n + pub fn solve_regular(&self, b: &DVector) -> Result, String> { + let q = self.qa_matrix().unwrap(); + let r = self.ra_matrix().unwrap(); + let perm_vec = self.permutation_a().unwrap(); + + let c = q.transpose() * b; + let mut y = DVector::zeros(self.n); + + let r_acc = r.columns(0, self.rank); + + match r_acc.solve_upper_triangular(&c) { + Some(y_main) => { + y.rows_mut(0, self.rank).copy_from(&y_main); + } + None => return Err("failed R solving".to_string()), + } + + let mut x = DVector::zeros(self.n); + + for i in 0..self.n { + x[perm_vec[i]] = y[i]; + } + + Ok(x) + } + + ///Solves the system for the underconstrained case when m < n; uses the precomputed QR of AT + pub fn solve_underconstrained(&self, b: &DVector) -> Result, String> { + let rank = self.rank(); + let q = self.qat_matrix().unwrap(); + let q_thin = q.columns(0, rank); + let r = self.rat_matrix().unwrap(); + let perm_vec = self.permutation_at().unwrap(); + + let mut c = DVector::zeros(self.m); + for i in 0..self.m { + c[i] = b[perm_vec[i]]; + } + + let r_acc = r.view((0, 0), (rank, rank)); + let c_main = c.rows(0, rank); + + let y = r_acc.transpose().solve_lower_triangular(&c_main).unwrap(); + + let x = q_thin * y; + + Ok(x) + } + + ///Returns the nullspace vectors of A. Uses the last n - r rows of Q from AT + pub fn get_nullspace(&self) -> Result, String> { + let q = &self.qat_matrix().unwrap(); + let null_space_vectors = q.columns(self.rank, self.n - self.rank).clone_owned(); + return Ok(null_space_vectors); + } +} + +impl Drop for SpqrFactorization { + fn drop(&mut self) { + unsafe { + if !self.q_a.is_null() { + ffi::cholmod_l_free_sparse(&mut self.q_a, &mut *self.cc_a); + } + if !self.q_at.is_null() { + ffi::cholmod_l_free_sparse(&mut self.q_at, &mut *self.cc_at); + } + if !self.r_a.is_null() { + ffi::cholmod_l_free_sparse(&mut self.r_a, &mut *self.cc_a); + } + if !self.r_at.is_null() { + ffi::cholmod_l_free_sparse(&mut self.r_at, &mut *self.cc_at); + } + + if !self.e_a.is_null() { + ffi::cholmod_l_free( + self.n, + std::mem::size_of::(), + self.e_a as *mut _, + &mut *self.cc_a, + ); + } + if !self.e_at.is_null() { + ffi::cholmod_l_free( + self.m, + std::mem::size_of::(), + self.e_at as *mut _, + &mut *self.cc_at, + ); + } + + if !self.cc_a.is_null() { + ffi::cholmod_l_finish(&mut *self.cc_a); + drop(Box::from_raw(self.cc_a)); + } + if !self.cc_at.is_null() { + ffi::cholmod_l_finish(&mut *self.cc_at); + drop(Box::from_raw(self.cc_at)); + } + } + } +} diff --git a/core/compiler/src/lib.rs b/core/compiler/src/lib.rs index de7c4c6..5fc0783 100644 --- a/core/compiler/src/lib.rs +++ b/core/compiler/src/lib.rs @@ -5,10 +5,12 @@ pub mod gds; pub mod layer; pub mod parse; pub mod solver; +pub mod spqr; #[cfg(test)] mod tests { + use gds21::GdsUnits; use std::path::PathBuf; use crate::{ @@ -19,7 +21,7 @@ mod tests { use approx::assert_relative_eq; use const_format::concatcp; - use crate::compile::{compile, CellArg, CompileInput}; + use crate::compile::{CellArg, CompileInput, compile}; const EPSILON: f64 = 1e-10; const EXAMPLES_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../../examples"); @@ -439,6 +441,7 @@ mod tests { cells .to_gds( GdsMap::from_lyp(SKY130_LYP).expect("failed to create GDS map"), + GdsUnits::new(1e-3, 1e-9), work_dir.join("layout.gds"), ) .expect("Failed to write to GDS"); @@ -681,6 +684,7 @@ mod tests { cells .to_gds( GdsMap::from_lyp(SKY130_LYP).expect("failed to create GDS map"), + GdsUnits::new(1e-3, 1e-9), work_dir.join("layout.gds"), ) .expect("Failed to write to GDS"); diff --git a/core/compiler/src/solver.rs b/core/compiler/src/solver.rs index 3ba5d1b..8df0d1b 100644 --- a/core/compiler/src/solver.rs +++ b/core/compiler/src/solver.rs @@ -1,7 +1,11 @@ -use approx::relative_eq; +use std::collections::HashMap; + use indexmap::{IndexMap, IndexSet}; use itertools::{Either, Itertools}; -use nalgebra::{DMatrix, DVector}; +use nalgebra::DVector; +use nalgebra_sparse::{CooMatrix, CsrMatrix}; +use rayon::iter::IntoParallelRefIterator; +use rayon::prelude::*; use serde::{Deserialize, Serialize}; const EPSILON: f64 = 1e-8; @@ -10,6 +14,7 @@ const INV_ROUND_STEP: f64 = 1. / ROUND_STEP; #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, Ord, PartialOrd)] pub struct Var(u64); +use crate::spqr::SpqrFactorization; #[derive(Clone, Default)] pub struct Solver { @@ -91,36 +96,103 @@ impl Solver { if n_vars == 0 || self.constraints.is_empty() { return; } - let a = DMatrix::from_row_iterator( - self.constraints.len(), - n_vars, - self.constraints - .iter() - .flat_map(|c| c.expr.coeff_vec(n_vars)), - ); - let b = DVector::from_iterator( - self.constraints.len(), - self.constraints.iter().map(|c| -c.expr.constant), - ); - - let svd = a.clone().svd(true, true); - let vt = svd.v_t.as_ref().expect("No V^T matrix"); - let r = svd.rank(EPSILON); - if r == 0 { - return; + + let old_triplets: Vec<(usize, usize, f64)> = self + .constraints + .par_iter() + .enumerate() + .flat_map(|(c_index, c)| { + c.expr + .coeff_vec(n_vars) + .into_par_iter() + .enumerate() + .filter_map(move |(v_index, v)| { + if v != 0.0 { + Some((c_index, v_index, v)) + } else { + None + } + }) + }) + .collect(); + + let mut used = vec![false; n_vars]; + for (_, v_index, _) in &old_triplets { + used[*v_index] = true; } - let vt_recons = vt.rows(0, r); - let sol = svd.solve(&b, EPSILON).unwrap(); - for i in 0..self.next_var { - let recons = (vt_recons.transpose() * vt_recons.column(i as usize))[((i as usize), 0)]; - if !self.solved_vars.contains_key(&Var(i)) - && relative_eq!(recons, 1., epsilon = EPSILON) - { - let val = round(sol[(i as usize, 0)]); - self.solved_vars.insert(Var(i), val); + let mut var_map = vec![usize::MAX; n_vars]; + let mut rev_var_map = Vec::with_capacity(n_vars); + let mut new_index = 0; + + for (old_index, &is_used) in used.iter().enumerate() { + if is_used { + var_map[old_index] = new_index; + rev_var_map.push(old_index); + new_index += 1; + } + } + + let n = new_index; + let m = self.constraints.len(); + + let triplets: Vec<(usize, usize, f64)> = old_triplets + .into_par_iter() + .map(|(c_index, v_index, val)| { + let new_index = var_map[v_index]; + (c_index, new_index, val) + }) + .collect(); + + let temp_b: Vec = self + .constraints + .par_iter() + .map(|c| -c.expr.constant) + .collect(); + + let b = DVector::from_iterator(self.constraints.len(), temp_b); + + let temp_a_constraind_ids: Vec = self.constraints.par_iter().map(|c| c.id).collect(); + let a_constraint_ids = Vec::from_iter(temp_a_constraind_ids); + + let mut a_coo = CooMatrix::new(m, n); + for (i, j, v) in triplets.iter() { + a_coo.push(*i, *j, *v); + } + let a_sparse: CsrMatrix = CsrMatrix::from(&a_coo); + + let qr = SpqrFactorization::from_triplets(&triplets, m, n).unwrap(); + + let rank = qr.rank(); + + let x = qr.solve(&b).unwrap(); + + let residual = &b - &a_sparse * &x; + + let tolerance = 1e-10; + + for i in 0..residual.nrows() { + let r = residual[(i, 0)]; + if r.abs() > tolerance { + self.inconsistent_constraints.insert(a_constraint_ids[i]); } } + + let ones_vector: DVector = DVector::from_element(n - rank, 1.0); + let null_space_components = qr.get_nspace_sparse().unwrap() * ones_vector; + + let par_solved_vars: HashMap = (0..n) + .into_par_iter() + .filter(|&i| null_space_components[i] < tolerance) + .map(|i| { + let actual_val = x[(i, 0)]; + let actual_var = rev_var_map[i]; + (Var(actual_var as u64), round(actual_val)) + }) + .collect(); + + self.solved_vars.extend(par_solved_vars); + for constraint in self.constraints.iter_mut() { substitute_expr(&self.solved_vars, &mut constraint.expr); if constraint.expr.coeffs.is_empty() @@ -300,4 +372,129 @@ mod tests { assert_relative_eq!(*solver.solved_vars.get(&y).unwrap(), 5., epsilon = EPSILON); assert!(!solver.solved_vars.contains_key(&z)); } + + #[test] + fn linear_constraints_solved_correctly_two() { + let mut solver = Solver::new(); + let x = solver.new_var(); + let y = solver.new_var(); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(1., x)], + constant: -3., + }); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(1., y)], + constant: -5., + }); + solver.solve(); + assert_relative_eq!(*solver.solved_vars.get(&x).unwrap(), 3., epsilon = EPSILON); + assert_relative_eq!(*solver.solved_vars.get(&y).unwrap(), 5., epsilon = EPSILON); + } + + #[test] + fn linear_constraints_solved_correctly_three() { + let mut solver = Solver::new(); + let x = solver.new_var(); + let y = solver.new_var(); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(1., x), (1., y)], + constant: -3., + }); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(1., y)], + constant: -5., + }); + solver.solve(); + assert_relative_eq!(*solver.solved_vars.get(&x).unwrap(), -2., epsilon = EPSILON); + assert_relative_eq!(*solver.solved_vars.get(&y).unwrap(), 5., epsilon = EPSILON); + } + + #[test] + fn linear_constraints_solved_correctly_four() { + let mut solver = Solver::new(); + let x = solver.new_var(); + let y = solver.new_var(); + let z = solver.new_var(); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(2., x), (1., y), (1., z)], + constant: -3., + }); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(1., x), (2., y), (1., z)], + constant: -5., + }); + solver.constrain_eq0(LinearExpr { + coeffs: vec![(1., x), (1., y), (2., z)], + constant: -8., + }); + solver.solve(); + assert_relative_eq!(*solver.solved_vars.get(&x).unwrap(), -1., epsilon = EPSILON); + assert_relative_eq!(*solver.solved_vars.get(&y).unwrap(), 1., epsilon = EPSILON); + assert_relative_eq!(*solver.solved_vars.get(&z).unwrap(), 4., epsilon = EPSILON); + } + + //big matrix + #[test] + fn big_num_stab_test() { + //make a graph laplacian for grid, want to see if QR = A + fn generate_bad_value() -> f64 { + use rand::Rng; + use rand::seq::SliceRandom; + use rand::thread_rng; + let mut rng = thread_rng(); + let exponents = [-8.0, 0.0, 8.0]; + let exp = exponents.choose(&mut rng).unwrap(); + let magnitude = 10.0_f64.powf(*exp); + let sign = if rng.gen_bool(0.5) { 1.0 } else { -1.0 }; + let noise = rng.gen_range(0.5..1.5); + magnitude * sign * noise + } + let size = 25; + + use nalgebra::DMatrix; + let mut a_dense = DMatrix::zeros(size * size, size * size); + let mut triplets: Vec<(usize, usize, f64)> = Vec::new(); + + for i in 0..(size) { + for j in 0..(size) { + let curr_node = i * size + j; + let mut bad_val = generate_bad_value(); + a_dense[(curr_node, curr_node)] = bad_val; + triplets.push((curr_node, curr_node, bad_val)); + if (j + 1) < size { + bad_val = generate_bad_value(); + a_dense[(curr_node, curr_node + 1)] = bad_val; + triplets.push((curr_node, curr_node + 1, bad_val)); + bad_val = generate_bad_value(); + a_dense[(curr_node + 1, curr_node)] = bad_val; + triplets.push((curr_node + 1, curr_node, bad_val)); + } + if (i + 1) < size { + let next_row_node = size * (i + 1) + j; + bad_val = generate_bad_value(); + a_dense[(curr_node, next_row_node)] = bad_val; + triplets.push((curr_node, next_row_node, bad_val)); + bad_val = generate_bad_value(); + a_dense[(next_row_node, curr_node)] = bad_val; + triplets.push((next_row_node, curr_node, bad_val)); + } + } + } + let qr = SpqrFactorization::from_triplets(&triplets, size * size, size * size).unwrap(); + let q = qr.qa_matrix().unwrap(); + let r = qr.ra_matrix().unwrap(); + + let p_indices = qr.permutation_a().unwrap(); + + let mut ap_dense = DMatrix::zeros(size * size, size * size); + for (new_col_idx, &old_col_idx) in p_indices.iter().enumerate() { + let col = a_dense.column(old_col_idx); + ap_dense.set_column(new_col_idx, &col); + } + let a_norm = ap_dense.norm(); + let resid = ap_dense - q * r; + let err = resid.norm(); + let relative_err = err / a_norm; + assert!(relative_err < 1e-12, "not num stab"); + } } diff --git a/core/compiler/wrapper.h b/core/compiler/wrapper.h new file mode 100644 index 0000000..1aa1c1c --- /dev/null +++ b/core/compiler/wrapper.h @@ -0,0 +1,3 @@ +#include +#include +#include diff --git a/examples/argon_workspace/lib.ar b/examples/argon_workspace/lib.ar index 09cd30b..b9699f2 100644 --- a/examples/argon_workspace/lib.ar +++ b/examples/argon_workspace/lib.ar @@ -2,5 +2,5 @@ mod utils; mod nested; cell test() { - rect("met1", x0=0., y0=0., x1=nested::nested::test(), y1=crate::utils::test()); + rect("met1", x0=0., y0=0., x1=#scope0 nested::nested::test(), y1=#scope1 crate::utils::test()); }