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 Resource Handle Interfaces. 109//===----------------------------------------------------------------------===// 110 111def TF_ResourceHandleAllocatorInterface : OpInterface<"ResourceHandleAllocatorInterface"> { 112 let description = [{ 113 A resource handle allocator operation is one that creates a resource handle, 114 or looks up and reuses an existing resource handle. 115 }]; 116 117 let methods = [ 118 InterfaceMethod< 119 /*desc=*/[{Returns resource handle values and the unique ids associated with 120 the resource handles for this op. The handles should be created 121 for only the resource tensors in the results of the op. If a 122 resource handle is reused, then an existing id will be 123 returned. The order of the resource handles in the returned 124 vector are the order of those resources in the results of the 125 op.}], 126 /*retTy=*/"llvm::SmallVector<ResourceHandleValueAndId, 4>", 127 /*methodName=*/"GetResourceHandleValueAndIdList", 128 /*args=*/(ins "llvm::SmallDenseMap<ResourceHandle, int64_t>&":$resource_handle_id_map, 129 "int64_t&":$next_id) 130 >, 131 ]; 132} 133 134def TF_GetResourceInstanceInterface : OpInterface<"GetResourceInstanceInterface"> { 135 let description = [{Returns a string corresponding to the resource instance 136 accessed by this op}]; 137 138 let methods = [ 139 InterfaceMethod< 140 /*desc=*/[{Returns a string corresponding to the resource instance 141 accessed by this op. The implementation must guarantee that the 142 mapping between resource instances and strings is bijective, 143 i.e., two op instances should return the same string if and 144 only if they access the same resource. The interface should 145 only be used for ops that access exactly one resource.}], 146 /*retTy=*/"std::string", 147 /*methodName=*/"GetResourceInstanceStr", 148 /*args=*/(ins) 149 >, 150 ]; 151} 152 153#endif // TF_OP_INTERFACES 154