Fix reshape concat with convert

Author: wsmoses

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -849,9 +849,9 @@ struct ConcatAppendingReshape final
     size_t frontSize = 0;
     for (auto v : op.getOperands()) {
       if (auto t = v.getDefiningOp<stablehlo::ConvertOp>()) {
-        converts.push_back(
-            t.getType().cast<RankedTensorType>().getElementType());
         v = t.getOperand();
+        converts.push_back(
+            v.getType().cast<RankedTensorType>().getElementType());
       } else
         converts.push_back(nullptr);
       if (auto t = v.getDefiningOp<stablehlo::ReshapeOp>()) {
@@ -906,11 +906,8 @@ struct ConcatAppendingReshape final
         lhs2);
 
     if (typeconvert)
-      res2 = rewriter.create<stablehlo::ConvertOp>(
-          op.getLoc(),
-          RankedTensorType::get(
-              res2.getType().cast<RankedTensorType>().getShape(), typeconvert),
-          res2);
+      res2 = rewriter.create<stablehlo::ConvertOp>(op.getLoc(), op.getType(),
+                                                   res2);
 
     rewriter.replaceOp(op, res2);
     return success();

test/lit_tests/concatreshape.mlir

@@ -9,8 +9,17 @@ module {
     return %concat : tensor<2x3x4xf32>
   }
 
+  func.func @main2(%a : tensor<3x4xf32>, %b : tensor<3x4xf32>) -> tensor<2x3x4xf64> {
+    %u = stablehlo.reshape %a : (tensor<3x4xf32>) -> tensor<1x3x4xf32>
+    %uc = stablehlo.convert %u : (tensor<1x3x4xf32>) -> tensor<1x3x4xf64>
+    %v = stablehlo.reshape %b : (tensor<3x4xf32>) -> tensor<1x3x4xf32>
+    %vc = stablehlo.convert %v : (tensor<1x3x4xf32>) -> tensor<1x3x4xf64>
+    %concat = stablehlo.concatenate %uc, %vc, dim=0 : (tensor<1x3x4xf64>, tensor<1x3x4xf64>) -> tensor<2x3x4xf64>
+    return %concat : tensor<2x3x4xf64>
+  }
+
   // TODO this opt
-  func.func @main2(%a : tensor<3x4xf32>, %b : tensor<3x4xf32>) -> tensor<3x2x4xf32> {
+  func.func @main3(%a : tensor<3x4xf32>, %b : tensor<3x4xf32>) -> tensor<3x2x4xf32> {
     %u = stablehlo.reshape %a : (tensor<3x4xf32>) -> tensor<3x1x4xf32>
     %v = stablehlo.reshape %b : (tensor<3x4xf32>) -> tensor<3x1x4xf32>
     %concat = stablehlo.concatenate %u, %v, dim=1 : (tensor<3x1x4xf32>, tensor<3x1x4xf32>) -> tensor<3x2x4xf32>
@@ -24,3 +33,11 @@ module {
 // CHECK-NEXT:    %1 = stablehlo.reshape %0 : (tensor<6x4xf32>) -> tensor<2x3x4xf32>
 // CHECK-NEXT:    return %1 : tensor<2x3x4xf32>
 // CHECK-NEXT:  }
+
+// CHECK:  func.func @main2(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x3x4xf64> {
+// CHECK-NEXT:    %0 = stablehlo.convert %arg0 : (tensor<3x4xf32>) -> tensor<3x4xf64>
+// CHECK-NEXT:    %1 = stablehlo.convert %arg1 : (tensor<3x4xf32>) -> tensor<3x4xf64>
+// CHECK-NEXT:    %2 = stablehlo.concatenate %0, %1, dim = 0 : (tensor<3x4xf64>, tensor<3x4xf64>) -> tensor<6x4xf64>
+// CHECK-NEXT:    %3 = stablehlo.reshape %2 : (tensor<6x4xf64>) -> tensor<2x3x4xf64>
+// CHECK-NEXT:    return %3 : tensor<2x3x4xf64>
+// CHECK-NEXT:  }