diff --git a/mlir/include/mlir/Dialect/Linalg/Analysis/ConstraintsSet.h b/mlir/include/mlir/Dialect/Linalg/Analysis/ConstraintsSet.h deleted file mode 100644 --- a/mlir/include/mlir/Dialect/Linalg/Analysis/ConstraintsSet.h +++ /dev/null @@ -1,67 +0,0 @@ -//===- ConstraintsSet.h - Ext. for FlatAffineValueConstraints ---*- C++ -*-===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// -// -// Linalg-specific constraints set extensions. -// -//===----------------------------------------------------------------------===// - -#ifndef MLIR_DIALECT_LINALG_ANALYSIS_CONSTRAINTS_SET_H_ -#define MLIR_DIALECT_LINALG_ANALYSIS_CONSTRAINTS_SET_H_ - -#include "mlir/Analysis/AffineStructures.h" -#include "mlir/IR/AffineMap.h" - -namespace mlir { -class ValueRange; - -/// Linalg-specific constraints set extensions. -class ConstraintsSet : public FlatAffineValueConstraints { -public: - ConstraintsSet() : FlatAffineValueConstraints() {} - - /// Assuming `val` is defined by `val = affine.min map (operands)`, introduce - /// all the constraints `val <= expr_i(operands)`, where expr_i are all the - /// results of `map`. - // This API avoids taking a dependence on the AffineMinOp definition. - LogicalResult composeMin(Value val, AffineMap map, ValueRange operands) { - return composeMinOrMaxMapAndOperands(val, map, operands, /*min=*/true); - } - - /// Assuming `val` is defined by `val = affine.max map (operands)`, introduce - /// all the constraints `val >= expr_i(operands)`, where expr_i are all the - /// results of `map`. - // This API avoids taking a dependence on the AffineMaxOp definition. - LogicalResult composeMax(Value val, AffineMap map, ValueRange operands) { - return composeMinOrMaxMapAndOperands(val, map, operands, /*min=*/false); - } - - /// Assuming `val` is defined by `val = affine.apply map (operands)`, call - /// composeMap. - // This API avoids taking a dependence on the AffineMApplyOp definition. - LogicalResult composeAffineApply(Value val, AffineMap map, - ValueRange operands); - - /// Asserts the identifier `id` is in the constraints set and returns it. - unsigned lookupPos(Value id) const; - - /// If v is not in the constraint set, insert it as a dim or symbol depending - /// on `asDim`. - /// Return success if v is of dim id type when `asDim` is true and of symbol - /// id type when `asDim` is false. - /// Return failure otherwise. - LogicalResult ensureIdOfType(Value v, bool asDim); - -private: - /// Implementation detail for composeMin/Max. - LogicalResult composeMinOrMaxMapAndOperands(Value val, AffineMap map, - ValueRange operands, bool min); -}; - -} // namespace mlir - -#endif // MLIR_DIALECT_LINALG_ANALYSIS_CONSTRAINTS_SET_H_ diff --git a/mlir/lib/Dialect/Linalg/Analysis/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Analysis/CMakeLists.txt --- a/mlir/lib/Dialect/Linalg/Analysis/CMakeLists.txt +++ b/mlir/lib/Dialect/Linalg/Analysis/CMakeLists.txt @@ -1,5 +1,4 @@ add_mlir_dialect_library(MLIRLinalgAnalysis - ConstraintsSet.cpp DependenceAnalysis.cpp ADDITIONAL_HEADER_DIRS diff --git a/mlir/lib/Dialect/Linalg/Analysis/ConstraintsSet.cpp b/mlir/lib/Dialect/Linalg/Analysis/ConstraintsSet.cpp deleted file mode 100644 --- a/mlir/lib/Dialect/Linalg/Analysis/ConstraintsSet.cpp +++ /dev/null @@ -1,87 +0,0 @@ -//===- ConstraintsSet.cpp - Extensions for FlatAffineConstraints ----------===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// -// -// Linalg-specific constraints set extensions. -// -//===----------------------------------------------------------------------===// - -#include "mlir/Dialect/Linalg/Analysis/ConstraintsSet.h" -#include "mlir/Dialect/Affine/IR/AffineValueMap.h" -#include "mlir/IR/AffineMap.h" - -using namespace mlir; - -unsigned ConstraintsSet::lookupPos(Value id) const { - unsigned pos; - if (!findId(id, &pos)) { - llvm::errs() << "Lookup failed: " << id << "\n"; - llvm_unreachable("Lookup failed"); - } - return pos; -} - -LogicalResult ConstraintsSet::ensureIdOfType(Value v, bool asDim) { - if (!containsId(v)) { - if (asDim) - addDimId(getNumDimIds(), v); - else - addSymbolId(getNumSymbolIds(), v); - return success(); - } - unsigned pos = lookupPos(v); - return success((asDim && pos < getNumDimIds()) || - (!asDim && getNumDimIds() <= pos && - pos < getNumDimIds() + getNumSymbolIds())); -} - -LogicalResult ConstraintsSet::composeAffineApply(Value val, AffineMap map, - ValueRange operands) { - AffineValueMap avm(map, operands, val); - return composeMap(&avm); -} - -LogicalResult ConstraintsSet::composeMinOrMaxMapAndOperands(Value val, - AffineMap map, - ValueRange operands, - bool min) { - ConstraintsSet localCst; - std::vector> flatExprs; - if (failed(getFlattenedAffineExprs(map, &flatExprs, &localCst))) - return failure(); - assert(flatExprs.size() == map.getNumResults() && - "incorrect number of flattened expressiosn"); - - // Local vars on a per-need basis. - if (localCst.getNumLocalIds() != 0) - return failure(); - - // Add one inequality for each result connecting `val` to the other ids in - // `operands`. For instance, uf the expression is: - // `16 * i0 + i1` and - // `min` is true - // add: - // -d_val + 16 * i0 + i1 >= 0. - for (const auto &flatExpr : flatExprs) { - assert(flatExpr.size() >= operands.size() + 1); - SmallVector ineq(getNumCols(), 0); - for (unsigned i = 0, e = operands.size(); i < e; i++) - ineq[lookupPos(operands[i])] = min ? flatExpr[i] : -flatExpr[i]; - - // Set the coefficient for `d_val`. - ineq[lookupPos(val)] = min ? -1 : 1; - - // Set the constant term (upper bound in flatExpr is exclusive). - ineq[getNumCols() - 1] = min ? flatExpr[flatExpr.size() - 1] - 1 - : -flatExpr[flatExpr.size() - 1]; - - // Add the inequality connecting the result of the map to the rest. - addInequality(ineq); - } - - return success(); -} diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp @@ -14,8 +14,8 @@ #include "mlir/Dialect/Linalg/Transforms/Hoisting.h" #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/SliceAnalysis.h" +#include "mlir/Dialect/Affine/IR/AffineValueMap.h" #include "mlir/Dialect/Affine/Utils.h" -#include "mlir/Dialect/Linalg/Analysis/ConstraintsSet.h" #include "mlir/Dialect/Linalg/IR/LinalgOps.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/Dialect/SCF/SCF.h" @@ -555,8 +555,9 @@ /// Given a set of loops, assumed to be scf::ForOp, create a constraint set /// containing the inequalities `iv - lb >= 0` and `-iv + ub - 1 >= 0` for each /// loop. -static ConstraintsSet initLoopIvsAndBounds(ArrayRef loops) { - ConstraintsSet constraints; +static FlatAffineValueConstraints +initLoopIvsAndBounds(ArrayRef loops) { + FlatAffineValueConstraints constraints; for (Operation *op : loops) constraints.addDimId(constraints.getNumDimIds(), cast(op).getInductionVar()); @@ -598,7 +599,7 @@ /// If any other operation is met, return failure. // TODO: extend on a per-need basis. static LogicalResult -foldUpperBoundsIntoConstraintsSet(ConstraintsSet &constraints, +foldUpperBoundsIntoConstraintsSet(FlatAffineValueConstraints &constraints, scf::ForOp outerLimit, ArrayRef loops) { SetVector toProjectOut; @@ -622,7 +623,13 @@ continue; // Ensure there is a auto ensureIdFailed = [&](Value v) { - return failed(constraints.ensureIdOfType(v, /*asDim=*/true)); + if (constraints.containsId(v)) { + unsigned pos; + constraints.findId(v, &pos); + return pos >= constraints.getNumDimIds(); + } + constraints.addDimId(constraints.getNumDimIds(), v); + return false; }; // Ensure all ids exist and add results for later projection. @@ -636,16 +643,22 @@ // Compose supported ops. if (auto affineApplyOp = dyn_cast(op)) { - if (failed(constraints.composeAffineApply(affineApplyOp.getResult(), - affineApplyOp.getAffineMap(), - affineApplyOp.getOperands()))) + AffineValueMap avm(affineApplyOp.getAffineMap(), + affineApplyOp.getOperands(), + affineApplyOp.getResult()); + if (failed(constraints.composeMap(&avm))) return failure(); continue; } auto affineMinOp = cast(op); - if (failed(constraints.composeMin(affineMinOp.getResult(), - affineMinOp.getAffineMap(), - affineMinOp.operands()))) + unsigned pos; + bool foundMinOp = constraints.findId(affineMinOp.getResult(), &pos); + (void)foundMinOp; + assert(foundMinOp); + AffineMap alignedMap = constraints.computeAlignedMap( + affineMinOp.getAffineMap(), affineMinOp.getOperands()); + if (failed( + constraints.addBound(FlatAffineConstraints::UB, pos, alignedMap))) return failure(); } } @@ -654,25 +667,6 @@ return success(); } -/// Compute dynamic tensor sizes, independent of any value defined inside -/// `outer` and such that every n-D iteration of the packingLoops has its own -/// space (so that each packed buffer has a storage location). This is achieved -/// by computing the extent for each of the packing loops. -static LogicalResult computeBounds(scf::ForOp outer, - ArrayRef packingLoops, - SmallVector &lbs, - SmallVector &ubs) { - // Packing loop IVs are introduced as the first positions. - ConstraintsSet constraints = initLoopIvsAndBounds(packingLoops); - if (failed( - foldUpperBoundsIntoConstraintsSet(constraints, outer, packingLoops))) - return failure(); - // Compute the bounds of the first positions, assuming the others are fixed. - constraints.getSliceBounds(/*pos=*/0, /*num=*/packingLoops.size(), - outer->getContext(), &lbs, &ubs); - return success(); -} - /// Ensure prerequisites that guarantee pad op hoisting can occur. /// Return failure in the cases when we cannot perform hoisting; i.e. if either: /// 1. There exists a use of `padTensorOp` that is not a linalg input operand. @@ -758,15 +752,17 @@ scf::ForOp outer = cast(outermostEnclosingForOp); - ConstraintsSet constraints = initLoopIvsAndBounds(packingLoops.getArrayRef()); + FlatAffineValueConstraints constraints = + initLoopIvsAndBounds(packingLoops.getArrayRef()); if (failed(foldUpperBoundsIntoConstraintsSet(constraints, outer, packingLoops.getArrayRef()))) return failure(); unsigned numLoops = packingLoops.size(); SmallVector lbs(numLoops), ubs(numLoops); - if (failed(computeBounds(outer, packingLoops.getArrayRef(), lbs, ubs))) - return failure(); + // Compute the bounds of the first positions, assuming the others are fixed. + constraints.getSliceBounds(/*pos=*/0, /*num=*/packingLoops.size(), + outer->getContext(), &lbs, &ubs); SmallVector allValues; constraints.getAllValues(&allValues); diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -6085,12 +6085,10 @@ "lib/Dialect/Linalg/Transforms/*.cpp", "lib/Dialect/Linalg/Transforms/*.h", ]) + [ - "lib/Dialect/Linalg/Analysis/ConstraintsSet.cpp", "lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp", "lib/Dialect/Linalg/Utils/Utils.cpp", ], hdrs = [ - "include/mlir/Dialect/Linalg/Analysis/ConstraintsSet.h", "include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h", "include/mlir/Dialect/Linalg/Passes.h", "include/mlir/Dialect/Linalg/Transforms/CodegenStrategy.h",