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 | #include "tensorflow/core/ir/interfaces.h" |
17 | |
18 | #include "llvm/ADT/SmallVector.h" |
19 | #include "mlir/IR/Operation.h" // from @llvm-project |
20 | #include "mlir/IR/Region.h" // from @llvm-project |
21 | #include "mlir/IR/Value.h" // from @llvm-project |
22 | #include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project |
23 | #include "tensorflow/core/ir/ops.h" |
24 | #include "tensorflow/core/ir/types/dialect.h" |
25 | |
26 | namespace mlir { |
27 | namespace tfg { |
28 | |
29 | LogicalResult ControlArgumentInterface::verifyRegion(Operation *op, |
30 | Region ®ion) { |
31 | unsigned num_ctl = 0, num_data = 0; |
32 | for (BlockArgument arg : region.getArguments()) { |
33 | bool is_ctl = arg.getType().isa<tf_type::ControlType>(); |
34 | num_ctl += is_ctl; |
35 | num_data += !is_ctl; |
36 | } |
37 | if (num_ctl != num_data) { |
38 | return op->emitOpError("region #") |
39 | << region.getRegionNumber() |
40 | << " expected same number of data values and control tokens (" |
41 | << num_data << " vs. "<< num_ctl << ")"; |
42 | } |
43 | return success(); |
44 | } |
45 | |
46 | void StatefulMemoryEffectInterface::getEffects( |
47 | Operation *op, |
48 | SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> |
49 | &effects) const { |
50 | auto registry = dyn_cast<TensorFlowRegistryInterface>(op); |
51 | // If the registry interface is not available, conservatively assume stateful. |
52 | // Otherwise, add a write effect if the operation is known to be stateful. |
53 | // FIXME: Prevent ops in GraphOp being pruned. Remove this when GraphToFunc |
54 | // and FuncToGraph land. |
55 | if (!registry || registry.isStateful() || op->getParentOfType<GraphOp>()) { |
56 | effects.emplace_back(MemoryEffects::Write::get()); |
57 | } |
58 | } |
59 | |
60 | } // namespace tfg |
61 | } // namespace mlir |
62 | |
63 | // Include the generated definitions. |
64 | #include "tensorflow/core/ir/interfaces.cc.inc" |
65 |