EnzymeAD / EnzymeAD/Enzyme-JAX
cudnn-hlo-opt: fusion functions collide on __cudnn_fused_elementwise_dot_0 (per-template static counter)
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
## Summary
`enzymexla-cudnn-hlo-opt` names its outlined fusion functions from a `static int
fusionCounter` inside `DotGeneralElementwiseToCuDNNFusion`. That is a class template, so
each `ElementwiseOpTy` instantiation gets its own counter and they all start at 0. A module
containing two fusions of different elementwise kinds -- an add-of-dot and a
multiply-of-dot, say -- gets two functions both named
`__cudnn_fused_elementwise_dot_0` and the pass fails.
## Minimal reproducer
```mlir
// repro.mlir
module {
func.func @main(%a: tensor<4x16x16xbf16>, %b: tensor<4x16x16xbf16>, %c: tensor<4x16x16xbf16>) -> (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) {
%1 = stablehlo.dot_general %a, %b, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
%2 = stablehlo.add %1, %c : tensor<4x16x16xbf16>
%3 = stablehlo.dot_general %a, %b, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
%4 = stablehlo.multiply %3, %c : tensor<4x16x16xbf16>
return %2, %4 : tensor<4x16x16xbf16>, tensor<4x16x16xbf16>
}
}
```
```
enzymexlamlir-opt --enzymexla-cudnn-hlo-opt repro.mlir
```
```
error: redefinition of symbol named '__cudnn_fused_elementwise_dot_0'
note: see current operation:
"func.func"() <{function_type = ..., no_inline, sym_name = "__cudnn_fused_elementwise_dot_0",
sym_visibility = "private"}> ({
^bb0(...):
%0 = "stablehlo.dot_general"(%arg0, %arg1) ...
%1 = "stablehlo.multiply"(%0, %arg2) ...
"func.return"(%1) ...
}) : () -> ()
note: see existing symbol definition here
```
(Run through Reactant.jl against `Reactant_jll` v0.0.405 -- the pass is plain MLIR, no GPU
needed to reproduce.)
## Cause
`CuDNNHLOOpt.cpp`:
```cpp
static int fusionCounter = 0;
std::string fnName = (kCuDNNFusionFuncPrefix + std::to_string(fusionCounter)).str();
auto fnSym = rewriter.getStringAttr(fnName);
fusionCounter++;
```
`patterns.add,
DotGeneralElementwiseToCuDNNFusion,
DotGeneralElementwiseToCuDNNFusion>` -- three instantiations, three
independent counters. The counter is also process-global rather than module-scoped, so the
names depend on how many fusions ran earlier in the process.
## Where it bites
Reactant.jl with `cudnn_hlo_optimize=true` on a training-step gradient graph: the forward
compiles, the gradient graph has both kinds of fused dot and the `cudnn-hlo-opt` pass fails.
(Separately: on bf16 gradient graphs where the pass did run, it fused nothing measurable --
throughput unchanged. Different topic, just noting it.)
PR to follow: look the candidate name up in the module's `SymbolTable` and take the first
free one, which also makes the pass robust against names already present in the input.
## Environment
Reactant.jl main (v0.2.283 source) with `Reactant_jll` v0.0.405, Julia 1.12.7, Linux x86_64.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in CuDNNHLOOpt.cpp, especially DotGeneralElementwiseToCuDNNFusion and the pattern registrations. Run enzymexlamlir-opt --enzymexla-cudnn-hlo-opt repro.mlir to reproduce the duplicate symbol error. Done means fused functions receive unique names across elementwise-operation instantiations and names already present in the input module.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100