diff --git a/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp b/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp index 0389e0803c..72bd625052 100644 --- a/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/CuDNNHLOOpt.cpp @@ -86,11 +86,24 @@ 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. + // + // 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); + 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); - 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..1bcc53d32c 100644 --- a/test/lit_tests/cudnn/fusion.mlir +++ b/test/lit_tests/cudnn/fusion.mlir @@ -36,3 +36,33 @@ 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: 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: }