CuDNNHLOOpt: give every fusion function an unused symbol name

[AsafManela]

Fixes #2985.

The outlined __cudnn_fused_elementwise_dot_N function got its name from a static int fusionCounter inside DotGeneralElementwiseToCuDNNFusion. That is a class template, so
each ElementwiseOpTy instantiation has its own counter and they all start at 0: as soon as
one module contains, say, an add-of-dot and a multiply-of-dot, both fusions are named
__cudnn_fused_elementwise_dot_0 and the pass fails with

error: redefinition of symbol named '__cudnn_fused_elementwise_dot_0'

Look the candidate name up in the module's SymbolTable and take the first free one. That
fixes the collision, drops the process-global state (the old counter also made the names
depend on how many fusions had run earlier in the process), and makes the pass robust
against a name already present in the input module.

Lit test appended to test/lit_tests/cudnn/fusion.mlir. The existing @dense1 / @dense2
cases are unaffected: they live in separate nested modules, so each still starts from _0.

Testing

The C++ is not built or run here -- no bazel in this environment. The reproducer and the
error message come from running the shipped pass (Reactant_jll v0.0.405) on the module in
the issue. The lit test is written against the current sources but has not been executed;
please let CI have it, and tell me if the CHECK lines need adjusting.

🤖 Generated with Claude Code

https://claude.ai/code/session_012k6KZhhCNSCvZFKzErHaHH

[AsafManela]

Update: this is now built and verified, so the "not built or run here" caveat in the
description no longer applies.

I compiled the patch into a JLL and compared it against an otherwise identical build from
the same pinned Enzyme-JAX commit, differing only by this change (confirmed by inspecting
each binary rather than trusting the build labels -- my first attempt at this comparison
was wrong because both bundles turned out to be patched).

Running the reproducer from the issue through enzymexla-cudnn-hlo-opt:

without this patch:  error: redefinition of symbol named '__cudnn_fused_elementwise_dot_0'
with this patch:     OK -- the two fusions get `__cudnn_fused_elementwise_dot_0` and `_1`

A seeded fp32/bf16 smoke test gives bit-identical numerical results on both builds, so the
change is behaviour-preserving outside the failing case.

The lit test in this PR has still not been executed -- I verified through the compiled
pass rather than through enzymexlamlir-opt -- so please let CI have it and tell me if
the CHECK lines need adjusting.

[AsafManela]

Important caveat, found by using this fix rather than just testing it: this patch is
necessary but not sufficient to make cudnn_hlo_optimize=true usable. Please don't
merge it expecting the feature to work end to end.

With the fix in a built JLL, the pass gets past the symbol collision and emits real work:
on a depth-12 bf16 training-step graph it creates __cudnn_fused_elementwise_dot_0..73,
74 fusions. XLA then refuses the result:

INTERNAL: during context [hlo verifier]: Instruction of fused computation does not match
expected instruction %fusion = bf16[768,16384]{1,0} fusion(%convert.514, %reshape.839,
%transpose.695), kind=kCustom, calls=%a_cudnn_fused_elementwise_dot_73.1.clone ...
Failed after computation-deduplicator

So the failure moves from MLIR to XLA. Before this patch the pass aborted at the third
fusion with the duplicate symbol, which masked whatever comes after it.

A hypothesis I have not confirmed, offered only as a lead: many of the 74 outlined bodies
are structurally identical, XLA's computation-deduplicator merges them, and the kCustom
fusion instruction then no longer matches the computation it names (note the .clone
suffix in the error). If that is right, the outlined computations may need to be marked
non-dedupable, or the custom call updated when its callee is merged.

The patch here still stands on its own terms — duplicate symbols are invalid MLIR whatever
happens downstream, and the earlier reproducer in #2985 fails without it — but the feature
needs more than this to work.

Two things I want to be straight about: this measurement is not mine, it came from a
colleague running my built JLL on an H100, and I have verified it only by reading their
log, not by reproducing it. And it contradicts something I wrote earlier in my own notes
("the pass runs but fuses nothing measurable on bf16 graphs") — the "fuses nothing" half
was wrong; it fuses plenty, the fusions just don't survive. Throughput remains untested,
since nothing gets far enough to time.

Happy to open this as a separate issue with the full log if you'd rather track it apart
from this PR.

[Antipath1]

