Create a copy of the loop with new bounds & partial copy of if-blocks
| 182 | |
| 183 | /// Create a copy of the loop with new bounds & partial copy of if-blocks |
| 184 | static ForOp copyLoop(RewriterBase &rewriter, ForOp forOp, CmpIOp cmpOp, |
| 185 | SmallPtrSet<Operation *, 4> &ifOps, |
| 186 | SmallVector<Operation *> &opsToClone, Value lowerBound, |
| 187 | Value upperBound, ValueRange iterArgs, bool cloneThen) { |
| 188 | Location loc = forOp.getLoc(); |
| 189 | Value step = forOp.getStep(); |
| 190 | IRMapping mapper; |
| 191 | auto newLoop = |
| 192 | ForOp::create(rewriter, loc, lowerBound, upperBound, step, iterArgs, |
| 193 | /*bodyBuilder=*/nullptr, forOp.getUnsignedCmp()); |
| 194 | rewriter.setInsertionPointToStart(newLoop.getBody()); |
| 195 | |
| 196 | for (auto [orig, repl] : llvm::zip(forOp.getBody()->getArguments(), |
| 197 | newLoop.getBody()->getArguments())) |
| 198 | mapper.map(orig, repl); |
| 199 | |
| 200 | // Process all operations selected for copy |
| 201 | for (Operation *op : opsToClone) { |
| 202 | if (cmpOp.getOperation() == op) { |
| 203 | // Replace CmpOp with constant value |
| 204 | TileType constType = TileType::get({}, rewriter.getI1Type()); |
| 205 | llvm::APInt val(1, cloneThen ? 1 : 0); |
| 206 | auto constAttr = DenseIntElementsAttr::get(constType, val); |
| 207 | auto cmpConst = |
| 208 | ConstantOp::create(rewriter, op->getLoc(), constType, constAttr); |
| 209 | mapper.map(cmpOp.getResult(), cmpConst); |
| 210 | continue; |
| 211 | } else if (!ifOps.contains(op)) { |
| 212 | rewriter.clone(*op, mapper); |
| 213 | continue; |
| 214 | } |
| 215 | // Current operation is IfOp that we split |
| 216 | IfOp ifOp = cast<IfOp>(op); |
| 217 | bool is_continue = false; |
| 218 | Region ®ion = cloneThen ? ifOp.getThenRegion() : ifOp.getElseRegion(); |
| 219 | for (Operation &subOp : region.front()) { |
| 220 | // Copy all operations from one of the regions |
| 221 | if (isa<ContinueOp>(subOp)) { |
| 222 | rewriter.clone(subOp, mapper); |
| 223 | // Stop cloning operations at ContinueOp |
| 224 | is_continue = true; |
| 225 | break; |
| 226 | } |
| 227 | if (isa<YieldOp>(subOp)) { |
| 228 | // Map ifResult to the YieldOp |
| 229 | auto yieldOp = cast<YieldOp>(subOp); |
| 230 | for (auto [ifResult, yieldArg] : |
| 231 | llvm::zip_equal(ifOp.getResults(), yieldOp.getOperands())) |
| 232 | mapper.map(ifResult, mapper.lookupOrDefault(yieldArg)); |
| 233 | } else { |
| 234 | // General operation |
| 235 | rewriter.clone(subOp, mapper); |
| 236 | } |
| 237 | } |
| 238 | // Continue was met inside if-block - don't need to copy operations below |
| 239 | if (is_continue) |
| 240 | break; |
| 241 | } |