[Mlir-commits] [mlir] [mlir][scf] Implement getSingle... of LoopLikeOpInterface for scf::ForallOp (PR #67883)

[Felix Schneider]

https://github.com/ubfx updated https://github.com/llvm/llvm-project/pull/67883

>From 8f1378be30b8dcbc738232187002a130403e8bee Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Sat, 30 Sep 2023 13:40:29 +0000
Subject: [PATCH 1/2] [mlir][scf] Implement getSingle... of LoopLikeOpInterface
 for scf::ForallOp

The `getSingle(UpperBound|LowerBound|Step)` methods of `LoopLikeOpInterface` are
useful to quickly query the iteration space of unidimensional loops. Until now,
`scf::ForallOp` always fell back to the default implementation of these methods,
returning `std::nullopt`.

This patch implements those methods, returning the respective bounds or steps
in the special case of `rank == 1`.
---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td |  3 ++-
 mlir/lib/Dialect/SCF/IR/SCF.cpp            | 21 +++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index e1a604a88715f0e..218d3b09d96fb05 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -331,7 +331,8 @@ def ForallOp : SCF_Op<"forall", [
        AttrSizedOperandSegments,
        AutomaticAllocationScope,
        DeclareOpInterfaceMethods<LoopLikeOpInterface,
-          ["promoteIfSingleIteration"]>,
+          ["promoteIfSingleIteration", "getSingleLowerBound",
+          "getSingleUpperBound", "getSingleStep"]>,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 8d8481421e18d57..bfe6ccd7b1169df 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1526,6 +1526,27 @@ InParallelOp ForallOp::getTerminator() {
   return cast<InParallelOp>(getBody()->getTerminator());
 }
 
+std::optional<OpFoldResult> ForallOp::getSingleLowerBound() {
+  if (getRank() != 1)
+    return std::nullopt;
+
+  return getMixedLowerBound()[0];
+}
+
+std::optional<OpFoldResult> ForallOp::getSingleUpperBound() {
+  if (getRank() != 1)
+    return std::nullopt;
+
+  return getMixedUpperBound()[0];
+}
+
+std::optional<OpFoldResult> ForallOp::getSingleStep() {
+  if (getRank() != 1)
+    return std::nullopt;
+
+  return getMixedStep()[0];
+}
+
 ForallOp mlir::scf::getForallOpThreadIndexOwner(Value val) {
   auto tidxArg = llvm::dyn_cast<BlockArgument>(val);
   if (!tidxArg)

>From 8062b56d1a18468c9e08da5d1483f20aa357ed0f Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Sat, 30 Sep 2023 16:23:34 +0200
Subject: [PATCH 2/2] also implement getSingleInductionVar()

---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 4 ++--
 mlir/lib/Dialect/SCF/IR/SCF.cpp            | 9 ++++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 218d3b09d96fb05..66552b0b387b750 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -331,8 +331,8 @@ def ForallOp : SCF_Op<"forall", [
        AttrSizedOperandSegments,
        AutomaticAllocationScope,
        DeclareOpInterfaceMethods<LoopLikeOpInterface,
-          ["promoteIfSingleIteration", "getSingleLowerBound",
-          "getSingleUpperBound", "getSingleStep"]>,
+          ["promoteIfSingleIteration", "getSingleInductionVar",
+          "getSingleLowerBound", "getSingleUpperBound", "getSingleStep"]>,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index bfe6ccd7b1169df..131e956a4799571 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1526,24 +1526,27 @@ InParallelOp ForallOp::getTerminator() {
   return cast<InParallelOp>(getBody()->getTerminator());
 }
 
-std::optional<OpFoldResult> ForallOp::getSingleLowerBound() {
+std::optional<Value> ForallOp::getSingleInductionVar() {
   if (getRank() != 1)
     return std::nullopt;
+  return getInductionVar(0);
+}
 
+std::optional<OpFoldResult> ForallOp::getSingleLowerBound() {
+  if (getRank() != 1)
+    return std::nullopt;
   return getMixedLowerBound()[0];
 }
 
 std::optional<OpFoldResult> ForallOp::getSingleUpperBound() {
   if (getRank() != 1)
     return std::nullopt;
-
   return getMixedUpperBound()[0];
 }
 
 std::optional<OpFoldResult> ForallOp::getSingleStep() {
   if (getRank() != 1)
     return std::nullopt;
-
   return getMixedStep()[0];
 }