Move XLA-independent transforms to the new MLIR-HLO directory

This is as straighforward as possible, more cleanup/rewrite to come.

PiperOrigin-RevId: 319849713

include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_COMPILER_MLIR_XLA_IR_CHLO_OPS_H_
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_CHLO_OPS_H_
-#define TENSORFLOW_COMPILER_MLIR_XLA_IR_CHLO_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_CHLO_OPS_H_
 
 #include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringRef.h"
 #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Dialect.h"
@@ -42,4 +42,4 @@ class XlaHloClientDialect : public Dialect {
 }  // namespace xla_chlo
 }  // namespace mlir
 
-#endif  // TENSORFLOW_COMPILER_MLIR_XLA_IR_CHLO_OPS_H_
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_CHLO_OPS_H_

include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h

@@ -15,8 +15,8 @@ limitations under the License.
 
 // This file defines the operations used in the XLA dialect.
 
-#ifndef TENSORFLOW_COMPILER_MLIR_XLA_IR_HLO_OPS_H_
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_HLO_OPS_H_
-#define TENSORFLOW_COMPILER_MLIR_XLA_IR_HLO_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_HLO_OPS_H_
 
 #include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringRef.h"
 #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
@@ -96,4 +96,4 @@ LogicalResult deriveShapeFromFirstOperand(
 }  // end namespace xla_hlo
 }  // end namespace mlir
 
-#endif  //  TENSORFLOW_COMPILER_MLIR_XLA_IR_HLO_OPS_H_
+#endif  //  TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_HLO_OPS_H_

include/mlir-hlo/Dialect/mhlo/IR/infer_fusibility_op_interface.h

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_COMPILER_MLIR_XLA_IR_INFER_FUSIBILITY_OP_INTERFACE_H_
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_INFER_FUSIBILITY_OP_INTERFACE_H_
-#define TENSORFLOW_COMPILER_MLIR_XLA_IR_INFER_FUSIBILITY_OP_INTERFACE_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_INFER_FUSIBILITY_OP_INTERFACE_H_
 
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
@@ -25,4 +25,4 @@ namespace mlir {
 
 }  // namespace mlir
 
-#endif  // TENSORFLOW_COMPILER_MLIR_XLA_IR_INFER_FUSIBILITY_OP_INTERFACE_H_
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_INFER_FUSIBILITY_OP_INTERFACE_H_

include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h

@@ -15,8 +15,8 @@ limitations under the License.
 
 // This file defines the operations used in the LXLA dialect.
 
-#ifndef TENSORFLOW_COMPILER_MLIR_XLA_IR_LHLO_OPS_H_
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_LHLO_OPS_H_
-#define TENSORFLOW_COMPILER_MLIR_XLA_IR_LHLO_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_LHLO_OPS_H_
 
 #include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringRef.h"
 #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
@@ -49,4 +49,4 @@ class XlaLhloDialect : public Dialect {
 }  // namespace xla_lhlo
 }  // end namespace mlir
 
-#endif  // TENSORFLOW_COMPILER_MLIR_XLA_IR_LHLO_OPS_H_
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_IR_LHLO_OPS_H_

include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h

