EnzymeAD / oxide-enzyme

Enzyme integration into Rust. Experimental, do not use.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

ODE example failing

ZuseZ4 opened this issue · comments

Using

 20     let fnc_j_0 = FncInfo::new("J_0", "unsafe_d_J0",
 21                              vec![DFT_DUP_ARG, DFT_CONSTANT, DFT_CONSTANT, DFT_CONSTANT,],
 22                              Some((DFT_OUT_DIFF, true)));

one the following test function:

  1 pub fn J_0(t: &mut Theta, y0: State, E: SigmaMat, input_true: FVec) -> f64 {
  2     let mut stepper = Dop853::new(t.clone(), 0.0, 100.0, 1e-3, y0, 1e-4, 1e-4);
  3     let res = stepper.integrate();
  4 
  5     if let Err(e) = res {
  6         panic!("An error occured: {}", e);
  7     }
  8     let rand_out = stepper.y_out().to_vec();
  9     let tmp = rand_out.iter().sum::<State>() / (rand_out.len() as f64);
 10     let input = f(tmp);
 11 
 12     let inner: FVec = input - input_true;
 13     let res: SMatrix<f64,1,1> = 0.5 * inner.transpose() * E * inner;
 14 
 15     res[0]
 16 }

will lead to Enzyme failing with the following message:

   %_2.i.i6.i.i = inttoptr i64 %new_layout.1 to i8*, !dbg !5582: {[-1]:Integer}, intvals: {}
    %20 = tail call i8* @__rust_alloc(i64 %new_layout.0, i64 %new_layout.1) #34, !dbg !5587: {}, intvals: {}
    %21 = tail call i8* @__rust_realloc(i8* nonnull %7, i64 %9, i64 %new_layout.1, i64 %new_layout.0) #34, !dbg !5601: {}, intvals: {}
  i8 48: {[-1]:Integer}, intvals: {48,}
  i8 127: {[-1]:Integer}, intvals: {127,}
  i8 49: {[-1]:Integer}, intvals: {49,}
  i8 53: {[-1]:Integer}, intvals: {53,}
  i8 40: {[-1]:Integer}, intvals: {40,}
  i8 32: {[-1]:Integer}, intvals: {32,}
  i8 56: {[-1]:Integer}, intvals: {56,}
  i8 46: {[-1]:Integer}, intvals: {46,}
  i8 47: {[-1]:Integer}, intvals: {47,}
  i8 45: {[-1]:Integer}, intvals: {45,}
  i8 99: {[-1]:Integer}, intvals: {99,}
  i8 114: {[-1]:Integer}, intvals: {114,}
  i8 109: {[-1]:Integer}, intvals: {109,}
  i8 101: {[-1]:Integer}, intvals: {101,}
  i8 111: {[-1]:Integer}, intvals: {111,}
  i8 116: {[-1]:Integer}, intvals: {116,}
  i8 117: {[-1]:Integer}, intvals: {117,}
  i8 110: {[-1]:Integer}, intvals: {110,}
  i8 108: {[-1]:Integer}, intvals: {108,}
  i8 105: {[-1]:Integer}, intvals: {105,}
  i8 103: {[-1]:Integer}, intvals: {103,}
  i8 100: {[-1]:Integer}, intvals: {100,}
  i8 115: {[-1]:Integer}, intvals: {115,}
  i8 104: {[-1]:Integer}, intvals: {104,}
  i8 122: {[-1]:Integer}, intvals: {122,}
  i8 121: {[-1]:Integer}, intvals: {121,}
  i8 54: {[-1]:Integer}, intvals: {54,}
  i8 98: {[-1]:Integer}, intvals: {98,}
  i8 120: {[-1]:Integer}, intvals: {120,}
  i8 95: {[-1]:Integer}, intvals: {95,}
  i8 52: {[-1]:Integer}, intvals: {52,}
  i8 107: {[-1]:Integer}, intvals: {107,}
  i8 119: {[-1]:Integer}, intvals: {119,}
  i8 118: {[-1]:Integer}, intvals: {118,}
  i8 -127: {[-1]:Integer}, intvals: {-127,}
  i8 113: {[-1]:Integer}, intvals: {113,}
  i8 96: {[-1]:Integer}, intvals: {96,}
  i8 -74: {[-1]:Integer}, intvals: {-74,}
  i8 -9: {[-1]:Integer}, intvals: {-9,}
  </analysis>
  Illegal updateAnalysis prev:{[-1]:Integer} new: {[-1]:Pointer}
  val:   %_2.i.i.i = inttoptr i64 %new_layout.1 to i8*, !dbg !5488 origin=  %.sroa.0.0.i.i.pn = phi i8* [ %_2.i.i.i, %bb3.i.i14 ], [ %13, %bb7.i.i15 ], [ %_2.i.i6.i.i, %bb3.i7.i.i ], [ %20, %bb7.i8.i.i ], [ %21, %bb10.i.i ]
  build-script-build: ../Enzyme/TypeAnalysis/TypeAnalysis.cpp:617: void TypeAnalyzer::updateAnalysis(llvm::Value*, TypeTree, llvm::Value*): Assertion `0 && "Performed illegal updateAnalysis"' failed.

I'm using an Enzyme fork with the Rust debug parser merged, but the failing code doesn't seem to be affected by that:
https://github.com/ZuseZ4/Enzyme/blob/08eb0646ee6f0ae7eca99fcc41f0c54fcc09cbc1/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp#L617

It could be due to the sloppy initialization of some values which (based on my understanding) shouldn't be seriously considered
by Enzyme: https://github.com/rust-ml/oxide-enzyme/blob/a1dfbde391b32cde6a594b9fe87f7e8b8940a9e2/src/enzyme/enzyme_wrapper.rs#L77

I guess this issue about allocating is outdated, correct? https://github.com/wsmoses/Enzyme.jl/blob/05c9b84f9051c5bd8863925e54a2fc4eb11fc784/src/Enzyme.jl#L126

This is a Type analysis bug, can you post a log, ideally with the EnzymePrintType flag on

Can you run with the following flag set to false: https://github.com/wsmoses/Enzyme/blob/f7c9fe57112bba48caee520b0dec28ec55e48285/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp#L68

It appears there is a potential type confusion as to whether something is an integer vs a pointer on std::option::Option<(std::ptr::NonNull<u8>, std::alloc::Layout)> which makes sense since presumably the ptr is a pointer and the layout is an integer.

Thanks for finding the origin.

Not sure how aware you are about the rust types, so the Option here doesn't read as either ptr or alloc, instead it's either Some((ptr, alloc)), which contains a tuple or None, where I'm not sure how it will be represented. Iirc, for cases where the T in Some can't be / begin with zero, they try to merge both. But that probably won't affect the TypeAnalysis anymore. Alloc is a struct { usize, NonZeroUsize } about the memory size and the alignment.

That's with strict-aliasing=true: https://drive.google.com/file/d/1PTqFNkxx7VyZzR9BZwe0nKq7RPJmn9IU/view?usp=sharing
That's with strict-aliasing=false: https://drive.google.com/file/d/1aRUcfKlT3TQKdeYdrjXCSwe5y3sj4aUS/view?usp=sharing
Unfortunately, I ran out of memory in the second case with 48GB (+50GB Swp). It also looks like it's missing the end.
I guess trying to trigger the issue with some simpler allocations in a smaller function, or starting to mark some rust functions as inactive could help here? What do you think?