diff --git a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td index 92ae4d2..13d5f02 100644 --- a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td +++ b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td @@ -360,6 +360,19 @@ def HLOClient_AcosOp : HLOClient_UnaryElementwiseOp<"acos", [], }]; } +def HLOClient_AtanOp : HLOClient_UnaryElementwiseOp<"atan", [], + HLO_FpOrComplexTensor> { + let summary = "Atan operator"; + + let description = [{ + Returns `Atan(operand)` element-wise. + + $$ + \atan(x) = \atan2(x, 1) + $$ + }]; +} + def HLOClient_SinhOp : HLOClient_UnaryElementwiseOp<"sinh", [], HLO_FpOrComplexTensor> { let summary = "Sinh operation"; diff --git a/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h b/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h index d2fb048..9cf1b6c 100644 --- a/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h +++ b/include/mlir-hlo/Dialect/mhlo/transforms/map_lmhlo_to_scalar_op.h @@ -149,6 +149,15 @@ inline Value MapLhloOpToStdScalarOp(Location loc, loc, result_types, args, b); } +template <> +inline Value MapLhloOpToStdScalarOp(Location loc, + ArrayRef result_types, + ArrayRef args, + OpBuilder* b) { + return MapLhloOpToStdScalarOpImpl{}( + loc, result_types, args, b); +} + template inline Optional getCmpPredicate(StringRef comparison_direction) { return llvm::None; diff --git a/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td b/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td index 9d9e6d9..fb4bf3a 100644 --- a/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td +++ b/lib/Dialect/mhlo/transforms/chlo_legalize_to_hlo_patterns.td @@ -49,6 +49,14 @@ def : Pat<(HLOClient_AcosOp $input), ), (HLO_ConstantLike<"M_PI"> $input))>; +// Express `atan` as +// atan(x) = atan2(x, 1) +def : Pat<(HLOClient_AtanOp $input), + (HLO_Atan2Op + $input, + (HLO_ConstantLike<"1"> $input) + )>; + // Express `sinh` as // sinh(x) = (e^x - e^-x) / 2 if |x| < 1 // = e^(x + log(1/2)) - e^(-x + log(1/2)) otherwise. @@ -95,4 +103,3 @@ def : Pat<(HLOClient_TanOp $input), (HLO_SinOp $input), (HLO_CosOp $input) )>; - diff --git a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc index 2d36b1c..57859b6 100644 --- a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc +++ b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc @@ -822,6 +822,7 @@ void populateLHLOToLinalgConversionPattern(MLIRContext* context, PointwiseToLinalgConverter, PointwiseToLinalgConverter, PointwiseToLinalgConverter, + PointwiseToLinalgConverter, PointwiseToLinalgConverter, PointwiseToLinalgConverter, PointwiseToLinalgConverter, @@ -932,6 +933,7 @@ void populateHLOToLinalgConversionPattern(MLIRContext* context, PointwiseToLinalgConverter, PointwiseToLinalgConverter, PointwiseToLinalgConverter, + PointwiseToLinalgConverter, PointwiseToLinalgConverter, PointwiseToLinalgConverter, PointwiseToLinalgConverter, diff --git a/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc b/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc index 0be4c68..7c01fa2 100644 --- a/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc +++ b/lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc @@ -48,7 +48,7 @@ namespace { // TODO(herhut): Generate these out of op definitions. #define MAP_CHLO_OPERATION_CWISE_UNARY(fn, sep) \ - fn(TanOp) sep fn(AcosOp) sep fn(SinhOp) + fn(AcosOp) sep fn(AtanOp) sep fn(SinhOp) sep fn(TanOp) template inline void AddLegalOpOnRankedTensor(ConversionTarget *target) {