1 /* Copyright 2020 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/ADT/SmallVector.h"
17 #include "mlir/IR/Attributes.h" // from @llvm-project
18 #include "mlir/IR/BuiltinTypes.h" // from @llvm-project
19 #include "mlir/Pass/Pass.h" // from @llvm-project
20 #include "mlir/Pass/PassRegistry.h" // from @llvm-project
21 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
22
23 namespace mlir {
24 namespace TF {
25
26 namespace {
27
28 // Rewrites RecvTPUEmbeddingActivationsOp and SendTPUEmbeddingGradients ops to
29 // internal variants by introducing _RecvTPUEmbeddingDeduplicationData op.
30 struct RewriteTPUEmbeddingOps
31 : public PassWrapper<RewriteTPUEmbeddingOps, FunctionPass> {
32 void runOnFunction() override;
33
getArgumentmlir::TF::__anone93152ab0111::RewriteTPUEmbeddingOps34 StringRef getArgument() const final { return "tf-rewrite-tpu-embedding-ops"; }
35
getDescriptionmlir::TF::__anone93152ab0111::RewriteTPUEmbeddingOps36 StringRef getDescription() const final {
37 return "Rewrites TPU embedding send/recv ops by adding TPU embedding "
38 "deduplication data";
39 }
40 };
41
42 // Rewrites the given op to `OpT` op after adding the given operand at the end.
43 template <typename OpT>
AddOperandAndRewriteAs(Operation * op,Value operand,OpBuilder * builder)44 OpT AddOperandAndRewriteAs(Operation* op, Value operand, OpBuilder* builder) {
45 builder->setInsertionPoint(op);
46 auto operands = llvm::to_vector<4>(op->getOperands());
47 operands.push_back(operand);
48 auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
49 operands, op->getAttrs());
50 op->replaceAllUsesWith(new_op.getOperation()->getResults());
51 op->erase();
52 return new_op;
53 }
54
55 // Returns success if the function has at most one op of the template type and
56 // assigns it to `result`, if present. If there are multiple such ops, returns
57 // failure.
58 template <typename OpT>
GetOp(Region * region,OpT * result)59 LogicalResult GetOp(Region* region, OpT* result) {
60 *result = {};
61 for (auto op : region->getOps<OpT>()) {
62 if (*result) return op.emitError("should be unique within a function");
63 *result = op;
64 }
65 return success();
66 }
67
RunOnRegion(Region * region)68 LogicalResult RunOnRegion(Region* region) {
69 RecvTPUEmbeddingActivationsOp recv_op;
70 if (failed(GetOp(region, &recv_op))) return failure();
71
72 SendTPUEmbeddingGradientsOp send_op;
73 if (failed(GetOp(region, &send_op))) return failure();
74
75 // No TPU embedding ops.
76 if (!recv_op && !send_op) return success();
77
78 Location loc = recv_op ? recv_op.getLoc() : send_op.getLoc();
79 StringRef config = recv_op ? recv_op.config() : send_op.config();
80
81 // Create _RecvTPUEmbeddingDeduplicationData op.
82 OpBuilder builder(region);
83 auto output_ty =
84 RankedTensorType::get({}, VariantType::get(region->getContext()));
85 auto dedup_op = builder.create<_RecvTPUEmbeddingDeduplicationDataOp>(
86 loc, output_ty, config);
87
88 // Rewrite RecvTPUEmbeddingActivations op to the corresponding internal op.
89 if (recv_op)
90 AddOperandAndRewriteAs<_RecvTPUEmbeddingActivationsOp>(recv_op, dedup_op,
91 &builder);
92
93 // Rewrite SendTPUEmbeddingGradients op to the corresponding internal op and
94 // then update the OperandSegmentSize attribute.
95 if (send_op) {
96 int32_t operand_sizes[] = {static_cast<int32_t>(send_op.N()),
97 static_cast<int32_t>(send_op.NN()), 1};
98 auto attr_ty = VectorType::get(3, builder.getI32Type());
99 auto operand_size_attr = DenseIntElementsAttr::get(attr_ty, operand_sizes);
100
101 auto new_send_op = AddOperandAndRewriteAs<_SendTPUEmbeddingGradientsOp>(
102 send_op, dedup_op, &builder);
103 new_send_op->setAttr(new_send_op.getOperandSegmentSizeAttr(),
104 operand_size_attr);
105 }
106 return success();
107 }
108
runOnFunction()109 void RewriteTPUEmbeddingOps::runOnFunction() {
110 FuncOp func = getFunction();
111 if (failed(RunOnRegion(&func.getBody()))) return signalPassFailure();
112
113 func.walk([&](Operation* op) {
114 for (Region& region : op->getRegions()) {
115 if (failed(RunOnRegion(®ion))) return signalPassFailure();
116 }
117 });
118 }
119
120 } // anonymous namespace
121
CreateRewriteTPUEmbeddingOpsPass()122 std::unique_ptr<OperationPass<FuncOp>> CreateRewriteTPUEmbeddingOpsPass() {
123 return std::make_unique<RewriteTPUEmbeddingOps>();
124 }
125
126 static PassRegistration<RewriteTPUEmbeddingOps> pass;
127
128 } // namespace TF
129 } // namespace mlir
130