================ @@ -0,0 +1,162 @@ +//===- StackToShared.cpp -------------------------------------------===// +// +// 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 file implements transforms to swap stack allocations on the target +// device with device shared memory where applicable. +// +//===----------------------------------------------------------------------===// + +#include "flang/Optimizer/Dialect/FIROps.h" +#include "flang/Optimizer/HLFIR/HLFIROps.h" +#include "flang/Optimizer/OpenMP/Passes.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/OpenMP/OpenMPDialect.h" +#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h" + +namespace flangomp { +#define GEN_PASS_DEF_STACKTOSHAREDPASS +#include "flang/Optimizer/OpenMP/Passes.h.inc" +} // namespace flangomp + +using namespace mlir; + +namespace { +class StackToSharedPass + : public flangomp::impl::StackToSharedPassBase<StackToSharedPass> { +public: + StackToSharedPass() = default; + + void runOnOperation() override { + MLIRContext *context = &getContext(); + OpBuilder builder(context); + + func::FuncOp funcOp = getOperation(); + auto offloadIface = funcOp->getParentOfType<omp::OffloadModuleInterface>(); + if (!offloadIface || !offloadIface.getIsTargetDevice()) + return; + + funcOp->walk([&](fir::AllocaOp allocaOp) { + if (!shouldReplaceAlloca(*allocaOp)) + return; + + // Replace fir.alloca with omp.alloc_shared_mem. + builder.setInsertionPoint(allocaOp); + auto sharedAllocOp = omp::AllocSharedMemOp::create( + builder, allocaOp->getLoc(), allocaOp.getResult().getType(), + allocaOp.getInType(), allocaOp.getUniqNameAttr(), + allocaOp.getBindcNameAttr(), allocaOp.getTypeparams(), + allocaOp.getShape()); + allocaOp.replaceAllUsesWith(sharedAllocOp.getOperation()); + allocaOp.erase(); + + // Create a new omp.free_shared_mem for the allocated buffer prior to + // exiting the region. + Block *allocaBlock = sharedAllocOp->getBlock(); + DominanceInfo domInfo; + for (Block &block : sharedAllocOp->getParentRegion()->getBlocks()) { + Operation *terminator = block.getTerminator(); + if (!terminator->hasSuccessors() && + domInfo.dominates(allocaBlock, &block)) { + builder.setInsertionPoint(terminator); + omp::FreeSharedMemOp::create(builder, sharedAllocOp.getLoc(), + sharedAllocOp); + } + } + }); + } + +private: + // TODO: Refactor the logic in `shouldReplaceAlloca` and `checkAllocaUses` to + // be reusable by the MLIR to LLVM IR translation stage, as something very + // similar is also implemented there to choose between allocas and device + // shared memory allocations when processing OpenMP reductions, mapping and + // privatization. + + // Decide whether to replace a fir.alloca with a pair of device shared memory ---------------- Meinersbur wrote:
```suggestion /// Decide whether to replace a fir.alloca with a pair of device shared memory ``` Use doxygen for describing functions https://github.com/llvm/llvm-project/pull/161863 _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
