Make return types optional and change names

srcarroll · srcarroll · commit a5fa3b3c4903 · 2024-06-05T20:56:02.000-05:00
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -118,8 +118,8 @@ def AffineForOp : Affine_Op<"for",
     [AttrSizedOperandSegments, AutomaticAllocationScope,
      ImplicitAffineTerminator, ConditionallySpeculatable,
      RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
-     ["getInductionVars", "getMixedLowerBound", "getMixedStep",
-      "getMixedUpperBound", "getYieldedValuesMutable",
+     ["getInductionVars", "getLowerBounds", "getSteps",
+      "getUpperBounds", "getYieldedValuesMutable",
       "replaceWithAdditionalYields"]>,
      DeclareOpInterfaceMethods<RegionBranchOpInterface,
      ["getEntrySuccessorOperands"]>]> {
@@ -671,7 +671,7 @@ def AffineParallelOp : Affine_Op<"parallel",
      I32ElementsAttr:$lowerBoundsGroups,
      AffineMapAttr:$upperBoundsMap,
      I32ElementsAttr:$upperBoundsGroups,
-     I64SmallVectorArrayAttr:$steps,
+     I64SmallVectorArrayAttr:$step,
      Variadic<Index>:$mapOperands);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
@@ -682,7 +682,7 @@ def AffineParallelOp : Affine_Op<"parallel",
     OpBuilder<(ins "TypeRange":$resultTypes,
       "ArrayRef<arith::AtomicRMWKind>":$reductions, "ArrayRef<AffineMap>":$lbMaps,
       "ValueRange":$lbArgs, "ArrayRef<AffineMap>":$ubMaps, "ValueRange":$ubArgs,
-      "ArrayRef<int64_t>":$steps)>
+      "ArrayRef<int64_t>":$step)>
   ];
 
   let extraClassDeclaration = [{
@@ -727,7 +727,7 @@ def AffineParallelOp : Affine_Op<"parallel",
     static StringRef getUpperBoundsGroupsAttrStrName() {
       return "upperBoundsGroups";
     }
-    static StringRef getStepsAttrStrName() { return "steps"; }
+    static StringRef getStepsAttrStrName() { return "step"; }
 
     /// Returns `true` if the loop bounds have min/max expressions.
     bool hasMinMaxBounds() {
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -136,8 +136,8 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
 def ForOp : SCF_Op<"for",
       [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
        ["getInitsMutable", "getLoopResults", "getRegionIterArgs",
-        "getInductionVars", "getMixedLowerBound", "getMixedStep",
-        "getMixedUpperBound", "getYieldedValuesMutable",
+        "getInductionVars", "getLowerBounds", "getSteps",
+        "getUpperBounds", "getYieldedValuesMutable",
         "promoteIfSingleIteration", "replaceWithAdditionalYields",
         "yieldTiledValuesAndReplace"]>,
        AllTypesMatch<["lowerBound", "upperBound", "step"]>,
@@ -302,7 +302,7 @@ def ForallOp : SCF_Op<"forall", [
        AutomaticAllocationScope,
        DeclareOpInterfaceMethods<LoopLikeOpInterface,
           ["getInitsMutable", "getRegionIterArgs", "getInductionVars", 
-           "getMixedLowerBound", "getMixedUpperBound", "getMixedStep",
+           "getLowerBounds", "getUpperBounds", "getSteps",
            "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
@@ -510,6 +510,27 @@ def ForallOp : SCF_Op<"forall", [
   ];
 
   let extraClassDeclaration = [{
+    // Get lower bounds as OpFoldResult.
+    SmallVector<OpFoldResult> getMixedLowerBound() {
+      auto maybeLowerBounds = getLowerBounds();
+      assert(maybeLowerBounds.has_value() && "expected values");
+      return *maybeLowerBounds;
+    }
+
+    // Get upper bounds as OpFoldResult.
+    SmallVector<OpFoldResult> getMixedUpperBound() {
+      auto maybeUpperBounds = getUpperBounds();
+      assert(maybeUpperBounds.has_value() && "expected values");
+      return *maybeUpperBounds;
+    }
+
+    // Get steps as OpFoldResult.
+    SmallVector<OpFoldResult> getMixedStep() {
+      auto maybeSteps = getSteps();
+      assert(maybeSteps.has_value() && "expected values");
+      return *maybeSteps;
+    }
+
     /// Get lower bounds as values.
     SmallVector<Value> getLowerBound(OpBuilder &b) {
       return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedLowerBound());
@@ -744,7 +765,7 @@ def ParallelOp : SCF_Op<"parallel",
     [AutomaticAllocationScope,
      AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getInductionVars",
-          "getMixedLowerBound", "getMixedUpperBound", "getMixedStep"]>,
+          "getLowerBounds", "getUpperBounds", "getSteps"]>,
      RecursiveMemoryEffects,
      DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      SingleBlockImplicitTerminator<"scf::ReduceOp">,
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -106,34 +106,34 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
     InterfaceMethod<[{
         Return all lower bounds.
       }],
-      /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
-      /*methodName=*/"getMixedLowerBound",
+      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+      /*methodName=*/"getLowerBounds",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return {};
+        return std::nullopt;
       }]
     >,
     InterfaceMethod<[{
         Return all steps.
       }],
-      /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
-      /*methodName=*/"getMixedStep",
+      /*retTy=*/"std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+      /*methodName=*/"getSteps",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return {};
+        return std::nullopt;
       }]
     >,
     InterfaceMethod<[{
         Return all upper bounds.
       }],
-      /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
-      /*methodName=*/"getMixedUpperBound",
+      /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+      /*methodName=*/"getUpperBounds",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return {};
+        return std::nullopt;
       }]
     >,
     InterfaceMethod<[{
@@ -242,26 +242,26 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
     /// Return the single lower bound value or attribute if it exists, otherwise
     /// return std::nullopt.
     ::std::optional<::mlir::OpFoldResult> getSingleLowerBound() {
-      auto lowerBounds = this->getMixedLowerBound();
-      if (lowerBounds.size() == 1)
-          return lowerBounds[0];
-        return std::nullopt;
+      auto lowerBounds = this->getLowerBounds();
+      if (lowerBounds.has_value() && (*lowerBounds).size() == 1)
+          return (*lowerBounds)[0];
+      return std::nullopt;
     }
     /// Return the single step value or attribute if it exists, otherwise
     /// return std::nullopt.
     ::std::optional<::mlir::OpFoldResult> getSingleStep() {
-      auto steps = this->getMixedStep(); 
-      if (steps.size() == 1)
-          return steps[0];
-        return std::nullopt;
+      auto steps = this->getSteps(); 
+      if (steps.has_value() && (*steps).size() == 1)
+          return (*steps)[0];
+      return std::nullopt;
     }
     /// Return the single upper bound value or attribute if it exists, otherwise
     /// return std::nullopt.
     ::std::optional<::mlir::OpFoldResult> getSingleUpperBound() {
-      auto upperBounds = this->getMixedUpperBound();
-      if (upperBounds.size() == 1)
-          return upperBounds[0];
-        return std::nullopt;
+      auto upperBounds = this->getUpperBounds();
+      if (upperBounds.has_value() && (*upperBounds).size() == 1)
+          return (*upperBounds)[0];
+      return std::nullopt;
     }
 
     /// Append the specified additional "init" operands: replace this loop with
diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -196,8 +196,8 @@ class AffineParallelLowering : public OpRewritePattern<AffineParallelOp> {
         return rewriter.notifyMatchFailure(op, "couldn't convert upper bounds");
       upperBoundTuple.push_back(upper);
     }
-    steps.reserve(op.getSteps().size());
-    for (int64_t step : op.getSteps())
+    steps.reserve(op.getStep().size());
+    for (int64_t step : op.getStep())
       steps.push_back(rewriter.create<arith::ConstantIndexOp>(loc, step));
 
     // Get the terminator op.
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -695,12 +695,9 @@ LogicalResult ForallLowering::matchAndRewrite(ForallOp forallOp,
         "only fully bufferized scf.forall ops can be lowered to scf.parallel");
 
   // Convert mixed bounds and steps to SSA values.
-  SmallVector<Value> lbs = getValueOrCreateConstantIndexOp(
-      rewriter, loc, forallOp.getMixedLowerBound());
-  SmallVector<Value> ubs = getValueOrCreateConstantIndexOp(
-      rewriter, loc, forallOp.getMixedUpperBound());
-  SmallVector<Value> steps =
-      getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep());
+  SmallVector<Value> lbs = forallOp.getLowerBound(rewriter);
+  SmallVector<Value> ubs = forallOp.getUpperBound(rewriter);
+  SmallVector<Value> steps = forallOp.getStep(rewriter);
 
   // Create empty scf.parallel op.
   auto parallelOp = rewriter.create<ParallelOp>(loc, lbs, ubs, steps);
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2456,23 +2456,26 @@ SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
 
 ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; }
 
-SmallVector<OpFoldResult> AffineForOp::getMixedLowerBound() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLowerBounds() {
   if (!hasConstantLowerBound())
-    return {};
+    return std::nullopt;
   OpBuilder b(getContext());
-  return {OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
+  return SmallVector<OpFoldResult>{
+      OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
 }
 
-SmallVector<OpFoldResult> AffineForOp::getMixedStep() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getSteps() {
   OpBuilder b(getContext());
-  return {OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
+  return SmallVector<OpFoldResult>{
+      OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
 }
 
-SmallVector<OpFoldResult> AffineForOp::getMixedUpperBound() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getUpperBounds() {
   if (!hasConstantUpperBound())
     return {};
   OpBuilder b(getContext());
-  return {OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
+  return SmallVector<OpFoldResult>{
+      OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
 }
 
 FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
@@ -3753,7 +3756,7 @@ SmallVector<Region *> AffineParallelOp::getLoopRegions() {
   return {&getRegion()};
 }
 
-unsigned AffineParallelOp::getNumDims() { return getSteps().size(); }
+unsigned AffineParallelOp::getNumDims() { return getStep().size(); }
 
 AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
   return getOperands().take_front(getLowerBoundsMap().getNumInputs());
@@ -3838,7 +3841,7 @@ void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) {
 }
 
 void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
-  setStepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
+  setStepAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
 }
 
 // check whether resultType match op or not in affine.parallel
@@ -3888,14 +3891,14 @@ LogicalResult AffineParallelOp::verify() {
   auto numDims = getNumDims();
   if (getLowerBoundsGroups().getNumElements() != numDims ||
       getUpperBoundsGroups().getNumElements() != numDims ||
-      getSteps().size() != numDims || getBody()->getNumArguments() != numDims) {
+      getStep().size() != numDims || getBody()->getNumArguments() != numDims) {
     return emitOpError() << "the number of region arguments ("
                          << getBody()->getNumArguments()
                          << ") and the number of map groups for lower ("
                          << getLowerBoundsGroups().getNumElements()
                          << ") and upper bound ("
                          << getUpperBoundsGroups().getNumElements()
-                         << "), and the number of steps (" << getSteps().size()
+                         << "), and the number of steps (" << getStep().size()
                          << ") must all match";
   }
 
@@ -4013,7 +4016,7 @@ void AffineParallelOp::print(OpAsmPrinter &p) {
   printMinMaxBound(p, getUpperBoundsMapAttr(), getUpperBoundsGroupsAttr(),
                    getUpperBoundsOperands(), "min");
   p << ')';
-  SmallVector<int64_t, 8> steps = getSteps();
+  SmallVector<int64_t, 8> steps = getStep();
   bool elideSteps = llvm::all_of(steps, [](int64_t step) { return step == 1; });
   if (!elideSteps) {
     p << " step (";
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -494,7 +494,7 @@ void mlir::affine::normalizeAffineParallel(AffineParallelOp op) {
     return;
 
   AffineMap lbMap = op.getLowerBoundsMap();
-  SmallVector<int64_t, 8> steps = op.getSteps();
+  SmallVector<int64_t, 8> steps = op.getStep();
   // No need to do any work if the parallel op is already normalized.
   bool isAlreadyNormalized =
       llvm::all_of(llvm::zip(steps, lbMap.getResults()), [](auto tuple) {
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -380,16 +380,16 @@ LogicalResult ForOp::verifyRegions() {
 
 ValueRange ForOp::getInductionVars() { return {getInductionVar()}; }
 
-SmallVector<OpFoldResult> ForOp::getMixedLowerBound() {
-  return {OpFoldResult(getLowerBound())};
+std::optional<SmallVector<OpFoldResult>> ForOp::getLowerBounds() {
+  return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
 }
 
-SmallVector<OpFoldResult> ForOp::getMixedStep() {
-  return {OpFoldResult(getStep())};
+std::optional<SmallVector<OpFoldResult>> ForOp::getSteps() {
+  return SmallVector<OpFoldResult, 1>{OpFoldResult(getStep())};
 }
 
-SmallVector<OpFoldResult> ForOp::getMixedUpperBound() {
-  return {OpFoldResult(getUpperBound())};
+std::optional<SmallVector<OpFoldResult>> ForOp::getUpperBounds() {
+  return SmallVector<OpFoldResult, 1>{OpFoldResult(getUpperBound())};
 }
 
 std::optional<ResultRange> ForOp::getLoopResults() { return getResults(); }
@@ -1431,19 +1431,19 @@ ValueRange ForallOp::getInductionVars() {
 }
 
 // Get lower bounds as OpFoldResult.
-SmallVector<OpFoldResult> ForallOp::getMixedLowerBound() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLowerBounds() {
   Builder b(getOperation()->getContext());
   return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
 }
 
 // Get upper bounds as OpFoldResult.
-SmallVector<OpFoldResult> ForallOp::getMixedUpperBound() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getUpperBounds() {
   Builder b(getOperation()->getContext());
   return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
 }
 
 // Get steps as OpFoldResult.
-SmallVector<OpFoldResult> ForallOp::getMixedStep() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getSteps() {
   Builder b(getOperation()->getContext());
   return getMixedValues(getStaticStep(), getDynamicStep(), b);
 }
@@ -3006,15 +3006,17 @@ SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
 
 ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); }
 
-SmallVector<OpFoldResult> ParallelOp::getMixedLowerBound() {
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLowerBounds() {
   return getLowerBound();
 }
 
-SmallVector<OpFoldResult> ParallelOp::getMixedUpperBound() {
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getUpperBounds() {
   return getUpperBound();
 }
 
-SmallVector<OpFoldResult> ParallelOp::getMixedStep() { return getStep(); }
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getSteps() {
+  return getStep();
+}
 
 ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto ivArg = llvm::dyn_cast<BlockArgument>(val);
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp b/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
@@ -34,12 +34,9 @@ mlir::scf::forallToForLoop(RewriterBase &rewriter, scf::ForallOp forallOp,
   rewriter.setInsertionPoint(forallOp);
 
   Location loc = forallOp.getLoc();
-  SmallVector<Value> lbs = getValueOrCreateConstantIndexOp(
-      rewriter, loc, forallOp.getMixedLowerBound());
-  SmallVector<Value> ubs = getValueOrCreateConstantIndexOp(
-      rewriter, loc, forallOp.getMixedUpperBound());
-  SmallVector<Value> steps =
-      getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep());
+  SmallVector<Value> lbs = forallOp.getLowerBound(rewriter);
+  SmallVector<Value> ubs = forallOp.getUpperBound(rewriter);
+  SmallVector<Value> steps = forallOp.getStep(rewriter);
   LoopNest loopNest = scf::buildLoopNest(rewriter, loc, lbs, ubs, steps);
 
   SmallVector<Value> ivs = llvm::map_to_vector(
diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
