167 lines · cpp
1//===-- CUFComputeSharedMemoryOffsetsAndSize.cpp --------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "flang/Optimizer/Builder/BoxValue.h"10#include "flang/Optimizer/Builder/CUFCommon.h"11#include "flang/Optimizer/Builder/FIRBuilder.h"12#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"13#include "flang/Optimizer/Builder/Todo.h"14#include "flang/Optimizer/CodeGen/Target.h"15#include "flang/Optimizer/CodeGen/TypeConverter.h"16#include "flang/Optimizer/Dialect/CUF/CUFOps.h"17#include "flang/Optimizer/Dialect/FIRAttr.h"18#include "flang/Optimizer/Dialect/FIRDialect.h"19#include "flang/Optimizer/Dialect/FIROps.h"20#include "flang/Optimizer/Dialect/FIROpsSupport.h"21#include "flang/Optimizer/Dialect/FIRType.h"22#include "flang/Optimizer/Support/DataLayout.h"23#include "flang/Runtime/CUDA/registration.h"24#include "flang/Runtime/entry-names.h"25#include "mlir/Dialect/DLTI/DLTI.h"26#include "mlir/Dialect/GPU/IR/GPUDialect.h"27#include "mlir/Dialect/LLVMIR/LLVMDialect.h"28#include "mlir/IR/Value.h"29#include "mlir/Pass/Pass.h"30#include "llvm/ADT/SmallVector.h"31 32namespace fir {33#define GEN_PASS_DEF_CUFCOMPUTESHAREDMEMORYOFFSETSANDSIZE34#include "flang/Optimizer/Transforms/Passes.h.inc"35} // namespace fir36 37using namespace Fortran::runtime::cuda;38 39namespace {40 41static bool isAssumedSize(mlir::ValueRange shape) {42 if (shape.size() != 1)43 return false;44 if (llvm::isa_and_nonnull<fir::AssumedSizeExtentOp>(shape[0].getDefiningOp()))45 return true;46 return false;47}48 49struct CUFComputeSharedMemoryOffsetsAndSize50 : public fir::impl::CUFComputeSharedMemoryOffsetsAndSizeBase<51 CUFComputeSharedMemoryOffsetsAndSize> {52 53 void runOnOperation() override {54 mlir::ModuleOp mod = getOperation();55 mlir::SymbolTable symTab(mod);56 mlir::OpBuilder opBuilder{mod.getBodyRegion()};57 fir::FirOpBuilder builder(opBuilder, mod);58 fir::KindMapping kindMap{fir::getKindMapping(mod)};59 std::optional<mlir::DataLayout> dl =60 fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/false);61 if (!dl) {62 mlir::emitError(mod.getLoc(),63 "data layout attribute is required to perform " +64 getName() + "pass");65 }66 67 auto gpuMod = cuf::getOrCreateGPUModule(mod, symTab);68 mlir::Type i8Ty = builder.getI8Type();69 mlir::Type i32Ty = builder.getI32Type();70 mlir::Type idxTy = builder.getIndexType();71 for (auto funcOp : gpuMod.getOps<mlir::gpu::GPUFuncOp>()) {72 unsigned nbDynamicSharedVariables = 0;73 unsigned nbStaticSharedVariables = 0;74 uint64_t sharedMemSize = 0;75 unsigned short alignment = 0;76 mlir::Value crtDynOffset;77 78 // Go over each shared memory operation and compute their start offset and79 // the size and alignment of the global to be generated if all variables80 // are static. If this is dynamic shared memory, then only the alignment81 // is computed.82 for (auto sharedOp : funcOp.getOps<cuf::SharedMemoryOp>()) {83 mlir::Location loc = sharedOp.getLoc();84 builder.setInsertionPoint(sharedOp);85 if (fir::hasDynamicSize(sharedOp.getInType())) {86 mlir::Type ty = sharedOp.getInType();87 if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(ty))88 ty = seqTy.getEleTy();89 unsigned short align = dl->getTypeABIAlignment(ty);90 alignment = std::max(alignment, align);91 uint64_t tySize = dl->getTypeSize(ty);92 ++nbDynamicSharedVariables;93 if (isAssumedSize(sharedOp.getShape()) || !crtDynOffset) {94 mlir::Value zero = builder.createIntegerConstant(loc, i32Ty, 0);95 sharedOp.getOffsetMutable().assign(zero);96 } else {97 sharedOp.getOffsetMutable().assign(98 builder.createConvert(loc, i32Ty, crtDynOffset));99 }100 101 mlir::Value dynSize =102 builder.createIntegerConstant(loc, idxTy, tySize);103 for (auto extent : sharedOp.getShape())104 dynSize =105 mlir::arith::MulIOp::create(builder, loc, dynSize, extent);106 if (crtDynOffset)107 crtDynOffset = mlir::arith::AddIOp::create(builder, loc,108 crtDynOffset, dynSize);109 else110 crtDynOffset = dynSize;111 112 continue;113 }114 auto [size, align] = fir::getTypeSizeAndAlignmentOrCrash(115 sharedOp.getLoc(), sharedOp.getInType(), *dl, kindMap);116 ++nbStaticSharedVariables;117 mlir::Value offset = builder.createIntegerConstant(118 loc, i32Ty, llvm::alignTo(sharedMemSize, align));119 sharedOp.getOffsetMutable().assign(offset);120 sharedMemSize =121 llvm::alignTo(sharedMemSize, align) + llvm::alignTo(size, align);122 alignment = std::max(alignment, align);123 }124 125 if (nbDynamicSharedVariables == 0 && nbStaticSharedVariables == 0)126 continue;127 128 if (nbDynamicSharedVariables > 0 && nbStaticSharedVariables > 0)129 mlir::emitError(130 funcOp.getLoc(),131 "static and dynamic shared variables in a single kernel");132 133 mlir::DenseElementsAttr init = {};134 if (sharedMemSize > 0) {135 auto vecTy = mlir::VectorType::get(sharedMemSize, i8Ty);136 mlir::Attribute zero = mlir::IntegerAttr::get(i8Ty, 0);137 init = mlir::DenseElementsAttr::get(vecTy, llvm::ArrayRef(zero));138 }139 140 // Create the shared memory global where each shared variable will point141 // to.142 auto sharedMemType = fir::SequenceType::get(sharedMemSize, i8Ty);143 std::string sharedMemGlobalName =144 (funcOp.getName() + llvm::Twine(cudaSharedMemSuffix)).str();145 // Dynamic shared memory needs an external linkage while static shared146 // memory needs an internal linkage.147 mlir::StringAttr linkage = nbDynamicSharedVariables > 0148 ? builder.createExternalLinkage()149 : builder.createInternalLinkage();150 builder.setInsertionPointToEnd(gpuMod.getBody());151 llvm::SmallVector<mlir::NamedAttribute> attrs;152 auto globalOpName = mlir::OperationName(fir::GlobalOp::getOperationName(),153 gpuMod.getContext());154 attrs.push_back(mlir::NamedAttribute(155 fir::GlobalOp::getDataAttrAttrName(globalOpName),156 cuf::DataAttributeAttr::get(gpuMod.getContext(),157 cuf::DataAttribute::Shared)));158 auto sharedMem = fir::GlobalOp::create(159 builder, funcOp.getLoc(), sharedMemGlobalName, false, false,160 sharedMemType, init, linkage, attrs);161 sharedMem.setAlignment(alignment);162 }163 }164};165 166} // end anonymous namespace167