ConvertConcat: don't rebuild the concat-to-batch input (non-terminating rewrite)

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -4180,6 +4180,30 @@ struct ConvertConcat final
     if (!concat)
       return failure();
 
+    // Leave `convert(concat(reshape_i(x_i)))` alone when every operand is a
+    // reshape-like op that inserts the concatenation dimension: that is exactly
+    // the form ConcatInsertDimToBatch produces out of
+    // `concat(reshape_i(convert(x_i)))`, so pushing the converts back into the
+    // operands here (after ElementwiseReshapeLike hoists the reshapes back out)
+    // rebuilds that pattern's input and the greedy rewriter cycles forever.
+    // Nothing is gained by the push in this shape anyway.
+    if (llvm::all_of(concat.getOperands(), [&](Value v) {
+          Operation *defOp = v.getDefiningOp();
+          if (auto reshape = dyn_cast_or_null<stablehlo::ReshapeOp>(defOp))
+            return llvm::is_contained(
+                findReshapeInsertionDims(
+                    cast<RankedTensorType>(reshape.getOperand().getType()),
+                    cast<RankedTensorType>(reshape.getType())),
+                (int64_t)concat.getDimension());
+          if (auto bcast = dyn_cast_or_null<stablehlo::BroadcastInDimOp>(defOp))
+            return stablehlo::OpIsReshapeLike(bcast) &&
+                   !llvm::is_contained(bcast.getBroadcastDimensions(),
+                                       (int64_t)concat.getDimension());
+          return false;
+        }))
+      return rewriter.notifyMatchFailure(
+          op, "concat of dimension-inserting reshapes (batchable form)");
+
     SmallVector<Value> newvals;
     for (auto v : concat.getOperands()) {
       newvals.push_back(stablehlo::ConvertOp::create(

test/lit_tests/convert_concat_insert_dim_cycle.mlir

@@ -0,0 +1,37 @@
+// RUN: enzymexlamlir-opt --transform-interpreter --enzyme-hlo-remove-transform %s | FileCheck %s
+
+// `convert_concat`, `elementwise_reshape_like` and `concat_insert_dim_elementwise`
+// used to form a rewrite cycle on `convert(concat(reshape_i(x_i)))`:
+//   convert_concat                 -> concat(convert(reshape(x_i)))
+//   elementwise_reshape_like       -> concat(reshape(convert(x_i)))
+//   concat_insert_dim_elementwise  -> convert(concat(reshape(x_i)))   (start)
+// and every trip through the batching rewrite left another unbatched wrapper
+// function behind, so the greedy driver never reached a fixed point. The convert
+// is not pushed into the batchable form any more; this module is already at its
+// fixed point and must come back unchanged.
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %0 {
+      transform.apply_patterns.enzyme_hlo.concat_insert_dim_elementwise
+      transform.apply_patterns.enzyme_hlo.elementwise_reshape_like
+      transform.apply_patterns.enzyme_hlo.convert_concat
+    } : !transform.any_op
+    transform.yield
+  }
+  func.func @cycle(%arg0: tensor<16x64x2xf32>, %arg1: tensor<16x64x2xf32>) -> tensor<2x16x64x2xbf16> {
+    %0 = stablehlo.reshape %arg0 : (tensor<16x64x2xf32>) -> tensor<1x16x64x2xf32>
+    %1 = stablehlo.reshape %arg1 : (tensor<16x64x2xf32>) -> tensor<1x16x64x2xf32>
+    %2 = stablehlo.concatenate %0, %1, dim = 0 : (tensor<1x16x64x2xf32>, tensor<1x16x64x2xf32>) -> tensor<2x16x64x2xf32>
+    %3 = stablehlo.convert %2 : (tensor<2x16x64x2xf32>) -> tensor<2x16x64x2xbf16>
+    return %3 : tensor<2x16x64x2xbf16>
+  }
+}
+
+// CHECK:      func.func @cycle(%arg0: tensor<16x64x2xf32>, %arg1: tensor<16x64x2xf32>) -> tensor<2x16x64x2xbf16> {
+// CHECK-NEXT:   %[[R0:.+]] = stablehlo.reshape %arg0 : (tensor<16x64x2xf32>) -> tensor<1x16x64x2xf32>
+// CHECK-NEXT:   %[[R1:.+]] = stablehlo.reshape %arg1 : (tensor<16x64x2xf32>) -> tensor<1x16x64x2xf32>
+// CHECK-NEXT:   %[[C:.+]] = stablehlo.concatenate %[[R0]], %[[R1]], dim = 0 : (tensor<1x16x64x2xf32>, tensor<1x16x64x2xf32>) -> tensor<2x16x64x2xf32>
+// CHECK-NEXT:   %[[CV:.+]] = stablehlo.convert %[[C]] : (tensor<2x16x64x2xf32>) -> tensor<2x16x64x2xbf16>
+// CHECK-NEXT:   return %[[CV]] : tensor<2x16x64x2xbf16>
+// CHECK-NEXT: }