133 lines · cpp
1//===- WrapInZeroTripCheck.cpp - Loop transforms to add zero-trip-check ---===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/SCF/IR/SCF.h"10#include "mlir/Dialect/SCF/Transforms/Transforms.h"11#include "mlir/IR/IRMapping.h"12#include "mlir/IR/PatternMatch.h"13 14using namespace mlir;15 16/// Create zero-trip-check around a `while` op and return the new loop op in the17/// check. The while loop is rotated to avoid evaluating the condition twice.18///19/// Given an example below:20///21/// scf.while (%arg0 = %init) : (i32) -> i64 {22/// %val = .., %arg0 : i6423/// %cond = arith.cmpi .., %arg0 : i3224/// scf.condition(%cond) %val : i6425/// } do {26/// ^bb0(%arg1: i64):27/// %next = .., %arg1 : i3228/// scf.yield %next : i3229/// }30///31/// First clone before block to the front of the loop:32///33/// %pre_val = .., %init : i6434/// %pre_cond = arith.cmpi .., %init : i3235/// scf.while (%arg0 = %init) : (i32) -> i64 {36/// %val = .., %arg0 : i6437/// %cond = arith.cmpi .., %arg0 : i3238/// scf.condition(%cond) %val : i6439/// } do {40/// ^bb0(%arg1: i64):41/// %next = .., %arg1 : i3242/// scf.yield %next : i3243/// }44///45/// Create `if` op with the condition, rotate and move the loop into the else46/// branch:47///48/// %pre_val = .., %init : i6449/// %pre_cond = arith.cmpi .., %init : i3250/// scf.if %pre_cond -> i64 {51/// %res = scf.while (%arg1 = %va0) : (i64) -> i64 {52/// // Original after block53/// %next = .., %arg1 : i3254/// // Original before block55/// %val = .., %next : i6456/// %cond = arith.cmpi .., %next : i3257/// scf.condition(%cond) %val : i6458/// } do {59/// ^bb0(%arg2: i64):60/// %scf.yield %arg2 : i3261/// }62/// scf.yield %res : i6463/// } else {64/// scf.yield %pre_val : i6465/// }66FailureOr<scf::WhileOp> mlir::scf::wrapWhileLoopInZeroTripCheck(67 scf::WhileOp whileOp, RewriterBase &rewriter, bool forceCreateCheck) {68 // If the loop is in do-while form (after block only passes through values),69 // there is no need to create a zero-trip-check as before block is always run.70 if (!forceCreateCheck && isa<scf::YieldOp>(whileOp.getAfterBody()->front())) {71 return whileOp;72 }73 74 OpBuilder::InsertionGuard insertion_guard(rewriter);75 76 IRMapping mapper;77 Block *beforeBlock = whileOp.getBeforeBody();78 // Clone before block before the loop for zero-trip-check.79 for (auto [arg, init] :80 llvm::zip_equal(beforeBlock->getArguments(), whileOp.getInits())) {81 mapper.map(arg, init);82 }83 rewriter.setInsertionPoint(whileOp);84 for (auto &op : *beforeBlock) {85 if (isa<scf::ConditionOp>(op)) {86 break;87 }88 // Safe to clone everything as in a single block all defs have been cloned89 // and added to mapper in order.90 rewriter.insert(op.clone(mapper));91 }92 93 scf::ConditionOp condOp = whileOp.getConditionOp();94 Value clonedCondition = mapper.lookupOrDefault(condOp.getCondition());95 SmallVector<Value> clonedCondArgs = llvm::map_to_vector(96 condOp.getArgs(), [&](Value arg) { return mapper.lookupOrDefault(arg); });97 98 // Create rotated while loop.99 auto newLoopOp = scf::WhileOp::create(100 rewriter, whileOp.getLoc(), whileOp.getResultTypes(), clonedCondArgs,101 [&](OpBuilder &builder, Location loc, ValueRange args) {102 // Rotate and move the loop body into before block.103 auto newBlock = builder.getBlock();104 rewriter.mergeBlocks(whileOp.getAfterBody(), newBlock, args);105 auto yieldOp = cast<scf::YieldOp>(newBlock->getTerminator());106 rewriter.mergeBlocks(whileOp.getBeforeBody(), newBlock,107 yieldOp.getResults());108 rewriter.eraseOp(yieldOp);109 },110 [&](OpBuilder &builder, Location loc, ValueRange args) {111 // Pass through values.112 scf::YieldOp::create(builder, loc, args);113 });114 115 // Create zero-trip-check and move the while loop in.116 auto ifOp = scf::IfOp::create(117 rewriter, whileOp.getLoc(), clonedCondition,118 [&](OpBuilder &builder, Location loc) {119 // Then runs the while loop.120 rewriter.moveOpBefore(newLoopOp, builder.getInsertionBlock(),121 builder.getInsertionPoint());122 scf::YieldOp::create(builder, loc, newLoopOp.getResults());123 },124 [&](OpBuilder &builder, Location loc) {125 // Else returns the results from precondition.126 scf::YieldOp::create(builder, loc, clonedCondArgs);127 });128 129 rewriter.replaceOp(whileOp, ifOp);130 131 return newLoopOp;132}133