Diff 368526

mlir/include/mlir/Dialect/SCF/Passes.h

	Show All 22 Lines
	/// Creates a pass that specializes for loop for unrolling and			/// Creates a pass that specializes for loop for unrolling and
	/// vectorization.			/// vectorization.
	std::unique_ptr<Pass> createForLoopSpecializationPass();			std::unique_ptr<Pass> createForLoopSpecializationPass();

	/// Creates a pass that peels for loops at their upper bounds for			/// Creates a pass that peels for loops at their upper bounds for
	/// better vectorization.			/// better vectorization.
	std::unique_ptr<Pass> createForLoopPeelingPass();			std::unique_ptr<Pass> createForLoopPeelingPass();

	/// Creates a pass that canonicalizes affine.min ops in scf.for loops with			/// Creates a pass that canonicalizes affine.min and affine.max operations
	/// known lower and upper bounds.			/// inside of scf.for loops with known lower and upper bounds.
	std::unique_ptr<Pass> createAffineMinSCFCanonicalizationPass();			std::unique_ptr<Pass> createSCFAffineOpCanonicalizationPass();

	/// Creates a loop fusion pass which fuses parallel loops.			/// Creates a loop fusion pass which fuses parallel loops.
	std::unique_ptr<Pass> createParallelLoopFusionPass();			std::unique_ptr<Pass> createParallelLoopFusionPass();

	/// Creates a pass that specializes parallel loop for unrolling and			/// Creates a pass that specializes parallel loop for unrolling and
	/// vectorization.			/// vectorization.
	std::unique_ptr<Pass> createParallelLoopSpecializationPass();			std::unique_ptr<Pass> createParallelLoopSpecializationPass();

	Show All 24 Lines

