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

[llvmlistbot at llvm.org]

Author: Felix Schneider
Date: 2023-10-20T07:16:21+02:00
New Revision: 833a8db5b1e3922d6a4a0a158614af884a49fc02

URL: https://github.com/llvm/llvm-project/commit/833a8db5b1e3922d6a4a0a158614af884a49fc02
DIFF: https://github.com/llvm/llvm-project/commit/833a8db5b1e3922d6a4a0a158614af884a49fc02.diff

LOG: [mlir][scf] Implement getSingle... of LoopLikeOpInterface for scf::ForallOp (#67883)

The `getSingle(IterationVar|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`.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
    mlir/lib/Dialect/SCF/IR/SCF.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index f2ea7dd868a377d..8ccbd9cfb159555 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -333,7 +333,8 @@ def ForallOp : SCF_Op<"forall", [
        AttrSizedOperandSegments,
        AutomaticAllocationScope,
        DeclareOpInterfaceMethods<LoopLikeOpInterface,
-          ["promoteIfSingleIteration"]>,
+          ["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 508227d6e7ce4b0..94e7dd4a0bf4482 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1527,6 +1527,30 @@ InParallelOp ForallOp::getTerminator() {
   return cast<InParallelOp>(getBody()->getTerminator());
 }
 
+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];
+}
+
 ForallOp mlir::scf::getForallOpThreadIndexOwner(Value val) {
   auto tidxArg = llvm::dyn_cast<BlockArgument>(val);
   if (!tidxArg)