1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_IR_INTERFACES_H_ 17 #define TENSORFLOW_CORE_IR_INTERFACES_H_ 18 19 #include "mlir/IR/Dialect.h" // from @llvm-project 20 #include "mlir/IR/DialectInterface.h" // from @llvm-project 21 #include "mlir/IR/OpDefinition.h" // from @llvm-project 22 #include "mlir/Interfaces/ControlFlowInterfaces.h" // from @llvm-project 23 #include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project 24 #include "mlir/Support/LLVM.h" // from @llvm-project 25 #include "tensorflow/core/ir/dialect.h" 26 27 // Include generated declarations. 28 #include "tensorflow/core/ir/interfaces.h.inc" 29 30 namespace mlir { 31 namespace tfg { 32 // The dialect fallback model for the TensorFlow registry interface. 33 class TensorFlowRegistryInterfaceBase 34 : public TensorFlowRegistryInterface::FallbackModel< 35 TensorFlowRegistryInterfaceBase>, 36 public DialectInterface::Base<TensorFlowRegistryInterfaceBase> { 37 public: TensorFlowRegistryInterfaceBase(Dialect * dialect)38 explicit TensorFlowRegistryInterfaceBase(Dialect *dialect) 39 : DialectInterface::Base<TensorFlowRegistryInterfaceBase>(dialect) {} 40 41 // Returns whether the operation is stateful. 42 virtual bool isStateful(Operation *op) const = 0; 43 }; 44 45 // This dialect fallback model implements memory effects for TensorFlow 46 // operations. 47 class StatefulMemoryEffectInterface 48 : public MemoryEffectOpInterface::FallbackModel< 49 StatefulMemoryEffectInterface>, 50 public DialectInterface::Base<StatefulMemoryEffectInterface> { 51 public: StatefulMemoryEffectInterface(Dialect * dialect)52 explicit StatefulMemoryEffectInterface(Dialect *dialect) 53 : DialectInterface::Base<StatefulMemoryEffectInterface>(dialect) {} 54 55 // Get the memory effects of a TensorFlow operation. If the operation is known 56 // to be stateless, then it has no memory effects. Otherwise, statefulness is 57 // modelled as `MemoryWrite`. 58 void getEffects( 59 Operation *op, 60 SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> 61 &effects) const; 62 }; 63 } // namespace tfg 64 65 namespace OpTrait { 66 // This trait marks intrinsic TFG operations, e.g. terminators, functions, 67 // and region control-flow operations. Any TFG operation that has this trait 68 // exists only in MLIR. 69 template <typename ConcreteType> 70 class IntrinsicOperation 71 : public mlir::OpTrait::TraitBase<ConcreteType, IntrinsicOperation> {}; 72 } // namespace OpTrait 73 } // namespace mlir 74 75 #endif // TENSORFLOW_CORE_IR_INTERFACES_H_ 76