diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h index ee401cca8f552..4b1802413f75f 100644 --- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h @@ -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 @@ -268,6 +269,13 @@ class IntegerRelation { 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, diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp index dd0b09f7f05d2..a6ed5c5b21e79 100644 --- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp +++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp @@ -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 actual; + for (unsigned var : r1.iterVarKind(VarKind::Range)) { + actual.push_back(var); + } + EXPECT_THAT(actual, ElementsAre(2, 3, 4)); +}