From 6d7812bbc5e8b63fdc09870363ef84146eb57a9a Mon Sep 17 00:00:00 2001
From: Tres Popp <tpopp@google.com>
Date: Tue, 15 Dec 2020 08:56:01 -0800
Subject: [PATCH] Correct non ConversionRewriter transformation in
 LegalizeToLinalg

PiperOrigin-RevId: 347622657
---
 lib/Dialect/mhlo/transforms/legalize_to_linalg.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
index ad4fe9a..cba8a52 100644
--- a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
+++ b/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();