1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3Licensed under the Apache License, Version 2.0 (the "License"); 4you may not use this file except in compliance with the License. 5You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9Unless required by applicable law or agreed to in writing, software 10distributed under the License is distributed on an "AS IS" BASIS, 11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12See the License for the specific language governing permissions and 13limitations under the License. 14==============================================================================*/ 15 16// This is the operation interface definition file for TensorFlow. 17 18#ifndef TF_OP_INTERFACES 19#define TF_OP_INTERFACES 20 21include "mlir/IR/OpBase.td" 22 23//===----------------------------------------------------------------------===// 24// TensorFlow Layout Optimization Interfaces. 25//===----------------------------------------------------------------------===// 26 27def TF_LayoutSensitiveInterface : OpInterface<"LayoutSensitiveInterface"> { 28 let description = [{ 29 A layout sensitive operation is one that depends on a `data_format` string 30 attribute, that gives a meaning to the data inside arguments and results. 31 32 Currently supported data formats (layouts): 33 - NHWC : channels last [batch, height, width, channels] 34 - NCHW : channels first [batch, channels, height, width] 35 36 Layout sensitive ops might have different preferred (and supported) layouts 37 depending on arguments shape/type and execution device (CPU or GPU). 38 }]; 39 40 let methods = [ 41 InterfaceMethod< 42 [{Returns current operation data format (data layout).}], 43 "StringRef", "data_format", (ins) 44 >, 45 InterfaceMethod< 46 [{Returns indices of layout dependent arguments.}], 47 "SmallVector<unsigned, 4>", "GetLayoutDependentArgs", (ins) 48 >, 49 InterfaceMethod< 50 [{Returns indices of layout dependent results.}], 51 "SmallVector<unsigned, 4>", "GetLayoutDependentResults", (ins) 52 >, 53 InterfaceMethod< 54 [{Returns the optimal data layout based on the available devices.}], 55 "StringRef", "GetOptimalLayout", (ins "const RuntimeDevices&":$devices) 56 >, 57 InterfaceMethod< 58 [{Updates operation attributes and operands to account for the updated 59 data format. If data format is not supported, must return failure.}], 60 "LogicalResult", "UpdateDataFormat", (ins "StringRef":$data_format) 61 >, 62 ]; 63 64 let verify = [{ 65 return VerifyLayoutSensitiveInterface($_op); 66 }]; 67} 68 69def TF_FoldOperandsTransposeInterface : OpInterface<"FoldOperandsTransposeInterface"> { 70 let description = [{ 71 Operation supports folding operand(s) transposes into the operation itself. 72 73 (1) Operation might have layout dependent operands and results... 74 75 Example: MaxPool(Transpose($arg, $perm)) 76 -> Transpose(MaxPool($arg, $perm)) 77 78 (2) ... or it might have only layout dependent operands: 79 80 Example: Mean(Transpose($arg, $reduction_dims)) 81 -> Mean($arg, Transpose($reduction_dims)) 82 }]; 83 84 let methods = [ 85 InterfaceMethod< 86 [{Returns indices of layout dependent arguments.}], 87 "SmallVector<unsigned, 4>", "GetLayoutDependentArgs", (ins) 88 >, 89 InterfaceMethod< 90 [{Returns indices of layout dependent results.}], 91 "SmallVector<unsigned, 4>", "GetLayoutDependentResults", (ins) 92 >, 93 InterfaceMethod< 94 [{Updates operation attributes and operands to account for the folded 95 permutation. If folding of permutation is not possible, must return 96 failure.}], 97 "LogicalResult", "FoldOperandsPermutation", 98 (ins "ArrayRef<int64_t>":$permutation) 99 >, 100 ]; 101 102 let verify = [{ 103 return VerifyFoldOperandsTransposeInterface($_op); 104 }]; 105} 106 107//===----------------------------------------------------------------------===// 108// TensorFlow Contraction Fusion Interfaces. 109//===----------------------------------------------------------------------===// 110 111def TF_ContractionFusableInterface : OpInterface<"ContractionFusableInterface"> { 112 let description = [{ 113 A contraction fusable operation is one that can be fused into the output of 114 a tensor contraction (MatMul, Conv2D, etc...) operation. 115 116 For example all element wise operations are trivially contraction fusable. 117 }]; 118 119 let methods = [ 120 InterfaceMethod< 121 [{Returns contraction fusion if the operation satisfies all the fusion 122 requirements. Otherwise returns empty optional.}], 123 "Optional<ContractionFusion>", "GetContractionFusion", (ins) 124 >, 125 ]; 126} 127 128//===----------------------------------------------------------------------===// 129// TensorFlow Resource Handle Interfaces. 130//===----------------------------------------------------------------------===// 131 132def TF_ResourceHandleAllocatorInterface : OpInterface<"ResourceHandleAllocatorInterface"> { 133 let description = [{ 134 A resource handle allocator operation is one that creates a resource handle, 135 or looks up and reuses an existing resource handle. 136 }]; 137 138 let methods = [ 139 InterfaceMethod< 140 /*desc=*/[{Returns the resource handle value and unique id associated with 141 the resource handle. If a resource handle is reused, then an 142 existing id will be returned.}], 143 /*retTy=*/"ResourceHandleValueAndId", 144 /*methodName=*/"GetResourceHandleValueAndId", 145 /*args=*/(ins "llvm::SmallDenseMap<ResourceHandle, int64_t>&":$resource_handle_id_map, 146 "int64_t&":$next_id) 147 >, 148 ]; 149} 150 151#endif // TF_OP_INTERFACES 152