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

[MLIR][KernelGen] Add `tf.Asinh` kernels and complete their lowerings 

PiperOrigin-RevId: 352604725
This commit is contained in:
A. Unique TensorFlower 2021-01-19 10:50:33 -08:00 committed by TensorFlow MLIR Team
parent 96fb617413
commit 0e85b4d511
6 changed files with 3 additions and 123 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h
View File

@ -66,8 +66,6 @@ static Value getConstantLike(OpBuilder& b, Location loc, T constant,
return b.create<ConstantLikeOp>(loc, getAttr(), val);
}
Value getConstantLikeMaxFiniteValue(OpBuilder& b, Location loc, Value val);
} // namespace chlo
} // namespace mlir

12
include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
View File

@ -372,18 +372,6 @@ def HLOClient_AsinOp : HLOClient_UnaryElementwiseOp<"asin", [],
}];
}
def HLOClient_AsinhOp : HLOClient_UnaryElementwiseOp<"asinh", [],
HLO_FpOrComplexTensor> {
let summary = "Asinh operation";
let description = [{
Returns `Asinh(operand)` element-wise.
$$
\asinh(x) = log(x + sqrt(x^2 + 1))
$$
}];
}
def HLOClient_AtanOp : HLOClient_UnaryElementwiseOp<"atan", [],
HLO_FpOrComplexTensor> {
let summary = "Atan operator";

3
include/mlir-hlo/Dialect/mhlo/IR/hlo_utils.td
View File

@ -30,9 +30,6 @@ class ConstantSplat<string value> : NativeCodeCall<
class HLO_ConstantLike<string value> : NativeCodeCall<
"chlo::getConstantLike($_builder, $_loc, " # value # ", $0)">;
def HLO_ConstantLikeMaxFiniteValue : NativeCodeCall<
"chlo::getConstantLikeMaxFiniteValue($_builder, $_loc, $0)">;
def NullDenseIntElementsAttr : NativeCodeCall<"DenseIntElementsAttr()">;
def BinBroadcastDimensions : NativeCodeCall<

14
lib/Dialect/mhlo/IR/chlo_ops.cc
View File

@ -32,20 +32,6 @@ static LogicalResult Verify(T op) {
return success();
}
static constexpr float kF16MaxFiniteValue = 0x1.ffcP15;
Value getConstantLikeMaxFiniteValue(OpBuilder& b, Location loc, Value val) {
Type ty = getElementTypeOrSelf(val.getType());
if (ty.isF16()) {
return getConstantLike(b, loc, kF16MaxFiniteValue, val);
} else if (ty.isF32()) {
return getConstantLike(b, loc, std::numeric_limits<float>::max(), val);
} else if (ty.isF64()) {
return getConstantLike(b, loc, std::numeric_limits<double>::max(), val);
}
llvm_unreachable("unhandled type");
}
//===----------------------------------------------------------------------===//
// BinaryOps
//===----------------------------------------------------------------------===//

88
lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td
View File

@ -79,94 +79,6 @@ def : Pat<(HLOClient_AsinOp NonComplexElementType:$input),
)
)>;
// Expand asinh to MHLO dialect as
// asinh(x) = log(x + sqrt(x^2 + 1))
//
// If x^2 will overflow and x is positive, we can approximate x + sqrt(x^2 + 1)
// as 2*x and return log(2) + log(x).
//
// For small x, sqrt(x^2 + 1) will evaluate to 1 due to floating point
// arithmetic. However, we would like to retain the low order term of this,
// which is around 0.5 * x^2 using a binomial expansion.
// Let z = sqrt(a^2 + 1)
// The following rewrite retains the lower order term.
// log(a + sqrt(a^2 + 1))
// = log((a + sqrt(a^2 + 1)) * (1 + sqrt(a^2 + 1)) / (1 + sqrt(a^2 + 1)))
// = log((a + a^2 + 1 + a * z + z) / (1 + z))
// = log(1 + a + a^2 / (1 + z))
// = log(1 + a + a^2 / (1 + sqrt(a^2 + 1)))
//
// If x is negative, the above would give us some trouble; we can't approximate
// the result as x + abs(x) = 0 but we are saved by the fact that asinh(-x) =
// -asinh(x).
def : Pat<(HLOClient_AsinhOp NonComplexElementType:$input),
(HLO_MulOp
(HLO_SignOp $input),
(HLO_SelectOp
(HLO_CompareOp
(HLO_AbsOp $input),
(HLO_SqrtOp
(HLO_ConstantLikeMaxFiniteValue $input)
),
HLO_COMPARISON_DIRECTION_GE,
(HLO_DEFAULT_COMPARISON_TYPE)
),
(HLO_AddOp
(HLO_LogOp
(HLO_AbsOp $input)
),
(HLO_LogOp
(HLO_ConstantLike<"2"> $input)
)
),
(HLO_SelectOp
(HLO_CompareOp
(HLO_AbsOp $input),
(HLO_ConstantLike<"1"> $input),
HLO_COMPARISON_DIRECTION_LE,
(HLO_DEFAULT_COMPARISON_TYPE)
),
(HLO_Log1pOp
(HLO_AddOp
(HLO_AbsOp $input),
(HLO_MulOp
(HLO_AbsOp $input),
(HLO_DivOp
(HLO_AbsOp $input),
(HLO_AddOp
(HLO_ConstantLike<"1"> $input),
(HLO_SqrtOp
(HLO_AddOp
(HLO_MulOp
(HLO_AbsOp $input),
(HLO_AbsOp $input)
),
(HLO_ConstantLike<"1"> $input)
)
)
)
)
)
)
),
(HLO_LogOp
(HLO_AddOp
(HLO_AbsOp $input),
(HLO_SqrtOp
(HLO_AddOp
(HLO_MulOp
(HLO_AbsOp $input),
(HLO_AbsOp $input)
),
(HLO_ConstantLike<"1"> $input)
)
)
)
)
)
)
)>;
// Express `atan` as
// atan(x) = atan2(x, 1)
def : Pat<(HLOClient_AtanOp $input),

5
lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc
View File

@ -51,9 +51,8 @@ namespace {
// TODO(herhut): Generate these out of op definitions.
#define MAP_CHLO_OPERATION_CWISE_UNARY(fn, sep) \
fn(AcosOp) sep fn(AsinOp) sep fn(AsinhOp) sep fn(AtanOp) sep fn(AtanhOp) \
sep fn(ConjOp) sep fn(CoshOp) sep fn(ErfOp) sep fn(ErfcOp) \
sep fn(SinhOp) sep fn(TanOp)
fn(AcosOp) sep fn(AsinOp) sep fn(AtanOp) sep fn(AtanhOp) sep fn(ConjOp) \
sep fn(CoshOp) sep fn(ErfOp) sep fn(ErfcOp) sep fn(SinhOp) sep fn(TanOp)
template <typename OpTy>
inline void AddLegalOpOnRankedTensor(ConversionTarget *target) {