Folder for mhlo.negate

PiperOrigin-RevId: 331592342
This commit is contained in:
Robert Suderman 2020-09-14 11:36:29 -07:00 committed by TensorFlow MLIR Team
parent 8bcca69845
commit dff98f8cd5
3 changed files with 66 additions and 2 deletions

View File

@ -224,11 +224,14 @@ def HLO_LogisticOp: HLO_UnaryElementwiseOp<"logistic",
def HLO_NotOp: HLO_UnaryElementwiseOp<"not",
[NoSideEffect, SameOperandsAndResultType], HLO_PredOrIntTensor>,
BASE_HLO_NotOp;
BASE_HLO_NotOp {
}
def HLO_NegOp: HLO_UnaryElementwiseOp<"negate",
[NoSideEffect, SameOperandsAndResultType], HLO_IntFpOrComplexTensor>,
BASE_HLO_NegOp;
BASE_HLO_NegOp {
let hasFolder = 1;
}
def HLO_PopulationCountOp: HLO_UnaryElementwiseOp<"popcnt",
[NoSideEffect, SameOperandsAndResultType], HLO_IntTensor>,

View File

@ -1821,6 +1821,50 @@ static LogicalResult Verify(CaseOp op) {
return success();
}
//===----------------------------------------------------------------------===//
// UnaryOps
//===----------------------------------------------------------------------===//
template <typename Op, typename ElementType = Type, typename ValType,
typename Convert>
static Attribute UnaryFolder(Op* op, ArrayRef<Attribute> attrs) {
if (!attrs[0]) return {};
DenseElementsAttr val = attrs[0].dyn_cast<DenseElementsAttr>();
if (!val) return {};
ShapedType type = op->getType().template cast<ShapedType>();
if (!type.hasStaticShape()) {
return {};
}
Type etype = type.getElementType();
// Evaluate for integer values.
if (!etype.isa<ElementType>()) {
return {};
}
SmallVector<ValType, 6> values;
values.reserve(val.getNumElements());
for (const auto v : val.getValues<ValType>()) {
values.push_back(Convert()(v));
}
return DenseElementsAttr::get(type, values);
}
#define UNARY_FOLDER(Op, Func) \
OpFoldResult Op::fold(ArrayRef<Attribute> attrs) { \
if (getElementTypeOrSelf(getType()).isa<FloatType>()) \
return UnaryFolder<Op, FloatType, APFloat, Func<APFloat>>(this, attrs); \
if (getElementTypeOrSelf(getType()).isa<IntegerType>()) \
return UnaryFolder<Op, IntegerType, APInt, Func<APInt>>(this, attrs); \
return {}; \
}
UNARY_FOLDER(NegOp, std::negate);
//===----------------------------------------------------------------------===//
// BinaryOps
//===----------------------------------------------------------------------===//

View File

@ -834,3 +834,20 @@ func @fold_xor_zeros_constants() -> tensor<4xi32> {
// CHECK: return %0
return %2 : tensor<4xi32>
}
// CHECK-LABEL: func @fold_negate_int
func @fold_negate_int() -> tensor<4xi32> {
%0 = mhlo.constant dense<[0, 1, 6, -3]> : tensor<4xi32>
// CHECK: mhlo.constant dense<[0, -1, -6, 3]>
%1 = "mhlo.negate"(%0) : (tensor<4xi32>) -> tensor<4xi32>
return %1 : tensor<4xi32>
}
// CHECK-LABEL: func @fold_negate_float
func @fold_negate_float() -> tensor<4xf32> {
%0 = mhlo.constant dense<[0., 1., 6., -3.]> : tensor<4xf32>
// CHECK: mhlo.constant dense<[-0.000000e+00, -1.000000e+00, -6.000000e+00, 3.000000e+00]>
%1 = "mhlo.negate"(%0) : (tensor<4xf32>) -> tensor<4xf32>
return %1 : tensor<4xf32>
}