mlir/include/mlir/Dialect/SCF/Passes.td

	Show All 13 Lines
	def SCFBufferize : FunctionPass<"scf-bufferize"> {			def SCFBufferize : FunctionPass<"scf-bufferize"> {
	let summary = "Bufferize the scf dialect.";			let summary = "Bufferize the scf dialect.";
	let constructor = "mlir::createSCFBufferizePass()";			let constructor = "mlir::createSCFBufferizePass()";
	let dependentDialects = ["memref::MemRefDialect"];			let dependentDialects = ["memref::MemRefDialect"];
	}			}

	// Note: Making this a canonicalization pattern would require a dependency			// Note: Making this a canonicalization pattern would require a dependency
	// of the SCF dialect on the Affine dialect or vice versa.			// of the SCF dialect on the Affine dialect or vice versa.
	def AffineMinSCFCanonicalization			def SCFAffineOpCanonicalization
	: FunctionPass<"canonicalize-scf-affine-min"> {			: FunctionPass<"canonicalize-scf-affine-op"> {
	let summary = "Canonicalize affine.min ops in the context of SCF loops with "			let summary = "Canonicalize affine.min and affine.max ops in the context of "
	"known bounds";			"SCF loops with known bounds";
	let constructor = "mlir::createAffineMinSCFCanonicalizationPass()";			let constructor = "mlir::createSCFAffineOpCanonicalizationPass()";
	let dependentDialects = ["AffineDialect"];			let dependentDialects = ["AffineDialect"];
	}			}

	def SCFForLoopPeeling			def SCFForLoopPeeling
	: FunctionPass<"for-loop-peeling"> {			: FunctionPass<"for-loop-peeling"> {
	let summary = "Peel `for` loops at their upper bounds.";			let summary = "Peel `for` loops at their upper bounds.";
	let constructor = "mlir::createForLoopPeelingPass()";			let constructor = "mlir::createForLoopPeelingPass()";
	let options = [			let options = [
	▲ Show 20 Lines • Show All 47 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/SCF/Transforms.h

Show All 12 Lines

#ifndef MLIR_DIALECT_SCF_TRANSFORMS_H_

#define MLIR_DIALECT_SCF_TRANSFORMS_H_

#include "mlir/Support/LLVM.h"

#include "llvm/ADT/ArrayRef.h"

namespace mlir {

class AffineMinOp;

class AffineMap;

class ConversionTarget;

struct LogicalResult;

class MLIRContext;

class Region;

class RewriterBase;

class TypeConverter;

class RewritePatternSet;

using OwningRewritePatternList = RewritePatternSet;

class Operation;

class Value;

class ValueRange;

namespace scf {

class IfOp;

class ForOp;

class ParallelOp;

class ForOp;

/// Try to canonicalize an affine.min operation in the context of `for` loops

/// Match "for loop"-like operations: If the first parameter is an iteration

/// with a known range.

/// variable, return lower/upper bounds via the second/third parameter and the

///

/// step size via the last parameter. The function should return `success` in

/// `loopMatcher` is used to retrieve loop bounds and step size for a given

/// that case. If the first parameter is not an iteration variable, return

/// iteration variable: If the first parameter is an iteration variable, return

/// `failure`.

/// lower/upper bounds via the second/third parameter and the step size via the

using LoopMatcherFn =

/// last parameter. The function should return `success` in that case. If the

function_ref<LogicalResult(Value, Value &, Value &, Value &)>;

/// first parameter is not an iteration variable, return `failure`.

/// Try to canonicalize an min/max operations in the context of for `loops` with

/// a known range.

///

/// `map` is the body of the min/max operation and `operands` are the SSA values

/// that the dimensions and symbols are bound to; dimensions are listed first.

/// If `isMin`, the operation is a min operation; otherwise, a max operation.

/// `loopMatcher` is used to retrieve loop bounds and the step size for a given

/// iteration variable.

ftynseUnsubmitted

Done

/// a known range.

///

- /// `map is the body of the min/max operation and `operands` are the SSA values

+ /// `map` is the body of the min/max operation and `operands` are the SSA values

/// that the dimensions and symbols are bound to; dimensions are listed first.

ftynse:

///

/// Note: `loopMatcher` allows this function to be used with any "for loop"-like

/// operation (scf.for, scf.parallel and even ops defined in other dialects).

LogicalResult canonicalizeAffineMinOpInLoop(

LogicalResult canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,

AffineMinOp minOp, RewriterBase &rewriter,

AffineMap map, ValueRange operands,

function_ref<LogicalResult(Value, Value &, Value &, Value &)> loopMatcher);

bool isMin, LoopMatcherFn loopMatcher);

/// Fuses all adjacent scf.parallel operations with identical bounds and step

/// into one scf.parallel operations. Uses a naive aliasing and dependency

/// analysis.

void naivelyFuseParallelOps(Region &region);

/// Rewrite a for loop with bounds/step that potentially do not divide evenly

/// into a for loop where the step divides the iteration space evenly, followed

Show All 18 Lines

/// (loop body)

/// }

/// scf.if %newUb < %ub {

/// (loop body)

/// }

/// ```

///

/// After loop peeling, this function tries to simplify/canonicalize affine.min

/// operations in the body of the loop and the scf.if, taking advantage of the

/// and affine.max ops in the body of the loop and the scf.if, taking advantage

/// fact that every iteration of the peeled loop is a "full" iteration. This

/// of the fact that the peeled loop has only "full" iterations and the scf.if

/// canonicalization is expected to enable further canonicalization

/// is always a partial iteration (if any). This canonicalization is expected to

/// opportunities through other patterns.

/// enable further canonicalization opportunities through other patterns.

///

/// The return value indicates whether the loop was rewritten or not. Loops are

/// not rewritten if:

/// * Loop step size is 1 or

/// * Loop bounds and step size are static, and step already divides the

/// iteration space evenly.

///

/// Note: This function rewrites the given scf.for loop in-place and creates a

▲ Show 20 Lines • Show All 63 Lines • ▼ Show 20 Lines

/// S0(I+2) S1(I+1) S2(I) // Pipelined kernel

/// }

/// S1(N) S2(N-1) // Epilogue

/// S2(N) // Epilogue

void populateSCFLoopPipeliningPatterns(RewritePatternSet &patterns,

const PipeliningOption &options);

/// Populate patterns for canonicalizing operations inside SCF loop bodies.

/// At the moment, only affine.min computations with iteration variables,

/// At the moment, only affine.min/max computations with iteration variables,

/// loop bounds and loop steps are canonicalized.

void populateSCFLoopBodyCanonicalizationPatterns(RewritePatternSet &patterns);

} // namespace scf

} // namespace mlir

#endif // MLIR_DIALECT_SCF_TRANSFORMS_H_

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

Show First 20 Lines • Show All 184 Lines • ▼ Show 20 Lines

static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints,

AffineMap alignedMap =

alignAffineMapWithValues(map, operands, dims, syms, &newSyms);

for (unsigned i = syms.size(); i < newSyms.size(); ++i)

constraints.addSymbolId(constraints.getNumSymbolIds(), newSyms[i]);

return constraints.addBound(type, pos, alignedMap);

}

/// This function tries to canonicalize affine.min operations by proving that

/// This function tries to canonicalize min/max operations by proving that their

/// its value is bounded by the same lower and upper bound. In that case, the

/// value is bounded by the same lower and upper bound. In that case, the

/// operation can be folded away.

///

/// Bounds are computed by FlatAffineValueConstraints. Invariants required for

/// finding/proving bounds should be supplied via `constraints`.

///

/// 1. Add dimensions for `minOp` and `minOpUb` (upper bound of `minOp`).

/// 1. Add dimensions for `op` and `opBound` (lower or upper bound of `op`).

/// 2. Compute an upper bound of `minOp` and bind it to `minOpUb`. SSA values

/// 2. Compute an upper bound of `op` (in case of `isMin`) or a lower bound (in

/// that are used in `minOp` but are not part of `dims`, are added as extra

/// case of `!isMin`) and bind it to `opBound`. SSA values that are used in

/// symbols to the constraint set.

/// `op` but are not part of `constraints`, are added as extra symbols.

/// 3. For each result of `minOp`: Add result as a dimension `r_i`. Prove that

/// 3. For each result of `op`: Add result as a dimension `r_i`. Prove that:

/// r_i >= minOpUb. If this is the case, ub(minOp) == lb(minOp) and `minOp`

/// * If `isMin`: r_i >= opBound

/// can be replaced with that bound.

/// * If `isMax`: r_i <= opBound

/// If this is the case, ub(op) == lb(op).

/// 4. Replace `op` with `opBound`.

ftynseUnsubmitted

Done

/// If this is the case, ub(op) == lb(op).

- /// 4. Replace `op` can be replaced with `opBound`.

+ /// 4. `op` can then be replaced with `opBound`.

///

/// In summary, the following constraints are added throughout this function.

ftynse:

///

/// In summary, the following constraints are added throughout this function.

/// Note: `invar` are dimensions added by the caller to express the invariants.

/// (Showing only the case where `isMin`.)

///

/// ------+-------+---------+-----+---------------+-------+-------------------

/// (various eq./ineq. constraining `invar`, added by the caller)

/// ... | 0 | 0 | 0 | 0 | ... | ...

/// ------+-------+---------+-----+---------------+-------+-------------------

/// (various ineq. constraining `minOp` in terms of `minOp` operands (`invar`

/// (various ineq. constraining `op` in terms of `op` operands (`invar` and

/// and extra `minOp` operands "extra syms" that are not in `invar`)).

/// extra `op` operands "extra syms" that are not in `invar`)).

/// ... | -1 | 0 | 0 | ... | ... | >= 0

/// ------+-------+---------+-----+---------------+-------+-------------------

/// (set `minOpUb` to `minOp` upper bound in terms of `invar` and extra syms)

/// (set `opBound` to `op` upper bound in terms of `invar` and "extra syms")

/// ... | 0 | -1 | 0 | ... | ... | = 0

/// ------+-------+---------+-----+---------------+-------+-------------------

/// (for each `minOp` map result r_i: copy previous constraints, set r_i to

/// (for each `op` map result r_i: set r_i to corresponding map result,

/// corresponding map result, prove r_i >= minOpUb via contradiction)

/// prove that r_i >= minOpUb via contradiction)

/// ... | 0 | 0 | -1 | ... | ... | = 0

/// 0 | 0 | 1 | -1 | 0 | -1 | >= 0

///

static LogicalResult

canonicalizeAffineMinOp(RewriterBase &rewriter, AffineMinOp minOp,

canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map,

ValueRange operands, bool isMin,

FlatAffineValueConstraints constraints) {

RewriterBase::InsertionGuard guard(rewriter);

AffineMap minOpMap = minOp.getAffineMap();

unsigned numResults = map.getNumResults();

unsigned numResults = minOpMap.getNumResults();

// Add a few extra dimensions.

unsigned dimMinOp = constraints.addDimId(); // `minOp`

unsigned dimOp = constraints.addDimId(); // `op`

unsigned dimMinOpUb = constraints.addDimId(); // `minOp` upper bound

unsigned dimOpBound = constraints.addDimId(); // `op` lower/upper bound

unsigned resultDimStart = constraints.getNumDimIds();

for (unsigned i = 0; i < numResults; ++i)

constraints.addDimId();

// Add an inequality for each result expr_i of minOpMap: minOp <= expr_i

// Add an inequality for each result expr_i of map:

if (failed(alignAndAddBound(constraints, FlatAffineConstraints::UB, dimMinOp,

// isMin: op <= expr_i, !isMin: op >= expr_i

minOpMap, minOp.operands())))

auto boundType =

isMin ? FlatAffineConstraints::UB : FlatAffineConstraints::LB;

if (failed(alignAndAddBound(constraints, boundType, dimOp, map, operands)))

return failure();

// Try to compute an upper bound for minOp, expressed in terms of the other

// Try to compute a lower/upper bound for op, expressed in terms of the other

// `dims` and extra symbols.

SmallVector<AffineMap> minOpValLb(1), minOpValUb(1);

SmallVector<AffineMap> opLb(1), opUb(1);

constraints.getSliceBounds(dimMinOp, 1, minOp.getContext(), &minOpValLb,

constraints.getSliceBounds(dimOp, 1, rewriter.getContext(), &opLb, &opUb);

&minOpValUb);

AffineMap boundMap = isMin ? opUb[0] : opLb[0];

// TODO: `getSliceBounds` may return multiple bounds at the moment. This is

// a TODO of `getSliceBounds` and not handled here.

if (!minOpValUb[0] || minOpValUb[0].getNumResults() != 1)

if (!boundMap || boundMap.getNumResults() != 1)

return failure(); // No or multiple upper bounds found.

return failure(); // No or multiple bounds found.

// Add an equality: dimMinOpUb = minOpValUb[0]

// Add an equality: Set dimOpBound to computed bound.

// Add back dimension for minOp. (Was removed by `getSliceBounds`.)

// Add back dimension for op. (Was removed by `getSliceBounds`.)

AffineMap alignedUbMap = minOpValUb[0].shiftDims(/*shift=*/1,

AffineMap alignedBoundMap = boundMap.shiftDims(/*shift=*/1, /*offset=*/dimOp);

/*offset=*/dimMinOp);

if (failed(constraints.addBound(FlatAffineConstraints::EQ, dimOpBound,

if (failed(constraints.addBound(FlatAffineConstraints::EQ, dimMinOpUb,

alignedBoundMap)))

alignedUbMap)))

return failure();

// If the constraint system is empty, there is an inconsistency. (E.g., this

// can happen if loop lb > ub.)

if (constraints.isEmpty())

return failure();

// Prove that each result of minOpMap has a lower bound that is equal to (or

// In the case of `isMin` (`!isMin` is inversed):

// greater than) the upper bound of minOp (`kDimMinOpUb`). In that case,

// Prove that each result of `map` has a lower bound that is equal to (or

// minOp can be replaced with the bound. I.e., prove that for each result

// greater than) the upper bound of `op` (`dimOpBound`). In that case, `op`

// can be replaced with the bound. I.e., prove that for each result

// expr_i (represented by dimension r_i):

// r_i >= minOpUb

// r_i >= opBound

// To prove this inequality, add its negation to the constraint set and prove

// that the constraint set is empty.

for (unsigned i = resultDimStart; i < resultDimStart + numResults; ++i) {

FlatAffineValueConstraints newConstr(constraints);

// Add an equality: r_i = expr_i

// Note: These equalities could have been added earlier and used to express

// minOp <= expr_i. However, then we run the risk that `getSliceBounds`

// computes minOpUb in terms of r_i dims, which is not desired.

if (failed(alignAndAddBound(newConstr, FlatAffineConstraints::EQ, i,

minOpMap.getSubMap({i - resultDimStart}),

map.getSubMap({i - resultDimStart}), operands)))

minOp.operands())))

return failure();

// Add inequality: r_i < minOpUb (equiv.: minOpUb - r_i - 1 >= 0)

// If `isMin`: Add inequality: r_i < opBound

// equiv.: opBound - r_i - 1 >= 0

// If `!isMin`: Add inequality: r_i > opBound

// equiv.: -opBound + r_i - 1 >= 0

ftynseUnsubmitted

Done

Nit: update minOpUb to opBound here? Maybe also spell out the max case?

ftynse: Nit: update minOpUb to opBound here? Maybe also spell out the max case?

SmallVector<int64_t> ineq(newConstr.getNumCols(), 0);

ineq[dimMinOpUb] = 1;

ineq[dimOpBound] = isMin ? 1 : -1;

ineq[i] = -1;

ineq[i] = isMin ? -1 : 1;

ineq[newConstr.getNumCols() - 1] = -1;

newConstr.addInequality(ineq);

if (!newConstr.isEmpty())

return failure();

}

