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 "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
17 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
18 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
19 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
20 #include "tensorflow/compiler/mlir/tensorflow/transforms/tf_device_passes_detail.h"
21 
22 namespace mlir {
23 
24 namespace {
25 // Outlines partitioned call ops with `_XlaMustCompile` to device clusters.
26 struct XlaClusterFormationPass
27     : public TFDevice::XlaClusterFormationPassBase<XlaClusterFormationPass> {
28   void runOnOperation() override;
29 };
30 
EncapsulatePartitionedCall(TF::StatefulPartitionedCallOp call_op)31 void EncapsulatePartitionedCall(TF::StatefulPartitionedCallOp call_op) {
32   mlir::OpBuilder builder(call_op);
33 
34   auto cluster = builder.create<mlir::tf_device::ClusterOp>(
35       call_op.getLoc(), call_op.getResultTypes());
36 
37   call_op.replaceAllUsesWith(cluster.getResults());
38 
39   cluster.body().push_back(new mlir::Block);
40 
41   call_op.getOperation()->moveBefore(&cluster.GetBody(),
42                                      cluster.GetBody().end());
43 
44   builder.setInsertionPointToEnd(&cluster.GetBody());
45   builder.create<mlir::tf_device::ReturnOp>(call_op.getLoc(),
46                                             call_op->getResults());
47 }
48 
runOnOperation()49 void XlaClusterFormationPass::runOnOperation() {
50   ModuleOp module = getOperation();
51 
52   llvm::SmallVector<TF::StatefulPartitionedCallOp, 4> ops;
53   module.walk([&](TF::StatefulPartitionedCallOp call_op) {
54     auto attr = call_op->getAttrOfType<BoolAttr>("_XlaMustCompile");
55     if (attr && attr.getValue()) {
56       ops.push_back(call_op);
57     }
58   });
59 
60   for (auto call_op : ops) {
61     EncapsulatePartitionedCall(call_op);
62   }
63 }
64 
65 }  // namespace
66 
67 namespace TFDevice {
CreateXlaClusterFormationPass()68 std::unique_ptr<OperationPass<ModuleOp>> CreateXlaClusterFormationPass() {
69   return std::make_unique<XlaClusterFormationPass>();
70 }
71 }  // namespace TFDevice
72 
73 }  // namespace mlir
74