diff --git a/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td b/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td index 7613f1e0ffc..f0c84dbc862 100644 --- a/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td +++ b/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td @@ -37,13 +37,12 @@ def LHLO_IntBuffer : MemRefOf<[HLO_Int]>; // Any floating-point tensor types def LHLO_FpBuffer : MemRefOf<[AnyFloat]>; - def LHLO_PredBuffer : MemRefOf<[HLO_Pred]>; // Any integer or floating-point tensor types def LHLO_IntOrFpBuffer : MemRefOf<[HLO_Int, AnyFloat]>; -def LHLO_Buffer : MemRefOf<[AnyFloat, AnySignlessInteger]>; +def LHLO_Buffer : MemRefOf<[AnyFloat, AnySignlessInteger, AnyComplex]>; def LHLO_TupleBuffer : NestedTupleOf<[LHLO_Buffer]>; @@ -105,9 +104,31 @@ def LHLO_SignOp: LHLO_UnaryElementwiseOp<"sign">, BASE_HLO_SignOp; def LHLO_TanhOp: LHLO_UnaryElementwiseOp<"tanh">, BASE_HLO_TanhOp; +//===----------------------------------------------------------------------===// +// XLA complex unary elementwise op definitions. +//===----------------------------------------------------------------------===// +// See https://www.tensorflow.org/xla/operation_semantics#element-wise_unary_functions + +def LHLO_ComplexOp: LHLO_Op<"complex", [SameOperandsShape]>, BASE_HLO_ComplexOp { + let arguments = (ins Arg:$lhs, + Arg:$rhs, + Arg:$output); +} + +def LHLO_ImagOp: LHLO_Op<"imag", [SameOperandsShape]>, BASE_HLO_ImagOp { + let arguments = (ins Arg:$input, + Arg:$output); +} + +def LHLO_RealOp: LHLO_Op<"real", [SameOperandsShape]>, BASE_HLO_RealOp { + let arguments = (ins Arg:$input, + Arg:$output); +} + //===----------------------------------------------------------------------===// // XLA binary elementwise op definitions. //===----------------------------------------------------------------------===// +// See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations class LHLO_BinaryElementwiseOp traits> : LHLO_Op {