[MLIR][KernelGen] Add asin kernels and tests

PiperOrigin-RevId: 351381423
This commit is contained in:
A. Unique TensorFlower 2021-01-12 09:01:28 -08:00 committed by TensorFlow MLIR Team
parent b0e0ca830c
commit 0b85d5c510
3 changed files with 34 additions and 2 deletions

View File

@ -359,6 +359,19 @@ def HLOClient_AcosOp : HLOClient_UnaryElementwiseOp<"acos", [],
}];
}
def HLOClient_AsinOp : HLOClient_UnaryElementwiseOp<"asin", [],
HLO_FpOrComplexTensor> {
let summary = "Asin operator";
let description = [{
Returns `Asin(operand)` element-wise.
$$
\asin(x) = 2 * atan(x / (1 + sqrt(1 - x^2)))
$$
}];
}
def HLOClient_AtanOp : HLOClient_UnaryElementwiseOp<"atan", [],
HLO_FpOrComplexTensor> {
let summary = "Atan operator";

View File

@ -60,6 +60,25 @@ def : Pat<(HLOClient_AcosOp NonComplexElementType:$input),
(HLO_ConstantLike<"M_PI"> $input)
)>;
// Expand asin to MHLO dialect as follows:
// asin(x) = 2 * atan(x / (1 + sqrt(1 - x^2)))
def : Pat<(HLOClient_AsinOp $input),
(HLO_MulOp
(HLO_ConstantLike<"2"> $input),
(HLO_Atan2Op
$input,
(HLO_AddOp
(HLO_ConstantLike<"1"> $input),
(HLO_SqrtOp
(HLO_SubOp
(HLO_ConstantLike<"1"> $input),
(HLO_MulOp $input, $input)
)
)
)
)
)>;
// Express `atan` as
// atan(x) = atan2(x, 1)
def : Pat<(HLOClient_AtanOp $input),

View File

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