1 //===- ConvertShapeConstraints.cpp - Conversion of shape constraints ------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
10
11 #include "../PassDetail.h"
12 #include "mlir/Dialect/SCF/SCF.h"
13 #include "mlir/Dialect/Shape/IR/Shape.h"
14 #include "mlir/Dialect/StandardOps/IR/Ops.h"
15 #include "mlir/Dialect/Tensor/IR/Tensor.h"
16 #include "mlir/IR/PatternMatch.h"
17 #include "mlir/Pass/Pass.h"
18 #include "mlir/Pass/PassRegistry.h"
19 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
20
21 using namespace mlir;
22 namespace {
23 #include "ShapeToStandard.cpp.inc"
24 } // namespace
25
26 namespace {
27 class ConvertCstrRequireOp : public OpRewritePattern<shape::CstrRequireOp> {
28 public:
29 using OpRewritePattern::OpRewritePattern;
matchAndRewrite(shape::CstrRequireOp op,PatternRewriter & rewriter) const30 LogicalResult matchAndRewrite(shape::CstrRequireOp op,
31 PatternRewriter &rewriter) const override {
32 rewriter.create<AssertOp>(op.getLoc(), op.pred(), op.msgAttr());
33 rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op, true);
34 return success();
35 }
36 };
37 } // namespace
38
populateConvertShapeConstraintsConversionPatterns(RewritePatternSet & patterns)39 void mlir::populateConvertShapeConstraintsConversionPatterns(
40 RewritePatternSet &patterns) {
41 patterns.add<CstrBroadcastableToRequire>(patterns.getContext());
42 patterns.add<CstrEqToRequire>(patterns.getContext());
43 patterns.add<ConvertCstrRequireOp>(patterns.getContext());
44 }
45
46 namespace {
47 // This pass eliminates shape constraints from the program, converting them to
48 // eager (side-effecting) error handling code. After eager error handling code
49 // is emitted, witnesses are satisfied, so they are replace with
50 // `shape.const_witness true`.
51 class ConvertShapeConstraints
52 : public ConvertShapeConstraintsBase<ConvertShapeConstraints> {
runOnOperation()53 void runOnOperation() override {
54 auto func = getOperation();
55 auto *context = &getContext();
56
57 RewritePatternSet patterns(context);
58 populateConvertShapeConstraintsConversionPatterns(patterns);
59
60 if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns))))
61 return signalPassFailure();
62 }
63 };
64 } // namespace
65
66 std::unique_ptr<OperationPass<FuncOp>>
createConvertShapeConstraintsPass()67 mlir::createConvertShapeConstraintsPass() {
68 return std::make_unique<ConvertShapeConstraints>();
69 }
70