CuDNNHLOOpt: give every fusion function an unused symbol name

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();

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]+}}]}