243 lines · cpp
1//===- MaskedloadToLoad.cpp - Lowers maskedload to load -------===//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 "mlir/Dialect/AMDGPU/Transforms/Passes.h"10 11#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"12#include "mlir/Dialect/Arith/IR/Arith.h"13#include "mlir/Dialect/MemRef/IR/MemRef.h"14#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"15#include "mlir/Dialect/SCF/IR/SCF.h"16#include "mlir/Dialect/Vector/IR/VectorOps.h"17#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"18#include "mlir/IR/BuiltinTypes.h"19#include "mlir/IR/OpDefinition.h"20#include "mlir/IR/PatternMatch.h"21#include "mlir/IR/TypeUtilities.h"22#include "mlir/Support/LogicalResult.h"23#include "mlir/Transforms/GreedyPatternRewriteDriver.h"24#include "llvm/Support/MathExtras.h"25 26namespace mlir::amdgpu {27#define GEN_PASS_DEF_AMDGPUMASKEDLOADTOLOADPASS28#include "mlir/Dialect/AMDGPU/Transforms/Passes.h.inc"29} // namespace mlir::amdgpu30 31using namespace mlir;32using namespace mlir::amdgpu;33 34/// This pattern supports lowering of: `vector.maskedload` to `vector.load`35/// and `arith.select` if the memref is in buffer address space.36static LogicalResult baseInBufferAddrSpace(PatternRewriter &rewriter,37 vector::MaskedLoadOp maskedOp) {38 auto memRefType = dyn_cast<MemRefType>(maskedOp.getBase().getType());39 if (!memRefType)40 return rewriter.notifyMatchFailure(maskedOp, "not a memref source");41 42 Attribute addrSpace = memRefType.getMemorySpace();43 if (!isa_and_nonnull<amdgpu::AddressSpaceAttr>(addrSpace))44 return rewriter.notifyMatchFailure(maskedOp, "no address space");45 46 if (dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace).getValue() !=47 amdgpu::AddressSpace::FatRawBuffer)48 return rewriter.notifyMatchFailure(maskedOp, "not in buffer address space");49 50 return success();51}52 53static Value createVectorLoadForMaskedLoad(OpBuilder &builder, Location loc,54 vector::MaskedLoadOp maskedOp,55 bool passthru) {56 VectorType vectorType = maskedOp.getVectorType();57 Value load = vector::LoadOp::create(58 builder, loc, vectorType, maskedOp.getBase(), maskedOp.getIndices());59 if (passthru)60 load = arith::SelectOp::create(builder, loc, vectorType, maskedOp.getMask(),61 load, maskedOp.getPassThru());62 return load;63}64 65/// Check if the given value comes from a broadcasted i1 condition.66static FailureOr<Value> matchFullMask(OpBuilder &b, Value val) {67 auto broadcastOp = val.getDefiningOp<vector::BroadcastOp>();68 if (!broadcastOp)69 return failure();70 if (isa<VectorType>(broadcastOp.getSourceType()))71 return failure();72 return broadcastOp.getSource();73}74 75static constexpr char kMaskedloadNeedsMask[] =76 "amdgpu.buffer_maskedload_needs_mask";77 78namespace {79 80struct MaskedLoadLowering final : OpRewritePattern<vector::MaskedLoadOp> {81 using OpRewritePattern::OpRewritePattern;82 83 LogicalResult matchAndRewrite(vector::MaskedLoadOp maskedOp,84 PatternRewriter &rewriter) const override {85 if (maskedOp->hasAttr(kMaskedloadNeedsMask))86 return failure();87 88 if (failed(baseInBufferAddrSpace(rewriter, maskedOp))) {89 return failure();90 }91 92 // Check if this is either a full inbounds load or an empty, oob load. If93 // so, take the fast path and don't generate an if condition, because we94 // know doing the oob load is always safe.95 if (succeeded(matchFullMask(rewriter, maskedOp.getMask()))) {96 Value load = createVectorLoadForMaskedLoad(rewriter, maskedOp.getLoc(),97 maskedOp, /*passthru=*/true);98 rewriter.replaceOp(maskedOp, load);99 return success();100 }101 102 Location loc = maskedOp.getLoc();103 Value src = maskedOp.getBase();104 105 VectorType vectorType = maskedOp.getVectorType();106 int64_t vectorSize = vectorType.getNumElements();107 int64_t elementBitWidth = vectorType.getElementTypeBitWidth();108 SmallVector<OpFoldResult> indices = maskedOp.getIndices();109 110 auto stridedMetadata =111 memref::ExtractStridedMetadataOp::create(rewriter, loc, src);112 SmallVector<OpFoldResult> strides =113 stridedMetadata.getConstifiedMixedStrides();114 SmallVector<OpFoldResult> sizes = stridedMetadata.getConstifiedMixedSizes();115 OpFoldResult offset = stridedMetadata.getConstifiedMixedOffset();116 memref::LinearizedMemRefInfo linearizedInfo;117 OpFoldResult linearizedIndices;118 std::tie(linearizedInfo, linearizedIndices) =119 memref::getLinearizedMemRefOffsetAndSize(rewriter, loc, elementBitWidth,120 elementBitWidth, offset, sizes,121 strides, indices);122 123 // delta = bufferSize - linearizedOffset124 Value vectorSizeOffset =125 arith::ConstantIndexOp::create(rewriter, loc, vectorSize);126 Value linearIndex =127 getValueOrCreateConstantIndexOp(rewriter, loc, linearizedIndices);128 Value totalSize = getValueOrCreateConstantIndexOp(129 rewriter, loc, linearizedInfo.linearizedSize);130 Value delta = arith::SubIOp::create(rewriter, loc, totalSize, linearIndex);131 132 // 1) check if delta < vectorSize133 Value isOutofBounds = arith::CmpIOp::create(134 rewriter, loc, arith::CmpIPredicate::ult, delta, vectorSizeOffset);135 136 // 2) check if (detla % elements_per_word != 0)137 Value elementsPerWord = arith::ConstantIndexOp::create(138 rewriter, loc, llvm::divideCeil(32, elementBitWidth));139 Value isNotWordAligned = arith::CmpIOp::create(140 rewriter, loc, arith::CmpIPredicate::ne,141 arith::RemUIOp::create(rewriter, loc, delta, elementsPerWord),142 arith::ConstantIndexOp::create(rewriter, loc, 0));143 144 // We take the fallback of maskedload default lowering only it is both145 // out-of-bounds and not word aligned. The fallback ensures correct results146 // when loading at the boundary of the buffer since buffer load returns147 // inconsistent zeros for the whole word when boundary is crossed.148 Value ifCondition =149 arith::AndIOp::create(rewriter, loc, isOutofBounds, isNotWordAligned);150 151 auto thenBuilder = [&](OpBuilder &builder, Location loc) {152 Operation *read = builder.clone(*maskedOp.getOperation());153 read->setAttr(kMaskedloadNeedsMask, builder.getUnitAttr());154 Value readResult = read->getResult(0);155 scf::YieldOp::create(builder, loc, readResult);156 };157 158 auto elseBuilder = [&](OpBuilder &builder, Location loc) {159 Value res = createVectorLoadForMaskedLoad(builder, loc, maskedOp,160 /*passthru=*/true);161 scf::YieldOp::create(rewriter, loc, res);162 };163 164 auto ifOp =165 scf::IfOp::create(rewriter, loc, ifCondition, thenBuilder, elseBuilder);166 167 rewriter.replaceOp(maskedOp, ifOp);168 169 return success();170 }171};172 173struct FullMaskedLoadToConditionalLoad174 : OpRewritePattern<vector::MaskedLoadOp> {175 using OpRewritePattern::OpRewritePattern;176 177 LogicalResult matchAndRewrite(vector::MaskedLoadOp loadOp,178 PatternRewriter &rewriter) const override {179 FailureOr<Value> maybeCond = matchFullMask(rewriter, loadOp.getMask());180 if (failed(maybeCond)) {181 return failure();182 }183 184 Value cond = maybeCond.value();185 auto trueBuilder = [&](OpBuilder &builder, Location loc) {186 Value res = createVectorLoadForMaskedLoad(builder, loc, loadOp,187 /*passthru=*/false);188 scf::YieldOp::create(rewriter, loc, res);189 };190 auto falseBuilder = [&](OpBuilder &builder, Location loc) {191 scf::YieldOp::create(rewriter, loc, loadOp.getPassThru());192 };193 auto ifOp = scf::IfOp::create(rewriter, loadOp.getLoc(), cond, trueBuilder,194 falseBuilder);195 rewriter.replaceOp(loadOp, ifOp);196 return success();197 }198};199 200struct FullMaskedStoreToConditionalStore201 : OpRewritePattern<vector::MaskedStoreOp> {202 using OpRewritePattern::OpRewritePattern;203 204 LogicalResult matchAndRewrite(vector::MaskedStoreOp storeOp,205 PatternRewriter &rewriter) const override {206 FailureOr<Value> maybeCond = matchFullMask(rewriter, storeOp.getMask());207 if (failed(maybeCond)) {208 return failure();209 }210 Value cond = maybeCond.value();211 212 auto trueBuilder = [&](OpBuilder &builder, Location loc) {213 vector::StoreOp::create(rewriter, loc, storeOp.getValueToStore(),214 storeOp.getBase(), storeOp.getIndices());215 scf::YieldOp::create(rewriter, loc);216 };217 auto ifOp =218 scf::IfOp::create(rewriter, storeOp.getLoc(), cond, trueBuilder);219 rewriter.replaceOp(storeOp, ifOp);220 return success();221 }222};223 224} // namespace225 226void mlir::amdgpu::populateAmdgpuMaskedloadToLoadPatterns(227 RewritePatternSet &patterns, PatternBenefit benefit) {228 patterns.add<MaskedLoadLowering, FullMaskedLoadToConditionalLoad,229 FullMaskedStoreToConditionalStore>(patterns.getContext(),230 benefit);231}232 233struct AmdgpuMaskedloadToLoadPass final234 : amdgpu::impl::AmdgpuMaskedloadToLoadPassBase<AmdgpuMaskedloadToLoadPass> {235 void runOnOperation() override {236 RewritePatternSet patterns(&getContext());237 populateAmdgpuMaskedloadToLoadPatterns(patterns);238 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {239 return signalPassFailure();240 }241 }242};243