EnzymeAD / EnzymeAD/Enzyme-JAX
rev diff of `stablehlo.dynamic_update_slice` is broken for complex tensors
- Dominant language
- MLIR
- Stars
- 133
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
it seems like #1121 broke the rev diff of `stablehlo.dynamic_update_slice` for complex tensors. proof of it is that the VQE benchmark (EnzymeAD/Reactant.jl#1416) worked with Reactant 0.2.138 but it doesn't on more recent versions.
this is the MLIR with `optimize = false` and `optimize = true` for Reactant 0.2.138
[vqe-diff-n30-d1-v0.2.138-opt_false.mlir.txt](https://github.com/user-attachments/files/21533093/vqe-diff-n30-d1-v0.2.138-opt_false.mlir.txt)
[vqe-diff-n30-d1-v0.2.138-opt_true.mlir.txt](https://github.com/user-attachments/files/21533094/vqe-diff-n30-d1-v0.2.138-opt_true.mlir.txt)
this is the erroring IR with `optimize = false` for Reactant 0.2.150
[vqe-diff-n30-d1-v0.2.150-opt_false.mlir.txt](https://github.com/user-attachments/files/21533076/vqe-diff-n30-d1-v0.2.150-opt_false.mlir.txt)
the error seems to be that the tablegen rev diff rule calls `mlir::FloatType::getFloatSemantics` on a complex type, which doesn't implement it?
## stacktrace
```
[16419] signal 11 (2): Segmentation fault: 11
in expression starting at REPL[39]:1
_ZNK4mlir9FloatType17getFloatSemanticsEv at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme15getConstantAttrENS_4TypeEN4llvm9StringRefE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme6detail41ReverseAutoDiffOpInterfaceInterfaceTraits13FallbackModelIN12_GLOBAL__N_133DynamicUpdateSliceOpRevDerivativeEE24createReverseModeAdjointEPKNS2_7ConceptEPNS_9OperationERNS_9OpBuilderEPNS0_21MGradientUtilsReverseEN4llvm11SmallVectorINS_5ValueELj6EEE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme26ReverseAutoDiffOpInterface24createReverseModeAdjointERNS_9OpBuilderEPNS0_21MGradientUtilsReverseEN4llvm11SmallVectorINS_5ValueELj6EEE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme12MEnzymeLogic10visitChildEPNS_9OperationERNS_9OpBuilderEPNS0_21MGradientUtilsReverseE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme12MEnzymeLogic13differentiateEPNS0_21MGradientUtilsReverseERNS_6RegionES5_N4llvm12function_refIFvRNS_9OpBuilderEPNS_5BlockEEEENSt3__18functionIFNSE_4pairINS_5ValueESH_EENS_4TypeEEEE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6enzyme12MEnzymeLogic17CreateReverseDiffENS_19FunctionOpInterfaceENSt3__16vectorI10DIFFE_TYPENS3_9allocatorIS5_EEEES8_RNS0_13MTypeAnalysisENS4_IbNS6_IbEEEESC_14DerivativeModebmNS_4TypeENS0_11MFnTypeInfoESC_PvbN4llvm9StringRefEbb at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN12_GLOBAL__N_117DifferentiatePass16lowerEnzymeCallsERN4mlir21SymbolTableCollectionENS1_19FunctionOpInterfaceE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6detail4walkINS_15ForwardIteratorEEEvPNS_9OperationEN4llvm12function_refIFvS4_EEENS_9WalkOrderE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN12_GLOBAL__N_117DifferentiatePass14runOnOperationEv at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir6detail17OpToOpPassAdaptor3runEPNS_4PassEPNS_9OperationENS_15AnalysisManagerEbj at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir11PassManager9runPassesEPNS_9OperationENS_15AnalysisManagerE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
_ZN4mlir11PassManager3runEPNS_9OperationE at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
mlirPassManagerRunOnOp at /Users/mofeing/.julia/artifacts/b67395cb97d7760a1150283970e04c23249f1230/lib/libReactantExtra.dylib (unknown line)
jlplt_mlirPassManagerRunOnOp_28272 at /Users/mofeing/.julia/compiled/v1.11/Reactant/p9PzF_rBHx1.dylib (unknown line)
mlirPassManagerRunOnOp at /Users/mofeing/Developer/Reactant.jl/src/mlir/libMLIR_h.jl:8876 [inlined]
run! at /Users/mofeing/Developer/Reactant.jl/src/mlir/IR/Pass.jl:154
#run_pass_pipeline!#2 at /Users/mofeing/Developer/Reactant.jl/src/Compiler.jl:1240
```
cc @wsmoses
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the TableGen reverse-diff rule for DynamicUpdateSliceOpRevDerivative and reproduce the failure using the provided Reactant 0.2.150 erroring MLIR, comparing it with the 0.2.138 files. Check the stacktrace path through getConstantAttr and FloatType::getFloatSemantics; done means reverse differentiation of complex tensors no longer crashes.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100