// Lower and upper bound of `minOp` are equal. Replace `minOp` with its bound.

// Lower and upper bound of `op` are equal. Replace `minOp` with its bound.

AffineMap newMap = alignedUbMap;

AffineMap newMap = alignedBoundMap;

SmallVector<Value> newOperands;

unpackOptionalValues(constraints.getMaybeDimAndSymbolValues(), newOperands);

mlir::canonicalizeMapAndOperands(&newMap, &newOperands);

rewriter.setInsertionPoint(minOp);

rewriter.setInsertionPoint(op);

rewriter.replaceOpWithNewOp<AffineApplyOp>(minOp, newMap, newOperands);

rewriter.replaceOpWithNewOp<AffineApplyOp>(op, newMap, newOperands);

return success();

}

/// Try to simplify an affine.min operation `minOp` after loop peeling. This

/// Try to simplify a min/max operation `op` after loop peeling. This function

/// function detects affine.min operations such as (ub is the previous upper

/// can simplify min/max operations such as (ub is the previous upper bound of

/// bound of the unpeeled loop):

/// the unpeeled loop):

/// ```

/// #map = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>

/// %r = affine.min #affine.min #map(%iv)[%step, %ub]

/// ```

/// and rewrites them into (in the case the peeled loop):

/// ```

/// %r = %step

