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 "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" // from @llvm-project
17 #include "mlir/Dialect/Func/IR/FuncOps.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 // Converts tf_framework.assert to a conditional branch that reports an error to
29 // OpKernelContext and creates a fake memref using NullMemRefOp.
30 class TFAssertOpConverter : public OpConversionPattern<TFAssertOp> {
31 public:
32 using OpConversionPattern<TFAssertOp>::OpConversionPattern;
33
matchAndRewrite(TFAssertOp op,OpAdaptor adaptor,ConversionPatternRewriter & rewriter) const34 LogicalResult matchAndRewrite(
35 TFAssertOp op, OpAdaptor adaptor,
36 ConversionPatternRewriter &rewriter) const override {
37 Location loc = op.getLoc();
38
39 // Split the block to insert CondBr.
40 OpBuilder::InsertPoint ip = rewriter.saveInsertionPoint();
41 Block *split_block = rewriter.splitBlock(
42 rewriter.getInsertionBlock(), std::next(rewriter.getInsertionPoint()));
43
44 auto func = op->getParentOfType<func::FuncOp>();
45 Block *error_reporting_block =
46 rewriter.createBlock(&func.getRegion(), {}, {});
47 rewriter.create<ReportErrorOp>(loc, adaptor.ctx(), adaptor.error_code(),
48 adaptor.msg());
49
50 SmallVector<Value, 2> null_memrefs;
51 for (auto type : func.getFunctionType().getResults()) {
52 null_memrefs.push_back(rewriter.create<NullMemRefOp>(loc, type));
53 }
54 rewriter.create<func::ReturnOp>(loc, null_memrefs);
55
56 rewriter.restoreInsertionPoint(ip);
57 rewriter.replaceOpWithNewOp<cf::CondBranchOp>(
58 op, adaptor.arg(), split_block, llvm::None, error_reporting_block,
59 llvm::None);
60 return success();
61 }
62 };
63
64 #define GEN_PASS_CLASSES
65 #include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
66
IsNotInsideTfEntryFunction(Operation * op)67 bool IsNotInsideTfEntryFunction(Operation *op) {
68 auto func = op->getParentOfType<func::FuncOp>();
69 return !func->hasAttrOfType<UnitAttr>(TFFrameworkDialect::kTFEntryAttrName);
70 }
71 // All contained `tf_framework.assert` operations are rewritten into calls to
72 // `tf_framework.report_error` and the required control flow to make
73 // execution of the function terminate.
74 class RewriteTFFrameworkAssertPass
75 : public RewriteTFFrameworkAssertBase<RewriteTFFrameworkAssertPass> {
getDependentDialects(DialectRegistry & registry) const76 void getDependentDialects(DialectRegistry ®istry) const override {
77 registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();
78 }
79
80 public:
runOnOperation()81 void runOnOperation() override {
82 ModuleOp m = getOperation();
83
84 // Populate patterns.
85 MLIRContext *context = &getContext();
86 RewritePatternSet patterns(context);
87 patterns.add<TFAssertOpConverter>(context);
88 PopulateEmbedTFFrameworkAssertPattern(&patterns);
89
90 // Set target.
91 ConversionTarget target(getContext());
92 target.addLegalDialect<tf_framework::TFFrameworkDialect, func::FuncDialect,
93 cf::ControlFlowDialect>();
94 target.addIllegalOp<TFAssertOp>();
95 target.addDynamicallyLegalOp<cf::AssertOp>(IsNotInsideTfEntryFunction);
96
97 if (failed(applyPartialConversion(m, target, std::move(patterns)))) {
98 signalPassFailure();
99 }
100 }
101 };
102
103 } // namespace
104
CreateRewriteTFFrameworkAssert()105 std::unique_ptr<OperationPass<ModuleOp> > CreateRewriteTFFrameworkAssert() {
106 return std::make_unique<RewriteTFFrameworkAssertPass>();
107 }
108
109 } // namespace tf_framework
110 } // namespace kernel_gen
111 } // namespace mlir
112