[Mlir-commits] [mlir] [mlir][scf] Implement getSingle... of LoopLikeOpinterface for scf::ParallelOp (PR #68511)
Felix Schneider
llvmlistbot at llvm.org
Thu Oct 19 23:30:38 PDT 2023
https://github.com/ubfx updated https://github.com/llvm/llvm-project/pull/68511
>From adf987004224a2900234a267ea7652529b72563e Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Sun, 8 Oct 2023 08:16:06 +0000
Subject: [PATCH 1/3] [mlir][scf] Implement getSingle... of LoopLikeOpinterface
for scf::ParallelOp
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: https://github.com/llvm/llvm-project/pull/67883
---
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 3 ++-
mlir/lib/Dialect/SCF/IR/SCF.cpp | 24 ++++++++++++++++++++++
2 files changed, 26 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 8ccbd9cfb159555..6ac0912f6f706c5 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -791,7 +791,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 94e7dd4a0bf4482..20a7b283c938d00 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -2936,6 +2936,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)
>From cb16424aa99ce7a0e9b800bc8dc8fa134d54fef2 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Thu, 19 Oct 2023 23:11:40 +0200
Subject: [PATCH 2/3] add unit test
---
mlir/unittests/Dialect/CMakeLists.txt | 1 +
mlir/unittests/Dialect/SCF/CMakeLists.txt | 8 ++
.../Dialect/SCF/LoopLikeSCFOpsTest.cpp | 91 +++++++++++++++++++
3 files changed, 100 insertions(+)
create mode 100644 mlir/unittests/Dialect/SCF/CMakeLists.txt
create mode 100644 mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index 2d2835c64b9844f..fbb73e8f499a35e 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -9,6 +9,7 @@ target_link_libraries(MLIRDialectTests
add_subdirectory(Index)
add_subdirectory(LLVMIR)
add_subdirectory(MemRef)
+add_subdirectory(SCF)
add_subdirectory(SparseTensor)
add_subdirectory(SPIRV)
add_subdirectory(Transform)
diff --git a/mlir/unittests/Dialect/SCF/CMakeLists.txt b/mlir/unittests/Dialect/SCF/CMakeLists.txt
new file mode 100644
index 000000000000000..4d23392af1f88dc
--- /dev/null
+++ b/mlir/unittests/Dialect/SCF/CMakeLists.txt
@@ -0,0 +1,8 @@
+add_mlir_unittest(MLIRSCFTests
+ LoopLikeSCFOpsTest.cpp
+)
+target_link_libraries(MLIRSCFTests
+ PRIVATE
+ MLIRIR
+ MLIRSCFDialect
+)
diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
new file mode 100644
index 000000000000000..209858b114668b0
--- /dev/null
+++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
@@ -0,0 +1,91 @@
+//===- LoopLikeSCFOpsTest.cpp - SCF LoopLikeOpInterface Tests -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/MLIRContext.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::scf;
+
+//===----------------------------------------------------------------------===//
+// Test Fixture
+//===----------------------------------------------------------------------===//
+
+class SCFLoopLikeTest : public ::testing::Test {
+protected:
+ SCFLoopLikeTest() : b(&context), loc(UnknownLoc::get(&context)) {
+ context.loadDialect<arith::ArithDialect, scf::SCFDialect>();
+ }
+
+ void checkUnidimensional(LoopLikeOpInterface loopLikeOp) {
+ std::optional<OpFoldResult> maybeLb = loopLikeOp.getSingleLowerBound();
+ EXPECT_TRUE(maybeLb.has_value());
+ std::optional<OpFoldResult> maybeUb = loopLikeOp.getSingleUpperBound();
+ EXPECT_TRUE(maybeUb.has_value());
+ std::optional<OpFoldResult> maybeStep = loopLikeOp.getSingleStep();
+ EXPECT_TRUE(maybeStep.has_value());
+ std::optional<OpFoldResult> maybeIndVar =
+ loopLikeOp.getSingleInductionVar();
+ EXPECT_TRUE(maybeIndVar.has_value());
+ }
+
+ void checkMultidimensional(LoopLikeOpInterface loopLikeOp) {
+ std::optional<OpFoldResult> maybeLb = loopLikeOp.getSingleLowerBound();
+ EXPECT_FALSE(maybeLb.has_value());
+ std::optional<OpFoldResult> maybeUb = loopLikeOp.getSingleUpperBound();
+ EXPECT_FALSE(maybeUb.has_value());
+ std::optional<OpFoldResult> maybeStep = loopLikeOp.getSingleStep();
+ EXPECT_FALSE(maybeStep.has_value());
+ std::optional<OpFoldResult> maybeIndVar =
+ loopLikeOp.getSingleInductionVar();
+ EXPECT_FALSE(maybeIndVar.has_value());
+ }
+
+ MLIRContext context;
+ OpBuilder b;
+ Location loc;
+};
+
+TEST_F(SCFLoopLikeTest, queryUnidimensionalLooplikes) {
+ Value lb = b.create<arith::ConstantIndexOp>(loc, 0);
+ Value ub = b.create<arith::ConstantIndexOp>(loc, 10);
+ Value step = b.create<arith::ConstantIndexOp>(loc, 2);
+
+ auto forOp = b.create<scf::ForOp>(loc, lb, ub, step);
+ checkUnidimensional(forOp);
+
+ auto forallOp = b.create<scf::ForallOp>(
+ loc, ArrayRef<OpFoldResult>(lb), ArrayRef<OpFoldResult>(ub),
+ ArrayRef<OpFoldResult>(step), ValueRange(), std::nullopt);
+ checkUnidimensional(forallOp);
+
+ auto parallelOp = b.create<scf::ParallelOp>(
+ loc, ValueRange(lb), ValueRange(ub), ValueRange(step), ValueRange());
+ checkUnidimensional(parallelOp);
+}
+
+TEST_F(SCFLoopLikeTest, queryMultidimensionalLooplikes) {
+ Value lb = b.create<arith::ConstantIndexOp>(loc, 0);
+ Value ub = b.create<arith::ConstantIndexOp>(loc, 10);
+ Value step = b.create<arith::ConstantIndexOp>(loc, 2);
+
+ auto forallOp = b.create<scf::ForallOp>(
+ loc, ArrayRef<OpFoldResult>({lb, lb}), ArrayRef<OpFoldResult>({ub, ub}),
+ ArrayRef<OpFoldResult>({step, step}), ValueRange(), std::nullopt);
+ checkMultidimensional(forallOp);
+
+ auto parallelOp = b.create<scf::ParallelOp>(
+ loc, ValueRange({lb, lb}), ValueRange({ub, ub}), ValueRange({step, step}), ValueRange());
+ checkMultidimensional(parallelOp);
+}
\ No newline at end of file
>From 6d71605d6f6df4aade153b416b72ad6300dce5ce Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Fri, 20 Oct 2023 06:12:35 +0000
Subject: [PATCH 3/3] code formatting
---
mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp | 10 ++++------
1 file changed, 4 insertions(+), 6 deletions(-)
diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
index 209858b114668b0..f75b84f12b6f1f1 100644
--- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
+++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
@@ -5,9 +5,6 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
-//
-//
-//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
@@ -85,7 +82,8 @@ TEST_F(SCFLoopLikeTest, queryMultidimensionalLooplikes) {
ArrayRef<OpFoldResult>({step, step}), ValueRange(), std::nullopt);
checkMultidimensional(forallOp);
- auto parallelOp = b.create<scf::ParallelOp>(
- loc, ValueRange({lb, lb}), ValueRange({ub, ub}), ValueRange({step, step}), ValueRange());
+ auto parallelOp =
+ b.create<scf::ParallelOp>(loc, ValueRange({lb, lb}), ValueRange({ub, ub}),
+ ValueRange({step, step}), ValueRange());
checkMultidimensional(parallelOp);
-}
\ No newline at end of file
+}
More information about the Mlir-commits
mailing list