diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 093b39d032..c4fd054a75 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -21918,6 +21918,10 @@ struct CommonCompareExpressionRewrite rewriter.replaceOp(op, negatedCondition); return success(); } else { + // The negation reads op's result, so it must sit after op — the + // rewriter's insertion point is before op here. + PatternRewriter::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(user); auto negatedCondition = stablehlo::NotOp::create( rewriter, userCompareOp.getLoc(), op.getResult()); rewriter.replaceOp(user, negatedCondition); diff --git a/test/lit_tests/commoncompare_neg_after.mlir b/test/lit_tests/commoncompare_neg_after.mlir new file mode 100644 index 0000000000..98592fd3f2 --- /dev/null +++ b/test/lit_tests/commoncompare_neg_after.mlir @@ -0,0 +1,15 @@ +// RUN: enzymexlamlir-opt %s --enzyme-hlo-opt | FileCheck %s + +// The negated twin appears after the matched compare: the rewrite replaces +// it with a negation of the earlier compare, which must be inserted after +// the compare it reads, not at the match point before it. +func.func @negpair(%a: tensor<4xi32>, %b: tensor<4xi32>) -> (tensor<4xi1>, tensor<4xi1>) { + %le = stablehlo.compare LE, %a, %b, SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> + %gt = stablehlo.compare GT, %a, %b, SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> + return %le, %gt : tensor<4xi1>, tensor<4xi1> +} + +// CHECK-LABEL: func.func @negpair( +// CHECK-NEXT: %[[LE:.+]] = stablehlo.compare LE, %arg0, %arg1, SIGNED +// CHECK-NEXT: %[[NOT:.+]] = stablehlo.not %[[LE]] +// CHECK-NEXT: return %[[LE]], %[[NOT]]