EnzymeAD / EnzymeAD/Enzyme-JAX

rev diff of `stablehlo.dynamic_update_slice` is broken for complex tensors

Open
#1,211 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.