diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 093b39d032..bd0e308261 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/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(defOp)) + return llvm::is_contained( + findReshapeInsertionDims( + cast(reshape.getOperand().getType()), + cast(reshape.getType())), + (int64_t)concat.getDimension()); + if (auto bcast = dyn_cast_or_null(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 newvals; for (auto v : concat.getOperands()) { newvals.push_back(stablehlo::ConvertOp::create( diff --git a/test/lit_tests/convert_concat_insert_dim_cycle.mlir b/test/lit_tests/convert_concat_insert_dim_cycle.mlir new file mode 100644 index 0000000000..5a50385380 --- /dev/null +++ b/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: }