/// ```

/// affine.min operations inside the generated scf.if operation are rewritten in

/// min/max operations inside the generated scf.if operation are rewritten in

/// a similar way.

///

/// This function builds up a set of constraints, capable of proving that:

/// * Inside the peeled loop: min(step, ub - iv) == step

/// * Inside the scf.if operation: min(step, ub - iv) == ub - iv

///

/// Returns `success` if the given operation was replaced by a new operation;

/// `failure` otherwise.

///

/// Note: `ub` is the previous upper bound of the loop (before peeling).

/// `insideLoop` must be true for affine.min ops inside the loop and false for

/// `insideLoop` must be true for min/max ops inside the loop and false for

/// affine.min ops inside the scf.for op.

/// affine.min ops inside the scf.for op. For an explanation of the other

static LogicalResult rewritePeeledAffineOp(RewriterBase &rewriter,

/// parameters, see comment of `canonicalizeMinMaxOpInLoop`.

AffineMinOp minOp, Value iv,

static LogicalResult rewritePeeledMinMaxOp(RewriterBase &rewriter,

Value ub, Value step,

Operation *op, AffineMap map,

ValueRange operands, bool isMin,

Value iv, Value ub, Value step,

bool insideLoop) {

FlatAffineValueConstraints constraints;

constraints.addDimId(0, iv);

constraints.addDimId(1, ub);

constraints.addDimId(2, step);

if (auto constUb = getConstantIntValue(ub))

constraints.addBound(FlatAffineConstraints::EQ, 1, *constUb);

if (auto constStep = getConstantIntValue(step))

Show All 10 Lines

if (insideLoop) {

// ub - iv < step (equiv.: iv + -ub + step - 1 >= 0)

// Intuitively: `iv` is the split bound here, i.e., the iteration variable

// value of the very last iteration (in the unpeeled loop). At that point,

// there are less than `step` elements remaining. (Otherwise, the peeled

// loop would run for at least one more iteration.)

constraints.addInequality({1, -1, 1, -1});

}

return canonicalizeAffineMinOp(rewriter, minOp, constraints);

return canonicalizeMinMaxOp(rewriter, op, map, operands, isMin, constraints);

}

template <typename OpTy, bool IsMin>

static void

rewriteAffineOpAfterPeeling(RewriterBase &rewriter, ForOp forOp, scf::IfOp ifOp,

Value iv, Value splitBound, Value ub, Value step) {

forOp.walk([&](OpTy affineOp) {

(void)rewritePeeledMinMaxOp(rewriter, affineOp, affineOp.getAffineMap(),

affineOp.operands(), IsMin, iv, ub, step,

/*insideLoop=*/true);

});

ifOp.walk([&](OpTy affineOp) {

(void)rewritePeeledMinMaxOp(rewriter, affineOp, affineOp.getAffineMap(),

affineOp.operands(), IsMin, splitBound, ub,

step, /*insideLoop=*/false);

});

}

LogicalResult mlir::scf::peelAndCanonicalizeForLoop(RewriterBase &rewriter,

ForOp forOp,

scf::IfOp &ifOp) {

Value ub = forOp.upperBound();

Value splitBound;

if (failed(peelForLoop(rewriter, forOp, ifOp, splitBound)))

return failure();

// Rewrite affine.min ops.

// Rewrite affine.min and affine.max ops.

forOp.walk([&](AffineMinOp minOp) {

Value iv = forOp.getInductionVar(), step = forOp.step();

(void)rewritePeeledAffineOp(rewriter, minOp, forOp.getInductionVar(), ub,

rewriteAffineOpAfterPeeling<AffineMinOp, /*IsMin=*/true>(

forOp.step(), /*insideLoop=*/true);

rewriter, forOp, ifOp, iv, splitBound, ub, step);

});

rewriteAffineOpAfterPeeling<AffineMaxOp, /*IsMin=*/false>(

ifOp.walk([&](AffineMinOp minOp) {

rewriter, forOp, ifOp, iv, splitBound, ub, step);

(void)rewritePeeledAffineOp(rewriter, minOp, splitBound, ub, forOp.step(),

/*insideLoop=*/false);

});

return success();

}

/// Canonicalize AffineMinOp operations in the context of for loops with a known

/// Canonicalize min/max operations in the context of for loops with a known

/// range. Call `canonicalizeAffineMinOp` and add the following constraints to

/// range. Call `canonicalizeMinMaxOp` and add the following constraints to

/// the constraint system (along with the missing dimensions):

///

/// * iv >= lb

/// * iv < lb + step * ((ub - lb - 1) floorDiv step) + 1

///

/// Note: Due to limitations of FlatAffineConstraints, only constant step sizes

/// are currently supported.

LogicalResult mlir::scf::canonicalizeAffineMinOpInLoop(

LogicalResult

AffineMinOp minOp, RewriterBase &rewriter,

mlir::scf::canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,

function_ref<LogicalResult(Value, Value &, Value &, Value &)> loopMatcher) {

AffineMap map, ValueRange operands,

bool isMin, LoopMatcherFn loopMatcher) {

FlatAffineValueConstraints constraints;

DenseSet<Value> allIvs;

// Find all iteration variables among `minOp`'s operands add constrain them.

for (Value operand : minOp.operands()) {

for (Value operand : operands) {

// Skip duplicate ivs.

if (llvm::find(allIvs, operand) != allIvs.end())

continue;

// If `operand` is an iteration variable: Find corresponding loop

// bounds and step.

Value iv = operand;

Value lb, ub, step;

Show All 35 Lines

for (Value operand : operands) {

auto map = AffineMap::get(

/*dimCount=*/constraints.getNumDimIds(),

/*symbolCount=*/constraints.getNumSymbolIds(), /*result=*/ivUb);

if (failed(constraints.addBound(FlatAffineConstraints::UB, dimIv, map)))

return failure();

}

return canonicalizeAffineMinOp(rewriter, minOp, constraints);

return canonicalizeMinMaxOp(rewriter, op, map, operands, isMin, constraints);

}

static constexpr char kPeeledLoopLabel[] = "__peeled_loop__";

static constexpr char kPartialIterationLabel[] = "__partial_iteration__";

namespace {

struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {

ForLoopPeelingPattern(MLIRContext *ctx, bool skipPartial)

Show All 28 Lines

struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {

/// If set to true, loops inside partial iterations of another peeled loop

/// are not peeled. This reduces the size of the generated code. Partial

/// iterations are not usually performance critical.

/// Note: Takes into account the entire chain of parent operations, not just

/// the direct parent.

bool skipPartial;

};

/// Canonicalize AffineMinOp operations in the context of scf.for and

/// Canonicalize AffineMinOp/AffineMaxOp operations in the context of scf.for

/// scf.parallel loops with a known range.

/// and scf.parallel loops with a known range.

struct AffineMinSCFCanonicalizationPattern

template <typename OpTy, bool IsMin>

: public OpRewritePattern<AffineMinOp> {

struct AffineOpSCFCanonicalizationPattern : public OpRewritePattern<OpTy> {

using OpRewritePattern<AffineMinOp>::OpRewritePattern;

using OpRewritePattern<OpTy>::OpRewritePattern;

LogicalResult matchAndRewrite(AffineMinOp minOp,

LogicalResult matchAndRewrite(OpTy op,

PatternRewriter &rewriter) const override {

auto loopMatcher = [](Value iv, Value &lb, Value &ub, Value &step) {

if (scf::ForOp forOp = scf::getForInductionVarOwner(iv)) {

lb = forOp.lowerBound();

ub = forOp.upperBound();

step = forOp.step();

return success();

}

if (scf::ParallelOp parOp = scf::getParallelForInductionVarOwner(iv)) {

for (unsigned idx = 0; idx < parOp.getNumLoops(); ++idx) {

if (parOp.getInductionVars()[idx] == iv) {

lb = parOp.lowerBound()[idx];

ub = parOp.upperBound()[idx];

step = parOp.step()[idx];

return success();

}

return failure();

}

return failure();

};

return scf::canonicalizeAffineMinOpInLoop(minOp, rewriter, loopMatcher);

return scf::canonicalizeMinMaxOpInLoop(rewriter, op, op.getAffineMap(),

op.operands(), IsMin, loopMatcher);

}

};

} // namespace

namespace {

struct ParallelLoopSpecialization

: public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {

void runOnFunction() override {

Show All 20 Lines

void runOnFunction() override {

// Drop the markers.

funcOp.walk([](Operation *op) {

op->removeAttr(kPeeledLoopLabel);

op->removeAttr(kPartialIterationLabel);

});

}

};

struct AffineMinSCFCanonicalization

struct SCFAffineOpCanonicalization

: public AffineMinSCFCanonicalizationBase<AffineMinSCFCanonicalization> {

: public SCFAffineOpCanonicalizationBase<SCFAffineOpCanonicalization> {

void runOnFunction() override {

FuncOp funcOp = getFunction();

MLIRContext *ctx = funcOp.getContext();

RewritePatternSet patterns(ctx);

patterns.add<AffineMinSCFCanonicalizationPattern>(ctx);

scf::populateSCFLoopBodyCanonicalizationPatterns(patterns);

if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))

signalPassFailure();

}

};

} // namespace

std::unique_ptr<Pass> mlir::createAffineMinSCFCanonicalizationPass() {

std::unique_ptr<Pass> mlir::createSCFAffineOpCanonicalizationPass() {

return std::make_unique<AffineMinSCFCanonicalization>();

return std::make_unique<SCFAffineOpCanonicalization>();

}

std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {

return std::make_unique<ParallelLoopSpecialization>();

}

std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() {

return std::make_unique<ForLoopSpecialization>();

}

std::unique_ptr<Pass> mlir::createForLoopPeelingPass() {

return std::make_unique<ForLoopPeeling>();

}

void mlir::scf::populateSCFLoopBodyCanonicalizationPatterns(

RewritePatternSet &patterns) {

patterns.insert<AffineMinSCFCanonicalizationPattern>(patterns.getContext());

MLIRContext *ctx = patterns.getContext();

patterns

.insert<AffineOpSCFCanonicalizationPattern<AffineMinOp, /*IsMin=*/true>,

AffineOpSCFCanonicalizationPattern<AffineMaxOp, /*IsMin=*/false>>(

ctx);

}

mlir/test/Dialect/SCF/canonicalize-affine-op.mlir

This file was moved from mlir/test/Dialect/SCF/canonicalize-scf-affine-min.mlir.

	// RUN: mlir-opt %s -canonicalize-scf-affine-min -split-input-file \| FileCheck %s			// RUN: mlir-opt %s -canonicalize-scf-affine-op -split-input-file \| FileCheck %s

	// CHECK-LABEL: func @scf_for_canonicalize_min			// CHECK-LABEL: func @scf_for_canonicalize_min
	// CHECK: %[[C2:.*]] = constant 2 : i64			// CHECK: %[[C2:.*]] = constant 2 : i64
	// CHECK: scf.for			// CHECK: scf.for
	// CHECK: memref.store %[[C2]], %{{.*}}[] : memref<i64>			// CHECK: memref.store %[[C2]], %{{.*}}[] : memref<i64>
	func @scf_for_canonicalize_min(%A : memref<i64>) {			func @scf_for_canonicalize_min(%A : memref<i64>) {
	%c0 = constant 0 : index			%c0 = constant 0 : index
	%c2 = constant 2 : index			%c2 = constant 2 : index
	%c4 = constant 4 : index			%c4 = constant 4 : index

	scf.for %i = %c0 to %c4 step %c2 {			scf.for %i = %c0 to %c4 step %c2 {
	%1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)			%1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
	%2 = index_cast %1: index to i64			%2 = index_cast %1: index to i64
	memref.store %2, %A[]: memref<i64>			memref.store %2, %A[]: memref<i64>
	}			}
	return			return
	}			}

	// -----			// -----

				// CHECK-LABEL: func @scf_for_canonicalize_max
				// CHECK: %[[Cneg2:.*]] = constant -2 : i64
				// CHECK: scf.for
				// CHECK: memref.store %[[Cneg2]], %{{.*}}[] : memref<i64>
				func @scf_for_canonicalize_max(%A : memref<i64>) {
				%c0 = constant 0 : index
				%c2 = constant 2 : index
				%c4 = constant 4 : index

				scf.for %i = %c0 to %c4 step %c2 {
				%1 = affine.max affine_map<(d0, d1)[] -> (-2, -(d1 - d0))> (%i, %c4)
				ftynseUnsubmitted Done Reply Inline Actions Could we also have a test where the canonicalization does not kick in, preferably an adjacent one, e.g., using %c3. ftynse: Could we also have a test where the canonicalization does not kick in, preferably an adjacent…
				%2 = index_cast %1: index to i64
				memref.store %2, %A[]: memref<i64>
				}
				return
				}

				// -----

				// CHECK-LABEL: func @scf_for_max_not_canonicalizable
				// CHECK: scf.for
				// CHECK: affine.max
				// CHECK: index_cast
				func @scf_for_max_not_canonicalizable(%A : memref<i64>) {
				%c0 = constant 0 : index
				%c2 = constant 2 : index
				%c3 = constant 3 : index
				%c4 = constant 4 : index

				scf.for %i = %c0 to %c4 step %c2 {
				%1 = affine.max affine_map<(d0, d1)[] -> (-2, -(d1 - d0))> (%i, %c3)
				%2 = index_cast %1: index to i64
				memref.store %2, %A[]: memref<i64>
				}
				return
				}

				// -----

	// CHECK-LABEL: func @scf_for_loop_nest_canonicalize_min			// CHECK-LABEL: func @scf_for_loop_nest_canonicalize_min
	// CHECK: %[[C5:.*]] = constant 5 : i64			// CHECK: %[[C5:.*]] = constant 5 : i64
	// CHECK: scf.for			// CHECK: scf.for
	// CHECK: scf.for			// CHECK: scf.for
	// CHECK: memref.store %[[C5]], %{{.*}}[] : memref<i64>			// CHECK: memref.store %[[C5]], %{{.*}}[] : memref<i64>
	func @scf_for_loop_nest_canonicalize_min(%A : memref<i64>) {			func @scf_for_loop_nest_canonicalize_min(%A : memref<i64>) {
	%c0 = constant 0 : index			%c0 = constant 0 : index
	%c2 = constant 2 : index			%c2 = constant 2 : index
	▲ Show 20 Lines • Show All 158 Lines • Show Last 20 Lines

mlir/test/Dialect/SCF/canonicalize-scf-affine-min.mlir

This file was moved to mlir/test/Dialect/SCF/canonicalize-affine-op.mlir.

mlir/test/Dialect/SCF/for-loop-peeling.mlir

Show First 20 Lines • Show All 143 Lines • ▼ Show 20 Lines	scf.for %iv = %lb to %ub step %step {
%0 = addi %r, %casted : i32		%0 = addi %r, %casted : i32
memref.store %0, %d[] : memref<i32>		memref.store %0, %d[] : memref<i32>
}		}
return		return
}		}

// -----		// -----

// Test rewriting of affine.min ops. Make sure that more general cases than		// Test rewriting of affine.min/max ops. Make sure that more general cases than
// the ones above are successfully rewritten. Also make sure that the pattern		// the ones above are successfully rewritten. Also make sure that the pattern
// does not rewrite affine.min ops that should not be rewritten.		// does not rewrite ops that should not be rewritten.

// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 1)>		// CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 1)>
// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1 - 1)>		// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0, s1] -> (s0, -d0 + s1 - 1)>
// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0, s1, s2] -> (s0, -d0 + s1, s2)>		// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0, s1, s2] -> (s0, -d0 + s1, s2)>
// CHECK-DAG: #[[MAP4:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0)>		// CHECK-DAG: #[[MAP4:.*]] = affine_map<()[s0] -> (-s0)>
// CHECK-DAG: #[[MAP5:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0 + 1)>		// CHECK-DAG: #[[MAP5:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0)>
// CHECK-DAG: #[[MAP6:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0 - 1)>		// CHECK-DAG: #[[MAP6:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0 + 1)>
// CHECK-DAG: #[[MAP7:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0, s2 - (s2 - (s2 - s1) mod s0), s3)>		// CHECK-DAG: #[[MAP7:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0 - 1)>
// CHECK: func @test_affine_min_rewrite(		// CHECK-DAG: #[[MAP8:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0, s2 - (s2 - (s2 - s1) mod s0), s3)>
		// CHECK-DAG: #[[MAP9:.*]] = affine_map<()[s0, s1, s2] -> (s0 - (s0 - s1) mod s2 - s0)>
		// CHECK: func @test_affine_op_rewrite(
// CHECK-SAME: %[[LB:.]]: index, %[[UB:.]]: index, %[[STEP:.*]]: index,		// CHECK-SAME: %[[LB:.]]: index, %[[UB:.]]: index, %[[STEP:.*]]: index,
// CHECK-SAME: %[[MEMREF:.]]: memref<?xindex>, %[[SOME_VAL:.]]: index		// CHECK-SAME: %[[MEMREF:.]]: memref<?xindex>, %[[SOME_VAL:.]]: index
// CHECK: scf.for %[[IV:.]] = %[[LB]] to %{{.}} step %[[STEP]] {		// CHECK: scf.for %[[IV:.]] = %[[LB]] to %{{.}} step %[[STEP]] {
// (affine.min folded away)		// (affine.min folded away)
// CHECK: memref.store %[[STEP]]		// CHECK: memref.store %[[STEP]]
// (affine.min folded away)		// (affine.min folded away)
// CHECK: memref.store %[[STEP]]		// CHECK: memref.store %[[STEP]]
// CHECK: %[[RES2:.*]] = affine.apply #[[MAP1]]()[%[[STEP]]]		// CHECK: %[[RES2:.*]] = affine.apply #[[MAP1]]()[%[[STEP]]]
// CHECK: memref.store %[[RES2]]		// CHECK: memref.store %[[RES2]]
// CHECK: %[[RES3:.*]] = affine.min #[[MAP2]](%[[IV]])[%[[STEP]], %[[UB]]]		// CHECK: %[[RES3:.*]] = affine.min #[[MAP2]](%[[IV]])[%[[STEP]], %[[UB]]]
// CHECK: memref.store %[[RES3]]		// CHECK: memref.store %[[RES3]]
// CHECK: %[[RES4:.*]] = affine.min #[[MAP3]](%[[IV]])[%[[STEP]], %[[UB]], %[[SOME_VAL]]]		// CHECK: %[[RES4:.*]] = affine.min #[[MAP3]](%[[IV]])[%[[STEP]], %[[UB]], %[[SOME_VAL]]]
// CHECK: memref.store %[[RES4]]		// CHECK: memref.store %[[RES4]]
		// CHECK: %[[RES5:.*]] = affine.apply #[[MAP4]]()[%[[STEP]]]
		// CHECK: memref.store %[[RES5]]
// CHECK: }		// CHECK: }
// CHECK: scf.if {{.*}} {		// CHECK: scf.if {{.*}} {
// CHECK: %[[RES_IF_0:.*]] = affine.apply #[[MAP4]]()[%[[UB]], %[[LB]], %[[STEP]]]		// CHECK: %[[RES_IF_0:.*]] = affine.apply #[[MAP5]]()[%[[UB]], %[[LB]], %[[STEP]]]
// CHECK: memref.store %[[RES_IF_0]]		// CHECK: memref.store %[[RES_IF_0]]
// CHECK: %[[RES_IF_1:.*]] = affine.apply #[[MAP5]]()[%[[UB]], %[[LB]], %[[STEP]]]		// CHECK: %[[RES_IF_1:.*]] = affine.apply #[[MAP6]]()[%[[UB]], %[[LB]], %[[STEP]]]
// CHECK: memref.store %[[RES_IF_1]]		// CHECK: memref.store %[[RES_IF_1]]
// CHECK: %[[RES_IF_2:.*]] = affine.apply #[[MAP5]]()[%[[UB]], %[[LB]], %[[STEP]]]		// CHECK: %[[RES_IF_2:.*]] = affine.apply #[[MAP6]]()[%[[UB]], %[[LB]], %[[STEP]]]
// CHECK: memref.store %[[RES_IF_2]]		// CHECK: memref.store %[[RES_IF_2]]
// CHECK: %[[RES_IF_3:.*]] = affine.apply #[[MAP6]]()[%[[UB]], %[[LB]], %[[STEP]]]		// CHECK: %[[RES_IF_3:.*]] = affine.apply #[[MAP7]]()[%[[UB]], %[[LB]], %[[STEP]]]
// CHECK: memref.store %[[RES_IF_3]]		// CHECK: memref.store %[[RES_IF_3]]
// CHECK: %[[RES_IF_4:.*]] = affine.min #[[MAP7]]()[%[[STEP]], %[[LB]], %[[UB]], %[[SOME_VAL]]]		// CHECK: %[[RES_IF_4:.*]] = affine.min #[[MAP8]]()[%[[STEP]], %[[LB]], %[[UB]], %[[SOME_VAL]]]
// CHECK: memref.store %[[RES_IF_4]]		// CHECK: memref.store %[[RES_IF_4]]
		// CHECK: %[[RES_IF_5:.*]] = affine.apply #[[MAP9]]()[%[[UB]], %[[LB]], %[[STEP]]]
		// CHECK: memref.store %[[RES_IF_5]]
#map0 = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>		#map0 = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
#map1 = affine_map<(d0, d1)[s0] -> (d0 - d1 + 1, s0)>		#map1 = affine_map<(d0, d1)[s0] -> (d0 - d1 + 1, s0)>
#map2 = affine_map<(d0, d1)[s0] -> (s0 + 1, d0 - d1 + 1)>		#map2 = affine_map<(d0, d1)[s0] -> (s0 + 1, d0 - d1 + 1)>
#map3 = affine_map<(d0, d1)[s0] -> (s0, d0 - d1 - 1)>		#map3 = affine_map<(d0, d1)[s0] -> (s0, d0 - d1 - 1)>
#map4 = affine_map<(d0, d1, d2)[s0] -> (s0, d0 - d1, d2)>		#map4 = affine_map<(d0, d1, d2)[s0] -> (s0, d0 - d1, d2)>
func @test_affine_min_rewrite(%lb : index, %ub: index,		#map5 = affine_map<(d0, d1)[s0] -> (-s0, -d0 + d1)>
		func @test_affine_op_rewrite(%lb : index, %ub: index,
%step: index, %d : memref<?xindex>,		%step: index, %d : memref<?xindex>,
%some_val: index) {		%some_val: index) {
%c0 = constant 0 : index		%c0 = constant 0 : index
%c1 = constant 1 : index		%c1 = constant 1 : index
%c2 = constant 2 : index		%c2 = constant 2 : index
%c3 = constant 3 : index		%c3 = constant 3 : index
%c4 = constant 4 : index		%c4 = constant 4 : index
		%c5 = constant 5 : index
scf.for %iv = %lb to %ub step %step {		scf.for %iv = %lb to %ub step %step {
// Most common case: Rewrite min(%ub - %iv, %step) to %step.		// Most common case: Rewrite min(%ub - %iv, %step) to %step.
%m0 = affine.min #map0(%ub, %iv)[%step]		%m0 = affine.min #map0(%ub, %iv)[%step]
memref.store %m0, %d[%c0] : memref<?xindex>		memref.store %m0, %d[%c0] : memref<?xindex>

// Increase %ub - %iv a little bit, pattern should still apply.		// Increase %ub - %iv a little bit, pattern should still apply.
%m1 = affine.min #map1(%ub, %iv)[%step]		%m1 = affine.min #map1(%ub, %iv)[%step]
memref.store %m1, %d[%c1] : memref<?xindex>		memref.store %m1, %d[%c1] : memref<?xindex>

// Rewrite min(%ub - %iv + 1, %step + 1) to %step + 1.		// Rewrite min(%ub - %iv + 1, %step + 1) to %step + 1.
%m2 = affine.min #map2(%ub, %iv)[%step]		%m2 = affine.min #map2(%ub, %iv)[%step]
memref.store %m2, %d[%c2] : memref<?xindex>		memref.store %m2, %d[%c2] : memref<?xindex>

// min(%ub - %iv - 1, %step) cannot be simplified because %ub - %iv - 1		// min(%ub - %iv - 1, %step) cannot be simplified because %ub - %iv - 1
// can be smaller than %step. (Can be simplified in if-statement.)		// can be smaller than %step. (Can be simplified in if-statement.)
%m3 = affine.min #map3(%ub, %iv)[%step]		%m3 = affine.min #map3(%ub, %iv)[%step]
memref.store %m3, %d[%c3] : memref<?xindex>		memref.store %m3, %d[%c3] : memref<?xindex>

// min(%ub - %iv, %step, %some_val) cannot be simplified because the range		// min(%ub - %iv, %step, %some_val) cannot be simplified because the range
// of %some_val is unknown.		// of %some_val is unknown.
%m4 = affine.min #map4(%ub, %iv, %some_val)[%step]		%m4 = affine.min #map4(%ub, %iv, %some_val)[%step]
memref.store %m4, %d[%c4] : memref<?xindex>		memref.store %m4, %d[%c4] : memref<?xindex>

		// Rewrite max(-%ub + %iv, -%step) to -%ub + %iv (and -%step in the scf.if).
		%m5 = affine.max #map5(%ub, %iv)[%step]
		memref.store %m5, %d[%c5] : memref<?xindex>
}		}
return		return
}		}

// -----		// -----

// CHECK: func @nested_loops		// CHECK: func @nested_loops
// CHECK: scf.for {{.*}} {		// CHECK: scf.for {{.*}} {
Show All 40 Lines

This is an archive of the discontinued LLVM Phabricator instance.

[mlir][SCF] Generalize AffineMinSCFCanonicalization to min/max ops
ClosedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 368526

mlir/include/mlir/Dialect/SCF/Passes.h

mlir/include/mlir/Dialect/SCF/Passes.td

mlir/include/mlir/Dialect/SCF/Transforms.h

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

mlir/test/Dialect/SCF/canonicalize-affine-op.mlir

mlir/test/Dialect/SCF/canonicalize-scf-affine-min.mlir

mlir/test/Dialect/SCF/for-loop-peeling.mlir

This is an archive of the discontinued LLVM Phabricator instance.

[mlir][SCF] Generalize AffineMinSCFCanonicalization to min/max opsClosedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 368526

mlir/include/mlir/Dialect/SCF/Passes.h

mlir/include/mlir/Dialect/SCF/Passes.td

mlir/include/mlir/Dialect/SCF/Transforms.h

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

mlir/test/Dialect/SCF/canonicalize-affine-op.mlir

mlir/test/Dialect/SCF/canonicalize-scf-affine-min.mlir

mlir/test/Dialect/SCF/for-loop-peeling.mlir

[mlir][SCF] Generalize AffineMinSCFCanonicalization to min/max ops
ClosedPublic