From d6c93c16a69ab900b53c2eea37d58e812712f9a6 Mon Sep 17 00:00:00 2001 From: Jeremy Kun Date: Mon, 4 Aug 2025 23:03:03 -0700 Subject: [PATCH 1/2] add iterVarKind for convenient iterating over variables --- .../mlir/Analysis/Presburger/IntegerRelation.h | 7 +++++++ .../Analysis/Presburger/IntegerRelationTest.cpp | 12 ++++++++++++ 2 files changed, 19 insertions(+) diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h index ee401cca8f552..fb895b1a6fb78 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,12 @@ class IntegerRelation { return space.getVarKindEnd(kind); } + /// Return an interator over the variables of the specified kind + /// starting at the relevant offset. + 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)); +} From f0d84e7d455f6d4f74571ad20ba7a8f10b4496aa Mon Sep 17 00:00:00 2001 From: Jeremy Kun Date: Tue, 5 Aug 2025 09:49:48 -0700 Subject: [PATCH 2/2] address comments --- mlir/include/mlir/Analysis/Presburger/IntegerRelation.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h index fb895b1a6fb78..4b1802413f75f 100644 --- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h @@ -270,7 +270,8 @@ class IntegerRelation { } /// Return an interator over the variables of the specified kind - /// starting at the relevant offset. + /// 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)); }