[mlir][scf] Implement getSingle... of LoopLikeOpinterface for scf::ParallelOp

[ubfx]

This adds implementations for getSingleIterationVar, getSingleLowerBound, getSingleUpperBound, getSingleStep of LoopLikeOpInterface to scf::ParallelOp. Until now, the implementations for these methods defaulted to returning std::nullopt, even in the special case where the parallel Op only has one dimension.

Related: #67883

[llvmbot]

@llvm/pr-subscribers-mlir-scf

@llvm/pr-subscribers-mlir

Changes

This adds implementations for getSingleIterationVar, getSingleLowerBound, getSingleUpperBound, getSingleStep of LoopLikeOpInterface to scf::ParallelOp. Until now, the implementations for these methods defaulted to returning std::nullopt, even in the special case where the parallel Op only has one dimension.

Related: #67883

---

Full diff: https://github.com/llvm/llvm-project/pull/68511.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/SCF/IR/SCFOps.td (+2-1)
• (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+24)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index e1a604a88715f0e..21847592aebfdcd 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -788,7 +788,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
 def ParallelOp : SCF_Op<"parallel",
     [AutomaticAllocationScope,
      AttrSizedOperandSegments,
-     DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+     DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getSingleInductionVar",
+          "getSingleLowerBound", "getSingleUpperBound", "getSingleStep"]>,
      RecursiveMemoryEffects,
      DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      SingleBlockImplicitTerminator<"scf::YieldOp">]> {
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 8d8481421e18d57..3cc1943d66c3c83 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2911,6 +2911,30 @@ void ParallelOp::print(OpAsmPrinter &p) {
 
 SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
 
+std::optional<Value> ParallelOp::getSingleInductionVar() {
+  if (getNumLoops() != 1)
+    return std::nullopt;
+  return getBody()->getArgument(0);
+}
+
+std::optional<OpFoldResult> ParallelOp::getSingleLowerBound() {
+  if (getNumLoops() != 1)
+    return std::nullopt;
+  return getLowerBound()[0];
+}
+
+std::optional<OpFoldResult> ParallelOp::getSingleUpperBound() {
+  if (getNumLoops() != 1)
+    return std::nullopt;
+  return getUpperBound()[0];
+}
+
+std::optional<OpFoldResult> ParallelOp::getSingleStep() {
+  if (getNumLoops() != 1)
+    return std::nullopt;
+  return getStep()[0];
+}
+
 ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto ivArg = llvm::dyn_cast<BlockArgument>(val);
   if (!ivArg)

[ftynse]

I don't suppose there is any chance to test this without writing a heavy unittest?

[ubfx]

I don't suppose there is any chance to test this without writing a heavy unittest?

I think a unittest for LoopLikeOpInterface would be a good idea, I'll take a look

[ubfx]

What do you think about something like this?
aa464d2

It might be a little overkill just for the getSingle... methods but maybe getting started on some unittests for the SCF dialect is useful beyond that.

[ftynse]

LGTM

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[joker-eph]

I don't suppose there is any chance to test this without writing a heavy unittest?

Actually I don't understand why this is a C++ unit-test. We have quite a track record of not doing so and making test passes instead.

Unless I'm missing something this could easily be done that way?

[ubfx]

I will take a look at this and rewrite it into a test pass if possible