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/STLExtras.h"
17 #include "mlir/IR/Attributes.h" // from @llvm-project
18 #include "mlir/IR/Builders.h" // from @llvm-project
19 #include "mlir/IR/BuiltinOps.h" // from @llvm-project
20 #include "mlir/IR/PatternMatch.h" // from @llvm-project
21 #include "mlir/Pass/Pass.h" // from @llvm-project
22 #include "mlir/Pass/PassManager.h" // from @llvm-project
23 #include "mlir/Pass/PassRegistry.h" // from @llvm-project
24 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
25 #include "mlir/Transforms/Passes.h" // from @llvm-project
26 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
27 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
28 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes_detail.h"
29
30 #define DEBUG_TYPE "tf-gpu-op-fusion"
31
32 namespace mlir {
33 namespace TF {
34
35 namespace {
36
37 // GpuOpFusionPass is a pass performing fusion specific to GPU targets.
38 // This is an ad-hoc pass for now, but should be integrated with some notion
39 // of "target" in the MLIR pipeline in the future.
40 class GpuOpFusionPass : public TensorflowGPUFusionBase<GpuOpFusionPass> {
41 public:
42 void runOnFunction() final;
43 };
44
45 // %y:6 = "tf.FusedBatchNormV3"(%x, %scale, %offset, %mean, %variance)
46 // %0 = "tf.Relu"(%y#0)
47 // ->
48 // %y:6 = "tf._FusedBatchNormEx"(%x, %scale, %offset, %mean, %variance)
49 //
50 // Or:
51 // %y:6 = "tf.FusedBatchNormV3"(%x, %scale, %offset, %mean, %variance)
52 // %0 = "tf.AddV2"(%y#0, %side_input)
53 // %1 = "tf.Relu"(%0)
54 // ->
55 // %y:6 = "tf._FusedBatchNormEx"(%x, %scale, %offset, %mean, %variance,
56 // %side_input)
57 // TODO(aminim): we should revisit this as a declarative pattern.
58 // For the second pattern, there is not good way in the framework to handle the
59 // commutativity of the AddV2: we want the FusedBatchNormV3 on any side.
60 // Also we need some native calls to handle the "hasOneUse" aspects and the
61 // optional extra operands for the AddV2 case.
62 struct ReluToFusedBatchNorm : public OpRewritePattern<ReluOp> {
63 using OpRewritePattern<ReluOp>::OpRewritePattern;
64
matchAndRewritemlir::TF::__anon81221d010111::ReluToFusedBatchNorm65 LogicalResult matchAndRewrite(ReluOp relu_op,
66 PatternRewriter &rewriter) const override {
67 Operation *relu_input = relu_op.features().getDefiningOp();
68 if (!relu_input) return failure();
69 auto batch_norm = dyn_cast_or_null<FusedBatchNormV3Op>(relu_input);
70 AddV2Op add_op;
71 Value side_input;
72 if (!batch_norm) {
73 // We don't have a FusedBatchNorm as input to the ReLu, but we can get
74 // through an AddV2 as well.
75 add_op = dyn_cast_or_null<AddV2Op>(relu_input);
76 if (!add_op) return failure();
77
78 batch_norm =
79 dyn_cast_or_null<FusedBatchNormV3Op>(add_op.x().getDefiningOp());
80 if (batch_norm) {
81 side_input = add_op.y();
82 } else {
83 // Didn't get a FusedBatchNorm on the LHS of the AddV2, try the RHS.
84 batch_norm =
85 dyn_cast_or_null<FusedBatchNormV3Op>(add_op.y().getDefiningOp());
86 if (!batch_norm) return failure();
87 side_input = add_op.x();
88 }
89 }
90 assert(batch_norm);
91 if (batch_norm.is_training()) return failure();
92 if (!batch_norm.y().hasOneUse()) return failure();
93
94 // Build the newly fused operation to replace the batch norm
95 OperationState state(batch_norm.getLoc(),
96 _FusedBatchNormExOp::getOperationName());
97 state.addOperands(batch_norm.getOperands());
98 if (side_input) state.operands.push_back(side_input);
99 state.addTypes(batch_norm.getResultTypes());
100 state.addAttributes(batch_norm->getAttrs());
101 Operation *op = rewriter.createOperation(state);
102 rewriter.replaceOp(batch_norm, op->getResults());
103
104 // Depending on the case, we may fuse the add, the relu, or both.
105 if (!add_op || add_op.z().hasOneUse()) {
106 // We fuse the Relu only if the add has a single use, otherwise we only
107 // fuse the add itself.
108 op->setAttr("activation_mode", rewriter.getStringAttr("Relu"));
109 rewriter.replaceOp(relu_op, op->getResult(0));
110 }
111 if (add_op) {
112 rewriter.replaceOp(add_op, op->getResult(0));
113 }
114
115 return success();
116 }
117 };
118
runOnFunction()119 void GpuOpFusionPass::runOnFunction() {
120 FuncOp func = getFunction();
121 OwningRewritePatternList patterns(&getContext());
122 patterns.insert<ReluToFusedBatchNorm>(&getContext());
123 (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
124 }
125
126 } // namespace
127
CreateGpuOpFusionPass()128 std::unique_ptr<OperationPass<FuncOp>> CreateGpuOpFusionPass() {
129 return std::make_unique<GpuOpFusionPass>();
130 }
131
132 static PassRegistration<GpuOpFusionPass> layout_assignment;
133
134 } // namespace TF
135 } // namespace mlir
136