EnzymeAD / EnzymeAD/Enzyme-JAX

cudnn-hlo-opt: fusion functions collide on __cudnn_fused_elementwise_dot_0 (per-template static counter)

Open
#2,985 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

## 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.