[Mlir-commits] [mlir] [mlir][acc] Add utilities for acc to gpu lowering (PR #209313)
Razvan Lupusoru
llvmlistbot at llvm.org
Mon Jul 13 14:52:13 PDT 2026
https://github.com/razvanlupusoru updated https://github.com/llvm/llvm-project/pull/209313
>From c686ac8dacf5f04e808cb746433d28e9a21460e8 Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Mon, 13 Jul 2026 14:25:06 -0700
Subject: [PATCH 1/3] [mlir][acc] Add utilities for acc to gpu lowering
In preparation for the pass that converts `acc.compute_region`
to GPU dialect, this PR adds several utilities which are used
in that pass. Doing so to simplify review and to ensure that
unit testing is added for each of the utilities.
---
.../mlir/Dialect/OpenACC/OpenACCUtilsCG.h | 55 +++
.../mlir/Dialect/OpenACC/OpenACCUtilsGPU.h | 12 +
.../Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp | 223 ++++++++++
.../Dialect/OpenACC/Utils/OpenACCUtilsGPU.cpp | 60 +++
.../Dialect/OpenACC/OpenACCUtilsCGTest.cpp | 418 +++++++++++++++++-
.../Dialect/OpenACC/OpenACCUtilsGPUTest.cpp | 59 ++-
6 files changed, 825 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
index 64de4eda7366a..cc854c3238c8a 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
@@ -15,14 +15,19 @@
#define MLIR_DIALECT_OPENACC_OPENACCUTILSCG_H_
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/OpenACC/OpenACCParMapping.h"
#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Value.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Support/LogicalResult.h"
#include "llvm/ADT/SmallVector.h"
#include <optional>
namespace mlir {
namespace acc {
+class OpenACCSupport;
+
/// Get the data layout for an operation.
///
/// Attempts to get the data layout from the operation or its parent module.
@@ -90,6 +95,56 @@ void updateParDimsAttr(Operation *op, GPUParallelDimsAttr attr);
/// Copy parallel dimensions from \p from to \p to.
void copyParDimsAttr(Operation *from, Operation *to);
+/// Tracks aligned byte consumption against a configurable shared memory cap.
+class SharedMemoryBudget {
+public:
+ SharedMemoryBudget(int64_t maxTotalBytes, int64_t initialBytesUsed = 0)
+ : bytesUsed_(initialBytesUsed), maxTotalBytes_(maxTotalBytes) {}
+
+ bool tryAllocate(int64_t bytes);
+ int64_t bytesUsed() const { return bytesUsed_; }
+ int64_t maxTotalBytes() const { return maxTotalBytes_; }
+ void setMaxTotalBytes(int64_t maxTotalBytes) {
+ maxTotalBytes_ = maxTotalBytes;
+ }
+
+ static int64_t alignOffset(int64_t offset);
+
+private:
+ int64_t bytesUsed_ = 0;
+ int64_t maxTotalBytes_ = 0;
+};
+
+/// Sum aligned static_upper_bound_bytes for all acc.gpu_shared_memory in \p
+/// region.
+int64_t sumExistingSharedMemoryBytes(Region ®ion);
+
+/// Resolve the acc.privatize operation associated with a private local.
+PrivatizeOp getPrivatizeOp(PrivateLocalOp privateLocal,
+ ComputeRegionOp computeRegion);
+
+/// Returns the ranked MemRef type used to allocate privatized storage.
+///
+/// \p baseTy is the `baseTy` parameter of `acc.private_type` (the privatized
+/// variable's type).
+MemRefType getPrivateBaseMemRefType(Type baseTy, ModuleOp module);
+
+/// Collect parallel dimensions that govern privatization of \p privateLocal.
+SmallVector<GPUParallelDimAttr>
+collectPrivateLocalParDims(PrivateLocalOp privateLocal,
+ ComputeRegionOp computeRegion);
+
+/// True when \p privateLocal may be placed in shared memory.
+FailureOr<bool> isPrivateLocalSharedMemoryCandidate(
+ PrivateLocalOp privateLocal, ComputeRegionOp computeRegion, ModuleOp module,
+ const ACCToGPUMappingPolicy &policy, OpenACCSupport *support = nullptr);
+
+/// Upper-bound byte size for a shared-memory private_local candidate, or
+/// std::nullopt when not eligible or not statically computable.
+std::optional<int64_t> getPrivateLocalSharedMemoryUpperBoundBytes(
+ PrivateLocalOp privateLocal, ComputeRegionOp computeRegion, ModuleOp module,
+ const ACCToGPUMappingPolicy &policy, OpenACCSupport *support = nullptr);
+
} // namespace acc
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsGPU.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsGPU.h
index aebd15daca060..71e14a78a6925 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsGPU.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsGPU.h
@@ -16,6 +16,8 @@
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Value.h"
+#include "llvm/ADT/DenseMap.h"
#include <optional>
namespace mlir {
@@ -39,6 +41,16 @@ std::optional<gpu::GPUModuleOp>
getOrCreateGPUModule(ModuleOp mod, bool create = true,
llvm::StringRef name = kDefaultGPUModuleName);
+/// Return the launch dimension for \p processor from \p launch, or from
+/// \p dimensionOps when \p launch is null.
+Value getGPUSize(gpu::Processor processor, gpu::LaunchOp launch,
+ const llvm::DenseMap<gpu::Processor, Value> &dimensionOps);
+
+/// Return the thread/block index for \p processor from \p launch, or from
+/// \p indexOps when \p launch is null.
+Value getGPUThreadId(gpu::Processor processor, gpu::LaunchOp launch,
+ const llvm::DenseMap<gpu::Processor, Value> &indexOps);
+
} // namespace acc
} // namespace mlir
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 06605371bdf42..06740c48a1eaf 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -11,10 +11,17 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/OpenACC/OpenACCUtilsCG.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenACC/OpenACCUtilsLoop.h"
+#include "mlir/Dialect/OpenACC/OpenACCUtilsReduction.h"
+#include "mlir/Dialect/OpenACC/OpenACCUtilsType.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/IRMapping.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -206,5 +213,221 @@ void copyParDimsAttr(Operation *from, Operation *to) {
setParDimsAttr(to, getParDimsAttr(from));
}
+int64_t SharedMemoryBudget::alignOffset(int64_t offset) {
+ return (offset + 15) & ~static_cast<int64_t>(15);
+}
+
+bool SharedMemoryBudget::tryAllocate(int64_t bytes) {
+ int64_t aligned = alignOffset(bytesUsed_);
+ if (aligned + bytes > maxTotalBytes_) {
+ return false;
+ }
+ bytesUsed_ = aligned + bytes;
+ return true;
+}
+
+int64_t sumExistingSharedMemoryBytes(Region ®ion) {
+ int64_t total = 0;
+ region.walk([&](GPUSharedMemoryOp op) {
+ int64_t upperBound = op.getStaticUpperBoundBytes();
+ total = SharedMemoryBudget::alignOffset(total) + upperBound;
+ });
+ return total;
+}
+
+PrivatizeOp getPrivatizeOp(PrivateLocalOp privateLocal,
+ ComputeRegionOp computeRegion) {
+ Value value = privateLocal.getPrivatized();
+ if (BlockArgument blockArg = dyn_cast<BlockArgument>(value)) {
+ auto owner = dyn_cast<ComputeRegionOp>(blockArg.getOwner()->getParentOp());
+ value = (owner ? owner : computeRegion).getOperand(blockArg);
+ }
+ PrivatizeOp privatizeOp = value.getDefiningOp<PrivatizeOp>();
+ assert(privatizeOp && "expected privatize op to be the defining op");
+ return privatizeOp;
+}
+
+static bool isThreadXPrivatize(PrivatizeOp privatize) {
+ if (GPUParallelDimsAttr parDimsAttr = privatize.getParDimsAttr())
+ return llvm::any_of(parDimsAttr.getArray(),
+ [](GPUParallelDimAttr d) { return d.isThreadX(); });
+ return false;
+}
+
+MemRefType getPrivateBaseMemRefType(Type baseTy, ModuleOp module) {
+ auto memrefTy = cast<PointerLikeType>(baseTy).getAsMemRefType(module);
+ assert(memrefTy && "private base type must be convertible to memref");
+ return memrefTy;
+}
+
+SmallVector<GPUParallelDimAttr>
+collectPrivateLocalParDims(PrivateLocalOp privateLocal,
+ ComputeRegionOp computeRegion) {
+ SmallVector<GPUParallelDimAttr> parDims;
+ auto parentLoop = privateLocal->getParentOfType<scf::ParallelOp>();
+ while (parentLoop) {
+ if (GPUParallelDimsAttr parDimsAttr = getParDimsAttr(parentLoop))
+ for (GPUParallelDimAttr parDim : parDimsAttr.getArray())
+ insertParDim(parDims, parDim);
+ parentLoop = parentLoop->getParentOfType<scf::ParallelOp>();
+ }
+ if (GPUParallelDimsAttr parDimsAttr = getParDimsAttr(computeRegion))
+ for (GPUParallelDimAttr parDim : parDimsAttr.getArray())
+ insertParDim(parDims, parDim);
+ if (parDims.empty()) {
+ for (GPUParallelDimAttr parDim : computeRegion.getLaunchParDims()) {
+ if (parDim.isAnyBlock())
+ insertParDim(parDims, parDim);
+ }
+ }
+
+ for (Operation *user : privateLocal.getResult().getUsers()) {
+ if (auto accumulateOp = dyn_cast<ReductionAccumulateOp>(user)) {
+ if (accumulateOp.getMemref() == privateLocal.getResult())
+ for (GPUParallelDimAttr parDim : accumulateOp.getParDims().getArray())
+ insertParDim(parDims, parDim);
+ }
+ if (auto combineOp = dyn_cast<ReductionCombineOp>(user)) {
+ if (combineOp.getSrcMemref() == privateLocal.getResult())
+ for (GPUParallelDimAttr parDim : getReductionCombineParDims(combineOp))
+ insertParDim(parDims, parDim);
+ }
+ if (auto combineRegionOp = dyn_cast<ReductionCombineRegionOp>(user)) {
+ if (combineRegionOp.getSrcVar() == privateLocal.getResult())
+ for (GPUParallelDimAttr parDim :
+ getReductionCombineParDims(combineRegionOp))
+ insertParDim(parDims, parDim);
+ }
+ }
+ return parDims;
+}
+
+static FailureOr<std::optional<int64_t>>
+getWorkerPrivateSharedMemoryNumCopies(PrivateLocalOp privateLocal,
+ ComputeRegionOp computeRegion,
+ bool isWorkerPrivate, bool isGangPrivate,
+ OpenACCSupport *support) {
+ if (!isWorkerPrivate || isGangPrivate)
+ return std::optional<int64_t>(1);
+
+ GPUParallelDimAttr threadY =
+ GPUParallelDimAttr::threadYDim(privateLocal.getContext());
+ std::optional<Value> workerArg = computeRegion.getKnownLaunchArg(threadY);
+ if (!workerArg)
+ return std::optional<int64_t>();
+
+ auto workerArgConst = workerArg->getDefiningOp<arith::ConstantIndexOp>();
+ if (!workerArgConst) {
+ if (support) {
+ (void)support->emitNYI(privateLocal.getLoc(),
+ "worker-private variables in shared memory "
+ "require compile-time constant num_workers");
+ return failure();
+ }
+ return std::optional<int64_t>();
+ }
+ return std::optional<int64_t>(workerArgConst.value());
+}
+
+static bool isInsideACCSpecializedRoutine(Operation *op) {
+ auto funcOp = op->getParentOfType<FunctionOpInterface>();
+ return funcOp && isSpecializedAccRoutine(funcOp);
+}
+
+FailureOr<bool>
+isPrivateLocalSharedMemoryCandidate(PrivateLocalOp privateLocal,
+ ComputeRegionOp computeRegion,
+ ModuleOp module,
+ const ACCToGPUMappingPolicy &policy,
+ OpenACCSupport *support) {
+ if (isInsideACCSpecializedRoutine(computeRegion))
+ return false;
+
+ if (isThreadXPrivatize(getPrivatizeOp(privateLocal, computeRegion)))
+ return false;
+
+ bool isReductionAccumulator =
+ llvm::any_of(privateLocal.getResult().getUsers(), [](Operation *user) {
+ return isa<ReductionAccumulateOp>(user);
+ });
+
+ SmallVector<GPUParallelDimAttr> parDims =
+ collectPrivateLocalParDims(privateLocal, computeRegion);
+ bool isGangPrivate =
+ llvm::any_of(parDims, [&](auto parDim) { return policy.isGang(parDim); });
+ bool isWorkerPrivate = llvm::any_of(
+ parDims, [&](auto parDim) { return policy.isWorker(parDim); });
+ bool isVectorPrivate = llvm::any_of(
+ parDims, [&](auto parDim) { return policy.isVector(parDim); });
+
+ auto baseTy = getPrivateBaseMemRefType(
+ cast<PrivateType>(privateLocal.getPrivatized().getType()).getBaseTy(),
+ module);
+
+ bool isBlockLevelPrivate =
+ !isVectorPrivate &&
+ (isGangPrivate ||
+ (isWorkerPrivate && baseTy.getRank() > 0 && !isReductionAccumulator));
+ if (!isBlockLevelPrivate)
+ return false;
+
+ for (int64_t dim : baseTy.getShape())
+ if (dim == ShapedType::kDynamic)
+ return false;
+
+ auto resultMemRefTy = dyn_cast<MemRefType>(privateLocal.getType());
+ if (!resultMemRefTy || !resultMemRefTy.getLayout().isIdentity() ||
+ resultMemRefTy.getMemorySpace())
+ return false;
+
+ if (isGangPrivate && isWorkerPrivate && !isReductionAccumulator)
+ return false;
+
+ FailureOr<std::optional<int64_t>> numCopies =
+ getWorkerPrivateSharedMemoryNumCopies(privateLocal, computeRegion,
+ isWorkerPrivate, isGangPrivate,
+ support);
+ if (failed(numCopies))
+ return failure();
+ return numCopies->has_value();
+}
+
+std::optional<int64_t> getPrivateLocalSharedMemoryUpperBoundBytes(
+ PrivateLocalOp privateLocal, ComputeRegionOp computeRegion, ModuleOp module,
+ const ACCToGPUMappingPolicy &policy, OpenACCSupport *support) {
+ FailureOr<bool> isCandidate = isPrivateLocalSharedMemoryCandidate(
+ privateLocal, computeRegion, module, policy);
+ if (failed(isCandidate) || !*isCandidate)
+ return std::nullopt;
+
+ SmallVector<GPUParallelDimAttr> parDims =
+ collectPrivateLocalParDims(privateLocal, computeRegion);
+ bool isGangPrivate =
+ llvm::any_of(parDims, [&](auto parDim) { return policy.isGang(parDim); });
+ bool isWorkerPrivate = llvm::any_of(
+ parDims, [&](auto parDim) { return policy.isWorker(parDim); });
+
+ FailureOr<std::optional<int64_t>> numCopies =
+ getWorkerPrivateSharedMemoryNumCopies(
+ privateLocal, computeRegion, isWorkerPrivate, isGangPrivate,
+ /*support=*/nullptr);
+ if (failed(numCopies) || !numCopies->has_value())
+ return std::nullopt;
+
+ auto baseTy = getPrivateBaseMemRefType(
+ cast<PrivateType>(privateLocal.getPrivatized().getType()).getBaseTy(),
+ module);
+ std::optional<TypeSizeAndAlignment> elementSizeAndAlignment =
+ getTypeSizeAndAlignment(baseTy.getElementType(), module, support);
+ if (!elementSizeAndAlignment)
+ return std::nullopt;
+
+ int64_t numElements = 1;
+ for (int64_t dim : baseTy.getShape())
+ numElements *= dim;
+ return elementSizeAndAlignment->first.getFixedValue() * numElements *
+ numCopies->value();
+}
+
} // namespace acc
} // namespace mlir
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsGPU.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsGPU.cpp
index fb1e34f514da9..d0f454302a2a8 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsGPU.cpp
@@ -43,5 +43,65 @@ std::optional<gpu::GPUModuleOp> getOrCreateGPUModule(ModuleOp mod, bool create,
return gpuMod;
}
+static Value getGPUSizeFromLaunch(gpu::LaunchOp launch,
+ gpu::Processor processor) {
+ gpu::KernelDim3 gridSize = launch.getGridSize();
+ gpu::KernelDim3 blockSize = launch.getBlockSize();
+ switch (processor) {
+ case gpu::Processor::ThreadX:
+ return blockSize.x;
+ case gpu::Processor::ThreadY:
+ return blockSize.y;
+ case gpu::Processor::ThreadZ:
+ return blockSize.z;
+ case gpu::Processor::BlockX:
+ return gridSize.x;
+ case gpu::Processor::BlockY:
+ return gridSize.y;
+ case gpu::Processor::BlockZ:
+ return gridSize.z;
+ default:
+ return {};
+ }
+}
+
+Value getGPUSize(gpu::Processor processor, gpu::LaunchOp launch,
+ const llvm::DenseMap<gpu::Processor, Value> &dimensionOps) {
+ if (launch)
+ return getGPUSizeFromLaunch(launch, processor);
+ assert(!dimensionOps.empty() && "dimension map is empty");
+ return dimensionOps.lookup(processor);
+}
+
+static Value getGPUThreadIdFromLaunch(gpu::LaunchOp launch,
+ gpu::Processor processor) {
+ gpu::KernelDim3 blockIds = launch.getBlockIds();
+ gpu::KernelDim3 threadIds = launch.getThreadIds();
+ switch (processor) {
+ case gpu::Processor::ThreadX:
+ return threadIds.x;
+ case gpu::Processor::ThreadY:
+ return threadIds.y;
+ case gpu::Processor::ThreadZ:
+ return threadIds.z;
+ case gpu::Processor::BlockX:
+ return blockIds.x;
+ case gpu::Processor::BlockY:
+ return blockIds.y;
+ case gpu::Processor::BlockZ:
+ return blockIds.z;
+ default:
+ return {};
+ }
+}
+
+Value getGPUThreadId(gpu::Processor processor, gpu::LaunchOp launch,
+ const llvm::DenseMap<gpu::Processor, Value> &indexOps) {
+ if (launch)
+ return getGPUThreadIdFromLaunch(launch, processor);
+ assert(!indexOps.empty() && "index map is empty");
+ return indexOps.lookup(processor);
+}
+
} // namespace acc
} // namespace mlir
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
index 75cf385af096f..244845e58e7e8 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
@@ -11,13 +11,17 @@
#include "mlir/Dialect/DLTI/DLTI.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/OpenACC/OpenACCParMapping.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/PatternMatch.h"
#include "gtest/gtest.h"
using namespace mlir;
@@ -32,7 +36,72 @@ class OpenACCUtilsCGTest : public ::testing::Test {
OpenACCUtilsCGTest() : b(&context), loc(UnknownLoc::get(&context)) {
context.loadDialect<acc::OpenACCDialect, arith::ArithDialect,
func::FuncDialect, scf::SCFDialect, gpu::GPUDialect,
- DLTIDialect>();
+ memref::MemRefDialect, DLTIDialect>();
+ }
+
+ static ComputeRegionOp buildComputeRegionWithPrivateLocal(
+ MLIRContext &context, OpBuilder &b, Location loc, ModuleOp module,
+ GPUParallelDimsAttr privatizeParDims, ValueRange launchArgs,
+ PrivateLocalOp &privateLocalOut, PrivatizeOp &privatizeOut) {
+ IRRewriter rewriter(&context);
+ rewriter.setInsertionPointToStart(module.getBody());
+
+ MemRefType memTy = MemRefType::get({4}, b.getI32Type());
+ Type privateTy = PrivateType::get(&context, memTy);
+ privatizeOut = PrivatizeOp::create(rewriter, loc, privateTy, ValueRange{},
+ privatizeParDims);
+
+ Region sourceRegion;
+ Block *srcBlock = new Block();
+ sourceRegion.push_back(srcBlock);
+ BlockArgument privArg = srcBlock->addArgument(privateTy, loc);
+ OpBuilder srcBuilder(&context);
+ srcBuilder.setInsertionPointToStart(srcBlock);
+
+ Value c0 = arith::ConstantIndexOp::create(srcBuilder, loc, 0);
+ Value c1 = arith::ConstantIndexOp::create(srcBuilder, loc, 1);
+ YieldOp::create(srcBuilder, loc);
+
+ // Nest one scf.parallel per privatization dimension so that
+ // collectPrivateLocalParDims recovers them from the parent loops.
+ srcBuilder.setInsertionPoint(srcBlock->getTerminator());
+ for (GPUParallelDimAttr parDim : privatizeParDims.getArray()) {
+ auto par = scf::ParallelOp::create(srcBuilder, loc, ValueRange{c0},
+ ValueRange{c1}, ValueRange{c1});
+ setParDimsAttr(par, GPUParallelDimsAttr::get(&context, {parDim}));
+ srcBuilder.setInsertionPoint(par.getBody()->getTerminator());
+ }
+ PrivateLocalOp::create(srcBuilder, loc, memTy, privArg);
+
+ IRMapping mapping;
+ auto cr = buildComputeRegion(
+ loc, launchArgs, ValueRange{privatizeOut.getResult()},
+ ParallelOp::getOperationName(), sourceRegion, rewriter, mapping,
+ /*output=*/{}, /*kernelFuncName=*/{}, /*kernelModuleName=*/{},
+ /*stream=*/{}, ValueRange{privArg});
+
+ privateLocalOut = {};
+ cr.walk([&](PrivateLocalOp op) { privateLocalOut = op; });
+ return cr;
+ }
+
+ static ComputeRegionOp createEmptyComputeRegion(MLIRContext &context,
+ Location loc,
+ ModuleOp module) {
+ IRRewriter rewriter(&context);
+ rewriter.setInsertionPointToStart(module.getBody());
+
+ Region sourceRegion;
+ Block *block = new Block();
+ sourceRegion.push_back(block);
+ OpBuilder regionBuilder(&context);
+ regionBuilder.setInsertionPointToStart(block);
+ YieldOp::create(regionBuilder, loc);
+
+ IRMapping mapping;
+ return buildComputeRegion(loc, ValueRange{}, ValueRange{},
+ ParallelOp::getOperationName(), sourceRegion,
+ rewriter, mapping);
}
MLIRContext context;
@@ -312,3 +381,350 @@ TEST_F(OpenACCUtilsCGTest, buildComputeRegionWithInputArgsToMap) {
func::ReturnOp::create(rewriter, loc);
}
+
+//===----------------------------------------------------------------------===//
+// SharedMemoryBudget Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, SharedMemoryBudgetAlignAndAllocate) {
+ EXPECT_EQ(SharedMemoryBudget::alignOffset(0), 0);
+ EXPECT_EQ(SharedMemoryBudget::alignOffset(1), 16);
+ EXPECT_EQ(SharedMemoryBudget::alignOffset(16), 16);
+
+ SharedMemoryBudget budget(100);
+ EXPECT_TRUE(budget.tryAllocate(50));
+ EXPECT_EQ(budget.bytesUsed(), 50);
+ EXPECT_TRUE(budget.tryAllocate(34));
+ EXPECT_EQ(budget.bytesUsed(), 98);
+ EXPECT_FALSE(budget.tryAllocate(10));
+}
+
+TEST_F(OpenACCUtilsCGTest, SharedMemoryBudgetInitialBytesUsed) {
+ SharedMemoryBudget budget(64, /*initialBytesUsed=*/48);
+ EXPECT_FALSE(budget.tryAllocate(32));
+ EXPECT_TRUE(budget.tryAllocate(16));
+ EXPECT_EQ(budget.bytesUsed(), 64);
+}
+
+//===----------------------------------------------------------------------===//
+// sumExistingSharedMemoryBytes Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, sumExistingSharedMemoryBytes) {
+ Region region;
+ Block *block = new Block();
+ region.push_back(block);
+ b.setInsertionPointToStart(block);
+
+ MemRefType ty = MemRefType::get({4}, b.getI32Type());
+ GPUSharedMemoryOp::create(b, loc, ty, b.getI64IntegerAttr(1),
+ b.getI64IntegerAttr(100), ValueRange{},
+ IntegerAttr{}, IntegerAttr{});
+ GPUSharedMemoryOp::create(b, loc, ty, b.getI64IntegerAttr(1),
+ b.getI64IntegerAttr(50), ValueRange{},
+ IntegerAttr{}, IntegerAttr{});
+
+ EXPECT_EQ(sumExistingSharedMemoryBytes(region), 162);
+}
+
+//===----------------------------------------------------------------------===//
+// getPrivateBaseMemRefType Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, getPrivateBaseMemRefType) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ MemRefType memTy = MemRefType::get({10}, b.getI64Type());
+
+ EXPECT_EQ(getPrivateBaseMemRefType(memTy, *module), memTy);
+}
+
+//===----------------------------------------------------------------------===//
+// getPrivatizeOp Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, getPrivatizeOpFromHandle) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+
+ MemRefType memTy = MemRefType::get({4}, b.getI32Type());
+ Type privateTy = PrivateType::get(&context, memTy);
+ auto privatize =
+ PrivatizeOp::create(b, loc, privateTy, /*dynamicSizes=*/ValueRange{});
+ auto privateLocal = PrivateLocalOp::create(b, loc, memTy, privatize);
+ auto computeRegion = createEmptyComputeRegion(context, loc, *module);
+
+ EXPECT_EQ(getPrivatizeOp(privateLocal, computeRegion), privatize);
+}
+
+TEST_F(OpenACCUtilsCGTest, getPrivatizeOpFromComputeRegionBlockArg) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr gangDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto bx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, gangDims, ValueRange{bx.getResult()},
+ privateLocal, privatize);
+
+ EXPECT_EQ(getPrivatizeOp(privateLocal, cr), privatize);
+}
+
+//===----------------------------------------------------------------------===//
+// collectPrivateLocalParDims Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, collectPrivateLocalParDimsFromParentLoops) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ IRRewriter rewriter(&context);
+ rewriter.setInsertionPointToStart(module->getBody());
+
+ auto c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ auto bx = ParWidthOp::create(rewriter, loc, c1,
+ GPUParallelDimAttr::blockXDim(&context));
+ auto tx = ParWidthOp::create(rewriter, loc, c1,
+ GPUParallelDimAttr::threadXDim(&context));
+
+ MemRefType memTy = MemRefType::get({4}, b.getI32Type());
+ Type privateTy = PrivateType::get(&context, memTy);
+ auto privatize = PrivatizeOp::create(rewriter, loc, privateTy, ValueRange{});
+
+ Region sourceRegion;
+ Block *srcBlock = new Block();
+ sourceRegion.push_back(srcBlock);
+ BlockArgument privArg = srcBlock->addArgument(privateTy, loc);
+ OpBuilder srcBuilder(&context);
+ srcBuilder.setInsertionPointToStart(srcBlock);
+ auto c0 = arith::ConstantIndexOp::create(srcBuilder, loc, 0);
+ auto c1Body = arith::ConstantIndexOp::create(srcBuilder, loc, 1);
+ YieldOp::create(srcBuilder, loc);
+ srcBuilder.setInsertionPoint(srcBlock->getTerminator());
+
+ // Outer gang (block_x) loop containing an inner vector (thread_x) loop.
+ auto gangLoop = scf::ParallelOp::create(srcBuilder, loc, ValueRange{c0},
+ ValueRange{c1Body}, ValueRange{c1Body});
+ setParDimsAttr(gangLoop, GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)}));
+ srcBuilder.setInsertionPoint(gangLoop.getBody()->getTerminator());
+ auto vectorLoop = scf::ParallelOp::create(
+ srcBuilder, loc, ValueRange{c0}, ValueRange{c1Body}, ValueRange{c1Body});
+ setParDimsAttr(vectorLoop,
+ GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::threadXDim(&context)}));
+ srcBuilder.setInsertionPoint(vectorLoop.getBody()->getTerminator());
+ PrivateLocalOp::create(srcBuilder, loc, memTy, privArg);
+
+ IRMapping mapping;
+ auto cr = buildComputeRegion(
+ loc, ValueRange{bx.getResult(), tx.getResult()},
+ ValueRange{privatize.getResult()}, ParallelOp::getOperationName(),
+ sourceRegion, rewriter, mapping, /*output=*/{}, /*kernelFuncName=*/{},
+ /*kernelModuleName=*/{}, /*stream=*/{}, ValueRange{privArg});
+
+ PrivateLocalOp clonedLocal;
+ cr.walk([&](PrivateLocalOp op) { clonedLocal = op; });
+
+ DefaultACCToGPUMappingPolicy policy;
+ SmallVector<GPUParallelDimAttr> parDims =
+ collectPrivateLocalParDims(clonedLocal, cr);
+ ASSERT_EQ(parDims.size(), 2u);
+ EXPECT_TRUE(policy.isGang(parDims[0]));
+ EXPECT_TRUE(policy.isVector(parDims[1]));
+}
+
+TEST_F(OpenACCUtilsCGTest, collectPrivateLocalParDimsFromLaunchFallback) {
+ // With no enclosing parallel loops, collectPrivateLocalParDims falls back to
+ // the block-level launch dimensions.
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr gangDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto bx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
+
+ MemRefType memTy = MemRefType::get({4}, b.getI32Type());
+ Type privateTy = PrivateType::get(&context, memTy);
+ auto privatize = PrivatizeOp::create(b, loc, privateTy, ValueRange{});
+
+ Region sourceRegion;
+ Block *srcBlock = new Block();
+ sourceRegion.push_back(srcBlock);
+ BlockArgument privArg = srcBlock->addArgument(privateTy, loc);
+ OpBuilder srcBuilder(&context);
+ srcBuilder.setInsertionPointToStart(srcBlock);
+ PrivateLocalOp::create(srcBuilder, loc, memTy, privArg);
+ YieldOp::create(srcBuilder, loc);
+
+ IRRewriter rewriter(&context);
+ rewriter.setInsertionPointToStart(module->getBody());
+ IRMapping mapping;
+ auto cr = buildComputeRegion(
+ loc, ValueRange{bx.getResult()}, ValueRange{privatize.getResult()},
+ ParallelOp::getOperationName(), sourceRegion, rewriter, mapping,
+ /*output=*/{}, /*kernelFuncName=*/{}, /*kernelModuleName=*/{},
+ /*stream=*/{}, ValueRange{privArg});
+ (void)gangDims;
+
+ PrivateLocalOp clonedLocal;
+ cr.walk([&](PrivateLocalOp op) { clonedLocal = op; });
+
+ DefaultACCToGPUMappingPolicy policy;
+ SmallVector<GPUParallelDimAttr> parDims =
+ collectPrivateLocalParDims(clonedLocal, cr);
+ ASSERT_EQ(parDims.size(), 1u);
+ EXPECT_TRUE(policy.isGang(parDims[0]));
+}
+
+TEST_F(OpenACCUtilsCGTest, collectPrivateLocalParDimsFromReductionUsers) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+
+ MemRefType memTy = MemRefType::get({}, b.getI32Type());
+ Type privateTy = PrivateType::get(&context, memTy);
+ auto privatize =
+ PrivatizeOp::create(b, loc, privateTy, /*dynamicSizes=*/ValueRange{});
+ auto privateLocal = PrivateLocalOp::create(b, loc, memTy, privatize);
+ auto computeRegion = createEmptyComputeRegion(context, loc, *module);
+
+ GPUParallelDimsAttr accDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context),
+ GPUParallelDimAttr::threadXDim(&context)});
+ Value partial = arith::ConstantIntOp::create(b, loc, b.getI32Type(), 1);
+ ReductionAccumulateOp::create(b, loc, partial, privateLocal.getResult(),
+ ReductionOperator::AccAdd, accDims);
+
+ DefaultACCToGPUMappingPolicy policy;
+ SmallVector<GPUParallelDimAttr> parDims =
+ collectPrivateLocalParDims(privateLocal, computeRegion);
+ ASSERT_EQ(parDims.size(), 2u);
+ EXPECT_TRUE(policy.isGang(parDims[0]));
+ EXPECT_TRUE(policy.isVector(parDims[1]));
+}
+
+//===----------------------------------------------------------------------===//
+// Shared-memory private_local eligibility Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsCGTest, isPrivateLocalSharedMemoryCandidateGangPrivate) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr gangDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto bx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, gangDims, ValueRange{bx.getResult()},
+ privateLocal, privatize);
+
+ DefaultACCToGPUMappingPolicy policy;
+ FailureOr<bool> isCandidate =
+ isPrivateLocalSharedMemoryCandidate(privateLocal, cr, *module, policy);
+ ASSERT_TRUE(succeeded(isCandidate));
+ EXPECT_TRUE(*isCandidate);
+}
+
+TEST_F(OpenACCUtilsCGTest, isPrivateLocalSharedMemoryCandidateThreadXPrivate) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr vectorDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::threadXDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto tx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::threadXDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, vectorDims, ValueRange{tx.getResult()},
+ privateLocal, privatize);
+
+ DefaultACCToGPUMappingPolicy policy;
+ FailureOr<bool> isCandidate =
+ isPrivateLocalSharedMemoryCandidate(privateLocal, cr, *module, policy);
+ ASSERT_TRUE(succeeded(isCandidate));
+ EXPECT_FALSE(*isCandidate);
+}
+
+TEST_F(OpenACCUtilsCGTest,
+ isPrivateLocalSharedMemoryCandidateWorkerPrivateConstant) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr workerDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::threadYDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto c4 = arith::ConstantIndexOp::create(b, loc, 4);
+ auto bx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
+ auto ty =
+ ParWidthOp::create(b, loc, c4, GPUParallelDimAttr::threadYDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, workerDims,
+ ValueRange{bx.getResult(), ty.getResult()}, privateLocal, privatize);
+
+ DefaultACCToGPUMappingPolicy policy;
+ FailureOr<bool> isCandidate =
+ isPrivateLocalSharedMemoryCandidate(privateLocal, cr, *module, policy);
+ ASSERT_TRUE(succeeded(isCandidate));
+ EXPECT_TRUE(*isCandidate);
+}
+
+TEST_F(OpenACCUtilsCGTest,
+ isPrivateLocalSharedMemoryCandidateWorkerPrivateDynamicFails) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr workerDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::threadYDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto bx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
+ // A non-constant num_workers: the launch operand exists but is not an
+ // arith.constant, which is what triggers the diagnostic / failure path.
+ auto dynNumWorkers = arith::AddIOp::create(b, loc, c1, c1);
+ auto ty = ParWidthOp::create(b, loc, dynNumWorkers,
+ GPUParallelDimAttr::threadYDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, workerDims,
+ ValueRange{bx.getResult(), ty.getResult()}, privateLocal, privatize);
+
+ DefaultACCToGPUMappingPolicy policy;
+ OpenACCSupport support;
+ FailureOr<bool> silent =
+ isPrivateLocalSharedMemoryCandidate(privateLocal, cr, *module, policy);
+ ASSERT_TRUE(succeeded(silent));
+ EXPECT_FALSE(*silent);
+
+ FailureOr<bool> diagnosed = isPrivateLocalSharedMemoryCandidate(
+ privateLocal, cr, *module, policy, &support);
+ EXPECT_TRUE(failed(diagnosed));
+}
+
+TEST_F(OpenACCUtilsCGTest, getPrivateLocalSharedMemoryUpperBoundBytes) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ GPUParallelDimsAttr gangDims = GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)});
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto bx =
+ ParWidthOp::create(b, loc, c1, GPUParallelDimAttr::blockXDim(&context));
+
+ PrivateLocalOp privateLocal;
+ PrivatizeOp privatize;
+ auto cr = buildComputeRegionWithPrivateLocal(
+ context, b, loc, *module, gangDims, ValueRange{bx.getResult()},
+ privateLocal, privatize);
+
+ DefaultACCToGPUMappingPolicy policy;
+ std::optional<int64_t> upperBound =
+ getPrivateLocalSharedMemoryUpperBoundBytes(privateLocal, cr, *module,
+ policy);
+ ASSERT_TRUE(upperBound.has_value());
+ EXPECT_EQ(*upperBound, 16);
+}
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsGPUTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsGPUTest.cpp
index e9d94edde3a92..49090f4d657d1 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsGPUTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsGPUTest.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/OpenACC/OpenACCUtilsGPU.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
@@ -23,7 +24,7 @@ using namespace mlir::acc;
class OpenACCUtilsGPUTest : public ::testing::Test {
protected:
OpenACCUtilsGPUTest() : b(&context), loc(UnknownLoc::get(&context)) {
- context.loadDialect<gpu::GPUDialect>();
+ context.loadDialect<arith::ArithDialect, gpu::GPUDialect>();
}
MLIRContext context;
@@ -87,3 +88,59 @@ TEST_F(OpenACCUtilsGPUTest, getOrCreateGPUModuleEmptyNameUsesDefault) {
ASSERT_TRUE(gpuMod.has_value());
EXPECT_EQ(gpuMod->getName(), kDefaultGPUModuleName);
}
+
+//===----------------------------------------------------------------------===//
+// getGPUSize / getGPUThreadId Tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(OpenACCUtilsGPUTest, getGPUSizeFromLaunch) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+
+ auto c128 = arith::ConstantIndexOp::create(b, loc, 128);
+ auto c4 = arith::ConstantIndexOp::create(b, loc, 4);
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto launch = gpu::LaunchOp::create(b, loc, c4, c1, c1, c128, c1, c1);
+ llvm::DenseMap<gpu::Processor, Value> dimensionOps;
+
+ // getGPUSize returns the in-kernel size SSA values (block arguments), which
+ // are the values exposed by getBlockSize()/getGridSize().
+ EXPECT_EQ(getGPUSize(gpu::Processor::ThreadX, launch, dimensionOps),
+ launch.getBlockSize().x);
+ EXPECT_EQ(getGPUSize(gpu::Processor::BlockX, launch, dimensionOps),
+ launch.getGridSize().x);
+}
+
+TEST_F(OpenACCUtilsGPUTest, getGPUSizeFromDimensionMap) {
+ auto c128 = arith::ConstantIndexOp::create(b, loc, 128);
+ llvm::DenseMap<gpu::Processor, Value> dimensionOps;
+ dimensionOps[gpu::Processor::ThreadX] = c128.getResult();
+
+ EXPECT_EQ(getGPUSize(gpu::Processor::ThreadX, nullptr, dimensionOps),
+ c128.getResult());
+}
+
+TEST_F(OpenACCUtilsGPUTest, getGPUThreadIdFromLaunch) {
+ OwningOpRef<ModuleOp> module = ModuleOp::create(b, loc);
+ b.setInsertionPointToStart(module->getBody());
+
+ auto c128 = arith::ConstantIndexOp::create(b, loc, 128);
+ auto c4 = arith::ConstantIndexOp::create(b, loc, 4);
+ auto c1 = arith::ConstantIndexOp::create(b, loc, 1);
+ auto launch = gpu::LaunchOp::create(b, loc, c4, c1, c1, c128, c1, c1);
+ llvm::DenseMap<gpu::Processor, Value> indexOps;
+
+ EXPECT_EQ(getGPUThreadId(gpu::Processor::ThreadX, launch, indexOps),
+ launch.getThreadIds().x);
+ EXPECT_EQ(getGPUThreadId(gpu::Processor::BlockX, launch, indexOps),
+ launch.getBlockIds().x);
+}
+
+TEST_F(OpenACCUtilsGPUTest, getGPUThreadIdFromIndexMap) {
+ auto c7 = arith::ConstantIndexOp::create(b, loc, 7);
+ llvm::DenseMap<gpu::Processor, Value> indexOps;
+ indexOps[gpu::Processor::BlockX] = c7.getResult();
+
+ EXPECT_EQ(getGPUThreadId(gpu::Processor::BlockX, nullptr, indexOps),
+ c7.getResult());
+}
>From 25652c08d21cceecb313ee92fcb5b753c1f9ca82 Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Mon, 13 Jul 2026 14:30:26 -0700
Subject: [PATCH 2/3] Fix formatting
---
.../Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp | 28 ++++++++-----------
.../Dialect/OpenACC/OpenACCUtilsCGTest.cpp | 9 +++---
2 files changed, 16 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 06740c48a1eaf..1c33f6fee48e8 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -302,11 +302,9 @@ collectPrivateLocalParDims(PrivateLocalOp privateLocal,
return parDims;
}
-static FailureOr<std::optional<int64_t>>
-getWorkerPrivateSharedMemoryNumCopies(PrivateLocalOp privateLocal,
- ComputeRegionOp computeRegion,
- bool isWorkerPrivate, bool isGangPrivate,
- OpenACCSupport *support) {
+static FailureOr<std::optional<int64_t>> getWorkerPrivateSharedMemoryNumCopies(
+ PrivateLocalOp privateLocal, ComputeRegionOp computeRegion,
+ bool isWorkerPrivate, bool isGangPrivate, OpenACCSupport *support) {
if (!isWorkerPrivate || isGangPrivate)
return std::optional<int64_t>(1);
@@ -334,12 +332,9 @@ static bool isInsideACCSpecializedRoutine(Operation *op) {
return funcOp && isSpecializedAccRoutine(funcOp);
}
-FailureOr<bool>
-isPrivateLocalSharedMemoryCandidate(PrivateLocalOp privateLocal,
- ComputeRegionOp computeRegion,
- ModuleOp module,
- const ACCToGPUMappingPolicy &policy,
- OpenACCSupport *support) {
+FailureOr<bool> isPrivateLocalSharedMemoryCandidate(
+ PrivateLocalOp privateLocal, ComputeRegionOp computeRegion, ModuleOp module,
+ const ACCToGPUMappingPolicy &policy, OpenACCSupport *support) {
if (isInsideACCSpecializedRoutine(computeRegion))
return false;
@@ -384,9 +379,8 @@ isPrivateLocalSharedMemoryCandidate(PrivateLocalOp privateLocal,
return false;
FailureOr<std::optional<int64_t>> numCopies =
- getWorkerPrivateSharedMemoryNumCopies(privateLocal, computeRegion,
- isWorkerPrivate, isGangPrivate,
- support);
+ getWorkerPrivateSharedMemoryNumCopies(
+ privateLocal, computeRegion, isWorkerPrivate, isGangPrivate, support);
if (failed(numCopies))
return failure();
return numCopies->has_value();
@@ -408,9 +402,9 @@ std::optional<int64_t> getPrivateLocalSharedMemoryUpperBoundBytes(
parDims, [&](auto parDim) { return policy.isWorker(parDim); });
FailureOr<std::optional<int64_t>> numCopies =
- getWorkerPrivateSharedMemoryNumCopies(
- privateLocal, computeRegion, isWorkerPrivate, isGangPrivate,
- /*support=*/nullptr);
+ getWorkerPrivateSharedMemoryNumCopies(privateLocal, computeRegion,
+ isWorkerPrivate, isGangPrivate,
+ /*support=*/nullptr);
if (failed(numCopies) || !numCopies->has_value())
return std::nullopt;
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
index 244845e58e7e8..abf2294db114d 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsCGTest.cpp
@@ -504,10 +504,11 @@ TEST_F(OpenACCUtilsCGTest, collectPrivateLocalParDimsFromParentLoops) {
srcBuilder.setInsertionPoint(srcBlock->getTerminator());
// Outer gang (block_x) loop containing an inner vector (thread_x) loop.
- auto gangLoop = scf::ParallelOp::create(srcBuilder, loc, ValueRange{c0},
- ValueRange{c1Body}, ValueRange{c1Body});
- setParDimsAttr(gangLoop, GPUParallelDimsAttr::get(
- &context, {GPUParallelDimAttr::blockXDim(&context)}));
+ auto gangLoop = scf::ParallelOp::create(
+ srcBuilder, loc, ValueRange{c0}, ValueRange{c1Body}, ValueRange{c1Body});
+ setParDimsAttr(gangLoop,
+ GPUParallelDimsAttr::get(
+ &context, {GPUParallelDimAttr::blockXDim(&context)}));
srcBuilder.setInsertionPoint(gangLoop.getBody()->getTerminator());
auto vectorLoop = scf::ParallelOp::create(
srcBuilder, loc, ValueRange{c0}, ValueRange{c1Body}, ValueRange{c1Body});
>From 09c9bcc80a5c5e944c11f13b8adafb168b2eecbf Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Mon, 13 Jul 2026 14:52:01 -0700
Subject: [PATCH 3/3] Avoid hardcoded alignment
---
mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h | 14 ++++++++++++--
mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp | 11 +++++++----
2 files changed, 19 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
index cc854c3238c8a..01be7760aad77 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsCG.h
@@ -98,17 +98,27 @@ void copyParDimsAttr(Operation *from, Operation *to);
/// Tracks aligned byte consumption against a configurable shared memory cap.
class SharedMemoryBudget {
public:
+ /// Default allocation alignment (bytes).
+ static constexpr int64_t kDefaultAlignmentBytes = 16;
+
SharedMemoryBudget(int64_t maxTotalBytes, int64_t initialBytesUsed = 0)
: bytesUsed_(initialBytesUsed), maxTotalBytes_(maxTotalBytes) {}
- bool tryAllocate(int64_t bytes);
+ /// Reserve \p bytes, rounding the current offset up to \p alignment first.
+ /// Returns false without mutating state if the reservation would exceed the
+ /// cap. \p alignment must be a power of two.
+ bool tryAllocate(int64_t bytes,
+ int64_t alignment = kDefaultAlignmentBytes);
int64_t bytesUsed() const { return bytesUsed_; }
int64_t maxTotalBytes() const { return maxTotalBytes_; }
void setMaxTotalBytes(int64_t maxTotalBytes) {
maxTotalBytes_ = maxTotalBytes;
}
- static int64_t alignOffset(int64_t offset);
+ /// Round \p offset up to the next multiple of \p alignment, which must be a
+ /// power of two.
+ static int64_t alignOffset(int64_t offset,
+ int64_t alignment = kDefaultAlignmentBytes);
private:
int64_t bytesUsed_ = 0;
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
index 1c33f6fee48e8..1bc79531523aa 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsCG.cpp
@@ -25,6 +25,7 @@
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/MathExtras.h"
namespace mlir {
namespace acc {
@@ -213,12 +214,14 @@ void copyParDimsAttr(Operation *from, Operation *to) {
setParDimsAttr(to, getParDimsAttr(from));
}
-int64_t SharedMemoryBudget::alignOffset(int64_t offset) {
- return (offset + 15) & ~static_cast<int64_t>(15);
+int64_t SharedMemoryBudget::alignOffset(int64_t offset, int64_t alignment) {
+ assert(alignment > 0 && llvm::isPowerOf2_64(alignment) &&
+ "alignment must be a power of two");
+ return (offset + alignment - 1) & ~(alignment - 1);
}
-bool SharedMemoryBudget::tryAllocate(int64_t bytes) {
- int64_t aligned = alignOffset(bytesUsed_);
+bool SharedMemoryBudget::tryAllocate(int64_t bytes, int64_t alignment) {
+ int64_t aligned = alignOffset(bytesUsed_, alignment);
if (aligned + bytes > maxTotalBytes_) {
return false;
}
More information about the Mlir-commits
mailing list