@@ -0,0 +1,80 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_MAP_HLO_TO_LHLO_OP_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_MAP_HLO_TO_LHLO_OP_H_
+
+#include <type_traits>
+
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+
+namespace mlir {
+namespace xla_hlo {
+
+template <typename HloOpTy>
+struct HloToLhloOpImpl {
+  using Type = std::false_type;
+};
+template <typename HloOpTy>
+using HloToLhloOp = typename HloToLhloOpImpl<HloOpTy>::Type;
+
+#define MAP_HLO_TO_LHLO(OpName)             \
+  template <>                               \
+  struct HloToLhloOpImpl<xla_hlo::OpName> { \
+    using Type = xla_lhlo::OpName;          \
+  }
+
+MAP_HLO_TO_LHLO(AbsOp);
+MAP_HLO_TO_LHLO(AddOp);
+MAP_HLO_TO_LHLO(AndOp);
+MAP_HLO_TO_LHLO(BroadcastInDimOp);
+MAP_HLO_TO_LHLO(CeilOp);
+MAP_HLO_TO_LHLO(ConstOp);
+MAP_HLO_TO_LHLO(CompareOp);
+MAP_HLO_TO_LHLO(ComplexOp);
+MAP_HLO_TO_LHLO(ConvOp);
+MAP_HLO_TO_LHLO(ConvertOp);
+MAP_HLO_TO_LHLO(CopyOp);
+MAP_HLO_TO_LHLO(CosOp);
+MAP_HLO_TO_LHLO(DivOp);
+MAP_HLO_TO_LHLO(DotOp);
+MAP_HLO_TO_LHLO(ExpOp);
+MAP_HLO_TO_LHLO(GatherOp);
+MAP_HLO_TO_LHLO(ImagOp);
+MAP_HLO_TO_LHLO(IotaOp);
+MAP_HLO_TO_LHLO(LogOp);
+MAP_HLO_TO_LHLO(MaxOp);
+MAP_HLO_TO_LHLO(MinOp);
+MAP_HLO_TO_LHLO(MulOp);
+MAP_HLO_TO_LHLO(NegOp);
+MAP_HLO_TO_LHLO(RealOp);
+MAP_HLO_TO_LHLO(ReduceOp);
+MAP_HLO_TO_LHLO(ReshapeOp);
+MAP_HLO_TO_LHLO(RemOp);
+MAP_HLO_TO_LHLO(RsqrtOp);
+MAP_HLO_TO_LHLO(SelectOp);
+MAP_HLO_TO_LHLO(SignOp);
+MAP_HLO_TO_LHLO(SinOp);
+MAP_HLO_TO_LHLO(SqrtOp);
+MAP_HLO_TO_LHLO(SubOp);
+MAP_HLO_TO_LHLO(TanhOp);
+
+#undef MAP_HLO_TO_LHLO
+
+}  // namespace xla_hlo
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_MAP_HLO_TO_LHLO_OP_H_

include/mlir-hlo/Dialect/mhlo/transforms/map_xla_to_scalar_op.h

@@ -0,0 +1,510 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_MAP_XLA_TO_SCALAR_OP_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_MAP_XLA_TO_SCALAR_OP_H_
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringRef.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringSwitch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace impl {
+
+// A struct to map LhloBinaryOpTy type to the corresponding floating-point and
+// integer scalar operation types.
+template <typename LhloBinaryOpTy>
+struct LhloToScalarOp;
+
+template <>
+struct LhloToScalarOp<xla_lhlo::AddOp> {
+  using FOp = ::mlir::AddFOp;
+  using IOp = ::mlir::AddIOp;
+};
+template <>
+struct LhloToScalarOp<xla_lhlo::CompareOp> {
+  using FOp = ::mlir::CmpFOp;
+  using IOp = ::mlir::CmpIOp;
+};
+template <>
+struct LhloToScalarOp<xla_lhlo::DivOp> {
+  using FOp = ::mlir::DivFOp;
+  using IOp = ::mlir::SignedDivIOp;
+};
+template <>
+struct LhloToScalarOp<xla_lhlo::MulOp> {
+  using FOp = ::mlir::MulFOp;
+  using IOp = ::mlir::MulIOp;
+};
+template <>
+struct LhloToScalarOp<xla_lhlo::RemOp> {
+  using FOp = ::mlir::RemFOp;
+  using IOp = ::mlir::SignedRemIOp;
+};
+template <>
+struct LhloToScalarOp<xla_lhlo::SubOp> {
+  using FOp = ::mlir::SubFOp;
+  using IOp = ::mlir::SubIOp;
+};
+
+template <typename LhloBinaryOpTy>
+struct ScalarOp {
+  using FOp = typename LhloToScalarOp<LhloBinaryOpTy>::FOp;
+  using IOp = typename LhloToScalarOp<LhloBinaryOpTy>::IOp;
+};
+
+// Alias for the map from LHLO binary op type to STD floating-point op type.
+template <typename LhloOp>
+using ScalarFOp = typename ScalarOp<LhloOp>::FOp;
+// Alias for the map from LHLO binary op type to STD integer op type.
+template <typename LhloOp>
+using ScalarIOp = typename ScalarOp<LhloOp>::IOp;
+
+template <typename... Args>
+struct MapLhloOpToStdScalarOpImpl {
+  Value operator()(Location loc, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b) {
+    return nullptr;
+  }
+};
+
+template <typename StdScalarOp>
+struct MapLhloOpToStdScalarOpImpl<StdScalarOp> {
+  Value operator()(Location loc, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b) {
+    return b->template create<StdScalarOp>(loc, result_types, args, mlir::None);
+  }
+};
+
+template <typename SupportedType, typename StdScalarOp, typename... Args>
+struct MapLhloOpToStdScalarOpImpl<SupportedType, StdScalarOp, Args...> {
+  Value operator()(Location loc, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b) {
+    Type element_type = args.front().getType();
+    if (element_type.isa<SupportedType>()) {
+      return b->template create<StdScalarOp>(loc, result_types, args,
+                                             mlir::None);
+    }
+    return MapLhloOpToStdScalarOpImpl<Args...>{}(loc, result_types, args, b);
+  }
+};
+
+// Inserts the computation that corresponds to the body of the loop for lowered
+// LHLO unary/binary op. Returns the value for the result.
+template <typename LhloOpTy>
+inline Value MapLhloOpToStdScalarOp(Location loc, ArrayRef<Type> result_types,
+                                    ArrayRef<Value> args, OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<IntegerType, ScalarIOp<LhloOpTy>, FloatType,
+                                    ScalarFOp<LhloOpTy>>{}(loc, result_types,
+                                                           args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::AbsOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  Type element_type = args.front().getType();
+  if (element_type.isa<FloatType>()) {
+    return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::AbsFOp>{}(
+        loc, result_types, args, b);
+  }
+  if (element_type.isa<IntegerType>()) {
+    // xla_lhlo.abs(x, result) ->  result = select((x > 0), x, sub(0, x))
+    Value lhs = args[0];
+    auto integer_type = element_type.dyn_cast<IntegerType>();
+
+    auto zero_intval =
+        b->create<::mlir::ConstantIntOp>(loc, 0, integer_type.getWidth());
+    auto lhs_gt_zero = b->create<ScalarIOp<CompareOp>>(loc, CmpIPredicate::sge,
+                                                       lhs, zero_intval);
+    auto neg_val = b->create<ScalarIOp<xla_lhlo::SubOp>>(loc, zero_intval, lhs);
+    return b->create<::mlir::SelectOp>(loc, lhs_gt_zero, lhs, neg_val);
+  }
+  return nullptr;
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::AndOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<IntegerType, ::mlir::AndOp>{}(
+      loc, result_types, args, b);
+}
+
+template <typename PredicateType>
+inline Optional<PredicateType> getCmpPredicate(
+    StringRef xla_comparison_direction) {
+  return llvm::None;
+}
+
+template <>
+inline Optional<CmpFPredicate> getCmpPredicate<CmpFPredicate>(
+    StringRef xla_comparison_direction) {
+  return llvm::StringSwitch<Optional<CmpFPredicate>>(xla_comparison_direction)
+      .Case("EQ", CmpFPredicate::OEQ)
+      .Case("NE", CmpFPredicate::ONE)
+      .Case("GE", CmpFPredicate::OGE)
+      .Case("GT", CmpFPredicate::OGT)
+      .Case("LE", CmpFPredicate::OLE)
+      .Case("LT", CmpFPredicate::OLT)
+      .Default(llvm::None);
+}
+
+template <>
+inline Optional<CmpIPredicate> getCmpPredicate<CmpIPredicate>(
+    StringRef xla_comparison_direction) {
+  return llvm::StringSwitch<Optional<CmpIPredicate>>(xla_comparison_direction)
+      .Case("EQ", CmpIPredicate::eq)
+      .Case("NE", CmpIPredicate::ne)
+      .Case("GE", CmpIPredicate::sge)
+      .Case("GT", CmpIPredicate::sgt)
+      .Case("LE", CmpIPredicate::sle)
+      .Case("LT", CmpIPredicate::slt)
+      .Default(llvm::None);
+}
+
+template <typename XLACompareOpTy>
+inline Value MapXlaCompareOpToStdScalarOp(Location loc,
+                                          StringRef comparison_direction,
+                                          ArrayRef<Type> result_types,
+                                          ArrayRef<Value> args, OpBuilder* b) {
+  const auto& lhs = args[0];
+  const auto& rhs = args[1];
+  Type element_type = lhs.getType();
+  if (element_type.isSignlessInteger()) {
+    Optional<CmpIPredicate> predicate =
+        getCmpPredicate<CmpIPredicate>(comparison_direction);
+    assert(predicate.hasValue() && "expected valid comparison direction");
+    return b->create<ScalarIOp<XLACompareOpTy>>(loc, predicate.getValue(), lhs,
+                                                rhs);
+  }
+  if (element_type.isa<FloatType>()) {
+    Optional<CmpFPredicate> predicate =
+        getCmpPredicate<CmpFPredicate>(comparison_direction);
+    assert(predicate.hasValue() && "expected valid comparison direction");
+    return b->create<ScalarFOp<XLACompareOpTy>>(loc, predicate.getValue(), lhs,
+                                                rhs);
+  }
+  return nullptr;
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::CopyOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return args.front();
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::ExpOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::ExpOp>{}(
+      loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::CeilOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::CeilFOp>{}(
+      loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::ComplexOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<CreateComplexOp>{}(loc, result_types, args,
+                                                       b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::RealOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<ReOp>{}(loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::ImagOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<ImOp>{}(loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::ConvertOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  Type sourceType = args.front().getType();
+  Type targetType = result_types.front();
+
+  if (mlir::SIToFPOp::areCastCompatible(sourceType, targetType)) {
+    return b->create<mlir::SIToFPOp>(loc, result_types, args, mlir::None);
+  } else if (sourceType.isa<FloatType>() && targetType.isa<FloatType>()) {
+    FloatType src = sourceType.cast<FloatType>();
+    FloatType res = targetType.cast<FloatType>();
+    if (src.getWidth() > res.getWidth()) {
+      return b->create<mlir::FPTruncOp>(loc, result_types, args, mlir::None);
+    } else if (src.getWidth() < res.getWidth()) {
+      return b->create<mlir::FPExtOp>(loc, result_types, args, mlir::None);
+    }
+    // No conversion is needed for the same width floats
+    return args.front();
+  }
+  if (sourceType.isSignlessInteger() && targetType.isSignlessInteger()) {
+    IntegerType src = sourceType.cast<IntegerType>();
+    IntegerType res = targetType.cast<IntegerType>();
+    if (src.getWidth() > res.getWidth()) {
+      return b->create<mlir::TruncateIOp>(loc, result_types, args, mlir::None);
+    } else if (src.getWidth() < res.getWidth()) {
+      return b->create<mlir::ZeroExtendIOp>(loc, result_types, args,
+                                            mlir::None);
+    }
+    // No conversion is needed for the same width integers
+    return args.front();
+  }
+  if (mlir::FPToSIOp::areCastCompatible(sourceType, targetType)) {
+    return b->create<mlir::FPToSIOp>(loc, result_types, args, mlir::None);
+  }
+  return nullptr;
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::DotOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  // Dot Op converter from lhlo to affine only accepts float and integer types.
+  const auto& lhs = args[0];
+  const auto& rhs = args[1];
+  const auto& result = args[2];
+  Type element_type = lhs.getType();
+  if (element_type.isa<FloatType>()) {
+    Value float_mul = MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::MulFOp>{}(
+        loc, result_types, {lhs, rhs}, b);
+    return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::AddFOp>{}(
+        loc, result_types, {float_mul, result}, b);
+  }
+  if (element_type.isa<IntegerType>()) {
+    Value int_mul = MapLhloOpToStdScalarOpImpl<IntegerType, ::mlir::MulIOp>{}(
+        loc, result_types, {lhs, rhs}, b);
+    return MapLhloOpToStdScalarOpImpl<IntegerType, ::mlir::AddIOp>{}(
+        loc, result_types, {int_mul, result}, b);
+  }
+  return nullptr;
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::CosOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::CosOp>{}(
+      loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::SinOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::SinOp>{}(
+      loc, result_types, args, b);
+}
+
+/// Implements the conversion of XLA op to scalar op (to use within region of a
+/// linalg.generic op) for compare-select style operations like min/max.
+template <typename... Args>
+struct XlaCompareSelectOpToStdScalarOp {
+  static Value map(Location loc, StringRef comparison_direction,
+                   ArrayRef<Type> result_types, ArrayRef<Value> args,
+                   OpBuilder* b) {
+    return nullptr;
+  }
+};
+
+/// Specialization which allows converting to a comparison operation in standard
+/// dialect with a given predicate based on the element type of the operand.
+template <typename SupportedType, typename StdCompareOp, typename Predicate,
+          typename... Args>
+struct XlaCompareSelectOpToStdScalarOp<SupportedType, StdCompareOp, Predicate,
+                                       Args...> {
+  static Value map(Location loc, StringRef comparison_direction,
+                   ArrayRef<Type> result_types, ArrayRef<Value> args,
+                   OpBuilder* b) {
+    Type element_type = args.front().getType();
+    if (element_type.isa<SupportedType>()) {
+      auto predicate = getCmpPredicate<Predicate>(comparison_direction);
+      assert(predicate.hasValue() && "expected valid comparison direction");
+      auto cmp = b->template create<StdCompareOp>(loc, predicate.getValue(),
+                                                  args[0], args[1]);
+      return b->create<::mlir::SelectOp>(loc, cmp, args[0], args[1]);
+    }
+    return XlaCompareSelectOpToStdScalarOp<Args...>::map(
+        loc, comparison_direction, result_types, args, b);
+  }
+};
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::LogOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::LogOp>{}(
+      loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::MaxOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return XlaCompareSelectOpToStdScalarOp<
+      IntegerType, ScalarIOp<xla_lhlo::CompareOp>, CmpIPredicate, FloatType,
+      ScalarFOp<xla_lhlo::CompareOp>, CmpFPredicate>::map(loc, "GT",
+                                                          result_types, args,
+                                                          b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::MinOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return XlaCompareSelectOpToStdScalarOp<
+      IntegerType, ScalarIOp<xla_lhlo::CompareOp>, CmpIPredicate, FloatType,
+      ScalarFOp<xla_lhlo::CompareOp>, CmpFPredicate>::map(loc, "LT",
+                                                          result_types, args,
+                                                          b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::NegOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  Type element_type = args.front().getType();
+  if (element_type.isa<FloatType>()) {
+    return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::NegFOp>{}(
+        loc, result_types, args, b);
+  }
+  if (element_type.isa<IntegerType>()) {
+    // xla_lhlo.neg(x, result) -> result = sub(0, x)
+    Value lhs = args[0];
+    auto integer_type = element_type.dyn_cast<IntegerType>();
+
+    auto zero_intval =
+        b->create<::mlir::ConstantIntOp>(loc, 0, integer_type.getWidth());
+    return b->create<ScalarIOp<xla_lhlo::SubOp>>(loc, zero_intval, lhs);
+  }
+  return nullptr;
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::RsqrtOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::RsqrtOp>{}(
+      loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::SelectOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<::mlir::SelectOp>{}(loc, result_types, args,
+                                                        b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::SignOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  Type element_type = args.front().getType();
+  if (element_type.isa<FloatType>()) {
+    FloatType float_type = element_type.cast<FloatType>();
+    APFloat const_value = float_type.isF32() ? APFloat(1.0f) : APFloat(1.0);
+    Value one = b->create<mlir::ConstantFloatOp>(loc, const_value, float_type);
+    return b->create<::mlir::CopySignOp>(loc, result_types, one, args[0]);
+  }
+  return nullptr;
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::SqrtOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::SqrtOp>{}(
+      loc, result_types, args, b);
+}
+
+template <>
+inline Value MapLhloOpToStdScalarOp<xla_lhlo::TanhOp>(
+    Location loc, ArrayRef<Type> result_types, ArrayRef<Value> args,
+    OpBuilder* b) {
+  return MapLhloOpToStdScalarOpImpl<FloatType, ::mlir::TanhOp>{}(
+      loc, result_types, args, b);
+}
+
+}  // namespace impl
+
+struct XlaOpToStdScalarOp {
+  // Implementation for LHLO ops except xla_lhlo::CompareOp.
+  template <typename XlaOpTy, typename LhloOpTy = XlaOpTy,
+            typename = std::enable_if_t<
+                !std::is_same<LhloOpTy, xla_lhlo::CompareOp>::value &&
+                std::is_same<typename xla_hlo::HloToLhloOp<LhloOpTy>,
+                             std::false_type>::value>>
+  static Value map(XlaOpTy op, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b, unsigned i = 0) {
+    return impl::MapLhloOpToStdScalarOp<LhloOpTy>(op.getLoc(), result_types,
+                                                  args, b);
+  }
+
+  // Implementation for HLO ops except xla_hlo::CompareOp.
+  template <typename XlaOpTy, typename LhloOpTy = xla_hlo::HloToLhloOp<XlaOpTy>,
+            typename = std::enable_if_t<
+                !std::is_same<LhloOpTy, xla_lhlo::CompareOp>::value &&
+                !std::is_same<LhloOpTy, std::false_type>::value>>
+  static Value map(XlaOpTy op, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b, int i = 0) {
+    return impl::MapLhloOpToStdScalarOp<LhloOpTy>(op.getLoc(), result_types,
+                                                  args, b);
+  }
+
+  // Implementation for xla_lhlo::CompareOp.
+  template <typename LhloOpTy, typename = std::enable_if_t<std::is_same<
+                                   LhloOpTy, xla_lhlo::CompareOp>::value>>
+  static Value map(xla_lhlo::CompareOp op, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b) {
+    auto comparison_direction = op.comparison_direction();
+    return impl::MapXlaCompareOpToStdScalarOp<xla_lhlo::CompareOp>(
+        op.getLoc(), comparison_direction, result_types, args, b);
+  }
+
+  // Implementation for xla_hlo::CompareOp.
+  template <typename HloOpTy, typename = std::enable_if_t<std::is_same<
+                                  HloOpTy, xla_hlo::CompareOp>::value>>
+  static Value map(xla_hlo::CompareOp op, ArrayRef<Type> result_types,
+                   ArrayRef<Value> args, OpBuilder* b) {
+    auto comparison_direction = op.comparison_direction();
+    return impl::MapXlaCompareOpToStdScalarOp<xla_lhlo::CompareOp>(
+        op.getLoc(), comparison_direction, result_types, args, b);
+  }
+};
+
+}  // namespace xla_lhlo
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_MAP_XLA_TO_SCALAR_OP_H_

include/mlir-hlo/Dialect/mhlo/transforms/passes.h

@@ -0,0 +1,105 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_PASSES_H_
+
+#include <memory>
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/ArrayRef.h"
+
+namespace mlir {
+
+class FuncOp;
+class ModuleOp;
+class Operation;
+template <typename T>
+class OperationPass;
+class Pass;
+
+namespace xla_hlo {
+
+/// Lowers HLO control flow ops to the Standard dialect.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeControlFlowPass();
+
+/// Lowers from HLO dialect to Standard dialect.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeToStdPass();
+
+/// Lowers from HLO dialect to LHLO dialect allocating/deallocating temporary
+/// buffers if necessary. If `results_escape_functions` is set to true,
+/// allocated buffers for function results will be returned and escape the
+/// function. Otherwise, the signature is rewritten with extra arguments for the
+/// buffers that are to be used for results.
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeToLhloPass(
+    bool results_escape_functions = false);
+
+// Lowers from HLO dialect to Linalg dialect.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeHloToLinalgPass();
+
+// Transforms unranked HLO operations to ranked ones where possible.
+std::unique_ptr<OperationPass<FuncOp>> createTransformUnrankedHloPass();
+
+// Sinks constants implicitly captured in control flow regions. This is
+// necessary to export to XLA.
+std::unique_ptr<OperationPass<FuncOp>> createSinkConstantsToControlFlowPass();
+
+// fuse xla_hlo ops to kLoop/kInput fusion patterns
+std::unique_ptr<OperationPass<FuncOp>> createXlaHloFusionPass();
+
+}  // namespace xla_hlo
+
+namespace xla_lhlo {
+
+// Lowers from LHLO dialect to Affine dialect.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeToAffinePass();
+
+// Lowers from LHLO dialect to Linalg dialect.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeLhloToLinalgPass();
+
+// Lowers from LHLO dialect to GPU dialect.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeToGpuPass();
+
+// Fuses linalg ops obtained after LHLO lowering. To enable fusion,
+// operations are first tiled.
+//
+// When 'use_parallel_loops' is set, the tiling will use scf.parallel
+// operations. Otherwise, scf.for operations are used.
+//
+// 'tile_sizes' provides the tile sizes to use for tiling. If the linalg
+// operation has more dimensions than tile sizes provided, 1 is used as
+// default.
+std::unique_ptr<OperationPass<FuncOp>> createLhloFuseLinalg(
+    bool use_parallel_loops = false, llvm::ArrayRef<unsigned> tile_sizes = {});
+
+// Removes unnecessary LHLO copies which copy from the allocated buffers to the
+// block arguments. The block arguments are used instead of all uses of these
+// buffers. The buffers are freed. This pass only works in regions that contain
+// a single block.
+std::unique_ptr<Pass> createLhloCopyRemovalPass();
+
+// Lowers from LHLO dialect to parallel loops.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeLhloToParallelLoopsPass();
+
+}  // namespace xla_lhlo
+
+namespace xla {
+
+/// Lowers the standard TanhOp to an approximation that does not use intrinsics.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeTanhToApproximationPass();
+
+}  // namespace xla
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_PASSES_H_

include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h

@@ -0,0 +1,106 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_REWRITERS_H_
+#define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_REWRITERS_H_
+
+#include <memory>
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+class LLVMTypeConverter;
+class LowerToLLVMOptions;
+class OwningRewritePatternList;
+class BufferAssignmentPlacer;
+namespace xla_hlo {
+
+// Collection of rewrite patterns for lowering a general dot product.
+void PopulateGeneralDotOpLoweringPatterns(OwningRewritePatternList *patterns,
+                                          MLIRContext *ctx);
+
+// Collection of rewrite patterns for lowering complex operations to equivalent
+// float operations.
+void PopulateComplexLoweringPatterns(MLIRContext *context,
+                                     OwningRewritePatternList *patterns);
+
+void PopulateXlaToStdPatterns(OwningRewritePatternList *patterns,
+                              MLIRContext *ctx);
+
+// Collection of rewrite patterns for lowering of HLO to LHLO dialect.
+void populateHLOToLHLOConversionPattern(
+    MLIRContext *context, BufferAssignmentPlacer *bufferAssignment,
+    TypeConverter *converter, OwningRewritePatternList *patterns);
+// Collection of rewrite patterns for lowering of HLO to Linalg dialect.
+void populateHLOToLinalgConversionPattern(MLIRContext *context,
+                                          OwningRewritePatternList *patterns);
+
+// Sets up legality definitions for materializing broadcasts.
+void SetupMaterializeBroadcastsLegality(MLIRContext *context,
+                                        ConversionTarget *conversionTarget);
+
+// Populates a collection of rewrite patterns for materializing broadcast
+// attributes to equivalent sequences of ops.
+void PopulateMaterializeBroadcastsPatterns(MLIRContext *context,
+                                           OwningRewritePatternList *patterns);
+
+// Sets up legality definitions for element-wise operations on ranked tensors.
+void SetupTransformUnrankedHloLegality(MLIRContext *context,
+                                       ConversionTarget *conversionTarget);
+
+// Populates a collection of rewrite patterns to realize element-wise operations
+// on ranked tensors where possible.
+void PopulateTransformUnrankedHloPatterns(MLIRContext *context,
+                                          OwningRewritePatternList *patterns);
+
+// Populate a collection of conversion patterns for un-fusing
+// batch_norm_inference and batch_norm_training into constituent HLO ops.
+// TODO(laurenzo): Implement un-fusing of batch_norm_training.
+void PopulateUnfuseBatchNormPatterns(MLIRContext *context,
+                                     OwningRewritePatternList *patterns);
+
+}  // namespace xla_hlo
+
+namespace xla_lhlo {
+
+/// Collect a set of patterns to convert from the LHLO dialect to LLVM.
+void PopulateLhloToLLVMConversionPatterns(const LowerToLLVMOptions &options,
+                                          LLVMTypeConverter *converter,
+                                          OwningRewritePatternList *patterns);
+
+}  // namespace xla_lhlo
+
+namespace xla_chlo {
+
+// Populates a collection of conversion patterns for legalizing client-HLO to
+// HLO.
+void PopulateLegalizeChloToHloPatterns(MLIRContext *context,
+                                       OwningRewritePatternList *patterns);
+
+}  // namespace xla_chlo
+
+namespace xla {
+
+// Populates a pattern that translates the standard TanhOp to an approximation
+// that does not use intrinsics.
+void PopulateTanhToApproximationPatterns(MLIRContext *context,
+                                         OwningRewritePatternList *patterns);
+
+}  // namespace xla
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_REWRITERS_H_

include/mlir-hlo/utils/cycle_detector.h

@@ -0,0 +1,165 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_XLA_TRANSFORMS_CYCLE_DETECTOR_H_
+#define TENSORFLOW_COMPILER_MLIR_XLA_TRANSFORMS_CYCLE_DETECTOR_H_
+
+#include <vector>
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/DenseMap.h"
+
+namespace mlir {
+
+// -------------------------------------------------------------------
+
+// This file contains a light version of GraphCycles implemented in
+// tensorflow/compiler/jit/graphcycles/graphcycles.h
+//
+// We re-implement it here because we do not want to rely
+// on TensorFlow data structures, and hence we can move
+// corresponding passes to llvm repo. easily in case necessnary.
+
+// --------------------------------------------------------------------
+
+// This is a set data structure that provides a deterministic iteration order.
+// The iteration order of elements only depends on the sequence of
+// inserts/deletes, so as long as the inserts/deletes happen in the same
+// sequence, the set will have the same iteration order.
+//
+// Assumes that T can be cheaply copied for simplicity.
+template <typename T>
+class OrderedSet {
+ public:
+  // Inserts `value` into the ordered set.  Returns true if the value was not
+  // present in the set before the insertion.
+  bool Insert(T value) {
+    bool new_insertion =
+        value_to_index_.insert({value, value_sequence_.size()}).second;
+    if (new_insertion) {
+      value_sequence_.push_back(value);
+    }
+    return new_insertion;
+  }
+
+  // Removes `value` from the set.  Assumes `value` is already present in the
+  // set.
+  void Erase(T value) {
+    auto it = value_to_index_.find(value);
+
+    // Since we don't want to move values around in `value_sequence_` we swap
+    // the value in the last position and with value to be deleted and then
+    // pop_back.
+    value_to_index_[value_sequence_.back()] = it->second;
+    std::swap(value_sequence_[it->second], value_sequence_.back());
+    value_sequence_.pop_back();
+    value_to_index_.erase(it);
+  }
+
+  void Reserve(size_t new_size) {
+    value_to_index_.reserve(new_size);
+    value_sequence_.reserve(new_size);
+  }
+
+  void Clear() {
+    value_to_index_.clear();
+    value_sequence_.clear();
+  }
+
+  bool Contains(T value) const { return value_to_index_.count(value); }
+  size_t Size() const { return value_sequence_.size(); }
+
+  const std::vector<T>& GetSequence() const { return value_sequence_; }
+
+ private:
+  // The stable order that we maintain through insertions and deletions.
+  std::vector<T> value_sequence_;
+
+  // Maps values to their indices in `value_sequence_`.
+  llvm::DenseMap<T, int> value_to_index_;
+};
+
+// ---------------------------------------------------------------------
+
+// GraphCycles detects the introduction of a cycle into a directed
+// graph that is being built up incrementally.
+//
+// Nodes are identified by small integers.  It is not possible to
+// record multiple edges with the same (source, destination) pair;
+// requests to add an edge where one already exists are silently
+// ignored.
+//
+// It is also not possible to introduce a cycle; an attempt to insert
+// an edge that would introduce a cycle fails and returns false.
+//
+// GraphCycles uses no internal locking; calls into it should be
+// serialized externally.
+
+// Performance considerations:
+//   Works well on sparse graphs, poorly on dense graphs.
+//   Extra information is maintained incrementally to detect cycles quickly.
+//   InsertEdge() is very fast when the edge already exists, and reasonably fast
+//   otherwise.
+//   FindPath() is linear in the size of the graph.
+// The current implementation uses O(|V|+|E|) space.
+
+class GraphCycles {
+ public:
+  explicit GraphCycles(int32_t num_nodes);
+  ~GraphCycles();
+
+  // Attempt to insert an edge from x to y.  If the
+  // edge would introduce a cycle, return false without making any
+  // changes. Otherwise add the edge and return true.
+  bool InsertEdge(int32_t x, int32_t y);
+
+  // Remove any edge that exists from x to y.
+  void RemoveEdge(int32_t x, int32_t y);
+
+  // Return whether there is an edge directly from x to y.
+  bool HasEdge(int32_t x, int32_t y) const;
+
+  // Contracts the edge from 'a' to node 'b', merging nodes 'a' and 'b'. One of
+  // the nodes is removed from the graph, and edges to/from it are added to
+  // the remaining one, which is returned. If contracting the edge would create
+  // a cycle, does nothing and return no value.
+  llvm::Optional<int32_t> ContractEdge(int32_t a, int32_t b);
+
+  // Return whether dest_node `y` is reachable from source_node `x`
+  // by following edges. This is non-thread-safe version.
+  bool IsReachable(int32_t x, int32_t y);
+
+  // Return a copy of the successors set. This is needed for code using the
+  // collection while modifying the GraphCycles.
+  std::vector<int32_t> SuccessorsCopy(int32_t node) const;
+
+  // Returns all nodes in post order.
+  //
+  // If there is a path from X to Y then X appears after Y in the
+  // returned vector.
+  std::vector<int32_t> AllNodesInPostOrder() const;
+
+  // ----------------------------------------------------
+  struct Rep;
+
+ private:
+  GraphCycles(const GraphCycles&) = delete;
+  GraphCycles& operator=(const GraphCycles&) = delete;
+
+  Rep* rep_;  // opaque representation
+};
+
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_XLA_TRANSFORMS_CYCLE_DETECTOR_H_

lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo.cc

@@ -0,0 +1,242 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Shape/IR/Shape.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/OperationSupport.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/utils/broadcast_utils.h"
+
+namespace mlir {
+namespace xla_chlo {
+
+namespace {
+
+// Converts binary ops that statically are determined to not broadcast directly
+// to the corresponding xla_hlo non-broadcasting op.
+template <typename ChloOpTy, typename HloOpTy, typename Adaptor>
+struct ConvertTrivialNonBroadcastBinaryOp : public OpRewritePattern<ChloOpTy> {
+  using OpRewritePattern<ChloOpTy>::OpRewritePattern;
+  LogicalResult matchAndRewrite(ChloOpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Only rewrite for statically determinable non-broadcasting cases.
+    auto lhs_type = op.lhs().getType().template dyn_cast<RankedTensorType>();
+    auto rhs_type = op.rhs().getType().template dyn_cast<RankedTensorType>();
+    if (!lhs_type || !rhs_type) return failure();
+
+    // Requires rank broadcast.
+    if (lhs_type.getRank() != rhs_type.getRank()) return failure();
+    // Any dynamic dimension may require broadcasting and requires more
+    // analysis.
+    if (!lhs_type.hasStaticShape() || !rhs_type.hasStaticShape())
+      return failure();
+
+    for (auto extents : llvm::zip(lhs_type.getShape(), rhs_type.getShape())) {
+      auto lhs_extent = std::get<0>(extents);
+      auto rhs_extent = std::get<1>(extents);
+      if (lhs_extent != rhs_extent) {
+        return failure();
+      }
+    }
+
+    rewriter.replaceOp(op, {Adaptor::CreateOp(op, op.getResult().getType(),
+                                              op.lhs(), op.rhs(), rewriter)});
+    return success();
+  }
+};
+
+// Converts a binary op with ranked broadcasting operands to explicitly
+// broadcast and invoke the corresponding xla_hlo non-broadcasting op.
+// Note that dynamic broadcasting supported by this pattern is only valid for
+// "numpy" broadcasting semantics as defined here:
+//   https://docs.scipy.org/doc/numpy/reference/ufuncs.html
+// Specifically, this includes the following cases:
+//   - Same rank broadcast (operands have the same static rank).
+//   - Different-rank broadcast, either without a broadcast_dims attribte or
+//     with the broadcast_dims attribute set to map to a prefix padding.
+//   - Legal combinations of degenerate (1-dim) implicit broadcasting.
+// The restriction on broadcast_dims derives from the definition of the
+// `shape.broadcast` op, which only supports prefix-padding.
+//
+// It may be possible to expand this pattern to operate on unranked tensors in
+// the future by emitting more code to dynamically differentiate based on rank.
+// Whether that is of any practical benefit remains to be seen.
+template <typename ChloOpTy, typename HloOpTy, typename Adaptor>
+struct ConvertRankedDynamicBroadcastBinaryOp
+    : public OpRewritePattern<ChloOpTy> {
+  using OpRewritePattern<ChloOpTy>::OpRewritePattern;
+  LogicalResult matchAndRewrite(ChloOpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Only support ranked operands.
+    Value lhs = op.lhs();
+    Value rhs = op.rhs();
+    auto lhs_type = lhs.getType().dyn_cast<RankedTensorType>();
+    auto rhs_type = rhs.getType().dyn_cast<RankedTensorType>();
+    auto result_type =
+        op.getResult().getType().template dyn_cast<RankedTensorType>();
+    if (!lhs_type || !rhs_type || !result_type) return failure();
+
+    // Check for "numpy"-style rank broadcast.
+    auto broadcast_dimensions = op.broadcast_dimensions();
+    if (broadcast_dimensions &&
+        !xla::IsLegalNumpyRankedBroadcast(lhs, rhs, *broadcast_dimensions)) {
+      // Note: It is unclear whether the general specification of explicit
+      // broadcast_dimensions on binary ops is a feature we want to carry
+      // forward. While it can technically be implemented for ranked-dynamic,
+      // it is incompatible with unranked inputs. If this warning is emitted
+      // in real programs, it is an indication that the feature should be
+      // implemented versus just falling back on the more standard definition
+      // of numpy-like prefix-padding.
+      op.emitWarning() << "unsupported non prefix-padded dynamic rank "
+                       << "broadcast_dimensions = " << *broadcast_dimensions;
+      return failure();
+    }
+
+    // Compute result shape.
+    auto loc = op.getLoc();
+
+    // Insert a constraint on the shapes being broadcastable and insert all
+    // future code into an assuming block reliant on the constraint.
+    Value lhs_shape = rewriter.create<shape::ShapeOfOp>(loc, lhs);
+    Value rhs_shape = rewriter.create<shape::ShapeOfOp>(loc, rhs);
+    auto broadcastable_cstr =
+        rewriter.create<shape::CstrBroadcastableOp>(loc, lhs_shape, rhs_shape);
+    auto assuming_op = rewriter.create<shape::AssumingOp>(
+        loc, ArrayRef<Type>{result_type}, broadcastable_cstr.result());
+
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.createBlock(&assuming_op.doRegion());
+
+    int64_t result_rank = std::max(lhs_type.getRank(), rhs_type.getRank());
+    Value result_extents =
+        xla::ComputeBinaryElementwiseBroadcastingResultExtents(loc, lhs, rhs,
+                                                               rewriter);
+
+    // Note that we unconditionally emit DynamicBroadcastInDim ops and let
+    // downstream canonicalizations fold them away if possible. This is
+    // because, in the dynamic case, there are many corner cases regarding
+    // when it is safe to omit, and some of them require analysis to prove
+    // properly.
+    auto lhs_broadcast_dimensions = llvm::to_vector<4>(
+        llvm::seq<int64_t>(result_rank - lhs_type.getRank(), result_rank));
+    Value broadcasted_lhs = rewriter.create<xla_hlo::DynamicBroadcastInDimOp>(
+        loc,
+        RankedTensorType::get(result_type.getShape(),
+                              lhs_type.getElementType()),
+        lhs, result_extents,
+        rewriter.getI64TensorAttr(lhs_broadcast_dimensions));
+    auto rhs_broadcast_dimensions = llvm::to_vector<4>(
+        llvm::seq<int64_t>(result_rank - rhs_type.getRank(), result_rank));
+    Value broadcasted_rhs = rewriter.create<xla_hlo::DynamicBroadcastInDimOp>(
+        loc,
+        RankedTensorType::get(result_type.getShape(),
+                              rhs_type.getElementType()),
+        rhs, result_extents,
+        rewriter.getI64TensorAttr(rhs_broadcast_dimensions));
+
+    // And generate the final non-broadcasted binary op.
+    Value final_result = Adaptor::CreateOp(op, result_type, broadcasted_lhs,
+                                           broadcasted_rhs, rewriter);
+    rewriter.create<shape::AssumingYieldOp>(loc, final_result);
+    rewriter.replaceOp(op, {assuming_op.getResult(0)});
+    return success();
+  }
+};
+
+template <typename ChloOpTy, typename HloOpTy, typename Adaptor>
+void PopulateForBinaryOp(MLIRContext *context,
+                         OwningRewritePatternList *patterns) {
+  patterns
+      ->insert<ConvertTrivialNonBroadcastBinaryOp<ChloOpTy, HloOpTy, Adaptor>>(
+          context, 10);
+  patterns->insert<
+      ConvertRankedDynamicBroadcastBinaryOp<ChloOpTy, HloOpTy, Adaptor>>(
+      context, 5);
+}
+
+template <typename FromOpTy, typename ToOpTy>
+struct HloBinaryElementwiseAdaptor {
+  static ToOpTy CreateOp(FromOpTy from_op, Type result_type,
+                         Value broadcasted_lhs, Value broadcasted_rhs,
+                         OpBuilder &builder) {
+    return builder.create<ToOpTy>(from_op.getLoc(), result_type,
+                                  broadcasted_lhs, broadcasted_rhs);
+  }
+};
+
+struct HloComplexAdaptor {
+  static xla_hlo::ComplexOp CreateOp(BroadcastComplexOp from_op,
+                                     Type result_type, Value broadcasted_lhs,
+                                     Value broadcasted_rhs,
+                                     OpBuilder &builder) {
+    return builder.create<xla_hlo::ComplexOp>(from_op.getLoc(), result_type,
+                                              broadcasted_lhs, broadcasted_rhs);
+  }
+};
+
+struct HloCompareAdaptor {
+  static xla_hlo::CompareOp CreateOp(BroadcastCompareOp from_op,
+                                     Type result_type, Value broadcasted_lhs,
+                                     Value broadcasted_rhs,
+                                     OpBuilder &builder) {
+    return builder.create<xla_hlo::CompareOp>(from_op.getLoc(), result_type,
+                                              broadcasted_lhs, broadcasted_rhs,
+                                              from_op.comparison_direction());
+  }
+};
+
+}  // namespace
+
+void PopulateLegalizeChloToHloPatterns(MLIRContext *context,
+                                       OwningRewritePatternList *patterns) {
+  // Instantiate conversion templates for conforming binary elementwise ops
+  // that do not have different dtypes between operands and results and do
+  // not have special attributes that need to be preserved.
+#define POPULATE_BCAST(ChloOp, HloOp)                                      \
+  PopulateForBinaryOp<ChloOp, HloOp,                                       \
+                      HloBinaryElementwiseAdaptor<ChloOp, HloOp>>(context, \
+                                                                  patterns);
+
+  POPULATE_BCAST(BroadcastAddOp, xla_hlo::AddOp);
+  POPULATE_BCAST(BroadcastAndOp, xla_hlo::AndOp);
+  POPULATE_BCAST(BroadcastAtan2Op, xla_hlo::Atan2Op);
+  POPULATE_BCAST(BroadcastDivOp, xla_hlo::DivOp);
+  POPULATE_BCAST(BroadcastMaxOp, xla_hlo::MaxOp);
+  POPULATE_BCAST(BroadcastMinOp, xla_hlo::MinOp);
+  POPULATE_BCAST(BroadcastMulOp, xla_hlo::MulOp);
+  POPULATE_BCAST(BroadcastOrOp, xla_hlo::OrOp);
+  POPULATE_BCAST(BroadcastPowOp, xla_hlo::PowOp);
+  POPULATE_BCAST(BroadcastRemOp, xla_hlo::RemOp);
+  POPULATE_BCAST(BroadcastShiftLeftOp, xla_hlo::ShiftLeftOp);
+  POPULATE_BCAST(BroadcastShiftRightArithmeticOp,
+                 xla_hlo::ShiftRightArithmeticOp);
+  POPULATE_BCAST(BroadcastShiftRightLogicalOp, xla_hlo::ShiftRightLogicalOp);
+  POPULATE_BCAST(BroadcastSubOp, xla_hlo::SubOp);
+  POPULATE_BCAST(BroadcastXorOp, xla_hlo::XorOp);
+
+  // Broadcasting ops requiring special construction.
+  PopulateForBinaryOp<BroadcastComplexOp, xla_hlo::ComplexOp,
+                      HloComplexAdaptor>(context, patterns);
+  PopulateForBinaryOp<BroadcastCompareOp, xla_hlo::CompareOp,
+                      HloCompareAdaptor>(context, patterns);
+}
+
+}  // namespace xla_chlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_pass.cc

@@ -0,0 +1,57 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Shape/IR/Shape.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla_chlo {
+
+namespace {
+
+struct TestChloLegalizeToHloPass
+    : public PassWrapper<TestChloLegalizeToHloPass, FunctionPass> {
+  void runOnFunction() override {
+    ConversionTarget conversionTarget(getContext());
+    OwningRewritePatternList conversionPatterns;
+
+    conversionTarget.addIllegalDialect<XlaHloClientDialect>();
+    // Consider the xla_hlo dialect legal for tests.
+    conversionTarget.addLegalDialect<xla_hlo::XlaHloDialect>();
+    // The conversion uses helpers from the Standard dialect.
+    conversionTarget.addLegalDialect<mlir::StandardOpsDialect>();
+    conversionTarget.addLegalDialect<mlir::shape::ShapeDialect>();
+
+    PopulateLegalizeChloToHloPatterns(&getContext(), &conversionPatterns);
+
+    if (failed(applyPartialConversion(getFunction(), conversionTarget,
+                                      conversionPatterns))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+}  // namespace xla_chlo
+}  // namespace mlir
+
+static mlir::PassRegistration<mlir::xla_chlo::TestChloLegalizeToHloPass> pass(
+    "test-xla-chlo-legalize-to-hlo",
+    "Test pass for applying chlo -> hlo legalization patterns");

lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc

@@ -0,0 +1,493 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering HLO dialect to LHLO dialect.
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/AffineMap.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BlockAndValueMapping.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Builders.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Location.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/BufferPlacement.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla_hlo {
+namespace {
+
+template <typename T>
+using BaseOpConversion = BufferAssignmentOpConversionPattern<T>;
+using StdReturnOpConverter =
+    detail::BufferAssignmentReturnOpConverter<mlir::ReturnOp, mlir::ReturnOp,
+                                              xla_lhlo::CopyOp, true>;
+
+Value InsertDynamicAllocAndDealloc(Location loc, Value result,
+                                   Value shape_operand,
+                                   ConversionPatternRewriter* rewriter) {
+  auto result_type = result.getType().dyn_cast<ShapedType>();
+  if (!result_type) {
+    result.getDefiningOp()->emitOpError()
+        << "tensor to buffer conversion expects ranked results";
+  }
+  auto memref_type =
+      MemRefType::get(result_type.getShape(), result_type.getElementType());
+
+  Operation* op = result.getDefiningOp();
+
+  // Extract the required element out of the vector.
+  SmallVector<Value, 4> dynamic_operands;
+  for (auto shape_element : llvm::enumerate(result_type.getShape())) {
+    if (shape_element.value() != ShapedType::kDynamicSize) continue;
+    Value index = rewriter->create<ConstantOp>(
+        loc, rewriter->getIntegerAttr(rewriter->getIndexType(),
+                                      shape_element.index()));
+    Value alloc_operand = rewriter->create<ExtractElementOp>(loc, shape_operand,
+                                                             ValueRange{index});
+    if (!alloc_operand.getType().isIndex()) {
+      alloc_operand = rewriter->create<IndexCastOp>(loc, alloc_operand,
+                                                    rewriter->getIndexType());
+    }
+    dynamic_operands.push_back(alloc_operand);
+  }
+
+  // Insert in front of op to ensure sizes are available.
+  OpBuilder allocBuilder(op);
+  auto alloc = allocBuilder.create<AllocOp>(loc, memref_type, dynamic_operands);
+  return alloc;
+}
+
+Value InsertAlloc(Location loc, OpResult result,
+                  BufferAssignmentPlacer* bufferAssignment,
+                  ConversionPatternRewriter* rewriter) {
+  auto result_type = result.getType().dyn_cast<ShapedType>();
+  if (!result_type || !result_type.hasStaticShape()) {
+    result.getDefiningOp()->emitOpError()
+        << "tensor to buffer conversion expects statically shaped results";
+  }
+  auto memref_type =
+      MemRefType::get(result_type.getShape(), result_type.getElementType());
+  OpBuilder::InsertionGuard guard(*rewriter);
+  rewriter->restoreInsertionPoint(
+      bufferAssignment->computeAllocPosition(result));
+  auto alloc = rewriter->create<AllocOp>(loc, memref_type);
+  return alloc;
+}
+
+template <typename HloOpTy>
+class HloToLhloOpConverter : public BaseOpConversion<HloOpTy> {
+ public:
+  using BaseOpConversion<HloOpTy>::BaseOpConversion;
+  LogicalResult matchAndRewrite(
+      HloOpTy hloOp, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const final {
+    Operation* op = hloOp.getOperation();
+    const auto& original_results = op->getResults();
+    SmallVector<Value, 4> buffer_args(operands.begin(), operands.end());
+    for (auto result : llvm::enumerate(original_results)) {
+      RankedTensorType resultType =
+          result.value().getType().dyn_cast<RankedTensorType>();
+      if (!resultType) {
+        return failure();
+      }
+      if (resultType.hasStaticShape()) {
+        buffer_args.push_back(InsertAlloc(op->getLoc(), result.value(),
+                                          this->bufferAssignment, &rewriter));
+      } else {
+        SmallVector<Value, 1> results_shape;
+        auto shape_type_op = dyn_cast<InferShapedTypeOpInterface>(op);
+        if (!shape_type_op) return failure();
+        if (failed(
+                shape_type_op.reifyReturnTypeShapes(rewriter, results_shape)))
+          return failure();
+        buffer_args.push_back(InsertDynamicAllocAndDealloc(
+            op->getLoc(), result.value(), results_shape.front(), &rewriter));
+      }
+    }
+    rewriter.create<xla_hlo::HloToLhloOp<HloOpTy>>(op->getLoc(), llvm::None,
+                                                   buffer_args, op->getAttrs());
+    rewriter.replaceOp(op, ArrayRef<Value>(buffer_args).slice(operands.size()));
+    return success();
+  }
+};
+
+struct HloToLhloDynamicBroadcastInDimOpConverter
+    : public BaseOpConversion<xla_hlo::DynamicBroadcastInDimOp> {
+ public:
+  using BaseOpConversion<xla_hlo::DynamicBroadcastInDimOp>::BaseOpConversion;
+
+  LogicalResult matchAndRewrite(
+      xla_hlo::DynamicBroadcastInDimOp op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = op.getLoc();
+    Value resultBuffer = InsertDynamicAllocAndDealloc(
+        loc, op.getResult(), op.output_dimensions(), &rewriter);
+
+    Value transformed_operand =
+        InsertDynamicMemrefCastOp(op, operands.front(), &rewriter);
+    rewriter.create<xla_lhlo::BroadcastInDimOp>(
+        loc, transformed_operand, resultBuffer, op.broadcast_dimensions());
+
+    rewriter.replaceOp(op, {resultBuffer});
+
+    return success();
+  }
+
+ private:
+  // Inserts dynamic memref to change the layout of the memref to put 0-stride
+  // and size of the target dimension if size-1 dimension expansion is
+  // necessary.
+  xla_lhlo::DynamicMemRefCastOp InsertDynamicMemrefCastOp(
+      xla_hlo::DynamicBroadcastInDimOp op, Value operand, OpBuilder* b) const {
+    auto loc = op.getLoc();
+    auto operand_type = operand.getType().cast<MemRefType>();
+    auto operand_shape = operand_type.getShape();
+
+    SmallVector<Value, 2> sizes, strides;
+    sizes.reserve(operand_shape.size());
+    strides.reserve(operand_shape.size());
+
+    Value zero = b->create<ConstantIndexOp>(loc, 0);
+    Value one = b->create<ConstantIndexOp>(loc, 1);
+    for (auto dim : llvm::enumerate(op.broadcast_dimensions())) {
+      Value broadcast_dim_value =
+          b->create<ConstantIndexOp>(loc, dim.value().getSExtValue());
+      Value result_dim_size = b->create<ExtractElementOp>(
+          loc, op.output_dimensions(), broadcast_dim_value);
+      Value operand_dim_size =
+          ShapedType::isDynamic(operand_shape[dim.index()])
+              ? b->create<DimOp>(loc, operand, dim.index()).getResult()
+              : b->create<ConstantIndexOp>(loc, operand_shape[dim.index()])
+                    .getResult();
+
+      // TODO(pifon): Revisit if this cast is needed. Maybe we can use
+      // tensor<index> for `output_dimensions` as well.
+      if (!result_dim_size.getType().isIndex()) {
+        result_dim_size =
+            b->create<IndexCastOp>(loc, result_dim_size, b->getIndexType());
+      }
+
+      // There can be two cases:
+      // 1) Operand dim == result dim => expansion is not needed => stride := 1.
+      // 2) Operand dim < result dim => expansion is needed => stride := 0.
+      Value is_expansion = b->create<CmpIOp>(loc, CmpIPredicate::slt,
+                                             operand_dim_size, result_dim_size);
+      strides.push_back(
+          b->create<mlir::SelectOp>(loc, is_expansion, zero, one));
+
+      // Size of input dim can be set to the size of the corresponding output
+      // dimension for both cases.
+      sizes.push_back(result_dim_size);
+    }
+
+    // Type-erased memref type with static rank, dynamic sizes and strides.
+    SmallVector<int64_t, 2> dynamic_layout(operand_shape.size(),
+                                           MemRefType::kDynamicStrideOrOffset);
+    SmallVector<int64_t, 2> dynamic_shape(operand_shape.size(),
+                                          MemRefType::kDynamicSize);
+    auto type_erased_memref_type = MemRefType::get(
+        dynamic_shape, operand_type.getElementType(),
+        makeStridedLinearLayoutMap(dynamic_layout,
+                                   /*offset=*/0, b->getContext()));
+
+    auto transformed_operand = b->create<xla_lhlo::DynamicMemRefCastOp>(
+        loc, type_erased_memref_type, operand, sizes, strides);
+    return transformed_operand;
+  }
+};
+
+struct HloToLhloReduceOpConverter : public BaseOpConversion<xla_hlo::ReduceOp> {
+ public:
+  using BaseOpConversion<xla_hlo::ReduceOp>::BaseOpConversion;
+
+  LogicalResult matchAndRewrite(
+      xla_hlo::ReduceOp op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = op.getLoc();
+    // TODO(b/137624192) Implement variadic reduce.
+    if (op.getNumResults() != 1) return failure();
+    if (!llvm::hasSingleElement(op.body())) {
+      return op.emitOpError()
+             << "tensor to buffer conversion expects a single block "
+                "in the region containing the operation";
+    }
+    const auto& original_results = op.getResults();
+    SmallVector<Value, 4> buffer_args(operands.begin(), operands.end());
+    for (auto result : original_results) {
+      buffer_args.push_back(
+          InsertAlloc(loc, result, this->bufferAssignment, &rewriter));
+    }
+    auto new_op = rewriter.create<xla_lhlo::ReduceOp>(
+        loc, llvm::None, buffer_args, op.getAttrs());
+
+    // Copy over the operations inside the region.
+    rewriter.inlineRegionBefore(op.body(), new_op.body(), new_op.body().end());
+
+    // Create new block arguments with correct type.
+    auto& entry_block = new_op.body().front();
+    int original_arg_count = entry_block.getNumArguments();
+    for (int i = 0; i < original_arg_count; ++i) {
+      auto old_arg = entry_block.getArgument(i);
+      auto old_type = old_arg.getType().cast<TensorType>();
+      auto new_type =
+          MemRefType::get(old_type.getShape(), old_type.getElementType());
+      auto new_arg = entry_block.addArgument(new_type);
+      rewriter.replaceUsesOfBlockArgument(old_arg, new_arg);
+    }
+    // Add an argument for the result.
+    entry_block.addArgument(
+        entry_block.getArgument(original_arg_count).getType());
+    // Remove the old arguments.
+    for (int i = original_arg_count - 1; i >= 0; --i) {
+      entry_block.eraseArgument(i);
+    }
+    // Insert terminator at the end.
+    rewriter.setInsertionPointToEnd(&entry_block);
+    rewriter.create<xla_lhlo::TerminatorOp>(loc);
+
+    rewriter.replaceOp(op, ArrayRef<Value>(buffer_args).slice(operands.size()));
+
+    return success();
+  }
+};
+
+class HloToLhloTensorLoadOpConverter
+    : public BaseOpConversion<mlir::TensorLoadOp> {
+ public:
+  using BaseOpConversion<mlir::TensorLoadOp>::BaseOpConversion;
+  LogicalResult matchAndRewrite(
+      mlir::TensorLoadOp op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const final {
+    rewriter.replaceOp(op, operands);
+    return success();
+  }
+};
+
+// TODO(b/137624192): Rewrite into a copy and elide copy if possible.
+class HloToLhloTensorStoreOpConverter
+    : public BaseOpConversion<mlir::TensorStoreOp> {
+ public:
+  using BaseOpConversion<mlir::TensorStoreOp>::BaseOpConversion;
+
+  LogicalResult matchAndRewrite(
+      mlir::TensorStoreOp op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const final {
+    rewriter.replaceOpWithNewOp<xla_lhlo::CopyOp>(
+        op, llvm::None, operands.front(), operands.back());
+    return success();
+  }
+};
+
+// Lowers from HLO dialect to LHLO dialect allocating/deallocating temporary
+// buffers if necessary.
+//
+// Example fusion with HLO ops.
+//
+// func @fusion(%arg0: memref<2x2xf32>,
+//              %arg1: memref<2x2xf32>,
+//              %arg2: memref<2x2xf32>,
+//              %arg3: memref<2x2xf32>) {
+//   "xla_lhlo.fusion"() ({
+//     %0 = tensor_load %arg1 : memref<2x2xf32>
+//     %1 = tensor_load %arg2 : memref<2x2xf32>
+//     %2 = "xla_hlo.add"(%0, %1) :
+//         (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+//     %3 = tensor_load %arg0 : memref<2x2xf32>
+//     %4 = "xla_hlo.multiply"(%2, %3) :
+//         (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+//     tensor_store %4, %arg3 : memref<2x2xf32>
+//     "xla_lhlo.terminator"() : () -> ()
+//   }) : () -> ()
+//   return
+// }
+//
+// Transformed fusion with LHLO ops.
+// func @fusion(%arg0: memref<2x2xf32>,
+//              %arg1: memref<2x2xf32>,
+//              %arg2: memref<2x2xf32>,
+//              %arg3: memref<2x2xf32>) {
+//   "xla_lhlo.fusion"() ( {
+//     %0 = alloc() : memref<2x2xf32>
+//     "xla_lhlo.add"(%arg1, %arg2, %0) :
+//         (memref<2x2xf32>, memref<2x2xf32>, memref<2x2xf32>) -> ()
+//     "xla_lhlo.multiply"(%0, %arg0, %arg3) :
+//         (memref<2x2xf32>, memref<2x2xf32>, memref<2x2xf32>) -> ()
+//     "xla_lhlo.terminator"() : () -> ()
+//   }) : () -> ()
+//   return
+// }
+//
+// FuncOp signature conversion example:
+//
+// func @func_op(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+//   %0 = "xla_hlo.maximum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) ->
+//   tensor<4xf32> %1 = "xla_hlo.add"(%arg0, %0)  : (tensor<4xf32>,
+//   tensor<4xf32>) -> tensor<4xf32> return %1 : tensor<4xf32>
+// }
+//
+// Transformed function with an extra argument for the result. The types have
+// been converted from tensor to memref.
+//
+// func @func_op(%arg0: memref<4xf32>,
+//               %arg1: memref<4xf32>,
+//               %arg2: memref<4xf32>) {
+//   %0 = alloc() : memref<4xf32>
+
+//   "xla_lhlo.maximum"(%arg0, %arg1, %0) :
+//         (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> ()
+//   %1 = alloc() : memref<4xf32>
+//   "xla_lhlo.add"(%arg0, %0, %1) :
+//         (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> ()
+//   "xla_lhlo.copy"(%1, %arg2) : (memref<4xf32>, memref<4xf32>) -> ()
+//   "xla_lhlo.terminator"() : () -> ()
+// }
+
+struct HloLegalizeToLhlo
+    : public PassWrapper<HloLegalizeToLhlo, OperationPass<ModuleOp>> {
+ public:
+  HloLegalizeToLhlo() = default;
+  HloLegalizeToLhlo(const HloLegalizeToLhlo& o) {
+    this->results_escape_function = o.results_escape_function.getValue();
+  }
+  explicit HloLegalizeToLhlo(bool results_escape_function) {
+    this->results_escape_function.setValue(results_escape_function);
+  }
+
+  void runOnOperation() override {
+    OwningRewritePatternList patterns;
+    auto& context = getContext();
+    ConversionTarget target(context);
+    target.addLegalDialect<xla_lhlo::XlaLhloDialect>();
+    target.addLegalDialect<StandardOpsDialect>();
+    target.addLegalOp<ModuleOp>();
+    target.addIllegalOp<mlir::TensorLoadOp>();
+    target.addIllegalOp<mlir::TensorStoreOp>();
+    target.addLegalOp<ModuleTerminatorOp>();
+    target.addLegalOp<TensorFromElementsOp>();
+    target.addIllegalDialect<xla_hlo::XlaHloDialect>();
+
+    BufferAssignmentTypeConverter converter;
+    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+      auto inputs = op.getType().getInputs();
+      return llvm::all_of(inputs,
+                          [](Type input) { return input.isa<MemRefType>(); }) &&
+             converter.isLegal(&op.getBody());
+    });
+    target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
+      return std::all_of(returnOp.operand_type_begin(),
+                         returnOp.operand_type_end(),
+                         [](Type type) { return type.isa<MemRefType>(); });
+    });
+
+    auto module = getOperation();
+    WalkResult result = module.walk([&](FuncOp func) -> WalkResult {
+      BufferAssignmentPlacer bufferAssignment(func);
+      OwningRewritePatternList patterns;
+      populateHLOToLHLOConversionPattern(func.getContext(), &bufferAssignment,
+                                         &converter, &patterns);
+      if (results_escape_function) {
+        populateWithBufferAssignmentOpConversionPatterns<
+            mlir::ReturnOp, mlir::ReturnOp, xla_lhlo::CopyOp,
+            /*allowMemrefFunctionResults=*/true>(&context, &bufferAssignment,
+                                                 &converter, &patterns);
+      } else {
+        populateWithBufferAssignmentOpConversionPatterns<
+            mlir::ReturnOp, mlir::ReturnOp, xla_lhlo::CopyOp,
+            /*allowMemrefFunctionResults=*/false>(&context, &bufferAssignment,
+                                                  &converter, &patterns);
+      }
+      return applyPartialConversion(func, target, patterns);
+    });
+    if (result.wasInterrupted()) {
+      signalPassFailure();
+    }
+  }
+
+ private:
+  Option<bool> results_escape_function{
+      *this, "results-escape-function",
+      llvm::cl::desc(
+          "Allocate the results of functions within the functions body"),
+      llvm::cl::init(false)};
+};
+}  // namespace
+
+void populateHLOToLHLOConversionPattern(
+    MLIRContext* context, BufferAssignmentPlacer* bufferAssignment,
+    TypeConverter* converter, OwningRewritePatternList* patterns) {
+  // clang-format off
+  patterns->insert<
+      HloToLhloDynamicBroadcastInDimOpConverter,
+      HloToLhloOpConverter<xla_hlo::AbsOp>,
+      HloToLhloOpConverter<xla_hlo::AddOp>,
+      HloToLhloOpConverter<xla_hlo::AndOp>,
+      HloToLhloOpConverter<xla_hlo::BroadcastInDimOp>,
+      HloToLhloOpConverter<xla_hlo::CeilOp>,
+      HloToLhloOpConverter<xla_hlo::CompareOp>,
+      HloToLhloOpConverter<xla_hlo::ComplexOp>,
+      HloToLhloOpConverter<xla_hlo::ConstOp>,
+      HloToLhloOpConverter<xla_hlo::ConvOp>,
+      HloToLhloOpConverter<xla_hlo::ConvertOp>,
+      HloToLhloOpConverter<xla_hlo::CopyOp>,
+      HloToLhloOpConverter<xla_hlo::CosOp>,
+      HloToLhloOpConverter<xla_hlo::DivOp>,
+      HloToLhloOpConverter<xla_hlo::DotOp>,
+      HloToLhloOpConverter<xla_hlo::ExpOp>,
+      HloToLhloOpConverter<xla_hlo::GatherOp>,
+      HloToLhloOpConverter<xla_hlo::ImagOp>,
+      HloToLhloOpConverter<xla_hlo::IotaOp>,
+      HloToLhloOpConverter<xla_hlo::LogOp>,
+      HloToLhloOpConverter<xla_hlo::MaxOp>,
+      HloToLhloOpConverter<xla_hlo::MinOp>,
+      HloToLhloOpConverter<xla_hlo::MulOp>,
+      HloToLhloOpConverter<xla_hlo::NegOp>,
+      HloToLhloOpConverter<xla_hlo::RealOp>,
+      HloToLhloOpConverter<xla_hlo::RemOp>,
+      HloToLhloOpConverter<xla_hlo::RsqrtOp>,
+      HloToLhloOpConverter<xla_hlo::ReshapeOp>,
+      HloToLhloOpConverter<xla_hlo::SelectOp>,
+      HloToLhloOpConverter<xla_hlo::SignOp>,
+      HloToLhloOpConverter<xla_hlo::SqrtOp>,
+      HloToLhloOpConverter<xla_hlo::SubOp>,
+      HloToLhloOpConverter<xla_hlo::TanhOp>,
+      HloToLhloReduceOpConverter,
+      HloToLhloTensorLoadOpConverter,
+      HloToLhloTensorStoreOpConverter
+  >(context, bufferAssignment, converter);
+  // clang-format on
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeToLhloPass(
+    bool results_escape_function) {
+  return absl::make_unique<HloLegalizeToLhlo>(results_escape_function);
+}
+
+static PassRegistration<HloLegalizeToLhlo> legalize_pass(
+    "hlo-legalize-to-lhlo", "Legalize from HLO dialect to LHLO dialect");
+
+}  // namespace xla_hlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/legalize_control_flow.cc

@@ -0,0 +1,237 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering XLA dialect to Standard dialect.
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringSwitch.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/Support/Casting.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Block.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BlockAndValueMapping.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Builders.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/TypeUtilities.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/PassRegistry.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Support/LogicalResult.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+
+using mlir::PassRegistration;
+
+namespace mlir {
+namespace xla_hlo {
+namespace {
+struct LegalizeControlFlow
+    : public mlir::PassWrapper<LegalizeControlFlow, FunctionPass> {
+  // Perform the lowering to MLIR control flow.
+  void runOnFunction() override;
+};
+
+// Replaces terminators for the newly created blocks from a targe region.
+// These terminators are replaced with branch operations to a target block.
+LogicalResult ReplaceTerminators(Region* region, Block* target_block,
+                                 Location loc,
+                                 const BlockAndValueMapping& mapper,
+                                 OpBuilder* builder) {
+  for (auto& old_block : region->getBlocks()) {
+    Block* block = mapper.lookup(&old_block);
+    auto return_op = dyn_cast<xla_hlo::ReturnOp>(block->getTerminator());
+    if (!return_op) continue;
+    builder->setInsertionPointToEnd(block);
+    builder->create<mlir::BranchOp>(loc, target_block, return_op.getOperands());
+    return_op.erase();
+  }
+
+  return success();
+}
+
+LogicalResult LowerIfOp(mlir::xla_hlo::IfOp if_op) {
+  Operation* op_inst = if_op.getOperation();
+  mlir::OpBuilder builder(if_op);
+  auto orig_block = op_inst->getBlock();
+  auto* tail_block = orig_block->splitBlock(op_inst);
+  auto loc = if_op.getLoc();
+
+  // Duplicate the true and false regions in the block between the sections
+  // before and after the conditional.
+  BlockAndValueMapping mapper;
+  if_op.true_branch().cloneInto(orig_block->getParent(),
+                                Region::iterator(tail_block), mapper);
+  if_op.false_branch().cloneInto(orig_block->getParent(),
+                                 Region::iterator(tail_block), mapper);
+
+  // Determine the blocks for the start of the true and false regions.
+  Block* true_block = mapper.lookup(&if_op.true_branch().front());
+  Block* false_block = mapper.lookup(&if_op.false_branch().front());
+
+  // Perform the conditional branch into the true/false cases.
+  builder.setInsertionPointToEnd(orig_block);
+
+  // Extract the predicate for checking branching, then branch to the true and
+  // false regions appropriately.
+  auto cond_value = builder.create<mlir::ExtractElementOp>(loc, if_op.pred());
+  builder.create<mlir::CondBranchOp>(loc, cond_value, true_block,
+                                     if_op.true_arg(), false_block,
+                                     if_op.false_arg());
+
+  // Replace the true case's return operations with a branch to the tail of
+  // the condition.
+  if (failed(ReplaceTerminators(&if_op.true_branch(), tail_block, loc, mapper,
+                                &builder)))
+    return failure();
+  if (failed(ReplaceTerminators(&if_op.false_branch(), tail_block, loc, mapper,
+                                &builder)))
+    return failure();
+
+  tail_block->addArguments(if_op.getResult().getType());
+  if_op.getResult().replaceAllUsesWith(tail_block->getArgument(0));
+
+  op_inst->erase();
+  return success();
+}
+
+LogicalResult LowerWhileOp(mlir::xla_hlo::WhileOp while_op) {
+  // Converts an XLA while loop into control flow. This generates a set of MLIR
+  // blocks and branches, along with inlining the regions provided by the XLA
+  // while loop. The structure should be similar to below:
+  //
+  //   <prior operations>
+  //   %0 = "xla_hlo.while"(%arg0) {^cond(...){...}, ^body(...){...}}
+  //   <post operations>
+  auto* op_inst = while_op.getOperation();
+  mlir::OpBuilder builder(while_op);
+  auto loc = while_op.getLoc();
+
+  // Break the block into four sections:
+  // orig_block - operations before the while and the branch into looping check.
+  // tail_block - operations after the while loop completes.
+  // cond_block - check the looping condition, then conditionally branch into
+  //              the loop or, if condition is false, jump to the tail branch.
+  // body_block - inlined loop body, then jump back to the condition block.
+  auto* orig_block = op_inst->getBlock();
+  auto* tail_block = orig_block->splitBlock(op_inst);
+
+  BlockAndValueMapping mapper;
+  while_op.cond().cloneInto(orig_block->getParent(),
+                            Region::iterator(tail_block), mapper);
+  while_op.body().cloneInto(orig_block->getParent(),
+                            Region::iterator(tail_block), mapper);
+
+  // Lookup the entry blocks for both condition and body.
+  auto* cond_block = mapper.lookup(&while_op.cond().front());
+  auto* body_block = mapper.lookup(&while_op.body().front());
+
+  // Setup the end of the original block:
+  //     <prior operations>
+  //     br ^cond(%arg0) // Jumps to the condition statement.
+  builder.setInsertionPointToEnd(orig_block);
+  builder.create<mlir::BranchOp>(loc, cond_block, while_op.getOperand());
+
+  // Updates the inlined condition blocks by replacing the return op with an
+  // extract_element and conditional branch. This changes the block below:
+  //   ^cond(%0):
+  //     <inlined conditional region>
+  //    "xla_hlo".return(%1)
+  //
+  //  Into:
+  //   ^cond(%0):
+  //     <inlined conditional region>
+  //     %2 = extract_element %1[] : tensor<i1> // Extract the condition value.
+  //     cond_br %2, ^body(%0), ^tail(%0) // Branch.
+  builder.setInsertionPointToStart(cond_block);
+
+  // Replace the xla_hlo::ReturnOp with a branch back to the condition block.
+  // This is required as the xla_hlo::ReturnOp is used to mark the end of a
+  // block for regions nested inside of a operations (MLIR ReturnOp cannot be
+  // nested within an non-function region).
+  for (auto& block : while_op.cond()) {
+    auto new_block = mapper.lookup(&block);
+
+    auto return_op = dyn_cast<xla_hlo::ReturnOp>(new_block->getTerminator());
+    if (!return_op) continue;
+    builder.setInsertionPointToEnd(new_block);
+
+    auto return_value = return_op.getOperand(0);
+    auto cond_value = builder.create<mlir::ExtractElementOp>(loc, return_value);
+
+    // Get the body block arguments.
+    llvm::SmallVector<Value, 4> successor_args(cond_block->args_begin(),
+                                               cond_block->args_end());
+    builder.create<mlir::CondBranchOp>(loc, cond_value, body_block,
+                                       successor_args, tail_block,
+                                       successor_args);
+    return_op.erase();
+  }
+
+  // Updates the body blocks by replace the return op with an branch to the
+  // conditional block. This changes the block below:
+  //   ^body(%0):
+  //     <inlined body block>
+  //    "xla_hlo".return(%1)
+  //
+  //  Into:
+  //   ^body(%0):
+  //     <inlined body block>
+  //     br ^cond(%0) // Branch.
+  for (auto& block : while_op.body()) {
+    auto new_block = mapper.lookup(&block);
+    auto return_op =
+        dyn_cast<mlir::xla_hlo::ReturnOp>(new_block->getTerminator());
+    if (!return_op) continue;
+    builder.setInsertionPointToEnd(new_block);
+    builder.create<mlir::BranchOp>(loc, cond_block, return_op.getOperands());
+    return_op.erase();
+  }
+
+  // Erase the original while loop.
+  tail_block->addArgument(while_op.getType());
+  while_op.getResult().replaceAllUsesWith(tail_block->getArgument(0));
+  op_inst->erase();
+
+  return success();
+}
+
+void LegalizeControlFlow::runOnFunction() {
+  auto func = getFunction();
+  llvm::SmallVector<IfOp, 4> if_ops;
+  func.walk([&](IfOp op) { if_ops.push_back(op); });
+
+  for (auto& op : if_ops) {
+    if (failed(LowerIfOp(op))) return signalPassFailure();
+  }
+
+  llvm::SmallVector<WhileOp, 4> while_ops;
+  func.walk([&](WhileOp op) { while_ops.push_back(op); });
+
+  for (auto& op : while_ops) {
+    if (failed(LowerWhileOp(op))) return signalPassFailure();
+  }
+}
+}  // namespace
+}  // namespace xla_hlo
+}  // namespace mlir
+
+std::unique_ptr<mlir::OperationPass<mlir::FuncOp>>
+mlir::xla_hlo::createLegalizeControlFlowPass() {
+  return std::make_unique<LegalizeControlFlow>();
+}
+
+static PassRegistration<mlir::xla_hlo::LegalizeControlFlow> legalize_cf_pass(
+    "xla-legalize-control-flow",
+    "Legalize from XLA control flow to MLIR control flow");

lib/Dialect/mhlo/transforms/legalize_tanh_to_approximation.cc

@@ -0,0 +1,156 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering the tanh standard ops to an
+// approximation.
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla {
+namespace {
+
+/// Emits the fast tanh approximation that is also used by XLA.
+Value EmitTanhApproximation(Value input, Location loc,
+                            PatternRewriter &rewriter) {
+  // For small values of x, we can approximate tanh(x)=x. For extremely small
+  // values of x (|x| < 1e-37), the other approximation would evaluate
+  // tanh(x) = 0.
+  constexpr float kCanUseApprox = 0.0004;
+  Value abs_value = rewriter.create<AbsFOp>(loc, input);
+  Value can_use_approx =
+      rewriter.create<ConstantOp>(loc, rewriter.getF32FloatAttr(kCanUseApprox));
+  Value return_input = rewriter.create<CmpFOp>(loc, CmpFPredicate::OLT,
+                                               abs_value, can_use_approx);
+  // Clamp the input to [-c, c].
+  Value max_clamp = rewriter.create<ConstantOp>(
+      loc, rewriter.getF32FloatAttr(7.90531110763549805f));
+  Value smaller_than_max =
+      rewriter.create<CmpFOp>(loc, CmpFPredicate::ULE, input, max_clamp);
+  Value clamped_half =
+      rewriter.create<SelectOp>(loc, smaller_than_max, input, max_clamp);
+  Value min_clamp = rewriter.create<ConstantOp>(
+      loc, rewriter.getF32FloatAttr(-7.90531110763549805f));
+  Value larger_than_min =
+      rewriter.create<CmpFOp>(loc, CmpFPredicate::UGE, clamped_half, min_clamp);
+  Value input_clamped =
+      rewriter.create<SelectOp>(loc, larger_than_min, clamped_half, min_clamp);
+
+  static constexpr std::array<float, 7> numerator_coeffs{
+      -2.76076847742355e-16f, 2.00018790482477e-13f, -8.60467152213735e-11f,
+      5.12229709037114e-08f,  1.48572235717979e-05f, 6.37261928875436e-04f,
+      4.89352455891786e-03f};
+
+  static constexpr std::array<float, 4> denominator_coeffs{
+      1.19825839466702e-06f, 1.18534705686654e-04f, 2.26843463243900e-03f,
+      4.89352518554385e-03f};
+
+  Value input_squared =
+      rewriter.create<MulFOp>(loc, input_clamped, input_clamped);
+  Value numerator = rewriter.create<ConstantOp>(
+      loc, rewriter.getF32FloatAttr(numerator_coeffs[0]));
+  for (int i = 1; i < numerator_coeffs.size(); i++) {
+    numerator = rewriter.create<AddFOp>(
+        loc, rewriter.create<MulFOp>(loc, input_squared, numerator),
+        rewriter.create<ConstantOp>(
+            loc, rewriter.getF32FloatAttr(numerator_coeffs[i])));
+  }
+
+  numerator = rewriter.create<MulFOp>(loc, input_clamped, numerator);
+
+  Value denominator = rewriter.create<ConstantOp>(
+      loc, rewriter.getF32FloatAttr(denominator_coeffs[0]));
+  for (int i = 1; i < denominator_coeffs.size(); i++) {
+    denominator = rewriter.create<AddFOp>(
+        loc, rewriter.create<MulFOp>(loc, input_squared, denominator),
+        rewriter.create<ConstantOp>(
+            loc, rewriter.getF32FloatAttr(denominator_coeffs[i])));
+  }
+
+  Value approx = rewriter.create<DivFOp>(loc, numerator, denominator);
+
+  return rewriter.create<SelectOp>(loc, return_input, input, approx);
+}
+
+class ApproximateTanhLowering : public OpRewritePattern<TanhOp> {
+ public:
+  explicit ApproximateTanhLowering(MLIRContext *ctx)
+      : OpRewritePattern<TanhOp>(ctx, 100) {}
+
+  LogicalResult matchAndRewrite(TanhOp tanhOp,
+                                PatternRewriter &rewriter) const override {
+    Type operand_type = tanhOp.getType();
+
+    if (operand_type.isF64()) {
+      // Similar to XLA, do not rewrite f64 as precision might matter.
+      return failure();
+    }
+
+    Location loc = tanhOp.getLoc();
+    Value input = tanhOp.operand();
+    if (operand_type.isF16()) {
+      input = rewriter.create<FPExtOp>(loc, input, rewriter.getF32Type());
+    }
+
+    // If we still do not have f32, fail.
+    if (!input.getType().isF32()) {
+      return failure();
+    }
+
+    Value result = EmitTanhApproximation(input, loc, rewriter);
+
+    // Truncate back if needed.
+    if (operand_type.isF16()) {
+      result = rewriter.create<FPTruncOp>(loc, result, rewriter.getF16Type());
+    }
+
+    rewriter.replaceOp(tanhOp, {result});
+    return success();
+  }
+};
+
+struct LegalizeTanhToApproximation
+    : public PassWrapper<LegalizeTanhToApproximation, FunctionPass> {
+  /// Perform the lowering of standard dialect operations to approximations.
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    PopulateTanhToApproximationPatterns(&getContext(), &patterns);
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::FuncOp>>
+createLegalizeTanhToApproximationPass() {
+  return std::make_unique<LegalizeTanhToApproximation>();
+}
+
+void PopulateTanhToApproximationPatterns(mlir::MLIRContext *context,
+                                         OwningRewritePatternList *patterns) {
+  patterns->insert<ApproximateTanhLowering>(context);
+}
+
+static PassRegistration<LegalizeTanhToApproximation> legalize_pass(
+    "xla-legalize-tanh-to-approximation",
+    "Legalize tanh from standard dialect to an approximation");
+
+}  // namespace xla
+}  // namespace mlir

lib/Dialect/mhlo/transforms/legalize_to_standard.cc

@@ -0,0 +1,208 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering XLA dialect to Standard dialect.
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringSwitch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace {
+#include "third_party/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/generated_legalize_to_standard.inc"
+}  // end anonymous namespace
+namespace xla_hlo {
+namespace {
+
+class CompareIConvert : public OpRewritePattern<xla_hlo::CompareOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(xla_hlo::CompareOp op,
+                                PatternRewriter &rewriter) const override {
+    auto lhs = op.lhs();
+    auto rhs = op.rhs();
+    auto lhs_type = lhs.getType().cast<TensorType>();
+    auto rhs_type = rhs.getType().cast<TensorType>();
+
+    // Broadcasting not supported by this rewrite.
+    if (lhs_type.getShape() != rhs_type.getShape()) return failure();
+
+    if (!lhs_type.getElementType().isSignlessInteger() ||
+        !rhs_type.getElementType().isSignlessInteger())
+      return failure();
+
+    auto comparison_direction = op.comparison_direction();
+    auto compare_predicate =
+        llvm::StringSwitch<Optional<CmpIPredicate>>(comparison_direction)
+            .Case("EQ", CmpIPredicate::eq)
+            .Case("NE", CmpIPredicate::ne)
+            .Case("LT", CmpIPredicate::slt)
+            .Case("LE", CmpIPredicate::sle)
+            .Case("GT", CmpIPredicate::sgt)
+            .Case("GE", CmpIPredicate::sge)
+            .Default(llvm::None);
+
+    if (!compare_predicate.hasValue()) return failure();
+
+    rewriter.replaceOpWithNewOp<CmpIOp>(op, compare_predicate.getValue(), lhs,
+                                        rhs);
+    return success();
+  }
+};
+
+class CompareFConvert : public OpRewritePattern<xla_hlo::CompareOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(xla_hlo::CompareOp op,
+                                PatternRewriter &rewriter) const override {
+    auto lhs = op.lhs();
+    auto rhs = op.rhs();
+    auto lhs_type = lhs.getType().cast<TensorType>();
+    auto rhs_type = rhs.getType().cast<TensorType>();
+
+    // Broadcasting not supported by this rewrite.
+    if (lhs_type.getShape() != rhs_type.getShape()) return failure();
+
+    if (!lhs_type.getElementType().isa<FloatType>() ||
+        !rhs_type.getElementType().isa<FloatType>())
+      return failure();
+
+    auto comparison_direction = op.comparison_direction();
+    auto compare_predicate =
+        llvm::StringSwitch<Optional<CmpFPredicate>>(comparison_direction)
+            .Case("EQ", CmpFPredicate::OEQ)
+            .Case("NE", CmpFPredicate::UNE)
+            .Case("LT", CmpFPredicate::OLT)
+            .Case("LE", CmpFPredicate::OLE)
+            .Case("GT", CmpFPredicate::OGT)
+            .Case("GE", CmpFPredicate::OGE)
+            .Default(llvm::None);
+
+    if (!compare_predicate.hasValue()) return failure();
+
+    rewriter.replaceOpWithNewOp<CmpFOp>(op, compare_predicate.getValue(), lhs,
+                                        rhs);
+    return success();
+  }
+};
+
+// Replace IotaOp with an integer constant. A ConvertOp is added to
+// convert the integer constant to iota result type. For complex types, the real
+// part is replaced with the generated constant and the imaginary part is
+// replaced with zero tensor.
+class ConvertIotaOp : public OpRewritePattern<xla_hlo::IotaOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(xla_hlo::IotaOp op,
+                                PatternRewriter &rewriter) const override {
+    auto output_type = op.getType().cast<ShapedType>();
+    auto output_size = output_type.getNumElements();
+    auto dimension = op.iota_dimension().getSExtValue();
+    auto max_dim_size = output_type.getDimSize(dimension);
+
+    auto element_type = output_type.getElementType();
+    int bitwidth;
+
+    auto complex_ty = element_type.dyn_cast<ComplexType>();
+    Type int_or_float_ty = element_type;
+    if (complex_ty) int_or_float_ty = complex_ty.getElementType();
+
+    bitwidth = int_or_float_ty.getIntOrFloatBitWidth();
+    llvm::SmallVector<APInt, 10> values;
+    values.reserve(output_size);
+
+    int64_t increase_stride = output_size;
+    for (int i = 0; i <= dimension; i++) {
+      increase_stride /= output_type.getDimSize(i);
+    }
+
+    int64_t current_value = 0;
+    for (int i = 0; i < output_size; i++) {
+      int64_t value = (current_value / increase_stride) % max_dim_size;
+      values.push_back(APInt(bitwidth, value));
+      ++current_value;
+    }
+
+    auto int_shape_type = RankedTensorType::get(
+        output_type.getShape(),
+        IntegerType::get(bitwidth, rewriter.getContext()));
+    auto loc = op.getLoc();
+    auto integer_const = rewriter.create<mlir::ConstantOp>(
+        loc, DenseIntElementsAttr::get(int_shape_type, values));
+
+    auto int_or_float_shape_ty =
+        RankedTensorType::get(output_type.getShape(), int_or_float_ty);
+
+    auto iota_const =
+        rewriter.create<ConvertOp>(loc, int_or_float_shape_ty, integer_const);
+
+    // For int/float types we are done, replace op and return.
+    if (!complex_ty) {
+      rewriter.replaceOp(op, iota_const.getResult());
+      return success();
+    }
+
+    // For complex types, generate a constant tensor of zeroes for the imaginary
+    // part and use iota_const for real part.
+    auto zeroes = rewriter.create<mlir::ConstantOp>(
+        loc, DenseIntElementsAttr::get(int_shape_type, APInt(bitwidth, 0)));
+    auto imag_zeroes =
+        rewriter.create<ConvertOp>(loc, int_or_float_shape_ty, zeroes);
+    rewriter.replaceOpWithNewOp<xla_hlo::ComplexOp>(op, iota_const,
+                                                    imag_zeroes);
+    return success();
+  }
+};
+
+}  // end anonymous namespace
+
+namespace {
+struct LegalizeToStandard
+    : public PassWrapper<LegalizeToStandard, FunctionPass> {
+  /// Perform the lowering to Standard dialect.
+  void runOnFunction() override;
+};
+}  // end anonymous namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::FuncOp>> createLegalizeToStdPass() {
+  return std::make_unique<LegalizeToStandard>();
+}
+
+void PopulateXlaToStdPatterns(OwningRewritePatternList *patterns,
+                              mlir::MLIRContext *ctx) {
+  mlir::populateWithGenerated(ctx, patterns);
+  patterns->insert<CompareFConvert, CompareIConvert, ConvertIotaOp>(ctx);
+}
+
+/// Perform the lowering to standard dialect.
+void LegalizeToStandard::runOnFunction() {
+  OwningRewritePatternList patterns;
+  mlir::xla_hlo::PopulateXlaToStdPatterns(&patterns, &getContext());
+  applyPatternsAndFoldGreedily(getFunction(), patterns);
+}
+
+static PassRegistration<LegalizeToStandard> legalize_pass(
+    "xla-legalize-to-std", "Legalize from XLA dialect to standard dialect");
+
+}  // end namespace xla_hlo
+}  // end namespace mlir

lib/Dialect/mhlo/transforms/legalize_to_standard_patterns.td

@@ -0,0 +1,71 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the legalization pattern definition file for XLA to StandardOps.
+
+include "third_party/llvm/llvm-project/mlir/include/mlir/IR/OpBase.td"
+include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td"
+include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td"
+
+//===----------------------------------------------------------------------===//
+// Nullary op patterns.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(HLO_ConstOp ElementsAttr:$value),
+          (ConstantOp $value)>;
+
+//===----------------------------------------------------------------------===//
+// Binary op patterns.
+//===----------------------------------------------------------------------===//
+
+def IsSameSizePred : CPred<
+    "$0.getType().cast<ShapedType>().getShape() "
+    "== $1.getType().cast<ShapedType>().getShape()">;
+def IsSameSizeConstraint : Constraint<IsSameSizePred, "inputs are same size">;
+
+
+def : Pat<(HLO_AndOp HLO_PredTensor:$l, HLO_PredTensor:$r),
+          (AndOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_AddOp HLO_FpTensor:$l, HLO_FpTensor:$r),
+          (AddFOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_SubOp HLO_FpTensor:$l, HLO_FpTensor:$r),
+          (SubFOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_MulOp HLO_FpTensor:$l, HLO_FpTensor:$r),
+          (MulFOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_DivOp HLO_FpTensor:$l, HLO_FpTensor:$r),
+          (DivFOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_RemOp HLO_FpTensor:$l, HLO_FpTensor:$r),
+          (RemFOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_AddOp HLO_IntTensor:$l, HLO_IntTensor:$r),
+          (AddIOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_SubOp HLO_IntTensor:$l, HLO_IntTensor:$r),
+          (SubIOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_MulOp HLO_IntTensor:$l, HLO_IntTensor:$r),
+          (MulIOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_DivOp HLO_IntTensor:$l, HLO_IntTensor:$r),
+          (SignedDivIOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;
+def : Pat<(HLO_RemOp HLO_IntTensor:$l, HLO_IntTensor:$r),
+          (SignedRemIOp $l, $r),
+          [(IsSameSizeConstraint $l, $r)]>;

lib/Dialect/mhlo/transforms/lhlo_copy_removal.cc

@@ -0,0 +1,105 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements a pass to remove redundant LHLO copy operations.
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+// Removes LHLO copy operations that copy from allocated buffers to block
+// arguments. All uses of each buffer are replaced with the corresponding block
+// argument and the buffer is freed. Note that this pass only works in regions
+// with a single block.
+struct LhloCopyRemoval : mlir::PassWrapper<LhloCopyRemoval, OperationPass<>> {
+  void runOnOperation() override {
+    llvm::SmallVector<mlir::Operation*, 2> eraseList;
+    auto operation = getOperation();
+    operation->walk([&](mlir::xla_lhlo::CopyOp copyOp) {
+      // If this region contains more than one block, then ignore this copy
+      // operation.
+      if (copyOp.getParentRegion()->getBlocks().size() > 1) {
+        return;
+      }
+
+      mlir::Value fromOperand = copyOp.operand();
+      mlir::Value toOperand = copyOp.output();
+
+      // If the fromOperand value is a block argument or the toOperand
+      // value is not a block argument, then ignore this copy operation.
+      if (!fromOperand.getDefiningOp() || toOperand.getDefiningOp()) {
+        return;
+      }
+
+      // The copy operation removal is illegal if there is at least a single use
+      // of toOperand value that lies between the first use of fromOperand value
+      // and the copy operation.
+      auto fromOperandUsers = fromOperand.getUsers();
+      auto firstUser = *fromOperandUsers.begin();
+      for (auto op : fromOperandUsers) {
+        if (op->isBeforeInBlock(firstUser)) firstUser = op;
+      }
+      for (auto op : toOperand.getUsers()) {
+        if (op->isBeforeInBlock(copyOp) && firstUser->isBeforeInBlock(op)) {
+          return;
+        }
+      }
+
+      // TODO(DFKI): Use live variable analysis to solve aliasing issues among
+      // block arguments.
+
+      // Remove the associated alloc operation.
+      auto allocOp = fromOperand.getDefiningOp();
+      eraseList.push_back(allocOp);
+
+      // Iterate over all uses of the fromOperand to find the associated
+      // deallocOp (if any).
+      for (auto op : fromOperandUsers) {
+        if (isa<mlir::DeallocOp>(op)) {
+          eraseList.push_back(op);
+          break;
+        }
+      }
+
+      // Replace all uses of the fromOperand with the toOperand. This rewires
+      // all references pointing to the original alloc operation to the new
+      // target operation in order to safely remove the copy op.
+      fromOperand.replaceAllUsesWith(toOperand);
+      copyOp.erase();
+    });
+    for (auto op : eraseList) {
+      op->erase();
+    }
+  };
+};
+
+}  // namespace
+
+std::unique_ptr<Pass> createLhloCopyRemovalPass() {
+  return absl::make_unique<LhloCopyRemoval>();
+}
+
+static PassRegistration<LhloCopyRemoval> copy_removal_pass(
+    "lhlo-copy-removal", "Removes redundant LHLO copy operations");
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lhlo_fuse_linalg.cc

@@ -0,0 +1,151 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for fusing linalg ops obtained after LHLO
+// lowering.
+
+#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/ArrayRef.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/FoldUtils.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+using linalg::LinalgOp;
+
+class LhloFuseLinalg : public PassWrapper<LhloFuseLinalg, FunctionPass> {
+ public:
+  LhloFuseLinalg() = default;
+  LhloFuseLinalg(const LhloFuseLinalg&) {}
+  LhloFuseLinalg(bool use_parallel_loops, llvm::ArrayRef<unsigned> tile_sizes) {
+    tile_sizes_ = tile_sizes;
+    use_parallel_loops_.setValue(use_parallel_loops);
+  }
+
+  void runOnFunction() override {
+    auto func = getFunction();
+
+    // TODO(pifon): Remove assumption that the function has a single block.
+    if (!llvm::hasSingleElement(func)) {
+      emitError(func.getLoc(), "The function needs to have a single block.");
+      signalPassFailure();
+      return;
+    }
+
+    // The fusion in Linalg is currently possible only when the consumer op is
+    // tiled. In order to greedily fuse the ops, we have to start from the tiled
+    // root linalg ops, i.e. linalg ops that write to output buffers of the
+    // function or are returned in case of escaping allocations.
+    llvm::SmallDenseSet<Value> result_buffers;
+    for (auto func_arg : func.getArguments()) {
+      result_buffers.insert(func_arg);
+    }
+    for (auto& block : func) {
+      auto returnOp = mlir::dyn_cast<mlir::ReturnOp>(block.getTerminator());
+      if (!returnOp) continue;
+      for (auto operand : returnOp.getOperands()) {
+        result_buffers.insert(operand);
+      }
+    }
+    MLIRContext* ctx = func.getContext();
+    OpBuilder b(func);
+    OperationFolder folder(ctx);
+    func.walk([&](linalg::GenericOp generic_op) {
+      SmallVector<int64_t, 2> tile_sizes(tile_sizes_.begin(),
+                                         tile_sizes_.end());
+      if (tile_sizes.empty()) {
+        tile_sizes = SmallVector<int64_t, 2>(generic_op.getNumLoops(), 1);
+      }
+      auto op = cast<LinalgOp>(generic_op.getOperation());
+      for (const Value result : op.getOutputBuffers()) {
+        if (!result_buffers.count(result)) continue;
+        if (tileGenericOp(op, tile_sizes, &b)) {
+          generic_op.erase();
+          return;
+        }
+      }
+    });
+    auto patterns = linalg::getLinalgTilingCanonicalizationPatterns(ctx);
+    applyPatternsAndFoldGreedily(func, patterns);
+
+    // Fuse producers of tiled linalg ops.
+    llvm::SmallDenseSet<Operation*> erase_set;
+    SmallVector<Operation*, 8> linalg_ops;
+    func.walk([&](LinalgOp op) { linalg_ops.push_back(op); });
+    for (auto* op : llvm::reverse(linalg_ops)) {
+      for (unsigned id = 0, e = LinalgOp(op).getNumInputs(); id < e; ++id) {
+        linalg::Aliases aliases;
+        linalg::LinalgDependenceGraph graph(aliases, linalg_ops);
+        if (auto info = fuseProducerOf(b, op, id, graph, &folder)) {
+          auto originalOp = info->originalProducer.getOperation();
+          erase_set.insert(originalOp);
+          auto originalOpInLinalgOpsVector = std::find_if(
+              linalg_ops.begin(), linalg_ops.end(),
+              [&](const Operation* op) { return op == originalOp; });
+          *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();
+        }
+      }
+
+      auto patterns = linalg::getLinalgTilingCanonicalizationPatterns(ctx);
+      applyPatternsAndFoldGreedily(func, patterns);
+    }
+    for (auto* e : erase_set) e->erase();
+  }
+
+ private:
+  bool tileGenericOp(LinalgOp op, ArrayRef<int64_t> tile_sizes, OpBuilder* b) {
+    auto loopType = use_parallel_loops_
+                        ? linalg::LinalgTilingLoopType::ParallelLoops
+                        : linalg::LinalgTilingLoopType::Loops;
+    auto tiled_generic_op = linalg::tileLinalgOp(*b, op,
+                                                 linalg::LinalgTilingOptions()
+                                                     .setTileSizes(tile_sizes)
+                                                     .setLoopType(loopType));
+    return tiled_generic_op.hasValue();
+  }
+
+  Option<bool> use_parallel_loops_{
+      *this, "use-parallel-loops",
+      llvm::cl::desc(
+          "Tiles GenericOp consumer to parallel loops before linalg fusion"),
+      llvm::cl::init(false)};
+
+  ListOption<unsigned> tile_sizes_{
+      *this, "tile-sizes",
+      llvm::cl::desc(
+          "Tile sizes by which to tile linalg generic before linalg fusion"),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createLhloFuseLinalg(
+    bool use_parallel_loops, ArrayRef<unsigned> tile_sizes) {
+  return absl::make_unique<LhloFuseLinalg>(use_parallel_loops, tile_sizes);
+}
+
+static PassRegistration<LhloFuseLinalg> legalize_pass(
+    "lhlo-fuse-linalg",
+    "Greedily fuse linalg ops obtained after LHLO lowering.");
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lhlo_legalize_to_affine.cc

@@ -0,0 +1,161 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering LHLO dialect to Affine dialect.
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Location.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/map_xla_to_scalar_op.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+// Builds an affine loop nest iterating from zeros to "upper_bounds" with unit
+// steps, and populates the body of the innermost loop using "body_builder".
+static void BuildBoundedAffineLoopNest(
+    OpBuilder& builder, Location location, ArrayRef<int64_t> upper_bounds,
+    function_ref<void(OpBuilder&, Location, ValueRange)> body_builder) {
+  SmallVector<int64_t, 3> lower_bounds(upper_bounds.size(), /*Value=*/0);
+  SmallVector<int64_t, 3> steps(upper_bounds.size(), /*Value=*/1);
+  buildAffineLoopNest(builder, location, lower_bounds, upper_bounds, steps,
+                      body_builder);
+}
+
+struct DotOpConverter : public OpRewritePattern<DotOp> {
+  using OpRewritePattern<DotOp>::OpRewritePattern;
+
+  // Supports only rank-2 tensors for LHS and RHS.
+  LogicalResult matchAndRewrite(DotOp op,
+                                PatternRewriter& rewriter) const override {
+    Value lhs = op.lhs();
+    Value rhs = op.rhs();
+    MemRefType lhs_type = lhs.getType().cast<MemRefType>();
+    MemRefType rhs_type = rhs.getType().cast<MemRefType>();
+    Type element_type = lhs_type.getElementType();
+    ArrayRef<int64_t> shape_lhs = lhs_type.getShape();
+    ArrayRef<int64_t> shape_rhs = rhs_type.getShape();
+
+    if ((lhs_type.getRank() != 2) || (rhs_type.getRank() != 2)) {
+      return failure();
+    }
+
+    LogicalResult map_status = success();
+    auto body_builder = [&](OpBuilder& builder, Location loc, ValueRange ivs) {
+      SmallVector<Value, 2> lhs_indices{ivs[0], ivs[2]},
+          rhs_indices{ivs[2], ivs[1]}, result_indices{ivs[0], ivs[1]};
+
+      auto l = builder.create<AffineLoadOp>(loc, lhs, lhs_indices);
+      auto r = builder.create<AffineLoadOp>(loc, rhs, rhs_indices);
+      auto result =
+          rewriter.create<AffineLoadOp>(loc, op.output(), result_indices);
+      Value op_result = xla_lhlo::XlaOpToStdScalarOp::map<DotOp>(
+          op, element_type, {l, r, result}, &builder);
+      map_status = success(op_result != nullptr);
+      if (failed(map_status)) return;
+      builder.create<AffineStoreOp>(loc, op_result, op.output(),
+                                    result_indices);
+    };
+
+    BuildBoundedAffineLoopNest(rewriter, op.getLoc(),
+                               {shape_lhs[0], shape_rhs[1], shape_rhs[0]},
+                               body_builder);
+    if (failed(map_status)) return failure();
+
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+template <typename LhloOpTy>
+struct BinaryOpConverter : public OpRewritePattern<LhloOpTy> {
+  using OpRewritePattern<LhloOpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(LhloOpTy op,
+                                PatternRewriter& rewriter) const override {
+    const auto& lhs = op.lhs();
+    const auto& rhs = op.rhs();
+    const auto& lhs_type = lhs.getType().template cast<MemRefType>();
+    const auto& rhs_type = rhs.getType().template cast<MemRefType>();
+    const auto& element_type = lhs_type.getElementType();
+
+    if (lhs_type.getShape() != rhs_type.getShape()) {
+      return failure();
+    }
+
+    LogicalResult map_status = success();
+    auto body_builder = [&](OpBuilder& builder, Location loc,
+                            ValueRange induction_vars) {
+      auto l = builder.create<AffineLoadOp>(loc, lhs, induction_vars);
+      auto r = builder.create<AffineLoadOp>(loc, rhs, induction_vars);
+      Value op_result = xla_lhlo::XlaOpToStdScalarOp::map<LhloOpTy>(
+          op, element_type, {l, r}, &builder);
+      map_status = success(op_result != nullptr);
+      if (failed(map_status)) return;
+      rewriter.create<AffineStoreOp>(loc, op_result, op.out(), induction_vars);
+    };
+
+    BuildBoundedAffineLoopNest(rewriter, op.getLoc(), lhs_type.getShape(),
+                               body_builder);
+    if (failed(map_status)) return failure();
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+void populateLHLOToAffineConversionPattern(MLIRContext* context,
+                                           OwningRewritePatternList* patterns) {
+  // clang-format off
+  patterns->insert<
+      BinaryOpConverter<xla_lhlo::AddOp>,
+      BinaryOpConverter<xla_lhlo::AndOp>,
+      BinaryOpConverter<xla_lhlo::DivOp>,
+      BinaryOpConverter<xla_lhlo::MaxOp>,
+      BinaryOpConverter<xla_lhlo::MinOp>,
+      BinaryOpConverter<xla_lhlo::MulOp>,
+      BinaryOpConverter<xla_lhlo::SubOp>,
+      DotOpConverter>(context);
+  // clang-format on
+}
+
+struct LhloLegalizeToAffine
+    : public PassWrapper<LhloLegalizeToAffine, FunctionPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    auto func = getFunction();
+    populateLHLOToAffineConversionPattern(func.getContext(), &patterns);
+    applyPatternsAndFoldGreedily(func, patterns);
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeToAffinePass() {
+  return absl::make_unique<LhloLegalizeToAffine>();
+}
+
+static PassRegistration<LhloLegalizeToAffine> legalize_pass(
+    "lhlo-legalize-to-affine", "Legalize from LHLO dialect to affine dialect");
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lhlo_legalize_to_gpu.cc

@@ -0,0 +1,196 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering LHLO dialect to GPU dialect.
+
+#include <cstdint>
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/ArrayRef.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/GPU/GPUDialect.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/SCF/SCF.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BlockAndValueMapping.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Builders.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Location.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/map_xla_to_scalar_op.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+// A simple translation of LHLO reduce operations to a corresponding gpu
+// launch operation. The transformation does no tiling and also only supports
+// 1d results.
+class LhloReduceToGPULaunchConverter : public OpConversionPattern<ReduceOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      ReduceOp reduce_op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = reduce_op.getLoc();
+    // Only support 1d reductions for now.
+    int64_t size = 0;
+    for (auto result : reduce_op.out()) {
+      auto shaped_type = result.getType().dyn_cast<ShapedType>();
+      if (!shaped_type || shaped_type.getRank() != 1) {
+        return failure();
+      }
+      auto dim_size = shaped_type.getDimSize(0);
+      if (size && size != dim_size) {
+        return failure();
+      }
+      size = dim_size;
+    }
+
+    auto reducing_dimension = *reduce_op.dimensions().int_value_begin();
+
+    // Require all inputs to have the same shape.
+    int64_t reduce_dim_size = 0;
+    for (auto input : reduce_op.operands()) {
+      auto shaped_type = input.getType().dyn_cast<ShapedType>();
+      if (!shaped_type || !shaped_type.hasStaticShape()) {
+        return failure();
+      }
+      reduce_dim_size =
+          shaped_type.getDimSize(reducing_dimension.getSExtValue());
+    }
+
+    // Create a launch that is parallel in the result dimension.
+    auto block_size_x = rewriter.create<mlir::ConstantOp>(
+        loc, rewriter.getIndexType(),
+        rewriter.getIntegerAttr(rewriter.getIndexType(), size));
+    auto one = rewriter.create<mlir::ConstantOp>(
+        loc, rewriter.getIndexType(),
+        rewriter.getIntegerAttr(rewriter.getIndexType(), 1));
+    auto launch_op = rewriter.create<mlir::gpu::LaunchOp>(
+        loc, one, one, one, block_size_x, one, one);
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPointToEnd(&launch_op.body().front());
+      auto index = launch_op.getThreadIds().x;
+
+      // Load the initial value and store it to the output.
+      for (auto pair : llvm::zip(reduce_op.init_values(), reduce_op.out())) {
+        auto init_value = rewriter.create<mlir::LoadOp>(loc, std::get<0>(pair));
+        rewriter.create<mlir::StoreOp>(loc, init_value, std::get<1>(pair),
+                                       ArrayRef<Value>{index});
+      }
+
+      // Insert a loop into the body to compute the reduction. The loop ranges
+      // from [0.dim).
+      auto zero = rewriter.create<mlir::ConstantOp>(
+          loc, rewriter.getIndexType(),
+          rewriter.getIntegerAttr(rewriter.getIndexType(), 0));
+      // TODO(b/137624192) Use dimOp to make it shape independent.
+      auto upper = rewriter.create<mlir::ConstantOp>(
+          loc, rewriter.getIndexType(),
+          rewriter.getIntegerAttr(rewriter.getIndexType(), reduce_dim_size));
+      auto step = rewriter.create<mlir::ConstantOp>(
+          loc, rewriter.getIndexType(),
+          rewriter.getIntegerAttr(rewriter.getIndexType(), 1));
+      auto loop = rewriter.create<mlir::scf::ForOp>(loc, zero, upper, step);
+
+      rewriter.setInsertionPointToStart(loop.getBody());
+      // Compute memrefs for the value to reduce. This makes it easier to just
+      // inline the body.
+      auto output = *reduce_op.out().begin();
+      // TODO(herhut) Move this to the SliceOp builder.
+      auto resType = MemRefType::get(
+          llvm::None, output.getType().cast<MemRefType>().getElementType(),
+          makeStridedLinearLayoutMap(llvm::None,
+                                     MemRefType::getDynamicStrideOrOffset(),
+                                     rewriter.getContext()));
+      auto accumulator = rewriter.create<mlir::linalg::SliceOp>(
+          loc, resType, output, ArrayRef<Value>{launch_op.getThreadIds().x});
+      llvm::SmallVector<Value, 4> indexings;
+      auto input_buffer = *reduce_op.operands().begin();
+      auto input_type = input_buffer.getType().cast<MemRefType>();
+      for (int64_t dim = 0; dim < input_type.getRank(); ++dim) {
+        indexings.push_back(dim == reducing_dimension
+                                ? loop.getInductionVar()
+                                : launch_op.getThreadIds().x);
+      }
+      // TODO(herhut) Move this to the SliceOp builder.
+      auto input = *reduce_op.operand_begin();
+      auto rhs = rewriter.create<mlir::linalg::SliceOp>(
+          loc,
+          MemRefType::get(
+              llvm::None, input_type.getElementType(),
+              makeStridedLinearLayoutMap(llvm::None,
+                                         MemRefType::getDynamicStrideOrOffset(),
+                                         rewriter.getContext())),
+          input, indexings);
+
+      // Now copy over the actual body of the reduction, leaving out the
+      // terminator.
+      BlockAndValueMapping mapping;
+      mapping.map(reduce_op.body().front().getArgument(0), accumulator);
+      mapping.map(reduce_op.body().front().getArgument(1), rhs);
+      mapping.map(reduce_op.body().front().getArgument(2), accumulator);
+      for (auto& nested : reduce_op.body().front().without_terminator()) {
+        auto clone = rewriter.clone(nested, mapping);
+        for (auto pair : llvm::zip(nested.getResults(), clone->getResults())) {
+          mapping.map(std::get<0>(pair), std::get<1>(pair));
+        }
+      }
+
+      // Finally, insert the terminator for the launchOp.
+      rewriter.setInsertionPointToEnd(&launch_op.body().front());
+      rewriter.create<mlir::gpu::TerminatorOp>(loc);
+    }
+
+    rewriter.eraseOp(reduce_op);
+    return success();
+  };
+};
+
+struct LhloLegalizeToGpu : public PassWrapper<LhloLegalizeToGpu, FunctionPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addLegalDialect<linalg::LinalgDialect, StandardOpsDialect,
+                           gpu::GPUDialect, scf::SCFDialect, XlaLhloDialect>();
+    target.addIllegalOp<ReduceOp>();
+    auto func = getFunction();
+    patterns.insert<LhloReduceToGPULaunchConverter>(func.getContext());
+    if (failed(applyPartialConversion(func, target, patterns))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeToGpuPass() {
+  return absl::make_unique<LhloLegalizeToGpu>();
+}
+
+static PassRegistration<LhloLegalizeToGpu> legalize_pass(
+    "lhlo-legalize-to-gpu", "Legalize from LHLO dialect to GPU dialect");
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lhlo_legalize_to_llvm.cc

@@ -0,0 +1,136 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+struct StaticMemRefCastOpConverter
+    : public ConvertOpToLLVMPattern<StaticMemRefCastOp> {
+  using ConvertOpToLLVMPattern<StaticMemRefCastOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult matchAndRewrite(
+      Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override {
+    auto loc = op->getLoc();
+    auto cast_op = cast<StaticMemRefCastOp>(op);
+
+    StaticMemRefCastOp::Adaptor operands_adaptor(operands);
+    MemRefDescriptor sourceMemRef(operands_adaptor.operand());
+
+    MemRefType targetMemRefType =
+        cast_op.getResult().getType().cast<MemRefType>();
+    auto llvmTargetDescriptorTy = typeConverter.convertType(targetMemRefType)
+                                      .dyn_cast_or_null<LLVM::LLVMType>();
+    if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy())
+      return failure();
+    // Create descriptor.
+    auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
+    Type llvmTargetElementTy = desc.getElementType();
+    // Set allocated ptr.
+    Value allocated = sourceMemRef.allocatedPtr(rewriter, loc);
+    allocated =
+        rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated);
+    desc.setAllocatedPtr(rewriter, loc, allocated);
+    // Set aligned ptr.
+    Value ptr = sourceMemRef.alignedPtr(rewriter, loc);
+    ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr);
+    desc.setAlignedPtr(rewriter, loc, ptr);
+
+    // Fill size and stride descriptors in memref.
+    auto target_sizes = targetMemRefType.getShape();
+    int64_t target_offset;
+    llvm::SmallVector<int64_t, 4> target_strides;
+    if (failed((getStridesAndOffset(targetMemRefType, target_strides,
+                                    target_offset))))
+      return failure();
+
+    // Copy offset of `targetMemRef`.
+    desc.setConstantOffset(rewriter, loc, target_offset);
+    for (int i = 0, e = targetMemRefType.getRank(); i < e; ++i) {
+      desc.setConstantSize(rewriter, loc, i, target_sizes[i]);
+      desc.setConstantStride(rewriter, loc, i, target_strides[i]);
+    }
+    rewriter.replaceOp(op, {desc});
+    return success();
+  }
+};
+
+struct DynamicMemRefCastOpConverter
+    : public ConvertOpToLLVMPattern<DynamicMemRefCastOp> {
+  using ConvertOpToLLVMPattern<DynamicMemRefCastOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult matchAndRewrite(
+      Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override {
+    auto loc = op->getLoc();
+    auto cast_op = cast<DynamicMemRefCastOp>(op);
+
+    DynamicMemRefCastOp::Adaptor operands_adaptor(operands);
+    MemRefDescriptor sourceMemRef(operands_adaptor.operand());
+
+    MemRefType targetMemRefType =
+        cast_op.getResult().getType().cast<MemRefType>();
+    auto llvmTargetDescriptorTy = typeConverter.convertType(targetMemRefType)
+                                      .dyn_cast_or_null<LLVM::LLVMType>();
+    if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy())
+      return failure();
+    // Create descriptor.
+    auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
+    Type llvmTargetElementTy = desc.getElementType();
+    // Set allocated ptr.
+    Value allocated = sourceMemRef.allocatedPtr(rewriter, loc);
+    allocated =
+        rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated);
+    desc.setAllocatedPtr(rewriter, loc, allocated);
+    // Set aligned ptr.
+    Value ptr = sourceMemRef.alignedPtr(rewriter, loc);
+    ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr);
+    desc.setAlignedPtr(rewriter, loc, ptr);
+    // Copy offset of `sourceMemRef`.
+    desc.setOffset(rewriter, loc, sourceMemRef.offset(rewriter, loc));
+
+    // Fill size and stride descriptors in memref.
+    if (!cast_op.sizes().empty()) {
+      auto sizes = operands_adaptor.sizes();
+      auto strides = operands_adaptor.strides();
+      for (int i = 0, e = targetMemRefType.getRank(); i < e; ++i) {
+        desc.setSize(rewriter, loc, i, sizes[i]);
+        desc.setStride(rewriter, loc, i, strides[i]);
+      }
+    }
+    rewriter.replaceOp(op, {desc});
+    return success();
+  }
+};
+
+}  // namespace
+
+void PopulateLhloToLLVMConversionPatterns(const LowerToLLVMOptions &options,
+                                          LLVMTypeConverter *converter,
+                                          OwningRewritePatternList *patterns) {
+  patterns->insert<DynamicMemRefCastOpConverter, StaticMemRefCastOpConverter>(
+      *converter, options);
+}
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lhlo_legalize_to_llvm_pass.cc

@@ -0,0 +1,59 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+class TestLhloToLLVMPass
+    : public ::mlir::PassWrapper<TestLhloToLLVMPass,
+                                 ::mlir::OperationPass<::mlir::ModuleOp>> {
+ public:
+  void runOnOperation() override {
+    ModuleOp m = getOperation();
+
+    OwningRewritePatternList patterns;
+    LLVMTypeConverter converter(m.getContext());
+    populateStdToLLVMConversionPatterns(converter, patterns);
+    PopulateLhloToLLVMConversionPatterns(
+        LowerToLLVMOptions::getDefaultOptions(), &converter, &patterns);
+
+    ConversionTarget target(getContext());
+    target.addLegalDialect<LLVM::LLVMDialect>();
+    target.addLegalOp<ModuleOp, ModuleTerminatorOp>();
+    target.addIllegalDialect<XlaLhloDialect>();
+
+    if (failed(applyFullConversion(m, target, patterns))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+static PassRegistration<TestLhloToLLVMPass> legalize_lhlo_pass(
+    "test-lhlo-legalize-to-llvm", "Legalize from LHLO dialect to LLVM.");
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lhlo_legalize_to_parallel_loops.cc

@@ -0,0 +1,731 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/ArrayRef.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/SmallVector.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/SCF/SCF.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+
+namespace mlir {
+namespace xla_lhlo {
+namespace {
+
+// Clones and adapts the code in `lhlo_block` that works on buffers and has a
+// single output buffer to make it compatible with `operands` that have element
+// types of the respective buffers. Returns the computed value.
+//
+// Example. For `operands` with (f32, i32) types and a block with LHLO ops and
+// with signature:
+//   ^bb(%lhs: memref<f32>, %rhs: memref<i32>, %res: memref<i1>):
+//     <LHLO_ops>
+//
+// inserts necessary alloc and store ops to compute and return result that has
+// `i1` type.
+Value ApplySingleResultLhloCode(Location loc, ValueRange operands,
+                                Block* lhlo_block, OpBuilder* b) {
+  SmallVector<Value, 2> arg_bufs;
+  for (auto arg_type : lhlo_block->getArgumentTypes()) {
+    arg_bufs.push_back(b->create<AllocOp>(loc, arg_type.cast<MemRefType>()));
+  }
+  for (auto operand : llvm::enumerate(operands)) {
+    b->create<StoreOp>(loc, operand.value(), arg_bufs[operand.index()]);
+  }
+  // Clone the ops from `lhlo_block`.
+  BlockAndValueMapping mapping;
+  mapping.map(lhlo_block->getArguments(), arg_bufs);
+  for (auto& nested : lhlo_block->without_terminator()) {
+    auto clone = b->clone(nested, mapping);
+    mapping.map(nested.getResults(), clone->getResults());
+  }
+  return b->create<LoadOp>(loc, arg_bufs.back());
+}
+
+// Converts a block with LHLO ops and with signature:
+//   ^bb(%lhs: memref<f32>, %rhs: memref<f32>, %res: memref<f32>):
+// into a reduction operator of scf.reduce by doing buffer allocation for
+// scalar arguments and the result of `scf.reduce` to make it compatible with
+// LHLO ops.
+void ConvertToReductionOperator(Location loc, scf::ReduceOp reduce_op,
+                                Block* lhlo_block, OpBuilder* b) {
+  Block& loop_reduce_op_body = reduce_op.reductionOperator().front();
+  OpBuilder::InsertionGuard guard(*b);
+  b->setInsertionPointToStart(&loop_reduce_op_body);
+  b->create<scf::ReduceReturnOp>(
+      loc, ApplySingleResultLhloCode(loc, loop_reduce_op_body.getArguments(),
+                                     lhlo_block, b));
+}
+
+// Returns result of ConstantOp if `dim` is static, otherwise uses DimOp to
+// extract dimension at runtime.
+Value GetStaticOrDynamicDim(mlir::Location loc, Value shaped_value,
+                            size_t dim_index, int64_t dim, OpBuilder* b) {
+  return dim == ShapedType::kDynamicSize
+             ? b->create<DimOp>(loc, shaped_value, dim_index).getResult()
+             : b->create<ConstantIndexOp>(loc, dim);
+}
+
+struct MappedIvs {
+  // False if the mapped indices are in the padding area, true otherwise.
+  Value in_bounds;
+  // Mapped indices.
+  SmallVector<Value, 2> ivs;
+};
+
+template <typename OpTy>
+MappedIvs MapWindowIvsToInput(OpTy op, ValueRange ivs, ValueRange window_ivs,
+                              OpBuilder* b) {
+  MappedIvs mapped_ivs;
+
+  if (!op.window_strides().hasValue()) {
+    op.emitOpError("No window strides specified.");
+  }
+  auto window_strides = op.window_strides().getValue();
+
+  if (!op.padding().hasValue()) {
+    op.emitOpError("No padding specified.");
+  }
+  auto padding = op.padding().getValue();
+
+  auto loc = op.getLoc();
+  auto operand = op.operand();
+  auto operand_shape = operand.getType().template cast<MemRefType>().getShape();
+
+  // `in_bounds` is false when the mapped indices are in the padding area.
+  mapped_ivs.in_bounds = b->create<mlir::ConstantOp>(
+      loc, b->getI1Type(), b->getIntegerAttr(b->getI1Type(), 1));
+  for (unsigned i = 0, e = ivs.size(); i < e; ++i) {
+    auto stride = window_strides.template getValue<llvm::APInt>(i);
+    auto pad_low = padding.template getValue<llvm::APInt>({i, 0});
+
+    Value stride_val = b->create<ConstantIndexOp>(loc, stride.getSExtValue());
+    Value pad_low_val = b->create<ConstantIndexOp>(loc, pad_low.getSExtValue());
+
+    Value center = b->create<MulIOp>(loc, ivs[i], stride_val);
+    Value offset = b->create<SubIOp>(loc, window_ivs[i], pad_low_val);
+    Value index = b->create<AddIOp>(loc, center, offset);
+    Value upper_bound =
+        GetStaticOrDynamicDim(loc, operand, i, operand_shape[i], b);
+    // We must check whether 0 <= index_i < shape_i, as otherwise we are in
+    // the pad and then we have to use the neutral element for reduction.
+    // Equivalently, it can be computed as the unsigned comparison index_i <
+    // shape_i, since a negative value wraps to a large positive value.
+    mapped_ivs.in_bounds = b->create<mlir::AndOp>(
+        loc, mapped_ivs.in_bounds,
+        b->create<CmpIOp>(loc, CmpIPredicate::ult, index, upper_bound));
+    mapped_ivs.ivs.push_back(index);
+  }
+  return mapped_ivs;
+}
+
+// Returns scf::Parallel over a shaped value with static or dynamic shape.
+scf::ParallelOp MakeLoopOverShape(Location loc, Value shaped_value,
+                                  OpBuilder* b) {
+  Value zero = b->create<ConstantIndexOp>(loc, 0);
+  Value one = b->create<ConstantIndexOp>(loc, 1);
+
+  ArrayRef<int64_t> shape =
+      shaped_value.getType().cast<ShapedType>().getShape();
+  SmallVector<Value, 2> lower, upper, step;
+  for (auto dim : llvm::enumerate(shape)) {
+    upper.push_back(
+        GetStaticOrDynamicDim(loc, shaped_value, dim.index(), dim.value(), b));
+    lower.push_back(zero);
+    step.push_back(one);
+  }
+  return b->create<scf::ParallelOp>(loc, lower, upper, step);
+}
+
+// Converts `xla_lhlo.ReduceOp` into two scf::ParallelOp and a scf::ReduceOp.
+// The outper `ParallelOp` refers to the parallel loops if there are
+// any. The inner `ParalleOp` refers to the reduction loops and `ReduceOp`
+// contains the reduction operator.
+//
+// Example:
+//
+//  "xla_lhlo.reduce"(%buffer, %init_buf, %result) ( {
+//    ^bb0(%lhs: memref<f32>, %rhs: memref<f32>, %res: memref<f32>):
+//      <LHLO ops>
+//    } ) {dimensions = dense<[1]> : tensor<1xi64>}
+//      : (memref<100x10x5xf32>, memref<f32>, memref<100x5xf32>) -> ()
+//
+//  is roughly converted into:
+//
+//  %init = load %init_buf[] : memref<f32>
+//  scf.parallel (%i, %k) = (%c0, %c0) to (%c100, %c5) step (%c1, %c1) {
+//    %result = scf.parallel (%j) = (%c0) to (%c10) step (%c1) init (%init) {
+//      %elem_to_reduce = load %buffer[%i, %j, %k] : memref<100x10x5xf32>
+//      scf.reduce(%elem_to_reduce)  {
+//        ^bb0(%elem: f32, %acc: f32):   // no predecessors
+//          elem_buf = alloc() : memref<f32>
+//          store %elem, elem_buf[] : memref<f32>
+//          acc_buf = alloc() : memref<f32>
+//          store %acc, acc_buf[] : memref<f32>
+//          <LHLO_ops>
+//          %acc_result = load acc_buf[] : memref<f32>
+//          scf.reduce.return %acc_result : f32
+//      } : f32
+//      scf.yield
+//    } : f32
+//    scf.yield
+//  }
+class ReduceOpConverter : public OpConversionPattern<xla_lhlo::ReduceOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::ReduceOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::ReduceOp xla_reduce_op, ArrayRef<Value> /*args*/,
+      ConversionPatternRewriter& rewriter) const final {
+    // TODO(b/137624192) Implement variadic reduce.
+    if (xla_reduce_op.out().size() != 1) return failure();
+
+    scf::ReduceOp reduce_op =
+        CreateReduceOpInNestedParallelLoops(xla_reduce_op, &rewriter);
+    ConvertToReductionOperator(xla_reduce_op.getLoc(), reduce_op,
+                               &xla_reduce_op.body().front(), &rewriter);
+    rewriter.replaceOp(xla_reduce_op, llvm::None);
+    return success();
+  }
+
+ private:
+  // Creates nested `scf.parallel` ops with `scf.reduce`. The outer ParallelOp
+  // refers to the parallel dimensions of `xla_reduce_op` if any and the inner
+  // ParallelOp refers to the reduction dimensions. The scf.reduce op is
+  // returned.
+  //
+  // If the reduction argument is a memref<100x10x5xf32> and the
+  // reduction is performed along dimension 1 then this method will generate
+  //
+  //  %init = load %init_buf[] : memref<f32>
+  //  scf.parallel (%i, %k) = (%c0, %c0) to (%c100, %c5) step (%c1, %c1) {
+  //    %result = scf.parallel (%j) = (%c0) to (%c10) step (%c1) init (%init) {
+  //      %elem_to_reduce = load %buffer[%i, %j, %k] : memref<100x10x5xf32>
+  //      scf.reduce(%elem_to_reduce)  {
+  //        <THE BLOCK PTR TO BE RETURNED>
+  //      } : f32
+  //      scf.yield
+  //    } : f32
+  //    scf.yield
+  //  }
+  scf::ReduceOp CreateReduceOpInNestedParallelLoops(
+      xla_lhlo::ReduceOp xla_reduce_op,
+      ConversionPatternRewriter* rewriter) const {
+    auto loc = xla_reduce_op.getLoc();
+    DenseSet<int> reducing_dims;
+    for (const auto& rdim : xla_reduce_op.dimensions().getIntValues()) {
+      reducing_dims.insert(rdim.getSExtValue());
+    }
+
+    Value operand = *xla_reduce_op.operands().begin();
+    Value out = *xla_reduce_op.out().begin();
+    SmallVector<Value, 2> parallel_lower, parallel_upper, parallel_step;
+    SmallVector<Value, 2> reduce_lower, reduce_upper, reduce_step;
+    auto operand_shape = operand.getType().cast<MemRefType>().getShape();
+    for (auto dim : llvm::enumerate(operand_shape)) {
+      const bool is_reducing_dim = reducing_dims.count(dim.index());
+
+      Value ub = GetStaticOrDynamicDim(loc, operand, dim.index(), dim.value(),
+                                       rewriter);
+      Value lb = rewriter->create<ConstantIndexOp>(loc, 0);
+      Value step = rewriter->create<ConstantIndexOp>(loc, 1);
+      (is_reducing_dim ? reduce_lower : parallel_lower).push_back(lb);
+      (is_reducing_dim ? reduce_upper : parallel_upper).push_back(ub);
+      (is_reducing_dim ? reduce_step : parallel_step).push_back(step);
+    }
+    // Load initial value from memref<element_type>.
+    SmallVector<Value, 1> init_value = {
+        rewriter->create<LoadOp>(loc, *xla_reduce_op.init_values().begin())};
+    // Outer ParallelOp is not needed if it is a reduction across all dims.
+    scf::ParallelOp outer;
+    if (!parallel_lower.empty()) {
+      outer = rewriter->create<scf::ParallelOp>(loc, parallel_lower,
+                                                parallel_upper, parallel_step);
+      rewriter->setInsertionPointToStart(outer.getBody());
+    }
+    scf::ParallelOp inner = rewriter->create<scf::ParallelOp>(
+        loc, reduce_lower, reduce_upper, reduce_step, ValueRange(init_value));
+    Value reduction_result = *inner.getResults().begin();
+
+    SmallVector<Value, 1> out_indices;
+    if (outer != nullptr) {
+      out_indices.reserve(outer.getNumLoops());
+      for (Value iv : outer.getInductionVars()) {
+        out_indices.push_back(iv);
+      }
+    } else {
+      out_indices.push_back(rewriter->create<ConstantIndexOp>(loc, 0));
+    }
+
+    rewriter->create<StoreOp>(loc, reduction_result, out, out_indices);
+
+    // Load the element to reduce.
+    SmallVector<Value, 2> indices;
+    indices.reserve(operand_shape.size());
+
+    if (outer) {
+      auto inner_ivs_it = inner.getInductionVars().begin();
+      auto outer_ivs_it = outer.getInductionVars().begin();
+      for (unsigned i = 0, e = operand_shape.size(); i < e; ++i) {
+        indices.push_back(reducing_dims.count(i) ? *inner_ivs_it++
+                                                 : *outer_ivs_it++);
+      }
+    } else {
+      indices = inner.getInductionVars();
+    }
+
+    rewriter->setInsertionPointToStart(inner.getBody());
+    Value elem = rewriter->create<mlir::LoadOp>(
+        loc, *xla_reduce_op.operands().begin(), indices);
+    return rewriter->create<scf::ReduceOp>(loc, elem);
+  }
+};
+
+// Pseudocode:
+// for each index O in output
+//   accumulator = neutral_value
+//   in_bounds = true
+//   for each index W in window
+//     for each dimension i from 0 to rank - 1
+//       index = O[i] * stride[i] + W[i] - pad_low[i]
+//       in_bounds = inbounds && (index `ult` shape[i])
+//       I[i] = index
+//     if (in_bounds)
+//       value = input[I]
+//     else
+//       value = neutral_value
+//     accumulator = reduction_operator(output[O], value)
+//   output[O] = accumulator
+//
+// Converts `xla_lhlo.ReduceWindowOp` into two scf::ParallelOp and a
+// scf::ReduceOp.
+// The outper `ParallelOp` refers to the parallel loops that traverese output
+// buffer. The inner `ParalleOp` refers to the reduction loops that traverse
+// reduction windows and `ReduceOp` contains the reduction operator.
+//
+// Example:
+//
+// func @reduce_window(%arg: memref<112x112xf32>,
+//              %init: memref<f32>,
+//              %result: memref<56x56xf32>) {
+//   "xla_lhlo.reduce_window"(%arg, %init, %result) ( {
+//     ^bb0(%lhs: memref<f32>, %rhs: memref<f32>, %res: memref<f32>):
+//       "xla_lhlo.maximum"(%lhs, %rhs, %res)
+//         : (memref<f32>, memref<f32>, memref<f32>) -> ()
+//       "xla_lhlo.terminator"() : () -> ()
+//     }) {
+//       padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>,
+//       window_dimensions = dense<[3, 3]> : tensor<2xi64>,
+//       window_strides = dense<[2, 2]> : tensor<2xi64>
+//     } : (memref<112x112xf32>, memref<f32>, memref<56x56xf32>) -> ()
+//   return
+// }
+//
+// is roughly converted into:
+//
+//    %neutral_elem = load %init_buf[] : memref<f32>
+//    scf.parallel (%i, %j) = (%c0, %c0) to (%c56, %c56) step (%c1, %c1) {
+//      %result = scf.parallel (%iw, %jw) = (%c0, %c0)
+//                  to (%c3, %c3) step (%c1, %c1) neutral_elem (%0) -> f32 {
+//        %in_bounds = <COMPUTE IF INDEX IS IN OPERAND'S pad>
+//        %elem = load %operand[%computed_i, %computed_j]
+//        %elem_or_neutral = select %in_bounds, %elem, %neutral_elem : f32
+//        scf.reduce(%elem_to_reduce)  : f32 {
+//          ^bb0(%arg7: f32, %arg8: f32):
+//            <LHLO ops>
+//        }
+//        scf.yield
+//      }
+//      store %result, %output_buffer[%i, %j] : memref<56x56xf32>
+//      scf.yield
+//    }
+//    return
+//  }
+class ReduceWindowOpConverter
+    : public OpConversionPattern<xla_lhlo::ReduceWindowOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::ReduceWindowOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::ReduceWindowOp xla_reduce_window_op, ArrayRef<Value> /*args*/,
+      ConversionPatternRewriter& rewriter) const final {
+    scf::ParallelOp output_loop, window_loop;
+    std::tie(output_loop, window_loop) =
+        CreateParallelLoopsToTraverseOutputAndWindow(xla_reduce_window_op,
+                                                     &rewriter);
+
+    scf::ReduceOp reduce_op = CreateReduceOpInNestedParallelLoops(
+        xla_reduce_window_op, output_loop, window_loop, &rewriter);
+
+    ConvertToReductionOperator(xla_reduce_window_op.getLoc(), reduce_op,
+                               &xla_reduce_window_op.body().front(), &rewriter);
+    rewriter.replaceOp(xla_reduce_window_op, llvm::None);
+    return success();
+  }
+
+ private:
+  std::pair<scf::ParallelOp, scf::ParallelOp>
+  CreateParallelLoopsToTraverseOutputAndWindow(
+      xla_lhlo::ReduceWindowOp xla_reduce_window_op,
+      ConversionPatternRewriter* rewriter) const {
+    auto loc = xla_reduce_window_op.getLoc();
+    Value init_value =
+        rewriter->create<LoadOp>(loc, xla_reduce_window_op.init_value());
+
+    Value zero = rewriter->create<ConstantIndexOp>(loc, 0);
+    Value one = rewriter->create<ConstantIndexOp>(loc, 1);
+
+    // Create an outer parallel loop that spans the output of ReduceWindowOp.
+    Value xla_output = xla_reduce_window_op.out();
+    auto output_loop = MakeLoopOverShape(loc, xla_output, rewriter);
+
+    // Create a nested loop that traverses the window.
+    SmallVector<Value, 2> window_lower, window_upper, window_step;
+    rewriter->setInsertionPointToStart(output_loop.getBody());
+    for (const auto& window_dim : xla_reduce_window_op.window_dimensions()) {
+      window_step.push_back(one);
+      window_lower.push_back(zero);
+      window_upper.push_back(
+          rewriter->create<ConstantIndexOp>(loc, window_dim.getSExtValue()));
+    }
+    auto window_loop = rewriter->create<scf::ParallelOp>(
+        loc, window_lower, window_upper, window_step, ValueRange(init_value));
+
+    Value reduction_result = *window_loop.getResults().begin();
+    auto output_ivs = output_loop.getInductionVars();
+    rewriter->create<StoreOp>(loc, reduction_result, xla_output, output_ivs);
+    return std::make_pair(output_loop, window_loop);
+  }
+
+  scf::ReduceOp CreateReduceOpInNestedParallelLoops(
+      xla_lhlo::ReduceWindowOp xla_reduce_window_op,
+      scf::ParallelOp output_loop, scf::ParallelOp window_loop,
+      ConversionPatternRewriter* rewriter) const {
+    rewriter->setInsertionPointToStart(window_loop.getBody());
+    auto loc = xla_reduce_window_op.getLoc();
+
+    if (xla_reduce_window_op.base_dilations().hasValue() ||
+        xla_reduce_window_op.window_dilations().hasValue()) {
+      xla_reduce_window_op.emitRemark(
+          "Lowering to parallel loops does not support `base_dilations` or "
+          "`window_dilations` attributes yet. The attributes will be ignored.");
+    }
+
+    Value xla_operand = xla_reduce_window_op.operand();
+    auto xla_operand_type = xla_operand.getType().cast<MemRefType>();
+
+    // Compute ivs in 'arg' buffer and whether these ivs are in pad area or not.
+    MappedIvs mapped_ivs = MapWindowIvsToInput(
+        xla_reduce_window_op, output_loop.getInductionVars(),
+        window_loop.getInductionVars(), rewriter);
+
+    auto elem_or_init = rewriter->create<scf::IfOp>(
+        loc, xla_operand_type.getElementType(), mapped_ivs.in_bounds,
+        /*withElseRegion=*/true);
+
+    OpBuilder then_builder = elem_or_init.getThenBodyBuilder();
+    Value elem = then_builder.create<mlir::LoadOp>(
+        loc, xla_reduce_window_op.operand(), mapped_ivs.ivs);
+    then_builder.create<scf::YieldOp>(loc, elem);
+
+    OpBuilder else_builder = elem_or_init.getElseBodyBuilder();
+    else_builder.create<scf::YieldOp>(loc, *window_loop.initVals().begin());
+
+    return rewriter->create<scf::ReduceOp>(loc,
+                                           *elem_or_init.results().begin());
+  }
+};
+
+// See the operation semantics in
+// https://www.tensorflow.org/xla/operation_semantics#selectandscatter
+//
+// Pseudocode:
+//  scf.parallel(coordinates O in the output):
+//    output[O] = init
+//  scf.parallel(coordinates S in the source):
+//    selected_ivs = 0
+//    selected_val = 0
+//    initialized_flag = false
+//    scf.for (first dim W_1 in the window)
+//         iter_args (selected_ivs, selected_val, initialized_flag):
+//    ...
+//      scf.for (last dim W_N in the window):
+//           iter_args (selected_ivs, selected_val, initialized_flag):
+//        I = S * stride + W - pad_low
+//        if I within bounds of operand:
+//          if (initialized_flag):
+//            pred = select(selected_value, operand(I))):
+//            if (pred)
+//              selected_value = operand(I)
+//              selected_index = I
+//          else
+//              selected_value = operand(I)
+//              selected_index = I
+//              initialized_flag = true
+//    output(selected_index) = scatter(output(selected_index), source(S))
+class SelectAndScatterOpConverter
+    : public OpConversionPattern<xla_lhlo::SelectAndScatterOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::SelectAndScatterOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::SelectAndScatterOp s_and_s_op, ArrayRef<Value> /*args*/,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = s_and_s_op.getLoc();
+    InitializeOutput(s_and_s_op, &rewriter);
+    scf::ParallelOp loop_over_src =
+        MakeLoopOverShape(loc, s_and_s_op.source(), &rewriter);
+    rewriter.setInsertionPointToStart(loop_over_src.getBody());
+
+    // Compute indices of the selected element in the window.
+    auto selected_ivs = SelectIvs(s_and_s_op, loop_over_src, &rewriter);
+
+    // Load `source[selected_ivs]`.
+    auto src_elem = rewriter.create<LoadOp>(loc, s_and_s_op.source(),
+                                            loop_over_src.getInductionVars());
+
+    // Compute `out[selected_ivs]` = scatter(out[selected_ivs], src_element)`.
+    auto rmw = rewriter.create<GenericAtomicRMWOp>(loc, s_and_s_op.out(),
+                                                   selected_ivs);
+    OpBuilder rmw_builder = OpBuilder::atBlockEnd(rmw.getBody());
+    auto acc_result =
+        ApplySingleResultLhloCode(loc, {src_elem, rmw.getCurrentValue()},
+                                  &s_and_s_op.scatter().front(), &rmw_builder);
+    rmw_builder.create<AtomicYieldOp>(loc, acc_result);
+
+    rewriter.replaceOp(s_and_s_op, llvm::None);
+    return success();
+  }
+
+ private:
+  void InitializeOutput(xla_lhlo::SelectAndScatterOp s_and_s_op,
+                        OpBuilder* b) const {
+    auto loc = s_and_s_op.getLoc();
+    Value init_value = b->create<LoadOp>(loc, s_and_s_op.init_value());
+
+    scf::ParallelOp loop_over_output =
+        MakeLoopOverShape(loc, s_and_s_op.out(), b);
+    OpBuilder::InsertionGuard guard(*b);
+    b->setInsertionPointToStart(loop_over_output.getBody());
+    b->create<StoreOp>(loc, init_value, s_and_s_op.out(),
+                       loop_over_output.getInductionVars());
+  }
+
+  struct WindowLoops {
+    SmallVector<Value, 2> selected_ivs;
+    SmallVector<Value, 2> window_ivs;
+    scf::ForOp inner_loop;
+  };
+  WindowLoops InsertWindowLoops(xla_lhlo::SelectAndScatterOp s_and_s_op,
+                                scf::ParallelOp loop_over_src,
+                                OpBuilder* b) const {
+    auto loc = s_and_s_op.getLoc();
+    Value zero = b->create<ConstantIndexOp>(loc, 0);
+    Value one = b->create<ConstantIndexOp>(loc, 1);
+
+    auto element_type =
+        s_and_s_op.out().getType().cast<MemRefType>().getElementType();
+    auto rank = loop_over_src.getNumLoops();
+
+    // `iter_args` = [iv_1, ..., iv_N, selected_value, is_initialized]
+    SmallVector<Value, 4> iter_args(rank, zero);
+    iter_args.push_back(b->create<mlir::ConstantOp>(
+        loc, element_type, b->getFloatAttr(element_type, 0)));
+    iter_args.push_back(b->create<mlir::ConstantOp>(
+        loc, b->getI1Type(), b->getIntegerAttr(b->getI1Type(), 0)));
+
+    // Create a nested loop that traverses the window.
+    OpBuilder::InsertPoint ip;
+    WindowLoops result;
+    for (const auto& window_dim :
+         s_and_s_op.window_dimensions()->getIntValues()) {
+      Value upper = b->create<ConstantIndexOp>(loc, window_dim.getSExtValue());
+      result.inner_loop =
+          b->create<scf::ForOp>(loc, zero, upper, one, iter_args);
+      if (b->getInsertionBlock() == loop_over_src.getBody()) {
+        ip = b->saveInsertionPoint();
+        result.selected_ivs = result.inner_loop.getResults().take_front(rank);
+      } else {
+        b->create<scf::YieldOp>(loc, result.inner_loop.getResults());
+      }
+      b->setInsertionPointToStart(result.inner_loop.getBody());
+      iter_args = ValueRange{result.inner_loop.getRegionIterArgs()};
+      result.window_ivs.push_back(result.inner_loop.getInductionVar());
+    }
+    b->restoreInsertionPoint(ip);
+    return result;
+  }
+
+  // Adapter to store iteration arguments of sequential loops that perform
+  // select in a window.
+  class IterArgs {
+   public:
+    explicit IterArgs(ValueRange ivs_val_flag) : ivs_val_flag_(ivs_val_flag) {}
+    IterArgs(ValueRange ivs, Value value, Value flag) {
+      ivs_val_flag_ = ivs;
+      ivs_val_flag_.push_back(value);
+      ivs_val_flag_.push_back(flag);
+    }
+
+    ArrayRef<Value> to_vector() const { return ivs_val_flag_; }
+
+    // Indices of the currently selected value.
+    ArrayRef<Value> ivs() const { return to_vector().drop_back(2); }
+    // Currently selected value w.r.t. select() function.
+    Value value() const { return ivs_val_flag_.end()[-2]; }
+    // i1 flag if value() and ivs() were initialized.
+    Value is_init() const { return ivs_val_flag_.back(); }
+
+   private:
+    // Vector that stores iv_1, ..., iv_N, value, init.
+    SmallVector<Value, 4> ivs_val_flag_;
+  };
+
+  SmallVector<Value, 2> SelectIvs(xla_lhlo::SelectAndScatterOp s_and_s_op,
+                                  scf::ParallelOp loop_over_src,
+                                  OpBuilder* b) const {
+    auto loc = s_and_s_op.getLoc();
+
+    WindowLoops window_loops = InsertWindowLoops(s_and_s_op, loop_over_src, b);
+    auto inner_loop_b =
+        OpBuilder::atBlockEnd(window_loops.inner_loop.getBody());
+
+    // Compute ivs in 'arg' buffer and whether these ivs are in the pad area.
+    MappedIvs mapped_ivs =
+        MapWindowIvsToInput(s_and_s_op, loop_over_src.getInductionVars(),
+                            window_loops.window_ivs, &inner_loop_b);
+
+    IterArgs ivs_val_flag(window_loops.inner_loop.getRegionIterArgs());
+
+    auto if_in_bounds = inner_loop_b.create<scf::IfOp>(
+        loc, window_loops.inner_loop.getResultTypes(), mapped_ivs.in_bounds,
+        /*withElseRegion=*/true);
+
+    // Case when we are inside boundaries of 'arg' and not in the pad area.
+    {
+      OpBuilder in_bounds_then_b = if_in_bounds.getThenBodyBuilder();
+      auto select_or_init_results = SelectOrInitialize(
+          s_and_s_op, mapped_ivs.ivs, &ivs_val_flag, &in_bounds_then_b);
+      in_bounds_then_b.create<scf::YieldOp>(loc, select_or_init_results);
+    }
+
+    // Case when we are in the pad.
+    {
+      OpBuilder in_bounds_else_b = if_in_bounds.getElseBodyBuilder();
+      in_bounds_else_b.create<scf::YieldOp>(loc, ivs_val_flag.to_vector());
+    }
+
+    inner_loop_b.create<scf::YieldOp>(loc, if_in_bounds.getResults());
+    return window_loops.selected_ivs;
+  }
+
+  SmallVector<Value, 4> SelectOrInitialize(
+      xla_lhlo::SelectAndScatterOp s_and_s_op, ArrayRef<Value> operand_ivs,
+      IterArgs* ivs_val_flag, OpBuilder* b) const {
+    auto loc = s_and_s_op.getLoc();
+    Value true_i1 = b->create<mlir::ConstantOp>(
+        loc, b->getI1Type(), b->getIntegerAttr(b->getI1Type(), 1));
+
+    TypeRange iter_arg_types{ivs_val_flag->to_vector()};
+    Value operand_elem =
+        b->create<LoadOp>(loc, s_and_s_op.operand(), operand_ivs);
+    auto if_init =
+        b->create<scf::IfOp>(loc, iter_arg_types, ivs_val_flag->is_init(),
+                             /*withElseRegion=*/true);
+    // Init == true, i.e. iter args are already initialized with a selected
+    // element in boundaries of the operand. Select function has to be computed
+    // here.
+    {
+      OpBuilder if_init_then_b = if_init.getThenBodyBuilder();
+
+      auto& lhlo_select = s_and_s_op.select().front();
+      Value pred =
+          ApplySingleResultLhloCode(loc, {operand_elem, ivs_val_flag->value()},
+                                    &lhlo_select, &if_init_then_b);
+
+      auto if_pred = if_init_then_b.create<scf::IfOp>(loc, iter_arg_types, pred,
+                                                      /*withElseRegion=*/true);
+
+      // Pred == true, therefore pack newly selected ivs, val and init flag back
+      // to iter_args and return.
+      {
+        OpBuilder if_pred_then_b = if_pred.getThenBodyBuilder();
+        if_pred_then_b.create<scf::YieldOp>(
+            loc, IterArgs{operand_ivs, operand_elem, true_i1}.to_vector());
+      }
+
+      // Pred == false, therefore return old iter_args.
+      {
+        OpBuilder if_pred_else_b = if_pred.getElseBodyBuilder();
+        if_pred_else_b.create<scf::YieldOp>(loc, ivs_val_flag->to_vector());
+      }
+
+      if_init_then_b.create<scf::YieldOp>(loc, if_pred.getResults());
+    }
+    // Init == false, i.e. only pad was visited before and this is the first
+    // element in the boundaries of the operand.
+    {
+      OpBuilder if_init_else_b = if_init.getElseBodyBuilder();
+
+      if_init_else_b.create<scf::YieldOp>(
+          loc, IterArgs{operand_ivs, operand_elem, true_i1}.to_vector());
+    }
+    return if_init.getResults();
+  }
+};
+
+struct LhloLegalizeToParallelLoops
+    : public PassWrapper<LhloLegalizeToParallelLoops, FunctionPass> {
+  void runOnFunction() override {
+    auto func = getFunction();
+
+    OwningRewritePatternList patterns;
+    // clang-format off
+    patterns.insert<
+        ReduceOpConverter,
+        ReduceWindowOpConverter,
+        SelectAndScatterOpConverter
+      >(func.getContext());
+    // clang-format on
+
+    ConversionTarget target(getContext());
+    target.addLegalDialect<linalg::LinalgDialect, StandardOpsDialect,
+                           scf::SCFDialect, XlaLhloDialect>();
+    target.addIllegalOp<xla_lhlo::ReduceOp, xla_lhlo::ReduceWindowOp,
+                        xla_lhlo::SelectAndScatterOp>();
+
+    if (failed(applyPartialConversion(func, target, patterns))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeLhloToParallelLoopsPass() {
+  return absl::make_unique<LhloLegalizeToParallelLoops>();
+}
+
+static PassRegistration<LhloLegalizeToParallelLoops> legalize_lhlo_pass(
+    "lhlo-legalize-to-parallel-loops",
+    "Legalize from LHLO dialect to parallel loops.");
+
+}  // namespace xla_lhlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/lower_complex.cc

@@ -0,0 +1,79 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Thsi file implements passes to convert complex operations to equivalent real
+// value operations. This does not include removing complex values from function
+// argument or return types.
+
+#include <cstddef>
+#include <cstdint>
+#include <iterator>
+#include <numeric>
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/TypeUtilities.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Types.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/PassRegistry.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/utils/hlo_utils.h"
+
+using mlir::FunctionPass;
+using mlir::OwningRewritePatternList;
+using mlir::PassRegistration;
+using mlir::PassWrapper;
+
+namespace {
+class LowerComplex : public PassWrapper<LowerComplex, FunctionPass> {
+ public:
+  explicit LowerComplex() : PassWrapper<LowerComplex, FunctionPass>() {}
+
+  /// Performs the lowering to XLA dialect.
+  void runOnFunction() override;
+};
+}  // end anonymous namespace
+
+namespace mlir {
+namespace xla {
+namespace {
+
+#include "third_party/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/generated_lower_complex.inc"
+
+}  // end anonymous namespace
+
+void PopulateComplexLoweringPatterns(MLIRContext* context,
+                                     OwningRewritePatternList* patterns) {
+  populateWithGenerated(context, patterns);
+}
+}  // end namespace xla
+}  // end namespace mlir
+
+// Lowers the complex operations that can be represented using other operations.
+void LowerComplex::runOnFunction() {
+  // Add lowering patterns to the list.
+  OwningRewritePatternList patterns;
+  mlir::xla::PopulateComplexLoweringPatterns(&getContext(), &patterns);
+
+  applyPatternsAndFoldGreedily(getFunction(), patterns);
+}
+
+static PassRegistration<LowerComplex> pass(
+    "test-xla-lower-complex",
+    "Lower complex operations into non-complex operations");

lib/Dialect/mhlo/transforms/lower_complex_patterns.td

@@ -0,0 +1,109 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the legalization pattern that converts complex operations into
+// equivalent real value operations.
+
+include "third_party/llvm/llvm-project/mlir/include/mlir/IR/OpBase.td"
+include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td"
+include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td"
+
+//===----------------------------------------------------------------------===//
+// Binary op patterns.
+//===----------------------------------------------------------------------===//
+
+// Add and subtraction are elementwise and can be distributed across the real
+// and imaginary components.
+foreach elementwiseOp = [HLO_AddOp, HLO_SubOp] in
+  def : Pat<(elementwiseOp HLO_ComplexTensor:$lhs,
+             HLO_ComplexTensor:$rhs),
+            (HLO_ComplexOp
+              (elementwiseOp (HLO_RealOp $lhs), (HLO_RealOp $rhs)),
+              (elementwiseOp (HLO_ImagOp $lhs), (HLO_ImagOp $rhs)))>;
+
+// Complex multiplication results in a cross product multiplication between the
+// real and imaginary components such that:
+//   result.real = lhs.real * rhs.real - lhs.imag * rhs.imag
+//   result.imag = lhs.imag * rhs.real + lhs.real * rhs.imag
+def : Pat<(HLO_MulOp HLO_ComplexTensor:$lhs,
+           HLO_ComplexTensor:$rhs),
+          (HLO_ComplexOp
+           (HLO_SubOp
+            (HLO_MulOp
+             (HLO_RealOp:$lhs_real $lhs),
+             (HLO_RealOp:$rhs_real $rhs)),
+            (HLO_MulOp
+             (HLO_ImagOp:$lhs_imag $lhs),
+             (HLO_ImagOp:$rhs_imag $rhs))),
+           (HLO_AddOp
+            (HLO_MulOp $lhs_real, $rhs_imag),
+            (HLO_MulOp $lhs_imag, $rhs_real)))>;
+
+// Multiplication between a complex and real tensor can be distributed by
+// applying the real multiplicant to both the real and complex component.
+//
+// Note that the sourcep pattern is not legal according to the HLO dialect but
+// instead handle intermediates generated by other patterns.
+def : Pat<(HLO_MulOp HLO_ComplexTensor:$lhs, HLO_IntOrFpTensor:$rhs),
+          (HLO_ComplexOp
+           (HLO_MulOp (HLO_RealOp $lhs), $rhs),
+           (HLO_MulOp (HLO_ImagOp $lhs), $rhs))>;
+
+def : Pat<(HLO_MulOp HLO_IntOrFpTensor:$lhs, HLO_ComplexTensor:$rhs),
+          (HLO_ComplexOp
+           (HLO_MulOp $lhs, (HLO_RealOp $rhs)),
+           (HLO_MulOp $lhs, (HLO_ImagOp $rhs)))>;
+
+
+// Division is performed by normalizing the denominator by multiplying by the
+// conjugate of the rhs.
+//   numerator = lhs * conj(rhs)
+//   denominator = rhs * conj(rhs)
+def : Pat<(HLO_DivOp HLO_ComplexTensor:$lhs, HLO_ComplexTensor:$rhs),
+            (HLO_DivOp
+             (HLO_MulOp:$num $lhs,
+              (HLO_ComplexOp:$conj
+               (HLO_RealOp $rhs),
+               (HLO_NegOp (HLO_ImagOp $rhs)))),
+             (HLO_RealOp:$den (HLO_MulOp $rhs, $conj)))>;
+
+
+def : Pat<(HLO_DivOp HLO_ComplexTensor:$lhs, HLO_IntOrFpTensor:$rhs),
+          (HLO_ComplexOp
+           (HLO_DivOp (HLO_RealOp $lhs), $rhs),
+           (HLO_DivOp (HLO_ImagOp $lhs), $rhs))>;
+
+
+// Absolute value is evaluated as:
+//   result = sqrt(val.real * val.real + val.imag * val.imag)
+def : Pat<(HLO_AbsOp HLO_ComplexTensor:$val),
+          (HLO_ComplexOp
+           (HLO_SqrtOp
+             (HLO_AddOp
+              (HLO_MulOp (HLO_RealOp:$real $val), $real),
+              (HLO_MulOp (HLO_ImagOp:$imag $val), $imag))),
+           (HLO_ConstOp (ConstantSplat<"0"> $real)))>;
+
+// Exponential can be lowered to an exponential on the real component and a
+// sum of sinusoids of the imaginary component, which equates to a normal
+// exponential operator multiplied by Euler's formula.
+//
+// Exp(a + ib) = Exp(a) * Exp(ib) = Exp(a) * (Cos(b) + iSin(b))
+def : Pat<(HLO_ExpOp HLO_ComplexTensor:$val),
+          (HLO_MulOp
+           (HLO_ExpOp (HLO_RealOp $val)),
+           (HLO_ComplexOp
+            (HLO_CosOp (HLO_ImagOp:$imag $val)),
+            (HLO_SinOp $imag)))>;

lib/Dialect/mhlo/transforms/lower_general_dot.cc

@@ -0,0 +1,194 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering XLA general dot to a regular dot.
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringSwitch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Location.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/TypeUtilities.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+using mlir::DenseIntElementsAttr;
+using mlir::ElementsAttr;
+using mlir::failure;
+using mlir::FunctionPass;
+using mlir::LogicalResult;
+using mlir::MLIRContext;
+using mlir::OpRewritePattern;
+using mlir::OwningRewritePatternList;
+using mlir::PassRegistration;
+using mlir::PassWrapper;
+using mlir::PatternRewriter;
+using mlir::RankedTensorType;
+using mlir::success;
+using mlir::Value;
+
+namespace {
+
+Value TransposeReshape(Value arg, mlir::Location loc,
+                       llvm::ArrayRef<int64_t> left_dims,
+                       llvm::ArrayRef<int64_t> right_dims,
+                       llvm::ArrayRef<int64_t> arg_shape,
+                       PatternRewriter *rewriter) {
+  auto element_type = mlir::getElementTypeOrSelf(arg.getType());
+
+  int64_t left_size = 1;
+  for (auto dim : left_dims) {
+    left_size *= arg_shape[dim];
+  }
+
+  int64_t right_size = 1;
+  for (auto dim : right_dims) {
+    right_size *= arg_shape[dim];
+  }
+
+  // Generate the transpose permutation attribute.
+  llvm::SmallVector<int64_t, 5> transpose_permutation(left_dims.begin(),
+                                                      left_dims.end());
+  transpose_permutation.append(right_dims.begin(), right_dims.end());
+
+  mlir::TensorType transpose_permutation_type = RankedTensorType::get(
+      {static_cast<int64_t>(transpose_permutation.size())},
+      rewriter->getIntegerType(64));
+
+  auto transpose_permutation_attr =
+      DenseIntElementsAttr::get(transpose_permutation_type,
+                                llvm::makeArrayRef(transpose_permutation))
+          .cast<DenseIntElementsAttr>();
+
+  // Compute the resulting shape.
+  llvm::SmallVector<int64_t, 5> transposed_shape;
+  for (auto val : transpose_permutation) {
+    transposed_shape.push_back(arg_shape[val]);
+  }
+  auto transpose_type = RankedTensorType::get(transposed_shape, element_type);
+  auto transpose_result = rewriter->create<mlir::xla_hlo::TransposeOp>(
+      loc, transpose_type, arg, transpose_permutation_attr);
+
+  // Return the final result.
+  auto reshaped_type =
+      RankedTensorType::get({left_size, right_size}, element_type);
+  return rewriter->create<mlir::xla_hlo::ReshapeOp>(loc, reshaped_type,
+                                                    transpose_result);
+}
+
+Value ProcessDotArg(Value arg, mlir::Location loc,
+                    ElementsAttr contract_dims_attr, bool outer_dims_first,
+                    PatternRewriter *rewriter) {
+  auto shape = arg.getType().cast<mlir::ShapedType>().getShape();
+
+  llvm::SmallVector<bool, 5> is_outer_dim;
+  is_outer_dim.resize(shape.size(), true);
+
+  // Compute the contract dimension ordering.
+  llvm::SmallVector<int64_t, 5> contract_dims;
+  for (auto dim : contract_dims_attr.getValues<int64_t>()) {
+    contract_dims.push_back(dim);
+    is_outer_dim[dim] = false;
+  }
+
+  // Compute the outer dimension orderings.
+  llvm::SmallVector<int64_t, 5> outer_dims;
+  for (auto it : llvm::enumerate(is_outer_dim)) {
+    if (it.value()) {
+      outer_dims.push_back(it.index());
+    }
+  }
+
+  if (outer_dims_first) {
+    return TransposeReshape(arg, loc, outer_dims, contract_dims, shape,
+                            rewriter);
+  }
+
+  return TransposeReshape(arg, loc, contract_dims, outer_dims, shape, rewriter);
+}
+
+struct GeneralDotConvert
+    : public OpRewritePattern<mlir::xla_hlo::DotGeneralOp> {
+  // Attempts to lower a General Dot operator to a standard Dot operator.
+  // General dots include batching dimensions and can have collapsing
+  // dimensions along any axis. Inserting correctly arrange transpose and
+  // reshape operators organizes the tensors and allows the General Dot to be
+  // replaced with the standard Dot operator.
+  //
+  // Note: This requires an empty list of batch dimensions.
+
+  explicit GeneralDotConvert(MLIRContext *context)
+      : OpRewritePattern(context) {}
+
+  LogicalResult matchAndRewrite(mlir::xla_hlo::DotGeneralOp op,
+                                PatternRewriter &rewriter) const override {
+    auto dot_element_type = mlir::getElementTypeOrSelf(op);
+
+    auto dot_numbers = op.dot_dimension_numbers();
+    if (dot_numbers.lhs_batching_dimensions().getNumElements() != 0 ||
+        dot_numbers.rhs_batching_dimensions().getNumElements() != 0) {
+      return failure();
+    }
+
+    auto lhs = ProcessDotArg(op.lhs(), op.getLoc(),
+                             dot_numbers.lhs_contracting_dimensions(),
+                             /*outer_dims_first=*/true, &rewriter);
+
+    auto rhs = ProcessDotArg(op.rhs(), op.getLoc(),
+                             dot_numbers.rhs_contracting_dimensions(),
+                             /*outer_dims_first=*/false, &rewriter);
+
+    // Dot resulting shape.
+    auto lhs_shape = lhs.getType().cast<mlir::ShapedType>().getShape();
+    auto rhs_shape = rhs.getType().cast<mlir::ShapedType>().getShape();
+    auto new_dot_type =
+        RankedTensorType::get({lhs_shape[0], rhs_shape[1]}, dot_element_type);
+
+    auto new_dot_op = rewriter.create<mlir::xla_hlo::DotOp>(
+        op.getLoc(), new_dot_type, lhs, rhs, *(op.precision_config()));
+
+    rewriter.replaceOpWithNewOp<mlir::xla_hlo::ReshapeOp>(op, op.getType(),
+                                                          new_dot_op);
+    return success();
+  }
+};
+
+struct LegalizeGeneralDot
+    : public PassWrapper<LegalizeGeneralDot, FunctionPass> {
+  /// Lower all general dots that can be represented as a non-batched matmul.
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    mlir::xla_hlo::PopulateGeneralDotOpLoweringPatterns(&patterns,
+                                                        &getContext());
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+  }
+};
+
+}  // namespace
+
+void mlir::xla_hlo::PopulateGeneralDotOpLoweringPatterns(
+    OwningRewritePatternList *patterns, MLIRContext *ctx) {
+  patterns->insert<GeneralDotConvert>(ctx);
+}
+
+static PassRegistration<LegalizeGeneralDot> legalize_pass(
+    "test-xla-lower-general-dot",
+    "Tests lowering general dot to a non-batched dot when possible");

lib/Dialect/mhlo/transforms/materialize_broadcasts.cc

@@ -0,0 +1,90 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <numeric>
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace xla_hlo {
+
+namespace {
+
+// Converts ClampOp with broadcast semantics. ClampOp requires "all three arrays
+// must be the same shape. Alternatively, as a restricted form of broadcasting,
+// min and/or max can be a scalar of type T."
+struct ClampWithBroadcastConvert : public OpRewritePattern<ClampOp> {
+  explicit ClampWithBroadcastConvert(MLIRContext *context)
+      : OpRewritePattern<ClampOp>(context) {}
+
+  LogicalResult matchAndRewrite(ClampOp op,
+                                PatternRewriter &rewriter) const override {
+    auto operand_type = op.operand().getType().dyn_cast<RankedTensorType>();
+    auto max_type = op.max().getType().dyn_cast<RankedTensorType>();
+    auto min_type = op.min().getType().dyn_cast<RankedTensorType>();
+    // Unrancked types are not supported.
+    if (!operand_type || !max_type || !min_type) return failure();
+    // Does not support operand with dynamic dimensions for now.
+    if (!operand_type.hasStaticShape()) return failure();
+
+    ArrayRef<int64_t> operand_shape = operand_type.getShape();
+
+    Value max_value = op.max();
+    if (max_type != operand_type) {
+      assert(max_type.getRank() == 0);
+      max_value = rewriter.createOrFold<BroadcastOp>(
+          op.getLoc(), operand_type, max_value,
+          rewriter.getI64TensorAttr(operand_shape));
+    }
+
+    Value min_value = op.min();
+    if (min_type != operand_type) {
+      assert(min_type.getRank() == 0);
+      min_value = rewriter.createOrFold<BroadcastOp>(
+          op.getLoc(), operand_type, min_value,
+          rewriter.getI64TensorAttr(operand_shape));
+    }
+
+    rewriter.replaceOpWithNewOp<ClampOp>(op, op.getType(), min_value,
+                                         op.operand(), max_value);
+    return success();
+  }
+};
+
+}  // namespace
+
+void SetupMaterializeBroadcastsLegality(MLIRContext *context,
+                                        ConversionTarget *conversionTarget) {
+  conversionTarget->addDynamicallyLegalOp<ClampOp>([](ClampOp op) {
+    return op.max().getType() == op.operand().getType() &&
+           op.min().getType() == op.operand().getType();
+  });
+}
+
+void PopulateMaterializeBroadcastsPatterns(MLIRContext *context,
+                                           OwningRewritePatternList *patterns) {
+  // ClampOp. This op has a special case where it accepts either same-shaped
+  // inputs or scalars (a restricted form of broadcasting). This makes the
+  // broadcast explicit.
+  patterns->insert<ClampWithBroadcastConvert>(context);
+}
+
+}  // namespace xla_hlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/materialize_broadcasts_pass.cc

@@ -0,0 +1,58 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla_hlo {
+
+namespace {
+
+struct TestMaterializeBroadcastsPass
+    : public PassWrapper<TestMaterializeBroadcastsPass, FunctionPass> {
+  void runOnFunction() override {
+    ConversionTarget conversionTarget(getContext());
+    OwningRewritePatternList conversionPatterns;
+
+    // Consider the xla_hlo dialect legal for tests.
+    conversionTarget.addLegalDialect<XlaHloDialect>();
+    // The conversion uses helpers from the Standard dialect.
+    conversionTarget.addLegalDialect<mlir::StandardOpsDialect>();
+
+    SetupMaterializeBroadcastsLegality(&getContext(), &conversionTarget);
+    PopulateMaterializeBroadcastsPatterns(&getContext(), &conversionPatterns);
+
+    if (failed(applyPartialConversion(getFunction(), conversionTarget,
+                                      conversionPatterns))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+}  // namespace xla_hlo
+}  // namespace mlir
+
+static mlir::PassRegistration<mlir::xla_hlo::TestMaterializeBroadcastsPass>
+    pass("test-xla-materialize-broadcasts",
+         "Test pass for materializing 'broadcast_dimensions' attributes");

lib/Dialect/mhlo/transforms/sink_constants_to_control_flow.cc

@@ -0,0 +1,85 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/DenseMap.h"
+#include "third_party/llvm/llvm-project/llvm/include/llvm/Support/Casting.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/PassManager.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Support/LLVM.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/RegionUtils.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace xla_hlo {
+
+namespace {
+
+// A pass that sinks constants implicitly captured in control flow regions. This
+// is necessary to export to XLA.
+class SinkConstantsToControlFlow
+    : public mlir::PassWrapper<SinkConstantsToControlFlow, FunctionPass> {
+  void runOnFunction() override {
+    getFunction().walk([](Operation* op) {
+      if (auto while_op = llvm::dyn_cast<WhileOp>(op)) {
+        SinkToRegion(&while_op.body());
+        SinkToRegion(&while_op.cond());
+      } else if (auto if_op = llvm::dyn_cast<IfOp>(op)) {
+        SinkToRegion(&if_op.true_branch());
+        SinkToRegion(&if_op.false_branch());
+      }
+    });
+  }
+
+ private:
+  // Performs constant sinking into a region.
+  static void SinkToRegion(Region* region) {
+    llvm::DenseMap<Value, ConstOp> sunk_constant;
+    visitUsedValuesDefinedAbove({*region}, [&](OpOperand* use) {
+      Value constant = use->get();
+      auto const_op = dyn_cast_or_null<ConstOp>(constant.getDefiningOp());
+      if (!const_op) return;
+      auto map_entry = sunk_constant.try_emplace(constant, nullptr);
+      if (!map_entry.second) {
+        // This constant has already been cloned into the region, reuse it.
+        use->set(map_entry.first->getSecond().getResult());
+        if (constant.use_empty()) const_op.erase();
+        return;
+      }
+      if (constant.hasOneUse()) {
+        const_op.getOperation()->moveBefore(&region->front().front());
+        return;
+      }
+      map_entry.first->getSecond() = const_op.clone();
+      region->front().getOperations().insert(region->front().begin(),
+                                             map_entry.first->getSecond());
+      use->set(map_entry.first->getSecond().getResult());
+    });
+  }
+};
+
+static mlir::PassRegistration<SinkConstantsToControlFlow> pass(
+    "xla-hlo-sink-constants-to-control-flow",
+    "Sink constants implicitly captured in control flow regions. This is "
+    "necessary to export to XLA.");
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createSinkConstantsToControlFlowPass() {
+  return std::make_unique<SinkConstantsToControlFlow>();
+}
+
+}  // namespace xla_hlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc

@@ -0,0 +1,100 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Identifier.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/OperationSupport.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Interfaces/InferTypeOpInterface.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace xla {
+namespace {
+
+struct InferReturnTypeComponentsPattern : public RewritePattern {
+  InferReturnTypeComponentsPattern(MLIRContext *context)
+      : RewritePattern("xla_test.get_return_type_components", 1, context) {}
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (op->getNumOperands() != 1) return failure();
+    auto defining_op = op->getOperand(0).getDefiningOp();
+    auto defining_op_int =
+        llvm::dyn_cast_or_null<InferShapedTypeOpInterface>(defining_op);
+    if (!defining_op_int) return failure();
+    SmallVector<ShapedTypeComponents, 4> components;
+    if (failed(defining_op_int.inferReturnTypeComponents(
+            op->getContext(), op->getLoc(), defining_op->getOperands(),
+            defining_op->getAttrDictionary(), defining_op->getRegions(),
+            components))) {
+      return failure();
+    }
+
+    // Replace the op with another pass-through op with attributes added.
+    OperationState state(op->getLoc(), "xla_test.return_type_components",
+                         op->getOperands(), op->getResultTypes(),
+                         op->getAttrs());
+    auto new_op = rewriter.createOperation(state);
+    for (auto it : llvm::enumerate(components)) {
+      if (it.value().hasRank()) {
+        new_op->setAttr((StringRef("dims") + Twine(it.index())).str(),
+                        rewriter.getI64ArrayAttr(it.value().getDims()));
+      }
+      if (it.value().getElementType()) {
+        new_op->setAttr((Twine("element_type") + Twine(it.index())).str(),
+                        TypeAttr::get(it.value().getElementType()));
+      }
+    }
+    rewriter.replaceOp(op, {new_op->getResults()});
+    return success();
+  }
+};
+
+struct ReifyReturnTypeShapesPattern : public RewritePattern {
+  ReifyReturnTypeShapesPattern(MLIRContext *context)
+      : RewritePattern("xla_test.reify_return_type_shapes", 1, context) {}
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (op->getNumOperands() != 1) return failure();
+    auto defining_op = llvm::dyn_cast_or_null<InferShapedTypeOpInterface>(
+        op->getOperand(0).getDefiningOp());
+    if (!defining_op) return failure();
+    SmallVector<Value, 4> return_shapes;
+    if (failed(defining_op.reifyReturnTypeShapes(rewriter, return_shapes))) {
+      return failure();
+    }
+    rewriter.replaceOp(op, return_shapes);
+    return success();
+  }
+};
+
+struct TestInferShapedTypeMethodsPass
+    : public PassWrapper<TestInferShapedTypeMethodsPass, FunctionPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    patterns.insert<ReifyReturnTypeShapesPattern>(&getContext());
+    patterns.insert<InferReturnTypeComponentsPattern>(&getContext());
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+  }
+};
+
+}  // namespace
+}  // namespace xla
+}  // namespace mlir
+
+static mlir::PassRegistration<mlir::xla::TestInferShapedTypeMethodsPass> pass(
+    "test-xla-infer-shaped-type-methods",
+    "Uses test ops to invoke InferShapedTypeOpInterface methods");

lib/Dialect/mhlo/transforms/unfuse_batch_norm.cc

@@ -0,0 +1,184 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/SmallVector.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Builders.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Types.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace xla_hlo {
+
+namespace {
+
+// Broadcasts the 1D value tensor 'value_1d' to the shape of 'result_type'. If
+// 'shape_value' is initialized, creates a dynamic broadcast, otherwise creates
+// a static broadcast.
+Value BroadcastToFeatureDim(Location loc, RankedTensorType result_type,
+                            Value value_1d, Value shape_value,
+                            int64_t feature_dim,
+                            PatternRewriter& rewriter) {  // NOLINT
+  Builder b(rewriter.getContext());
+  auto dims_type = RankedTensorType::get({1}, b.getIntegerType(64));
+  auto dims = DenseIntElementsAttr::get(dims_type, {feature_dim});
+  if (shape_value) {
+    return rewriter.createOrFold<xla_hlo::DynamicBroadcastInDimOp>(
+        loc, result_type, value_1d, shape_value, dims);
+  }
+  assert(result_type.hasStaticShape());
+  return rewriter.create<xla_hlo::BroadcastInDimOp>(loc, result_type, value_1d,
+                                                    dims);
+}
+
+// Calculate the shape value of operand, assuming it is a dynamic shape with
+// static rank.
+Value CalculateShapeValue(Location loc, Value operand,
+                          PatternRewriter& rewriter) {  // NOLINT
+  RankedTensorType result_type = operand.getType().dyn_cast<RankedTensorType>();
+  llvm::SmallVector<Value, 4> shape_values;
+  int64_t rank = result_type.getRank();
+  shape_values.reserve(rank);
+  for (int64_t i = 0; i < rank; ++i) {
+    shape_values.push_back(rewriter.create<mlir::DimOp>(loc, operand, i));
+  }
+  return rewriter.create<TensorFromElementsOp>(loc, shape_values);
+}
+
+Value MaterializeEpsilon(Operation* op, FloatAttr epsilon_attr,
+                         FloatType fp_type, Value variance,
+                         RankedTensorType broadcast_to_type,
+                         PatternRewriter& rewriter) {  // NOLINT
+  Builder b(rewriter.getContext());
+  if (epsilon_attr.getType() != fp_type) {
+    // Need to convert.
+    bool loses_info;
+    APFloat epsilon_float = epsilon_attr.getValue();
+    auto status = epsilon_float.convert(
+        fp_type.getFloatSemantics(), APFloat::rmNearestTiesToEven, &loses_info);
+    if ((status & (~APFloat::opInexact)) != APFloat::opOK) {
+      op->emitWarning() << "Could not convert batch_norm epsilon to target fp "
+                           "type: opStatus = "
+                        << static_cast<int>(status);
+      return nullptr;
+    }
+    if (loses_info) {
+      op->emitWarning("Conversion of epsilon loses precision");
+    }
+    epsilon_attr = b.getFloatAttr(fp_type, epsilon_float);
+  }
+
+  auto scalar_type = RankedTensorType::get({}, fp_type);
+  auto epsilon_tensor_attr =
+      DenseElementsAttr::get(scalar_type, {epsilon_attr.cast<Attribute>()});
+  Value epsilon =
+      rewriter.create<xla_hlo::ConstOp>(op->getLoc(), epsilon_tensor_attr);
+  auto dims_type = RankedTensorType::get({0}, b.getIntegerType(64));
+  auto dims = DenseIntElementsAttr::get(dims_type, SmallVector<int64_t, 1>{});
+  if (broadcast_to_type.hasStaticShape()) {
+    return rewriter.create<xla_hlo::BroadcastInDimOp>(
+        op->getLoc(), broadcast_to_type, epsilon, /*broadcast_dims=*/dims);
+  }
+  Value shape_value = CalculateShapeValue(op->getLoc(), variance, rewriter);
+  return rewriter.createOrFold<xla_hlo::DynamicBroadcastInDimOp>(
+      op->getLoc(), broadcast_to_type, epsilon, shape_value,
+      /*broadcast_dims=*/dims);
+}
+
+class UnfuseBatchNormInferencePattern
+    : public OpRewritePattern<xla_hlo::BatchNormInferenceOp> {
+ public:
+  using OpRewritePattern<xla_hlo::BatchNormInferenceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(xla_hlo::BatchNormInferenceOp bn_op,
+                                PatternRewriter& rewriter) const override {
+    // Enforce type invariants.
+    // Note that we deduce the actual element type from the variance,
+    // which should not be subject to quantization at a higher level.
+    auto input_type = bn_op.operand().getType().dyn_cast<RankedTensorType>();
+    auto variance_type =
+        bn_op.variance().getType().dyn_cast<RankedTensorType>();
+    if (!input_type || !variance_type) {
+      return failure();
+    }
+    auto fp_type = variance_type.getElementType().dyn_cast<FloatType>();
+    if (!fp_type) {
+      return failure();
+    }
+    int64_t feature_dim = bn_op.feature_index().getSExtValue();
+
+    // Add epsilon to the variance and sqrt to get stddev:
+    // stddev = sqrt(variance + epsilon)
+    auto epsilon =
+        MaterializeEpsilon(bn_op.getOperation(), bn_op.epsilonAttr(), fp_type,
+                           bn_op.variance(), variance_type, rewriter);
+    if (!epsilon) {
+      return failure();
+    }
+    Value stddev = rewriter.create<xla_hlo::AddOp>(bn_op.getLoc(),
+                                                   bn_op.variance(), epsilon);
+    stddev = rewriter.create<xla_hlo::SqrtOp>(bn_op.getLoc(), stddev);
+
+    // Broadcast all terms.
+    Value shape_value;
+    if (!input_type.hasStaticShape()) {
+      shape_value =
+          CalculateShapeValue(bn_op.getLoc(), bn_op.operand(), rewriter);
+    }
+    auto broadcast_scale =
+        BroadcastToFeatureDim(bn_op.getLoc(), input_type, bn_op.scale(),
+                              shape_value, feature_dim, rewriter);
+    auto broadcast_offset =
+        BroadcastToFeatureDim(bn_op.getLoc(), input_type, bn_op.offset(),
+                              shape_value, feature_dim, rewriter);
+    auto broadcast_mean =
+        BroadcastToFeatureDim(bn_op.getLoc(), input_type, bn_op.mean(),
+                              shape_value, feature_dim, rewriter);
+    auto broadcast_stddev = BroadcastToFeatureDim(
+        bn_op.getLoc(), input_type, stddev, shape_value, feature_dim, rewriter);
+
+    // Compute:
+    // scale * (input - mean) / stddev + offset
+    Value result = rewriter.create<xla_hlo::SubOp>(
+        bn_op.getLoc(), bn_op.operand(), broadcast_mean);
+    result = rewriter.create<xla_hlo::MulOp>(bn_op.getLoc(), result,
+                                             broadcast_scale);
+    result = rewriter.create<xla_hlo::DivOp>(bn_op.getLoc(), result,
+                                             broadcast_stddev);
+    rewriter.replaceOpWithNewOp<xla_hlo::AddOp>(bn_op, result,
+                                                broadcast_offset);
+
+    return success();
+  }
+};
+
+}  // namespace
+
+// Populates conversion patterns to unfuse batch normalization operations.
+// In combination with marking such ops as illegal, this allows backends that
+// do not have special support for fused batchnorm to use simpler arithmetic
+// primitives.
+void PopulateUnfuseBatchNormPatterns(MLIRContext* context,
+                                     OwningRewritePatternList* patterns) {
+  patterns->insert<UnfuseBatchNormInferencePattern>(context);
+}
+
+}  // namespace xla_hlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/unfuse_batch_norm_pass.cc

@@ -0,0 +1,46 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla_hlo {
+
+namespace {
+
+struct TestUnfuseBatchNormPass
+    : public PassWrapper<TestUnfuseBatchNormPass, OperationPass<>> {
+  void runOnOperation() override {
+    OwningRewritePatternList patterns;
+    PopulateUnfuseBatchNormPatterns(&getContext(), &patterns);
+    applyPatternsAndFoldGreedily(getOperation(), patterns);
+  }
+};
+
+}  // namespace
+
+}  // namespace xla_hlo
+}  // namespace mlir
+
+static mlir::PassRegistration<mlir::xla_hlo::TestUnfuseBatchNormPass> pass(
+    "test-xla-unfuse-batch-norm",
+    "Test pass for materializing 'broadcast_dimensions' attributes");

lib/Dialect/mhlo/transforms/xla_hlo_fusion.cc

@@ -0,0 +1,579 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "mlir/Dialect/StandardOps/IR/Ops.h"  // TF:llvm-project
+#include "mlir/IR/MLIRContext.h"              // TF:llvm-project
+#include "mlir/IR/Matchers.h"
+#include "mlir/Pass/Pass.h"               // TF:local_config_mlir
+#include "mlir/Transforms/RegionUtils.h"  // TF:llvm-project
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/EquivalenceClasses.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/utils/cycle_detector.h"
+
+// This pass has similar functionality of the fusion pass in XLA stack.
+// However, unlike XLA, it targets the fully dynamic shape scenario.
+// Currently, it implements the kLoop and kInput fusion templates.
+// During conversion, it tries to greedily find kLoop/kInput fusion
+// patterns.
+//
+// Similar to XLA, this pass supports fusion pattern having multiple outputs
+// if all the shape of outputs are consistent. Following are some examples.
+//
+//        kLoop                          kInput
+// +----+  +----+  +----+    +----+    +----+    +----+
+// |elem|  |elem|  |elem|    |elem<----+elem+---->elem+----+
+// +-+--+  +-+--+  +-+--+    +-+--+    +----+    +-+--+    |
+//   |       |       |         |                   |       |
+//   |               |         |                   |       |
+// +-v--+    |     +-v--+   +--v---+            +--v---+   |
+// |elem+<---+----<+elem|   |reduce|            |reduce|   |
+// +-+--+          +-+--+   +--+---+            +--+---+   |
+//   |               |         |                   |       |
+//   |               |         |                   |       |
+//   v               v         v                   v       v
+//
+// To this end, we also add an simple shape constraint analysis phase.
+// For kLoop fusion template, it requires all the outputs of the fused
+// pattern have the same shape. However, we don't know the actual value
+// of the shape at the compile time in the dynamic shape world.
+// Fortunately, we could still infer the relationship among different ops
+// according to their shape constrain traits. Currently, We only consider
+// shape equality propagation for elementwise ops (assuming that implicit
+// shape broadcast is forbidden). The above process could be built on the
+// shape dialect once it is ready.
+
+namespace mlir {
+namespace xla_hlo {
+namespace {
+
+using llvm::EquivalenceClasses;
+using FusionPattern = std::vector<Operation*>;
+using FusionPlan = std::vector<FusionPattern>;
+
+// To support using EquivalenceClasses for Value
+class ValueWrapper {
+ public:
+  explicit ValueWrapper(Value value) : value_(std::move(value)) {}
+
+  Value getValue() const { return value_; }
+
+  bool operator==(const ValueWrapper& rhs) const {
+    return getValue() == rhs.getValue();
+  }
+
+ private:
+  Value value_;
+};
+
+bool operator<(const ValueWrapper& lhs, const ValueWrapper& rhs) {
+  auto lhs_value = lhs.getValue().getAsOpaquePointer();
+  auto rhs_value = rhs.getValue().getAsOpaquePointer();
+  return lhs_value < rhs_value;
+}
+
+bool IsFusible(Operation* op) {
+  if (matchPattern(op, m_Constant())) {
+    return true;
+  }
+  auto op_fusibility = dyn_cast<InferFusibilityOpInterface>(op);
+  return op_fusibility && (op_fusibility.isFusibleWithOperand() ||
+                           op_fusibility.isFusibleWithConsumer());
+}
+
+SmallVector<Value, 4> GetInputsOfFusionPattern(const FusionPattern& pattern) {
+  SmallVector<Value, 4> inputs;
+  DenseSet<Value> input_set;
+  DenseSet<Operation*> op_set;
+  for (Operation* op : pattern) {
+    bool inserted = op_set.insert(op).second;
+    (void)inserted;
+    assert(inserted && "FusionPattern contains duplicate operations");
+  }
+
+  for (Operation* op : pattern) {
+    for (Value operand : op->getOperands()) {
+      Operation* operand_op = operand.getDefiningOp();
+      if (op_set.find(operand_op) != op_set.end()) {
+        // skip if defining op is in the pattern
+        continue;
+      }
+      if (input_set.insert(operand).second) {
+        inputs.push_back(operand);
+      }
+    }
+  }
+  return inputs;
+}
+
+SmallVector<Value, 4> GetOutputsOfFusionPattern(const FusionPattern& pattern) {
+  SmallVector<Value, 4> outputs;
+  DenseSet<Operation*> op_set;
+  for (Operation* op : pattern) {
+    bool inserted = op_set.insert(op).second;
+    (void)inserted;
+    assert(inserted && "FusionPattern contains duplicate operations");
+  }
+
+  for (Operation* op : pattern) {
+    for (Value result : op->getResults()) {
+      bool has_external_user = llvm::any_of(
+          result.getUses(),
+          [&](OpOperand& use) { return !op_set.count(use.getOwner()); });
+      if (has_external_user) {
+        outputs.push_back(result);
+      }
+    }
+  }
+  return outputs;
+}
+
+FusionPattern MergeFusionPattern(const FusionPattern& lhs,
+                                 const FusionPattern& rhs) {
+  FusionPattern pattern(lhs);
+  pattern.insert(pattern.end(), rhs.begin(), rhs.end());
+  return pattern;
+}
+
+inline int EffectiveSize(const FusionPattern& pattern) {
+  return llvm::count_if(
+      pattern, [](Operation* op) { return !matchPattern(op, m_Constant()); });
+}
+
+// This is an simple shape constraint analysis, which is used to
+// guide fusion decision (e.g. we only fuse shape-compatible ops).
+//
+// Currently, We only consider shape equality propagation based
+// on the shape constrain traits of elementwise ops (assuming that
+// implicit shape broadcast is forbidden).
+class ShapeConstraintAnalysis {
+ public:
+  explicit ShapeConstraintAnalysis(const SmallVectorImpl<Operation*>& op_list) {
+    PropagateEquality(op_list);
+  }
+
+  // Returns true is `lhs` and `rhs` are supposed to have same shape.
+  bool HasSameShape(Value lhs, Value rhs) {
+    return impl_.isEquivalent(ValueWrapper(lhs), ValueWrapper(rhs));
+  }
+
+ private:
+  // shape equality propagation based on the shape constrains of
+  // elementwise ops.
+  void PropagateEquality(const SmallVectorImpl<Operation*>& op_list) {
+    bool converged = true;
+    do {
+      converged = true;
+      auto update = [&](Value lhs, Value rhs) {
+        if (!impl_.isEquivalent(ValueWrapper(lhs), ValueWrapper(rhs))) {
+          converged = false;
+          impl_.unionSets(ValueWrapper(lhs), ValueWrapper(rhs));
+        }
+      };
+      for (Operation* op : op_list) {
+        auto op_fusibility = dyn_cast<InferFusibilityOpInterface>(op);
+        if (!op_fusibility) continue;
+        int numInput = op->getNumOperands();
+        int numOutput = op->getNumResults();
+        // shape equality propagation between inputs.
+        for (int input1 = 0; input1 < numInput; ++input1)
+          for (int input2 = input1 + 1; input2 < numInput; ++input2)
+            if (op_fusibility.inferInputsShapeEquality(input1, input2))
+              update(op->getOperand(input1), op->getOperand(input2));
+
+        // shape equality propagation between outputs.
+        for (int output1 = 0; output1 < numOutput; ++output1)
+          for (int output2 = output1 + 1; output2 < numOutput; ++output2)
+            if (op_fusibility.inferOutputsShapeEquality(output1, output2))
+              update(op->getResult(output1), op->getResult(output2));
+
+        // shape equality propagation between input and output.
+        for (int input = 0; input < numInput; ++input)
+          for (int output = 0; output < numOutput; ++output)
+            if (op_fusibility.inferInputOutputShapeEquality(input, output))
+              update(op->getOperand(input), op->getResult(output));
+      }
+    } while (!converged);
+  }
+
+  // a UnionFind set
+  EquivalenceClasses<ValueWrapper> impl_;
+};
+
+// A fusion planner that can propose a fusion plan for a block of ops.
+// The fusion plan is consisted of a group of fusion patterns.
+//
+// Currently all proposed patterns followed xla kLoop/kInput like fusion
+// templates while are adapted to the fully dynamic shape world.
+//
+// kLoop fusion template satifies:
+//   - all ops in the fusion pattern are element-wise.
+//   - all the shapes of outputs of fusion pattern are same, and thus can
+//     fit into a same parallel loop.
+//
+// kInput fusion template satifies:
+//   - any op in the fusion pattern is either element-wise or a reduction.
+//   - if a op is a reduction, its output cannot be consumered by other
+//     ops in the same fusion pattern.
+//   - all the effective shapes of outputs of fusion pattern are same.
+//     - For element-wise op, its effective shape is its output shape.
+//     - For reduction op, its effective shape is its operand shape.
+class FusionPlanner {
+ public:
+  explicit FusionPlanner(const SmallVectorImpl<Operation*>& op_list)
+      : op_list_(op_list),
+        shape_analysis_(op_list),
+        cycle_detector_(op_list.size()) {
+    BuildNodeMap();
+  }
+
+  // Returns a fusion plan if success, otherwise none.
+  llvm::Optional<FusionPlan> Run() {
+    // Greedily search connected fusible pattern, and ops belonging to
+    // a same fusion pattern are grouped into a cluster.
+    RunEdgeContractionLoop();
+
+    // After doing edge contraction, each unique cluster having size
+    // more than one represents a potential fusion pattern.
+    // We collect all these clusters and construct a fusion plan.
+    //
+    // Note that the ops in a fusion pattern are in topological ordering.
+    FusionPlan plan;
+    DenseMap<int, int> pattern_ids;
+    for (Operation* op : op_list_) {
+      Cluster* cluster = GetClusterForNode(op);
+      int node_id = cluster->cycles_graph_node_id();
+      if (!IsFusible(op_list_[node_id]) ||
+          EffectiveSize(GetClusterForNode(op)->fused_pattern()) <= 1) {
+        continue;
+      }
+      if (!pattern_ids.count(node_id)) {
+        int pattern_id = pattern_ids.size();
+        pattern_ids[node_id] = pattern_id;
+        plan.emplace_back();
+      }
+      plan[pattern_ids[node_id]].push_back(op);
+    }
+    return plan;
+  }
+
+  // Returns the op_list this planner operates on.
+  const SmallVectorImpl<Operation*>& op_list() const { return op_list_; }
+
+ private:
+  // Represent a (partial) fused pattern
+  class Cluster {
+   public:
+    Cluster(int node_id, FusionPlanner* planner) : node_id_(node_id) {
+      const SmallVectorImpl<Operation*>& op_list = planner->op_list();
+      pattern_.push_back(op_list[node_id]);
+    }
+
+    // Merges `other` into this cluster, and clears `other`.
+    void Merge(Cluster* other) {
+      pattern_.insert(pattern_.end(), other->pattern_.begin(),
+                      other->pattern_.end());
+      other->pattern_.clear();
+    }
+
+    // The number of nodes in this cluster.
+    int cluster_size() const { return pattern_.size(); }
+
+    // The ID of the cluster as represented in `cycle_detector_`.
+    int cycles_graph_node_id() const { return node_id_; }
+
+    // Sets the ID of the cluster as represented in `cycle_detector_`.
+    void set_cycles_graph_node_id(int cycles_graph_node_id) {
+      node_id_ = cycles_graph_node_id;
+    }
+
+    // Currently the fused pattern this cluster holds.
+    const FusionPattern& fused_pattern() { return pattern_; }
+
+   private:
+    // ID of the representative node of this cluster.
+    int node_id_;
+
+    // the fused pattern this cluster holds.
+    FusionPattern pattern_;
+  };
+
+ private:
+  Cluster* MakeCluster(int cycles_graph_node_id) {
+    cluster_storage_.emplace_back(new Cluster(cycles_graph_node_id, this));
+    return cluster_storage_.back().get();
+  }
+
+  void BuildNodeMap() {
+    int num_nodes = op_list_.size();
+    for (int node_id = 0; node_id < num_nodes; ++node_id) {
+      Operation* op = op_list_[node_id];
+      MakeCluster(node_id);
+      op_to_node_id_[op] = node_id;
+      leader_for_node_.insert(node_id);
+      for (Value operand : op->getOperands()) {
+        Operation* operand_op = operand.getDefiningOp();
+        if (operand_op == nullptr) {
+          // skip block argument
+          continue;
+        }
+        auto iter = op_to_node_id_.find(operand_op);
+        assert(iter != op_to_node_id_.end());
+        cycle_detector_.InsertEdge(iter->second, node_id);
+      }
+    }
+  }
+
+  // Returns the cluster contains this op.
+  Cluster* GetClusterForNode(Operation* n) {
+    int id = op_to_node_id_.at(n);
+    id = leader_for_node_.getLeaderValue(id);
+    return cluster_storage_[id].get();
+  }
+
+  // Returns the cluster contains the op having `node_id`.
+  Cluster* GetClusterForCyclesGraphNode(int node_id) {
+    return cluster_storage_[leader_for_node_.getLeaderValue(node_id)].get();
+  }
+
+  // Merges the clusters `cluster_from` and `cluster_to`.
+  bool MergeClusters(Cluster* cluster_from, Cluster* cluster_to) {
+    int from = cluster_from->cycles_graph_node_id();
+    int to = cluster_to->cycles_graph_node_id();
+
+    auto optional_merged_node = cycle_detector_.ContractEdge(from, to);
+    if (!optional_merged_node.hasValue()) {
+      llvm::dbgs() << "Could not contract " << from << " -> " << to
+                   << " because contracting the edge would create a cycle.";
+      return false;
+    }
+
+    // Merge the clusters.
+    cluster_from->Merge(cluster_to);
+    cluster_from->set_cycles_graph_node_id(*optional_merged_node);
+
+    // Merge the UnionFind Set.
+    leader_for_node_.unionSets(from, to);
+    return true;
+  }
+
+  template <typename FnTy>
+  bool ForEachEdgeInPostOrder(FnTy fn) {
+    bool changed = false;
+    for (int32_t node : cycle_detector_.AllNodesInPostOrder()) {
+      Cluster* cluster_from = GetClusterForCyclesGraphNode(node);
+      // Make a copy of the set of successors because we may modify the graph in
+      // TryToContractEdge.
+      std::vector<int32_t> successors_copy =
+          cycle_detector_.SuccessorsCopy(cluster_from->cycles_graph_node_id());
+
+      for (int to : successors_copy) {
+        Cluster* cluster_to = GetClusterForCyclesGraphNode(to);
+        bool contracted_edge = fn(cluster_from, cluster_to);
+        changed |= contracted_edge;
+      }
+    }
+
+    return changed;
+  }
+
+  // returns the outputs if two cluster were merged
+  SmallVector<Value, 4> GetResultsOfFusedPattern(Cluster* from, Cluster* to) {
+    FusionPattern fused_pattern =
+        MergeFusionPattern(from->fused_pattern(), to->fused_pattern());
+    return GetOutputsOfFusionPattern(fused_pattern);
+  }
+
+  // This function check if fusing `from` with `to` is valid and if so perform
+  // the merge. The validity is based on the operations in the clusters and
+  // the compatibility of the shapes of the outputs of the would-be fused
+  // clusters.
+  // Returns true is the merge was performed.
+  bool TryToContractEdge(Cluster* from, Cluster* to) {
+    int node_to = to->cycles_graph_node_id();
+    int node_from = from->cycles_graph_node_id();
+
+    // Both node_to and node_from should be fusible
+    if (!IsFusible(op_list_[node_to]) || !IsFusible(op_list_[node_from])) {
+      return false;
+    }
+
+    auto op_from_fusibility =
+        dyn_cast<InferFusibilityOpInterface>(op_list_[node_from]);
+    if (op_from_fusibility && !op_from_fusibility.isFusibleWithConsumer()) {
+      // This op cannot be fused with its consumers.
+      return false;
+    }
+
+    auto op_to_fusibility =
+        dyn_cast<InferFusibilityOpInterface>(op_list_[node_to]);
+    if (op_to_fusibility && !op_to_fusibility.isFusibleWithOperand()) {
+      // This op cannot be fused with its operands.
+      return false;
+    }
+
+    // Output shapes of a fusion pattern should be compatible as described in
+    // the document of this class.
+    SmallVector<Value, 4> results = GetResultsOfFusedPattern(from, to);
+    auto get_workload_shape = [](Value v) {
+      Operation* op = v.getDefiningOp();
+      // Block argument
+      if (!op) return v;
+      auto op_fusibility = dyn_cast<InferFusibilityOpInterface>(op);
+      // Const value
+      if (!op_fusibility) return v;
+      llvm::Optional<Value> workload =
+          op_fusibility.inferEffectiveWorkloadShape();
+      return workload.hasValue() ? *workload : v;
+    };
+
+    Value ref = get_workload_shape(results[0]);
+    if (!llvm::all_of(results, [&](Value result) {
+          Value val = get_workload_shape(result);
+          return shape_analysis_.HasSameShape(ref, val);
+        })) {
+      return false;
+    }
+
+    return MergeClusters(from, to);
+  }
+
+  // Greedily fuse connected node.
+  bool RunEdgeContractionLoop() {
+    using std::placeholders::_1;
+    using std::placeholders::_2;
+    return ForEachEdgeInPostOrder(
+        std::bind(&FusionPlanner::TryToContractEdge, this, _1, _2));
+  }
+
+  const SmallVectorImpl<Operation*>& op_list_;
+
+  // Shape equality checker
+  ShapeConstraintAnalysis shape_analysis_;
+
+  // op -> node_id
+  std::unordered_map<Operation*, int> op_to_node_id_;
+
+  // make sure not introduce cycle after fusion
+  GraphCycles cycle_detector_;
+  std::vector<std::unique_ptr<Cluster>> cluster_storage_;
+
+  // a UnionFind set. Each set represents a (partial) fused pattern
+  // and has a leader as representation.
+  EquivalenceClasses<int32_t> leader_for_node_;
+};
+
+struct XlaHloFusion : public mlir::PassWrapper<XlaHloFusion, FunctionPass> {
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+    if (!IsTargetFunc(func)) {
+      return;
+    }
+
+    // process each block and do fusion within a block.
+    for (Block& block : func) {
+      SmallVector<Operation*, 4> op_list;
+      for (Operation& op : block) {
+        op_list.push_back(&op);
+      }
+
+      FusionPlanner planner(op_list);
+      llvm::Optional<FusionPlan> plan = planner.Run();
+      if (!plan) {
+        emitError(func.getLoc(), "can't find a fusion plan");
+        signalPassFailure();
+        return;
+      }
+      if (!ApplyFusionPlan(*plan)) {
+        emitError(func.getLoc(), "apply fusion plan failed");
+        signalPassFailure();
+        return;
+      }
+    }
+  }
+
+  bool IsTargetFunc(FuncOp func) {
+    int num_fusible_ops = 0;
+    bool is_target_func = false;
+    // We only process the function having enough candidates
+    func.walk([&](Operation* op) {
+      num_fusible_ops +=
+          static_cast<int>(dyn_cast<InferFusibilityOpInterface>(op) != nullptr);
+      is_target_func = (num_fusible_ops > 1);
+      // early stop
+      if (is_target_func) return WalkResult::interrupt();
+      return WalkResult::advance();
+    });
+    return is_target_func;
+  }
+
+  bool ApplyFusionPlan(const FusionPlan& plan) {
+    for (const FusionPattern& pattern : plan) {
+      OpBuilder b(pattern.back());
+
+      SmallVector<Location, 4> locations;
+      locations.reserve(pattern.size());
+      for (Operation* op : pattern) {
+        locations.push_back(op->getLoc());
+      }
+      Location fused_loc =
+          FusedLoc::get(locations, pattern.back()->getContext());
+
+      SmallVector<Value, 4> inputs = GetInputsOfFusionPattern(pattern);
+      SmallVector<Value, 4> outputs = GetOutputsOfFusionPattern(pattern);
+      SmallVector<Type, 4> output_types;
+      output_types.reserve(outputs.size());
+      for (Value v : outputs) {
+        output_types.push_back(v.getType());
+      }
+
+      FusionOp fusion =
+          b.create<xla_hlo::FusionOp>(fused_loc, output_types, inputs);
+      Region& region = fusion.fused_computation();
+      region.push_back(new Block);
+      Block& block = region.front();
+      for (Operation* op : pattern) {
+        op->moveBefore(&block, block.end());
+      }
+      b.setInsertionPoint(&block, block.end());
+      b.create<xla_hlo::ReturnOp>(fused_loc, outputs);
+
+      for (auto output_and_result : llvm::zip(outputs, fusion.getResults())) {
+        Value output = std::get<0>(output_and_result);
+        Value fusion_result = std::get<1>(output_and_result);
+        for (OpOperand& use : llvm::make_early_inc_range(output.getUses())) {
+          if (use.getOwner()->getBlock() != &block) use.set(fusion_result);
+        }
+      }
+    }
+    return true;
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> createXlaHloFusion() {
+  return std::make_unique<XlaHloFusion>();
+}
+
+static PassRegistration<XlaHloFusion> xla_hlo_fusion_pass(
+    "xla-hlo-fusion", "fuse xla_hlo ops to kLoop/kInput fusion patterns.");
+
+}  // namespace xla_hlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/xla_legalize_to_linalg.cc

@@ -0,0 +1,909 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for lowering HLO/LHLO dialect to Linalg dialect.
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/AffineExpr.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Builders.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Location.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/map_xla_to_scalar_op.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace {
+
+SmallVector<StringRef, 3> GetNParallelLoopsAttrs(unsigned nParallelLoops) {
+  static constexpr StringRef kParallelIterType = "parallel";
+  return SmallVector<StringRef, 3>(nParallelLoops, kParallelIterType);
+}
+
+template <bool isLHLO = true>
+Value getResultValue(Operation* op) {
+  return isLHLO ? op->getOperand(op->getNumOperands() - 1) : op->getResult(0);
+}
+
+template <bool isLHLO = true>
+ShapedType getXLAOpResultType(Operation* op) {
+  return getResultValue<isLHLO>(op).getType().template cast<ShapedType>();
+}
+
+template <bool isLHLO = true>
+bool verifyXLAOpBufferOrTensorSemantics(Operation* op) {
+  auto verifyType = [&](Value val) -> bool {
+    return (isLHLO && val.getType().isa<MemRefType>()) ||
+           (!isLHLO && val.getType().isa<RankedTensorType>());
+  };
+  if (!llvm::all_of(op->getOperands(), verifyType)) return false;
+  return isLHLO ? op->getResults().empty()
+                : llvm::all_of(op->getResults(), verifyType);
+}
+
+template <typename OpTy, bool isLHLO = true>
+class PointwiseToLinalgConverter : public OpConversionPattern<OpTy> {
+ public:
+  using OpConversionPattern<OpTy>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      OpTy op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = op.getLoc();
+    auto argType =
+        op.getOperation()->getOperand(0).getType().template cast<ShapedType>();
+    if (!argType.hasRank()) {
+      emitError(loc, "lhlo to linalg conversion expects ranked args");
+      return failure();
+    }
+    auto elemTy = argType.getElementType();
+    if (!elemTy.isSignlessIntOrFloat() && !elemTy.template isa<ComplexType>()) {
+      return failure();
+    }
+
+    // Construct the indexing maps needed for linalg.generic ops.
+    SmallVector<AffineMap, 2> indexing_maps;
+    SmallVector<Type, 4> bodyArgTypes, bodyResultTypes, opResultTypes;
+
+    // This doesnt account for implicit broadcast, but the working assumption
+    // here is that are broadcasts have been made explicit.
+    unsigned nloops = argType.getRank();
+
+    if (isLHLO && !nloops) return failure();
+
+    int operandCount = (isLHLO ? args.size() - 1 : args.size());
+    auto verifyArgOrResultType = [&](Value val) -> ShapedType {
+      auto shapedType = val.getType().dyn_cast<ShapedType>();
+      if (!shapedType ||
+          (!shapedType.isa<MemRefType>() &&
+           !shapedType.isa<RankedTensorType>()) ||
+          shapedType.getRank() != nloops)
+        return nullptr;
+      indexing_maps.emplace_back(
+          nloops ? rewriter.getMultiDimIdentityMap(nloops)
+                 : AffineMap::get(nloops, 0, rewriter.getContext()));
+      return shapedType;
+    };
+    for (const auto& arg : llvm::enumerate(args)) {
+      auto shapedType = verifyArgOrResultType(arg.value());
+      if (!shapedType) return failure();
+      auto& result_or_body_arg =
+          arg.index() < operandCount ? bodyArgTypes : bodyResultTypes;
+      result_or_body_arg.emplace_back(shapedType.getElementType());
+    }
+    if (!isLHLO) {
+      // HLO operations have return as tensor types.
+      assert(bodyResultTypes.empty() &&
+             "When lowering HLO ops result can't be part of arguments");
+      Value result = op.getOperation()->getResult(0);
+      auto shapedType = verifyArgOrResultType(result);
+      if (!shapedType) return failure();
+      bodyResultTypes.push_back(shapedType.getElementType());
+      opResultTypes.push_back(shapedType);
+    }
+
+    int64_t args_count = bodyArgTypes.size();
+    int64_t results_count = bodyResultTypes.size();
+    auto linalgOp = rewriter.create<linalg::GenericOp>(
+        loc, opResultTypes, args, args_count, results_count, indexing_maps,
+        GetNParallelLoopsAttrs(nloops),
+        [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+          // TODO(ravishankarm) : For now use the method in xla_lhlo namespace.
+          // That method needs to be moved out of there.
+          Value opResult = xla_lhlo::XlaOpToStdScalarOp::map<OpTy>(
+              op, bodyResultTypes,
+              llvm::to_vector<2>(args.take_front(args_count)), &rewriter);
+          nestedBuilder.create<linalg::YieldOp>(loc, opResult);
+        });
+    rewriter.replaceOp(op, linalgOp.getOperation()->getResults());
+    return success();
+  }
+};
+
+template <typename LhloOp>
+class ScalarPointwiseToStandardConverter : public OpConversionPattern<LhloOp> {
+ public:
+  using OpConversionPattern<LhloOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      LhloOp lhlo_op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = lhlo_op.getLoc();
+    auto argType =
+        lhlo_op.getOperand(0).getType().template dyn_cast<ShapedType>();
+    if (!argType || !argType.getElementType().isSignlessIntOrFloat() ||
+        (argType.getRank() != 0)) {
+      return failure();
+    }
+
+    // Create two loads from the input.
+    auto lhs = rewriter.create<LoadOp>(loc, lhlo_op.lhs());
+    auto rhs = rewriter.create<LoadOp>(loc, lhlo_op.rhs());
+    // TODO(ravishankarm) : Move this method out of xla_lhlo namespace.
+    Value opResult = xla_lhlo::XlaOpToStdScalarOp::map<LhloOp>(
+        lhlo_op, argType.getElementType(), llvm::ArrayRef<Value>{lhs, rhs},
+        &rewriter);
+    rewriter.create<StoreOp>(loc, opResult, lhlo_op.out());
+    rewriter.eraseOp(lhlo_op);
+    return success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// xla_lhlo.convolution conversion pattern.
+//===----------------------------------------------------------------------===//
+
+/// Converts xla_lhlo.convolution operation to a linalg.conv op.
+struct ConvToLinalgConverter : public OpConversionPattern<xla_lhlo::ConvOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::ConvOp>::OpConversionPattern;
+
+  //  This code has been adapted from IREE's
+  //  (https://github.com/google/iree/) xla_hlo -> linalg conversion.
+  LogicalResult matchAndRewrite(
+      xla_lhlo::ConvOp op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    // Check validity of dimension information.
+    if (const xla_lhlo::ConvDimensionNumbers& dimensionNumbers =
+            op.dimension_numbers()) {
+      const int inputSpatialRank =
+          llvm::size(dimensionNumbers.input_spatial_dimensions());
+      // The dimensions for input should follow the order of
+      // batch_count, spatial_dims..., input_feature_count.
+      if (dimensionNumbers.input_batch_dimension().getInt() != 0 ||
+          dimensionNumbers.input_feature_dimension().getInt() !=
+              (inputSpatialRank + 1))
+        return failure();
+
+      const int kernelSpatialRank =
+          llvm::size(dimensionNumbers.kernel_spatial_dimensions());
+      // The dimensions for filter should follow the order of
+      // spatial_dims..., input_feature_count, num_output_feature_count.
+      if (dimensionNumbers.kernel_input_feature_dimension().getInt() !=
+              kernelSpatialRank ||
+          dimensionNumbers.kernel_output_feature_dimension().getInt() !=
+              (kernelSpatialRank + 1))
+        return failure();
+
+      const int outputSpatialRank =
+          llvm::size(dimensionNumbers.output_spatial_dimensions());
+      // The dimensions for output should follow the order of
+      // batch_count, spatial_dims.., output_feature_count.
+      if (dimensionNumbers.output_batch_dimension().getInt() != 0 ||
+          dimensionNumbers.output_feature_dimension().getInt() !=
+              (outputSpatialRank + 1))
+        return failure();
+
+      if (inputSpatialRank != outputSpatialRank ||
+          inputSpatialRank != kernelSpatialRank)
+        return failure();
+
+      auto inputSpatialDim =
+          dimensionNumbers.input_spatial_dimensions().begin();
+      auto kernelSpatialDim =
+          dimensionNumbers.kernel_spatial_dimensions().begin();
+      auto outputSpatialDim =
+          dimensionNumbers.output_spatial_dimensions().begin();
+      // Check if spatial dims are ordered correctly.
+      for (int i = 0; i < inputSpatialRank; ++i) {
+        const int dim = i + 1;
+        if ((*inputSpatialDim++).getZExtValue() != dim ||
+            (*outputSpatialDim++).getZExtValue() != dim ||
+            (*kernelSpatialDim++).getZExtValue() != i)
+          return failure();
+      }
+    }
+
+    // TODO: LHS dilation for deconvolution not supported yet.
+    if (op.lhs_dilation()) {
+      return failure();
+    }
+
+    llvm::SmallVector<Attribute, 4> strides;
+    if (auto windowStrides = op.window_strides()) {
+      auto range = windowStrides->getAttributeValues();
+      strides.assign(range.begin(), range.end());
+    }
+    auto stridesArg = ArrayAttr::get(strides, op.getContext());
+
+    llvm::SmallVector<Attribute, 2> dilation;
+    if (auto rhsDilation = op.rhs_dilation()) {
+      auto range = rhsDilation->getAttributeValues();
+      dilation.assign(range.begin(), range.end());
+    } else {
+      // Default dilation of 1.
+      dilation.resize(2, IntegerAttr::get(rewriter.getIntegerType(64), 1));
+    }
+    auto dilationArg = ArrayAttr::get(dilation, op.getContext());
+
+    // Set padding only if it is non-zero.
+    DenseIntElementsAttr padding = op.paddingAttr();
+    if (!padding || !llvm::any_of(padding.getValues<APInt>(), [](APInt intVal) {
+          return !intVal.isNullValue();
+        })) {
+      padding = nullptr;
+    }
+
+    // The order of input and filter are switched with linalg.conv.
+    rewriter.replaceOpWithNewOp<linalg::ConvOp>(
+        op, args[1], args[0], args[2], stridesArg, dilationArg, padding);
+    return success();
+  }
+};
+
+/// Base class for lowering xla operations that have one operand and one result,
+/// and are semantically equivalent to a copy of the input to the output (like
+/// transpose, some reshape, etc.). The derived classes need to provide a method
+/// `getIndexingMaps` that returns AffineMaps for the index maps of the input
+/// and the output.
+template <typename Derived, typename OpTy, bool isLHLO = true>
+class DataMovementOpConverter : public OpConversionPattern<OpTy> {
+ public:
+  using OpConversionPattern<OpTy>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      OpTy op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    if (!verifyXLAOpBufferOrTensorSemantics<isLHLO>(op)) return failure();
+    auto resultType = getXLAOpResultType<isLHLO>(op);
+
+    SmallVector<AffineMap, 2> indexing_maps =
+        Derived::getIndexingMaps(op, &rewriter);
+    if (indexing_maps.empty()) return failure();
+
+    auto nloops = resultType.getRank();
+    auto loc = op.getLoc();
+    auto linalgOp = rewriter.create<linalg::GenericOp>(
+        loc, isLHLO ? ArrayRef<Type>{} : resultType, args, /*inputCount=*/1,
+        /*outputCount=*/1, indexing_maps, GetNParallelLoopsAttrs(nloops),
+        [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+          nestedBuilder.create<linalg::YieldOp>(loc, *args.begin());
+        });
+
+    rewriter.replaceOp(op, linalgOp.getOperation()->getResults());
+    return success();
+  }
+};
+
+/// Pattern to convert BroadcastOp to Linalg ops.
+template <typename OpTy, bool isLHLO = true>
+class BroadcastConverter
+    : public DataMovementOpConverter<BroadcastConverter<OpTy, isLHLO>, OpTy,
+                                     isLHLO> {
+ public:
+  using DataMovementOpConverter<BroadcastConverter, OpTy,
+                                isLHLO>::DataMovementOpConverter;
+
+  static SmallVector<AffineMap, 2> getIndexingMaps(OpTy broadcastOp,
+                                                   Builder* b) {
+    ShapedType inputType =
+        broadcastOp.operand().getType().template cast<ShapedType>();
+    unsigned inputRank = inputType.getRank();
+    unsigned nloops = getXLAOpResultType<isLHLO>(broadcastOp).getRank();
+
+    // BroadcastOp prepends the dimensions in the `broadcast_sizes` attribute to
+    // the input's dimensions.
+    unsigned numPrependedDims = llvm::size(broadcastOp.broadcast_sizes());
+    SmallVector<AffineExpr, 4> inputDimExprs;
+    inputDimExprs.reserve(inputRank);
+    for (int i = 0; i < inputRank; ++i) {
+      inputDimExprs.push_back(b->getAffineDimExpr(numPrependedDims + i));
+    }
+
+    AffineMap inputMap;
+    MLIRContext* context = b->getContext();
+    if (inputDimExprs.empty()) {
+      // The input is a scalar, i.e. this is a scalar broadcast op.
+      inputMap = AffineMap::get(nloops, /*symbolCount=*/0, context);
+    } else {
+      inputMap =
+          AffineMap::get(nloops, /*symbolCount=*/0, inputDimExprs, context);
+    }
+    return {inputMap, b->getMultiDimIdentityMap(nloops)};
+  }
+};
+
+class HloBroadcastInDimConverter
+    : public DataMovementOpConverter<HloBroadcastInDimConverter,
+                                     xla_hlo::BroadcastInDimOp, false> {
+ public:
+  using DataMovementOpConverter<HloBroadcastInDimConverter,
+                                xla_hlo::BroadcastInDimOp,
+                                false>::DataMovementOpConverter;
+
+  static SmallVector<AffineMap, 2> getIndexingMaps(
+      xla_hlo::BroadcastInDimOp broadcastOp, Builder* b) {
+    auto resultType = getXLAOpResultType<false>(broadcastOp);
+    auto operandType =
+        broadcastOp.operand().getType().template cast<ShapedType>();
+    unsigned nloops = resultType.getRank();
+
+    // The input is a scalar, i.e. this is a scalar broadcast op.
+    if (operandType.getRank() == 0) {
+      return {AffineMap::get(nloops, /*symbolCount=*/0, b->getContext()),
+              b->getMultiDimIdentityMap(nloops)};
+    }
+
+    auto operandShape = operandType.getShape();
+    SmallVector<AffineExpr, 4> dimExprs;
+    dimExprs.reserve(nloops);
+
+    if (broadcastOp.broadcast_dimensions()) {
+      for (const auto& broadcastDim :
+           enumerate(broadcastOp.broadcast_dimensions().getIntValues())) {
+        int size = broadcastDim.value().getSExtValue();
+        bool expansion_needed = operandShape[broadcastDim.index()] == 1 &&
+                                resultType.getShape()[size] != 1;
+        dimExprs.push_back(expansion_needed ? b->getAffineConstantExpr(0)
+                                            : b->getAffineDimExpr(size));
+      }
+    }
+    return {
+        AffineMap::get(nloops, /*symbolCount=*/0, dimExprs, b->getContext()),
+        b->getMultiDimIdentityMap(nloops)};
+  }
+};
+
+class LhloBroadcastInDimConverter
+    : public OpConversionPattern<xla_lhlo::BroadcastInDimOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::BroadcastInDimOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::BroadcastInDimOp op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    xla_lhlo::BroadcastInDimOp::Adaptor operand_adaptor(args);
+    auto result_type = operand_adaptor.output().getType().cast<MemRefType>();
+    auto result_shape = result_type.getShape();
+
+    auto operand_and_dims = InsertReshapeIfNecessary(op, args, rewriter);
+
+    Value operand = std::get<0>(operand_and_dims);
+    auto broadcast_dims = std::get<1>(operand_and_dims);
+
+    auto loc = op.getLoc();
+    auto nloops = result_type.getRank();
+    auto operand_type = operand.getType().cast<MemRefType>();
+
+    // For a degenerate case, i.e. broadcasting with expansion of
+    // memref<1xELEMENT_TYPE>, the operand is not passed to `linalg.generic`.
+    // Instead the value is loaded and used directly in `linalg.yield`.
+    if (operand_type.getRank() == 1 &&
+        operand_type.getDimSize(0) <
+            result_type.getDimSize(broadcast_dims.front())) {
+      Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+      Value val =
+          rewriter.create<LoadOp>(loc, operand, llvm::makeArrayRef({zero}));
+      rewriter.create<linalg::GenericOp>(
+          loc, llvm::None, llvm::makeArrayRef(operand_adaptor.output()),
+          /*inputCount=*/0, /*outputCount=*/1,
+          llvm::makeArrayRef(rewriter.getMultiDimIdentityMap(nloops)),
+          GetNParallelLoopsAttrs(nloops),
+          [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+            nestedBuilder.create<linalg::YieldOp>(loc, val);
+          });
+
+    } else {
+      auto indexing_maps = getIndexingMaps(op, broadcast_dims, result_shape,
+                                           operand_type, &rewriter);
+      rewriter.create<linalg::GenericOp>(
+          loc, llvm::None,
+          llvm::makeArrayRef({operand, operand_adaptor.output()}),
+          /*inputCount=*/1, /*outputCount=*/1, indexing_maps,
+          GetNParallelLoopsAttrs(nloops),
+          [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+            nestedBuilder.create<linalg::YieldOp>(loc, *args.begin());
+          });
+    }
+    rewriter.replaceOp(op, llvm::None);
+    return success();
+  }
+
+  // Inserts 'linalg.reshape' if there is a size-1 dim expansion.
+  std::pair<Value, SmallVector<int64_t, 2>> InsertReshapeIfNecessary(
+      xla_lhlo::BroadcastInDimOp op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const {
+    xla_lhlo::BroadcastInDimOp::Adaptor operand_adaptor(args);
+    Value operand = operand_adaptor.operand();
+    auto operand_type = operand_adaptor.operand().getType().cast<MemRefType>();
+    auto operand_shape = operand_type.getShape();
+
+    Value result = operand_adaptor.output();
+    auto result_type = result.getType().cast<MemRefType>();
+    auto result_shape = result_type.getShape();
+
+    SmallVector<int64_t, 2> operand_strides;
+    int64_t operand_offset;
+    if (failed(getStridesAndOffset(operand_type, operand_strides,
+                                   operand_offset))) {
+      op.emitOpError() << "Failed to get offset and strides.";
+    }
+
+    SmallVector<int64_t, 2> new_shape, new_strides, broadcast_dims;
+    SmallVector<linalg::ReassociationIndices, 4> collapsed_dims_list;
+    linalg::ReassociationIndices collapsed_dims;
+    for (const auto& item :
+         enumerate(op.broadcast_dimensions().getIntValues())) {
+      size_t index = item.index();
+      int dim = item.value().getSExtValue();
+
+      collapsed_dims.push_back(index);
+
+      bool expansion_needed =
+          operand_shape[index] == 1 && result_shape[dim] != 1;
+      if (expansion_needed) {
+        continue;
+      }
+      new_shape.push_back(operand_shape[index]);
+      new_strides.push_back(operand_strides[index]);
+      broadcast_dims.push_back(dim);
+
+      collapsed_dims_list.push_back(collapsed_dims);
+      collapsed_dims.clear();
+    }
+    // If `collapsed_dims_list` is empty, then the memref has shape [1, ..., 1]
+    // and all dimensions need expansion. Such memref will be reshaped to a 1D
+    // memref with a single element. New shape and strides needs to be updated
+    // accordingly.
+    if (collapsed_dims_list.empty()) {
+      collapsed_dims_list.push_back({});
+      new_shape.push_back(1);
+      new_strides.push_back(1);
+      broadcast_dims.push_back(0);
+    }
+    for (const auto& dims : collapsed_dims) {
+      collapsed_dims_list.back().push_back(dims);
+    }
+
+    // `linalg.reshape` is inserted only if necessary, i.e. when the rank can be
+    // reduced.
+    if (new_shape.size() < operand_shape.size()) {
+      auto new_memref_type = MemRefType::get(
+          new_shape, operand_type.getElementType(),
+          makeStridedLinearLayoutMap(new_strides, operand_offset,
+                                     rewriter.getContext()));
+      operand = rewriter.create<linalg::ReshapeOp>(op.getLoc(), new_memref_type,
+                                                   operand_adaptor.operand(),
+                                                   collapsed_dims_list);
+    }
+    return std::make_pair(operand, broadcast_dims);
+  }
+
+  SmallVector<AffineMap, 2> getIndexingMaps(xla_lhlo::BroadcastInDimOp op,
+                                            ArrayRef<int64_t> broadcastDims,
+                                            ArrayRef<int64_t> resultShape,
+                                            MemRefType operandType,
+                                            Builder* b) const {
+    unsigned nloops = resultShape.size();
+
+    // The input is a scalar, i.e. this is a scalar broadcast op.
+    if (operandType.getRank() == 0) {
+      return {AffineMap::get(nloops, /*symbolCount=*/0, b->getContext()),
+              b->getMultiDimIdentityMap(nloops)};
+    }
+
+    auto operandShape = operandType.getShape();
+    SmallVector<AffineExpr, 4> dimExprs;
+    dimExprs.reserve(nloops);
+
+    for (const auto& broadcastDim : llvm::enumerate(broadcastDims)) {
+      int size = broadcastDim.value();
+      bool expansion_needed =
+          operandShape[broadcastDim.index()] == 1 && resultShape[size] != 1;
+      if (expansion_needed) {
+        op.emitOpError(
+            "BroadcastInDimOp lowering to Linalg does not support size-1 "
+            "dimensions expansion.");
+      }
+      dimExprs.push_back(b->getAffineDimExpr(size));
+    }
+    return {
+        AffineMap::get(nloops, /*symbolCount=*/0, dimExprs, b->getContext()),
+        b->getMultiDimIdentityMap(nloops)};
+  }
+};
+
+template <typename OpTy, bool isLHLO = true>
+class TransposeConverter
+    : public DataMovementOpConverter<TransposeConverter<OpTy, isLHLO>, OpTy,
+                                     isLHLO> {
+ public:
+  using DataMovementOpConverter<TransposeConverter<OpTy, isLHLO>, OpTy,
+                                isLHLO>::DataMovementOpConverter;
+  static SmallVector<AffineMap, 2> getIndexingMaps(OpTy op, Builder* b) {
+    auto resultType =
+        getXLAOpResultType<isLHLO>(op).template cast<ShapedType>();
+    auto nloops = resultType.getRank();
+    SmallVector<AffineExpr, 2> inputExprs;
+    inputExprs.resize(resultType.getRank());
+    for (auto permutation : llvm::enumerate(op.permutation())) {
+      inputExprs[permutation.value().getZExtValue()] =
+          b->getAffineDimExpr(permutation.index());
+    }
+    return {
+        AffineMap::get(nloops, /*symbolCount=*/0, inputExprs, b->getContext()),
+        b->getMultiDimIdentityMap(nloops)};
+  }
+};
+
+// Converts reshape ops that can be proven to be either a collapse of dimensions
+// or expansion of dimensions of the operand.
+template <typename OpTy, bool isLHLO = true>
+class ReshapeOpConverter : public OpConversionPattern<OpTy> {
+ public:
+  using OpConversionPattern<OpTy>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      OpTy reshapeOp, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    if (!verifyXLAOpBufferOrTensorSemantics<isLHLO>(reshapeOp))
+      return failure();
+    ShapedType operandType =
+        reshapeOp.operand().getType().template cast<ShapedType>();
+    ShapedType resultType = getXLAOpResultType<isLHLO>(reshapeOp);
+
+    if (!operandType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+
+    // Compute the reassociation maps for the linalg operation.
+    ArrayRef<int64_t> srcShape =
+        (operandType.getRank() > resultType.getRank() ? operandType.getShape()
+                                                      : resultType.getShape());
+    ArrayRef<int64_t> dstShape =
+        (operandType.getRank() > resultType.getRank() ? resultType.getShape()
+                                                      : operandType.getShape());
+    unsigned currSrcDim = 0, currDstDim = 0;
+    SmallVector<linalg::ReassociationExprs, 4> reassociationMap(
+        dstShape.size());
+    while (currSrcDim < srcShape.size() && currDstDim < dstShape.size()) {
+      int64_t dstSize = dstShape[currDstDim];
+      int64_t srcSize = srcShape[currSrcDim];
+      while (srcSize < dstSize && currSrcDim < srcShape.size()) {
+        reassociationMap[currDstDim].push_back(
+            rewriter.getAffineDimExpr(currSrcDim++));
+        srcSize *= srcShape[currSrcDim];
+      }
+      if (srcSize == dstSize) {
+        reassociationMap[currDstDim].push_back(
+            rewriter.getAffineDimExpr(currSrcDim++));
+        // If the next dim in dstShape is not 1, treat subsequent dims in
+        // srcShape which are 1 to be collapsed.
+        if (currDstDim == dstShape.size() - 1 ||
+            dstShape[currDstDim + 1] != 1) {
+          while (currSrcDim < srcShape.size() && srcShape[currSrcDim] == 1) {
+            reassociationMap[currDstDim].push_back(
+                rewriter.getAffineDimExpr(currSrcDim++));
+          }
+        }
+      } else {
+        return failure();
+      }
+      currDstDim++;
+    }
+    if (currSrcDim != srcShape.size()) return failure();
+
+    if (isLHLO) {
+      Value reshapeBuffer = rewriter.create<linalg::ReshapeOp>(
+          reshapeOp.getLoc(), resultType, args[0], reassociationMap);
+      rewriter.replaceOpWithNewOp<linalg::CopyOp>(
+          reshapeOp, reshapeBuffer, args[1], /*inputPermutation =*/nullptr,
+          /*outputPermutation =*/nullptr);
+    } else {
+      rewriter.replaceOpWithNewOp<linalg::TensorReshapeOp>(
+          reshapeOp, resultType, args[0], reassociationMap);
+    }
+    return success();
+  }
+};
+
+class IotaConverter : public OpConversionPattern<xla_lhlo::IotaOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::IotaOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::IotaOp iotaOp, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto resultMemrefType =
+        iotaOp.getOperand().getType().dyn_cast<MemRefType>();
+    if (!resultMemrefType) return failure();
+
+    auto resultElementType = resultMemrefType.getElementType();
+    if (!resultElementType.isSignlessIntOrFloat()) return failure();
+
+    // Construct the indexing maps needed for linalg.generic ops.
+    unsigned nloops = resultMemrefType.getRank();
+
+    rewriter.create<linalg::IndexedGenericOp>(
+        iotaOp.getLoc(), ArrayRef<Type>{}, args,
+        0,  // args_in
+        1,  // args_out
+        llvm::makeArrayRef(rewriter.getMultiDimIdentityMap(nloops)),
+        GetNParallelLoopsAttrs(nloops),
+        [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange ivs,
+            ValueRange args) {
+          Value castOp = nestedBuilder.create<IndexCastOp>(
+              nestedLoc, ivs[iotaOp.iota_dimension().getZExtValue()],
+              nestedBuilder.getIntegerType(
+                  resultElementType.getIntOrFloatBitWidth()));
+          if (resultElementType.isa<FloatType>()) {
+            castOp = nestedBuilder.create<SIToFPOp>(nestedLoc, castOp,
+                                                    resultElementType);
+          }
+          nestedBuilder.create<linalg::YieldOp>(nestedLoc, castOp);
+        });
+
+    rewriter.replaceOp(iotaOp, llvm::None);
+    return success();
+  }
+};
+
+class ConstConverter : public OpConversionPattern<xla_lhlo::ConstOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::ConstOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::ConstOp constOp, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = constOp.getLoc();
+    auto valueAttr = constOp.value().cast<DenseElementsAttr>();
+    if (valueAttr.getType().getRank() != 0) return failure();
+    auto stdConstOp =
+        rewriter.create<mlir::ConstantOp>(loc, valueAttr.getValue({}));
+    rewriter.create<mlir::StoreOp>(loc, stdConstOp, constOp.getOperand());
+    rewriter.eraseOp(constOp);
+    return success();
+  }
+};
+
+// TODO(b/156787842): Support the lowering for dynamic shapes.
+template <typename OpTy, bool isLHLO = true>
+class ReverseConverter
+    : public DataMovementOpConverter<ReverseConverter<OpTy, isLHLO>, OpTy,
+                                     isLHLO> {
+ public:
+  using DataMovementOpConverter<ReverseConverter<OpTy, isLHLO>, OpTy,
+                                isLHLO>::DataMovementOpConverter;
+  static SmallVector<AffineMap, 2> getIndexingMaps(OpTy op, Builder* b) {
+    auto resultType =
+        getXLAOpResultType<isLHLO>(op).template cast<ShapedType>();
+    auto nloops = resultType.getRank();
+    SmallVector<AffineExpr, 2> inputExprs;
+    inputExprs.reserve(nloops);
+    for (int i = 0; i < nloops; ++i)
+      inputExprs.push_back(b->getAffineDimExpr(i));
+    for (auto dim : op.dimensions()) {
+      int i = dim.getZExtValue();
+      if (resultType.isDynamicDim(i)) return {};
+      int n = resultType.getShape()[i];
+      inputExprs[i] = b->getAffineConstantExpr(n - 1) - inputExprs[i];
+    }
+    return {
+        AffineMap::get(nloops, /*symbolCount=*/0, inputExprs, b->getContext()),
+        b->getMultiDimIdentityMap(nloops)};
+  }
+};
+
+class SliceConverter : public OpConversionPattern<xla_lhlo::SliceOp> {
+ public:
+  using OpConversionPattern<xla_lhlo::SliceOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_lhlo::SliceOp sliceOp, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = sliceOp.getLoc();
+    auto argType =
+        sliceOp.getOperand(0).getType().template dyn_cast<ShapedType>();
+    if (!argType || !argType.hasRank()) {
+      emitError(loc, "lhlo to linalg conversion expects known-rank args");
+      return failure();
+    }
+
+    SmallVector<Value, 3> ranges;
+    for (int i = 0, e = argType.getRank(); i < e; ++i) {
+      Value start_index = rewriter.create<ConstantIndexOp>(
+          loc, sliceOp.start_indices().getValue<int64_t>(i));
+      Value limit_index = rewriter.create<ConstantIndexOp>(
+          loc, sliceOp.limit_indices().getValue<int64_t>(i));
+      Value stride = rewriter.create<ConstantIndexOp>(
+          loc, sliceOp.strides().getValue<int64_t>(i));
+      ranges.push_back(rewriter.create<linalg::RangeOp>(loc, start_index,
+                                                        limit_index, stride));
+    }
+    auto linalg_slice =
+        rewriter.create<linalg::SliceOp>(loc, sliceOp.getOperand(0), ranges);
+    rewriter.create<linalg::CopyOp>(loc, linalg_slice, sliceOp.getOperand(1));
+    rewriter.eraseOp(sliceOp);
+    return success();
+  }
+};
+
+void populateLHLOToLinalgConversionPattern(MLIRContext* context,
+                                           OwningRewritePatternList* patterns) {
+  // clang-format off
+  patterns->insert<BroadcastConverter<xla_lhlo::BroadcastOp>,
+                   ConstConverter,
+                   ConvToLinalgConverter,
+                   IotaConverter,
+                   LhloBroadcastInDimConverter,
+                   PointwiseToLinalgConverter<xla_lhlo::AbsOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::AddOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::AndOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::CeilOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::CompareOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::ComplexOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::ConvertOp>,
+                   // TODO(ataei): Remove this pattern, CopyOp is folded away.
+                   PointwiseToLinalgConverter<xla_lhlo::CopyOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::CosOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::DivOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::ExpOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::ImagOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::LogOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::MaxOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::MinOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::MulOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::NegOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::RealOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::RemOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::RsqrtOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::SelectOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::SignOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::SinOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::SqrtOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::SubOp>,
+                   PointwiseToLinalgConverter<xla_lhlo::TanhOp>,
+                   ReshapeOpConverter<xla_lhlo::ReshapeOp>,
+                   ReverseConverter<xla_lhlo::ReverseOp>,
+                   ScalarPointwiseToStandardConverter<xla_lhlo::AddOp>,
+                   SliceConverter
+                  >(context);
+  // clang-format on
+}
+
+// Converts LHLO ops to Linalg generic.
+// Sample result for xla_lhlo::AddOp.
+//
+// "xla_lhlo.add"(%arg1, %arg2, %out) :
+//      (memref<2x2xf32>, memref<2x2xf32>, memref<2x2xf32>) -> ()
+//
+// will be converted to
+//
+// #map0 = (d0, d1) -> (d0, d1)
+// "linalg.generic"(%arg1, %arg2, %out) ( {
+//   ^bb0(%arg4: f32, %arg5: f32):
+//     %0 = addf %arg4, %arg5 : f32
+//     "linalg.yield"(%0) : (f32) -> ()
+// }) {
+//     args_in = 2,
+//     args_out = 1,
+//     indexing_maps = [#map0, #map0, #map0],
+//     iterator_types = ["parallel", "parallel"],
+// } : (memref<2x2xf32>, memref<2x2xf32>, memref<2x2xf32>) -> ()
+struct LhloLegalizeToLinalg
+    : public PassWrapper<LhloLegalizeToLinalg, FunctionPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addLegalDialect<linalg::LinalgDialect, StandardOpsDialect>();
+
+    auto func = getFunction();
+    populateLHLOToLinalgConversionPattern(func.getContext(), &patterns);
+    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
+      signalPassFailure();
+    }
+  }
+};
+
+struct HloLegalizeToLinalg
+    : public PassWrapper<HloLegalizeToLinalg, FunctionPass> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addLegalDialect<linalg::LinalgDialect, StandardOpsDialect>();
+
+    auto func = getFunction();
+    xla_hlo::populateHLOToLinalgConversionPattern(func.getContext(), &patterns);
+    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+namespace xla_lhlo {
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeLhloToLinalgPass() {
+  return absl::make_unique<LhloLegalizeToLinalg>();
+}
+
+static PassRegistration<LhloLegalizeToLinalg> legalize_lhlo_pass(
+    "lhlo-legalize-to-linalg", "Legalize from LHLO dialect to Linalg dialect");
+}  // namespace xla_lhlo
+
+namespace xla_hlo {
+
+void populateHLOToLinalgConversionPattern(MLIRContext* context,
+                                          OwningRewritePatternList* patterns) {
+  patterns->insert<BroadcastConverter<xla_hlo::BroadcastOp, false>,
+                   HloBroadcastInDimConverter,
+                   PointwiseToLinalgConverter<xla_hlo::AbsOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::AddOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::AndOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::CeilOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::CompareOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::ComplexOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::ConvertOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::CopyOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::CosOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::DivOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::ExpOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::ImagOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::LogOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::MaxOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::MinOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::MulOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::NegOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::RealOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::RemOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::RsqrtOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::SelectOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::SinOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::SqrtOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::SubOp, false>,
+                   PointwiseToLinalgConverter<xla_hlo::TanhOp, false>,
+                   ReshapeOpConverter<xla_hlo::ReshapeOp, false>,
+                   ReverseConverter<xla_hlo::ReverseOp, false>,
+                   TransposeConverter<xla_hlo::TransposeOp, false>>(context);
+}
+
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeHloToLinalgPass() {
+  return absl::make_unique<HloLegalizeToLinalg>();
+}
+
+static PassRegistration<HloLegalizeToLinalg> legalize_hlo_pass(
+    "hlo-legalize-to-linalg", "Legalize from HLO dialect to Linalg dialect");
+}  // namespace xla_hlo
+}  // namespace mlir

lib/Dialect/mhlo/transforms/xla_transform_unranked_hlo.cc

@@ -0,0 +1,188 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+==============================================================================*/
+
+#include "third_party/absl/memory/memory.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/Shape/IR/Shape.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Dialect/StandardOps/IR/Ops.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Function.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/MLIRContext.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/PatternMatch.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/IR/StandardTypes.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Pass/Pass.h"
+#include "third_party/llvm/llvm-project/mlir/include/mlir/Transforms/DialectConversion.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla_hlo {
+namespace {
+
+// TODO(frgossen): Make it variadic.
+template <typename OpTy>
+inline void AddLegalOpOnRankedTensor(ConversionTarget *target) {
+  target->addDynamicallyLegalOp<OpTy>([](OpTy op) {
+    return llvm::all_of((op.getOperation())->getOperandTypes(),
+                        [&](Type t) { return t.isa<RankedTensorType>(); });
+  });
+}
+
+/// Unary element-wise operations on unranked tensors can be applied to the
+/// flattened tensor with the same effect.
+/// This pattern rewrites every such operation to
+///   (i)   flatten the input tensor,
+///   (ii)  apply the unary operation, and
+///   (iii) restore the original shape.
+template <typename OpTy>
+struct UnaryElementwiseOpConversion : public OpRewritePattern<OpTy> {
+  explicit UnaryElementwiseOpConversion(MLIRContext *context)
+      : OpRewritePattern<OpTy>(context) {}
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Don't apply conversion to ops with statically shaped operands.
+    Value operand = op.getOperand();
+    auto operandTy = operand.getType().dyn_cast<TensorType>();
+    if (operandTy.hasRank()) return failure();
+
+    // Generate IR to flatten the operand.
+    auto loc = op.getLoc();
+    Value shape = rewriter.create<shape::ShapeOfOp>(loc, operand);
+    Value numElements = rewriter.create<shape::NumElementsOp>(
+        loc, rewriter.getType<shape::SizeType>(), shape);
+    Value numElementsAsIndex = rewriter.create<shape::SizeToIndexOp>(
+        loc, rewriter.getIndexType(), numElements);
+    Value flatShapeAsDimTensor =
+        rewriter.create<TensorFromElementsOp>(loc, numElementsAsIndex);
+    auto flatTensorTy = RankedTensorType::get({ShapedType::kDynamicSize},
+                                              operandTy.getElementType());
+    Value flatOperand = rewriter.create<xla_hlo::DynamicReshapeOp>(
+        loc, flatTensorTy, operand, flatShapeAsDimTensor);
+
+    // Generate IR for the actual operation.
+    Value flatResult = rewriter.create<OpTy>(loc, flatTensorTy, flatOperand);
+
+    // Generate IR to restore the original shape.
+    auto extentTensorTy = RankedTensorType::get({ShapedType::kDynamicSize},
+                                                rewriter.getIndexType());
+    Value shapeAsExtentTensor =
+        rewriter.create<shape::ToExtentTensorOp>(loc, extentTensorTy, shape);
+    Value result = rewriter.create<xla_hlo::DynamicReshapeOp>(
+        loc, operandTy, flatResult, shapeAsExtentTensor);
+    rewriter.replaceOp(op, result);
+
+    return success();
+  }
+};
+
+/// Binary element-wise operation on unranked tensors can be applied to the
+/// flattened operand tensors with the same effect.
+/// This pattern rewrites every such operation to
+///   (i)   flatten the operand tensors,
+///   (ii)  apply the binary operation, and
+//    (iii) restore the original shape.
+template <typename OpTy>
+struct BinaryElementwiseOpConversion : public OpRewritePattern<OpTy> {
+  explicit BinaryElementwiseOpConversion(MLIRContext *context)
+      : OpRewritePattern<OpTy>(context) {}
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Don't apply conversion unless both operands are unranked.
+    if (op.lhs().getType().template isa<RankedTensorType>() ||
+        op.rhs().getType().template isa<RankedTensorType>()) {
+      return failure();
+    }
+
+    // Flatten operands.
+    Type shapeTy = shape::ShapeType::get(rewriter.getContext());
+    auto loc = op.getLoc();
+    Value shapeLhs = rewriter.create<shape::ShapeOfOp>(loc, op.lhs());
+    Value shapeRhs = rewriter.create<shape::ShapeOfOp>(loc, op.rhs());
+    Value shape = rewriter.create<shape::AnyOp>(loc, shapeTy,
+                                                ValueRange{shapeLhs, shapeRhs});
+    Value numElements = rewriter.create<shape::NumElementsOp>(loc, shape);
+    Value numElementsAsIndex =
+        rewriter.create<shape::SizeToIndexOp>(loc, numElements);
+    Value flatShape =
+        rewriter.create<TensorFromElementsOp>(loc, numElementsAsIndex);
+    TensorType lhsTy = op.lhs().getType().template cast<TensorType>();
+    Type flatLhsTy = RankedTensorType::get({ShapedType::kDynamicSize},
+                                           lhsTy.getElementType());
+    Value flatLhs =
+        rewriter.create<DynamicReshapeOp>(loc, flatLhsTy, op.lhs(), flatShape);
+    TensorType rhsTy = op.rhs().getType().template cast<TensorType>();
+    Type flatRhsTy = RankedTensorType::get({ShapedType::kDynamicSize},
+                                           rhsTy.getElementType());
+    Value flatRhs =
+        rewriter.create<DynamicReshapeOp>(loc, flatRhsTy, op.rhs(), flatShape);
+
+    // Apply actual operation to flattened operands.
+    Value flatResult = rewriter.create<OpTy>(loc, flatLhs, flatRhs);
+
+    // Restore original shape.
+    auto extentTensorTy = RankedTensorType::get({ShapedType::kDynamicSize},
+                                                rewriter.getIndexType());
+    Value shapeAsExtentTensor =
+        rewriter.create<shape::ToExtentTensorOp>(loc, extentTensorTy, shape);
+    Value result = rewriter.create<DynamicReshapeOp>(
+        loc, op.getType(), flatResult, shapeAsExtentTensor);
+    rewriter.replaceOp(op, result);
+
+    return success();
+  }
+};
+
+struct TransformUnrankedHloPass
+    : public PassWrapper<TransformUnrankedHloPass, FunctionPass> {
+  void runOnFunction() override {
+    // Setup conversion target.
+    MLIRContext &ctx = getContext();
+    ConversionTarget target(ctx);
+    target.addLegalDialect<XlaHloDialect, StandardOpsDialect,
+                           shape::ShapeDialect>();
+    target.addLegalOp<FuncOp>();
+    AddLegalOpOnRankedTensor<SqrtOp>(&target);
+    AddLegalOpOnRankedTensor<AddOp>(&target);
+
+    // Populate rewrite patterns.
+    OwningRewritePatternList patterns;
+    PopulateTransformUnrankedHloPatterns(&ctx, &patterns);
+
+    // Apply transformation.
+    if (failed(applyFullConversion(getFunction(), target, patterns)))
+      return signalPassFailure();
+  }
+};
+
+}  // namespace
+
+void PopulateTransformUnrankedHloPatterns(MLIRContext *context,
+                                          OwningRewritePatternList *patterns) {
+  // TODO(frgossen): Populate all unary and binary operations.
+  // clang-format off
+  patterns->insert<
+      BinaryElementwiseOpConversion<AddOp>,
+      UnaryElementwiseOpConversion<SqrtOp>>(context);
+  // clang-format on
+}
+
+static PassRegistration<TransformUnrankedHloPass> transform_unranked_hlo_pass(
+    "transform-unranked-hlo",
+    "Realize element-wise operations on ranked tensors where possible");
+
+}  // namespace xla_hlo
+}  // namespace mlir

