1 //===- TosaInferShapes.cpp ------------------------------------------===//
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 // Propogate shapes forward along TOSA operations to resolve dynamic shape
10 // operations.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Analysis/DataFlowAnalysis.h"
15 #include "mlir/Dialect/StandardOps/IR/Ops.h"
16 #include "mlir/Dialect/Tensor/IR/Tensor.h"
17 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
18 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
19 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
20 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
21 #include "mlir/IR/BlockAndValueMapping.h"
22 #include "mlir/IR/Builders.h"
23 #include "mlir/IR/BuiltinOps.h"
24 #include "mlir/IR/Matchers.h"
25 #include "mlir/Pass/Pass.h"
26 #include "mlir/Transforms/DialectConversion.h"
27 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
28 #include "llvm/Support/FormatVariadic.h"
29 
30 using namespace mlir;
31 using namespace mlir::tosa;
32 
33 namespace {
34 
35 void propagateShapesInRegion(Region &region);
36 
37 void propagateShapesToTosaIf(Operation &op) {
38   tosa::IfOp ifOp = dyn_cast<tosa::IfOp>(op);
39   if (!ifOp)
40     return;
41 
42   for (auto &region : op.getRegions()) {
43     Block &frontBlock = region.front();
44     if (frontBlock.getNumArguments() + 1 != ifOp.getNumOperands())
45       return;
46 
47     for (int i = 0, e = frontBlock.getNumArguments(); i < e; i++) {
48       ValueKnowledge operandKnowledge = ValueKnowledge::getKnowledgeFromType(
49           ifOp.getOperand(i + 1).getType());
50       ValueKnowledge blockKnowledge = ValueKnowledge::getKnowledgeFromType(
51           frontBlock.getArgument(i).getType());
52       ValueKnowledge joinedKnowledge =
53           ValueKnowledge::join(operandKnowledge, blockKnowledge);
54       if (!joinedKnowledge)
55         continue;
56       frontBlock.getArgument(i).setType(joinedKnowledge.getType());
57     }
58 
59     propagateShapesInRegion(region);
60   }
61 
62   return;
63 }
64 
65 void propagateShapesInRegion(Region &region) {
66   DenseMap<Value, ShapedTypeComponents> shapesStorage;
67   auto setShapes = [&](Value val, Type t) {
68     if (auto st = t.dyn_cast<ShapedType>())
69       shapesStorage[val] = st;
70     else
71       shapesStorage[val] = t;
72   };
73   auto operandShape = [&](Value val) -> ShapeAdaptor {
74     // Query the WIP mapping rather than the type if set.
75     auto it = shapesStorage.find(val);
76     if (it == shapesStorage.end())
77       return nullptr;
78     return it->second;
79   };
80 
81   for (auto &block : region) {
82     for (Operation &op : block) {
83       if (op.getDialect()->getNamespace() !=
84           tosa::TosaDialect::getDialectNamespace())
85         continue;
86 
87       propagateShapesToTosaIf(op);
88 
89       InferShapedTypeOpInterface shapeInterface =
90           dyn_cast<InferShapedTypeOpInterface>(op);
91       if (!shapeInterface)
92         continue;
93 
94       SmallVector<ShapedTypeComponents> returnedShapes;
95 
96       ValueShapeRange range(op.getOperands(), operandShape);
97       if (shapeInterface
98               .inferReturnTypeComponents(op.getContext(), op.getLoc(), range,
99                                          op.getAttrDictionary(),
100                                          op.getRegions(), returnedShapes)
101               .succeeded()) {
102         for (auto it : llvm::zip(op.getResults(), returnedShapes)) {
103           Value result = std::get<0>(it);
104           ShapedTypeComponents predictedShape = std::get<1>(it);
105 
106           // Check whether this use case is replaceable. We define an op as
107           // being replaceable if it is used by a ReturnOp or a TosaOp.
108           bool replaceable = true;
109           for (auto user : result.getUsers()) {
110             if (isa<ReturnOp>(user))
111               continue;
112             if (user->getDialect()->getNamespace() ==
113                 tosa::TosaDialect::getDialectNamespace())
114               continue;
115 
116             replaceable = false;
117           }
118 
119           // Determine the knowledge based on the output type.
120           // TODO: should also query WIP type probably
121           Type resultTy = result.getType();
122           auto currentKnowledge =
123               ValueKnowledge::getKnowledgeFromType(resultTy);
124 
125           // Compute the knowledge based on the inferred type.
126           auto inferredKnowledge = ValueKnowledge::getPessimisticValueState();
127           inferredKnowledge.dtype =
128               resultTy.cast<ShapedType>().getElementType();
129           inferredKnowledge.hasRank = predictedShape.hasRank();
130           if (predictedShape.hasRank()) {
131             for (auto dim : predictedShape.getDims()) {
132               inferredKnowledge.sizes.push_back(dim);
133             }
134           }
135 
136           if (!replaceable)
137             continue;
138 
139           // Compute the new type based on the joined version.
140           auto newKnowledge =
141               ValueKnowledge::join(currentKnowledge, inferredKnowledge);
142           if (!newKnowledge)
143             continue;
144           setShapes(result, newKnowledge.getType());
145         }
146       }
147     }
148   }
149 
150   // Actually update types with updated shape knowledge.
151   for (auto it : shapesStorage) {
152     auto result = it.second;
153     if (result.hasRank()) {
154       Type t = it.first.getType().cast<ShapedType>().clone(result.getDims());
155       it.first.setType(t);
156     }
157   }
158 }
159 
160 /// Pass that performs shape propagation across TOSA operations. This includes
161 /// migrating to within the regions of if/while operations.
162 struct TosaInferShapes : public TosaInferShapesBase<TosaInferShapes> {
163 public:
164   void runOnFunction() override {
165     FuncOp func = getOperation();
166 
167     IRRewriter rewriter(func.getContext());
168 
169     propagateShapesInRegion(func.body());
170 
171     // Insert UnrealizedConversionCasts to guarantee ReturnOp agress with
172     // the FuncOp type.
173     func.walk([&](ReturnOp op) {
174       FuncOp parent = dyn_cast<FuncOp>(op->getParentOp());
175       if (!parent)
176         return;
177 
178       rewriter.setInsertionPoint(op);
179       FunctionType funcTy = func.getType();
180       auto resultTys = funcTy.getResults();
181 
182       bool castAdded = false;
183       SmallVector<Value> castedValues;
184       for (auto it : llvm::zip(op->getOperands(), resultTys)) {
185         auto operand = std::get<0>(it);
186         auto currentTy = operand.getType();
187         auto castTy = std::get<1>(it);
188         if (currentTy == castTy) {
189           castedValues.push_back(operand);
190           continue;
191         }
192 
193         castedValues.push_back(
194             rewriter.create<tensor::CastOp>(op.getLoc(), castTy, operand)
195                 .getResult());
196 
197         castAdded = true;
198       }
199 
200       if (castAdded) {
201         rewriter.replaceOpWithNewOp<ReturnOp>(op, castedValues);
202       }
203     });
204   }
205 };
206 } // end anonymous namespace
207 
208 std::unique_ptr<Pass> mlir::tosa::createTosaInferShapesPass() {
209   return std::make_unique<TosaInferShapes>();
210 }
211