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 };
34
35 // Rewrites the given op to `OpT` op after adding the given operand at the end.
36 template <typename OpT>
AddOperandAndRewriteAs(Operation * op,Value operand,OpBuilder * builder)37 OpT AddOperandAndRewriteAs(Operation* op, Value operand, OpBuilder* builder) {
38 builder->setInsertionPoint(op);
39 auto operands = llvm::to_vector<4>(op->getOperands());
40 operands.push_back(operand);
41 auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
42 operands, op->getAttrs());
43 op->replaceAllUsesWith(new_op.getOperation()->getResults());
44 op->erase();
45 return new_op;
46 }
47
48 // Returns success if the function has at most one op of the template type and
49 // assigns it to `result`, if present. If there are multiple such ops, returns
50 // failure.
51 template <typename OpT>
GetOp(FuncOp func,OpT * result)52 LogicalResult GetOp(FuncOp func, OpT* result) {
53 *result = {};
54 for (auto op : func.getOps<OpT>()) {
55 if (*result) return op.emitError("should be unique within a function");
56 *result = op;
57 }
58 return success();
59 }
60
runOnFunction()61 void RewriteTPUEmbeddingOps::runOnFunction() {
62 FuncOp func = getFunction();
63
64 RecvTPUEmbeddingActivationsOp recv_op;
65 if (failed(GetOp(func, &recv_op))) return signalPassFailure();
66
67 SendTPUEmbeddingGradientsOp send_op;
68 if (failed(GetOp(func, &send_op))) return signalPassFailure();
69
70 // No TPU embedding ops.
71 if (!recv_op && !send_op) return;
72
73 Location loc = recv_op ? recv_op.getLoc() : send_op.getLoc();
74 StringRef config = recv_op ? recv_op.config() : send_op.config();
75
76 // Create _RecvTPUEmbeddingDeduplicationData op.
77 OpBuilder builder(func.getBody());
78 auto output_ty = RankedTensorType::get({}, VariantType::get(&getContext()));
79 auto dedup_op = builder.create<_RecvTPUEmbeddingDeduplicationDataOp>(
80 loc, output_ty, config);
81
82 // Rewrite RecvTPUEmbeddingActivations op to the corresponding internal op.
83 if (recv_op)
84 AddOperandAndRewriteAs<_RecvTPUEmbeddingActivationsOp>(recv_op, dedup_op,
85 &builder);
86
87 // Rewrite SendTPUEmbeddingGradients op to the corresponding internal op and
88 // then update the OperandSegmentSize attribute.
89 if (send_op) {
90 int32_t operand_sizes[] = {static_cast<int32_t>(send_op.N()),
91 static_cast<int32_t>(send_op.NN()), 1};
92 auto attr_ty = VectorType::get(3, builder.getI32Type());
93 auto operand_size_attr = DenseIntElementsAttr::get(attr_ty, operand_sizes);
94
95 auto new_send_op = AddOperandAndRewriteAs<_SendTPUEmbeddingGradientsOp>(
96 send_op, dedup_op, &builder);
97 new_send_op->setAttr(new_send_op.getOperandSegmentSizeAttr(),
98 operand_size_attr);
99 }
100 }
101
102 } // anonymous namespace
103
CreateRewriteTPUEmbeddingOpsPass()104 std::unique_ptr<OperationPass<FuncOp>> CreateRewriteTPUEmbeddingOpsPass() {
105 return std::make_unique<RewriteTPUEmbeddingOps>();
106 }
107
108 static PassRegistration<RewriteTPUEmbeddingOps> pass(
109 "tf-rewrite-tpu-embedding-ops",
110 "Rewrites TPU embedding send/recv ops by adding TPU embedding "
111 "deduplication data");
112
113 } // namespace TF
114 } // namespace mlir
115