Correct non ConversionRewriter transformation in LegalizeToLinalg

PiperOrigin-RevId: 347622657 Change-Id: Id171aa6454f72cc0a57eee14953b1354d62c162f
2020-12-15 08:56:01 -08:00 · 2020-12-15 08:56:01 -08:00 · ad82a249b4
commit ad82a249b4
parent 282c83de94
1 changed files with 2 additions and 1 deletions
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
@ -813,7 +813,8 @@ class ReduceConverter : public OpConversionPattern<lmhlo::ReduceOp> {
        loc, /*resultTensorTypes=*/ArrayRef<Type>{},
        /*inputs=*/adaptor.operands(), /*outputBuffers=*/adaptor.out(),
        /*initTensors=*/ValueRange{}, maps, types);
-    linalg_op.region().takeBody(reduce_op.body());
+    rewriter.inlineRegionBefore(reduce_op.body(), linalg_op.region(),
+                                linalg_op.region().end());
    {
      OpBuilder::InsertionGuard region_guard(rewriter);
      Block* block = linalg_op.getBody();