From 14f39fefe747a8cd4288f61cd0f7d31f54d727e5 Mon Sep 17 00:00:00 2001 From: Andi Drebes Date: Tue, 19 Sep 2023 12:19:33 +0200 Subject: [PATCH] refactor(compiler): Separate code for static loop analysis from batching --- .../concretelang/Analysis/StaticLoops.h | 65 +++ .../compiler/lib/Analysis/CMakeLists.txt | 1 + .../compiler/lib/Analysis/StaticLoops.cpp | 437 +++++++++++++++++ .../compiler/lib/Transforms/Batching.cpp | 451 +----------------- 4 files changed, 511 insertions(+), 443 deletions(-) create mode 100644 compilers/concrete-compiler/compiler/include/concretelang/Analysis/StaticLoops.h create mode 100644 compilers/concrete-compiler/compiler/lib/Analysis/StaticLoops.cpp diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Analysis/StaticLoops.h b/compilers/concrete-compiler/compiler/include/concretelang/Analysis/StaticLoops.h new file mode 100644 index 000000000..ad45519e1 --- /dev/null +++ b/compilers/concrete-compiler/compiler/include/concretelang/Analysis/StaticLoops.h @@ -0,0 +1,65 @@ +// Part of the Concrete Compiler Project, under the BSD3 License with Zama +// Exceptions. See +// https://github.com/zama-ai/concrete-compiler-internal/blob/main/LICENSE.txt +// for license information. + +#ifndef CONCRETELANG_ANALYSIS_STATIC_LOOPS_H +#define CONCRETELANG_ANALYSIS_STATIC_LOOPS_H + +#include + +namespace mlir { +namespace concretelang { + +/// Convenience class that holds all parameters of a loop +struct LoopsBoundsAndStep { + int64_t lb; + int64_t ub; + int64_t step; + + LoopsBoundsAndStep operator+(const LoopsBoundsAndStep &other) { + return LoopsBoundsAndStep{lb + other.lb, ub + other.ub, step + other.step}; + } + LoopsBoundsAndStep operator-(const LoopsBoundsAndStep &other) { + return LoopsBoundsAndStep{lb - other.lb, ub - other.ub, step - other.step}; + } + LoopsBoundsAndStep operator*(const LoopsBoundsAndStep &other) { + return LoopsBoundsAndStep{lb * other.lb, ub * other.ub, step * other.step}; + } + LoopsBoundsAndStep operator/(int64_t d) { + return LoopsBoundsAndStep{lb / d, ub / d, step / d}; + } +}; + +bool isQuasiAffineIVExpression(mlir::Value expr, + mlir::scf::ForOp *owningForOp = nullptr); + +bool isQuasiAffineIVExpression(mlir::OpFoldResult expr, + mlir::scf::ForOp *owningForOp = nullptr); + +bool isQuasiAffineIVExpressionWithConstantStep( + mlir::OpFoldResult expr, mlir::scf::ForOp *forOp = nullptr, + LoopsBoundsAndStep *basOut = nullptr); + +std::optional +getBoundsOfQuasiAffineIVExpression(mlir::Value expr, mlir::scf::ForOp forOp); + +std::optional +getBoundsOfQuasiAffineIVExpression(mlir::OpFoldResult expr, + mlir::scf::ForOp forOp); + +int64_t getStaticTripCount(int64_t lb, int64_t ub, int64_t step); +int64_t getStaticTripCount(const LoopsBoundsAndStep &bas); +int64_t getStaticTripCount(mlir::scf::ForOp forOp); +int64_t getNestedStaticTripCount(llvm::ArrayRef nest); +bool isStaticLoop(mlir::scf::ForOp forOp, int64_t *ilb = nullptr, + int64_t *iub = nullptr, int64_t *istep = nullptr); + +bool isConstantIndexValue(mlir::Value v); +int64_t getConstantIndexValue(mlir::Value v); +bool isConstantIndexValue(mlir::Value v, int64_t i); + +} // namespace concretelang +} // namespace mlir + +#endif diff --git a/compilers/concrete-compiler/compiler/lib/Analysis/CMakeLists.txt b/compilers/concrete-compiler/compiler/lib/Analysis/CMakeLists.txt index b05f53093..c9839e08b 100644 --- a/compilers/concrete-compiler/compiler/lib/Analysis/CMakeLists.txt +++ b/compilers/concrete-compiler/compiler/lib/Analysis/CMakeLists.txt @@ -2,6 +2,7 @@ add_compile_options(-fsized-deallocation) add_mlir_library( AnalysisUtils + StaticLoops.cpp Utils.cpp DEPENDS mlir-headers diff --git a/compilers/concrete-compiler/compiler/lib/Analysis/StaticLoops.cpp b/compilers/concrete-compiler/compiler/lib/Analysis/StaticLoops.cpp new file mode 100644 index 000000000..953427f89 --- /dev/null +++ b/compilers/concrete-compiler/compiler/lib/Analysis/StaticLoops.cpp @@ -0,0 +1,437 @@ +// Part of the Concrete Compiler Project, under the BSD3 License with Zama +// Exceptions. See +// https://github.com/zama-ai/concrete-compiler-internal/blob/main/LICENSE.txt +// for license information. + +#include +#include + +#include + +namespace mlir { +namespace concretelang { + +/// Checks whether the expression `expr` is a quasi-affine expression +/// on a single induction variable. If an induction variable is +/// referenced, the owning for loop is returned in `*owningForOp`. +bool isQuasiAffineIVExpression(mlir::Value expr, + mlir::scf::ForOp *owningForOp) { + if (mlir::Operation *op = expr.getDefiningOp()) { + if (llvm::isa(op)) { + return true; + } else if (llvm::isa(op)) { + mlir::scf::ForOp forLHS; + mlir::scf::ForOp forRHS; + + if (!isQuasiAffineIVExpression(op->getOperand(0), &forLHS) || + !isQuasiAffineIVExpression(op->getOperand(1), &forRHS)) { + return false; + } else { + // Check that appearances of IVs refer to the same IV + if (forLHS && forRHS && forLHS != forRHS) + return false; + } + + // Assume that the expression is already canonicalized, so that + // IVs appear only in numerators and on one side of a + // multiplication subexpression + if ((llvm::isa(op) && forLHS && forRHS) || + (llvm::isa(op) && forRHS)) + return false; + + if (owningForOp != nullptr) { + if (forLHS) + *owningForOp = forLHS; + else if (forRHS) + *owningForOp = forRHS; + } + + return true; + } else if (mlir::AffineApplyOp applyOp = + llvm::dyn_cast(op)) { + // Affine.apply: make sure that all operands are either constant + // expressions or using IVs of the same loop + mlir::scf::ForOp ivOwner; + + for (mlir::Value operand : applyOp->getOperands()) { + mlir::scf::ForOp thisOwner; + if (!isQuasiAffineIVExpression(operand, &thisOwner)) + return false; + + if (thisOwner) { + if (!ivOwner) { + ivOwner = thisOwner; + } else { + if (thisOwner != ivOwner) + return false; + } + } + } + + if (owningForOp != nullptr) + *owningForOp = ivOwner; + } + + return false; + } + // Base case: Expression is an induction variable + else if (mlir::scf::ForOp forOp = scf::getForInductionVarOwner(expr)) { + if (owningForOp != nullptr) + *owningForOp = forOp; + + return true; + } + + return false; +} + +bool isQuasiAffineIVExpression(mlir::OpFoldResult expr, + mlir::scf::ForOp *owningForOp) { + if (mlir::Value dynExpr = expr.dyn_cast()) + return isQuasiAffineIVExpression(dynExpr, owningForOp); + + return true; +} + +/// Checks if `expr` is a quasi affine expression on a single +/// induction variable, for which the increment of the induction +/// variable with the step of the associated for loop results in a +/// constant incrementation when evaluating the expression. +/// +/// E.g., this is true for the expression `i+1` for any constant step +/// size, since `((i+step)+1) - (i+1)` is constant. This is also true +/// for `(i+5)/7` for a step size that is a multiple of `7`, but false +/// for any other step size. +bool isQuasiAffineIVExpressionWithConstantStep(mlir::OpFoldResult expr, + mlir::scf::ForOp *forOp, + LoopsBoundsAndStep *basOut) { + mlir::scf::ForOp tmpForOp; + + if (isQuasiAffineIVExpression(expr, &tmpForOp)) { + std::optional bas = + getBoundsOfQuasiAffineIVExpression(expr, tmpForOp); + + if (bas.has_value()) { + if (forOp != nullptr) + *forOp = tmpForOp; + if (basOut != nullptr) + *basOut = *bas; + + return true; + } + } + + return false; +} + +static std::optional +getBoundsOfAffineExpression(mlir::AffineExpr expr, + llvm::ArrayRef dimBounds); + +// Returns the static bounds of an affine binary expression `expr` +// given the bounds `dimBounds` for any dimension expression appearing +// in the affine expression by determining the bounds for the left +// hand side and right hand side separately and applying `combinator` +// on them. +static std::optional getBoundsOfAffineBinaryExpression( + mlir::AffineExpr expr, llvm::ArrayRef dimBounds, + llvm::function_ref + combinator) { + mlir::AffineBinaryOpExpr binExpr = expr.cast(); + + std::optional lhs = + getBoundsOfAffineExpression(binExpr.getLHS(), dimBounds); + + if (!lhs.has_value()) + return std::nullopt; + + std::optional rhs = + getBoundsOfAffineExpression(binExpr.getRHS(), dimBounds); + + if (!rhs.has_value()) + return std::nullopt; + + return combinator(lhs.value(), rhs.value()); +} + +// Returns the static bounds of an affine expression given the bounds +// `dimBounds` for any dimension expression appearing in the affine +// expression. +static std::optional +getBoundsOfAffineExpression(mlir::AffineExpr expr, + llvm::ArrayRef dimBounds) { + // Cannot just use AffineExpr::compose() due to the check on + // division + switch (expr.getKind()) { + case mlir::AffineExprKind::SymbolId: + assert(false && + "Symbol found in affine expression that should not contain sumbols"); + break; + case mlir::AffineExprKind::Constant: { + int64_t cstVal = expr.cast().getValue(); + return LoopsBoundsAndStep{cstVal, cstVal, 0}; + } + case mlir::AffineExprKind::DimId: { + unsigned dimId = expr.cast().getPosition(); + assert(dimId < dimBounds.size()); + return dimBounds[dimId]; + } + case AffineExprKind::Add: + return getBoundsOfAffineBinaryExpression( + expr, dimBounds, [](LoopsBoundsAndStep lhs, LoopsBoundsAndStep rhs) { + return lhs + rhs; + }); + case AffineExprKind::Mul: + return getBoundsOfAffineBinaryExpression( + expr, dimBounds, [](LoopsBoundsAndStep lhs, LoopsBoundsAndStep rhs) { + return lhs * rhs; + }); + + case AffineExprKind::Mod: + case AffineExprKind::CeilDiv: + case AffineExprKind::FloorDiv: { + mlir::AffineBinaryOpExpr binExpr = expr.cast(); + + std::optional lhs = + getBoundsOfAffineExpression(binExpr.getLHS(), dimBounds); + std::optional rhs = + getBoundsOfAffineExpression(binExpr.getRHS(), dimBounds); + + assert(rhs->ub == rhs->lb && rhs->step == 0 && + "Expression for divisor references IV"); + int64_t rhsVal = rhs->ub; + + assert(rhsVal != 0 && "Division by zero"); + + // If the step value of the subexpression is not a multiple of + // the divisor, there may be two iterations with the same + // value. Conservatively bail out. + if (lhs->step % rhsVal != 0) + return std::nullopt; + + return *lhs / rhsVal; + } + } + + llvm_unreachable("Unknown affine expression kind"); +} + +// Returns the static bounds for the affine map `map` given the static +// bounds for all operands on which the map is applied. The map must +// not contain any symbols, all of its expressions must be pure affine +// expressions and the number of results must be one. +static std::optional +getBoundsOfAffineMap(mlir::AffineMap map, + llvm::ArrayRef mapOperandBounds) { + assert(map.getNumResults() == 1 && + "Attempting to get bounds for map with multiple result dimensions"); + assert(map.getNumSymbols() == 0 && + "Attempting to get bounds for map with symbols"); + assert(map.getResult(0).isPureAffine() && + "Attempting to get bounds for non-pure affine expression"); + + return getBoundsOfAffineExpression(map.getResult(0), mapOperandBounds); +} + +/// Returns the lower bound, upper bound and step of the quasi-affine +/// expression `expr` on the the induction variable from a for +/// operation. +std::optional +getBoundsOfQuasiAffineIVExpression(mlir::Value expr, mlir::scf::ForOp forOp) { + // Base case: expression is the induction variable itself -> check + // if the bounds are static and return them + if (forOp && expr == forOp.getInductionVar() && + isConstantIndexValue(forOp.getLowerBound()) && + isConstantIndexValue(forOp.getUpperBound()) && + isConstantIndexValue(forOp.getStep())) { + return LoopsBoundsAndStep{getConstantIndexValue(forOp.getLowerBound()), + getConstantIndexValue(forOp.getUpperBound()), + getConstantIndexValue(forOp.getStep())}; + } + // Arithmetic expression + else if (mlir::Operation *op = expr.getDefiningOp()) { + if (llvm::isa(op)) { + + std::optional lhs = + getBoundsOfQuasiAffineIVExpression(op->getOperand(0), forOp); + std::optional rhs = + getBoundsOfQuasiAffineIVExpression(op->getOperand(1), forOp); + + if (!lhs.has_value() || !rhs.has_value()) + return std::nullopt; + + if (llvm::isa(op)) + return *lhs + *rhs; + else if (llvm::isa(op)) + return *lhs - *rhs; + else if (llvm::isa(op)) + return (*lhs) * (*rhs); + else if (llvm::isa(op)) { + assert(rhs->ub == rhs->lb && rhs->step == 0 && + "Expression for divisor references IV"); + int64_t rhsVal = rhs->ub; + + assert(rhsVal != 0 && "Division by zero"); + + // If the step value of the subexpression is not a multiple of + // the divisor, there may be two iterations with the same + // value. Conservatively bail out. + if (lhs->step % rhsVal != 0) + return std::nullopt; + + return *lhs / rhsVal; + } + } + // Affine.apply + if (mlir::AffineApplyOp applyOp = llvm::dyn_cast(op)) { + if (applyOp.getMap().getNumResults() != 1 || + applyOp.getMap().getNumSymbols() != 0 || + !applyOp.getMap().getResult(0).isPureAffine()) + return std::nullopt; + + llvm::SmallVector bounds; + + for (mlir::Value operand : applyOp.getMapOperands()) { + std::optional operatorBounds = + getBoundsOfQuasiAffineIVExpression(operand, forOp); + + if (!operatorBounds.has_value()) + return std::nullopt; + + bounds.push_back(operatorBounds.value()); + } + + return getBoundsOfAffineMap(applyOp.getMap(), bounds); + } + // Base case: constant -> return constant value + else if (llvm::isa(expr.getDefiningOp())) { + mlir::arith::ConstantIndexOp cst = + llvm::dyn_cast(expr.getDefiningOp()); + return LoopsBoundsAndStep{cst.value(), cst.value(), 0}; + } + } + + return std::nullopt; +} + +std::optional +getBoundsOfQuasiAffineIVExpression(mlir::OpFoldResult expr, + mlir::scf::ForOp forOp) { + if (mlir::Value dynExpr = expr.dyn_cast()) + return getBoundsOfQuasiAffineIVExpression(dynExpr, forOp); + + mlir::IntegerAttr exprAttr = + expr.dyn_cast().dyn_cast_or_null(); + + assert(exprAttr && "Expected OpFoldResult to contain either a Value or an " + "integer attribute"); + + return LoopsBoundsAndStep{exprAttr.getInt(), exprAttr.getInt(), 0}; +} + +/// Checks if `forOp` has constant bounds and a constant step +/// resulting from quasi affine expressions. +bool isStaticLoop(mlir::scf::ForOp forOp, int64_t *ilb, int64_t *iub, + int64_t *istep) { + std::optional basLB = + getBoundsOfQuasiAffineIVExpression(forOp.getLowerBound(), nullptr); + std::optional basUB = + getBoundsOfQuasiAffineIVExpression(forOp.getUpperBound(), nullptr); + std::optional basStep = + getBoundsOfQuasiAffineIVExpression(forOp.getStep(), nullptr); + + if (!basLB.has_value() || !basUB.has_value() || !basStep.has_value()) + return false; + + if ((basLB->lb != basLB->ub || basLB->step != 0) || + (basUB->lb != basUB->ub || basUB->step != 0) || + (basStep->lb != basStep->ub || basStep->step != 0)) + return false; + + if (ilb) + *ilb = basLB->lb; + + if (iub) + *iub = basUB->lb; + + if (istep) + *istep = basStep->lb; + + return true; +} + +int64_t getStaticTripCount(int64_t lb, int64_t ub, int64_t step) { + assert((step == 0 && lb == ub) || (step >= 0 && lb <= ub) || + (step < 0 && lb > ub)); + + if (lb == ub) + return 0; + + if (lb > ub) + return getStaticTripCount(ub, lb, -step); + + assert(ub - lb < std::numeric_limits::max() - step); + + return (ub - lb + step - 1) / step; +} + +int64_t getStaticTripCount(const LoopsBoundsAndStep &bas) { + return getStaticTripCount(bas.lb, bas.ub, bas.step); +} + +// Returns the number of iterations of a static loop +int64_t getStaticTripCount(mlir::scf::ForOp forOp) { + int64_t lb; + int64_t ub; + int64_t step; + + bool isStatic = isStaticLoop(forOp, &lb, &ub, &step); + + assert(isStatic && "Loop must be static"); + + return getStaticTripCount(lb, ub, step); +} + +// Returns the total number of executions of the body of the innermost +// loop of a nest of static loops +int64_t getNestedStaticTripCount(llvm::ArrayRef nest) { + int64_t tripCount = 1; + + for (mlir::scf::ForOp forOp : nest) { + int64_t thisCount = getStaticTripCount(forOp); + + if (thisCount == 0) + return 0; + + assert(std::numeric_limits::max() / thisCount >= tripCount); + tripCount *= thisCount; + } + + return tripCount; +} + +// Checks whether `v` is a constant value of type index +bool isConstantIndexValue(mlir::Value v) { + return v.getDefiningOp() && + llvm::isa(*v.getDefiningOp()); +} + +/// Assumes that `v` is a constant index operation and returns the +/// constant value as an `int64_t`. +int64_t getConstantIndexValue(mlir::Value v) { + assert(isConstantIndexValue(v)); + + return llvm::dyn_cast(*v.getDefiningOp()) + .value(); +} + +// Checks whether `v` is a constant value of type index and its values is `i` +bool isConstantIndexValue(mlir::Value v, int64_t i) { + return isConstantIndexValue(v) && getConstantIndexValue(v) == i; +} + +} // namespace concretelang +} // namespace mlir diff --git a/compilers/concrete-compiler/compiler/lib/Transforms/Batching.cpp b/compilers/concrete-compiler/compiler/lib/Transforms/Batching.cpp index ef881205b..01bec55e5 100644 --- a/compilers/concrete-compiler/compiler/lib/Transforms/Batching.cpp +++ b/compilers/concrete-compiler/compiler/lib/Transforms/Batching.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -504,60 +505,6 @@ static bool isHoistable(mlir::Value v, mlir::scf::ForOp loop) { })); } -/// Checks if `forOp` has constant bounds and a constant step. -static bool isStaticLoop(mlir::scf::ForOp forOp, int64_t *ilb = nullptr, - int64_t *iub = nullptr, int64_t *istep = nullptr) { - - mlir::Operation *lbOp = forOp.getLowerBound().getDefiningOp(); - mlir::Operation *ubOp = forOp.getUpperBound().getDefiningOp(); - mlir::Operation *stepOp = forOp.getStep().getDefiningOp(); - - if (!lbOp || !ubOp || !stepOp) - return false; - - mlir::arith::ConstantIndexOp lb = - llvm::dyn_cast(lbOp); - mlir::arith::ConstantIndexOp ub = - llvm::dyn_cast(ubOp); - mlir::arith::ConstantIndexOp step = - llvm::dyn_cast(stepOp); - - if (lb && ub && step) { - if (ilb) - *ilb = lb.value(); - - if (iub) - *iub = ub.value(); - - if (istep) - *istep = step.value(); - - return true; - } - - return false; -} - -// Checks whether `v` is a constant value of type index -static bool isConstantIndexValue(mlir::Value v) { - return v.getDefiningOp() && - llvm::isa(*v.getDefiningOp()); -} - -/// Assumes that `v` is a constant index operation and returns the -/// constant value as an `int64_t`. -static int64_t getConstantIndexValue(mlir::Value v) { - assert(isConstantIndexValue(v)); - - return llvm::dyn_cast(*v.getDefiningOp()) - .value(); -} - -// Checks whether `v` is a constant value of type index and its values is `i` -static bool isConstantIndexValue(mlir::Value v, int64_t i) { - return isConstantIndexValue(v) && getConstantIndexValue(v) == i; -} - llvm::SmallVector buildNormalizedIndexes(mlir::PatternRewriter &rewriter, llvm::ArrayRef nest) { @@ -679,388 +626,6 @@ mlir::OpFoldResult opFoldExpr(mlir::ImplicitLocOpBuilder &builder, } } -/// Convenience class that holds all parameters of a loop -struct BoundsAndStep { - int64_t lb; - int64_t ub; - int64_t step; - - BoundsAndStep operator+(const BoundsAndStep &other) { - return BoundsAndStep{lb + other.lb, ub + other.ub, step + other.step}; - } - BoundsAndStep operator-(const BoundsAndStep &other) { - return BoundsAndStep{lb - other.lb, ub - other.ub, step - other.step}; - } - BoundsAndStep operator*(const BoundsAndStep &other) { - return BoundsAndStep{lb * other.lb, ub * other.ub, step * other.step}; - } - BoundsAndStep operator/(int64_t d) { - return BoundsAndStep{lb / d, ub / d, step / d}; - } -}; - -static std::optional -getBoundsOfAffineExpression(mlir::AffineExpr expr, - llvm::ArrayRef dimBounds); - -// Returns the static bounds of an affine binary expression `expr` -// given the bounds `dimBounds` for any dimension expression appearing -// in the affine expression by determining the bounds for the left -// hand side and right hand side separately and applying `combinator` -// on them. -static std::optional getBoundsOfAffineBinaryExpression( - mlir::AffineExpr expr, llvm::ArrayRef dimBounds, - llvm::function_ref - combinator) { - mlir::AffineBinaryOpExpr binExpr = expr.cast(); - - std::optional lhs = - getBoundsOfAffineExpression(binExpr.getLHS(), dimBounds); - - if (!lhs.has_value()) - return std::nullopt; - - std::optional rhs = - getBoundsOfAffineExpression(binExpr.getRHS(), dimBounds); - - if (!rhs.has_value()) - return std::nullopt; - - return combinator(lhs.value(), rhs.value()); -} - -// Returns the static bounds of an affine expression given the bounds -// `dimBounds` for any dimension expression appearing in the affine -// expression. -static std::optional -getBoundsOfAffineExpression(mlir::AffineExpr expr, - llvm::ArrayRef dimBounds) { - // Cannot just use AffineExpr::compose() due to the check on - // division - switch (expr.getKind()) { - case mlir::AffineExprKind::SymbolId: - assert(false && - "Symbol found in affine expression that should not contain sumbols"); - break; - case mlir::AffineExprKind::Constant: { - int64_t cstVal = expr.cast().getValue(); - return BoundsAndStep{cstVal, cstVal, 0}; - } - case mlir::AffineExprKind::DimId: { - unsigned dimId = expr.cast().getPosition(); - assert(dimId < dimBounds.size()); - return dimBounds[dimId]; - } - case AffineExprKind::Add: - return getBoundsOfAffineBinaryExpression( - expr, dimBounds, - [](BoundsAndStep lhs, BoundsAndStep rhs) { return lhs + rhs; }); - case AffineExprKind::Mul: - return getBoundsOfAffineBinaryExpression( - expr, dimBounds, - [](BoundsAndStep lhs, BoundsAndStep rhs) { return lhs * rhs; }); - - case AffineExprKind::Mod: - case AffineExprKind::CeilDiv: - case AffineExprKind::FloorDiv: { - mlir::AffineBinaryOpExpr binExpr = expr.cast(); - - std::optional lhs = - getBoundsOfAffineExpression(binExpr.getLHS(), dimBounds); - std::optional rhs = - getBoundsOfAffineExpression(binExpr.getRHS(), dimBounds); - - assert(rhs->ub == rhs->lb && rhs->step == 0 && - "Expression for divisor references IV"); - int64_t rhsVal = rhs->ub; - - assert(rhsVal != 0 && "Division by zero"); - - // If the step value of the subexpression is not a multiple of - // the divisor, there may be two iterations with the same - // value. Conservatively bail out. - if (lhs->step % rhsVal != 0) - return std::nullopt; - - return *lhs / rhsVal; - } - } - - llvm_unreachable("Unknown affine expression kind"); -} - -// Returns the static bounds for the affine map `map` given the static -// bounds for all operands on which the map is applied. The map must -// not contain any symbols, all of its expressions must be pure affine -// expressions and the number of results must be one. -static std::optional -getBoundsOfAffineMap(mlir::AffineMap map, - llvm::ArrayRef mapOperandBounds) { - assert(map.getNumResults() == 1 && - "Attempting to get bounds for map with multiple result dimensions"); - assert(map.getNumSymbols() == 0 && - "Attempting to get bounds for map with symbols"); - assert(map.getResult(0).isPureAffine() && - "Attempting to get bounds for non-pure affine expression"); - - return getBoundsOfAffineExpression(map.getResult(0), mapOperandBounds); -} - -/// Returns the lower bound, upper bound and step of the quasi-affine -/// expression `expr` on the the induction variable from a for -/// operation. -static std::optional -getBoundsOfQuasiAffineIVExpression(mlir::Value expr, mlir::scf::ForOp forOp) { - // Base case: expression is the induction variable itself -> check - // if the bounds are static and return them - if (forOp && expr == forOp.getInductionVar() && - isConstantIndexValue(forOp.getLowerBound()) && - isConstantIndexValue(forOp.getUpperBound()) && - isConstantIndexValue(forOp.getStep())) { - return BoundsAndStep{getConstantIndexValue(forOp.getLowerBound()), - getConstantIndexValue(forOp.getUpperBound()), - getConstantIndexValue(forOp.getStep())}; - } - // Arithmetic expression - else if (mlir::Operation *op = expr.getDefiningOp()) { - if (llvm::isa(op)) { - - std::optional lhs = - getBoundsOfQuasiAffineIVExpression(op->getOperand(0), forOp); - std::optional rhs = - getBoundsOfQuasiAffineIVExpression(op->getOperand(1), forOp); - - if (!lhs.has_value() || !rhs.has_value()) - return std::nullopt; - - if (llvm::isa(op)) - return *lhs + *rhs; - else if (llvm::isa(op)) - return *lhs - *rhs; - else if (llvm::isa(op)) - return (*lhs) * (*rhs); - else if (llvm::isa(op)) { - assert(rhs->ub == rhs->lb && rhs->step == 0 && - "Expression for divisor references IV"); - int64_t rhsVal = rhs->ub; - - assert(rhsVal != 0 && "Division by zero"); - - // If the step value of the subexpression is not a multiple of - // the divisor, there may be two iterations with the same - // value. Conservatively bail out. - if (lhs->step % rhsVal != 0) - return std::nullopt; - - return *lhs / rhsVal; - } - } - // Affine.apply - if (mlir::AffineApplyOp applyOp = llvm::dyn_cast(op)) { - if (applyOp.getMap().getNumResults() != 1 || - applyOp.getMap().getNumSymbols() != 0 || - !applyOp.getMap().getResult(0).isPureAffine()) - return std::nullopt; - - llvm::SmallVector bounds; - - for (mlir::Value operand : applyOp.getMapOperands()) { - std::optional operatorBounds = - getBoundsOfQuasiAffineIVExpression(operand, forOp); - - if (!operatorBounds.has_value()) - return std::nullopt; - - bounds.push_back(operatorBounds.value()); - } - - return getBoundsOfAffineMap(applyOp.getMap(), bounds); - } - // Base case: constant -> return constant value - else if (llvm::isa(expr.getDefiningOp())) { - mlir::arith::ConstantIndexOp cst = - llvm::dyn_cast(expr.getDefiningOp()); - return BoundsAndStep{cst.value(), cst.value(), 0}; - } - } - - return std::nullopt; -} - -static std::optional -getBoundsOfQuasiAffineIVExpression(mlir::OpFoldResult expr, - mlir::scf::ForOp forOp) { - if (mlir::Value dynExpr = expr.dyn_cast()) - return getBoundsOfQuasiAffineIVExpression(dynExpr, forOp); - - mlir::IntegerAttr exprAttr = - expr.dyn_cast().dyn_cast_or_null(); - - assert(exprAttr && "Expected OpFoldResult to contain either a Value or an " - "integer attribute"); - - return BoundsAndStep{exprAttr.getInt(), exprAttr.getInt(), 0}; -} - -static int64_t getStaticTripCount(int64_t lb, int64_t ub, int64_t step) { - assert(ub > lb && "Upper bound must be greater than lower bound"); - assert(step > 0 && "Step must be positive"); - assert(ub - lb < std::numeric_limits::max() - step); - - return (ub - lb + step - 1) / step; -} - -static int64_t getStaticTripCount(const BoundsAndStep &bas) { - return getStaticTripCount(bas.lb, bas.ub, bas.step); -} - -// Returns the number of iterations of a static loop -static int64_t getStaticTripCount(mlir::scf::ForOp forOp) { - int64_t lb; - int64_t ub; - int64_t step; - - bool isStatic = isStaticLoop(forOp, &lb, &ub, &step); - - assert(isStatic && "Loop must be static"); - - return getStaticTripCount(lb, ub, step); -} - -// Returns the total number of executions of the body of the innermost -// loop of a nest of static loops -static int64_t getNestedStaticTripCount(llvm::ArrayRef nest) { - int64_t tripCount = 1; - - for (mlir::scf::ForOp forOp : nest) { - int64_t thisCount = getStaticTripCount(forOp); - - if (thisCount == 0) - return 0; - - assert(std::numeric_limits::max() / thisCount >= tripCount); - tripCount *= thisCount; - } - - return tripCount; -} - -/// Checks whether the expression `expr` is a quasi-affine expression -/// on a single induction variable. If an induction variable is -/// referenced, the owning for loop is returned in `*owningForOp`. -static bool isQuasiAffineIVExpression(mlir::Value expr, - mlir::scf::ForOp *owningForOp = nullptr) { - if (mlir::Operation *op = expr.getDefiningOp()) { - if (llvm::isa(op)) { - return true; - } else if (llvm::isa(op)) { - mlir::scf::ForOp forLHS; - mlir::scf::ForOp forRHS; - - if (!isQuasiAffineIVExpression(op->getOperand(0), &forLHS) || - !isQuasiAffineIVExpression(op->getOperand(1), &forRHS)) { - return false; - } else { - // Check that appearances of IVs refer to the same IV - if (forLHS && forRHS && forLHS != forRHS) - return false; - } - - // Assume that the expression is already canonicalized, so that - // IVs appear only in numerators and on one side of a - // multiplication subexpression - if ((llvm::isa(op) && forLHS && forRHS) || - (llvm::isa(op) && forRHS)) - return false; - - if (owningForOp != nullptr) { - if (forLHS) - *owningForOp = forLHS; - else if (forRHS) - *owningForOp = forRHS; - } - - return true; - } else if (mlir::AffineApplyOp applyOp = - llvm::dyn_cast(op)) { - // Affine.apply: make sure that all operands are either constant - // expressions or using IVs of the same loop - mlir::scf::ForOp ivOwner; - - for (mlir::Value operand : applyOp->getOperands()) { - mlir::scf::ForOp thisOwner; - if (!isQuasiAffineIVExpression(operand, &thisOwner)) - return false; - - if (thisOwner) { - if (!ivOwner) { - ivOwner = thisOwner; - } else { - if (thisOwner != ivOwner) - return false; - } - } - } - - if (owningForOp != nullptr) - *owningForOp = ivOwner; - } - - return false; - } - // Base case: Expression is an induction variable - else if (mlir::scf::ForOp forOp = scf::getForInductionVarOwner(expr)) { - if (owningForOp != nullptr) - *owningForOp = forOp; - - return true; - } - - return false; -} - -static bool isQuasiAffineIVExpression(mlir::OpFoldResult expr, - mlir::scf::ForOp *owningForOp = nullptr) { - if (mlir::Value dynExpr = expr.dyn_cast()) - return isQuasiAffineIVExpression(dynExpr, owningForOp); - - return true; -} - -/// Checks if `expr` is a quasi affine expression on a single -/// induction variable, for which the increment of the induction -/// variable with the step of the associated for loop results in a -/// constant incrementation of when evaluating the expression. -/// -/// E.g., this is true for the expression `i+1` for any constant step -/// size, since `((i+step)+1) - (i+1)` is constant. This is also true -/// for `(i+5)/7` for a step size that is a multiple of `7`, but false -/// for any other step size. -static bool -isQuasiAffineIVExpressionWithConstantStep(mlir::OpFoldResult expr, - mlir::scf::ForOp *forOp = nullptr, - BoundsAndStep *basOut = nullptr) { - mlir::scf::ForOp tmpForOp; - - if (isQuasiAffineIVExpression(expr, &tmpForOp)) { - std::optional bas = - getBoundsOfQuasiAffineIVExpression(expr, tmpForOp); - - if (bas.has_value()) { - if (forOp != nullptr) - *forOp = tmpForOp; - if (basOut != nullptr) - *basOut = *bas; - - return true; - } - } - - return false; -} - /// Hoists the pure operation producing the value `v` out of /// `outermostFor` recursively. All newly created mappings are /// collected in `mapping`. @@ -1141,7 +706,7 @@ mlir::Value hoistIndexedOp( if (isAffine && forOp && (forOp == outermostFor || outermostFor->isAncestor(forOp))) { - std::optional bas = + std::optional bas = getBoundsOfQuasiAffineIVExpression(idxExpr, forOp); assert(bas.has_value()); @@ -1873,7 +1438,7 @@ getLoopsForCandidateIndexes(IndexedOpTy op) { } mlir::scf::ForOp qaLoop; - BoundsAndStep bas; + LoopsBoundsAndStep bas; if (isQuasiAffineIVExpressionWithConstantStep(expr, &qaLoop, &bas)) { if (qaLoop) { @@ -2206,7 +1771,7 @@ public: mlir::Operation *targetOp = nullptr; llvm::SmallVector nest; llvm::SmallVector idxMap; - llvm::SmallVector idxBounds; + llvm::SmallVector idxBounds; mlir::arith::ConstantOp cdo; mlir::RankedTensorType constantType; mlir::Type origElementType; @@ -2242,14 +1807,14 @@ public: llvm::DenseSet nestUnsorted; llvm::SmallVector currIdxMap; - llvm::SmallVector currIdxBounds; + llvm::SmallVector currIdxBounds; // Make sure that the extract operation uses only quasi affine // expressions on IVs, where each index uses at most a single // IV. for (mlir::Value idx : currExtractOp.getIndices()) { mlir::scf::ForOp forOp; - BoundsAndStep bas; + LoopsBoundsAndStep bas; if (!isQuasiAffineIVExpressionWithConstantStep(idx, &forOp, &bas)) return mlir::WalkResult::skip(); @@ -2304,7 +1869,7 @@ public: // there are no out-of-bounds accesses. for (auto it : llvm::enumerate(currExtractOp.getIndices())) { mlir::scf::ForOp forOp; - BoundsAndStep bas; + LoopsBoundsAndStep bas; mlir::Value idx = it.value(); size_t i = it.index(); @@ -2360,7 +1925,7 @@ public: // current index mlir::SmallVector idx = - map(idxBounds, [](BoundsAndStep &bas) { return bas.lb; }); + map(idxBounds, [](LoopsBoundsAndStep &bas) { return bas.lb; }); // Maps the index of each IV in the loop nest to the indexes of // the extract operation