From e1df9f1d48be48ce76d97d0de35227cfa42ba2e1 Mon Sep 17 00:00:00 2001 From: Asaf Manela Date: Wed, 26 Aug 2026 16:02:02 -0500 Subject: [PATCH 1/2] CuDNNHLOOpt: give every fusion function an unused symbol name The name of the outlined `__cudnn_fused_elementwise_dot_N` function came 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' (hit by `cudnn_hlo_optimize=true` in Reactant.jl on a training-step graph.) Look the candidate name up in the module's SymbolTable instead and take the first one that is free, which is also robust against names already present in the input module. Co-Authored-By: Claude Opus 5 --- src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp | 15 +++++++++++---- test/lit_tests/cudnn/fusion.mlir | 20 ++++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp b/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp index 0389e0803c..9ed9b1769f 100644 --- a/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp @@ -86,11 +86,18 @@ struct DotGeneralElementwiseToCuDNNFusion if (!mod) return rewriter.notifyMatchFailure(elemOp, "No module found"); - static int fusionCounter = 0; - std::string fnName = - (kCuDNNFusionFuncPrefix + std::to_string(fusionCounter)).str(); + // Pick a name that is unused in this module. A counter local to the pattern + // is not enough: it is `static` inside a class template, so every + // ElementwiseOpTy instantiation gets its own and they all restart at 0, + // redefining each other's symbols. + SymbolTable symbolTable(mod); + std::string fnName; + for (unsigned i = 0;; ++i) { + fnName = (kCuDNNFusionFuncPrefix + std::to_string(i)).str(); + if (!symbolTable.lookup(fnName)) + break; + } auto fnSym = rewriter.getStringAttr(fnName); - fusionCounter++; // Input Types auto dotGeneralLhsTy = dotGeneral.getLhs().getType(); diff --git a/test/lit_tests/cudnn/fusion.mlir b/test/lit_tests/cudnn/fusion.mlir index 4c2b7a67cf..55856b6cfa 100644 --- a/test/lit_tests/cudnn/fusion.mlir +++ b/test/lit_tests/cudnn/fusion.mlir @@ -36,3 +36,23 @@ module { // CHECK-NEXT: %0 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_0]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> // CHECK-NEXT: return %0 : tensor<4x16x16xbf16> // CHECK-NEXT: } + +// Two fusions in the same module must get distinct symbols. The counter used to +// be a `static` inside the class template, so the add and multiply +// instantiations both handed out `__cudnn_fused_elementwise_dot_0` and the +// module failed to verify with "redefinition of symbol named ...". +module { + func.func @dense3(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) { + %1 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> + %2 = stablehlo.add %1, %arg2 : tensor<4x16x16xbf16> + %3 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1], precision = [DEFAULT, DEFAULT] : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> + %4 = stablehlo.multiply %3, %arg2 : tensor<4x16x16xbf16> + return %2, %4 : tensor<4x16x16xbf16>, tensor<4x16x16xbf16> + } +} + +// CHECK-DAG: func.func private @__cudnn_fused_elementwise_dot_0(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} { +// CHECK-DAG: func.func private @__cudnn_fused_elementwise_dot_1(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} { +// CHECK: func.func @dense3 +// CHECK: stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_{{[0-9]+}}]} +// CHECK: stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_{{[0-9]+}}]} From 6fcf0605f50ee2953893c2beee6ab579e4518df7 Mon Sep 17 00:00:00 2001 From: Asaf Manela Date: Mon, 31 Aug 2026 11:45:09 -0500 Subject: [PATCH 2/2] CuDNNHLOOpt: use SymbolTable::generateSymbolName; tighten the lit test Review feedback: - Replace the hand-rolled "first free name" loop with SymbolTable::generateSymbolName, which is what SymbolTable::insert itself delegates to. It carries a uniquing counter rather than rescanning from 0 on every match, and emits the same __cudnn_fused_elementwise_dot_0, _1, ... names. insert() was the other candidate, but it uniquifies by appending to the whole existing name, so the second fusion would come out as ..._dot_0_0 rather than ..._dot_1. generateSymbolName appends "_" 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. - Check the dense3 fusion bodies with CHECK-NEXT like the other cases in the file, and name both called_computations explicitly instead of matching {{[0-9]+}}. Verified against a local build: //test/lit_tests/... is 1214/1214, and putting the old static counter back makes cudnn/fusion.mlir fail with "redefinition of symbol named '__cudnn_fused_elementwise_dot_0'". Co-Authored-By: Claude Opus 5 --- src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp | 18 ++++++++++++------ test/lit_tests/cudnn/fusion.mlir | 20 +++++++++++++++----- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp b/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp index 9ed9b1769f..72bd625052 100644 --- a/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp @@ -90,13 +90,19 @@ struct DotGeneralElementwiseToCuDNNFusion // is not enough: it is `static` inside a class template, so every // ElementwiseOpTy instantiation gets its own and they all restart at 0, // redefining each other's symbols. + // + // generateSymbolName appends "_", so it gets the prefix without its + // trailing underscore and keeps the existing _0, _1, ... names. Trimming + // the underscore from kCuDNNFusionFuncPrefix itself would instead break the + // starts_with() guard above. SymbolTable symbolTable(mod); - std::string fnName; - for (unsigned i = 0;; ++i) { - fnName = (kCuDNNFusionFuncPrefix + std::to_string(i)).str(); - if (!symbolTable.lookup(fnName)) - break; - } + unsigned uniquingCounter = 0; + SmallString<128> fnName = SymbolTable::generateSymbolName<128>( + kCuDNNFusionFuncPrefix.drop_back(), + [&](llvm::StringRef candidate) { + return symbolTable.lookup(candidate) != nullptr; + }, + uniquingCounter); auto fnSym = rewriter.getStringAttr(fnName); // Input Types diff --git a/test/lit_tests/cudnn/fusion.mlir b/test/lit_tests/cudnn/fusion.mlir index 55856b6cfa..1bcc53d32c 100644 --- a/test/lit_tests/cudnn/fusion.mlir +++ b/test/lit_tests/cudnn/fusion.mlir @@ -51,8 +51,18 @@ module { } } -// CHECK-DAG: func.func private @__cudnn_fused_elementwise_dot_0(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} { -// CHECK-DAG: func.func private @__cudnn_fused_elementwise_dot_1(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} { -// CHECK: func.func @dense3 -// CHECK: stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_{{[0-9]+}}]} -// CHECK: stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_{{[0-9]+}}]} +// CHECK: func.func private @__cudnn_fused_elementwise_dot_1(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} { +// CHECK-NEXT: %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> +// CHECK-NEXT: %1 = stablehlo.add %0, %arg2 : tensor<4x16x16xbf16> +// CHECK-NEXT: return %1 : tensor<4x16x16xbf16> +// CHECK-NEXT: } +// CHECK-NEXT: func.func private @__cudnn_fused_elementwise_dot_0(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> attributes {no_inline} { +// CHECK-NEXT: %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> +// CHECK-NEXT: %1 = stablehlo.multiply %0, %arg2 : tensor<4x16x16xbf16> +// CHECK-NEXT: return %1 : tensor<4x16x16xbf16> +// CHECK-NEXT: } +// CHECK-NEXT: func.func @dense3(%arg0: tensor<4x16x16xbf16>, %arg1: tensor<4x16x16xbf16>, %arg2: tensor<4x16x16xbf16>) -> (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) { +// CHECK-NEXT: %0 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_1]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> +// CHECK-NEXT: %1 = stablehlo.custom_call @__cudnn$fusion(%arg0, %arg1, %arg2) {called_computations = [@__cudnn_fused_elementwise_dot_0]} : (tensor<4x16x16xbf16>, tensor<4x16x16xbf16>, tensor<4x16x16xbf16>) -> tensor<4x16x16xbf16> +// CHECK-NEXT: return %0, %1 : tensor<4x16x16xbf16>, tensor<4x16x16xbf16> +// CHECK-NEXT: }