[MLIR][Presburger] add iterVarKind for convenient iterating over variables (#152091)
I find myself doing this alot
```
for (unsigned varIndex = rel.getVarKindOffset(VarKind::Domain);
varIndex < rel.getVarKindEnd(VarKind::Domain); ++varIndex) {
...
}
```
Adding this convenience method so I can instead do
```
for (unsigned varIndex : rel.iterVarKind(VarKind::Domain)) {
...
}
```
---------
Co-authored-by: Jeremy Kun <j2kun@users.noreply.github.com>
This commit is contained in:
parent
d478502a42
commit
97dee32445
@ -20,6 +20,7 @@
|
||||
#include "mlir/Analysis/Presburger/PresburgerSpace.h"
|
||||
#include "mlir/Analysis/Presburger/Utils.h"
|
||||
#include "llvm/ADT/DynamicAPInt.h"
|
||||
#include "llvm/ADT/Sequence.h"
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
#include "llvm/Support/LogicalResult.h"
|
||||
#include <optional>
|
||||
@ -268,6 +269,13 @@ public:
|
||||
return space.getVarKindEnd(kind);
|
||||
}
|
||||
|
||||
/// Return an interator over the variables of the specified kind
|
||||
/// starting at the relevant offset. The return type is auto in
|
||||
/// keeping with the convention for iterators.
|
||||
auto iterVarKind(VarKind kind) {
|
||||
return llvm::seq(getVarKindOffset(kind), getVarKindEnd(kind));
|
||||
}
|
||||
|
||||
/// Get the number of elements of the specified kind in the range
|
||||
/// [varStart, varLimit).
|
||||
unsigned getVarKindOverlap(VarKind kind, unsigned varStart,
|
||||
|
||||
@ -16,6 +16,7 @@
|
||||
|
||||
using namespace mlir;
|
||||
using namespace presburger;
|
||||
using ::testing::ElementsAre;
|
||||
|
||||
TEST(IntegerRelationTest, getDomainAndRangeSet) {
|
||||
IntegerRelation rel = parseRelationFromSet(
|
||||
@ -702,3 +703,14 @@ TEST(IntegerRelationTest, rangeProductSymbols) {
|
||||
|
||||
EXPECT_TRUE(expected.isEqual(rangeProd));
|
||||
}
|
||||
|
||||
TEST(IntegerRelationTest, getVarKindRange) {
|
||||
IntegerRelation r1 = parseRelationFromSet(
|
||||
"(i1, i2, i3, i4, i5) : (i1 >= 0, i2 >= 0, i3 >= 0, i4 >= 0, i5 >= 0)",
|
||||
2);
|
||||
SmallVector<unsigned> actual;
|
||||
for (unsigned var : r1.iterVarKind(VarKind::Range)) {
|
||||
actual.push_back(var);
|
||||
}
|
||||
EXPECT_THAT(actual, ElementsAre(2, 3, 4));
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user