1 /* Copyright 2022 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 #include <algorithm>
16 #include <string>
17 #include <tuple>
18 #include <utility>
19
20 #include "llvm/ADT/StringRef.h"
21 #include "llvm/Support/SourceMgr.h"
22 #include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project
23 #include "mlir/IR/Builders.h" // from @llvm-project
24 #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
25 #include "mlir/IR/BuiltinTypes.h" // from @llvm-project
26 #include "mlir/IR/PatternMatch.h" // from @llvm-project
27 #include "mlir/Pass/Pass.h" // from @llvm-project
28 #include "mlir/Support/LogicalResult.h" // from @llvm-project
29 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
30 #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
31 #include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
32 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
33 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
34
35 namespace mlir {
36 namespace quant {
37 namespace {
38
39 class ConvertCustomAggregationOpToQuantStatsPass
40 : public PassWrapper<ConvertCustomAggregationOpToQuantStatsPass,
41 OperationPass<func::FuncOp>> {
42 public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertCustomAggregationOpToQuantStatsPass)43 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
44 ConvertCustomAggregationOpToQuantStatsPass)
45
46 StringRef getArgument() const final {
47 // This is the argument used to refer to the pass in the textual format (on
48 // the commandline for example).
49 return "quant-convert-tf-custom-aggregator-op-to-quant-stats";
50 }
51
getDescription() const52 StringRef getDescription() const final {
53 // This is a brief description of the pass.
54 return "Convert tf.CustomAggregator op to quant.Stats";
55 }
56
getDependentDialects(DialectRegistry & registry) const57 void getDependentDialects(DialectRegistry ®istry) const override {
58 registry.insert<TF::TensorFlowDialect>();
59 registry.insert<quant::QuantizationDialect>();
60 registry.insert<quantfork::QuantizationForkDialect>();
61 }
62
63 void runOnOperation() override;
64 };
65
66 class ConvertCustomAggregationOpToQuantStats : public RewritePattern {
67 public:
68 // Does not take ownership of context, which must refer to a valid value that
69 // outlives this object.
ConvertCustomAggregationOpToQuantStats(MLIRContext * context)70 explicit ConvertCustomAggregationOpToQuantStats(MLIRContext *context)
71 : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
72
matchAndRewrite(Operation * op,PatternRewriter & rewriter) const73 LogicalResult matchAndRewrite(Operation *op,
74 PatternRewriter &rewriter) const override {
75 // Return early if the given operator isn't the custom aggregator op.
76 if (op->getName().getStringRef() != "tf.CustomAggregator") return failure();
77
78 FloatAttr min = op->getAttr("min").dyn_cast_or_null<FloatAttr>();
79 FloatAttr max = op->getAttr("max").dyn_cast_or_null<FloatAttr>();
80
81 // When there are no min and max attributes, remove op.
82 if (min == nullptr || max == nullptr) {
83 op->replaceAllUsesWith(op->getOperands());
84 rewriter.eraseOp(op);
85 return success();
86 }
87
88 // The layer stats contain only the first min/max pairs.
89 ElementsAttr layer_stats = DenseFPElementsAttr::get(
90 RankedTensorType::get({2}, rewriter.getF32Type()),
91 {static_cast<float>(min.getValueAsDouble()),
92 static_cast<float>(max.getValueAsDouble())});
93 ElementsAttr axis_stats;
94 IntegerAttr axis;
95
96 rewriter.replaceOpWithNewOp<quantfork::StatisticsOp>(
97 op, op->getOperand(0), layer_stats, axis_stats, axis);
98 return success();
99 }
100 };
101
102 static PassRegistration<ConvertCustomAggregationOpToQuantStatsPass> pass;
103
runOnOperation()104 void ConvertCustomAggregationOpToQuantStatsPass::runOnOperation() {
105 MLIRContext *ctx = &getContext();
106 RewritePatternSet patterns(ctx);
107 func::FuncOp func = getOperation();
108
109 patterns.add<ConvertCustomAggregationOpToQuantStats>(ctx);
110 if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
111 func.emitError()
112 << "quant-convert-tf-custom-aggregator-op-to-quant-stats failed.";
113 signalPassFailure();
114 }
115 }
116
117 } // namespace
118
119 std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertCustomAggregationOpToQuantStatsPass()120 CreateConvertCustomAggregationOpToQuantStatsPass() {
121 return std::make_unique<ConvertCustomAggregationOpToQuantStatsPass>();
122 }
123
124 } // namespace quant
125 } // namespace mlir
126