[Mlir-commits] [mlir] 7468b56 - [mlir][acc] Add utilities for working with acc par dims (#208120)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 8 08:01:39 PDT 2026
Author: Razvan Lupusoru
Date: 2026-07-08T08:01:34-07:00
New Revision: 7468b5634b038ba915fc326a8fb2f3fe6c3ed098
URL: https://github.com/llvm/llvm-project/commit/7468b5634b038ba915fc326a8fb2f3fe6c3ed098
DIFF: https://github.com/llvm/llvm-project/commit/7468b5634b038ba915fc326a8fb2f3fe6c3ed098.diff
LOG: [mlir][acc] Add utilities for working with acc par dims (#208120)
Add shared helpers for reading, setting, and manipulating parallel
dimensions on operations, covering both discardable and inherent
attributes.
Added:
Modified:
mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
index bb132f5d02e8c..64de4eda7366a 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
@@ -17,6 +17,7 @@
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "llvm/ADT/SmallVector.h"
#include <optional>
namespace mlir {
@@ -62,6 +63,33 @@ ComputeRegionOp buildComputeRegion(Location loc, ValueRange launchArgs,
Value stream = {},
ValueRange inputArgsToMap = {});
+/// Insert \p parDim into \p parDims while preserving dimension ordering. If the
+/// dimension is already present, this is a no-op.
+void insertParDim(llvm::SmallVector<GPUParallelDimAttr> &parDims,
+ GPUParallelDimAttr parDim);
+
+/// Remove \p parDim from \p parDims if present.
+void removeParDim(llvm::SmallVector<GPUParallelDimAttr> &parDims,
+ GPUParallelDimAttr parDim);
+
+/// Obtain the parallel dimensions carried by \p op, if any.
+GPUParallelDimsAttr getParDimsAttr(Operation *op);
+
+/// Return whether \p op carries parallel dimensions.
+bool hasParDimsAttr(Operation *op);
+
+/// Return whether \p op carries sequential parallel dimensions.
+bool hasSeqParDims(Operation *op);
+
+/// Set parallel dimensions on \p op.
+void setParDimsAttr(Operation *op, GPUParallelDimsAttr attr);
+
+/// Update parallel dimensions on \p op.
+void updateParDimsAttr(Operation *op, GPUParallelDimsAttr attr);
+
+/// Copy parallel dimensions from \p from to \p to.
+void copyParDimsAttr(Operation *from, Operation *to);
+
} // namespace acc
} // namespace mlir
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
index 1c5f302f23192..b3d54d2e05697 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
@@ -138,10 +138,6 @@ static bool isOpInSerialRegion(Operation *op) {
return false;
}
-static void setParDimsAttr(Operation *op, GPUParallelDimsAttr attr) {
- op->setAttr(GPUParallelDimsAttr::name, attr);
-}
-
/// Clone defining ops of constant live-in values into `region`, rewrite uses
/// inside the region to the clones, and remove those values from
/// `liveInValues` so they are not threaded through `acc.compute_region` ins.
@@ -172,19 +168,6 @@ static void materializeConstantLiveInsIntoRegion(Region ®ion,
}
}
-/// Insert a parallel dimension into the list, maintaining order by
-/// GPUParallelDimAttr::getOrder (descending).
-static void insertParDim(SmallVectorImpl<GPUParallelDimAttr> &parDims,
- GPUParallelDimAttr parDim) {
- GPUParallelDimAttr *lb = llvm::lower_bound(
- parDims, parDim,
- [](const GPUParallelDimAttr &a, const GPUParallelDimAttr &b) {
- return a.getOrder() > b.getOrder();
- });
- if (lb == parDims.end() || *lb != parDim)
- parDims.insert(lb, parDim);
-}
-
/// Return the device type from which gang/worker/vector clauses should be read.
/// If the requested device type has any such clauses, use that exclusively;
/// otherwise fall back to the default (DeviceType::None).
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index d18522c3c440f..06605371bdf42 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -16,6 +16,8 @@
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/IRMapping.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
namespace mlir {
namespace acc {
@@ -128,5 +130,81 @@ ComputeRegionOp buildComputeRegion(Location loc, ValueRange launchArgs,
return computeRegion;
}
+static SmallVector<GPUParallelDimAttr>::iterator
+findParDim(SmallVector<GPUParallelDimAttr> &parDims,
+ GPUParallelDimAttr parDim) {
+ return llvm::lower_bound(
+ parDims, parDim,
+ [](const GPUParallelDimAttr &lhs, const GPUParallelDimAttr &rhs) {
+ return lhs.getOrder() > rhs.getOrder();
+ });
+}
+
+void insertParDim(SmallVector<GPUParallelDimAttr> &parDims,
+ GPUParallelDimAttr parDim) {
+ SmallVector<GPUParallelDimAttr>::iterator lb = findParDim(parDims, parDim);
+ if (lb == parDims.end() || *lb != parDim)
+ parDims.insert(lb, parDim);
+}
+
+void removeParDim(SmallVector<GPUParallelDimAttr> &parDims,
+ GPUParallelDimAttr parDim) {
+ SmallVector<GPUParallelDimAttr>::iterator lb = findParDim(parDims, parDim);
+ if (lb != parDims.end() && *lb == parDim)
+ parDims.erase(lb);
+}
+
+#define ACC_OP_WITH_PAR_DIMS_LIST \
+ PrivatizeOp, ReductionAccumulateOp, ReductionAccumulateArrayOp
+
+GPUParallelDimsAttr getParDimsAttr(Operation *op) {
+ return llvm::TypeSwitch<Operation *, GPUParallelDimsAttr>(op)
+ .Case<ACC_OP_WITH_PAR_DIMS_LIST>(
+ [](auto parOp) { return parOp.getParDimsAttr(); })
+ .Default([](Operation *op) -> GPUParallelDimsAttr {
+ if (Attribute attr = op->getAttr(GPUParallelDimsAttr::name)) {
+ GPUParallelDimsAttr parDimsAttr = dyn_cast<GPUParallelDimsAttr>(attr);
+ assert(parDimsAttr && "acc.par_dims must be a GPUParallelDimsAttr");
+ return parDimsAttr;
+ }
+ return nullptr;
+ });
+}
+
+bool hasParDimsAttr(Operation *op) { return getParDimsAttr(op) != nullptr; }
+
+bool hasSeqParDims(Operation *op) {
+ if (GPUParallelDimsAttr parDimsAttr = getParDimsAttr(op))
+ return parDimsAttr.isSeq();
+ return false;
+}
+
+void setParDimsAttr(Operation *op, GPUParallelDimsAttr attr) {
+ assert(!hasParDimsAttr(op) && "parallel dimensions attribute is already set");
+ llvm::TypeSwitch<Operation *>(op)
+ .Case<ACC_OP_WITH_PAR_DIMS_LIST>(
+ [&](auto parOp) { parOp.setParDimsAttr(attr); })
+ .Default(
+ [&](Operation *op) { op->setAttr(GPUParallelDimsAttr::name, attr); });
+}
+
+void updateParDimsAttr(Operation *op, GPUParallelDimsAttr attr) {
+ assert(hasParDimsAttr(op) &&
+ "expected parallel dimensions attribute to already be set");
+ llvm::TypeSwitch<Operation *>(op)
+ .Case<ACC_OP_WITH_PAR_DIMS_LIST>(
+ [&](auto parOp) { parOp.setParDimsAttr(attr); })
+ .Default(
+ [&](Operation *op) { op->setAttr(GPUParallelDimsAttr::name, attr); });
+}
+
+#undef ACC_OP_WITH_PAR_DIMS_LIST
+
+void copyParDimsAttr(Operation *from, Operation *to) {
+ assert(hasParDimsAttr(from) &&
+ "expected parallel dimensions attribute to already be set");
+ setParDimsAttr(to, getParDimsAttr(from));
+}
+
} // namespace acc
} // namespace mlir
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
index 6fe0ffb2d54fe..75cf385af096f 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
@@ -83,6 +83,95 @@ TEST_F(OpenACCUtilsCGTest, getDataLayoutWithSpec) {
EXPECT_TRUE(dl2.has_value());
}
+//===----------------------------------------------------------------------===//
+// ParDim utilities Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, insertParDimOrdersAndDeduplicates) {
+ SmallVector<GPUParallelDimAttr> parDims;
+ GPUParallelDimAttr threadX = GPUParallelDimAttr::threadXDim(&context);
+ GPUParallelDimAttr blockX = GPUParallelDimAttr::blockXDim(&context);
+ GPUParallelDimAttr threadY = GPUParallelDimAttr::threadYDim(&context);
+
+ insertParDim(parDims, threadX);
+ insertParDim(parDims, blockX);
+ insertParDim(parDims, threadY);
+ insertParDim(parDims, threadX);
+
+ ASSERT_EQ(parDims.size(), 3u);
+ EXPECT_EQ(parDims[0], blockX);
+ EXPECT_EQ(parDims[1], threadY);
+ EXPECT_EQ(parDims[2], threadX);
+}
+
+TEST_F(OpenACCUtilsCGTest, removeParDimRemovesOnlyMatchingDim) {
+ GPUParallelDimAttr threadX = GPUParallelDimAttr::threadXDim(&context);
+ GPUParallelDimAttr blockX = GPUParallelDimAttr::blockXDim(&context);
+ GPUParallelDimAttr threadY = GPUParallelDimAttr::threadYDim(&context);
+ SmallVector<GPUParallelDimAttr> parDims{blockX, threadY, threadX};
+
+ removeParDim(parDims, threadX);
+ removeParDim(parDims, GPUParallelDimAttr::blockYDim(&context));
+
+ ASSERT_EQ(parDims.size(), 2u);
+ EXPECT_EQ(parDims[0], blockX);
+ EXPECT_EQ(parDims[1], threadY);
+}
+
+TEST_F(OpenACCUtilsCGTest, parDimsOperationAttributes) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ OwningOpRef<ModuleOp> otherModule = ModuleOp::create(b, loc);
+ Operation *op = module->getOperation();
+ Operation *otherOp = otherModule->getOperation();
+ GPUParallelDimsAttr seqAttr = GPUParallelDimsAttr::seq(&context);
+ GPUParallelDimsAttr blockAttr = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+
+ EXPECT_FALSE(hasParDimsAttr(op));
+ setParDimsAttr(op, seqAttr);
+ EXPECT_TRUE(hasParDimsAttr(op));
+ EXPECT_TRUE(hasSeqParDims(op));
+ EXPECT_EQ(getParDimsAttr(op), seqAttr);
+
+ updateParDimsAttr(op, blockAttr);
+ EXPECT_FALSE(hasSeqParDims(op));
+ EXPECT_EQ(getParDimsAttr(op), blockAttr);
+
+ copyParDimsAttr(op, otherOp);
+ EXPECT_EQ(getParDimsAttr(otherOp), blockAttr);
+}
+
+TEST_F(OpenACCUtilsCGTest, getParDimsAttrReadsInherentAttribute) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+
+ GPUParallelDimsAttr blockAttr = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+ Type privateTy = PrivateType::get(&context, b.getI32Type());
+ auto privatize =
+ PrivatizeOp::create(b, loc, privateTy,
+ /*dynamicSizes=*/ValueRange{}, blockAttr);
+
+ EXPECT_TRUE(hasParDimsAttr(privatize));
+ EXPECT_EQ(getParDimsAttr(privatize), blockAttr);
+ EXPECT_EQ(privatize.getParDimsAttr(), blockAttr);
+}
+
+TEST_F(OpenACCUtilsCGTest, setParDimsAttrSetsInherentAttribute) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+
+ Type privateTy = PrivateType::get(&context, b.getI32Type());
+ auto privatize =
+ PrivatizeOp::create(b, loc, privateTy, /*dynamicSizes=*/ValueRange{});
+ GPUParallelDimsAttr blockAttr = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+
+ setParDimsAttr(privatize, blockAttr);
+ EXPECT_EQ(getParDimsAttr(privatize), blockAttr);
+ EXPECT_EQ(privatize.getParDimsAttr(), blockAttr);
+}
+
//===----------------------------------------------------------------------===//
// buildComputeRegion Tests
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list