[MLIR][KernelGen] Add asin kernels and tests
PiperOrigin-RevId: 351381423
This commit is contained in:
parent
b0e0ca830c
commit
0b85d5c510
|
@ -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";
|
||||
|
|
|
@ -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),
|
||||
|
|
|
@ -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) {
|
||||
|
|
Loading…
Reference in New Issue