EnzymeAD / EnzymeAD/Enzyme-JAX
cudnn-hlo-opt fuses dot_generals with no contracting dimensions, which cuDNN cannot compile
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
`DotGeneralElementwiseToCuDNNFusion` matches any `dot_general` feeding an elementwise op. A `dot_general` with empty `contracting_dims` is a batched outer product, not a matmul, and cuDNN has no candidate for it. The fusion is built happily and the GPU compiler then fails the whole module:
```
INTERNAL: Failed to get configs for: 2 out of 168 instructions. See logs for all failures.
Example failure: ... kind=kCustom, calls=%a_cudnn_fused_elementwise_dot_60.1.clone,
backend_config={... "fusion_backend_config":{"kind":"__cudnn$fusion"} ...}.
Status: INTERNAL: No candidates could be compiled.
```
The failing fused computations look like this — note the empty contracting dims:
```
%dot.0 = f32[2048,8,768]{2,1,0} dot(%Arg_0.0, %Arg_1.0),
lhs_batch_dims={0,1}, lhs_contracting_dims={},
rhs_batch_dims={1,2}, rhs_contracting_dims={}
ROOT %multiply.0 = f32[2048,8,768]{2,1,0} multiply(%dot.0, %Arg_2.0)
```
On a depth-12 bf16 training step, grouped by output shape (the autotuner deduplicates by shape, which is why "2 out of 168 instructions" covers 49 fusions):
| output shape | count | dot | result |
| --- | --- | --- | --- |
| `f32[2048,8,768]` | 25 | `contracting_dims={}` | fails |
| `f32[6,2048,8,128]` | 24 | `contracting_dims={}` | fails |
| `bf16[16384,768]` | 12 | `contracting_dims={1}` | compiles |
Every fusion with a real contraction compiled; every one without failed.
Reproduces with `cudnn_hlo_optimize=true` from Reactant.jl once the symbol collision (#2985) is out of the way; it is the next failure after that one. Fix in a PR.
Contributor guide
No contributing guide indexed for this repository
Research direction
Locate the DotGeneralElementwiseToCuDNNFusion pattern and reproduce the failure with cudnn_hlo_optimize=true from Reactant.jl, after resolving the symbol collision in #2985. Verify that dot_generals with empty contracting_dims are no longer fused, while the real-contraction case still compiles.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100