return option induction vars

llvm · srcarroll · Jul 3, 2024 · Jun 5, 2024 · Jun 4, 2024 · Jun 5, 2024
commit e0e526210a5ab6ce28fbc5fa5ee24f79cb1ee9a8
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -511,7 +511,9 @@ def ForallOp : SCF_Op<"forall", [
 
   let extraClassDeclaration = [{
     SmallVector<Value> getInductionVars() {
-      return getLoopInductionVars();
+      auto maybeInductionVars = getLoopInductionVars();;
+      assert(maybeInductionVars.has_value() && "expected values");
+      return *maybeInductionVars;
     }
     // Get lower bounds as OpFoldResult.
     SmallVector<OpFoldResult> getMixedLowerBound() {
@@ -591,7 +593,7 @@ def ForallOp : SCF_Op<"forall", [
     }
 
     ::mlir::Value getInductionVar(int64_t idx) {
-      return getLoopInductionVars()[idx];
+      return getInductionVars()[idx];
     }
 
     ::mlir::Block::BlockArgListType getRegionOutArgs() {
@@ -850,7 +852,9 @@ def ParallelOp : SCF_Op<"parallel",
   let extraClassDeclaration = [{
     // Get induction variables.
     SmallVector<Value> getInductionVars() {
-      return getLoopInductionVars();
+      auto maybeInductionVars = getLoopInductionVars();;
+      assert(maybeInductionVars.has_value() && "expected values");
+      return *maybeInductionVars;
     }
     unsigned getNumLoops() { return getStep().size(); }
     unsigned getNumReductions() { return getInitVals().size(); }

diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -95,12 +95,12 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
     InterfaceMethod<[{
         Return all induction variables.
       }],
-      /*retTy=*/"::llvm::SmallVector<::mlir::Value>",
+      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::Value>>",
       /*methodName=*/"getLoopInductionVars",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return {};
+        return std::nullopt;
       }]
     >,
     InterfaceMethod<[{
@@ -235,8 +235,8 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
     /// std::nullopt.
     ::std::optional<::mlir::Value> getSingleInductionVar() {
       auto inductionVars = this->getLoopInductionVars();
-      if (inductionVars.size() == 1)
-          return inductionVars[0];
+      if (inductionVars.has_value() && (*inductionVars).size() == 1)
+          return (*inductionVars)[0];
         return std::nullopt;
     }
     /// Return the single lower bound value or attribute if it exists, otherwise

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2454,8 +2454,8 @@ bool AffineForOp::matchingBoundOperandList() {
 
 SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
 
-SmallVector<Value> AffineForOp::getLoopInductionVars() {
-  return {getInductionVar()};
+std::optional<SmallVector<Value>> AffineForOp::getLoopInductionVars() {
+  return SmallVector<Value>{getInductionVar()};
 }
 
 std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopLowerBounds() {

diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -378,7 +378,9 @@ LogicalResult ForOp::verifyRegions() {
   return success();
 }
 
-SmallVector<Value> ForOp::getLoopInductionVars() { return {getInductionVar()}; }
+std::optional<SmallVector<Value>> ForOp::getLoopInductionVars() {
+  return SmallVector<Value, 1>{getInductionVar()};
+}
 
 std::optional<SmallVector<OpFoldResult>> ForOp::getLoopLowerBounds() {
   return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
@@ -1426,7 +1428,7 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
   return storeOps;
 }
 
-SmallVector<Value> ForallOp::getLoopInductionVars() {
+std::optional<SmallVector<Value>> ForallOp::getLoopInductionVars() {
   return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
 }
 
@@ -3004,7 +3006,7 @@ void ParallelOp::print(OpAsmPrinter &p) {
 
 SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
 
-SmallVector<Value> ParallelOp::getLoopInductionVars() {
+std::optional<SmallVector<Value>> ParallelOp::getLoopInductionVars() {
   return SmallVector<Value>{getBody()->getArguments()};
 }
 

diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
@@ -51,7 +51,10 @@ class SCFLoopLikeTest : public ::testing::Test {
         loopLikeOp.getLoopSteps();
     EXPECT_TRUE(maybeStep.has_value());
     EXPECT_EQ((*maybeStep).size(), 1u);
-    EXPECT_EQ(loopLikeOp.getLoopInductionVars().size(), 1u);
+    std::optional<SmallVector<Value>> maybeInductionVars =
+        loopLikeOp.getLoopInductionVars();
+    EXPECT_TRUE(maybeInductionVars.has_value());
+    EXPECT_EQ((*maybeInductionVars).size(), 1u);
   }
 
   void checkMultidimensional(LoopLikeOpInterface loopLikeOp) {
@@ -79,7 +82,10 @@ class SCFLoopLikeTest : public ::testing::Test {
         loopLikeOp.getLoopSteps();
     EXPECT_TRUE(maybeStep.has_value());
     EXPECT_EQ((*maybeStep).size(), 2u);
-    EXPECT_EQ(loopLikeOp.getLoopInductionVars().size(), 2u);
+    std::optional<SmallVector<Value>> maybeInductionVars =
+        loopLikeOp.getLoopInductionVars();
+    EXPECT_TRUE(maybeInductionVars.has_value());
+    EXPECT_EQ((*maybeInductionVars).size(), 2u);
   }
 
   MLIRContext context;