lib/utils/cycle_detector.cc

@@ -0,0 +1,340 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/utils/cycle_detector.h"
+
+#include <algorithm>
+
+#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/DenseSet.h"
+
+namespace mlir {
+
+namespace {
+
+using NodeSet = llvm::DenseSet<int32_t>;
+using OrderedNodeSet = OrderedSet<int32_t>;
+
+template <typename T>
+struct VecStruct {
+  typedef llvm::SmallVector<T, 4> type;
+};
+template <typename T>
+using Vec = typename VecStruct<T>::type;
+
+struct Node {
+  // rank number assigned by Pearce-Kelly algorithm
+  int32_t rank;
+  // Temporary marker used by depth-first-search
+  bool visited;
+  // User-supplied data
+  void* data;
+  // List of immediate predecessor nodes in graph
+  OrderedNodeSet in;
+  // List of immediate successor nodes in graph
+  OrderedNodeSet out;
+};
+
+}  // namespace
+
+struct GraphCycles::Rep {
+  Vec<Node*> nodes;
+  // Indices for unused entries in nodes
+  Vec<int32_t> free_nodes;
+
+  // Temporary state.
+  // Results of forward DFS
+  Vec<int32_t> deltaf;
+  // Results of backward DFS
+  Vec<int32_t> deltab;
+  // All nodes to reprocess
+  Vec<int32_t> list;
+  // Rank values to assign to list entries
+  Vec<int32_t> merged;
+  // Emulates recursion stack when doing depth first search
+  Vec<int32_t> stack;
+};
+
+GraphCycles::GraphCycles(int32_t num_nodes) : rep_(new Rep) {
+  rep_->nodes.reserve(num_nodes);
+  for (int32_t i = 0; i < num_nodes; ++i) {
+    Node* n = new Node;
+    n->visited = false;
+    n->data = nullptr;
+    n->rank = rep_->nodes.size();
+    rep_->nodes.push_back(n);
+  }
+}
+
+GraphCycles::~GraphCycles() {
+  for (Vec<Node*>::size_type i = 0, e = rep_->nodes.size(); i < e; ++i) {
+    delete rep_->nodes[i];
+  }
+  delete rep_;
+}
+
+bool GraphCycles::HasEdge(int32_t x, int32_t y) const {
+  return rep_->nodes[x]->out.Contains(y);
+}
+
+void GraphCycles::RemoveEdge(int32_t x, int32_t y) {
+  rep_->nodes[x]->out.Erase(y);
+  rep_->nodes[y]->in.Erase(x);
+  // No need to update the rank assignment since a previous valid
+  // rank assignment remains valid after an edge deletion.
+}
+
+static bool ForwardDFS(GraphCycles::Rep* r, int32_t n, int32_t upper_bound);
+static void BackwardDFS(GraphCycles::Rep* r, int32_t n, int32_t lower_bound);
+static void Reorder(GraphCycles::Rep* r);
+static void Sort(const Vec<Node*>&, Vec<int32_t>* delta);
+static void MoveToList(GraphCycles::Rep* r, Vec<int32_t>* src,
+                       Vec<int32_t>* dst);
+static void ClearVisitedBits(GraphCycles::Rep* r, const Vec<int32_t>& nodes);
+
+bool GraphCycles::InsertEdge(int32_t x, int32_t y) {
+  if (x == y) return false;
+  Rep* r = rep_;
+  Node* nx = r->nodes[x];
+  if (!nx->out.Insert(y)) {
+    // Edge already exists.
+    return true;
+  }
+
+  Node* ny = r->nodes[y];
+  ny->in.Insert(x);
+
+  if (nx->rank <= ny->rank) {
+    // New edge is consistent with existing rank assignment.
+    return true;
+  }
+
+  // Current rank assignments are incompatible with the new edge.  Recompute.
+  // We only need to consider nodes that fall in the range [ny->rank,nx->rank].
+  if (ForwardDFS(r, y, nx->rank)) {
+    // Found a cycle.  Undo the insertion and tell caller.
+    nx->out.Erase(y);
+    ny->in.Erase(x);
+    // Since we do not call Reorder() on this path, clear any visited
+    // markers left by ForwardDFS.
+    ClearVisitedBits(r, r->deltaf);
+    return false;
+  }
+  BackwardDFS(r, x, ny->rank);
+  Reorder(r);
+  return true;
+}
+
+// Follows the edges from producer to consumer and searchs if the node having
+// rank `n` can reach the node having rank `upper_bound` using a DFS search.
+// When doing DFS search, We only consider the pathes that satisfy the ranks
+// of the nodes of the path are all smaller than `upper_bound`.
+//
+// Returns true if such path exists.
+static bool ForwardDFS(GraphCycles::Rep* r, int32_t n, int32_t upper_bound) {
+  // Avoid recursion since stack space might be limited.
+  // We instead keep a stack of nodes to visit.
+  r->deltaf.clear();
+  r->stack.clear();
+  r->stack.push_back(n);
+  while (!r->stack.empty()) {
+    n = r->stack.back();
+    r->stack.pop_back();
+    Node* nn = r->nodes[n];
+    if (nn->visited) continue;
+
+    nn->visited = true;
+    r->deltaf.push_back(n);
+
+    for (auto w : nn->out.GetSequence()) {
+      Node* nw = r->nodes[w];
+      if (nw->rank == upper_bound) {
+        return true;
+      }
+      if (!nw->visited && nw->rank < upper_bound) {
+        r->stack.push_back(w);
+      }
+    }
+  }
+  return false;
+}
+
+// Follows the edges from consumer to producer and visit all the nodes that
+// is reachable from node `n` and have rank larger than `lower_bound`.
+static void BackwardDFS(GraphCycles::Rep* r, int32_t n, int32_t lower_bound) {
+  r->deltab.clear();
+  r->stack.clear();
+  r->stack.push_back(n);
+  while (!r->stack.empty()) {
+    n = r->stack.back();
+    r->stack.pop_back();
+    Node* nn = r->nodes[n];
+    if (nn->visited) continue;
+
+    nn->visited = true;
+    r->deltab.push_back(n);
+
+    for (auto w : nn->in.GetSequence()) {
+      Node* nw = r->nodes[w];
+      if (!nw->visited && lower_bound < nw->rank) {
+        r->stack.push_back(w);
+      }
+    }
+  }
+}
+
+// Recomputes rank assignments to make them compatible with the edges (producer
+// has smaller rank than its consumer)
+static void Reorder(GraphCycles::Rep* r) {
+  Sort(r->nodes, &r->deltab);
+  Sort(r->nodes, &r->deltaf);
+
+  // Adds contents of delta lists to list (backwards deltas first).
+  r->list.clear();
+  MoveToList(r, &r->deltab, &r->list);
+  MoveToList(r, &r->deltaf, &r->list);
+
+  // Produce sorted list of all ranks that will be reassigned.
+  r->merged.resize(r->deltab.size() + r->deltaf.size());
+  std::merge(r->deltab.begin(), r->deltab.end(), r->deltaf.begin(),
+             r->deltaf.end(), r->merged.begin());
+
+  // Assign the ranks in order to the collected list.
+  for (Vec<int32_t>::size_type i = 0, e = r->list.size(); i < e; ++i) {
+    r->nodes[r->list[i]]->rank = r->merged[i];
+  }
+}
+
+// Sorts nodes in the vector according to their ranks. Small rank first.
+static void Sort(const Vec<Node*>& nodes, Vec<int32_t>* delta) {
+  struct ByRank {
+    const Vec<Node*>* nodes;
+    bool operator()(int32_t a, int32_t b) const {
+      return (*nodes)[a]->rank < (*nodes)[b]->rank;
+    }
+  };
+  ByRank cmp;
+  cmp.nodes = &nodes;
+  std::sort(delta->begin(), delta->end(), cmp);
+}
+
+// Collects ranks of nodes in vector `src` to vector `dst`
+static void MoveToList(GraphCycles::Rep* r, Vec<int32_t>* src,
+                       Vec<int32_t>* dst) {
+  for (Vec<int32_t>::size_type i = 0, e = src->size(); i < e; i++) {
+    int32_t w = (*src)[i];
+    // Replace src entry with its rank
+    (*src)[i] = r->nodes[w]->rank;
+    // Prepare for future DFS calls
+    r->nodes[w]->visited = false;
+    dst->push_back(w);
+  }
+}
+
+// Clears bookkeeping fileds used during the last DFS process.
+static void ClearVisitedBits(GraphCycles::Rep* r, const Vec<int32_t>& nodes) {
+  for (Vec<int32_t>::size_type i = 0, e = nodes.size(); i < e; i++) {
+    r->nodes[nodes[i]]->visited = false;
+  }
+}
+
+bool GraphCycles::IsReachable(int32_t x, int32_t y) {
+  if (x == y) return true;
+  Rep* r = rep_;
+  Node* nx = r->nodes[x];
+  Node* ny = r->nodes[y];
+
+  if (nx->rank >= ny->rank) {
+    // x cannot reach y since it is after it in the topological ordering
+    return false;
+  }
+
+  // See if x can reach y using a DFS search that is limited to y's rank
+  bool reachable = ForwardDFS(r, x, ny->rank);
+
+  // Clear any visited markers left by ForwardDFS.
+  ClearVisitedBits(r, r->deltaf);
+  return reachable;
+}
+
+llvm::Optional<int32_t> GraphCycles::ContractEdge(int32_t a, int32_t b) {
+  assert(HasEdge(a, b));
+  RemoveEdge(a, b);
+
+  if (IsReachable(a, b)) {
+    // Restore the graph to its original state.
+    InsertEdge(a, b);
+    return {};
+  }
+
+  if (rep_->nodes[b]->in.Size() + rep_->nodes[b]->out.Size() >
+      rep_->nodes[a]->in.Size() + rep_->nodes[a]->out.Size()) {
+    // Swap "a" and "b" to minimize copying.
+    std::swap(a, b);
+  }
+
+  Node* nb = rep_->nodes[b];
+  OrderedNodeSet out = std::move(nb->out);
+  OrderedNodeSet in = std::move(nb->in);
+  for (int32_t y : out.GetSequence()) {
+    rep_->nodes[y]->in.Erase(b);
+  }
+  for (int32_t y : in.GetSequence()) {
+    rep_->nodes[y]->out.Erase(b);
+  }
+  rep_->free_nodes.push_back(b);
+
+  rep_->nodes[a]->out.Reserve(rep_->nodes[a]->out.Size() + out.Size());
+  for (int32_t y : out.GetSequence()) {
+    InsertEdge(a, y);
+  }
+
+  rep_->nodes[a]->in.Reserve(rep_->nodes[a]->in.Size() + in.Size());
+  for (int32_t y : in.GetSequence()) {
+    InsertEdge(y, a);
+  }
+
+  // Note, if the swap happened it might be what originally was called "b".
+  return a;
+}
+
+std::vector<int32_t> GraphCycles::SuccessorsCopy(int32_t node) const {
+  return rep_->nodes[node]->out.GetSequence();
+}
+
+namespace {
+void SortInPostOrder(const Vec<Node*>& nodes, std::vector<int32_t>* to_sort) {
+  std::sort(to_sort->begin(), to_sort->end(), [&](int32_t a, int32_t b) {
+    return nodes[a]->rank > nodes[b]->rank;
+  });
+}
+}  // namespace
+
+std::vector<int32_t> GraphCycles::AllNodesInPostOrder() const {
+  llvm::DenseSet<int32_t> free_nodes_set;
+  for (int32_t n : rep_->free_nodes) free_nodes_set.insert(n);
+
+  std::vector<int32_t> all_nodes;
+  all_nodes.reserve(rep_->nodes.size() - free_nodes_set.size());
+  for (size_t i = 0, e = rep_->nodes.size(); i < e; i++) {
+    if (!free_nodes_set.count(i)) {
+      all_nodes.push_back(i);
+    }
+  }
+
+  SortInPostOrder(rep_->nodes, &all_nodes);
+  return all_nodes;
+}
+
+}  // namespace mlir

