EnzymeAD / rust

A rust fork to work towards Enzyme integration
https://www.rust-lang.org
Other
53 stars 7 forks source link

examples/neohookean.rs: small linear algebra and math of hyperelasticity #31

Closed jedbrown closed 3 months ago

jedbrown commented 1 year ago

This checks several identities. The nested forward mode is failing to compile currently. The matrix calculus and symmetries are documented in Ratel and appear (sometimes indirectly) in textbooks on nonlinear mechanics, such as Wriggers or Holzapfel.

https://ratel.micromorph.org/doc/modeling/models/neo-hookean/

jedbrown commented 1 year ago

This may be a simple usage thing, but I get an error if I try to move the constitutive model into a method on the material instead of a free function as now.

impl NH {
    #[autodiff(d_psi, Reverse, Active)]
    fn psi(&self, #[dup] e: &KM) -> f64 {
        let mu = self.mu;
        let lambda = self.lambda;
        let J = e.cauchy_green().det().sqrt();
        0.25 * lambda * (J * J - 1.0 - 2.0 * J.ln()) // convex
      + mu * (e.trace() - J.ln())
    }

generates

error[E0308]: mismatched types
   --> examples/neohookean.rs:228:5
    |
228 |     #[autodiff(d_psi, Reverse, Active)]
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected struct `KM`, found struct `NH`
    |
    = note: expected reference `&KM`
               found reference `&NH`
    = note: this error originates in the attribute macro `autodiff` (in Nightly builds, run with -Z macro-backtrace for more info)
jedbrown commented 1 year ago

Note: I'd like to be able to make an ElasticMaterial trait with required methods for stress and d_stress and have enzyme provide those inside impl ElasticMaterial for NH. Perhaps I should try the other syntax or a change to the macro is needed. But all in all, this was pretty painless.

jedbrown commented 1 year ago

I added testing stubs in all the examples and retargeted to branch rebase. This is the current output from cargo +enzyme nextest run --examples. image

@bytesnake Feel free to use this as you're fleshing out functionality in the rebase. I'm also happy to rebase and re-run this branch when you think you've made progress.

ZuseZ4 commented 10 months ago

@jedbrown I just pushed an update which allows using env variables for debugging. https://github.com/EnzymeAD/rust#enzyme-config

For your biquad.rs however it is not printing any information which means that we currently don't forward the function to Enzyme. I'll be gone over the weekend, but have another look afterwards.

jedbrown commented 9 months ago

I rebased this branch and tried two variants of forward-mode syntax in the sin demo. Each works independently, but when both are active (as in the commit), we get this panic:

$ cargo +enzyme run --example=sin --release
   Compiling autodiff v0.1.0 (/home/jed/src/rust-enzyme/library/autodiff)
thread '<unnamed>' panicked at 'called `Option::unwrap()` on a `None` value', compiler/rustc_codegen_llvm/src/back/write.rs:655:79
stack backtrace:
   0:     0x7fec302cdf72 - <std::sys_common::backtrace::_print::DisplayBacktrace as core::fmt::Display>::fmt::h703f56af3e3a10d3
   1:     0x7fec3035c03f - core::fmt::write::ha3889d0a7e6d59ce
   2:     0x7fec302e6d35 - std::io::Write::write_fmt::h7f649fa7294c887e
   3:     0x7fec302cdd35 - std::sys_common::backtrace::print::he0f0d8320645afde
   4:     0x7fec302dcc1e - std::panicking::default_hook::{{closure}}::h35c5cff5de4a449b
   5:     0x7fec302dc8d9 - std::panicking::default_hook::hde5044267d1c5297
   6:     0x7fec30b8332c - rustc_driver_impl[21666f9b46a8347f]::install_ice_hook::{closure#0}
   7:     0x7fec302dd437 - std::panicking::rust_panic_with_hook::ha4cff8c4fbe47b3f
   8:     0x7fec302ce262 - std::panicking::begin_panic_handler::{{closure}}::h79813a75119e03c4
   9:     0x7fec302ce086 - std::sys_common::backtrace::__rust_end_short_backtrace::h0da6022a2a159ccd
  10:     0x7fec302dcfa2 - rust_begin_unwind
  11:     0x7fec302b0bc3 - core::panicking::panic_fmt::h4f93a341da2498c8
  12:     0x7fec302b0c5d - core::panicking::panic::heb3f6474ea499c12
  13:     0x7fec30dc404c - rustc_codegen_llvm[60f7880173d7aea7]::back::write::differentiate
  14:     0x7fec30d113e8 - rustc_codegen_ssa[7fa8a61fe150c5a6]::back::write::generate_lto_work::<rustc_codegen_llvm[60f7880173d7aea7]::LlvmCodegenBackend>
  15:     0x7fec30defdb4 - std[b8b380443920fe86]::sys_common::backtrace::__rust_begin_short_backtrace::<<rustc_codegen_llvm[60f7880173d7aea7]::LlvmCodegenBackend as rustc_codegen_ssa[7fa8a61f
e150c5a6]::traits::backend::ExtraBackendMethods>::spawn_thread<rustc_codegen_ssa[7fa8a61fe150c5a6]::back::write::start_executing_work<rustc_codegen_llvm[60f7880173d7aea7]::LlvmCodegenBackend>
::{closure#4}, core[ba18fb509857a0d3]::result::Result<rustc_codegen_ssa[7fa8a61fe150c5a6]::back::write::CompiledModules, ()>>::{closure#0}, core[ba18fb509857a0d3]::result::Result<rustc_codege
n_ssa[7fa8a61fe150c5a6]::back::write::CompiledModules, ()>>
  16:     0x7fec30df4c07 - <<std[b8b380443920fe86]::thread::Builder>::spawn_unchecked_<<rustc_codegen_llvm[60f7880173d7aea7]::LlvmCodegenBackend as rustc_codegen_ssa[7fa8a61fe150c5a6]::traits
::backend::ExtraBackendMethods>::spawn_thread<rustc_codegen_ssa[7fa8a61fe150c5a6]::back::write::start_executing_work<rustc_codegen_llvm[60f7880173d7aea7]::LlvmCodegenBackend>::{closure#4}, core[ba18fb509857a0d3]::result::Result<rustc_codegen_ssa[7fa8a61fe150c5a6]::back::write::CompiledModules, ()>>::{closure#0}, core[ba18fb509857a0d3]::result::Result<rustc_codegen_ssa[7fa8a61fe15
0c5a6]::back::write::CompiledModules, ()>>::{closure#1} as core[ba18fb509857a0d3]::ops::function::FnOnce<()>>::call_once::{shim:vtable#0}
  17:     0x7fec30312f85 - std::sys::unix::thread::Thread::new::thread_start::h7f5f13471f2522bc
  18:     0x7fec300eb9eb - <unknown>
  19:     0x7fec3016f7cc - <unknown>
  20:                0x0 - <unknown>

error: the compiler unexpectedly panicked. this is a bug.
jedbrown commented 8 months ago

Rebased. My example runs with --release, but one of the unit tests triggers an Enzyme error.

ZuseZ4 commented 8 months ago

Sounds good. For Enzyme errors you can just use ENZYME_PRINT_MOD and put the LLVMModule together with the Rust testcase into a new gh issue: https://github.com/EnzymeAD/rust#enzyme-config

ZuseZ4 commented 5 months ago

@jedbrown Can you rebase it please when you have more time at some point, so I can check if it fails Rust or Enzyme?

ZuseZ4 commented 3 months ago

@jedbrown Can we close this? I think we have it merged and working in the other repo, right?

jedbrown commented 3 months ago

Closing since further activity will be in https://github.com/EnzymeAD/rustbook