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(op)) { SinkToRegion(&if_op.true_branch()); SinkToRegion(&if_op.false_branch()); + } else if (auto reduce_window_op = llvm::dyn_cast(op)) { + SinkToRegion(&reduce_window_op.body()); } else if (auto sort_op = llvm::dyn_cast(op)) { SinkToRegion(&sort_op.comparator()); }