[mlir] Add helpers to set lbs, ubs, steps for linalg.tiled_loop.

Differential Revision: https://reviews.llvm.org/D100579
This commit is contained in:
Alexander Belyaev 2021-04-15 19:48:38 +02:00
parent 94662ee0c1
commit 3c982a6e54
1 changed files with 25 additions and 0 deletions

View File

@ -591,6 +591,31 @@ def Linalg_TiledLoopOp : Linalg_Op<"tiled_loop", [
return getBody()->getArguments();
}
void setLowerBounds(ValueRange lowerBounds) {
unsigned numLoops = getNumLoops();
assert(lowerBounds.size() == numLoops &&
"expected lower bounds for every loop dimension");
for (unsigned i = 0; i < numLoops; ++i)
setOperand(i, lowerBounds[i]);
}
void setUpperBounds(ValueRange upperBounds) {
unsigned numLoops = getNumLoops();
assert(upperBounds.size() == numLoops &&
"expected upper bounds for every loop dimension");
for (unsigned i = 0, pos = numLoops; i < numLoops; ++i, ++pos)
setOperand(pos, upperBounds[i]);
}
void setSteps(ValueRange steps) {
unsigned numLoops = getNumLoops();
assert(steps.size() == numLoops &&
"expected upper bounds for every loop dimension");
for (unsigned i = 0, pos = 2 * numLoops; i < numLoops; ++i, ++pos)
setOperand(pos, steps[i]);
}
/// Result that corresponds to the `outputs` argument of tensor type.
OpResult getTiedOpResult(OpOperand& opOperand) {
// No result can correspond to a memref argument.