From d0901a83e50d3b67a688b11fa25ce36208b77c09 Mon Sep 17 00:00:00 2001
From: Smit Hinsu <hinsu@google.com>
Date: Fri, 13 Nov 2020 13:48:17 -0800
Subject: [PATCH] Enable fallback legalization for MaxPoolGradGrad and
 MaxPool3DGradGrad ops

Requires,
* Override for ReducePrecision in HloMlirBuilder
* Sinking of constants for ReduceWindow op

PiperOrigin-RevId: 342330848
---
 lib/Dialect/mhlo/transforms/sink_constants_to_control_flow.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/lib/Dialect/mhlo/transforms/sink_constants_to_control_flow.cc b/lib/Dialect/mhlo/transforms/sink_constants_to_control_flow.cc
index d863d82..ae57fd5 100644
--- a/lib/Dialect/mhlo/transforms/sink_constants_to_control_flow.cc
+++ b/lib/Dialect/mhlo/transforms/sink_constants_to_control_flow.cc
@@ -50,6 +50,8 @@ class SinkConstantsToControlFlowPass
       } else if (auto if_op = llvm::dyn_cast<IfOp>(op)) {
         SinkToRegion(&if_op.true_branch());
         SinkToRegion(&if_op.false_branch());
+      } else if (auto reduce_window_op = llvm::dyn_cast<ReduceWindowOp>(op)) {
+        SinkToRegion(&reduce_window_op.body());
       } else if (auto sort_op = llvm::dyn_cast<SortOp>(op)) {
         SinkToRegion(&sort_op.comparator());
       }