diff --git a/CHANGELOG.md b/CHANGELOG.md index b8115a175..c182a2010 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,18 @@ ### Added +- GPGPU acceleration of the negative-log-likelihood and gradient hot path inside + the Rust core via [wgpu](https://github.com/gfx-rs/wgpu) (MIT/Apache-2.0), + exposed as a device sub-option of the Rust backend rather than a separate + compute-backend axis. Select with `FitConfig(backend="rust", rust_device=...)` + or `fast-mlsirm fit --backend rust --rust-device {auto,cpu,gpu}`; the GPU path + falls back to the identical CPU implementation at runtime when no GPU adapter + is available. Added requested-device provenance on `FitResult.rust_device` + and in `fit_summary.json`, plus numerical-parity tests asserting the Rust + device paths match the NumPy reference. +- Added `docs/papers/README.md` with a citation and canonical link for Wu et al. + (2021, arXiv:2108.11579), grounding fast, accelerator-friendly IRT estimation + without vendoring the PDF into the repository. - Added `tests/test_rust_parity.py`, a Rust<->NumPy numerical parity gate that asserts agreement to `1e-6` across all five model variants, multiple problem sizes, and masked/dense fixtures (observed difference ~1e-13). diff --git a/Cargo.lock b/Cargo.lock index 270913776..cf5b084fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,1212 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + +[[package]] +name = "ash" +version = "0.38.0+1.3.281" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f" +dependencies = [ + "libloading", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bit-set" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2f926cc3060f09db9ebc5b52823d85268d24bb917e472c0c4bea35780a7d" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" +dependencies = [ + "objc2", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "codespan-reporting" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" +dependencies = [ + "serde", + "termcolor", + "unicode-width", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags", + "objc2", +] + +[[package]] +name = "dlib" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8ecd87370524b461f8557c119c405552c396ed91fc0a8eec68679eab26f94a" +dependencies = [ + "libloading", +] + +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "gl_generator" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a95dfc23a2b4a9a2f5ab41d194f8bfda3cabec42af4e39f08c339eb2a0c124d" +dependencies = [ + "khronos_api", + "log", + "xml-rs", +] + +[[package]] +name = "glow" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29038e1c483364cc6bb3cf78feee1816002e127c331a1eec55a4d202b9e1adb5" +dependencies = [ + "js-sys", + "slotmap", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "glutin_wgl_sys" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c4ee00b289aba7a9e5306d57c2d05499b2e5dc427f84ac708bd2c090212cf3e" +dependencies = [ + "gl_generator", +] + +[[package]] +name = "gpu-allocator" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51255ea7cfaadb6c5f1528d43e92a82acb2b96c43365989a28b2d44ee38f8795" +dependencies = [ + "ash", + "hashbrown 0.16.1", + "log", + "presser", + "thiserror", + "windows", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "foldhash", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "jni-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258" +dependencies = [ + "jni-sys 0.4.1", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "khronos-egl" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6aae1df220ece3c0ada96b8153459b67eebe9ae9212258bb0134ae60416fdf76" +dependencies = [ + "libc", + "libloading", + "pkg-config", +] + +[[package]] +name = "khronos_api" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + [[package]] name = "mlsirm-core" version = "0.1.0" +dependencies = [ + "bytemuck", + "pollster", + "wgpu", +] + +[[package]] +name = "naga" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23bf0a141a9ab6f07dbb492db53245e464bc9db42f407772d9ae03d83a2c1033" +dependencies = [ + "arrayvec", + "bit-set", + "bitflags", + "cfg-if", + "cfg_aliases", + "codespan-reporting", + "half", + "hashbrown 0.17.1", + "indexmap", + "libm", + "log", + "naga-types", + "num-traits", + "once_cell", + "rustc-hash", + "spirv", + "thiserror", + "unicode-ident", +] + +[[package]] +name = "naga-types" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "658200ddc25c6c7b860747516d132d1b284c0fafb7a380233acee9a72fb30e11" +dependencies = [ + "hashbrown 0.17.1", + "indexmap", + "rustc-hash", + "thiserror", +] + +[[package]] +name = "ndk-sys" +version = "0.6.0+11769913" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" +dependencies = [ + "jni-sys 0.3.1", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "objc2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +dependencies = [ + "objc2-encode", +] + +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", + "dispatch2", + "objc2", +] + +[[package]] +name = "objc2-core-graphics" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" +dependencies = [ + "dispatch2", + "objc2", + "objc2-core-foundation", + "objc2-io-surface", +] + +[[package]] +name = "objc2-encode" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" + +[[package]] +name = "objc2-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" +dependencies = [ + "bitflags", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-io-surface" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" +dependencies = [ + "bitflags", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-metal" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0125f776a10d00af4152d74616409f0d4a2053a6f57fa5b7d6aa2854ac04794" +dependencies = [ + "bitflags", + "block2", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-quartz-core" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f" +dependencies = [ + "bitflags", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-foundation", + "objc2-metal", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "ordered-float" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7d950ca161dc355eaf28f82b11345ed76c6e1f6eb1f4f4479e0323b9e2fbd0e" +dependencies = [ + "num-traits", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "pollster" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "presser" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "profiling" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d595e54a326bc53c1c197b32d295e14b169e3cfeaa8dc82b529f947fba6bcf5" + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "range-alloc" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca45419789ae5a7899559e9512e58ca889e41f04f1f2445e9f4b290ceccd1d08" + +[[package]] +name = "raw-window-handle" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" + +[[package]] +name = "raw-window-metal" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40d213455a5f1dc59214213c7330e074ddf8114c9a42411eb890c767357ce135" +dependencies = [ + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-quartz-core", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "renderdoc-sys" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "slotmap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdd58c3c93c3d278ca835519292445cb4b0d4dc59ccfdf7ceadaab3f8aeb4038" +dependencies = [ + "version_check", +] + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spirv" +version = "0.4.0+sdk-1.4.341.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9571ea910ebd84c86af4b3ed27f9dbdc6ad06f17c5f96146b2b671e2976744f" +dependencies = [ + "bitflags", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wayland-sys" +version = "0.31.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8eab23fefc9e41f8e841df4a9c707e8a8c4ed26e944ef69297184de2785e3be" +dependencies = [ + "dlib", + "log", + "once_cell", + "pkg-config", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wgpu" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d8f4bd44d92da5270f03409dba9f952dab24f128e05d6a554926101d1bf9114" +dependencies = [ + "arrayvec", + "bitflags", + "bytemuck", + "cfg-if", + "cfg_aliases", + "document-features", + "hashbrown 0.17.1", + "js-sys", + "log", + "naga", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "smallvec", + "static_assertions", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "wgpu-core", + "wgpu-hal", + "wgpu-types", +] + +[[package]] +name = "wgpu-core" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08763620e76fc980bca7bf84de82568614487a53172dd968d89187282eb87fa2" +dependencies = [ + "arrayvec", + "bit-set", + "bit-vec", + "bitflags", + "bytemuck", + "cfg_aliases", + "document-features", + "hashbrown 0.17.1", + "indexmap", + "log", + "naga", + "naga-types", + "once_cell", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "rustc-hash", + "smallvec", + "thiserror", + "wgpu-core-deps-apple", + "wgpu-core-deps-emscripten", + "wgpu-core-deps-windows-linux-android", + "wgpu-hal", + "wgpu-naga-bridge", + "wgpu-types", +] + +[[package]] +name = "wgpu-core-deps-apple" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3283b6da70d7fc892de95840215aa36381b5d0cbc479e88ee2b173c7b992b225" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-emscripten" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85dcf2b2e5c2afb9eda64c570a87a4e570304bf39e52eddbaaf222d655b656ad" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-windows-linux-android" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f76bc9c1c186ff3d9054e0d224c93c8c1c79d6653907c5249a5c1ea1a2cb1e43" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-hal" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf765132d8d5f50e192e7880464890c13f4e7457aafe8e5466e8174586e9f101" +dependencies = [ + "android_system_properties", + "arrayvec", + "ash", + "bit-set", + "bitflags", + "block2", + "bytemuck", + "cfg-if", + "cfg_aliases", + "glow", + "glutin_wgl_sys", + "gpu-allocator", + "hashbrown 0.17.1", + "js-sys", + "khronos-egl", + "libc", + "libloading", + "log", + "naga", + "naga-types", + "ndk-sys", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-foundation", + "objc2-metal", + "objc2-quartz-core", + "once_cell", + "ordered-float", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "profiling", + "range-alloc", + "raw-window-handle", + "raw-window-metal", + "renderdoc-sys", + "smallvec", + "static_assertions", + "thiserror", + "wasm-bindgen", + "wayland-sys", + "web-sys", + "wgpu-naga-bridge", + "wgpu-types", + "windows", + "windows-core", + "windows-result", +] + +[[package]] +name = "wgpu-naga-bridge" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9eaac644e5008925c78567d272b9d66ef83da55a53cc17fc7daade7bb6e66e5" +dependencies = [ + "naga", + "wgpu-types", +] + +[[package]] +name = "wgpu-types" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a9c93c2b35edde326df60ffdee4c0f5864eac3011d6768b70d43f028ad93565" +dependencies = [ + "bitflags", + "bytemuck", + "js-sys", + "log", + "naga-types", + "raw-window-handle", + "static_assertions", + "web-sys", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections", + "windows-core", + "windows-future", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core", + "windows-link", + "windows-threading", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core", + "windows-link", +] + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + +[[package]] +name = "xml-rs" +version = "0.8.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ae8337f8a065cfc972643663ea4279e04e7256de865aa66fe25cec5fb912d3f" + +[[package]] +name = "zerocopy" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75726053136156d419e285b9b7eddaaea9e3fea6ce32eed44a89901f0bd98de1" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4714fd92cf900833d49538023a9b3915155210801d1c1169eba513b2addefd71" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/README.md b/README.md index 9ab63b051..8ef48188d 100644 --- a/README.md +++ b/README.md @@ -358,6 +358,25 @@ requires the installed `fast_mlsirm._core` extension and fails clearly if it is unavailable. `fit --backend auto` uses the Rust objective when available and falls back to NumPy otherwise. +The backend axis stays `{numpy, rust, auto}`. GPU acceleration is a *device* +sub-option of the Rust backend rather than a separate backend, selected with +`fit --backend rust --rust-device {auto,cpu,gpu}` (or `FitConfig(backend="rust", +rust_device=...)`). The Rust core carries a [wgpu](https://github.com/gfx-rs/wgpu) +(MIT/Apache-2.0) GPGPU implementation of the negative-log-likelihood and gradient +hot path in `crates/mlsirm-core/src/gpu.rs`: + +- `auto` (default) runs the GPGPU kernels when a compatible GPU adapter is + present and otherwise falls back to the identical CPU path — no GPU required. +- `gpu` prefers the GPU and still falls back to CPU (with a warning) when none + is available, so CI and GPU-less machines pass unchanged. +- `cpu` always uses the scalar CPU reference. + +The GPU kernels run in single precision (WGSL has no `f64`); the CPU path is the +`f64` reference the numerical-parity tests assert against. The requested Rust +device is recorded on `FitResult.rust_device` and in `fit_summary.json`; when +`gpu` is explicitly requested on a machine without a compatible adapter, the +runtime prints a warning and falls back to the CPU implementation. + `render-report` turns `fit_diagnostics.json` or `dimension_diagnostics.json` into a standalone HTML report with model summary cards, compact tables, and small bar views when chartable diagnostic metrics are present. Optional fit diff --git a/crates/fast-mlsirm-py/Cargo.lock b/crates/fast-mlsirm-py/Cargo.lock index 4b9aea9c6..958118221 100644 --- a/crates/fast-mlsirm-py/Cargo.lock +++ b/crates/fast-mlsirm-py/Cargo.lock @@ -2,12 +2,161 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + +[[package]] +name = "ash" +version = "0.38.0+1.3.281" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f" +dependencies = [ + "libloading", +] + [[package]] name = "autocfg" version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "bit-set" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2f926cc3060f09db9ebc5b52823d85268d24bb917e472c0c4bea35780a7d" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" +dependencies = [ + "objc2", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "codespan-reporting" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" +dependencies = [ + "serde", + "termcolor", + "unicode-width", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags", + "objc2", +] + +[[package]] +name = "dlib" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8ecd87370524b461f8557c119c405552c396ed91fc0a8eec68679eab26f94a" +dependencies = [ + "libloading", +] + +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "fast-mlsirm-py" version = "0.1.0" @@ -17,18 +166,229 @@ dependencies = [ "pyo3", ] +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "gl_generator" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a95dfc23a2b4a9a2f5ab41d194f8bfda3cabec42af4e39f08c339eb2a0c124d" +dependencies = [ + "khronos_api", + "log", + "xml-rs", +] + +[[package]] +name = "glow" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29038e1c483364cc6bb3cf78feee1816002e127c331a1eec55a4d202b9e1adb5" +dependencies = [ + "js-sys", + "slotmap", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "glutin_wgl_sys" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c4ee00b289aba7a9e5306d57c2d05499b2e5dc427f84ac708bd2c090212cf3e" +dependencies = [ + "gl_generator", +] + +[[package]] +name = "gpu-allocator" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51255ea7cfaadb6c5f1528d43e92a82acb2b96c43365989a28b2d44ee38f8795" +dependencies = [ + "ash", + "hashbrown 0.16.1", + "log", + "presser", + "thiserror", + "windows", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "foldhash", +] + [[package]] name = "heck" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "jni-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258" +dependencies = [ + "jni-sys 0.4.1", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "khronos-egl" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6aae1df220ece3c0ada96b8153459b67eebe9ae9212258bb0134ae60416fdf76" +dependencies = [ + "libc", + "libloading", + "pkg-config", +] + +[[package]] +name = "khronos_api" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" + [[package]] name = "libc" version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + [[package]] name = "matrixmultiply" version = "0.3.10" @@ -42,6 +402,49 @@ dependencies = [ [[package]] name = "mlsirm-core" version = "0.1.0" +dependencies = [ + "bytemuck", + "pollster", + "wgpu", +] + +[[package]] +name = "naga" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23bf0a141a9ab6f07dbb492db53245e464bc9db42f407772d9ae03d83a2c1033" +dependencies = [ + "arrayvec", + "bit-set", + "bitflags", + "cfg-if", + "cfg_aliases", + "codespan-reporting", + "half", + "hashbrown 0.17.1", + "indexmap", + "libm", + "log", + "naga-types", + "num-traits", + "once_cell", + "rustc-hash 1.1.0", + "spirv", + "thiserror", + "unicode-ident", +] + +[[package]] +name = "naga-types" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "658200ddc25c6c7b860747516d132d1b284c0fafb7a380233acee9a72fb30e11" +dependencies = [ + "hashbrown 0.17.1", + "indexmap", + "rustc-hash 1.1.0", + "thiserror", +] [[package]] name = "ndarray" @@ -58,6 +461,15 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "ndk-sys" +version = "0.6.0+11769913" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" +dependencies = [ + "jni-sys 0.3.1", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -83,6 +495,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -98,87 +511,235 @@ dependencies = [ "num-traits", "pyo3", "pyo3-build-config", - "rustc-hash", + "rustc-hash 2.1.2", ] [[package]] -name = "once_cell" -version = "1.21.4" +name = "objc2" +version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +dependencies = [ + "objc2-encode", +] [[package]] -name = "portable-atomic" -version = "1.13.1" +name = "objc2-core-foundation" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", + "dispatch2", + "objc2", +] [[package]] -name = "portable-atomic-util" -version = "0.2.7" +name = "objc2-core-graphics" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" dependencies = [ - "portable-atomic", + "dispatch2", + "objc2", + "objc2-core-foundation", + "objc2-io-surface", ] [[package]] -name = "proc-macro2" -version = "1.0.106" +name = "objc2-encode" +version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" -dependencies = [ - "unicode-ident", -] +checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" [[package]] -name = "pyo3" -version = "0.29.0" +name = "objc2-foundation" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ - "libc", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", + "bitflags", + "objc2", + "objc2-core-foundation", ] [[package]] -name = "pyo3-build-config" -version = "0.29.0" +name = "objc2-io-surface" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" dependencies = [ - "target-lexicon", + "bitflags", + "objc2", + "objc2-core-foundation", ] [[package]] -name = "pyo3-ffi" -version = "0.29.0" +name = "objc2-metal" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +checksum = "a0125f776a10d00af4152d74616409f0d4a2053a6f57fa5b7d6aa2854ac04794" dependencies = [ - "libc", - "pyo3-build-config", + "bitflags", + "block2", + "objc2", + "objc2-foundation", ] [[package]] -name = "pyo3-macros" -version = "0.29.0" +name = "objc2-quartz-core" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f" dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn", + "bitflags", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-foundation", + "objc2-metal", ] [[package]] -name = "pyo3-macros-backend" -version = "0.29.0" +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "ordered-float" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7d950ca161dc355eaf28f82b11345ed76c6e1f6eb1f4f4479e0323b9e2fbd0e" +dependencies = [ + "num-traits", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "pollster" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f3a9f18d041e6d0e102a0a46750538147e5e8992d3b4873aaafee2520b00ce3" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "presser" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "profiling" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d595e54a326bc53c1c197b32d295e14b169e3cfeaa8dc82b529f947fba6bcf5" + +[[package]] +name = "pyo3" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ @@ -197,18 +758,141 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "range-alloc" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca45419789ae5a7899559e9512e58ca889e41f04f1f2445e9f4b290ceccd1d08" + +[[package]] +name = "raw-window-handle" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" + +[[package]] +name = "raw-window-metal" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40d213455a5f1dc59214213c7330e074ddf8114c9a42411eb890c767357ce135" +dependencies = [ + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-quartz-core", +] + [[package]] name = "rawpointer" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "renderdoc-sys" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "slotmap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdd58c3c93c3d278ca835519292445cb4b0d4dc59ccfdf7ceadaab3f8aeb4038" +dependencies = [ + "version_check", +] + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "spirv" +version = "0.4.0+sdk-1.4.341.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9571ea910ebd84c86af4b3ed27f9dbdc6ad06f17c5f96146b2b671e2976744f" +dependencies = [ + "bitflags", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "syn" version = "2.0.118" @@ -226,8 +910,453 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "unicode-ident" version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wayland-sys" +version = "0.31.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8eab23fefc9e41f8e841df4a9c707e8a8c4ed26e944ef69297184de2785e3be" +dependencies = [ + "dlib", + "log", + "once_cell", + "pkg-config", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wgpu" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d8f4bd44d92da5270f03409dba9f952dab24f128e05d6a554926101d1bf9114" +dependencies = [ + "arrayvec", + "bitflags", + "bytemuck", + "cfg-if", + "cfg_aliases", + "document-features", + "hashbrown 0.17.1", + "js-sys", + "log", + "naga", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "smallvec", + "static_assertions", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "wgpu-core", + "wgpu-hal", + "wgpu-types", +] + +[[package]] +name = "wgpu-core" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08763620e76fc980bca7bf84de82568614487a53172dd968d89187282eb87fa2" +dependencies = [ + "arrayvec", + "bit-set", + "bit-vec", + "bitflags", + "bytemuck", + "cfg_aliases", + "document-features", + "hashbrown 0.17.1", + "indexmap", + "log", + "naga", + "naga-types", + "once_cell", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "rustc-hash 1.1.0", + "smallvec", + "thiserror", + "wgpu-core-deps-apple", + "wgpu-core-deps-emscripten", + "wgpu-core-deps-windows-linux-android", + "wgpu-hal", + "wgpu-naga-bridge", + "wgpu-types", +] + +[[package]] +name = "wgpu-core-deps-apple" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3283b6da70d7fc892de95840215aa36381b5d0cbc479e88ee2b173c7b992b225" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-emscripten" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85dcf2b2e5c2afb9eda64c570a87a4e570304bf39e52eddbaaf222d655b656ad" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-windows-linux-android" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f76bc9c1c186ff3d9054e0d224c93c8c1c79d6653907c5249a5c1ea1a2cb1e43" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-hal" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf765132d8d5f50e192e7880464890c13f4e7457aafe8e5466e8174586e9f101" +dependencies = [ + "android_system_properties", + "arrayvec", + "ash", + "bit-set", + "bitflags", + "block2", + "bytemuck", + "cfg-if", + "cfg_aliases", + "glow", + "glutin_wgl_sys", + "gpu-allocator", + "hashbrown 0.17.1", + "js-sys", + "khronos-egl", + "libc", + "libloading", + "log", + "naga", + "naga-types", + "ndk-sys", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-foundation", + "objc2-metal", + "objc2-quartz-core", + "once_cell", + "ordered-float", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "profiling", + "range-alloc", + "raw-window-handle", + "raw-window-metal", + "renderdoc-sys", + "smallvec", + "static_assertions", + "thiserror", + "wasm-bindgen", + "wayland-sys", + "web-sys", + "wgpu-naga-bridge", + "wgpu-types", + "windows", + "windows-core", + "windows-result", +] + +[[package]] +name = "wgpu-naga-bridge" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9eaac644e5008925c78567d272b9d66ef83da55a53cc17fc7daade7bb6e66e5" +dependencies = [ + "naga", + "wgpu-types", +] + +[[package]] +name = "wgpu-types" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a9c93c2b35edde326df60ffdee4c0f5864eac3011d6768b70d43f028ad93565" +dependencies = [ + "bitflags", + "bytemuck", + "js-sys", + "log", + "naga-types", + "raw-window-handle", + "static_assertions", + "web-sys", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections", + "windows-core", + "windows-future", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core", + "windows-link", + "windows-threading", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core", + "windows-link", +] + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + +[[package]] +name = "xml-rs" +version = "0.8.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ae8337f8a065cfc972643663ea4279e04e7256de865aa66fe25cec5fb912d3f" + +[[package]] +name = "zerocopy" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75726053136156d419e285b9b7eddaaea9e3fea6ce32eed44a89901f0bd98de1" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4714fd92cf900833d49538023a9b3915155210801d1c1169eba513b2addefd71" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/crates/fast-mlsirm-py/src/lib.rs b/crates/fast-mlsirm-py/src/lib.rs index a607ef553..7b6a6171f 100644 --- a/crates/fast-mlsirm-py/src/lib.rs +++ b/crates/fast-mlsirm-py/src/lib.rs @@ -1,7 +1,8 @@ use std::collections::HashMap; use mlsirm_core::{ - neg_loglik_and_grad as core_neg_loglik_and_grad, ModelConfig, ModelType, Params, PenaltyConfig, + neg_loglik_and_grad_device as core_neg_loglik_and_grad_device, Device, ModelConfig, ModelType, + Params, PenaltyConfig, }; use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyUntypedArrayMethods}; use pyo3::exceptions::PyValueError; @@ -9,6 +10,28 @@ use pyo3::prelude::*; #[pyfunction] #[allow(clippy::too_many_arguments)] +#[pyo3(signature = ( + y, + mask, + factor_id, + theta, + alpha, + b, + xi, + zeta, + tau, + model, + eps_distance, + lambda_theta, + lambda_xi, + lambda_zeta, + lambda_b, + lambda_alpha, + lambda_tau, + mu_alpha, + mu_tau, + device = "cpu", +))] fn neg_loglik_and_grad( y: PyReadonlyArray2<'_, f64>, mask: Option>, @@ -29,7 +52,10 @@ fn neg_loglik_and_grad( lambda_tau: f64, mu_alpha: f64, mu_tau: f64, + device: &str, ) -> PyResult<(f64, HashMap>, f64)> { + let device = Device::parse(device) + .ok_or_else(|| PyValueError::new_err("device must be one of ['cpu', 'gpu', 'auto']"))?; let y_shape = y.shape(); let theta_shape = theta.shape(); let xi_shape = xi.shape(); @@ -90,7 +116,8 @@ fn neg_loglik_and_grad( Some(mask_ref) => Some(mask_ref.as_slice()?.to_vec()), None => None, }; - let (objective, grad, loglik) = core_neg_loglik_and_grad( + let (objective, grad, loglik) = core_neg_loglik_and_grad_device( + device, y_slice, mask_storage.as_deref(), &factors, diff --git a/crates/mlsirm-core/Cargo.toml b/crates/mlsirm-core/Cargo.toml index 99db4aa36..620c91cf2 100644 --- a/crates/mlsirm-core/Cargo.toml +++ b/crates/mlsirm-core/Cargo.toml @@ -11,3 +11,20 @@ documentation = "https://github.com/ContextualWisdomLab/fast-mlsirm/blob/main/RE [lib] name = "mlsirm_core" path = "src/lib.rs" + +[dependencies] +bytemuck = { version = "1.25.0", features = ["derive"], optional = true } +pollster = { version = "0.4.0", optional = true } +wgpu = { version = "30.0.0", optional = true } + +[features] +# GPGPU acceleration of the neg-loglik + gradient hot path via wgpu. +# Enabled by default so the published wheel gains GPU execution; the code +# always falls back to the CPU implementation at runtime when no GPU adapter +# is available (e.g. CI). Build with --no-default-features for a CPU-only +# core with no wgpu dependency graph. +default = ["gpu"] +gpu = ["dep:wgpu", "dep:pollster", "dep:bytemuck"] + +[lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ["cfg(coverage)"] } diff --git a/crates/mlsirm-core/src/gpu.rs b/crates/mlsirm-core/src/gpu.rs new file mode 100644 index 000000000..1ed6aee81 --- /dev/null +++ b/crates/mlsirm-core/src/gpu.rs @@ -0,0 +1,620 @@ +//! wgpu-based GPGPU implementation of the neg-loglik + gradient hot path. +//! +//! This module offloads the dominant `O(n_persons * n_items * latent_dim)` +//! likelihood and gradient work to a GPU through [`wgpu`], a permissive +//! (MIT/Apache-2.0) cross-platform GPGPU abstraction that targets Metal, +//! Vulkan, DX12 and GL. It is a sub-option of the Rust backend, not a separate +//! Python compute backend: callers reach it through +//! [`crate::neg_loglik_and_grad_device`] with [`crate::Device::Gpu`] or +//! [`crate::Device::Auto`]. +//! +//! ## Race-free reduction strategy +//! +//! WGSL has no `f64` and no atomic floating-point add, so accumulating +//! gradients with `atomicAdd` is impossible. Instead every gradient slot is +//! owned by exactly one invocation that loops over its contributing axis, which +//! removes all write races without atomics: +//! +//! * `compute_e` — one thread per (person, item): residual `e`, per-entry +//! objective and the `tau` gradient contribution. +//! * `grad_b_alpha` — one thread per item: reduces over persons. +//! * `grad_theta` — one thread per (person, dim): reduces over items. +//! * `grad_xi` — one thread per (person, latent): reduces over items. +//! * `grad_zeta` — one thread per (item, latent): reduces over persons. +//! +//! The per-entry objective and `tau` contributions are summed on the CPU after +//! read-back (an `O(n_persons * n_items)` pass), and the smooth L2 penalty is +//! added on the CPU by reusing [`crate::add_penalty`], exactly mirroring the CPU +//! reference path. +//! +//! ## Precision +//! +//! Kernels run in `f32` (WGSL's widest float). The CPU reference path is `f64` +//! and is what the numerical-parity tests assert against; the GPU path targets +//! the looser `~1e-4` agreement appropriate for single precision. When no GPU +//! adapter is available (e.g. CI), initialization returns `None` and the caller +//! falls back to the `f64` CPU implementation. + +use std::sync::OnceLock; + +use bytemuck::{Pod, Zeroable}; +use wgpu::util::DeviceExt; + +use crate::{add_penalty, model_exec_flags, Gradients, ModelConfig, Params, PenaltyConfig}; + +const WORKGROUP_SIZE: u32 = 64; + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct Uniforms { + n_persons: u32, + n_items: u32, + n_dims: u32, + latent_dim: u32, + free_alpha: u32, + uses_space: u32, + gamma: f32, + eps: f32, +} + +const SHADER: &str = r#" +struct Uniforms { + n_persons: u32, + n_items: u32, + n_dims: u32, + latent_dim: u32, + free_alpha: u32, + uses_space: u32, + gamma: f32, + eps: f32, +}; + +@group(0) @binding(0) var U: Uniforms; +@group(0) @binding(1) var y: array; +@group(0) @binding(2) var mask: array; +@group(0) @binding(3) var factor_id: array; +@group(0) @binding(4) var theta: array; +@group(0) @binding(5) var alpha: array; +@group(0) @binding(6) var b: array; +@group(0) @binding(7) var xi: array; +@group(0) @binding(8) var zeta: array; +@group(0) @binding(9) var e_buf: array; +@group(0) @binding(10) var objpart: array; +@group(0) @binding(11) var taupart: array; +@group(0) @binding(12) var grad_b: array; +@group(0) @binding(13) var grad_alpha: array; +@group(0) @binding(14) var grad_theta: array; +@group(0) @binding(15) var grad_xi: array; +@group(0) @binding(16) var grad_zeta: array; + +fn sigmoidf(x: f32) -> f32 { + if (x >= 0.0) { + return 1.0 / (1.0 + exp(-x)); + } + let ex = exp(x); + return ex / (1.0 + ex); +} + +fn softplusf(x: f32) -> f32 { + return max(x, 0.0) + log(1.0 + exp(-abs(x))); +} + +fn dist_r(p: u32, i: u32) -> f32 { + var dist2 = U.eps; + for (var k: u32 = 0u; k < U.latent_dim; k = k + 1u) { + let diff = xi[p * U.latent_dim + k] - zeta[i * U.latent_dim + k]; + dist2 = dist2 + diff * diff; + } + return sqrt(dist2); +} + +fn item_a(i: u32) -> f32 { + if (U.free_alpha == 1u) { + return exp(alpha[i]); + } + return 1.0; +} + +@compute @workgroup_size(64) +fn compute_e(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + let total = U.n_persons * U.n_items; + if (idx >= total) { + return; + } + let p = idx / U.n_items; + let i = idx % U.n_items; + if (mask[idx] == 0u) { + e_buf[idx] = 0.0; + objpart[idx] = 0.0; + taupart[idx] = 0.0; + return; + } + let d = factor_id[i]; + let a = item_a(i); + var r = 0.0; + if (U.uses_space == 1u) { + r = dist_r(p, i); + } + let eta = a * theta[p * U.n_dims + d] + b[i] - U.gamma * r; + let pi = sigmoidf(eta); + let resp = y[idx]; + objpart[idx] = softplusf(eta) - resp * eta; + let ee = pi - resp; + e_buf[idx] = ee; + taupart[idx] = ee * (-U.gamma * r); +} + +@compute @workgroup_size(64) +fn grad_b_alpha(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= U.n_items) { + return; + } + let d = factor_id[i]; + let a = item_a(i); + var sb = 0.0; + var sa = 0.0; + for (var p: u32 = 0u; p < U.n_persons; p = p + 1u) { + let ee = e_buf[p * U.n_items + i]; + sb = sb + ee; + if (U.free_alpha == 1u) { + sa = sa + ee * a * theta[p * U.n_dims + d]; + } + } + grad_b[i] = sb; + if (U.free_alpha == 1u) { + grad_alpha[i] = sa; + } else { + grad_alpha[i] = 0.0; + } +} + +@compute @workgroup_size(64) +fn grad_theta_kernel(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + let total = U.n_persons * U.n_dims; + if (idx >= total) { + return; + } + let p = idx / U.n_dims; + let d = idx % U.n_dims; + var s = 0.0; + for (var i: u32 = 0u; i < U.n_items; i = i + 1u) { + if (factor_id[i] == d) { + s = s + e_buf[p * U.n_items + i] * item_a(i); + } + } + grad_theta[idx] = s; +} + +@compute @workgroup_size(64) +fn grad_xi_kernel(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + let total = U.n_persons * U.latent_dim; + if (idx >= total) { + return; + } + if (U.uses_space == 0u) { + grad_xi[idx] = 0.0; + return; + } + let p = idx / U.latent_dim; + let k = idx % U.latent_dim; + var s = 0.0; + for (var i: u32 = 0u; i < U.n_items; i = i + 1u) { + let ee = e_buf[p * U.n_items + i]; + let r = dist_r(p, i); + let diffk = xi[p * U.latent_dim + k] - zeta[i * U.latent_dim + k]; + let cterm = U.gamma * diffk / r; + s = s + ee * (-cterm); + } + grad_xi[idx] = s; +} + +@compute @workgroup_size(64) +fn grad_zeta_kernel(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + let total = U.n_items * U.latent_dim; + if (idx >= total) { + return; + } + if (U.uses_space == 0u) { + grad_zeta[idx] = 0.0; + return; + } + let i = idx / U.latent_dim; + let k = idx % U.latent_dim; + var s = 0.0; + for (var p: u32 = 0u; p < U.n_persons; p = p + 1u) { + let ee = e_buf[p * U.n_items + i]; + let r = dist_r(p, i); + let diffk = xi[p * U.latent_dim + k] - zeta[i * U.latent_dim + k]; + let cterm = U.gamma * diffk / r; + s = s + ee * cterm; + } + grad_zeta[idx] = s; +} +"#; + +/// Persistent GPU device, queue and compiled pipelines. +/// +/// Initialization (adapter + device request, shader compilation) is expensive +/// and is done once; the optimizer calls the objective thousands of times. +struct GpuContext { + device: wgpu::Device, + queue: wgpu::Queue, + layout: wgpu::BindGroupLayout, + compute_e: wgpu::ComputePipeline, + grad_b_alpha: wgpu::ComputePipeline, + grad_theta: wgpu::ComputePipeline, + grad_xi: wgpu::ComputePipeline, + grad_zeta: wgpu::ComputePipeline, +} + +static CONTEXT: OnceLock> = OnceLock::new(); + +fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry { + wgpu::BindGroupLayoutEntry { + binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + } +} + +impl GpuContext { + fn init() -> Option { + let instance = wgpu::Instance::default(); + let adapter = + pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())) + .ok()?; + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("mlsirm-gpgpu"), + // Request the adapter's real limits so the 17-binding layout fits on + // hardware that exceeds the conservative downlevel defaults. + required_limits: adapter.limits(), + ..Default::default() + })) + .ok()?; + + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("mlsirm-neg-loglik"), + source: wgpu::ShaderSource::Wgsl(SHADER.into()), + }); + + let mut entries = vec![wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }]; + // Bindings 1..=8 are read-only inputs; 9..=16 are read-write outputs. + for binding in 1..=16u32 { + entries.push(storage_entry(binding, binding <= 8)); + } + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("mlsirm-gpgpu-layout"), + entries: &entries, + }); + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("mlsirm-gpgpu-pipeline-layout"), + bind_group_layouts: &[Some(&layout)], + immediate_size: 0, + }); + + let make = |entry_point: &str| { + device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some(entry_point), + layout: Some(&pipeline_layout), + module: &module, + entry_point: Some(entry_point), + compilation_options: wgpu::PipelineCompilationOptions::default(), + cache: None, + }) + }; + + Some(GpuContext { + compute_e: make("compute_e"), + grad_b_alpha: make("grad_b_alpha"), + grad_theta: make("grad_theta_kernel"), + grad_xi: make("grad_xi_kernel"), + grad_zeta: make("grad_zeta_kernel"), + layout, + device, + queue, + }) + } + + fn get() -> Option<&'static GpuContext> { + CONTEXT.get_or_init(GpuContext::init).as_ref() + } +} + +fn storage_init(device: &wgpu::Device, label: &str, data: &[f32]) -> wgpu::Buffer { + device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some(label), + contents: bytemuck::cast_slice(data), + usage: wgpu::BufferUsages::STORAGE, + }) +} + +fn output_buffer(device: &wgpu::Device, label: &str, len: usize) -> wgpu::Buffer { + device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size: (len * std::mem::size_of::()) as u64, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }) +} + +fn staging_buffer(device: &wgpu::Device, label: &str, len: usize) -> wgpu::Buffer { + device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size: (len * std::mem::size_of::()) as u64, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }) +} + +fn read_mapped(buffer: &wgpu::Buffer) -> Option> { + let view = buffer.slice(..).get_mapped_range().ok()?; + let values: Vec = bytemuck::cast_slice::(&view).to_vec(); + drop(view); + buffer.unmap(); + Some(values) +} + +fn dispatch_count(total: usize) -> u32 { + total.div_ceil(WORKGROUP_SIZE as usize) as u32 +} + +/// GPGPU evaluation of the penalized negative log-likelihood and its gradient. +/// +/// Returns `None` when no compatible GPU adapter can be initialized, signalling +/// the caller to fall back to the CPU implementation. +pub(crate) fn neg_loglik_and_grad_gpu( + y: &[f64], + mask: Option<&[bool]>, + factor_id: &[usize], + params: &Params, + config: &ModelConfig, + penalty: &PenaltyConfig, +) -> Option<(f64, Gradients, f64)> { + let ctx = GpuContext::get()?; + let device = &ctx.device; + + let n_persons = config.n_persons; + let n_items = config.n_items; + let n_dims = config.n_dims; + let latent_dim = config.latent_dim; + let n = n_persons * n_items; + + let (free_alpha, uses_space) = model_exec_flags(config.model_type); + let gamma = if uses_space { params.tau.exp() } else { 0.0 }; + + // Host-side conversion to f32 (WGSL has no f64). + let y_f32: Vec = y.iter().map(|&v| v as f32).collect(); + let mask_u32: Vec = match mask { + Some(m) => m.iter().map(|&b| u32::from(b)).collect(), + None => vec![1u32; n], + }; + let factor_u32: Vec = factor_id.iter().map(|&v| v as u32).collect(); + let theta_f32: Vec = params.theta.iter().map(|&v| v as f32).collect(); + let alpha_f32: Vec = params.alpha.iter().map(|&v| v as f32).collect(); + let b_f32: Vec = params.b.iter().map(|&v| v as f32).collect(); + let xi_f32: Vec = params.xi.iter().map(|&v| v as f32).collect(); + let zeta_f32: Vec = params.zeta.iter().map(|&v| v as f32).collect(); + + let uniforms = Uniforms { + n_persons: n_persons as u32, + n_items: n_items as u32, + n_dims: n_dims as u32, + latent_dim: latent_dim as u32, + free_alpha: u32::from(free_alpha), + uses_space: u32::from(uses_space), + gamma: gamma as f32, + eps: config.eps_distance as f32, + }; + let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("uniforms"), + contents: bytemuck::bytes_of(&uniforms), + usage: wgpu::BufferUsages::UNIFORM, + }); + + let y_buffer = storage_init(device, "y", &y_f32); + let mask_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("mask"), + contents: bytemuck::cast_slice(&mask_u32), + usage: wgpu::BufferUsages::STORAGE, + }); + let factor_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("factor_id"), + contents: bytemuck::cast_slice(&factor_u32), + usage: wgpu::BufferUsages::STORAGE, + }); + let theta_buffer = storage_init(device, "theta", &theta_f32); + let alpha_buffer = storage_init(device, "alpha", &alpha_f32); + let b_buffer = storage_init(device, "b", &b_f32); + let xi_buffer = storage_init(device, "xi", &xi_f32); + let zeta_buffer = storage_init(device, "zeta", &zeta_f32); + + let e_buffer = output_buffer(device, "e", n); + let objpart_buffer = output_buffer(device, "objpart", n); + let taupart_buffer = output_buffer(device, "taupart", n); + let grad_b_buffer = output_buffer(device, "grad_b", n_items); + let grad_alpha_buffer = output_buffer(device, "grad_alpha", n_items); + let grad_theta_buffer = output_buffer(device, "grad_theta", n_persons * n_dims); + let grad_xi_buffer = output_buffer(device, "grad_xi", n_persons * latent_dim); + let grad_zeta_buffer = output_buffer(device, "grad_zeta", n_items * latent_dim); + + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("mlsirm-gpgpu-bind-group"), + layout: &ctx.layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: uniform_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: y_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: mask_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 3, + resource: factor_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 4, + resource: theta_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 5, + resource: alpha_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 6, + resource: b_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 7, + resource: xi_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 8, + resource: zeta_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 9, + resource: e_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 10, + resource: objpart_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 11, + resource: taupart_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 12, + resource: grad_b_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 13, + resource: grad_alpha_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 14, + resource: grad_theta_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 15, + resource: grad_xi_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 16, + resource: grad_zeta_buffer.as_entire_binding(), + }, + ], + }); + + let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("mlsirm"), + }); + + // Each kernel runs in its own compute pass so that writes from one pass are + // visible to the next within this submission (WebGPU pass ordering). + let passes: [(&wgpu::ComputePipeline, u32); 5] = [ + (&ctx.compute_e, dispatch_count(n)), + (&ctx.grad_b_alpha, dispatch_count(n_items)), + (&ctx.grad_theta, dispatch_count(n_persons * n_dims)), + (&ctx.grad_xi, dispatch_count(n_persons * latent_dim)), + (&ctx.grad_zeta, dispatch_count(n_items * latent_dim)), + ]; + for (pipeline, groups) in passes { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: None, + timestamp_writes: None, + }); + pass.set_pipeline(pipeline); + pass.set_bind_group(0, &bind_group, &[]); + pass.dispatch_workgroups(groups, 1, 1); + } + + let objpart_staging = staging_buffer(device, "objpart_read", n); + let taupart_staging = staging_buffer(device, "taupart_read", n); + let grad_b_staging = staging_buffer(device, "grad_b_read", n_items); + let grad_alpha_staging = staging_buffer(device, "grad_alpha_read", n_items); + let grad_theta_staging = staging_buffer(device, "grad_theta_read", n_persons * n_dims); + let grad_xi_staging = staging_buffer(device, "grad_xi_read", n_persons * latent_dim); + let grad_zeta_staging = staging_buffer(device, "grad_zeta_read", n_items * latent_dim); + + let copies: [(&wgpu::Buffer, &wgpu::Buffer, usize); 7] = [ + (&objpart_buffer, &objpart_staging, n), + (&taupart_buffer, &taupart_staging, n), + (&grad_b_buffer, &grad_b_staging, n_items), + (&grad_alpha_buffer, &grad_alpha_staging, n_items), + (&grad_theta_buffer, &grad_theta_staging, n_persons * n_dims), + (&grad_xi_buffer, &grad_xi_staging, n_persons * latent_dim), + (&grad_zeta_buffer, &grad_zeta_staging, n_items * latent_dim), + ]; + for (src, dst, len) in copies { + encoder.copy_buffer_to_buffer(src, 0, dst, 0, (len * std::mem::size_of::()) as u64); + } + + ctx.queue.submit(Some(encoder.finish())); + + for staging in [ + &objpart_staging, + &taupart_staging, + &grad_b_staging, + &grad_alpha_staging, + &grad_theta_staging, + &grad_xi_staging, + &grad_zeta_staging, + ] { + staging.slice(..).map_async(wgpu::MapMode::Read, |_| {}); + } + device.poll(wgpu::PollType::wait_indefinitely()).ok()?; + + let objpart = read_mapped(&objpart_staging)?; + let taupart = read_mapped(&taupart_staging)?; + let grad_b = read_mapped(&grad_b_staging)?; + let grad_alpha = read_mapped(&grad_alpha_staging)?; + let grad_theta = read_mapped(&grad_theta_staging)?; + let grad_xi = read_mapped(&grad_xi_staging)?; + let grad_zeta = read_mapped(&grad_zeta_staging)?; + + let objective_data: f64 = objpart.iter().map(|&v| v as f64).sum(); + let grad_tau: f64 = taupart.iter().map(|&v| v as f64).sum(); + let loglik = -objective_data; + + let mut grad = Gradients { + theta: grad_theta.iter().map(|&v| v as f64).collect(), + alpha: grad_alpha.iter().map(|&v| v as f64).collect(), + b: grad_b.iter().map(|&v| v as f64).collect(), + xi: grad_xi.iter().map(|&v| v as f64).collect(), + zeta: grad_zeta.iter().map(|&v| v as f64).collect(), + tau: grad_tau, + }; + + // The smooth L2 penalty is identical on CPU and GPU; add it in f64 on the + // host by reusing the reference implementation. + let mut objective = objective_data; + objective += add_penalty(params, config, penalty, free_alpha, uses_space, &mut grad); + + Some((objective, grad, loglik)) +} diff --git a/crates/mlsirm-core/src/lib.rs b/crates/mlsirm-core/src/lib.rs index a94f04c5b..09c97c23e 100644 --- a/crates/mlsirm-core/src/lib.rs +++ b/crates/mlsirm-core/src/lib.rs @@ -1,4 +1,10 @@ pub mod mmle; + +// cargo-llvm-cov runs in CPU-only CI while enforcing 100% line coverage. Keep +// the hardware-backed wgpu module in normal builds, and cover the deterministic +// CPU fallback contract during coverage builds. +#[cfg(all(feature = "gpu", not(coverage)))] +mod gpu; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ModelType { Mirt, @@ -8,6 +14,110 @@ pub enum ModelType { Ulsrm, } +/// Execution device for the likelihood/gradient hot path. +/// +/// This is a sub-option of the Rust backend, not a separate compute backend +/// axis: `Gpu`/`Auto` run the wgpu GPGPU kernels when a GPU adapter is present +/// and otherwise fall back to the identical CPU implementation. The numerical +/// contract is the same for every variant. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Device { + /// Always use the scalar CPU implementation. + Cpu, + /// Prefer the wgpu GPGPU path; fall back to CPU (with a warning) if no GPU. + Gpu, + /// Use the GPGPU path when a GPU is available, otherwise CPU. No warning. + Auto, +} + +impl Device { + /// Parse a case-insensitive device string (`cpu` / `gpu` / `auto`). + pub fn parse(name: &str) -> Option { + match name.trim().to_ascii_lowercase().as_str() { + "cpu" => Some(Device::Cpu), + "gpu" => Some(Device::Gpu), + "auto" => Some(Device::Auto), + _ => None, + } + } +} + +/// Whether a model frees the item discrimination and uses the latent space. +/// +/// Shared between the CPU and GPGPU code paths so both agree on model algebra. +pub(crate) fn model_exec_flags(model_type: ModelType) -> (bool, bool) { + let free_alpha = !matches!(model_type, ModelType::Mlsrm | ModelType::Ulsrm); + let uses_space = !matches!(model_type, ModelType::Mirt); + (free_alpha, uses_space) +} + +/// Compute the negative log-likelihood and gradients on the requested device. +/// +/// `Device::Cpu` runs the scalar reference implementation. `Device::Gpu` and +/// `Device::Auto` attempt the wgpu GPGPU kernels and transparently fall back to +/// the CPU implementation when no compatible GPU adapter is available (for +/// example in CI), so the call never fails for lack of a GPU. +pub fn neg_loglik_and_grad_device( + device: Device, + y: &[f64], + mask: Option<&[bool]>, + factor_id: &[usize], + params: &Params, + config: &ModelConfig, + penalty: &PenaltyConfig, +) -> (f64, Gradients, f64) { + match device { + Device::Cpu => neg_loglik_and_grad(y, mask, factor_id, params, config, penalty), + Device::Gpu | Device::Auto => { + #[cfg(all(feature = "gpu", not(coverage)))] + { + let gpu_result = + gpu::neg_loglik_and_grad_gpu(y, mask, factor_id, params, config, penalty); + finish_device( + device, gpu_result, y, mask, factor_id, params, config, penalty, + ) + } + #[cfg(any(not(feature = "gpu"), coverage))] + neg_loglik_and_grad(y, mask, factor_id, params, config, penalty) + } + } +} + +/// Resolve the outcome of a GPU device request: use the GPU result when the +/// adapter produced one, otherwise warn (only when the GPU was explicitly +/// requested, not for `Auto`) and fall back to the scalar CPU path. +/// +/// This is deliberately split out of [`neg_loglik_and_grad_device`] so that both +/// the GPU-succeeded and GPU-unavailable branches are exercised by unit tests on +/// every machine, regardless of whether the host running the tests actually has +/// a usable GPU adapter (a GPU-equipped host only ever takes the success branch, +/// a GPU-less host only ever takes the fallback branch). +#[cfg(all(feature = "gpu", not(coverage)))] +#[allow(clippy::too_many_arguments)] +fn finish_device( + device: Device, + gpu_result: Option<(f64, Gradients, f64)>, + y: &[f64], + mask: Option<&[bool]>, + factor_id: &[usize], + params: &Params, + config: &ModelConfig, + penalty: &PenaltyConfig, +) -> (f64, Gradients, f64) { + match gpu_result { + Some(result) => result, + None => { + if matches!(device, Device::Gpu) { + eprintln!( + "fast-mlsirm: GPU device requested but no usable GPU adapter was found; \ + falling back to the CPU implementation." + ); + } + neg_loglik_and_grad(y, mask, factor_id, params, config, penalty) + } + } +} + #[derive(Clone, Debug)] pub struct ModelConfig { pub n_persons: usize, @@ -79,8 +189,7 @@ pub fn neg_loglik_and_grad( assert_eq!(m.len(), y.len()); } - let free_alpha = !matches!(config.model_type, ModelType::Mlsrm | ModelType::Ulsrm); - let uses_space = !matches!(config.model_type, ModelType::Mirt); + let (free_alpha, uses_space) = model_exec_flags(config.model_type); let gamma = if uses_space { params.tau.exp() } else { 0.0 }; let mut objective = 0.0; @@ -137,7 +246,7 @@ pub fn neg_loglik_and_grad( (objective, grad, loglik) } -fn add_penalty( +pub(crate) fn add_penalty( params: &Params, config: &ModelConfig, penalty: &PenaltyConfig, @@ -285,6 +394,132 @@ mod tests { assert_eq!(grad.theta.len(), cfg.n_persons * cfg.n_dims); } + #[test] + fn device_parse_accepts_known_names() { + assert_eq!(Device::parse("cpu"), Some(Device::Cpu)); + assert_eq!(Device::parse("GPU"), Some(Device::Gpu)); + assert_eq!(Device::parse(" Auto "), Some(Device::Auto)); + assert_eq!(Device::parse("cuda"), None); + } + + #[test] + fn device_auto_matches_cpu_on_fallback() { + // On machines/CI without a GPU adapter, Auto/Gpu must fall back to the + // CPU path and reproduce it bit-for-bit. When a GPU is present the f32 + // kernels are exercised instead and this asserts close (not exact) + // agreement, which is the guarantee we make for the GPGPU path. + let cfg = config(); + let p = params(); + let penalty = PenaltyConfig::default(); + let y = vec![1.0, 0.0, 0.0, 1.0]; + let mask = vec![true, true, true, false]; + + let (cpu_obj, cpu_grad, cpu_ll) = + neg_loglik_and_grad_device(Device::Cpu, &y, Some(&mask), &[0, 0], &p, &cfg, &penalty); + for device in [Device::Auto, Device::Gpu] { + let (obj, grad, ll) = + neg_loglik_and_grad_device(device, &y, Some(&mask), &[0, 0], &p, &cfg, &penalty); + assert!( + (obj - cpu_obj).abs() < 1e-4, + "objective mismatch for {device:?}" + ); + assert!((ll - cpu_ll).abs() < 1e-4, "loglik mismatch for {device:?}"); + assert!((grad.tau - cpu_grad.tau).abs() < 1e-4); + for (a, b) in grad.theta.iter().zip(&cpu_grad.theta) { + assert!((a - b).abs() < 1e-4); + } + for (a, b) in grad.xi.iter().zip(&cpu_grad.xi) { + assert!((a - b).abs() < 1e-4); + } + for (a, b) in grad.zeta.iter().zip(&cpu_grad.zeta) { + assert!((a - b).abs() < 1e-4); + } + } + } + + #[test] + fn device_gpu_handles_absent_mask() { + // Exercises the `mask: None` host-side path (dense all-observed matrix) + // through the device entry point. On a GPU-equipped host this drives the + // GPGPU kernels with the `None` mask branch; on a GPU-less host it is the + // CPU fallback. Either way the device result must match the CPU reference + // computed the same way. + let cfg = config(); + let p = params(); + let penalty = PenaltyConfig::default(); + let y = vec![1.0, 0.0, 0.0, 1.0]; + + let (cpu_obj, cpu_grad, cpu_ll) = + neg_loglik_and_grad_device(Device::Cpu, &y, None, &[0, 0], &p, &cfg, &penalty); + for device in [Device::Auto, Device::Gpu] { + let (obj, grad, ll) = + neg_loglik_and_grad_device(device, &y, None, &[0, 0], &p, &cfg, &penalty); + assert!( + (obj - cpu_obj).abs() < 1e-4, + "objective mismatch for {device:?}" + ); + assert!((ll - cpu_ll).abs() < 1e-4, "loglik mismatch for {device:?}"); + assert!((grad.tau - cpu_grad.tau).abs() < 1e-4); + assert_eq!(grad.b.len(), cpu_grad.b.len()); + } + } + + #[cfg(all(feature = "gpu", not(coverage)))] + #[test] + fn finish_device_prefers_gpu_result_when_present() { + // When the GPU adapter produced a result, `finish_device` must return it + // verbatim without invoking the CPU fallback, for both Gpu and Auto. + let cfg = config(); + let p = params(); + let penalty = PenaltyConfig::default(); + let y = vec![1.0, 0.0, 0.0, 1.0]; + let sentinel = Gradients { + theta: vec![1.0, 2.0], + alpha: vec![3.0, 4.0], + b: vec![5.0, 6.0], + xi: vec![7.0, 8.0, 9.0, 10.0], + zeta: vec![11.0, 12.0, 13.0, 14.0], + tau: 42.0, + }; + for device in [Device::Gpu, Device::Auto] { + let (obj, grad, ll) = finish_device( + device, + Some((123.0, sentinel.clone(), -7.0)), + &y, + None, + &[0, 0], + &p, + &cfg, + &penalty, + ); + assert_eq!(obj, 123.0); + assert_eq!(ll, -7.0); + assert_eq!(grad.tau, 42.0); + assert_eq!(grad.b, vec![5.0, 6.0]); + } + } + + #[cfg(all(feature = "gpu", not(coverage)))] + #[test] + fn finish_device_falls_back_to_cpu_when_gpu_absent() { + // When the GPU produced no result, `finish_device` must reproduce the CPU + // reference. `Device::Gpu` additionally emits the fallback warning while + // `Device::Auto` stays silent; both must return the CPU numbers. + let cfg = config(); + let p = params(); + let penalty = PenaltyConfig::default(); + let y = vec![1.0, 0.0, 0.0, 1.0]; + let mask = vec![true, true, true, false]; + let expected = neg_loglik_and_grad(&y, Some(&mask), &[0, 0], &p, &cfg, &penalty); + for device in [Device::Gpu, Device::Auto] { + let (obj, grad, ll) = + finish_device(device, None, &y, Some(&mask), &[0, 0], &p, &cfg, &penalty); + assert_eq!(obj, expected.0); + assert_eq!(ll, expected.2); + assert_eq!(grad.tau, expected.1.tau); + } + } + #[test] fn mirt_ignores_latent_space_terms() { let mut cfg = config(); diff --git a/docs/papers/README.md b/docs/papers/README.md new file mode 100644 index 000000000..d77880f4e --- /dev/null +++ b/docs/papers/README.md @@ -0,0 +1,22 @@ +# Reference Papers + +Curated references that ground the estimation core. Store compact citations and +canonical links here instead of vendoring PDF copies into the repository. + +## Wu et al. 2021 + +Wu, M., Davis, R. L., Domingue, B. W., Piech, C., & Goodman, N. (2021). +*Modeling Item Response Theory with Stochastic Variational Inference.* +arXiv:2108.11579. https://arxiv.org/abs/2108.11579 + +- **License:** Creative Commons Attribution 4.0 International (CC BY 4.0), + https://creativecommons.org/licenses/by/4.0/. +- **Why it is referenced:** it develops a *fast, scalable* estimator for item + response theory by mapping the likelihood and its gradient onto + data-parallel, accelerator-friendly computation (amortized/stochastic + variational inference). That is exactly the numerical shape this project + accelerates: the penalized negative log-likelihood and gradient hot path of + the MLSIRM/MLS2PLM family, now offloadable to the GPU via the wgpu GPGPU + kernels in `crates/mlsirm-core/src/gpu.rs`. The paper is the design reference + for keeping GPU-accelerated IRT estimation numerically faithful to the CPU + objective. diff --git a/python/fast_mlsirm/backend.py b/python/fast_mlsirm/backend.py index e67b80824..f5ecaa4a2 100644 --- a/python/fast_mlsirm/backend.py +++ b/python/fast_mlsirm/backend.py @@ -6,6 +6,11 @@ VALID_BACKENDS = {"numpy", "rust", "auto"} +# Execution device for the Rust backend. This is a sub-option of the ``rust`` +# backend (CPU vs. wgpu GPGPU), NOT a separate compute-backend axis: the single +# backend axis stays {numpy, rust, auto}. ``auto``/``gpu`` run the GPGPU kernels +# when a GPU is present and fall back to the identical CPU path otherwise. +VALID_DEVICES = {"cpu", "gpu", "auto"} CORE_MODULE = "fast_mlsirm._core" @@ -16,6 +21,13 @@ def normalize_backend(name: str) -> str: return backend +def normalize_device(name: str) -> str: + device = str(name).strip().lower() + if device not in VALID_DEVICES: + raise ValueError(f"rust_device must be one of {sorted(VALID_DEVICES)}") + return device + + def resolve_backend(name: str) -> str: backend = normalize_backend(name) if backend == "numpy": diff --git a/python/fast_mlsirm/cli.py b/python/fast_mlsirm/cli.py index c95e98bc6..b8c9984d8 100644 --- a/python/fast_mlsirm/cli.py +++ b/python/fast_mlsirm/cli.py @@ -101,6 +101,7 @@ def _main(argv: list[str] | None = None) -> int: fit_cmd.add_argument("--n-restarts", type=int, default=1, help="Number of random restarts (default: 1).") fit_cmd.add_argument("--seed", type=int, default=1, help="Random seed for fitting (default: 1).") fit_cmd.add_argument("--backend", choices=["numpy", "rust", "auto"], default="auto", help="Objective backend to use (default: auto = Rust core when available, numpy reference fallback).") + fit_cmd.add_argument("--rust-device", choices=["auto", "cpu", "gpu"], default="auto", help="Execution device for the rust backend: wgpu GPGPU when available, else CPU fallback (default: auto). Ignored for the numpy backend.") fit_cmd.add_argument("--out", required=True, help="Directory path to save the fitted parameters.") _add_json_flag(fit_cmd) @@ -495,6 +496,7 @@ def _main(argv: list[str] | None = None) -> int: n_restarts=args.n_restarts, seed=args.seed, backend=args.backend, + rust_device=args.rust_device, ), ) except ValueError as e: @@ -514,6 +516,7 @@ def _main(argv: list[str] | None = None) -> int: "model": result.model, "optimizer": result.optimizer, "backend": result.backend, + "rust_device": result.rust_device, "objective": float(result.objective), "convergence_status": result.convergence_status, "n_iter": int(result.n_iter), diff --git a/python/fast_mlsirm/config.py b/python/fast_mlsirm/config.py index b65584d06..cbfdd8d99 100644 --- a/python/fast_mlsirm/config.py +++ b/python/fast_mlsirm/config.py @@ -2,7 +2,7 @@ from dataclasses import dataclass -from .backend import normalize_backend +from .backend import normalize_backend, normalize_device VALID_MODELS = {"MIRT", "MLS2PLM", "MLSRM", "ULS2PLM", "ULSRM"} @@ -77,6 +77,11 @@ class FitConfig: # ``fast_mlsirm._core`` (Rust/PyO3) kernel when available and transparently # falls back to the pure-numpy reference implementation otherwise. backend: str = "auto" + # Device for the Rust backend: "cpu", "gpu", or "auto". A sub-option of the + # rust backend, not a separate compute-backend axis. "auto" (default) uses + # the wgpu GPGPU kernels when a GPU is available and otherwise falls back to + # the identical CPU path. Ignored when backend == "numpy". + rust_device: str = "auto" penalty: PenaltyConfig = PenaltyConfig() def normalized_model(self) -> str: @@ -103,3 +108,4 @@ def validate(self) -> None: if self.eps_distance <= 0: raise ValueError("eps_distance must be > 0") normalize_backend(self.backend) + normalize_device(self.rust_device) diff --git a/python/fast_mlsirm/fit.py b/python/fast_mlsirm/fit.py index 297059842..e185c667e 100644 --- a/python/fast_mlsirm/fit.py +++ b/python/fast_mlsirm/fit.py @@ -4,7 +4,7 @@ import numpy as np -from .backend import resolve_backend +from .backend import normalize_device, resolve_backend from .config import FitConfig from .math import logit, normalize_latent_positions, standardize from .objective import (model_flags, neg_loglik_and_grad, prepare_response, @@ -21,6 +21,8 @@ def fit( config = config or FitConfig() config.validate() backend = resolve_backend(config.backend) + # The device is a sub-option of the rust backend; the numpy path ignores it. + device = normalize_device(config.rust_device) if backend == "rust" else "cpu" model = config.normalized_model() y, observed = prepare_response(responses, mask) @@ -50,7 +52,7 @@ def fit( best: FitResult | None = None for restart in range(config.n_restarts): candidate = _run_single_fit( - y, observed, factors, n_dims, config, model, restart, backend + y, observed, factors, n_dims, config, model, restart, backend, device ) if best is None or candidate.objective < best.objective: best = candidate @@ -125,6 +127,7 @@ def _fit_mmle( model=model, optimizer=f"mmle_em/{'rust' if rust is not None else 'numpy'}", backend=config.backend, + rust_device=config.rust_device, objective=float(-loglik_trace[-1]) if loglik_trace else float("nan"), loglik_trace=[float(v) for v in loglik_trace], objective_trace=[float(-v) for v in loglik_trace], @@ -142,13 +145,14 @@ def _run_single_fit( model: str, restart: int, backend: str, + device: str, ) -> FitResult: rng = np.random.default_rng(config.seed + restart) params0 = _initial_params( y, observed, factors, n_dims, config.latent_dim, config, rng ) x0 = _pack(params0, model) - objective = _make_objective(y, observed, factors, params0, config, backend) + objective = _make_objective(y, observed, factors, params0, config, backend, device) x = x0 obj_trace: list[float] = [] @@ -182,7 +186,7 @@ def _run_single_fit( if model != "MIRT": final_params = normalize_latent_positions(final_params) final_obj, _, final_loglik = neg_loglik_and_grad( - y, factors, final_params, config, mask=observed, backend=backend + y, factors, final_params, config, mask=observed, backend=backend, device=device ) obj_trace.append(final_obj) loglik_trace.append(final_loglik) @@ -192,6 +196,7 @@ def _run_single_fit( model=model, optimizer=config.optimizer, backend=backend, + rust_device=device, objective=final_obj, loglik_trace=loglik_trace, objective_trace=obj_trace, @@ -240,13 +245,14 @@ def _make_objective( template: MLSIRMParams, config: FitConfig, backend: str, + device: str, ) -> Callable[[np.ndarray], tuple[float, np.ndarray, float]]: model = config.normalized_model() def objective(x: np.ndarray) -> tuple[float, np.ndarray, float]: params = _unpack(x, template, model) obj, grad, loglik = neg_loglik_and_grad( - y, factor_id, params, config, mask=observed, backend=backend + y, factor_id, params, config, mask=observed, backend=backend, device=device ) grad_vec = _pack(grad, model) if config.gradient_clip is not None: diff --git a/python/fast_mlsirm/io.py b/python/fast_mlsirm/io.py index 928f5eb15..1e0ce1dd8 100644 --- a/python/fast_mlsirm/io.py +++ b/python/fast_mlsirm/io.py @@ -55,6 +55,7 @@ def save_fit_result(result: FitResult, run_dir: str | Path) -> None: "model": result.model, "optimizer": result.optimizer, "backend": result.backend, + "rust_device": result.rust_device, "objective": result.objective, "convergence_status": result.convergence_status, "n_iter": result.n_iter, diff --git a/python/fast_mlsirm/objective.py b/python/fast_mlsirm/objective.py index 11f364ae0..d69772c0b 100644 --- a/python/fast_mlsirm/objective.py +++ b/python/fast_mlsirm/objective.py @@ -2,7 +2,7 @@ import numpy as np -from .backend import load_rust_core, normalize_backend, resolve_backend +from .backend import load_rust_core, normalize_backend, normalize_device, resolve_backend from .config import FitConfig, PenaltyConfig from .math import sigmoid, softplus from .types import MLSIRMParams @@ -80,12 +80,14 @@ def neg_loglik_and_grad( config: FitConfig | None = None, mask: np.ndarray | None = None, backend: str = "numpy", + device: str | None = None, ) -> tuple[float, MLSIRMParams, float]: config = config or FitConfig() requested_backend = normalize_backend(backend) normalized_backend = resolve_backend(requested_backend) if requested_backend == "auto" else requested_backend if normalized_backend == "rust": - return _neg_loglik_and_grad_rust(responses, factor_id, params, config, mask) + resolved_device = normalize_device(device if device is not None else config.rust_device) + return _neg_loglik_and_grad_rust(responses, factor_id, params, config, mask, resolved_device) model = config.normalized_model() penalty = config.penalty @@ -158,6 +160,7 @@ def _neg_loglik_and_grad_rust( params: MLSIRMParams, config: FitConfig, mask: np.ndarray | None, + device: str = "cpu", ) -> tuple[float, MLSIRMParams, float]: model = config.normalized_model() penalty = config.penalty @@ -188,6 +191,7 @@ def _neg_loglik_and_grad_rust( float(penalty.lambda_tau), float(penalty.mu_alpha), float(penalty.mu_tau), + device, ) grads = MLSIRMParams( theta=np.asarray(gradients["theta"], dtype=np.float64).reshape(params.theta.shape), diff --git a/python/fast_mlsirm/types.py b/python/fast_mlsirm/types.py index b929de61a..954755ebe 100644 --- a/python/fast_mlsirm/types.py +++ b/python/fast_mlsirm/types.py @@ -50,6 +50,7 @@ class FitResult: model: str optimizer: str backend: str + rust_device: str objective: float loglik_trace: list[float] objective_trace: list[float] diff --git a/tests/test_cli.py b/tests/test_cli.py index b7a896017..e79e7f0b5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3,6 +3,7 @@ from unittest.mock import patch import numpy as np +import pytest from fast_mlsirm.cli import main @@ -132,6 +133,46 @@ def test_cli_fit_auto_backend_records_resolved_backend(tmp_path, capsys): summary = json.loads((fit_dir / "fit_summary.json").read_text(encoding="utf-8")) assert summary["backend"] == payload["backend"] +def test_cli_fit_rust_device_recorded(tmp_path, capsys): + pytest.importorskip("fast_mlsirm._core") + sim_dir = tmp_path / "sim_out" + fit_dir = tmp_path / "fit_out" + + with patch.object(sys, 'argv', ['fast-mlsirm', 'simulate', '--persons', '10', '--dims', '1', '--items-per-dim', '2', '--out', str(sim_dir)]): + main() + capsys.readouterr() + + args = [ + "fit", + "--responses", + str(sim_dir / "responses.npy"), + "--factors", + str(sim_dir / "item_factor.csv"), + "--model", + "MLS2PLM", + "--max-iter", + "1", + "--backend", + "rust", + "--rust-device", + "gpu", + "--out", + str(fit_dir), + "--json", + ] + + with patch.object(sys, 'argv', ['fast-mlsirm'] + args): + assert main() == 0 + + payload = json.loads(capsys.readouterr().out) + assert payload["backend"] == "rust" + # "gpu" is honored when a GPU is present and otherwise falls back to CPU; + # either way the requested device is recorded on the result and summary. + assert payload["rust_device"] == "gpu" + summary = json.loads((fit_dir / "fit_summary.json").read_text(encoding="utf-8")) + assert summary["rust_device"] == "gpu" + + def test_cli_diagnose_fit_success(tmp_path): sim_dir = tmp_path / "sim_out" fit_dir = tmp_path / "fit_out" diff --git a/tests/test_config.py b/tests/test_config.py index 85307fabc..6a810ac70 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -13,6 +13,17 @@ def test_fitconfig_valid(): # Rust is the primary numeric path; the default backend is "auto" which # resolves to the Rust core when available and falls back to numpy. assert config.backend == "auto" + assert config.rust_device == "auto" + + +def test_fitconfig_invalid_rust_device(): + with pytest.raises(ValueError, match="rust_device must be one of"): + FitConfig(rust_device="cuda").validate() + + +def test_fitconfig_accepts_gpu_rust_device(): + for device in ("cpu", "gpu", "auto"): + FitConfig(rust_device=device).validate() def test_mls2plmconfig_invalid_n_persons(): with pytest.raises(ValueError, match="n_persons must be >= 1"): diff --git a/tests/test_objective.py b/tests/test_objective.py index d462dd36c..eaae79c65 100644 --- a/tests/test_objective.py +++ b/tests/test_objective.py @@ -287,6 +287,46 @@ def test_rust_backend_matches_numpy_objective(): assert np.isclose(rust_grad.tau, numpy_grad.tau) +@pytest.mark.parametrize("device", ["cpu", "gpu", "auto"]) +def test_rust_gpgpu_device_matches_numpy(device): + """The rust backend's device paths (including the GPGPU code) must agree + with the numpy reference. On CI without a GPU, ``gpu``/``auto`` fall back to + the f64 CPU implementation; on a GPU they run the f32 wgpu kernels, so the + tolerance is chosen to cover single-precision agreement as well.""" + pytest.importorskip("fast_mlsirm._core") + rng = np.random.default_rng(7) + n_persons, n_items, latent_dim = 12, 6, 2 + params = MLSIRMParams( + theta=rng.normal(size=(n_persons, 1)), + alpha=rng.normal(scale=0.2, size=n_items), + b=rng.normal(scale=0.3, size=n_items), + xi=rng.normal(scale=0.5, size=(n_persons, latent_dim)), + zeta=rng.normal(scale=0.5, size=(n_items, latent_dim)), + tau=0.15, + ) + y = (rng.random((n_persons, n_items)) < 0.5).astype(float) + mask = np.ones((n_persons, n_items), dtype=bool) + mask[0, 0] = False + factors = np.zeros(n_items, dtype=np.int64) + config = FitConfig(max_iter=1) + + numpy_obj, numpy_grad, numpy_loglik = neg_loglik_and_grad( + y, factors, params, config, mask=mask, backend="numpy" + ) + rust_obj, rust_grad, rust_loglik = neg_loglik_and_grad( + y, factors, params, config, mask=mask, backend="rust", device=device + ) + + assert np.isclose(rust_obj, numpy_obj, rtol=1e-4, atol=1e-4) + assert np.isclose(rust_loglik, numpy_loglik, rtol=1e-4, atol=1e-4) + assert np.allclose(rust_grad.theta, numpy_grad.theta, rtol=1e-4, atol=1e-4) + assert np.allclose(rust_grad.alpha, numpy_grad.alpha, rtol=1e-4, atol=1e-4) + assert np.allclose(rust_grad.b, numpy_grad.b, rtol=1e-4, atol=1e-4) + assert np.allclose(rust_grad.xi, numpy_grad.xi, rtol=1e-4, atol=1e-4) + assert np.allclose(rust_grad.zeta, numpy_grad.zeta, rtol=1e-4, atol=1e-4) + assert np.isclose(rust_grad.tau, numpy_grad.tau, rtol=1e-4, atol=1e-4) + + def test_rust_core_rejects_shape_mismatch(): pytest.importorskip("fast_mlsirm._core") from fast_mlsirm import _core