42526d240c
This commit adds support for casting memrefs into fat raw buffer pointers to the AMDGPU dialect. Fat raw buffer pointers - or, in LLVM terms, ptr addrspcae(7), allow encapsulating a buffer descriptor (as produced by the make.buffer.rsrc intrinsic or provided from some API) into a pointer that supports ordinary pointer operations like load or store. This allows people to take advantage of the additional semantics that buffer_load and similar instructions provide without forcing the use of entirely separate amdgpu.raw_buffer_* operations. Operations on fat raw buffer pointers are translated to the corresponding LLVM intrinsics by the backend. This commit also goes and and defines a #amdgpu.address_space<> attribute so that AMDGPU-specific memory spaces can be represented. Only #amdgpu.address_space<fat_raw_buffer> will work correctly with the memref dialect, but the other possible address spaces are included for completeness. --------- Co-authored-by: Jakub Kuderski <kubakuderski@gmail.com> Co-authored-by: Prashant Kumar <pk5561@gmail.com>
80 lines
3.0 KiB
C++
80 lines
3.0 KiB
C++
//===- ResolveStridedMetadata.cpp - AMDGPU expand_strided_metadata ------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/AMDGPU/Transforms/Passes.h"
|
|
|
|
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
namespace mlir::amdgpu {
|
|
#define GEN_PASS_DEF_AMDGPURESOLVESTRIDEDMETADATAPASS
|
|
#include "mlir/Dialect/AMDGPU/Transforms/Passes.h.inc"
|
|
} // namespace mlir::amdgpu
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::amdgpu;
|
|
|
|
namespace {
|
|
struct AmdgpuResolveStridedMetadataPass
|
|
: public amdgpu::impl::AmdgpuResolveStridedMetadataPassBase<
|
|
AmdgpuResolveStridedMetadataPass> {
|
|
void runOnOperation() override;
|
|
};
|
|
|
|
struct ExtractStridedMetadataOnFatRawBufferCastFolder final
|
|
: public OpRewritePattern<memref::ExtractStridedMetadataOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp metadataOp,
|
|
PatternRewriter &rewriter) const override {
|
|
auto castOp = metadataOp.getSource().getDefiningOp<FatRawBufferCastOp>();
|
|
if (!castOp)
|
|
return rewriter.notifyMatchFailure(metadataOp,
|
|
"not a fat raw buffer cast");
|
|
Location loc = castOp.getLoc();
|
|
auto sourceMetadata = rewriter.create<memref::ExtractStridedMetadataOp>(
|
|
loc, castOp.getSource());
|
|
SmallVector<Value> results;
|
|
if (metadataOp.getBaseBuffer().use_empty()) {
|
|
results.push_back(nullptr);
|
|
} else {
|
|
auto baseBufferType =
|
|
cast<MemRefType>(metadataOp.getBaseBuffer().getType());
|
|
if (baseBufferType == castOp.getResult().getType()) {
|
|
results.push_back(castOp.getResult());
|
|
} else {
|
|
results.push_back(rewriter.create<memref::ReinterpretCastOp>(
|
|
loc, baseBufferType, castOp.getResult(), /*offset=*/0,
|
|
/*sizes=*/ArrayRef<int64_t>{}, /*strides=*/ArrayRef<int64_t>{}));
|
|
}
|
|
}
|
|
if (castOp.getResetOffset())
|
|
results.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
|
|
else
|
|
results.push_back(sourceMetadata.getOffset());
|
|
llvm::append_range(results, sourceMetadata.getSizes());
|
|
llvm::append_range(results, sourceMetadata.getStrides());
|
|
rewriter.replaceOp(metadataOp, results);
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void mlir::amdgpu::populateAmdgpuResolveStridedMetadataPatterns(
|
|
RewritePatternSet &patterns) {
|
|
patterns.add<ExtractStridedMetadataOnFatRawBufferCastFolder>(
|
|
patterns.getContext());
|
|
}
|
|
|
|
void AmdgpuResolveStridedMetadataPass::runOnOperation() {
|
|
RewritePatternSet patterns(&getContext());
|
|
populateAmdgpuResolveStridedMetadataPatterns(patterns);
|
|
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
|
|
signalPassFailure();
|
|
}
|