148 lines · cpp
1//===- DistributionUtils.cpp - Distribution tools for GPUOps --------------===//2//3// Part of the MLIR 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 distribution utility methods.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/GPU/Utils/DistributionUtils.h"14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Arith/IR/Arith.h"16#include "mlir/IR/Value.h"17#include "llvm/ADT/DenseMap.h"18#include "llvm/ADT/STLExtras.h"19 20#include <numeric>21 22using namespace mlir;23using namespace mlir::gpu;24 25WarpExecuteOnLane0Op26WarpDistributionPattern::moveRegionToNewWarpOpAndReplaceReturns(27 RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,28 ValueRange newYieldedValues, TypeRange newReturnTypes) const {29 // Create a new op before the existing one, with the extra operands.30 OpBuilder::InsertionGuard g(rewriter);31 rewriter.setInsertionPoint(warpOp);32 auto newWarpOp = WarpExecuteOnLane0Op::create(33 rewriter, warpOp.getLoc(), newReturnTypes, warpOp.getLaneid(),34 warpOp.getWarpSize(), warpOp.getArgs(),35 warpOp.getBody()->getArgumentTypes());36 37 Region &opBody = warpOp.getBodyRegion();38 Region &newOpBody = newWarpOp.getBodyRegion();39 Block &newOpFirstBlock = newOpBody.front();40 rewriter.inlineRegionBefore(opBody, newOpBody, newOpBody.begin());41 rewriter.eraseBlock(&newOpFirstBlock);42 assert(newWarpOp.getWarpRegion().hasOneBlock() &&43 "expected WarpOp with single block");44 45 auto yield =46 cast<gpu::YieldOp>(newOpBody.getBlocks().begin()->getTerminator());47 48 rewriter.modifyOpInPlace(49 yield, [&]() { yield.getValuesMutable().assign(newYieldedValues); });50 return newWarpOp;51}52 53WarpExecuteOnLane0Op54WarpDistributionPattern::moveRegionToNewWarpOpAndAppendReturns(55 RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,56 ValueRange newYieldedValues, TypeRange newReturnTypes,57 SmallVector<size_t> &indices) const {58 SmallVector<Type> types(warpOp.getResultTypes().begin(),59 warpOp.getResultTypes().end());60 gpu::YieldOp yield = warpOp.getTerminator();61 SmallVector<Value> yieldValues(yield.getOperands().begin(),62 yield.getOperands().end());63 llvm::SmallDenseMap<Value, unsigned> indexLookup;64 // Record the value -> first index mapping for faster lookup.65 for (auto [i, v] : llvm::enumerate(yieldValues)) {66 if (!indexLookup.count(v))67 indexLookup[v] = i;68 }69 70 for (auto [value, type] : llvm::zip_equal(newYieldedValues, newReturnTypes)) {71 // If the value already exists in the yield, don't create a new output.72 if (indexLookup.count(value)) {73 indices.push_back(indexLookup[value]);74 } else {75 // If the value is new, add it to the yield and to the types.76 yieldValues.push_back(value);77 types.push_back(type);78 indices.push_back(yieldValues.size() - 1);79 }80 }81 82 WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndReplaceReturns(83 rewriter, warpOp, yieldValues, types);84 rewriter.replaceOp(warpOp,85 newWarpOp.getResults().take_front(warpOp.getNumResults()));86 return newWarpOp;87}88 89OpOperand *WarpDistributionPattern::getWarpResult(90 WarpExecuteOnLane0Op warpOp,91 llvm::function_ref<bool(Operation *)> fn) const {92 gpu::YieldOp yield = warpOp.getTerminator();93 for (OpOperand &yieldOperand : yield->getOpOperands()) {94 Value yieldValues = yieldOperand.get();95 Operation *definedOp = yieldValues.getDefiningOp();96 if (definedOp && fn(definedOp)) {97 if (!warpOp.getResult(yieldOperand.getOperandNumber()).use_empty())98 return &yieldOperand;99 }100 }101 return nullptr;102}103 104bool WarpDistributionPattern::delinearizeLaneId(105 OpBuilder &builder, Location loc, ArrayRef<int64_t> originalShape,106 ArrayRef<int64_t> distributedShape, int64_t warpSize, Value laneId,107 SmallVectorImpl<Value> &delinearizedIds) const {108 // If the original shape and the distributed shape is the same, we don't109 // distribute at all--every thread is handling the whole. For such case, we110 // should not rely on lane IDs later. So just return an empty lane ID vector.111 if (originalShape == distributedShape) {112 delinearizedIds.clear();113 return true;114 }115 116 SmallVector<int64_t> sizes;117 for (auto [large, small] : llvm::zip_equal(originalShape, distributedShape)) {118 if (large % small != 0)119 return false;120 sizes.push_back(large / small);121 }122 if (llvm::product_of(sizes) != warpSize)123 return false;124 125 AffineExpr s0, s1;126 bindSymbols(builder.getContext(), s0, s1);127 128 int64_t usedThreads = 1;129 130 Value zero = arith::ConstantIndexOp::create(builder, loc, 0);131 delinearizedIds.assign(sizes.size(), zero);132 133 for (int i = sizes.size() - 1; i >= 0; --i) {134 usedThreads *= sizes[i];135 if (usedThreads == warpSize) {136 // We've used up all available threads. Don't need to perform modulo137 // anymore. And we can stop the calculation for further dimensions.138 delinearizedIds[i] = laneId;139 break;140 }141 delinearizedIds[i] =142 affine::makeComposedAffineApply(builder, loc, s0 % sizes[i], {laneId});143 laneId = affine::makeComposedAffineApply(144 builder, loc, s0.floorDiv(usedThreads), {laneId});145 }146 return true;147}148