Split the loop at the correct threshold based on predicate.
| 259 | |
| 260 | /// Split the loop at the correct threshold based on predicate. |
| 261 | static void performLoopSplit(RewriterBase &rewriter, ForOp forOp, |
| 262 | ComparisonPredicate pred, Value splitValue, |
| 263 | SmallPtrSet<Operation *, 4> &ifOps, CmpIOp cmp, |
| 264 | bool secondThen, bool copyCmp) { |
| 265 | Location loc = forOp.getLoc(); |
| 266 | Value step = forOp.getStep(); |
| 267 | auto constStep = llvm::dyn_cast<ConstantOp>(step.getDefiningOp()); |
| 268 | ValueRange iterArgs = forOp.getInitValues(); |
| 269 | Value lb = forOp.getLowerBound(); |
| 270 | Value ub = forOp.getUpperBound(); |
| 271 | |
| 272 | // Determine ForOp signedness based on unsignedCmp attribute |
| 273 | Signedness forOpSignedness = |
| 274 | forOp.getUnsignedCmp() ? Signedness::Unsigned : Signedness::Signed; |
| 275 | |
| 276 | // Compute split point depending on predicate. |
| 277 | // Increase splitPoint by 1 in the case of GT or LTE |
| 278 | Value splitPoint = splitValue; |
| 279 | if (pred == ComparisonPredicate::GREATER_THAN || |
| 280 | pred == ComparisonPredicate::LESS_THAN_OR_EQUAL) { |
| 281 | TileType constType = llvm::dyn_cast<TileType>(step.getType()); |
| 282 | auto intType = llvm::dyn_cast<IntegerType>(constType.getElementType()); |
| 283 | llvm::APInt val(intType.getWidth(), 1); |
| 284 | auto constAttr = DenseIntElementsAttr::get(constType, val); |
| 285 | auto constOp = ConstantOp::create(rewriter, loc, constType, constAttr); |
| 286 | splitPoint = AddIOp::create(rewriter, loc, splitValue, constOp); |
| 287 | } |
| 288 | if (!constStep || !isConstOne(constStep)) { |
| 289 | // Step is not equal to one (or dynamic) |
| 290 | // Need special handling, so that loop split point is aligned (i.e. == lb + |
| 291 | // k * step) So, splitPoint = start + Ceil(splitPoint - lb, step) * step |
| 292 | Value diff = SubIOp::create(rewriter, loc, splitPoint, lb); |
| 293 | Value k = DivIOp::create(rewriter, loc, diff, step, forOpSignedness, |
| 294 | RoundingMode::POSITIVE_INF); |
| 295 | Value kstep = MulIOp::create(rewriter, loc, k, step); |
| 296 | splitPoint = AddIOp::create(rewriter, loc, lb, kstep); |
| 297 | } |
| 298 | |
| 299 | Value minSplitPoint = |
| 300 | MinIOp::create(rewriter, loc, splitPoint, ub, forOpSignedness); |
| 301 | Value maxSplitPoint = |
| 302 | MaxIOp::create(rewriter, loc, splitPoint, lb, forOpSignedness); |
| 303 | |
| 304 | ForOp firstLoop, secondLoop; |
| 305 | Block *originalBody = forOp.getBody(); |
| 306 | SmallVector<Operation *> opsToClone; |
| 307 | // Collect operations for cloning |
| 308 | for (Operation &op : *originalBody) { |
| 309 | if (copyCmp || (&op != cmp.getOperation())) |
| 310 | opsToClone.push_back(&op); |
| 311 | } |
| 312 | |
| 313 | // First loop: before the condition flips true |
| 314 | rewriter.setInsertionPoint(forOp); |
| 315 | firstLoop = copyLoop(rewriter, forOp, cmp, ifOps, opsToClone, lb, |
| 316 | minSplitPoint, iterArgs, !secondThen); |
| 317 | // Second loop: after the condition is true |
| 318 | rewriter.setInsertionPointAfter(firstLoop); |
no test coverage detected