Instead of relying on adhoc bounds calculations, use a projection-based implementation. This simplifies the implementation and finds more static constant sizes than previously/ Differential Revision: https://reviews.llvm.org/D106054
88 lines
3.0 KiB
C++
88 lines
3.0 KiB
C++
//===- 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<SmallVector<int64_t, 8>> 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<int64_t, 8> 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();
|
|
}
|