[Mlir-commits] [mlir] [mlir][acc] Introduce acc to gpu codegen pass (PR #209606)
Valentin Clement バレンタイン クレメン
llvmlistbot at llvm.org
Tue Jul 14 13:46:54 PDT 2026
================
@@ -0,0 +1,3757 @@
+//===- ACCCGToGPU.cpp - Lower acc.compute_region to gpu.launch ------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass lowers `acc.compute_region` to the GPU dialect. For host-side
+// kernels it wraps the region in `gpu.launch`; for specialized acc routines
+// already inside a `gpu.func`, the body is lowered in place without emitting
+// a launch.
+//
+// Overview:
+// ---------
+// `acc.compute_region` is the compute-body representation produced after
+// OpenACC compute constructs are decomposed and parallelism has been assigned.
+// This pass is the final ACC-to-GPU lowering step for that body: it converts
+// nested `scf.parallel` / `scf.for` loops marked with `acc.par_dims` into GPU
+// block and thread parallelism, materializes privatization and reductions for
+// the device, inserts synchronization where shared state is observed across
+// threads, and erases the ACC scaffolding (`acc.compute_region`,
+// `acc.par_width`).
+//
+// Transformations:
+// ----------------
+// 1. Launch creation: outside a `gpu.func`, each `acc.compute_region` becomes a
+// `gpu.launch` whose grid and block sizes come from `acc.par_width` launch
+// operands (defaulting to 1). Kernel/module name attributes are preserved.
+// Inside a `gpu.func` (specialized acc routine), no launch is emitted.
+//
+// 2. Parallel loops: `scf.parallel` with a single `acc.par_dims` entry is
+// mapped to the corresponding GPU dimension (`block_*` or `thread_*`).
+// Sequential dimensions remain as `scf.parallel`/`scf.for` loops in the
+// generated kernel body.
+//
+// 3. Privatization: `acc.privatize` / `acc.private_local` storage is
+// materialized as one of: a per-thread `memref.alloca` (thread-private
+// arrays within the stack budget), an `acc.gpu_shared_memory` buffer
+// (gang-/worker-private arrays that fit the shared-memory budget), or a
+// `memref.alloc` whose pointer is broadcast to the block through a small
+// shared-memory slot (the data lives in global memory; shared memory only
+// holds the broadcast pointer).
+//
+// 4. Predication: `acc.predicate_region` becomes `scf.if` guarded by active
+// thread/block indices derived from `acc.par_dims` and launch dimensions.
+//
+// 5. Reductions: `acc.reduction_*` ops are lowered to GPU reduction and
+// synchronization primitives according to each reduction's parallel
+// dimensions and accumulator storage class.
+//
+// Example:
+// --------
+// Before:
+// %c128 = arith.constant 128 : index
+// %tx = acc.par_width %c128 {par_dim = #acc.par_dim<thread_x>}
+// acc.compute_region launch(%arg0 = %tx) {
+// %c0 = arith.constant 0 : index
+// %c1 = arith.constant 1 : index
+// scf.parallel (%iv) = (%c0) to (%c128) step (%c1) {
+// ...
+// scf.reduce
+// } {acc.par_dims = #acc<par_dims[thread_x]>}
+// acc.yield
+// } {origin = "acc.parallel"}
+//
+// After:
+// gpu.launch blocks(%bidx, %bidy, %bidz) in (%gdimx = %c1, ...)
+// threads(%tidx, %tidy, %tidz) in (%bdimx = %c128, ...) {
+// ...
+// }
+//
+// Requirements:
+// -------------
+// - Must run on a GPU device type (`device-type` option); host and multicore
+// targets are rejected.
+// - Input must already be in the `acc.compute_region` form: nested SCF loops
+// carry `acc.par_dims`, privatization is expressed via `acc.privatize` /
+// `acc.private_local`, and reductions use the `acc.reduction_*` ops.
+// - Each `scf.parallel` processed by this pass is expected to have exactly
+// one parallel dimension and one induction variable.
+// - For acc routines, the `acc.compute_region` must live inside a `gpu.func`
+// in the GPU module.
+// - Uses `acc::OpenACCSupport` for NYI reporting and compiler remarks.
+// - Pass options: `max-workgroup-shared-memory`, `max-thread-private-stack`,
+// and `subgroup-size` (used for reductions and block-dimension alignment).
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/GPU/Utils/GPUUtils.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.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/OpenACC/OpenACCUtilsCG.h"
+#include "mlir/Dialect/OpenACC/OpenACCUtilsGPU.h"
+#include "mlir/Dialect/OpenACC/OpenACCUtilsReduction.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/CallInterfaces.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/Debug.h"
+#include <algorithm>
+#include <optional>
+#include <utility>
+
+namespace mlir {
+namespace acc {
+#define GEN_PASS_DEF_ACCCGTOGPU
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
+} // namespace acc
+} // namespace mlir
+
+#define DEBUG_TYPE "acc-cg-to-gpu"
+
+namespace {
+using namespace mlir;
+using namespace mlir::acc;
+
+enum class PrivateMemScope { Thread, Worker, Gang, None };
+
+/// Device label used in compiler remarks (e.g. "NVIDIA GPU").
+static std::string getDeviceRemarkQualifier(DeviceType deviceType) {
+ switch (deviceType) {
+ case DeviceType::None:
+ case DeviceType::Star:
+ case DeviceType::Default:
+ return "GPU";
+ default: {
+ std::string name;
+ llvm::StringRef deviceName = stringifyDeviceType(deviceType);
+ name.reserve(deviceName.size());
+ for (char c : deviceName) {
+ name.push_back(llvm::toUpper(c));
+ }
+ return name + " GPU";
+ }
+ }
+}
+
+/// True when \p op is inside a specialized acc routine function.
+static bool isInsideACCSpecializedRoutine(Operation *op) {
+ FunctionOpInterface funcOp = op->getParentOfType<FunctionOpInterface>();
+ return funcOp && acc::isSpecializedAccRoutine(funcOp);
+}
+
+/// Maps an acc.routine's parallelism clauses to a GPU parallel dimension.
+static GPUParallelDimAttr
+getAccRoutineParDim(RoutineOp routineOp, MLIRContext *ctx,
+ const ACCToGPUMappingPolicy &policy) {
+ if (routineOp.getGangDimValue() ||
+ routineOp.getGangDimValue(DeviceType::Nvidia)) {
+ int64_t gangDimValue = routineOp.getGangDimValue(DeviceType::Nvidia)
+ ? *routineOp.getGangDimValue(DeviceType::Nvidia)
+ : *routineOp.getGangDimValue();
+ ParLevel gangLevel = getGangParLevel(gangDimValue);
+ return policy.gangDim(ctx, gangLevel);
+ }
+ if (routineOp.hasGang() || routineOp.hasGang(DeviceType::Nvidia))
+ return policy.gangDim(ctx, ParLevel::gang_dim1);
+ if (routineOp.hasWorker() || routineOp.hasWorker(DeviceType::Nvidia))
+ return policy.workerDim(ctx);
+ if (routineOp.hasVector() || routineOp.hasVector(DeviceType::Nvidia))
+ return policy.vectorDim(ctx);
+ return policy.seqDim(ctx);
+}
+
+/// Looks up the acc.routine symbol associated with \p funcOp.
+static RoutineOp getRoutineOpForAccRoutineFunction(FunctionOpInterface funcOp,
+ const SymbolTable &symTab) {
+ if (isSpecializedAccRoutine(funcOp)) {
+ SpecializedRoutineAttr attr = funcOp->getAttrOfType<SpecializedRoutineAttr>(
+ getSpecializedRoutineAttrName());
+ return symTab.lookup<RoutineOp>(attr.getRoutine().getLeafReference());
+ }
+ RoutineInfoAttr routineInfo =
+ funcOp->getAttrOfType<RoutineInfoAttr>(getRoutineInfoAttrName());
+ if (!routineInfo || routineInfo.getAccRoutines().empty())
+ return nullptr;
+ return symTab.lookup<RoutineOp>(
+ routineInfo.getAccRoutines().front().getLeafReference());
+}
+
+/// Returns the parallelism level of a specialized acc routine function.
+static GPUParallelDimAttr
+getSpecializedRoutineDim(FunctionOpInterface funcOp,
+ const ACCToGPUMappingPolicy &policy) {
+ SpecializedRoutineAttr specAttr =
+ funcOp->getAttrOfType<SpecializedRoutineAttr>(
+ getSpecializedRoutineAttrName());
+ assert(specAttr && "expected specialized routine attribute");
+ return policy.map(funcOp->getContext(), specAttr.getLevel().getValue());
+}
+
+/// Returns the parallelism dimension of a callee acc routine, if any.
+static GPUParallelDimAttr
+getAccRoutineCallParDim(CallOpInterface callOp,
+ const ACCToGPUMappingPolicy &policy) {
+ std::optional<CallInterfaceCallable> callee = callOp.getCallableForCallee();
+ if (!callee)
+ return nullptr;
+ SymbolRefAttr calleeSymbolRef = dyn_cast<SymbolRefAttr>(*callee);
+ if (!calleeSymbolRef)
+ return nullptr;
+ ModuleOp moduleOp = callOp->getParentOfType<ModuleOp>();
+ if (!moduleOp)
+ return nullptr;
+
+ SymbolTable symTab(moduleOp);
+ FunctionOpInterface funcOp =
+ symTab.lookup<FunctionOpInterface>(calleeSymbolRef.getLeafReference());
+ if (!funcOp)
+ return nullptr;
+
+ if (isSpecializedAccRoutine(funcOp))
+ return getSpecializedRoutineDim(funcOp, policy);
+ if (RoutineOp routineOp = getRoutineOpForAccRoutineFunction(funcOp, symTab))
+ return getAccRoutineParDim(routineOp, funcOp.getContext(), policy);
+ return nullptr;
+}
+
+/// Collects parallel dimensions from enclosing loops and the compute region.
+static SmallVector<GPUParallelDimAttr> getAncestorParDims(Operation *op) {
+ SmallVector<GPUParallelDimAttr> parDimsArray;
+ scf::ParallelOp parentLoop = op->getParentOfType<scf::ParallelOp>();
+ while (parentLoop) {
+ if (GPUParallelDimsAttr parDimsAttr = getParDimsAttr(parentLoop))
+ for (GPUParallelDimAttr parDim : parDimsAttr.getArray())
+ insertParDim(parDimsArray, parDim);
+ parentLoop = parentLoop->getParentOfType<scf::ParallelOp>();
+ }
+
+ ComputeRegionOp computeRegion = op->getParentOfType<ComputeRegionOp>();
+ assert(computeRegion && "missing enclosing acc.compute_region");
+ if (GPUParallelDimsAttr parDimsAttr = getParDimsAttr(computeRegion))
+ for (GPUParallelDimAttr parDim : parDimsAttr.getArray())
+ insertParDim(parDimsArray, parDim);
+ return parDimsArray;
+}
+
+/// Strips index casts to reach the underlying defining value.
+static Value stripIndexCastsFromValue(Value x) {
+ Operation *op = x.getDefiningOp();
+ if (!op)
+ return x;
+ while (arith::IndexCastOp castOp = dyn_cast<arith::IndexCastOp>(op)) {
+ op = castOp->getOperand(0).getDefiningOp();
+ if (!op)
+ return x;
+ }
+ return op->getResult(0);
+}
+
+/// Extracts a compile-time integer constant from \p x, when known.
+static FailureOr<int64_t> extractIntConst(Value x,
+ bool stripIndexCasts = false) {
+ if (stripIndexCasts)
+ x = stripIndexCastsFromValue(x);
+ Operation *op = x.getDefiningOp();
+ if (op) {
+ if (arith::ConstantIntOp constOp = dyn_cast<arith::ConstantIntOp>(op)) {
+ assert(constOp.getType().getIntOrFloatBitWidth() <= 64);
+ return constOp.value();
+ }
+ if (arith::ConstantIndexOp constOp = dyn_cast<arith::ConstantIndexOp>(op))
+ return constOp.value();
+ }
+ return failure();
+}
+
+/// True when \p x is a constant equal to \p y (modulo index casts).
+static bool sameEffectiveValue(Value x, int64_t y) {
+ x = stripIndexCastsFromValue(x);
+ FailureOr<int64_t> conX = extractIntConst(x);
+ if (failed(conX))
+ return false;
+ return *conX == y;
+}
+
+/// Continues tracking a memref through view-like and partial-access ops.
+static bool getPassThroughResults(Operation *userOp, Value trackedOperand,
+ SmallVectorImpl<Value> &passThroughResults) {
+ if (ViewLikeOpInterface viewLikeOp = dyn_cast<ViewLikeOpInterface>(userOp)) {
+ if (viewLikeOp.getViewSource() == trackedOperand) {
+ passThroughResults.push_back(viewLikeOp.getViewDest());
+ return true;
+ }
+ return false;
+ }
+
+ // Partial-entity accesses (e.g. array element or field access) forward the
+ // base entity through to their results, so treat them as pass-through when
+ // the base entity is the value being tracked.
+ if (acc::PartialEntityAccessOpInterface partialAccess =
+ dyn_cast<acc::PartialEntityAccessOpInterface>(userOp)) {
+ if (partialAccess.getBaseEntity() == trackedOperand) {
+ passThroughResults.append(userOp->result_begin(), userOp->result_end());
+ return true;
+ }
+ return false;
+ }
+ return false;
+}
+
+/// Skips memref view/cast chains to reach the underlying buffer.
+static Value unwrapMemRefConversion(Value v) {
+ while (Operation *op = v.getDefiningOp()) {
+ if (ViewLikeOpInterface viewLike = dyn_cast<ViewLikeOpInterface>(op)) {
+ if (isa<MemRefType>(viewLike.getViewSource().getType()) ||
+ isa<MemRefType>(viewLike.getViewDest().getType())) {
+ v = viewLike.getViewSource();
+ continue;
+ }
+ }
+ break;
+ }
+ return v;
+}
+
+/// Casts between pointer-like private types when lowering requires it.
+static Value castPointerLikeTypeIfNeeded(OpBuilder &builder, Location loc,
+ Value value, Type resultType) {
+ if (value.getType() == resultType)
+ return value;
+ if (PointerLikeType ptrLike = dyn_cast<PointerLikeType>(value.getType())) {
+ if (Value casted = ptrLike.genCast(builder, loc, value, resultType))
+ return casted;
+ }
+ if (PointerLikeType ptrLike = dyn_cast<PointerLikeType>(resultType)) {
+ if (Value casted = ptrLike.genCast(builder, loc, value, resultType))
+ return casted;
+ }
+ emitError(loc) << "unsupported pointer-like type cast from "
+ << value.getType() << " to " << resultType;
+ return value;
+}
+
+/// Returns the sole user of \p v, or null if it has zero or multiple uses.
+static Operation *getOnlyUser(Value v) {
+ if (!v.hasOneUse())
+ return nullptr;
+ return *v.user_begin();
+}
+
+/// True when \p privatize is privatized at thread_x parallelism.
+static bool isThreadXPrivatize(PrivatizeOp privatize) {
+ if (GPUParallelDimsAttr parDimsAttr = privatize.getParDimsAttr())
+ return llvm::any_of(parDimsAttr.getArray(),
+ [](GPUParallelDimAttr d) { return d.isThreadX(); });
+ return false;
+}
+
+/// Emits a workgroup-wide GPU barrier.
+static void emitGPUBarrierWorkgroup(OpBuilder &builder, Location loc) {
+ gpu::BarrierOp::create(builder, loc);
+}
+
+/// Emits a subgroup-scoped GPU barrier.
+static void emitGPUBarrierSubgroup(OpBuilder &builder, Location loc) {
+ gpu::BarrierOp::create(builder, loc, /*address_spaces=*/ArrayAttr{},
+ /*named_barrier=*/Value{},
+ gpu::BarrierScope::Subgroup);
+}
+
+/// Lowers a single `acc.compute_region` to GPU dialect IR.
+class ACCCGToGPULowering {
+public:
+ explicit ACCCGToGPULowering(acc::ComputeRegionOp computeRegion,
+ RewriterBase &rewriter,
+ acc::OpenACCSupport &accSupport,
+ const ACCCGToGPUOptions &options)
+ : rewriter(rewriter), computeRegion(computeRegion),
+ accSupport(accSupport), options(options),
+ sharedMemBudget(
+ options.maxWorkgroupSharedMemory,
+ sumExistingSharedMemoryBytes(computeRegion.getRegion())) {}
+
+ /// Main entry point: emit launch (if needed) and lower the region body.
+ LogicalResult rewrite();
+
+ gpu::LaunchOp getLaunch() const { return launch; }
+
+ bool hasFailed = false;
+ bool insideAccumulateGridStride = false;
+ Value reductionSharedBuf;
+ // Reduction-accumulator slot (memref) -> the block-reduced value stored into
+ // it; lets a block combine use the register instead of reloading.
+ llvm::DenseMap<Value, Value> reductionAccumValue;
+ // Combine reloads recorded before accumulates are lowered, patched up after.
+ llvm::SmallVector<std::pair<Value, memref::LoadOp>> pendingCombineReloads;
+
+private:
+ /// Lower a parallel loop to the GPU dimension given by its `acc.par_dims`.
+ void processParallelOp(scf::ParallelOp parallelOp);
+ /// Lower a sequential loop, including any required post-loop barriers.
+ template <typename LoopOp>
+ void processSeqLoop(LoopOp loopOp);
+ /// Lower an `acc.predicate_region` to a predicated `scf.if`.
+ void processPredicateRegion(acc::PredicateRegionOp interOp);
+ /// Materialize storage for an `acc.private_local`.
+ void
+ processPrivateLocal(acc::PrivateLocalOp privateLocal,
+ std::optional<int64_t> sharedMemCopies = std::nullopt);
+ /// Lower an `acc.privatize` to device storage.
+ Value processPrivatize(acc::PrivatizeOp privatize);
+ /// Clone and lower an `scf.execute_region`.
+ void processExecuteRegion(scf::ExecuteRegionOp op);
+ /// Lower `acc.reduction_accumulate`.
+ void processAccumulateOp(acc::ReductionAccumulateOp op);
+ /// Lower `acc.reduction_accumulate_array`.
+ void processAccumulateArrayOp(acc::ReductionAccumulateArrayOp op);
+ /// Lower `acc.reduction_init`.
+ void processReductionOp(acc::ReductionInitOp op);
+ /// Lower `acc.reduction_combine`.
+ void processReductionCombineOp(acc::ReductionCombineOp op);
+ /// Lower `acc.reduction_combine_region`.
+ void processCombineRegionOp(acc::ReductionCombineRegionOp op);
+ /// Clone a leaf operation into the lowered region.
+ void processGenericOp(Operation *op);
+ /// Clone and recursively lower an operation with nested regions.
+ void processGenericOpWithRegions(Operation *op);
+ /// Dispatch lowering for one operation in the compute-region body.
+ void processOp(Operation *op);
+
+ /// Emit an atomic reduction update to \p memref.
+ void constructAtomicAccumulation(Location loc, Value memref,
+ ValueRange indices, Value input,
+ arith::AtomicRMWKind kind);
+
+ /// Map an ACC reduction operator to an atomic RMW kind.
+ FailureOr<arith::AtomicRMWKind> getReductionKind(acc::ReductionOperator redOp,
+ Type type, Location loc);
+
+ /// Split launch dimensions into those that execute \p op and those that do
+ /// not, for predication and barrier placement.
+ std::pair<SmallVector<mlir::acc::GPUParallelDimAttr>,
+ SmallVector<mlir::acc::GPUParallelDimAttr>>
+ computeActiveAndInactiveParDims(Operation *op, Block *block);
+
+ /// Build a predicate that is true only on inactive parallel dimensions.
+ Value
+ emitPredicate(Location loc,
+ SmallVector<mlir::acc::GPUParallelDimAttr> &inactiveParDims);
+
+ /// True when \p privateLocal may be placed in shared memory; returns the
+ /// number of copies needed, or nullopt if ineligible.
+ std::optional<int64_t>
+ isEligibleForSharedMemory(acc::PrivateLocalOp privateLocal,
+ MemRefType baseTy);
+
+ /// Reserve \p bytes from the shared-memory budget.
+ bool tryAllocateSharedMemory(int64_t bytes);
+
+ /// Element size in bytes for \p elementType .
+ int64_t getElementSizeInBytes(Location loc, Type elementType) const;
+
+ /// True when a static privatization fits in the per-thread stack budget.
+ bool canUseStackAlloca(MemRefType baseTy, Location loc,
+ int64_t maxThreadPrivateStack) const;
+
+ /// Emit a barrier scoped to the parallel dimensions in \p parDimsAttr.
+ void createBarrier(Location loc, mlir::acc::GPUParallelDimsAttr parDimsAttr);
+
+ /// Emit a per-row (per-worker) barrier.
+ /// Runtime branch on blockDim.y == 1 (workgroup-wide); compile-time choice
+ /// between gpu.barrier scope<subgroup> (staticBlockDimX <= subgroupSize)
+ /// and a named gpu.barrier (staticBlockDimX > subgroupSize) with tid.y+1.
+ void createPerRowBarrier(Location loc);
+
+ /// Insert barriers after a sequential loop when shared private state must be
+ /// visible to later loops.
+ void createBarrierAfterSeqLoop(Operation *loopOp);
+
+ /// Flush any deferred post-loop barriers that precede \p beforeOp.
+ void flushDeferredBarriersBefore(Operation *beforeOp);
+
+ /// True when \p loopOp may write shared memory read by a later sibling loop.
+ bool mayWriteSharedMemory(Operation *loopOp);
+
+ /// Parallelism scope (thread, worker, or gang) of a privatized variable.
+ PrivateMemScope getPrivateMemScope(acc::PrivatizeOp privatizeOp);
+
+ /// Parallelism scope of the private buffer backing \p memref.
+ PrivateMemScope getPrivateScopeForMemref(Value memref);
+
+ /// `acc.privatize` that materialized the private buffer for \p memref.
+ acc::PrivatizeOp getPrivatizeForMemref(Value memref);
+
+ /// Whether a predicate region needs a barrier before stores that will be read
+ /// by a later parallel loop over the same private memory.
+ PrivateMemScope needsPreStoreReuseBarrier(acc::PredicateRegionOp interOp);
+
+ /// Emit `gpu.all_reduce` for a reduction partial.
+ void createGPUAllReduceOp(Location loc, Value input, Value memref,
+ arith::AtomicRMWKind kind,
+ mlir::acc::GPUParallelDimsAttr parDimsAttr,
+ ValueRange indices = {});
+
+ /// Finish lowering a deferred `acc.reduction_accumulate`.
+ void postprocessAccumulateOp(acc::ReductionAccumulateOp op);
+
+ /// Finish lowering reductions attached to a parallel loop.
+ void postprocessLoopReduction(scf::ParallelOp parLoop);
+
+ /// Populate block/thread id and grid/block dimension maps for device
+ /// routines.
+ static void
+ createForAllDimensions(RewriterBase &rewriter, Location loc,
+ llvm::DenseMap<gpu::Processor, Value> &ids,
+ llvm::DenseMap<gpu::Processor, Value> &dims) {
+ ids[gpu::Processor::BlockX] = gpu::BlockIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::x);
+ ids[gpu::Processor::BlockY] = gpu::BlockIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::y);
+ ids[gpu::Processor::BlockZ] = gpu::BlockIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::z);
+ ids[gpu::Processor::ThreadX] = gpu::ThreadIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::x);
+ ids[gpu::Processor::ThreadY] = gpu::ThreadIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::y);
+ ids[gpu::Processor::ThreadZ] = gpu::ThreadIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::z);
+ dims[gpu::Processor::BlockX] = gpu::GridDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::x);
+ dims[gpu::Processor::BlockY] = gpu::GridDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::y);
+ dims[gpu::Processor::BlockZ] = gpu::GridDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::z);
+ dims[gpu::Processor::ThreadX] = gpu::BlockDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::x);
+ dims[gpu::Processor::ThreadY] = gpu::BlockDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::y);
+ dims[gpu::Processor::ThreadZ] = gpu::BlockDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::z);
+ }
+
+ /// Return the compute-region block argument for \p outside, adding an `ins`
+ /// operand when needed.
+ BlockArgument getOrAppendInsBlockArg(Value outside) {
+ if (std::optional<BlockArgument> blockArg =
+ computeRegion.getBlockArg(outside)) {
+ return *blockArg;
+ }
+ return computeRegion.appendInputArg(outside);
+ }
+
+ /// Wire dynamic privatization extents into the compute region as `ins` args.
+ void preparePrivatizeExtentInsOperands() {
+ computeRegion.walk([&](acc::PrivateLocalOp privateLocal) {
+ acc::PrivatizeOp privatizeOp =
+ getPrivatizeOp(privateLocal, computeRegion);
+ if (privatizeOp->getParentOfType<acc::ComputeRegionOp>() ==
+ computeRegion) {
+ return;
+ }
+ for (Value extent : privatizeOp.getDynamicSizes()) {
+ getOrAppendInsBlockArg(extent);
+ }
+ });
+ }
+
+ /// Resolve dynamic size operands for a privatized array.
+ SmallVector<Value>
+ resolvePrivateLocalDynamicExtents(acc::PrivateLocalOp privateLocal) {
+ acc::PrivatizeOp privatizeOp = getPrivatizeOp(privateLocal, computeRegion);
+ SmallVector<Value> extents;
+ for (Value extent : privatizeOp.getDynamicSizes()) {
+ if (std::optional<BlockArgument> blockArg =
+ computeRegion.getBlockArg(extent)) {
+ extents.push_back(mapping.lookupOrDefault(*blockArg));
+ continue;
+ }
+ extents.push_back(mapping.lookupOrDefault(extent));
+ }
+ return extents;
+ }
+
+ RewriterBase &rewriter;
+ acc::ComputeRegionOp computeRegion;
+
+ acc::OpenACCSupport &accSupport;
+ const ACCCGToGPUOptions &options;
+ gpu::LaunchOp launch;
+ IRMapping mapping;
+ llvm::SmallVector<scf::ParallelOp> loopReductions;
+ llvm::DenseMap<gpu::Processor, Value> threadIdMap;
+ llvm::DenseMap<gpu::Processor, Value> dimensionMap;
+ // True if ThreadY reduction exists, which triggers subgroup alignment
+ bool hasThreadYReduction = false;
+ // True if any ThreadX routine call exists in the kernel
+ bool hasThreadLevelRoutineCall = false;
+ // True when a per-row ThreadY barrier is emitted
+ bool hasThreadYBarrier = false;
+
+ // Reusable privatize broadcast slots per type; disabled for kernels.
+ llvm::DenseMap<Type, Value> privatizeBroadcastCache;
+
+ int64_t staticBlockDimX = 1024;
+ acc::DefaultACCToGPUMappingPolicy defaultPolicy;
+ SharedMemoryBudget sharedMemBudget;
+ SmallVector<std::string> sharedMemPrivateVarNames;
+ llvm::SmallVector<Operation *, 4> deferredBarrierSeqLoops;
+
+ Value getThreadId(Location loc, gpu::Dimension dim) {
+ return gpu::ThreadIdOp::create(rewriter, loc, rewriter.getIndexType(), dim);
+ }
+
+ Value getBlockDim(Location loc, gpu::Dimension dim) {
+ return gpu::BlockDimOp::create(rewriter, loc, rewriter.getIndexType(), dim);
+ }
+
+ /// Thread id for \p proc, from the launch op or the routine context map.
+ Value getGPUThreadIdFor(gpu::Processor proc) {
+ return getGPUThreadId(proc, getLaunch(), threadIdMap);
+ }
+
+ /// Grid/block dimension for \p proc, from the launch op or routine map.
+ Value getGPUSizeFor(gpu::Processor proc) {
+ return getGPUSize(proc, getLaunch(), dimensionMap);
+ }
+};
+
+int64_t ACCCGToGPULowering::getElementSizeInBytes(Location loc,
+ Type elementType) const {
+ ModuleOp module = computeRegion->getParentOfType<ModuleOp>();
+ if (std::optional<acc::TypeSizeAndAlignment> sizeAndAlignment =
+ accSupport.getTypeSizeAndAlignment(elementType, module)) {
+ return sizeAndAlignment->first.getFixedValue();
+ }
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ os << "element size computation for unsupported type: " << elementType;
+ (void)accSupport.emitNYI(loc, os.str());
+ return 0;
+}
+
+bool ACCCGToGPULowering::canUseStackAlloca(
+ MemRefType baseTy, Location loc, int64_t maxThreadPrivateStack) const {
+ for (int64_t dim : baseTy.getShape()) {
+ if (dim == ShapedType::kDynamic) {
+ return false;
+ }
+ }
+ int64_t elementSize = getElementSizeInBytes(loc, baseTy.getElementType());
+ int64_t numElements = 1;
+ for (int64_t dim : baseTy.getShape()) {
+ if (numElements > maxThreadPrivateStack / std::max<int64_t>(dim, 1)) {
+ return false;
+ }
+ numElements *= dim;
+ }
+ return elementSize * numElements < maxThreadPrivateStack;
+}
+
+/// True if the accumulate spans a block dim or is nested in a block-mapped
+/// loop, i.e. each block owns the elements it reduces across threads. A
+/// thread-only accumulate with no block context grid-strides its element loop
+/// onto blocks, so per-thread partials would be dropped; such reductions must
+/// stay shared.
+static bool reductionHasBlockContext(acc::ReductionAccumulateArrayOp accArr) {
+ auto hasBlock = [](mlir::acc::GPUParallelDimsAttr parDims) {
+ return parDims && llvm::any_of(parDims.getArray(),
+ [](auto pd) { return pd.isAnyBlock(); });
+ };
+ if (hasBlock(accArr.getParDimsAttr())) {
+ return true;
+ }
+ for (scf::ParallelOp loop = accArr->getParentOfType<scf::ParallelOp>(); loop;
+ loop = loop->getParentOfType<scf::ParallelOp>()) {
+ if (hasBlock(mlir::acc::getParDimsAttr(loop))) {
+ return true;
+ }
+ }
+ return false;
+}
+
+/// Returns the array reduction accumulate (through cast/view ops) that \p v
+/// feeds if it needs per-thread storage: its par_dims include a thread dim
+/// and it has block context so the cross-thread all_reduce is well defined.
+static acc::ReductionAccumulateArrayOp perThreadArrayReductionAccum(Value v) {
+ SmallVector<Value> worklist{v};
+ DenseSet<Value> seen;
+ while (!worklist.empty()) {
+ Value cur = worklist.pop_back_val();
+ if (!seen.insert(cur).second) {
+ continue;
+ }
+ for (Operation *user : cur.getUsers()) {
+ if (acc::ReductionAccumulateArrayOp accArr =
+ dyn_cast<acc::ReductionAccumulateArrayOp>(user)) {
+ bool hasThread = false;
+ for (auto pd : accArr.getParDims().getArray()) {
+ hasThread |= pd.isAnyThread();
+ }
+ if (hasThread && reductionHasBlockContext(accArr)) {
+ return accArr;
+ }
+ continue;
+ }
+ SmallVector<Value> through;
+ if (getPassThroughResults(user, cur, through)) {
+ worklist.append(through.begin(), through.end());
+ } else if (isa<ViewLikeOpInterface>(user)) {
+ worklist.append(user->result_begin(), user->result_end());
+ }
+ }
+ }
+ return nullptr;
+}
+
+/// Store the reduction identity to every element of a freshly allocated
+/// per-thread array accumulator so all lanes start from identity (the original
+/// init loop may only run on one lane).
+static void initPerThreadArrayAccum(OpBuilder &b, Location loc, Value alloca,
+ MemRefType baseTy,
+ arith::AtomicRMWKind kind) {
+ assert(baseTy.getRank() == 1 && baseTy.hasStaticShape() &&
+ "per-thread array reduction accumulator must be static rank-1");
+ Value ident = createIdentityValue(b, loc, baseTy.getElementType(), kind,
+ /*useOnlyFiniteValue=*/true);
+ Value lb = arith::ConstantIndexOp::create(b, loc, 0);
+ Value ub = arith::ConstantIndexOp::create(b, loc, baseTy.getShape()[0]);
+ Value step = arith::ConstantIndexOp::create(b, loc, 1);
+ auto forOp = scf::ForOp::create(b, loc, lb, ub, step);
+ OpBuilder::InsertionGuard g(b);
+ b.setInsertionPoint(forOp.getBody()->getTerminator());
+ memref::StoreOp::create(b, loc, ident, alloca, forOp.getInductionVar());
+}
+
+std::optional<int64_t>
+ACCCGToGPULowering::isEligibleForSharedMemory(acc::PrivateLocalOp privateLocal,
+ MemRefType baseTy) {
+ // Cross-thread array reduction accumulators must stay per-thread.
+ if (perThreadArrayReductionAccum(privateLocal.getResult())) {
+ return std::nullopt;
+ }
+ ModuleOp module = computeRegion->getParentOfType<ModuleOp>();
+ FailureOr<bool> isCandidate = isPrivateLocalSharedMemoryCandidate(
+ privateLocal, computeRegion, module, defaultPolicy, &accSupport);
+ if (failed(isCandidate)) {
+ hasFailed = true;
+ return std::nullopt;
+ }
+ if (!isCandidate.value()) {
+ return std::nullopt;
+ }
+ std::optional<int64_t> upperBound =
+ getPrivateLocalSharedMemoryUpperBoundBytes(privateLocal, computeRegion,
+ module, defaultPolicy);
+ assert(upperBound && "candidate private_local must have an upper bound");
+ int64_t elementSize =
+ getElementSizeInBytes(privateLocal.getLoc(), baseTy.getElementType());
+ int64_t numElements = 1;
+ for (int64_t dim : baseTy.getShape()) {
+ numElements *= dim;
+ }
+ return *upperBound / (elementSize * numElements);
+}
+
+bool ACCCGToGPULowering::tryAllocateSharedMemory(int64_t bytes) {
+ return sharedMemBudget.tryAllocate(bytes);
+}
+
+FailureOr<arith::AtomicRMWKind>
+ACCCGToGPULowering::getReductionKind(acc::ReductionOperator redOp, Type type,
+ Location loc) {
+ if (std::optional<arith::AtomicRMWKind> kind =
+ translateACCReductionOperator(redOp, type))
+ return *kind;
+
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ os << "reduction operator (" << redOp << ") for type " << type;
+ (void)accSupport.emitNYI(loc, os.str());
+ return failure();
+}
+
+LogicalResult ACCCGToGPULowering::rewrite() {
+
+ // Pre-compute if thread-level reductions exist. ThreadY reduction generates
+ // shuffles which require subgroup alignment (blockDim.x = subgroupSize),
+ // meaning ThreadX lanes exist even without explicit ThreadX parallelism.
+ computeRegion->walk([&](acc::ReductionAccumulateOp op) -> WalkResult {
+ for (auto parDim : op.getParDimsAttr().getArray()) {
+ if (parDim.isThreadY()) {
+ hasThreadYReduction = true;
+ return WalkResult::interrupt();
+ }
+ }
+ return WalkResult::advance();
+ });
+
+ // Pre-compute if any thread-level (vector or worker) routine call exists.
+ // Such routines partition work across ThreadX/ThreadY and emit workgroup-wide
+ // barriers internally (e.g. for shared memory alloca synchronization), so all
+ // workgroup threads must reach the call site for those barriers to converge.
+ computeRegion->walk([&](CallOpInterface callOp) -> WalkResult {
+ if (mlir::acc::GPUParallelDimAttr parDim =
+ getAccRoutineCallParDim(callOp, defaultPolicy)) {
+ if (parDim.isThreadX() || parDim.isThreadY()) {
+ hasThreadLevelRoutineCall = true;
+ return WalkResult::interrupt();
+ }
+ }
+ return WalkResult::advance();
+ });
+
+ Location loc = computeRegion->getLoc();
+ Value constantOne = arith::ConstantIndexOp::create(rewriter, loc, 1);
+
+ auto launchArgument = [&](gpu::Processor processor) -> Value {
+ mlir::acc::GPUParallelDimAttr parDim = mlir::acc::GPUParallelDimAttr::get(
+ computeRegion->getContext(), processor);
+ std::optional<Value> maybeLaunchArg =
+ computeRegion.getKnownLaunchArg(parDim);
+ LLVM_DEBUG(llvm::dbgs() << "ACCCGToGPU: launch-arg: "
+ << " parDim: " << parDim << " gpu: " << processor
+ << " widthValue: "
+ << maybeLaunchArg.value_or(constantOne) << "\n");
+
+ return getValueOrCreateCastToIndexLike(
+ rewriter, loc, rewriter.getIndexType(),
+ maybeLaunchArg.value_or(constantOne));
+ };
+ LLVM_DEBUG(llvm::dbgs() << "ACCCGToGPU: creating gpu launch op: \n");
+
+ // acc.compute_region keeps launch argument as block argument, for rewriting
+ // we now replace these with gpu.launch dimensions.
+ auto mapLaunchArguments = [&](gpu::Processor processor, Value launchArg) {
+ mlir::acc::GPUParallelDimAttr parDim = mlir::acc::GPUParallelDimAttr::get(
+ computeRegion->getContext(), processor);
+ std::optional<Value> kernelArg = computeRegion.getLaunchArg(parDim);
+ if (kernelArg) {
+ mapping.map(computeRegion.gpuParWidth(processor), launchArg);
+ }
+ };
+
+ llvm::StringRef blockDimXName = "blockDim.x";
+ llvm::StringRef blockDimYName = "blockDim.y";
+ std::string deviceLabel = getDeviceRemarkQualifier(options.deviceType);
+
+ if (!computeRegion->getParentOfType<gpu::GPUFuncOp>()) {
+ Value blockDimX = launchArgument(gpu::Processor::ThreadX);
+ APInt bdxVal;
+ if (matchPattern(blockDimX, m_ConstantInt(&bdxVal))) {
+ staticBlockDimX = bdxVal.getSExtValue();
+ }
+ Value blockDimY = launchArgument(gpu::Processor::ThreadY);
+ Value blockDimZ = launchArgument(gpu::Processor::ThreadZ);
+ Value gridDimX = launchArgument(gpu::Processor::BlockX);
+ Value gridDimY = launchArgument(gpu::Processor::BlockY);
+ Value gridDimZ = launchArgument(gpu::Processor::BlockZ);
+
+ // The format of the message is:
+ // Generating [serial] {deviceLabel} code with gridDim=32x1x1
+ // blockDim=256x1x1
+ accSupport.emitRemark(computeRegion, [&]() {
+ auto getName = [&](Value val) -> std::string {
+ std::string name = accSupport.getVariableName(val);
+ return name.empty() ? "(*)" : name;
+ };
+ bool isEffectivelySerial =
+ sameEffectiveValue(blockDimX, 1) &&
+ sameEffectiveValue(blockDimY, 1) &&
+ sameEffectiveValue(blockDimZ, 1) && sameEffectiveValue(gridDimX, 1) &&
+ sameEffectiveValue(gridDimY, 1) && sameEffectiveValue(gridDimZ, 1);
+ return (llvm::Twine("Generating ") +
+ llvm::Twine(isEffectivelySerial ? "serial " : "") + deviceLabel +
+ " code with gridDim=" + getName(gridDimX) + "x" +
+ getName(gridDimY) + "x" + getName(gridDimZ) +
+ " blockDim=" + getName(blockDimX) + "x" + getName(blockDimY) +
+ "x" + getName(blockDimZ))
+ .str();
+ });
+
+ // Check if kernel has a stream operand for async execution
+ if (mlir::Value streamValue = computeRegion.getStream()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "\nDEBUG: Creating async gpu.launch with stream: "
+ << streamValue << "\n");
+ launch = gpu::LaunchOp::create(
+ rewriter, loc, gridDimX, gridDimY, gridDimZ, blockDimX, blockDimY,
+ blockDimZ,
+ /*dynamicSharedMemorySize=*/mlir::Value{},
+ /*asyncTokenType=*/
+ mlir::gpu::AsyncTokenType::get(rewriter.getContext()));
+ // Add the stream as an async dependency
+ launch.getAsyncDependenciesMutable().append(streamValue);
+ } else {
+ LLVM_DEBUG(llvm::dbgs()
+ << "\nDEBUG: No stream, creating sync gpu.launch\n");
+ launch = gpu::LaunchOp::create(rewriter, loc, gridDimX, gridDimY,
+ gridDimZ, blockDimX, blockDimY, blockDimZ);
+ }
+
+ // Transfer kernel function name and module name from acc.compute_region to
+ // gpu.launch if present
+ if (auto kernelFuncName = computeRegion.getKernelFuncNameAttr()) {
+ launch.setFunctionAttr(kernelFuncName);
+ }
+ if (auto kernelModuleName = computeRegion.getKernelModuleNameAttr()) {
+ launch.setModuleAttr(kernelModuleName);
+ }
+
+ rewriter.setInsertionPointToEnd(&launch.getBody().front());
+ gpu::TerminatorOp::create(rewriter, loc);
+ rewriter.setInsertionPointToStart(&launch.getBody().front());
+ mapLaunchArguments(gpu::Processor::BlockX,
+ gpu::GridDimOp::create(rewriter, loc,
+ rewriter.getIndexType(),
+ gpu::Dimension::x));
+ mapLaunchArguments(gpu::Processor::BlockY,
+ gpu::GridDimOp::create(rewriter, loc,
+ rewriter.getIndexType(),
+ gpu::Dimension::y));
+ mapLaunchArguments(gpu::Processor::BlockZ,
+ gpu::GridDimOp::create(rewriter, loc,
+ rewriter.getIndexType(),
+ gpu::Dimension::z));
+ mapLaunchArguments(gpu::Processor::ThreadX,
+ gpu::BlockDimOp::create(rewriter, loc,
+ rewriter.getIndexType(),
+ gpu::Dimension::x));
+ mapLaunchArguments(gpu::Processor::ThreadY,
+ gpu::BlockDimOp::create(rewriter, loc,
+ rewriter.getIndexType(),
+ gpu::Dimension::y));
+ mapLaunchArguments(gpu::Processor::ThreadZ,
+ gpu::BlockDimOp::create(rewriter, loc,
+ rewriter.getIndexType(),
+ gpu::Dimension::z));
+ } else {
+ // Do not create gpu.launch for acc routine and map
+ // to block/thread index and block/grid size instead
+ // of launch arguments, using created maps.
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointToStart(computeRegion->getBlock());
+ createForAllDimensions(rewriter, loc, threadIdMap, dimensionMap);
+ mapLaunchArguments(gpu::Processor::BlockX,
+ dimensionMap[gpu::Processor::BlockX]);
+ mapLaunchArguments(gpu::Processor::BlockY,
+ dimensionMap[gpu::Processor::BlockY]);
+ mapLaunchArguments(gpu::Processor::BlockZ,
+ dimensionMap[gpu::Processor::BlockZ]);
+ mapLaunchArguments(gpu::Processor::ThreadX,
+ dimensionMap[gpu::Processor::ThreadX]);
+ mapLaunchArguments(gpu::Processor::ThreadY,
+ dimensionMap[gpu::Processor::ThreadY]);
+ mapLaunchArguments(gpu::Processor::ThreadZ,
+ dimensionMap[gpu::Processor::ThreadZ]);
+ }
+
+ // Map input arguments for compute region; we go from an IsolatedFromAbove
+ // operation to gpu.launch which is not IsolatedFromAbove.
+ preparePrivatizeExtentInsOperands();
+ Block *body = computeRegion.getBody();
+ unsigned numLaunchArgs = computeRegion.getLaunchArgs().size();
+ ValueRange inputArgs = computeRegion.getInputArgs();
+ for (unsigned i = numLaunchArgs; i < body->getNumArguments(); ++i) {
+ mapping.map(body->getArgument(i), inputArgs[i - numLaunchArgs]);
+ }
+
+ assert(computeRegion.getRegion().hasOneBlock() &&
+ "compute region only supports one block region for now");
+ // process all operations inside kernel region
+ for (auto &op :
+ computeRegion.getRegion().getBlocks().front().getOperations()) {
+ processOp(&op);
+ }
+
+ for (auto &parLoop : loopReductions) {
+ postprocessLoopReduction(parLoop);
+ }
+
+ // Replace combine reloads of a reduction slot with the block-reduced value.
+ // Only when it dominates the reload; otherwise keep the reload.
+ if (!pendingCombineReloads.empty() && launch) {
+ DominanceInfo domInfo(launch);
+ for (auto &[slot, loadOp] : pendingCombineReloads) {
+ llvm::DenseMap<Value, Value>::iterator it =
+ reductionAccumValue.find(slot);
+ if (it == reductionAccumValue.end()) {
+ continue;
+ }
+ if (!domInfo.dominates(it->second, loadOp.getOperation())) {
+ continue;
+ }
+ rewriter.replaceOp(loadOp, ValueRange{it->second});
+ }
+ }
+
+ if (launch) {
+ const int64_t subgroupSize = options.subgroupSize;
+ const int64_t subgroupAlignMask = subgroupSize - 1;
+
+ // Adjust blockDim.x to be a multiple of subgroupSize. This is required
+ // because:
+ // - Subgroup reductions (gpu.all_reduce) require full subgroups
+ // - Per-row workgroup barriers require blockDim.x aligned to subgroupSize
+ bool isShuffleEnabled = false;
+
+ launch.walk([&](gpu::AllReduceOp allReduce) -> WalkResult {
+ ArrayRef<mlir::acc::GPUParallelDimAttr> parDims =
+ mlir::acc::getParDimsAttr(allReduce).getArray();
+ for (auto parDim : parDims) {
+ if (parDim.isThreadX() || parDim.isThreadY()) {
+ // Shuffle are enabled. Need to adjust the ThreadX length.
+ isShuffleEnabled = true;
+ return WalkResult::interrupt();
+ }
+ }
+ return WalkResult::advance();
+ });
+ // Also check if called routines have ThreadY reductions
+ if (!isShuffleEnabled) {
+ launch.walk([&](func::CallOp callOp) -> WalkResult {
+ if (gpu::GPUFuncOp callee =
+ callOp->getParentOfType<ModuleOp>()
+ .lookupSymbol<gpu::GPUFuncOp>(callOp.getCallee())) {
+ callee.walk([&](gpu::AllReduceOp allReduce) -> WalkResult {
+ ArrayRef<mlir::acc::GPUParallelDimAttr> parDims =
+ mlir::acc::getParDimsAttr(allReduce).getArray();
+ for (auto parDim : parDims) {
+ if (parDim.isThreadX() || parDim.isThreadY()) {
+ isShuffleEnabled = true;
+ return WalkResult::interrupt();
+ }
+ }
+ return WalkResult::advance();
+ });
+ }
+ return isShuffleEnabled ? WalkResult::interrupt()
+ : WalkResult::advance();
+ });
+ }
+
+ if (isShuffleEnabled || hasThreadYBarrier) {
+ rewriter.setInsertionPoint(launch);
+
+ Value curBlockDimX = launch.getBlockSizeX();
+ Value curBlockDimY = launch.getBlockSizeY();
+
+ // Emit a report on changing parallelism.
+ accSupport.emitRemark(computeRegion, [&]() {
+ auto getName = [&](Value val) -> std::string {
+ std::string name = accSupport.getVariableName(val);
+ return name.empty() ? "(*)" : name;
+ };
+ std::string blockDimXValStr = getName(curBlockDimX);
+ std::string blockDimYValStr = getName(curBlockDimY);
+ llvm::StringRef kind =
+ isShuffleEnabled ? "Shuffle reduction" : "ThreadY barrier";
+ return (llvm::Twine(kind) +
+ " is generated while adjusting the number of threads into "
+ "groups of " +
+ llvm::Twine(subgroupSize) + ".\n\t" + blockDimXName + ": `" +
+ blockDimXValStr + "` to `((" + blockDimXValStr + " + " +
+ llvm::Twine(subgroupAlignMask) + ") / " +
+ llvm::Twine(subgroupSize) + ") * " + llvm::Twine(subgroupSize) +
+ "`\n" + "\t" + blockDimYName + ": `" + blockDimYValStr +
+ "` to `max(1, (new-" + blockDimXName + " * " + blockDimYValStr +
+ ") / new-" + blockDimXName + ")`")
+ .str();
+ });
+
+ std::optional<int64_t> constBlockDimX = getConstantIntValue(curBlockDimX);
+ std::optional<int64_t> constBlockDimY = getConstantIntValue(curBlockDimY);
+
+ // Skip subgroup alignment only when the total thread count is already
+ // below a subgroup (constant blockDim.x in 2..subgroupSize-1 and
+ // constant blockDim.y == 1). If blockDim.y > 1 or is unknown, padding
+ // blockDim.x to a subgroup is still required so subgroups don't cross
+ // row boundaries for row-local shuffle/ThreadY-barrier reductions.
+ bool skipAlign = false;
+ if (constBlockDimX && constBlockDimY && *constBlockDimX > 1 &&
+ *constBlockDimX < subgroupSize && *constBlockDimY == 1) {
+ skipAlign = true;
+ }
+
+ // Update both the ThreadX length and the number of ThreadY.
+ // When the original blockDim.x and blockDim.y are compile-time
+ // constants, compute the adjusted dimensions as constants directly so
+ // that the GpuKernelOutliningPass can set `known_block_size` on the
+ // outlined gpu.func.
+ Value newBlockDimX, newBlockDimY;
+ if (constBlockDimX && constBlockDimY) {
+ int64_t bdx = *constBlockDimX;
+ int64_t bdy = *constBlockDimY;
+ int64_t alignedBdx =
+ ((bdx + subgroupAlignMask) / subgroupSize) * subgroupSize;
+ int64_t numThreads = bdx * bdy;
+ int64_t newBdy = std::max<int64_t>(1, numThreads / alignedBdx);
+ newBlockDimX =
+ arith::ConstantIndexOp::create(rewriter, loc, alignedBdx);
+ newBlockDimY = arith::ConstantIndexOp::create(rewriter, loc, newBdy);
+ } else {
+ // numThreads = blockDim.x * blockDim.y
+ Value numThreads =
+ arith::MulIOp::create(rewriter, loc, curBlockDimX, curBlockDimY);
+ // blockDim.x = ((blockDim.x + mask) / subgroupSize) * subgroupSize
+ Value cstMask =
+ arith::ConstantIndexOp::create(rewriter, loc, subgroupAlignMask);
+ Value cstSubgroupSize =
+ arith::ConstantIndexOp::create(rewriter, loc, subgroupSize);
+ Value padded =
+ arith::AddIOp::create(rewriter, loc, curBlockDimX, cstMask);
+ Value subgroupsRequired =
+ arith::DivUIOp::create(rewriter, loc, padded, cstSubgroupSize);
+ newBlockDimX = arith::MulIOp::create(rewriter, loc, subgroupsRequired,
+ cstSubgroupSize);
+ // blockDim.y = max(1, numThreads / blockDim.x)
+ Value quotient =
+ arith::DivUIOp::create(rewriter, loc, numThreads, newBlockDimX);
+ Value cst1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ newBlockDimY = arith::MaxUIOp::create(rewriter, loc, cst1, quotient);
+ }
+
+ if (!skipAlign) {
+ launch.getBlockSizeXMutable().assign(newBlockDimX);
+ launch.getBlockSizeYMutable().assign(newBlockDimY);
+ }
+ }
+ }
+
+ if (hasFailed) {
+ return failure();
+ }
+
+ if (!sharedMemPrivateVarNames.empty()) {
+ accSupport.emitRemark(computeRegion, [&]() {
+ return (llvm::Twine("GPU shared memory used for ") +
+ llvm::join(sharedMemPrivateVarNames, ","))
+ .str();
+ });
+ }
+
+ rewriter.eraseOp(computeRegion);
+ return success();
+}
+
+/// True when this accumulate is redundant in a nested reduction chain: the
+/// value is a load of the destination memref and a sibling
+/// acc.reduction_combine with block par_dims has already reduced %M across
+/// threads in the block.
+///
+/// %v = memref.load %M[]
+/// acc.reduction_accumulate %v to %M ...
+/// acc.reduction_combine %M into %parent ... {block par_dims}
+///
+/// Lowering the accumulate again would double-count. Detection is structural;
+/// nested reductions into per-thread privates do not match because their
+/// combines are not block-scoped.
+static bool isRedundantChainAccumulate(acc::ReductionAccumulateOp op) {
+ Value memref = op.getMemref();
+ memref::LoadOp loadOp = op.getValue().getDefiningOp<memref::LoadOp>();
+ if (!loadOp || loadOp.getMemRef() != memref) {
+ return false;
+ }
+ for (Operation *user : memref.getUsers()) {
+ acc::ReductionCombineOp combineOp = dyn_cast<acc::ReductionCombineOp>(user);
+ if (!combineOp || combineOp.getDestMemref() != memref) {
+ continue;
+ }
+ SmallVector<mlir::acc::GPUParallelDimAttr> parDims =
+ getReductionCombineParDims(combineOp);
+ if (llvm::any_of(parDims, [](mlir::acc::GPUParallelDimAttr d) {
+ return d.isAnyBlock();
+ })) {
+ return true;
+ }
+ }
+ return false;
+}
+
+std::pair<SmallVector<mlir::acc::GPUParallelDimAttr>,
+ SmallVector<mlir::acc::GPUParallelDimAttr>>
+ACCCGToGPULowering::computeActiveAndInactiveParDims(Operation *op,
+ Block *block) {
+ MLIRContext *ctx = computeRegion->getContext();
+ SmallVector<mlir::acc::GPUParallelDimAttr> ancestorParDims =
+ getAncestorParDims(op);
+ // Preserve whether there were any structural ancestor par-dims before
+ // we start augmenting them based on inner uses (e.g. private_local).
+ // This is needed for gang redundancy check - stores to worker-indexed
+ // private_local should not disable redundant gang execution.
+ bool noStructuralAncestorParDims =
+ llvm::none_of(ancestorParDims, [](auto pd) { return !pd.isSeq(); });
+
+ mlir::acc::GPUParallelDimAttr routineParDim;
+ if (isInsideACCSpecializedRoutine(computeRegion)) {
+ FunctionOpInterface funcOp =
+ computeRegion->getParentOfType<FunctionOpInterface>();
+ routineParDim = getSpecializedRoutineDim(funcOp, defaultPolicy);
+ if (routineParDim.isThreadX()) {
+ mlir::acc::insertParDim(ancestorParDims,
+ mlir::acc::GPUParallelDimAttr::threadYDim(ctx));
+ }
+ mlir::acc::insertParDim(ancestorParDims,
+ mlir::acc::GPUParallelDimAttr::blockXDim(ctx));
+ }
+
+ // acc.private_local should use the same par_dims as acc.reduction_accumulate.
+ if (acc::PrivateLocalOp privateLocalOp = dyn_cast<acc::PrivateLocalOp>(op)) {
+ for (Operation *user : privateLocalOp.getResult().getUsers()) {
+ if (acc::ReductionAccumulateOp accumulateOp =
+ dyn_cast<acc::ReductionAccumulateOp>(user)) {
+ if (accumulateOp.getMemref() == privateLocalOp.getResult()) {
+ for (mlir::acc::GPUParallelDimAttr parDim :
+ accumulateOp.getParDims().getArray()) {
+ mlir::acc::insertParDim(ancestorParDims, parDim);
+ }
+ }
+ }
+ // For decomposed complex reductions, the private_local is consumed
+ // by an acc.reduction_combine{,_region} (no acc.reduction_accumulate
+ // user). Mirror the par_dims so this private_local is treated as the
+ // accumulator at the same parallelism level as a scalar reduction
+ // would be (per-thread, not block-shared).
+ if (acc::ReductionCombineOp combineOp =
+ dyn_cast<acc::ReductionCombineOp>(user)) {
+ if (combineOp.getSrcMemref() == privateLocalOp.getResult()) {
+ for (mlir::acc::GPUParallelDimAttr parDim :
+ getReductionCombineParDims(combineOp)) {
+ mlir::acc::insertParDim(ancestorParDims, parDim);
+ }
+ }
+ }
+ if (auto combineRegionOp =
+ dyn_cast<acc::ReductionCombineRegionOp>(user)) {
+ if (combineRegionOp.getSrcVar() == privateLocalOp.getResult()) {
+ for (mlir::acc::GPUParallelDimAttr parDim :
+ getReductionCombineParDims(combineRegionOp)) {
+ mlir::acc::insertParDim(ancestorParDims, parDim);
+ }
+ }
+ }
+ }
+ }
+
+ bool hasBlock = false;
+ for (mlir::acc::GPUParallelDimAttr parDim : ancestorParDims) {
+ if (parDim.isAnyBlock()) {
+ hasBlock = true;
+ }
+ }
+
+ mlir::acc::GPUParallelDimAttr lowestParDim =
+ mlir::acc::GPUParallelDimAttr::threadXDim(ctx);
+ if (block) {
+ block->walk([&](Operation *op) {
+ // Check stores to acc.private_local - add the privatize's par_dims
+ // as active dims so predication is correct for per-worker/gang memory.
+ if (memref::StoreOp storeOp = dyn_cast<memref::StoreOp>(op)) {
+ if (auto privateLocalOp =
+ storeOp.getMemref().getDefiningOp<acc::PrivateLocalOp>()) {
+ acc::PrivatizeOp privatizeOp =
+ getPrivatizeOp(privateLocalOp, computeRegion);
+ if (mlir::acc::GPUParallelDimsAttr parDimsAttr =
+ privatizeOp.getParDimsAttr()) {
+ for (auto parDim : parDimsAttr.getArray()) {
+ mlir::acc::insertParDim(ancestorParDims, parDim);
+ }
+ }
+ }
+ }
+ // Consider ACC routine calls; routine calls should be predicated up to
+ // one level above the parallel dimension of the callee.
+ if (CallOpInterface callOp = dyn_cast<CallOpInterface>(op)) {
+ if (mlir::acc::GPUParallelDimAttr parDim =
+ getAccRoutineCallParDim(callOp, defaultPolicy)) {
+ if (parDim.isBlockZ()) {
+ lowestParDim = parDim;
+ } else {
+ lowestParDim = parDim.getOneHigher();
+ }
+ }
+ }
+ // acc.reduction_combine_region should be predicated with the par_dims of
+ // acc.reduction_accumulate. This is required when using combine between
+ // kernel and loop in combined constructs.
+ if (acc::ReductionCombineOp reductionCombineOp =
+ dyn_cast<acc::ReductionCombineOp>(op)) {
+ for (mlir::acc::GPUParallelDimAttr parDim :
+ getReductionCombineParDims(reductionCombineOp)) {
+ mlir::acc::removeParDim(ancestorParDims, parDim);
+ }
+ }
+ if (acc::ReductionCombineRegionOp combineRegionOp =
+ dyn_cast<acc::ReductionCombineRegionOp>(op)) {
+ for (mlir::acc::GPUParallelDimAttr parDim :
+ getReductionCombineParDims(combineRegionOp)) {
+ mlir::acc::removeParDim(ancestorParDims, parDim);
+ }
+ }
+ // An array accumulate reduces across its par_dims via gpu.all_reduce, so
+ // all those threads must execute it - treat them as active (unlike the
+ // scalar accumulate, which is active through its enclosing scf.parallel).
+ if (acc::ReductionAccumulateArrayOp accArrayOp =
+ dyn_cast<acc::ReductionAccumulateArrayOp>(op)) {
+ for (mlir::acc::GPUParallelDimAttr parDim :
+ accArrayOp.getParDims().getArray()) {
+ mlir::acc::insertParDim(ancestorParDims, parDim);
+ }
+ }
+ return WalkResult::advance();
+ });
+ }
+
+ // Obtain launch dimensions
+ SmallVector<mlir::acc::GPUParallelDimAttr> launchParDims;
+ if (routineParDim) {
+ for (mlir::acc::GPUParallelDimAttr parDim = routineParDim;
+ parDim.getOrder() >= lowestParDim.getOrder();
+ parDim = parDim.getOneLower()) {
+ mlir::acc::insertParDim(launchParDims, parDim);
+ }
+ } else {
+ launchParDims = computeRegion.getLaunchParDims();
+ }
+
+ // Compute dimensions that execute op
+ SmallVector<mlir::acc::GPUParallelDimAttr> activeParDims, inactiveParDims;
+ for (mlir::acc::GPUParallelDimAttr launchParDim : launchParDims) {
+ if (launchParDim.getOrder() < lowestParDim.getOrder()) {
+ break;
+ }
+ if (llvm::find(ancestorParDims, launchParDim) != ancestorParDims.end() ||
+ (launchParDim.isAnyBlock() &&
+ (noStructuralAncestorParDims || hasBlock))) {
+ activeParDims.push_back(launchParDim);
+ } else {
+ inactiveParDims.push_back(launchParDim);
+ }
+ }
+
+ return std::pair{activeParDims, inactiveParDims};
+}
+
+Value ACCCGToGPULowering::emitPredicate(
+ Location loc, SmallVector<mlir::acc::GPUParallelDimAttr> &inactiveParDims) {
+ Value predicate;
+ for (mlir::acc::GPUParallelDimAttr inactiveParDim : inactiveParDims) {
+ Value threadId = getGPUThreadIdFor(inactiveParDim.getProcessor());
+ TypedAttr zeroAttr = rewriter.getZeroAttr(threadId.getType());
+ Value zero = arith::ConstantOp::create(rewriter, loc, zeroAttr);
+ Value cmp = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq,
+ threadId, zero);
+ if (predicate) {
+ predicate = arith::AndIOp::create(rewriter, loc, cmp, predicate);
+ } else {
+ predicate = cmp;
+ }
+ }
+ return predicate;
+}
+
+void ACCCGToGPULowering::createBarrier(
+ Location loc, mlir::acc::GPUParallelDimsAttr parDimsAttr) {
+ bool hasAnyBlock = false, hasThreadY = false, hasThreadX = false;
+ for (auto parDim : parDimsAttr.getArray()) {
+ if (parDim.isAnyBlock()) {
+ hasAnyBlock = true;
+ }
+ if (parDim.isThreadY()) {
+ hasThreadY = true;
+ }
+ if (parDim.isThreadX()) {
+ hasThreadX = true;
+ }
+ }
+
+ if (hasAnyBlock || hasThreadY) {
+ emitGPUBarrierWorkgroup(rewriter, loc);
+ } else if (hasThreadX) {
+ createPerRowBarrier(loc);
+ }
+}
+
+void ACCCGToGPULowering::createPerRowBarrier(Location loc) {
+ hasThreadYBarrier = true;
+
+ if (staticBlockDimX <= options.subgroupSize) {
+ emitGPUBarrierSubgroup(rewriter, loc);
+ return;
+ }
+
+ if (options.deviceType != mlir::acc::DeviceType::Nvidia) {
+ (void)accSupport.emitNYI(
+ loc,
+ "per-row barrier to support worker parallelism on non-NVIDIA device");
+ }
+
+ // Per-row barrier with fully runtime branching.
+ // Three mutually exclusive paths:
+ // blockDim.y == 1 -> gpu.barrier (workgroup-wide, only one worker)
+ // blockDim.x <= subgroupSize -> gpu.barrier scope<subgroup>
+ // blockDim.x > subgroupSize -> nvvm.barrier (tid.y + 1), blockDim.x
+ // (named)
+ //
+ // Per-row barriers use tid.y+1 so IDs start at 1, avoiding clash with
+ // barrier 0. When blockDim.y >= 16, worker 15's ID (16) wraps to
+ // physical barrier 0; this is safe because named barriers are reusable
+ // resources - workgroup-wide and per-row barriers on the same physical
+ // barrier execute at different program points and never overlap.
+ Value blockDimX = gpu::BlockDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::x);
+ Value blockDimY = gpu::BlockDimOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::y);
+ Value cst1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ Value isSingleWorker = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, blockDimY, cst1);
+
+ auto outerIf = scf::IfOp::create(rewriter, loc, isSingleWorker,
+ /*withElseRegion=*/true);
+
+ // Then: blockDim.y == 1 -> workgroup-wide barrier (safe, only one worker)
+ rewriter.setInsertionPointToStart(&outerIf.getThenRegion().front());
+ emitGPUBarrierWorkgroup(rewriter, loc);
+
+ // Else: blockDim.y > 1 - choose between subgroup sync and named barrier
+ rewriter.setInsertionPointToStart(&outerIf.getElseRegion().front());
+ Value cstSubgroupSize =
+ arith::ConstantIndexOp::create(rewriter, loc, options.subgroupSize);
+ Value isSubgroupSized = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ule, blockDimX, cstSubgroupSize);
+
+ auto innerIf = scf::IfOp::create(rewriter, loc, isSubgroupSized,
+ /*withElseRegion=*/true);
+
+ // Then: blockDim.x <= subgroupSize -> subgroup barrier (one worker per
+ // subgroup)
+ rewriter.setInsertionPointToStart(&innerIf.getThenRegion().front());
+ emitGPUBarrierSubgroup(rewriter, loc);
+
+ // Else: blockDim.x > subgroupSize -> per-row named barrier with tid.y + 1.
+ // The 1024-thread-per-block hardware limit with subgroup-aligned blockDim.x
+ // (>= 64 here) guarantees blockDim.y <= 16, so IDs span at most 16
+ // physical barriers (0-15) with no aliasing across workers.
+ rewriter.setInsertionPointToStart(&innerIf.getElseRegion().front());
+ Value threadYId = gpu::ThreadIdOp::create(
+ rewriter, loc, rewriter.getIndexType(), gpu::Dimension::y);
+ Value barrierId = arith::AddIOp::create(rewriter, loc, threadYId, cst1);
+ Type i32Ty = rewriter.getI32Type();
+ Value barrierId32 =
+ arith::IndexCastOp::create(rewriter, loc, i32Ty, barrierId);
+ Value numberOfThreads32 =
+ arith::IndexCastOp::create(rewriter, loc, i32Ty, blockDimX);
+
+ // GPU dialect named barriers do not have a means to create a custom barrier
+ // id. Thus use nvvm directly.
+ assert(options.deviceType == mlir::acc::DeviceType::Nvidia);
+ NVVM::BarrierOp::create(rewriter, loc, barrierId32, numberOfThreads32);
+
+ rewriter.setInsertionPointAfter(outerIf);
+}
+
+/// Whether any later sibling of \p loopOp (or a loop nested inside one) is a
+/// loop, i.e. whether some subsequent loop in the same region may read what
+/// \p loopOp wrote. Used to skip a barrier after the last loop, where nothing
+/// reads the data afterward.
+static bool hasSubsequentLoopSibling(Operation *loopOp) {
+ for (Operation *next = loopOp->getNextNode(); next;
+ next = next->getNextNode()) {
+ if (isa<scf::ParallelOp, scf::ForOp>(next)) {
+ return true;
+ }
+ bool nested = false;
+ next->walk([&](Operation *op) {
+ if (isa<scf::ParallelOp, scf::ForOp>(op)) {
+ nested = true;
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ });
+ if (nested) {
+ return true;
+ }
+ }
+ return false;
+}
+
+/// Nearest enclosing sequential loop ancestor of \p op.
+static LoopLikeOpInterface findFirstSequentialLoop(Operation *op) {
+ auto isAllSequentialParDims = [](scf::ParallelOp par) -> bool {
+ mlir::acc::GPUParallelDimsAttr pd = mlir::acc::getParDimsAttr(par);
+ if (!pd || pd.getArray().empty()) {
+ return false;
+ }
+ return llvm::all_of(pd.getArray(), [](mlir::acc::GPUParallelDimAttr d) {
+ return d.isSeq();
+ });
+ };
+
+ for (Operation *p = op->getParentOp(); p; p = p->getParentOp()) {
+ // Do not need to check scf.for op's parents
+ if (isa<scf::ForOp>(p)) {
+ return cast<LoopLikeOpInterface>(p);
+ }
+ if (scf::ParallelOp parOp = dyn_cast<scf::ParallelOp>(p)) {
+ if (isAllSequentialParDims(parOp)) {
+ return cast<LoopLikeOpInterface>(p);
+ }
+ }
+ }
+ return nullptr;
+}
+
+// A sequential loop that uses gang-private shared memory needs a
+// workgroup-wide barrier afterward so every thread in the block observes the
+// same state. When more work still lies between the loop and the next
+// thread-reconvergence point in the same block, that barrier must follow that
+// work; not sit immediately after the loop.
+//
+// The helpers below mark loop-body closure and other reconvergence points
+// where any postponed barrier must be inserted.
+
+/// Marks the end of a loop body's iteration in the current block.
+static bool isLoopBodyClosureOp(Operation *op) {
+ return isa<scf::ReduceOp, scf::YieldOp, acc::YieldOp>(op);
+}
+
+/// Thread-reconvergence point where any postponed post-loop barrier for earlier
+/// loops in this block must be inserted before proceeding.
+static bool isDeferredBarrierFlushPoint(Operation *op) {
+ if (isLoopBodyClosureOp(op)) {
+ return true;
+ }
+ // The next sequential loop may consume shared state produced by the prior
+ // one.
+ if (isa<scf::ForOp>(op)) {
+ return true;
+ }
+ if (scf::ParallelOp parallelOp = dyn_cast<scf::ParallelOp>(op)) {
+ if (mlir::acc::hasParDimsAttr(parallelOp)) {
+ if (mlir::acc::GPUParallelDimsAttr parDims =
+ mlir::acc::getParDimsAttr(parallelOp)) {
+ if (parDims.getArray().size() == 1 &&
+ parDims.getArray().front().isSeq()) {
+ return true;
+ }
+ }
+ }
+ }
+ return false;
+}
+
+/// True when a loop is followed by other work in the same block before the
+/// loop body closes; the post-loop barrier must wait for that reconvergence
+/// point instead of being placed right after the loop.
+static bool hasTrailingSideEffectSiblings(Operation *loopOp) {
+ for (Operation *next = loopOp->getNextNode(); next;
+ next = next->getNextNode()) {
+ return !isLoopBodyClosureOp(next);
+ }
----------------
clementval wrote:
no braces
https://github.com/llvm/llvm-project/pull/209606
More information about the Mlir-commits
mailing list