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 "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project
17 #include "mlir/Pass/Pass.h" // from @llvm-project
18 #include "mlir/Transforms/DialectConversion.h" // from @llvm-project
19 #include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
20 #include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
21 #include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
22
23 namespace mlir {
24 namespace kernel_gen {
25 namespace tf_framework {
26 namespace {
27
28 #define GEN_PASS_CLASSES
29 #include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
30
IsNotInsideTfEntryFunction(Operation * op)31 bool IsNotInsideTfEntryFunction(Operation* op) {
32 auto func = op->getParentOfType<FuncOp>();
33 return !func->hasAttrOfType<UnitAttr>(TFFrameworkDialect::kTFEntryAttrName);
34 }
35
36 // The pass rewrites the function marked with `tf_entry` attribute.
37 // * adds tf_framework::OpKernelContextType argument to the function,
38 // * std.alloc becomes tf_framework.alloc_raw,
39 // * std.dealloc becomes tf_framework.dealloc_raw.
40 class EmbedTFFrameworkFunctionAndAllocPass
41 : public EmbedTFFrameworkFunctionAndAllocPassBase<
42 EmbedTFFrameworkFunctionAndAllocPass> {
getDependentDialects(DialectRegistry & registry) const43 void getDependentDialects(DialectRegistry& registry) const override {
44 registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();
45 }
46
47 public:
runOnOperation()48 void runOnOperation() override {
49 ModuleOp m = getOperation();
50
51 // Populate patterns.
52 OwningRewritePatternList patterns;
53 PopulateEmbedTFFrameworkFunctionAndAllocConversionPatterns(m.getContext(),
54 &patterns);
55
56 // Set target.
57 ConversionTarget target(getContext());
58 target.addLegalDialect<tf_framework::TFFrameworkDialect>();
59
60 target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
61 if (!op->hasAttrOfType<UnitAttr>(TFFrameworkDialect::kTFEntryAttrName)) {
62 return true;
63 }
64 FunctionType func_type = op.getType();
65 return func_type.getNumInputs() > 0 &&
66 func_type.getInput(0).isa<OpKernelContextType>();
67 });
68 target.addDynamicallyLegalOp<AllocOp, DeallocOp>(
69 IsNotInsideTfEntryFunction);
70
71 if (failed(applyPartialConversion(m, target, std::move(patterns)))) {
72 signalPassFailure();
73 }
74 }
75 };
76
77 // The pass rewrites the function marked with `tf_entry` attribute.
78 // All contained `std.assert` operations are rewritten into calls to
79 // `tf_framework.report_error` and the required control flow to make
80 // execution of the function terminate.
81
82 class EmbedTFFrameworkAssertPass
83 : public EmbedTFFrameworkAssertPassBase<EmbedTFFrameworkAssertPass> {
getDependentDialects(DialectRegistry & registry) const84 void getDependentDialects(DialectRegistry& registry) const override {
85 registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();
86 }
87
88 public:
runOnOperation()89 void runOnOperation() override {
90 ModuleOp m = getOperation();
91
92 // Populate patterns.
93 OwningRewritePatternList patterns;
94 PopulateEmbedTFFrameworkAssertConversionPatterns(m.getContext(), &patterns);
95
96 // Set target.
97 ConversionTarget target(getContext());
98 target.addLegalDialect<tf_framework::TFFrameworkDialect,
99 StandardOpsDialect>();
100
101 target.addDynamicallyLegalOp<AssertOp>(IsNotInsideTfEntryFunction);
102
103 if (failed(applyPartialConversion(m, target, std::move(patterns)))) {
104 signalPassFailure();
105 }
106 }
107 };
108
109 } // namespace
110
111 std::unique_ptr<OperationPass<ModuleOp> >
CreateEmbedTFFrameworkFunctionAndAllocPass()112 CreateEmbedTFFrameworkFunctionAndAllocPass() {
113 return std::make_unique<EmbedTFFrameworkFunctionAndAllocPass>();
114 }
115
CreateEmbedTFFrameworkAssertPass()116 std::unique_ptr<OperationPass<ModuleOp> > CreateEmbedTFFrameworkAssertPass() {
117 return std::make_unique<EmbedTFFrameworkAssertPass>();
118 }
119
120 } // namespace tf_framework
121 } // namespace kernel_gen
122 } // namespace mlir
123