1 /* Copyright 2019 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 "llvm/Support/Debug.h"
17 #include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project
18 #include "mlir/IR/Builders.h" // from @llvm-project
19 #include "mlir/IR/Operation.h" // from @llvm-project
20 #include "mlir/Pass/Pass.h" // from @llvm-project
21 #include "mlir/Pass/PassRegistry.h" // from @llvm-project
22 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
23
24 #define DEBUG_TYPE "tf-functional-to-executor"
25
26 namespace mlir {
27
28 namespace {
29 // This pass converts mlir functions consisting of mlir ops into a tf_executor
30 // dialect as a single island.
31 // Result like so:
32 // func @my_fn(%argi...) -> (result_t) {
33 // %results:[[n_args]] = tf_executor.graph {
34 // %island_results:[[nargs + 1]] = tf_executor.island {
35 // ... original ops ...
36 // tf_executor.yield %results...
37 // }
38 // tf_executor.fetch %island_results#...
39 // }
40 // return %graph_results#...
41 // }
42 struct FunctionalToExecutorDialectConversion
43 : public PassWrapper<FunctionalToExecutorDialectConversion, FunctionPass> {
getArgumentmlir::__anon373a277d0111::FunctionalToExecutorDialectConversion44 StringRef getArgument() const final {
45 // This is the argument used to refer to the pass in
46 // the textual format (on the commandline for example).
47 return "tf-functional-to-executor-conversion";
48 }
getDescriptionmlir::__anon373a277d0111::FunctionalToExecutorDialectConversion49 StringRef getDescription() const final {
50 // This is a brief description of the pass.
51 return "Transform from func op to TF executor dialect.";
52 }
53 void runOnFunction() override;
54
getDependentDialectsmlir::__anon373a277d0111::FunctionalToExecutorDialectConversion55 void getDependentDialects(DialectRegistry& registry) const override {
56 registry.insert<mlir::tf_executor::TensorFlowExecutorDialect>();
57 }
58 };
59 } // end anonymous namespace
60
runOnFunction()61 void FunctionalToExecutorDialectConversion::runOnFunction() {
62 if (!llvm::hasSingleElement(getFunction())) {
63 LLVM_DEBUG(llvm::dbgs() << "Expect single block function, skip conversion "
64 "to tf_executor dialect\n");
65 return;
66 }
67 auto loc = getFunction().getLoc();
68 mlir::Block& body = getFunction().front();
69 // Find region of interest and ReturnOp.
70 auto copy_range = body.without_terminator();
71 if (copy_range.begin() != copy_range.end() &&
72 std::next(copy_range.begin()) == copy_range.end() &&
73 isa<tf_executor::GraphOp>(*copy_range.begin())) {
74 // Already a graph.
75 return;
76 }
77
78 auto return_op = dyn_cast<ReturnOp>(body.getTerminator());
79 if (!return_op) {
80 LLVM_DEBUG(llvm::dbgs() << "Expect function to end with return\n");
81 return;
82 }
83 // Build GraphOp.
84 OpBuilder builder(&body, body.begin());
85 auto graph_op = builder.create<tf_executor::GraphOp>(
86 loc, getFunction().getType().getResults());
87 graph_op.body().push_back(new Block);
88 builder.setInsertionPointToEnd(&graph_op.GetBody());
89 auto island = builder.create<tf_executor::IslandOp>(
90 loc, getFunction().getType().getResults(),
91 tf_executor::ControlType::get(&getContext()), ArrayRef<Value>());
92 // Create Fetch.
93 ValueRange to_fetch = island.getResults();
94 if (to_fetch.size() != 1) {
95 // Drop control result for fetch.
96 to_fetch = to_fetch.drop_back();
97 }
98 builder.create<tf_executor::FetchOp>(loc, to_fetch);
99 // Build Island.
100 island.body().push_back(new Block);
101 island.body().front().getOperations().splice(
102 island.body().front().begin(), body.getOperations(), copy_range.begin(),
103 copy_range.end());
104 builder.setInsertionPointToEnd(&island.body().front());
105 builder.create<tf_executor::YieldOp>(loc, return_op.getOperands());
106 for (auto item : llvm::enumerate(graph_op.getResults())) {
107 return_op.setOperand(item.index(), item.value());
108 }
109 }
110
111 std::unique_ptr<OperationPass<FuncOp>>
CreateFunctionalToExecutorDialectConversionPass()112 CreateFunctionalToExecutorDialectConversionPass() {
113 return std::make_unique<FunctionalToExecutorDialectConversion>();
114 }
115
116 } // namespace mlir
117
118 static mlir::PassRegistration<mlir::FunctionalToExecutorDialectConversion> pass;
119