048922edbb
When converting a `memref.load` from the image address space to a `spirv.ImageFetch` ensure that we correctly map the load indices to width, height and depth. The lowering currently assumes a linear image tiling, that is row-major memory layout. This allows us to support any memref layout that is a permutation of the dimensions, more complex layouts are not currently supported. Because the ordering of the dimensions in the vector passed to image fetch is the opposite to that in the memref directions a final reversal of the mapped dimensions is always required. --------- Signed-off-by: Jack Frankland <jack.frankland@arm.com>
1115 lines
45 KiB
C++
1115 lines
45 KiB
C++
//===- MemRefToSPIRV.cpp - MemRef to SPIR-V Patterns ----------------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements patterns to convert MemRef dialect to SPIR-V dialect.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
|
|
#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
|
|
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
|
|
#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
|
|
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/MLIRContext.h"
|
|
#include "mlir/IR/Visitors.h"
|
|
#include <cassert>
|
|
#include <limits>
|
|
#include <optional>
|
|
|
|
#define DEBUG_TYPE "memref-to-spirv-pattern"
|
|
|
|
using namespace mlir;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Utility functions
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Returns the offset of the value in `targetBits` representation.
|
|
///
|
|
/// `srcIdx` is an index into a 1-D array with each element having `sourceBits`.
|
|
/// It's assumed to be non-negative.
|
|
///
|
|
/// When accessing an element in the array treating as having elements of
|
|
/// `targetBits`, multiple values are loaded in the same time. The method
|
|
/// returns the offset where the `srcIdx` locates in the value. For example, if
|
|
/// `sourceBits` equals to 8 and `targetBits` equals to 32, the x-th element is
|
|
/// located at (x % 4) * 8. Because there are four elements in one i32, and one
|
|
/// element has 8 bits.
|
|
static Value getOffsetForBitwidth(Location loc, Value srcIdx, int sourceBits,
|
|
int targetBits, OpBuilder &builder) {
|
|
assert(targetBits % sourceBits == 0);
|
|
Type type = srcIdx.getType();
|
|
IntegerAttr idxAttr = builder.getIntegerAttr(type, targetBits / sourceBits);
|
|
auto idx = builder.createOrFold<spirv::ConstantOp>(loc, type, idxAttr);
|
|
IntegerAttr srcBitsAttr = builder.getIntegerAttr(type, sourceBits);
|
|
auto srcBitsValue =
|
|
builder.createOrFold<spirv::ConstantOp>(loc, type, srcBitsAttr);
|
|
auto m = builder.createOrFold<spirv::UModOp>(loc, srcIdx, idx);
|
|
return builder.createOrFold<spirv::IMulOp>(loc, type, m, srcBitsValue);
|
|
}
|
|
|
|
/// Returns an adjusted spirv::AccessChainOp. Based on the
|
|
/// extension/capabilities, certain integer bitwidths `sourceBits` might not be
|
|
/// supported. During conversion if a memref of an unsupported type is used,
|
|
/// load/stores to this memref need to be modified to use a supported higher
|
|
/// bitwidth `targetBits` and extracting the required bits. For an accessing a
|
|
/// 1D array (spirv.array or spirv.rtarray), the last index is modified to load
|
|
/// the bits needed. The extraction of the actual bits needed are handled
|
|
/// separately. Note that this only works for a 1-D tensor.
|
|
static Value
|
|
adjustAccessChainForBitwidth(const SPIRVTypeConverter &typeConverter,
|
|
spirv::AccessChainOp op, int sourceBits,
|
|
int targetBits, OpBuilder &builder) {
|
|
assert(targetBits % sourceBits == 0);
|
|
const auto loc = op.getLoc();
|
|
Value lastDim = op->getOperand(op.getNumOperands() - 1);
|
|
Type type = lastDim.getType();
|
|
IntegerAttr attr = builder.getIntegerAttr(type, targetBits / sourceBits);
|
|
auto idx = builder.createOrFold<spirv::ConstantOp>(loc, type, attr);
|
|
auto indices = llvm::to_vector<4>(op.getIndices());
|
|
// There are two elements if this is a 1-D tensor.
|
|
assert(indices.size() == 2);
|
|
indices.back() = builder.createOrFold<spirv::SDivOp>(loc, lastDim, idx);
|
|
Type t = typeConverter.convertType(op.getComponentPtr().getType());
|
|
return spirv::AccessChainOp::create(builder, loc, t, op.getBasePtr(),
|
|
indices);
|
|
}
|
|
|
|
/// Casts the given `srcBool` into an integer of `dstType`.
|
|
static Value castBoolToIntN(Location loc, Value srcBool, Type dstType,
|
|
OpBuilder &builder) {
|
|
assert(srcBool.getType().isInteger(1));
|
|
if (dstType.isInteger(1))
|
|
return srcBool;
|
|
Value zero = spirv::ConstantOp::getZero(dstType, loc, builder);
|
|
Value one = spirv::ConstantOp::getOne(dstType, loc, builder);
|
|
return builder.createOrFold<spirv::SelectOp>(loc, dstType, srcBool, one,
|
|
zero);
|
|
}
|
|
|
|
/// Returns the `targetBits`-bit value shifted by the given `offset`, and cast
|
|
/// to the type destination type, and masked.
|
|
static Value shiftValue(Location loc, Value value, Value offset, Value mask,
|
|
OpBuilder &builder) {
|
|
IntegerType dstType = cast<IntegerType>(mask.getType());
|
|
int targetBits = static_cast<int>(dstType.getWidth());
|
|
int valueBits = value.getType().getIntOrFloatBitWidth();
|
|
assert(valueBits <= targetBits);
|
|
|
|
if (valueBits == 1) {
|
|
value = castBoolToIntN(loc, value, dstType, builder);
|
|
} else {
|
|
if (valueBits < targetBits) {
|
|
value = spirv::UConvertOp::create(
|
|
builder, loc, builder.getIntegerType(targetBits), value);
|
|
}
|
|
|
|
value = builder.createOrFold<spirv::BitwiseAndOp>(loc, value, mask);
|
|
}
|
|
return builder.createOrFold<spirv::ShiftLeftLogicalOp>(loc, value.getType(),
|
|
value, offset);
|
|
}
|
|
|
|
/// Returns true if the allocations of memref `type` generated from `allocOp`
|
|
/// can be lowered to SPIR-V.
|
|
static bool isAllocationSupported(Operation *allocOp, MemRefType type) {
|
|
if (isa<memref::AllocOp, memref::DeallocOp>(allocOp)) {
|
|
auto sc = dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
|
|
if (!sc || sc.getValue() != spirv::StorageClass::Workgroup)
|
|
return false;
|
|
} else if (isa<memref::AllocaOp>(allocOp)) {
|
|
auto sc = dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
|
|
if (!sc || sc.getValue() != spirv::StorageClass::Function)
|
|
return false;
|
|
} else {
|
|
return false;
|
|
}
|
|
|
|
// Currently only support static shape and int or float or vector of int or
|
|
// float element type.
|
|
if (!type.hasStaticShape())
|
|
return false;
|
|
|
|
Type elementType = type.getElementType();
|
|
if (auto vecType = dyn_cast<VectorType>(elementType))
|
|
elementType = vecType.getElementType();
|
|
return elementType.isIntOrFloat();
|
|
}
|
|
|
|
/// Returns the scope to use for atomic operations use for emulating store
|
|
/// operations of unsupported integer bitwidths, based on the memref
|
|
/// type. Returns std::nullopt on failure.
|
|
static std::optional<spirv::Scope> getAtomicOpScope(MemRefType type) {
|
|
auto sc = dyn_cast_or_null<spirv::StorageClassAttr>(type.getMemorySpace());
|
|
switch (sc.getValue()) {
|
|
case spirv::StorageClass::StorageBuffer:
|
|
return spirv::Scope::Device;
|
|
case spirv::StorageClass::Workgroup:
|
|
return spirv::Scope::Workgroup;
|
|
default:
|
|
break;
|
|
}
|
|
return {};
|
|
}
|
|
|
|
/// Casts the given `srcInt` into a boolean value.
|
|
static Value castIntNToBool(Location loc, Value srcInt, OpBuilder &builder) {
|
|
if (srcInt.getType().isInteger(1))
|
|
return srcInt;
|
|
|
|
auto one = spirv::ConstantOp::getZero(srcInt.getType(), loc, builder);
|
|
return builder.createOrFold<spirv::INotEqualOp>(loc, srcInt, one);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Operation conversion
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// Note that DRR cannot be used for the patterns in this file: we may need to
|
|
// convert type along the way, which requires ConversionPattern. DRR generates
|
|
// normal RewritePattern.
|
|
|
|
namespace {
|
|
|
|
/// Converts memref.alloca to SPIR-V Function variables.
|
|
class AllocaOpPattern final : public OpConversionPattern<memref::AllocaOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts an allocation operation to SPIR-V. Currently only supports lowering
|
|
/// to Workgroup memory when the size is constant. Note that this pattern needs
|
|
/// to be applied in a pass that runs at least at spirv.module scope since it
|
|
/// wil ladd global variables into the spirv.module.
|
|
class AllocOpPattern final : public OpConversionPattern<memref::AllocOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.automic_rmw operations to SPIR-V atomic operations.
|
|
class AtomicRMWOpPattern final
|
|
: public OpConversionPattern<memref::AtomicRMWOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::AtomicRMWOp atomicOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Removed a deallocation if it is a supported allocation. Currently only
|
|
/// removes deallocation if the memory space is workgroup memory.
|
|
class DeallocOpPattern final : public OpConversionPattern<memref::DeallocOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::DeallocOp operation, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.load to spirv.Load + spirv.AccessChain on integers.
|
|
class IntLoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.load to spirv.Load + spirv.AccessChain.
|
|
class LoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.load to spirv.Image + spirv.ImageFetch
|
|
class ImageLoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.store to spirv.Store on integers.
|
|
class IntStoreOpPattern final : public OpConversionPattern<memref::StoreOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.memory_space_cast to the appropriate spirv cast operations.
|
|
class MemorySpaceCastOpPattern final
|
|
: public OpConversionPattern<memref::MemorySpaceCastOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::MemorySpaceCastOp addrCastOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Converts memref.store to spirv.Store.
|
|
class StoreOpPattern final : public OpConversionPattern<memref::StoreOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
class ReinterpretCastPattern final
|
|
: public OpConversionPattern<memref::ReinterpretCastOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::ReinterpretCastOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
class CastPattern final : public OpConversionPattern<memref::CastOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::CastOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value src = adaptor.getSource();
|
|
Type srcType = src.getType();
|
|
|
|
const TypeConverter *converter = getTypeConverter();
|
|
Type dstType = converter->convertType(op.getType());
|
|
if (srcType != dstType)
|
|
return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
|
|
diag << "types doesn't match: " << srcType << " and " << dstType;
|
|
});
|
|
|
|
rewriter.replaceOp(op, src);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
/// Converts memref.extract_aligned_pointer_as_index to spirv.ConvertPtrToU.
|
|
class ExtractAlignedPointerAsIndexOpPattern final
|
|
: public OpConversionPattern<memref::ExtractAlignedPointerAsIndexOp> {
|
|
public:
|
|
using Base::Base;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(memref::ExtractAlignedPointerAsIndexOp extractOp,
|
|
OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// AllocaOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
AllocaOpPattern::matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
MemRefType allocType = allocaOp.getType();
|
|
if (!isAllocationSupported(allocaOp, allocType))
|
|
return rewriter.notifyMatchFailure(allocaOp, "unhandled allocation type");
|
|
|
|
// Get the SPIR-V type for the allocation.
|
|
Type spirvType = getTypeConverter()->convertType(allocType);
|
|
if (!spirvType)
|
|
return rewriter.notifyMatchFailure(allocaOp, "type conversion failed");
|
|
|
|
rewriter.replaceOpWithNewOp<spirv::VariableOp>(allocaOp, spirvType,
|
|
spirv::StorageClass::Function,
|
|
/*initializer=*/nullptr);
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// AllocOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
AllocOpPattern::matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
MemRefType allocType = operation.getType();
|
|
if (!isAllocationSupported(operation, allocType))
|
|
return rewriter.notifyMatchFailure(operation, "unhandled allocation type");
|
|
|
|
// Get the SPIR-V type for the allocation.
|
|
Type spirvType = getTypeConverter()->convertType(allocType);
|
|
if (!spirvType)
|
|
return rewriter.notifyMatchFailure(operation, "type conversion failed");
|
|
|
|
// Insert spirv.GlobalVariable for this allocation.
|
|
Operation *parent =
|
|
SymbolTable::getNearestSymbolTable(operation->getParentOp());
|
|
if (!parent)
|
|
return failure();
|
|
Location loc = operation.getLoc();
|
|
spirv::GlobalVariableOp varOp;
|
|
{
|
|
OpBuilder::InsertionGuard guard(rewriter);
|
|
Block &entryBlock = *parent->getRegion(0).begin();
|
|
rewriter.setInsertionPointToStart(&entryBlock);
|
|
auto varOps = entryBlock.getOps<spirv::GlobalVariableOp>();
|
|
std::string varName =
|
|
std::string("__workgroup_mem__") +
|
|
std::to_string(std::distance(varOps.begin(), varOps.end()));
|
|
varOp = spirv::GlobalVariableOp::create(rewriter, loc, spirvType, varName,
|
|
/*initializer=*/nullptr);
|
|
}
|
|
|
|
// Get pointer to global variable at the current scope.
|
|
rewriter.replaceOpWithNewOp<spirv::AddressOfOp>(operation, varOp);
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// AllocOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
AtomicRMWOpPattern::matchAndRewrite(memref::AtomicRMWOp atomicOp,
|
|
OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
if (isa<FloatType>(atomicOp.getType()))
|
|
return rewriter.notifyMatchFailure(atomicOp,
|
|
"unimplemented floating-point case");
|
|
|
|
auto memrefType = cast<MemRefType>(atomicOp.getMemref().getType());
|
|
std::optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
|
|
if (!scope)
|
|
return rewriter.notifyMatchFailure(atomicOp,
|
|
"unsupported memref memory space");
|
|
|
|
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
|
|
Type resultType = typeConverter.convertType(atomicOp.getType());
|
|
if (!resultType)
|
|
return rewriter.notifyMatchFailure(atomicOp,
|
|
"failed to convert result type");
|
|
|
|
auto loc = atomicOp.getLoc();
|
|
Value ptr =
|
|
spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
|
|
adaptor.getIndices(), loc, rewriter);
|
|
|
|
if (!ptr)
|
|
return failure();
|
|
|
|
#define ATOMIC_CASE(kind, spirvOp) \
|
|
case arith::AtomicRMWKind::kind: \
|
|
rewriter.replaceOpWithNewOp<spirv::spirvOp>( \
|
|
atomicOp, resultType, ptr, *scope, \
|
|
spirv::MemorySemantics::AcquireRelease, adaptor.getValue()); \
|
|
break
|
|
|
|
switch (atomicOp.getKind()) {
|
|
ATOMIC_CASE(addi, AtomicIAddOp);
|
|
ATOMIC_CASE(maxs, AtomicSMaxOp);
|
|
ATOMIC_CASE(maxu, AtomicUMaxOp);
|
|
ATOMIC_CASE(mins, AtomicSMinOp);
|
|
ATOMIC_CASE(minu, AtomicUMinOp);
|
|
ATOMIC_CASE(ori, AtomicOrOp);
|
|
ATOMIC_CASE(andi, AtomicAndOp);
|
|
default:
|
|
return rewriter.notifyMatchFailure(atomicOp, "unimplemented atomic kind");
|
|
}
|
|
|
|
#undef ATOMIC_CASE
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// DeallocOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
DeallocOpPattern::matchAndRewrite(memref::DeallocOp operation,
|
|
OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
MemRefType deallocType = cast<MemRefType>(operation.getMemref().getType());
|
|
if (!isAllocationSupported(operation, deallocType))
|
|
return rewriter.notifyMatchFailure(operation, "unhandled allocation type");
|
|
rewriter.eraseOp(operation);
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// LoadOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
struct MemoryRequirements {
|
|
spirv::MemoryAccessAttr memoryAccess;
|
|
IntegerAttr alignment;
|
|
};
|
|
|
|
/// Given an accessed SPIR-V pointer, calculates its alignment requirements, if
|
|
/// any.
|
|
static FailureOr<MemoryRequirements>
|
|
calculateMemoryRequirements(Value accessedPtr, bool isNontemporal,
|
|
uint64_t preferredAlignment) {
|
|
if (preferredAlignment >= std::numeric_limits<uint32_t>::max()) {
|
|
return failure();
|
|
}
|
|
|
|
MLIRContext *ctx = accessedPtr.getContext();
|
|
|
|
auto memoryAccess = spirv::MemoryAccess::None;
|
|
if (isNontemporal) {
|
|
memoryAccess = spirv::MemoryAccess::Nontemporal;
|
|
}
|
|
|
|
auto ptrType = cast<spirv::PointerType>(accessedPtr.getType());
|
|
bool mayOmitAlignment =
|
|
!preferredAlignment &&
|
|
ptrType.getStorageClass() != spirv::StorageClass::PhysicalStorageBuffer;
|
|
if (mayOmitAlignment) {
|
|
if (memoryAccess == spirv::MemoryAccess::None) {
|
|
return MemoryRequirements{spirv::MemoryAccessAttr{}, IntegerAttr{}};
|
|
}
|
|
return MemoryRequirements{spirv::MemoryAccessAttr::get(ctx, memoryAccess),
|
|
IntegerAttr{}};
|
|
}
|
|
|
|
// PhysicalStorageBuffers require the `Aligned` attribute.
|
|
// Other storage types may show an `Aligned` attribute.
|
|
auto pointeeType = dyn_cast<spirv::ScalarType>(ptrType.getPointeeType());
|
|
if (!pointeeType)
|
|
return failure();
|
|
|
|
// For scalar types, the alignment is determined by their size.
|
|
std::optional<int64_t> sizeInBytes = pointeeType.getSizeInBytes();
|
|
if (!sizeInBytes.has_value())
|
|
return failure();
|
|
|
|
memoryAccess |= spirv::MemoryAccess::Aligned;
|
|
auto memAccessAttr = spirv::MemoryAccessAttr::get(ctx, memoryAccess);
|
|
auto alignmentValue = preferredAlignment ? preferredAlignment : *sizeInBytes;
|
|
auto alignment = IntegerAttr::get(IntegerType::get(ctx, 32), alignmentValue);
|
|
return MemoryRequirements{memAccessAttr, alignment};
|
|
}
|
|
|
|
/// Given an accessed SPIR-V pointer and the original memref load/store
|
|
/// `memAccess` op, calculates the alignment requirements, if any. Takes into
|
|
/// account the alignment attributes applied to the load/store op.
|
|
template <class LoadOrStoreOp>
|
|
static FailureOr<MemoryRequirements>
|
|
calculateMemoryRequirements(Value accessedPtr, LoadOrStoreOp loadOrStoreOp) {
|
|
static_assert(
|
|
llvm::is_one_of<LoadOrStoreOp, memref::LoadOp, memref::StoreOp>::value,
|
|
"Must be called on either memref::LoadOp or memref::StoreOp");
|
|
|
|
return calculateMemoryRequirements(accessedPtr,
|
|
loadOrStoreOp.getNontemporal(),
|
|
loadOrStoreOp.getAlignment().value_or(0));
|
|
}
|
|
|
|
LogicalResult
|
|
IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto loc = loadOp.getLoc();
|
|
auto memrefType = cast<MemRefType>(loadOp.getMemref().getType());
|
|
if (!memrefType.getElementType().isSignlessInteger())
|
|
return failure();
|
|
|
|
auto memorySpaceAttr =
|
|
dyn_cast_if_present<spirv::StorageClassAttr>(memrefType.getMemorySpace());
|
|
if (!memorySpaceAttr)
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "missing memory space SPIR-V storage class attribute");
|
|
|
|
if (memorySpaceAttr.getValue() == spirv::StorageClass::Image)
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp,
|
|
"failed to lower memref in image storage class to storage buffer");
|
|
|
|
const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
|
|
Value accessChain =
|
|
spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
|
|
adaptor.getIndices(), loc, rewriter);
|
|
|
|
if (!accessChain)
|
|
return failure();
|
|
|
|
int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
|
|
bool isBool = srcBits == 1;
|
|
if (isBool)
|
|
srcBits = typeConverter.getOptions().boolNumBits;
|
|
|
|
auto pointerType = typeConverter.convertType<spirv::PointerType>(memrefType);
|
|
if (!pointerType)
|
|
return rewriter.notifyMatchFailure(loadOp, "failed to convert memref type");
|
|
|
|
Type pointeeType = pointerType.getPointeeType();
|
|
Type dstType;
|
|
if (typeConverter.allows(spirv::Capability::Kernel)) {
|
|
if (auto arrayType = dyn_cast<spirv::ArrayType>(pointeeType))
|
|
dstType = arrayType.getElementType();
|
|
else
|
|
dstType = pointeeType;
|
|
} else {
|
|
// For Vulkan we need to extract element from wrapping struct and array.
|
|
Type structElemType =
|
|
cast<spirv::StructType>(pointeeType).getElementType(0);
|
|
if (auto arrayType = dyn_cast<spirv::ArrayType>(structElemType))
|
|
dstType = arrayType.getElementType();
|
|
else
|
|
dstType = cast<spirv::RuntimeArrayType>(structElemType).getElementType();
|
|
}
|
|
int dstBits = dstType.getIntOrFloatBitWidth();
|
|
assert(dstBits % srcBits == 0);
|
|
|
|
// If the rewritten load op has the same bit width, use the loading value
|
|
// directly.
|
|
if (srcBits == dstBits) {
|
|
auto memoryRequirements = calculateMemoryRequirements(accessChain, loadOp);
|
|
if (failed(memoryRequirements))
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "failed to determine memory requirements");
|
|
|
|
auto [memoryAccess, alignment] = *memoryRequirements;
|
|
Value loadVal = spirv::LoadOp::create(rewriter, loc, accessChain,
|
|
memoryAccess, alignment);
|
|
if (isBool)
|
|
loadVal = castIntNToBool(loc, loadVal, rewriter);
|
|
rewriter.replaceOp(loadOp, loadVal);
|
|
return success();
|
|
}
|
|
|
|
// Bitcasting is currently unsupported for Kernel capability /
|
|
// spirv.PtrAccessChain.
|
|
if (typeConverter.allows(spirv::Capability::Kernel))
|
|
return failure();
|
|
|
|
auto accessChainOp = accessChain.getDefiningOp<spirv::AccessChainOp>();
|
|
if (!accessChainOp)
|
|
return failure();
|
|
|
|
// Assume that getElementPtr() works linearizely. If it's a scalar, the method
|
|
// still returns a linearized accessing. If the accessing is not linearized,
|
|
// there will be offset issues.
|
|
assert(accessChainOp.getIndices().size() == 2);
|
|
Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
|
|
srcBits, dstBits, rewriter);
|
|
auto memoryRequirements = calculateMemoryRequirements(adjustedPtr, loadOp);
|
|
if (failed(memoryRequirements))
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "failed to determine memory requirements");
|
|
|
|
auto [memoryAccess, alignment] = *memoryRequirements;
|
|
Value spvLoadOp = spirv::LoadOp::create(rewriter, loc, dstType, adjustedPtr,
|
|
memoryAccess, alignment);
|
|
|
|
// Shift the bits to the rightmost.
|
|
// ____XXXX________ -> ____________XXXX
|
|
Value lastDim = accessChainOp->getOperand(accessChainOp.getNumOperands() - 1);
|
|
Value offset = getOffsetForBitwidth(loc, lastDim, srcBits, dstBits, rewriter);
|
|
Value result = rewriter.createOrFold<spirv::ShiftRightArithmeticOp>(
|
|
loc, spvLoadOp.getType(), spvLoadOp, offset);
|
|
|
|
// Apply the mask to extract corresponding bits.
|
|
Value mask = rewriter.createOrFold<spirv::ConstantOp>(
|
|
loc, dstType, rewriter.getIntegerAttr(dstType, (1 << srcBits) - 1));
|
|
result =
|
|
rewriter.createOrFold<spirv::BitwiseAndOp>(loc, dstType, result, mask);
|
|
|
|
// Apply sign extension on the loading value unconditionally. The signedness
|
|
// semantic is carried in the operator itself, we relies other pattern to
|
|
// handle the casting.
|
|
IntegerAttr shiftValueAttr =
|
|
rewriter.getIntegerAttr(dstType, dstBits - srcBits);
|
|
Value shiftValue =
|
|
rewriter.createOrFold<spirv::ConstantOp>(loc, dstType, shiftValueAttr);
|
|
result = rewriter.createOrFold<spirv::ShiftLeftLogicalOp>(loc, dstType,
|
|
result, shiftValue);
|
|
result = rewriter.createOrFold<spirv::ShiftRightArithmeticOp>(
|
|
loc, dstType, result, shiftValue);
|
|
|
|
rewriter.replaceOp(loadOp, result);
|
|
|
|
assert(accessChainOp.use_empty());
|
|
rewriter.eraseOp(accessChainOp);
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
LoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto memrefType = cast<MemRefType>(loadOp.getMemref().getType());
|
|
if (memrefType.getElementType().isSignlessInteger())
|
|
return failure();
|
|
|
|
auto memorySpaceAttr =
|
|
dyn_cast_if_present<spirv::StorageClassAttr>(memrefType.getMemorySpace());
|
|
if (!memorySpaceAttr)
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "missing memory space SPIR-V storage class attribute");
|
|
|
|
if (memorySpaceAttr.getValue() == spirv::StorageClass::Image)
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp,
|
|
"failed to lower memref in image storage class to storage buffer");
|
|
|
|
Value loadPtr = spirv::getElementPtr(
|
|
*getTypeConverter<SPIRVTypeConverter>(), memrefType, adaptor.getMemref(),
|
|
adaptor.getIndices(), loadOp.getLoc(), rewriter);
|
|
|
|
if (!loadPtr)
|
|
return failure();
|
|
|
|
auto memoryRequirements = calculateMemoryRequirements(loadPtr, loadOp);
|
|
if (failed(memoryRequirements))
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "failed to determine memory requirements");
|
|
|
|
auto [memoryAccess, alignment] = *memoryRequirements;
|
|
rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp, loadPtr, memoryAccess,
|
|
alignment);
|
|
return success();
|
|
}
|
|
|
|
template <typename OpAdaptor>
|
|
static FailureOr<SmallVector<Value>>
|
|
extractLoadCoordsForComposite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) {
|
|
// At present we only support linear "tiling" as specified in Vulkan, this
|
|
// means that texels are assumed to be laid out in memory in a row-major
|
|
// order. This allows us to support any memref layout that is a permutation of
|
|
// the dimensions. Future work will pass an optional image layout to the
|
|
// rewrite pattern so that we can support optimized target specific tilings.
|
|
SmallVector<Value> indices = adaptor.getIndices();
|
|
AffineMap map = loadOp.getMemRefType().getLayout().getAffineMap();
|
|
if (!map.isPermutation())
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp,
|
|
"Cannot lower memrefs with memory layout which is not a permutation");
|
|
|
|
// The memrefs layout determines the dimension ordering so we need to follow
|
|
// the map to get the ordering of the dimensions/indices.
|
|
const unsigned dimCount = map.getNumDims();
|
|
SmallVector<Value, 3> coords(dimCount);
|
|
for (unsigned dim = 0; dim < dimCount; ++dim)
|
|
coords[map.getDimPosition(dim)] = indices[dim];
|
|
|
|
// We need to reverse the coordinates because the memref layout is slowest to
|
|
// fastest moving and the vector coordinates for the image op is fastest to
|
|
// slowest moving.
|
|
return llvm::to_vector(llvm::reverse(coords));
|
|
}
|
|
|
|
LogicalResult
|
|
ImageLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto memrefType = cast<MemRefType>(loadOp.getMemref().getType());
|
|
|
|
auto memorySpaceAttr =
|
|
dyn_cast_if_present<spirv::StorageClassAttr>(memrefType.getMemorySpace());
|
|
if (!memorySpaceAttr)
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "missing memory space SPIR-V storage class attribute");
|
|
|
|
if (memorySpaceAttr.getValue() != spirv::StorageClass::Image)
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "failed to lower memref in non-image storage class to image");
|
|
|
|
Value loadPtr = adaptor.getMemref();
|
|
auto memoryRequirements = calculateMemoryRequirements(loadPtr, loadOp);
|
|
if (failed(memoryRequirements))
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "failed to determine memory requirements");
|
|
|
|
const auto [memoryAccess, alignment] = *memoryRequirements;
|
|
|
|
if (!loadOp.getMemRefType().hasRank())
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp, "cannot lower unranked memrefs to SPIR-V images");
|
|
|
|
// We currently only support lowering of scalar memref elements to texels in
|
|
// the R[16|32][f|i|ui] formats. Future work will enable lowering of vector
|
|
// elements to texels in richer formats.
|
|
if (!isa<spirv::ScalarType>(loadOp.getMemRefType().getElementType()))
|
|
return rewriter.notifyMatchFailure(
|
|
loadOp,
|
|
"cannot lower memrefs who's element type is not a SPIR-V scalar type"
|
|
"to SPIR-V images");
|
|
|
|
// We currently only support sampled images since OpImageFetch does not work
|
|
// for plain images and the OpImageRead instruction needs to be materialized
|
|
// instead or texels need to be accessed via atomics through a texel pointer.
|
|
// Future work will generalize support to plain images.
|
|
auto convertedPointeeType = cast<spirv::PointerType>(
|
|
getTypeConverter()->convertType(loadOp.getMemRefType()));
|
|
if (!isa<spirv::SampledImageType>(convertedPointeeType.getPointeeType()))
|
|
return rewriter.notifyMatchFailure(loadOp,
|
|
"cannot lower memrefs which do not "
|
|
"convert to SPIR-V sampled images");
|
|
|
|
// Materialize the lowering.
|
|
Location loc = loadOp->getLoc();
|
|
auto imageLoadOp =
|
|
spirv::LoadOp::create(rewriter, loc, loadPtr, memoryAccess, alignment);
|
|
// Extract the image from the sampled image.
|
|
auto imageOp = spirv::ImageOp::create(rewriter, loc, imageLoadOp);
|
|
|
|
// Build a vector of coordinates or just a scalar index if we have a 1D image.
|
|
Value coords;
|
|
if (memrefType.getRank() == 1) {
|
|
coords = adaptor.getIndices()[0];
|
|
} else {
|
|
FailureOr<SmallVector<Value>> maybeCoords =
|
|
extractLoadCoordsForComposite(loadOp, adaptor, rewriter);
|
|
if (failed(maybeCoords))
|
|
return failure();
|
|
auto coordVectorType = VectorType::get({loadOp.getMemRefType().getRank()},
|
|
adaptor.getIndices().getType()[0]);
|
|
coords = spirv::CompositeConstructOp::create(rewriter, loc, coordVectorType,
|
|
maybeCoords.value());
|
|
}
|
|
|
|
// Fetch the value out of the image.
|
|
auto resultVectorType = VectorType::get({4}, loadOp.getType());
|
|
auto fetchOp = spirv::ImageFetchOp::create(
|
|
rewriter, loc, resultVectorType, imageOp, coords,
|
|
mlir::spirv::ImageOperandsAttr{}, ValueRange{});
|
|
|
|
// Note that because OpImageFetch returns a rank 4 vector we need to extract
|
|
// the elements corresponding to the load which will since we only support the
|
|
// R[16|32][f|i|ui] formats will always be the R(red) 0th vector element.
|
|
auto compositeExtractOp =
|
|
spirv::CompositeExtractOp::create(rewriter, loc, fetchOp, 0);
|
|
|
|
rewriter.replaceOp(loadOp, compositeExtractOp);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto memrefType = cast<MemRefType>(storeOp.getMemref().getType());
|
|
if (!memrefType.getElementType().isSignlessInteger())
|
|
return rewriter.notifyMatchFailure(storeOp,
|
|
"element type is not a signless int");
|
|
|
|
auto loc = storeOp.getLoc();
|
|
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
|
|
Value accessChain =
|
|
spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
|
|
adaptor.getIndices(), loc, rewriter);
|
|
|
|
if (!accessChain)
|
|
return rewriter.notifyMatchFailure(
|
|
storeOp, "failed to convert element pointer type");
|
|
|
|
int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
|
|
|
|
bool isBool = srcBits == 1;
|
|
if (isBool)
|
|
srcBits = typeConverter.getOptions().boolNumBits;
|
|
|
|
auto pointerType = typeConverter.convertType<spirv::PointerType>(memrefType);
|
|
if (!pointerType)
|
|
return rewriter.notifyMatchFailure(storeOp,
|
|
"failed to convert memref type");
|
|
|
|
Type pointeeType = pointerType.getPointeeType();
|
|
IntegerType dstType;
|
|
if (typeConverter.allows(spirv::Capability::Kernel)) {
|
|
if (auto arrayType = dyn_cast<spirv::ArrayType>(pointeeType))
|
|
dstType = dyn_cast<IntegerType>(arrayType.getElementType());
|
|
else
|
|
dstType = dyn_cast<IntegerType>(pointeeType);
|
|
} else {
|
|
// For Vulkan we need to extract element from wrapping struct and array.
|
|
Type structElemType =
|
|
cast<spirv::StructType>(pointeeType).getElementType(0);
|
|
if (auto arrayType = dyn_cast<spirv::ArrayType>(structElemType))
|
|
dstType = dyn_cast<IntegerType>(arrayType.getElementType());
|
|
else
|
|
dstType = dyn_cast<IntegerType>(
|
|
cast<spirv::RuntimeArrayType>(structElemType).getElementType());
|
|
}
|
|
|
|
if (!dstType)
|
|
return rewriter.notifyMatchFailure(
|
|
storeOp, "failed to determine destination element type");
|
|
|
|
int dstBits = static_cast<int>(dstType.getWidth());
|
|
assert(dstBits % srcBits == 0);
|
|
|
|
if (srcBits == dstBits) {
|
|
auto memoryRequirements = calculateMemoryRequirements(accessChain, storeOp);
|
|
if (failed(memoryRequirements))
|
|
return rewriter.notifyMatchFailure(
|
|
storeOp, "failed to determine memory requirements");
|
|
|
|
auto [memoryAccess, alignment] = *memoryRequirements;
|
|
Value storeVal = adaptor.getValue();
|
|
if (isBool)
|
|
storeVal = castBoolToIntN(loc, storeVal, dstType, rewriter);
|
|
rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, accessChain, storeVal,
|
|
memoryAccess, alignment);
|
|
return success();
|
|
}
|
|
|
|
// Bitcasting is currently unsupported for Kernel capability /
|
|
// spirv.PtrAccessChain.
|
|
if (typeConverter.allows(spirv::Capability::Kernel))
|
|
return failure();
|
|
|
|
auto accessChainOp = accessChain.getDefiningOp<spirv::AccessChainOp>();
|
|
if (!accessChainOp)
|
|
return failure();
|
|
|
|
// Since there are multiple threads in the processing, the emulation will be
|
|
// done with atomic operations. E.g., if the stored value is i8, rewrite the
|
|
// StoreOp to:
|
|
// 1) load a 32-bit integer
|
|
// 2) clear 8 bits in the loaded value
|
|
// 3) set 8 bits in the loaded value
|
|
// 4) store 32-bit value back
|
|
//
|
|
// Step 2 is done with AtomicAnd, and step 3 is done with AtomicOr (of the
|
|
// loaded 32-bit value and the shifted 8-bit store value) as another atomic
|
|
// step.
|
|
assert(accessChainOp.getIndices().size() == 2);
|
|
Value lastDim = accessChainOp->getOperand(accessChainOp.getNumOperands() - 1);
|
|
Value offset = getOffsetForBitwidth(loc, lastDim, srcBits, dstBits, rewriter);
|
|
|
|
// Create a mask to clear the destination. E.g., if it is the second i8 in
|
|
// i32, 0xFFFF00FF is created.
|
|
Value mask = rewriter.createOrFold<spirv::ConstantOp>(
|
|
loc, dstType, rewriter.getIntegerAttr(dstType, (1 << srcBits) - 1));
|
|
Value clearBitsMask = rewriter.createOrFold<spirv::ShiftLeftLogicalOp>(
|
|
loc, dstType, mask, offset);
|
|
clearBitsMask =
|
|
rewriter.createOrFold<spirv::NotOp>(loc, dstType, clearBitsMask);
|
|
|
|
Value storeVal = shiftValue(loc, adaptor.getValue(), offset, mask, rewriter);
|
|
Value adjustedPtr = adjustAccessChainForBitwidth(typeConverter, accessChainOp,
|
|
srcBits, dstBits, rewriter);
|
|
std::optional<spirv::Scope> scope = getAtomicOpScope(memrefType);
|
|
if (!scope)
|
|
return rewriter.notifyMatchFailure(storeOp, "atomic scope not available");
|
|
|
|
Value result = spirv::AtomicAndOp::create(
|
|
rewriter, loc, dstType, adjustedPtr, *scope,
|
|
spirv::MemorySemantics::AcquireRelease, clearBitsMask);
|
|
result = spirv::AtomicOrOp::create(
|
|
rewriter, loc, dstType, adjustedPtr, *scope,
|
|
spirv::MemorySemantics::AcquireRelease, storeVal);
|
|
|
|
// The AtomicOrOp has no side effect. Since it is already inserted, we can
|
|
// just remove the original StoreOp. Note that rewriter.replaceOp()
|
|
// doesn't work because it only accepts that the numbers of result are the
|
|
// same.
|
|
rewriter.eraseOp(storeOp);
|
|
|
|
assert(accessChainOp.use_empty());
|
|
rewriter.eraseOp(accessChainOp);
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// MemorySpaceCastOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult MemorySpaceCastOpPattern::matchAndRewrite(
|
|
memref::MemorySpaceCastOp addrCastOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Location loc = addrCastOp.getLoc();
|
|
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
|
|
if (!typeConverter.allows(spirv::Capability::Kernel))
|
|
return rewriter.notifyMatchFailure(
|
|
loc, "address space casts require kernel capability");
|
|
|
|
auto sourceType = dyn_cast<MemRefType>(addrCastOp.getSource().getType());
|
|
if (!sourceType)
|
|
return rewriter.notifyMatchFailure(
|
|
loc, "SPIR-V lowering requires ranked memref types");
|
|
auto resultType = cast<MemRefType>(addrCastOp.getResult().getType());
|
|
|
|
auto sourceStorageClassAttr =
|
|
dyn_cast_or_null<spirv::StorageClassAttr>(sourceType.getMemorySpace());
|
|
if (!sourceStorageClassAttr)
|
|
return rewriter.notifyMatchFailure(loc, [sourceType](Diagnostic &diag) {
|
|
diag << "source address space " << sourceType.getMemorySpace()
|
|
<< " must be a SPIR-V storage class";
|
|
});
|
|
auto resultStorageClassAttr =
|
|
dyn_cast_or_null<spirv::StorageClassAttr>(resultType.getMemorySpace());
|
|
if (!resultStorageClassAttr)
|
|
return rewriter.notifyMatchFailure(loc, [resultType](Diagnostic &diag) {
|
|
diag << "result address space " << resultType.getMemorySpace()
|
|
<< " must be a SPIR-V storage class";
|
|
});
|
|
|
|
spirv::StorageClass sourceSc = sourceStorageClassAttr.getValue();
|
|
spirv::StorageClass resultSc = resultStorageClassAttr.getValue();
|
|
|
|
Value result = adaptor.getSource();
|
|
Type resultPtrType = typeConverter.convertType(resultType);
|
|
if (!resultPtrType)
|
|
return rewriter.notifyMatchFailure(addrCastOp,
|
|
"failed to convert memref type");
|
|
|
|
Type genericPtrType = resultPtrType;
|
|
// SPIR-V doesn't have a general address space cast operation. Instead, it has
|
|
// conversions to and from generic pointers. To implement the general case,
|
|
// we use specific-to-generic conversions when the source class is not
|
|
// generic. Then when the result storage class is not generic, we convert the
|
|
// generic pointer (either the input on ar intermediate result) to that
|
|
// class. This also means that we'll need the intermediate generic pointer
|
|
// type if neither the source or destination have it.
|
|
if (sourceSc != spirv::StorageClass::Generic &&
|
|
resultSc != spirv::StorageClass::Generic) {
|
|
Type intermediateType =
|
|
MemRefType::get(sourceType.getShape(), sourceType.getElementType(),
|
|
sourceType.getLayout(),
|
|
rewriter.getAttr<spirv::StorageClassAttr>(
|
|
spirv::StorageClass::Generic));
|
|
genericPtrType = typeConverter.convertType(intermediateType);
|
|
}
|
|
if (sourceSc != spirv::StorageClass::Generic) {
|
|
result = spirv::PtrCastToGenericOp::create(rewriter, loc, genericPtrType,
|
|
result);
|
|
}
|
|
if (resultSc != spirv::StorageClass::Generic) {
|
|
result =
|
|
spirv::GenericCastToPtrOp::create(rewriter, loc, resultPtrType, result);
|
|
}
|
|
rewriter.replaceOp(addrCastOp, result);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto memrefType = cast<MemRefType>(storeOp.getMemref().getType());
|
|
if (memrefType.getElementType().isSignlessInteger())
|
|
return rewriter.notifyMatchFailure(storeOp, "signless int");
|
|
auto storePtr = spirv::getElementPtr(
|
|
*getTypeConverter<SPIRVTypeConverter>(), memrefType, adaptor.getMemref(),
|
|
adaptor.getIndices(), storeOp.getLoc(), rewriter);
|
|
|
|
if (!storePtr)
|
|
return rewriter.notifyMatchFailure(storeOp, "type conversion failed");
|
|
|
|
auto memoryRequirements = calculateMemoryRequirements(storePtr, storeOp);
|
|
if (failed(memoryRequirements))
|
|
return rewriter.notifyMatchFailure(
|
|
storeOp, "failed to determine memory requirements");
|
|
|
|
auto [memoryAccess, alignment] = *memoryRequirements;
|
|
rewriter.replaceOpWithNewOp<spirv::StoreOp>(
|
|
storeOp, storePtr, adaptor.getValue(), memoryAccess, alignment);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult ReinterpretCastPattern::matchAndRewrite(
|
|
memref::ReinterpretCastOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value src = adaptor.getSource();
|
|
auto srcType = dyn_cast<spirv::PointerType>(src.getType());
|
|
|
|
if (!srcType)
|
|
return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
|
|
diag << "invalid src type " << src.getType();
|
|
});
|
|
|
|
const TypeConverter *converter = getTypeConverter();
|
|
|
|
auto dstType = converter->convertType<spirv::PointerType>(op.getType());
|
|
if (dstType != srcType)
|
|
return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
|
|
diag << "invalid dst type " << op.getType();
|
|
});
|
|
|
|
OpFoldResult offset =
|
|
getMixedValues(adaptor.getStaticOffsets(), adaptor.getOffsets(), rewriter)
|
|
.front();
|
|
if (isZeroInteger(offset)) {
|
|
rewriter.replaceOp(op, src);
|
|
return success();
|
|
}
|
|
|
|
Type intType = converter->convertType(rewriter.getIndexType());
|
|
if (!intType)
|
|
return rewriter.notifyMatchFailure(op, "failed to convert index type");
|
|
|
|
Location loc = op.getLoc();
|
|
auto offsetValue = [&]() -> Value {
|
|
if (auto val = dyn_cast<Value>(offset))
|
|
return val;
|
|
|
|
int64_t attrVal = cast<IntegerAttr>(cast<Attribute>(offset)).getInt();
|
|
Attribute attr = rewriter.getIntegerAttr(intType, attrVal);
|
|
return rewriter.createOrFold<spirv::ConstantOp>(loc, intType, attr);
|
|
}();
|
|
|
|
rewriter.replaceOpWithNewOp<spirv::InBoundsPtrAccessChainOp>(
|
|
op, src, offsetValue, ValueRange());
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ExtractAlignedPointerAsIndexOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult ExtractAlignedPointerAsIndexOpPattern::matchAndRewrite(
|
|
memref::ExtractAlignedPointerAsIndexOp extractOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
|
|
Type indexType = typeConverter.getIndexType();
|
|
rewriter.replaceOpWithNewOp<spirv::ConvertPtrToUOp>(extractOp, indexType,
|
|
adaptor.getSource());
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Pattern population
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace mlir {
|
|
void populateMemRefToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
|
|
RewritePatternSet &patterns) {
|
|
patterns.add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
|
|
DeallocOpPattern, IntLoadOpPattern, ImageLoadOpPattern,
|
|
IntStoreOpPattern, LoadOpPattern, MemorySpaceCastOpPattern,
|
|
StoreOpPattern, ReinterpretCastPattern, CastPattern,
|
|
ExtractAlignedPointerAsIndexOpPattern>(typeConverter,
|
|
patterns.getContext());
|
|
}
|
|
} // namespace mlir
|