experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA][MLIR] Add complex unary elementwise op definitions to LHLO dialect.

Browse Source 
PiperOrigin-RevId: 307828261
Change-Id: I62c10002cd79209449a2c7b00122b59df829678a
This commit is contained in:
A. Unique TensorFlower 2020-04-22 09:04:41 -07:00 committed by TensorFlower Gardener
parent 126f665f39
commit e4c6f288d4
1 changed files with 23 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
tensorflow/compiler/mlir/xla/ir/lhlo_ops.td
View File

@ -37,13 +37,12 @@ def LHLO_IntBuffer : MemRefOf<[HLO_Int]>;
// Any floating-point tensor types // Any floating-point tensor types
def LHLO_FpBuffer : MemRefOf<[AnyFloat]>; def LHLO_FpBuffer : MemRefOf<[AnyFloat]>;
def LHLO_PredBuffer : MemRefOf<[HLO_Pred]>; def LHLO_PredBuffer : MemRefOf<[HLO_Pred]>;
// Any integer or floating-point tensor types // Any integer or floating-point tensor types
def LHLO_IntOrFpBuffer : MemRefOf<[HLO_Int, AnyFloat]>; 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]>; 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; 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<LHLO_Buffer, "", [MemRead]>:$lhs,
Arg<LHLO_Buffer, "", [MemRead]>:$rhs,
Arg<LHLO_Buffer, "", [MemWrite]>:$output);
}
def LHLO_ImagOp: LHLO_Op<"imag", [SameOperandsShape]>, BASE_HLO_ImagOp {
let arguments = (ins Arg<LHLO_Buffer, "", [MemRead]>:$input,
Arg<LHLO_Buffer, "", [MemWrite]>:$output);
}
def LHLO_RealOp: LHLO_Op<"real", [SameOperandsShape]>, BASE_HLO_RealOp {
let arguments = (ins Arg<LHLO_Buffer, "", [MemRead]>:$input,
Arg<LHLO_Buffer, "", [MemWrite]>:$output);
}
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// XLA binary elementwise op definitions. // XLA binary elementwise op definitions.
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// See https://www.tensorflow.org/xla/operation_semantics#element-wise_binary_arithmetic_operations
class LHLO_BinaryElementwiseOp<string mnemonic, list<OpTrait> traits> : class LHLO_BinaryElementwiseOp<string mnemonic, list<OpTrait> traits> :
LHLO_Op<mnemonic, traits> { LHLO_Op<mnemonic, traits> {