lib/utils/cycle_detector_test.cc

@@ -0,0 +1,89 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "third_party/tensorflow/compiler/mlir/hlo/include/mlir-hlo/utils/cycle_detector.h"
+
+#include "third_party/tensorflow/compiler/xla/test.h"
+
+class GraphCyclesTest : public ::testing::Test {
+ public:
+  GraphCyclesTest() : g_(100) {}
+
+  bool AddEdge(int x, int y) { return g_.InsertEdge(x, y); }
+
+  void AddMultiples() {
+    // For every node x > 0: add edge to 2*x, 3*x
+    for (int x = 1; x < 25; x++) {
+      EXPECT_TRUE(AddEdge(x, 2 * x)) << x;
+      EXPECT_TRUE(AddEdge(x, 3 * x)) << x;
+    }
+  }
+
+  mlir::GraphCycles g_;
+};
+
+TEST_F(GraphCyclesTest, NoCycle) { AddMultiples(); }
+
+TEST_F(GraphCyclesTest, SimpleCycle) {
+  AddMultiples();
+  EXPECT_FALSE(AddEdge(8, 4));
+}
+
+TEST_F(GraphCyclesTest, IndirectCycle) {
+  AddMultiples();
+  EXPECT_TRUE(AddEdge(16, 9));
+  EXPECT_FALSE(AddEdge(9, 2));
+}
+
+TEST_F(GraphCyclesTest, RemoveEdge) {
+  EXPECT_TRUE(AddEdge(1, 2));
+  EXPECT_TRUE(AddEdge(2, 3));
+  EXPECT_TRUE(AddEdge(3, 4));
+  EXPECT_TRUE(AddEdge(4, 5));
+  g_.RemoveEdge(2, 3);
+  EXPECT_FALSE(g_.HasEdge(2, 3));
+}
+
+TEST_F(GraphCyclesTest, IsReachable) {
+  EXPECT_TRUE(AddEdge(1, 2));
+  EXPECT_TRUE(AddEdge(2, 3));
+  EXPECT_TRUE(AddEdge(3, 4));
+  EXPECT_TRUE(AddEdge(4, 5));
+
+  EXPECT_TRUE(g_.IsReachable(1, 5));
+  EXPECT_FALSE(g_.IsReachable(5, 1));
+}
+
+TEST_F(GraphCyclesTest, ContractEdge) {
+  ASSERT_TRUE(AddEdge(1, 2));
+  ASSERT_TRUE(AddEdge(1, 3));
+  ASSERT_TRUE(AddEdge(2, 3));
+  ASSERT_TRUE(AddEdge(2, 4));
+  ASSERT_TRUE(AddEdge(3, 4));
+
+  // It will introduce a cycle if the edge is contracted
+  EXPECT_FALSE(g_.ContractEdge(1, 3).hasValue());
+  EXPECT_TRUE(g_.HasEdge(1, 3));
+
+  // Node (2) has more edges.
+  EXPECT_EQ(*g_.ContractEdge(1, 2), 2);
+  EXPECT_TRUE(g_.HasEdge(2, 3));
+  EXPECT_TRUE(g_.HasEdge(2, 4));
+  EXPECT_TRUE(g_.HasEdge(3, 4));
+
+  // Node (2) has more edges.
+  EXPECT_EQ(*g_.ContractEdge(2, 3), 2);
+  EXPECT_TRUE(g_.HasEdge(2, 4));
+}