colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add folder for HLO NotOp 

PiperOrigin-RevId: 364989658
This commit is contained in:
Geoffrey Martin-Noble 2021-03-25 02:07:29 -07:00 committed by TensorFlow MLIR Team
parent cbef26c6a8
commit a2b6060c0c
3 changed files with 35 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
View File

@ -217,6 +217,7 @@ def HLO_LogisticOp: HLO_UnaryElementwiseOp<"logistic",
def HLO_NotOp: HLO_UnaryElementwiseOp<"not",
[NoSideEffect, SameOperandsAndResultType], HLO_PredOrIntTensor>,
BASE_HLO_NotOp {
let hasFolder = 1;
}
def HLO_NegOp: HLO_UnaryElementwiseOp<"negate",

18
lib/Dialect/mhlo/IR/hlo_ops.cc
View File

@ -2195,6 +2195,12 @@ struct round {
}
};
struct logical_not {
APInt operator()(const APInt& i) {
return APInt(i.getBitWidth(), static_cast<uint64_t>(!i));
}
};
#define UNARY_FOLDER(Op, Func) \
OpFoldResult Op::fold(ArrayRef<Attribute> attrs) { \
if (getElementTypeOrSelf(getType()).isa<FloatType>()) \
@ -2204,6 +2210,13 @@ struct round {
return {}; \
}
#define UNARY_FOLDER_INT(Op, Func) \
OpFoldResult Op::fold(ArrayRef<Attribute> attrs) { \
if (getElementTypeOrSelf(getType()).isa<IntegerType>()) \
return UnaryFolder<Op, IntegerType, APInt, Func>(this, attrs); \
return {}; \
}
#define UNARY_FOLDER_FLOAT(Op, Func) \
OpFoldResult Op::fold(ArrayRef<Attribute> attrs) { \
if (getElementTypeOrSelf(getType()).isa<FloatType>()) \
@ -2212,8 +2225,13 @@ struct round {
}
UNARY_FOLDER(NegOp, std::negate);
UNARY_FOLDER_INT(NotOp, logical_not);
UNARY_FOLDER_FLOAT(RoundOp, round);
#undef UNARY_FOLDER
#undef UNARY_FOLDER_INT
#undef UNARY_FOLDER_FLOAT
//===----------------------------------------------------------------------===//
// BinaryOps
//===----------------------------------------------------------------------===//

16
tests/canonicalize.mlir
View File

@ -1237,6 +1237,22 @@ func @fold_negate_float() -> tensor<4xf32> {
return %1 : tensor<4xf32>
}
// CHECK-LABEL func @fold_not()
func @fold_not() -> tensor<2x2xi1> {
%0 = mhlo.constant dense<[[true, false], [true, false]]> : tensor<2x2xi1>
// CHECK{LITERAL}: mhlo.constant dense<[[false, true], [false, true]]> : tensor<2x2xi1>
%1 = "mhlo.not"(%0) : (tensor<2x2xi1>) -> tensor<2x2xi1>
return %1 : tensor<2x2xi1>
}
// CHECK-LABEL func @fold_not_i32()
func @fold_not_i32() -> tensor<2x2xi32> {
%0 = mhlo.constant dense<[[42, -12], [1, 0]]> : tensor<2x2xi32>
// CHECK-LITERAL: mhlo.constant dense<[[0, 0], [0, 1]]> : tensor<2x2xi32>
%1 = "mhlo.not"(%0) : (tensor<2x2xi32>) -> tensor<2x2xi32>
return %1 : tensor<2x2xi32>
}
// CHECK-LABEL: func @fold_sqrt_f32_constants
func @fold_sqrt_f32_constants() -> tensor<4xf32> {
%0 = mhlo.constant dense<1.0> : tensor<4xf32>