EnzymeAD / EnzymeAD/Enzyme-JAX

cudnn-hlo-opt fuses dot_generals with no contracting dimensions, which cuDNN cannot compile

Open
#3,058 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.