From e9077ee0471383a00d1e0867761ecc8c752ab0c6 Mon Sep 17 00:00:00 2001 From: Natasha Kononenko Date: Sat, 30 May 2020 02:28:57 -0700 Subject: [PATCH] Add division, min and max constant folding Note: For integers, these operations are performed in their signed versions. PiperOrigin-RevId: 313918602 Change-Id: Icc1b68219351985f537ad69b45f7328f65d5d223 --- tensorflow/compiler/mlir/xla/ir/hlo_ops.cc | 39 ++++++++++++++ tensorflow/compiler/mlir/xla/ir/hlo_ops.td | 3 ++ .../compiler/mlir/xla/tests/canonicalize.mlir | 54 +++++++++++++++++++ 3 files changed, 96 insertions(+) diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc index 22f1ffe1d3a..c66b8f12332 100644 --- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc +++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc @@ -1474,6 +1474,38 @@ static Attribute BinaryFolder(Op* op, ArrayRef attrs) { return DenseElementsAttr::get(type, values); } +template +struct divide : std::divides {}; + +template <> +struct divide { + APInt operator()(const APInt& a, const APInt& b) const { return a.sdiv(b); } +}; + +template +struct max { + T operator()(const T& a, const T& b) const { return std::max(a, b); } +}; + +template <> +struct max { + APInt operator()(const APInt& a, const APInt& b) const { + return llvm::APIntOps::smax(a, b); + } +}; + +template +struct min { + T operator()(const T& a, const T& b) const { return std::min(a, b); } +}; + +template <> +struct min { + APInt operator()(const APInt& a, const APInt& b) const { + return llvm::APIntOps::smin(a, b); + } +}; + #define BINARY_FOLDER(Op, Func) \ OpFoldResult Op::fold(ArrayRef attrs) { \ if (getElementTypeOrSelf(getType()).isa()) \ @@ -1483,9 +1515,16 @@ static Attribute BinaryFolder(Op* op, ArrayRef attrs) { return {}; \ } +// Addition, subtraction and multiplication use the std:: versions of the ops. +// Due to the other ops behaving differently in signed vs unsigned integers, +// APInts need a special implementation. Currently, it replicates signed int +// op behavior. BINARY_FOLDER(AddOp, std::plus); BINARY_FOLDER(SubOp, std::minus); BINARY_FOLDER(MulOp, std::multiplies); +BINARY_FOLDER(DivOp, divide); +BINARY_FOLDER(MaxOp, max); +BINARY_FOLDER(MinOp, min); #undef BINARY_FOLDER diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops.td index be8c86a3a0b..97b8e1c1863 100644 --- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.td +++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.td @@ -303,14 +303,17 @@ def HLO_ComplexOp: HLO_Op<"complex", def HLO_DivOp : HLO_BinaryElementwiseOp<"divide", [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_DivOp { + let hasFolder = 1; } def HLO_MaxOp : HLO_BinaryElementwiseOp<"maximum", [Commutative, NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_MaxOp { + let hasFolder = 1; } def HLO_MinOp : HLO_BinaryElementwiseOp<"minimum", [Commutative, NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_MinOp { + let hasFolder = 1; } def HLO_MulOp : HLO_BinaryElementwiseOp<"multiply", diff --git a/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir b/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir index ed9f1661616..6b9ed3e463a 100644 --- a/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir +++ b/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir @@ -45,6 +45,60 @@ func @multiply_scalar_fold() -> tensor<4xi64> { return %2 : tensor<4xi64> } +// CHECK-LABEL: divide_scalar_fold +func @divide_scalar_fold() -> tensor<4xi64> { + %0 = xla_hlo.constant dense<7> : tensor<4xi64> + %1 = xla_hlo.constant dense<5> : tensor<4xi64> + // CHECK: xla_hlo.constant dense<1> + %2 = "xla_hlo.divide"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) + return %2 : tensor<4xi64> +} + +// CHECK-LABEL: divide_fold_float +func @divide_fold_float() -> tensor<4xf64> { + %0 = xla_hlo.constant dense<[5.0, 66.0, 5.0, 1.0]> : tensor<4xf64> + %1 = xla_hlo.constant dense<[5.0, 3.0, 2.0, 4.0]> : tensor<4xf64> + // CHECK: xla_hlo.constant dense<[1.000000e+00, 2.200000e+01, 2.500000e+00, 2.500000e-01]> + %2 = "xla_hlo.divide"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) + return %2 : tensor<4xf64> +} + +// CHECK-LABEL: max_scalar_fold +func @max_scalar_fold() -> tensor<4xi64> { + %0 = xla_hlo.constant dense<7> : tensor<4xi64> + %1 = xla_hlo.constant dense<5> : tensor<4xi64> + // CHECK: xla_hlo.constant dense<7> + %2 = "xla_hlo.maximum"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) + return %2 : tensor<4xi64> +} + +// CHECK-LABEL: max_fold_float +func @max_fold_float() -> tensor<4xf64> { + %0 = xla_hlo.constant dense<[5.0, 66.0, 5.0, 1.0]> : tensor<4xf64> + %1 = xla_hlo.constant dense<[5.0, 3.0, 2.0, 4.0]> : tensor<4xf64> + // CHECK: xla_hlo.constant dense<[5.000000e+00, 6.600000e+01, 5.000000e+00, 4.000000e+00]> + %2 = "xla_hlo.maximum"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) + return %2 : tensor<4xf64> +} + +// CHECK-LABEL: min_scalar_fold +func @min_scalar_fold() -> tensor<4xi64> { + %0 = xla_hlo.constant dense<7> : tensor<4xi64> + %1 = xla_hlo.constant dense<-5> : tensor<4xi64> + // CHECK: xla_hlo.constant dense<-5> + %2 = "xla_hlo.minimum"(%0, %1) : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>) + return %2 : tensor<4xi64> +} + +// CHECK-LABEL: min_fold_float +func @min_fold_float() -> tensor<4xf64> { + %0 = xla_hlo.constant dense<[5.0, 66.0, 5.0, 1.0]> : tensor<4xf64> + %1 = xla_hlo.constant dense<[5.0, 3.0, 2.0, 4.0]> : tensor<4xf64> + // CHECK: xla_hlo.constant dense<[5.000000e+00, 3.000000e+00, 2.000000e+00, 1.000000e+00]> + %2 = "xla_hlo.minimum"(%0, %1) : (tensor<4xf64>, tensor<4xf64>) -> (tensor<4xf64>) + return %2 : tensor<4xf64> +} + // CHECK-LABEL: concatenate_noop func @concatenate_noop(%arg0: tensor<4xi32>) -> tensor<4xi32> { // CHECK-SAME: [[ARG:%.+]]: tensor<4xi32>