Thank you for the PR it definitely is an improvement. CI is green, and I built the unpatched and patched version of this on the test file and could reproduce the error without this PR, and I made a suggestion that contains the actual symbol names instead of wildcards because I think wildcards are just plain wrong choice in this case. Also, I would feel a lot better about this if you ran enzymexlamlir-opt locally I mean you can build it locally already and running this could have produced a better test as well as running bazel test //test/lit_tests/... on your machine before CI. I also think the shape of the newly added test case is not great because it lackes all the information that the other tests have.
Here is the rest of the information of running enzymexlamlir-opt on my machine

module {
  module {
    func.func private @__cudnn_fused_elementwise_dot_0(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} {
      %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      %1 = stablehlo.add %0, %arg2 : tensor<4x16x16xbf16>
      return %1 : tensor<4x16x16xbf16>
    }
    func.func @dense1(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> {
      %0 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_0]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      return %0 : tensor<4x16x16xbf16>
    }
  }
  module {
    func.func private @__cudnn_fused_elementwise_dot_0(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} {
      %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      %1 = stablehlo.subtract %arg2, %0 : tensor<4x16x16xbf16>
      return %1 : tensor<4x16x16xbf16>
    }
    func.func @dense2(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> {
      %0 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_0]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      return %0 : tensor<4x16x16xbf16>
    }
  }
  module {
    func.func private @__cudnn_fused_elementwise_dot_1(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} {
      %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      %1 = stablehlo.add %0, %arg2 : tensor<4x16x16xbf16>
      return %1 : tensor<4x16x16xbf16>
    }
    func.func private @__cudnn_fused_elementwise_dot_0(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} {
      %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      %1 = stablehlo.multiply %0, %arg2 : tensor<4x16x16xbf16>
      return %1 : tensor<4x16x16xbf16>
    }
    func.func @dense3(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) {
      %0 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_1]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      %1 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_0]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16>
      return %0, %1 : tensor<4x16x16xbf16>, tensor<4x16x16xbf16>
    }
  }
}

that should allow for writing a better test for you.

You can install Bazel via Bazelisk quite easily if you do not have it already
https://bazel.build/install/bazelisk

Happy to open this as a separate issue with the full log if you'd rather track it apart

Yeah sounds like a good idea

I have to say that I had to put your text body in Claude to turn it into something that I wanted to actually read. Most of what you wrote can be significantly condensed and far easier for a human to read. Maybe next time, prompt your LLM to be more terse, relevant and on-topic. I mean, I used an LLM to condense the information, so it is definitely possible.

[AsafManela]

All three taken, and you were right that I should have built this locally first. Done now; new commit pushed.

Wildcards + test shape. dense3 now checks the fusion bodies with CHECK-NEXT like dense1/dense2, and names both called_computations explicitly (_1, then _0). Your pasted output matches mine byte for byte.

SymbolTable::insert. It uniquifies by appending to the whole existing name, so starting from ..._dot_0 the second fusion comes out __cudnn_fused_elementwise_dot_0_0, not ..._dot_1. So I used SymbolTable::generateSymbolName instead — the API insert itself delegates to. Same _0, _1, ... names, and it carries a uniquing counter instead of rescanning from 0 on every match. It appends "_<n>" itself, hence the drop_back() on the prefix: trimming the underscore from kCuDNNFusionFuncPrefix instead would break the starts_with() guard that stops the pass re-fusing its own output.

Local runs (bazel 7.7.0, --config=public_cache):

• bazel test //test/lit_tests/... -> Executed 1213 out of 1214 tests: 1214 tests pass.
• Negative control, old static int fusionCounter put back: cudnn/fusion.mlir FAILS with redefinition of symbol named '__cudnn_fused_elementwise_dot_0'. Restored, PASSES.

The XLA verifier failure is unrelated to this patch; I'll track it separately rather than here.

And point taken on length.

[AsafManela]

Tracked: openxla/xla#48147.

Root cause is in XLA, not here. CuDnnCustomCallConverter turns our __cudnn$fusion custom calls into kCustom fusions pre-SPMD (gpu_compiler.cc:686), and the mutating computation-deduplicator at :1021 then merges the identical fusion computations, leaving two fusion instructions sharing one — which CheckFusionInstruction rejects. Reduced to a 25-line HLO reproducer with no GPU, plus a one-line candidate fix and a regression test that fails without it.

So nothing further is needed on this PR.

[Antipath1]

Thank you a lot. LGTM