158 lines · cpp
1//===- MemoryPromotion.cpp - Utilities for moving data across GPU memories ===//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// This file implements utilities that allow one to create IR moving the data10// across different levels of the GPU memory hierarchy.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/GPU/Transforms/MemoryPromotion.h"15 16#include "mlir/Dialect/Affine/LoopUtils.h"17#include "mlir/Dialect/Arith/IR/Arith.h"18#include "mlir/Dialect/GPU/IR/GPUDialect.h"19#include "mlir/Dialect/MemRef/IR/MemRef.h"20#include "mlir/Dialect/SCF/IR/SCF.h"21 22using namespace mlir;23using namespace mlir::gpu;24 25/// Emits the (imperfect) loop nest performing the copy between "from" and "to"26/// values using the bounds derived from the "from" value. Emits at least27/// GPUDialect::getNumWorkgroupDimensions() loops, completing the nest with28/// single-iteration loops. Maps the innermost loops to thread dimensions, in29/// reverse order to enable access coalescing in the innermost loop.30static void insertCopyLoops(ImplicitLocOpBuilder &b, Value from, Value to) {31 auto memRefType = cast<MemRefType>(from.getType());32 auto rank = memRefType.getRank();33 34 SmallVector<Value, 4> lbs, ubs, steps;35 Value zero = arith::ConstantIndexOp::create(b, 0);36 Value one = arith::ConstantIndexOp::create(b, 1);37 38 // Make sure we have enough loops to use all thread dimensions, these trivial39 // loops should be outermost and therefore inserted first.40 if (rank < GPUDialect::getNumWorkgroupDimensions()) {41 unsigned extraLoops = GPUDialect::getNumWorkgroupDimensions() - rank;42 lbs.resize(extraLoops, zero);43 ubs.resize(extraLoops, one);44 steps.resize(extraLoops, one);45 }46 47 // Add existing bounds.48 lbs.append(rank, zero);49 ubs.reserve(lbs.size());50 steps.reserve(lbs.size());51 for (auto idx = 0; idx < rank; ++idx) {52 ubs.push_back(b.createOrFold<memref::DimOp>(from, idx));53 steps.push_back(one);54 }55 56 // Obtain thread identifiers and block sizes, necessary to map to them.57 auto indexType = b.getIndexType();58 SmallVector<Value, 3> threadIds, blockDims;59 for (auto dim : {gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z}) {60 threadIds.push_back(gpu::ThreadIdOp::create(b, indexType, dim));61 blockDims.push_back(gpu::BlockDimOp::create(b, indexType, dim));62 }63 64 // Produce the loop nest with copies.65 SmallVector<Value, 8> ivs(lbs.size());66 mlir::scf::buildLoopNest(67 b, b.getLoc(), lbs, ubs, steps,68 [&](OpBuilder &b, Location loc, ValueRange loopIvs) {69 ivs.assign(loopIvs.begin(), loopIvs.end());70 auto activeIvs = llvm::ArrayRef(ivs).take_back(rank);71 Value loaded = memref::LoadOp::create(b, loc, from, activeIvs);72 memref::StoreOp::create(b, loc, loaded, to, activeIvs);73 });74 75 // Map the innermost loops to threads in reverse order.76 for (const auto &en :77 llvm::enumerate(llvm::reverse(llvm::ArrayRef(ivs).take_back(78 GPUDialect::getNumWorkgroupDimensions())))) {79 Value v = en.value();80 auto loop = cast<scf::ForOp>(v.getParentRegion()->getParentOp());81 affine::mapLoopToProcessorIds(loop, {threadIds[en.index()]},82 {blockDims[en.index()]});83 }84}85 86/// Emits the loop nests performing the copy to the designated location in the87/// beginning of the region, and from the designated location immediately before88/// the terminator of the first block of the region. The region is expected to89/// have one block. This boils down to the following structure90///91/// ^bb(...):92/// <loop-bound-computation>93/// for %arg0 = ... to ... step ... {94/// ...95/// for %argN = <thread-id-x> to ... step <block-dim-x> {96/// %0 = load %from[%arg0, ..., %argN]97/// store %0, %to[%arg0, ..., %argN]98/// }99/// ...100/// }101/// gpu.barrier102/// <... original body ...>103/// gpu.barrier104/// for %arg0 = ... to ... step ... {105/// ...106/// for %argN = <thread-id-x> to ... step <block-dim-x> {107/// %1 = load %to[%arg0, ..., %argN]108/// store %1, %from[%arg0, ..., %argN]109/// }110/// ...111/// }112///113/// Inserts the barriers unconditionally since different threads may be copying114/// values and reading them. An analysis would be required to eliminate barriers115/// in case where value is only used by the thread that copies it. Both copies116/// are inserted unconditionally, an analysis would be required to only copy117/// live-in and live-out values when necessary. This copies the entire memref118/// pointed to by "from". In case a smaller block would be sufficient, the119/// caller can create a subview of the memref and promote it instead.120static void insertCopies(Region ®ion, Location loc, Value from, Value to) {121 auto fromType = cast<MemRefType>(from.getType());122 auto toType = cast<MemRefType>(to.getType());123 (void)fromType;124 (void)toType;125 assert(fromType.getShape() == toType.getShape());126 assert(fromType.getRank() != 0);127 assert(region.hasOneBlock() && "unstructured control flow not supported");128 129 auto b = ImplicitLocOpBuilder::atBlockBegin(loc, ®ion.front());130 insertCopyLoops(b, from, to);131 gpu::BarrierOp::create(b);132 133 b.setInsertionPoint(®ion.front().back());134 gpu::BarrierOp::create(b);135 insertCopyLoops(b, to, from);136}137 138/// Promotes a function argument to workgroup memory in the given function. The139/// copies will be inserted in the beginning and in the end of the function.140void mlir::promoteToWorkgroupMemory(GPUFuncOp op, unsigned arg) {141 Value value = op.getArgument(arg);142 auto type = dyn_cast<MemRefType>(value.getType());143 assert(type && type.hasStaticShape() && "can only promote memrefs");144 145 // Get the type of the buffer in the workgroup memory.146 auto workgroupMemoryAddressSpace = gpu::AddressSpaceAttr::get(147 op->getContext(), gpu::AddressSpace::Workgroup);148 auto bufferType = MemRefType::get(type.getShape(), type.getElementType(),149 MemRefLayoutAttrInterface{},150 Attribute(workgroupMemoryAddressSpace));151 Value attribution = op.addWorkgroupAttribution(bufferType, value.getLoc());152 153 // Replace the uses first since only the original uses are currently present.154 // Then insert the copies.155 value.replaceAllUsesWith(attribution);156 insertCopies(op.getBody(), op.getLoc(), value, attribution);157}158