Fix concat fuse

EnzymeAD · Mar 10, 2024 · c339731 · c339731
1 parent 3cd0cac
commit c339731
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -578,8 +578,8 @@ struct ConcatFuse final : OpRewritePattern<mlir::stablehlo::ConcatenateOp> {
     }
     if (!changed)
       return failure();
-    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(op, op.getType(),
-                                                          vals);
+    rewriter.replaceOpWithNewOp<stablehlo::ConcatenateOp>(
+        op, op.getType(), vals, op.getDimensionAttr());
     return success();
   }
 };

diff --git a/test/lit_tests/concatfuse.mlir b/test/lit_tests/concatfuse.mlir
@@ -0,0 +1,15 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+
+  func.func @main(%a : tensor<2xf32>, %b : tensor<1xf32>, %c : tensor<1xf32>) -> tensor<4xf32> {
+    %concat = stablehlo.concatenate %a, %b, dim=0 : (tensor<2xf32>, tensor<1xf32>) -> tensor<3xf32>
+    %concat2 = stablehlo.concatenate %concat, %c, dim=0 : (tensor<3xf32>, tensor<1xf32>) -> tensor<4xf32>
+    return %concat2 : tensor<4xf32>
+  }
+}
+
+// CHECK:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>) -> tensor<4xf32> {
+// CHECK-NEXT:    %[[concat:.+]] = stablehlo.concatenate %arg0, %arg1, %arg2, dim = 0 : (tensor<2xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4xf32>
+// CHECK-NEXT:    return %[[concat]] : tensor<4xf32>
+// CHECK-NEXT:  }