1 /* Copyright 2021 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 <algorithm>
17 #include <string>
18 #include <vector>
19
20 #include "llvm/ADT/StringRef.h"
21 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
22 #include "mlir/IR/Builders.h" // from @llvm-project
23 #include "mlir/IR/BuiltinOps.h" // from @llvm-project
24 #include "mlir/IR/UseDefLists.h" // from @llvm-project
25 #include "mlir/Pass/Pass.h" // from @llvm-project
26 #include "mlir/Support/LLVM.h" // from @llvm-project
27 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
28 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
29 #include "tensorflow/compiler/mlir/tensorflow/transforms/savedmodel_passes_detail.h"
30 #include "tensorflow/core/platform/path.h"
31
32 namespace mlir {
33 namespace tf_saved_model {
34 namespace {
35
36 // This pass will replace a func's saved model asset bound inputs which are
37 // bound to tf.InitializeTableFromTextFileV2Op ops with tf.Const ops inside the
38 // func's body.
39 struct FreezeAssetsPass : public FreezeAssetsPassBase<FreezeAssetsPass> {
40 FreezeAssetsPass() = default;
41
FreezeAssetsPassmlir::tf_saved_model::__anon186ba2e10111::FreezeAssetsPass42 FreezeAssetsPass(const FreezeAssetsPass& pass) {}
FreezeAssetsPassmlir::tf_saved_model::__anon186ba2e10111::FreezeAssetsPass43 explicit FreezeAssetsPass(std::string saved_model_dir) {
44 this->saved_model_dir = saved_model_dir;
45 }
46 void runOnOperation() override;
47
48 private:
49 // TODO(team): should be a pass option.
50 std::string saved_model_dir;
51 };
52
runOnOperation()53 void FreezeAssetsPass::runOnOperation() {
54 auto module = getOperation();
55 if (!tf_saved_model::HasTfSavedModelSemantics(module)) {
56 return;
57 }
58 SymbolTable symbol_table(module);
59
60 for (auto func : module.getOps<func::FuncOp>()) {
61 llvm::BitVector args_to_erase(func.getNumArguments());
62 OpBuilder builder(func.getBody());
63
64 for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
65 SmallVector<TF::InitializeTableFromTextFileV2Op, 4>
66 init_table_from_text_file_ops_to_erase;
67 auto asset = LookupBoundInputOfType<AssetOp>(func, i, symbol_table);
68
69 if (!asset) continue;
70
71 auto arg = func.getArgument(i);
72 bool arg_is_deletable = true;
73 for (auto user : arg.getUsers()) {
74 if (auto read_op =
75 llvm::dyn_cast<TF::InitializeTableFromTextFileV2Op>(user)) {
76 init_table_from_text_file_ops_to_erase.push_back(read_op);
77 } else {
78 arg_is_deletable = false;
79 continue;
80 }
81 }
82 if (arg_is_deletable) {
83 args_to_erase.set(i);
84 }
85
86 // Replace the arg with a tf.Const op in the function body.
87 builder.setInsertionPointToStart(&func.getBody().front());
88
89 std::string asset_filename = asset.filename().str();
90 std::string filename =
91 tensorflow::io::JoinPath(saved_model_dir, asset_filename);
92 ShapedType shaped_type =
93 RankedTensorType::get({1}, TF::StringType::get(builder.getContext()));
94 auto const_op = builder.create<TF::ConstOp>(
95 asset.getLoc(),
96 DenseStringElementsAttr::get(shaped_type, {filename}));
97 for (auto init_op : init_table_from_text_file_ops_to_erase) {
98 // Replace the InitializeTableFromTextFileV2Op to use the saved model's
99 // asset filepath.
100 builder.setInsertionPoint(init_op);
101 builder.create<TF::InitializeTableFromTextFileV2Op>(
102 init_op.getLoc(), init_op.table_handle(), const_op.getResult(),
103 init_op.key_index(), init_op.value_index(), init_op.vocab_size(),
104 init_op.delimiter());
105 init_op.erase();
106 }
107 }
108 func.eraseArguments(args_to_erase);
109 }
110 }
111
112 } // namespace
113
CreateFreezeAssetsPass(std::string saved_model_dir)114 std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeAssetsPass(
115 std::string saved_model_dir) {
116 return std::make_unique<FreezeAssetsPass>(saved_model_dir);
117 }
118
119 } // namespace tf_saved_model
120 } // namespace mlir
121