Retain address space during MLIR > LLVM conversion.

PiperOrigin-RevId: 267206460
This commit is contained in:
MLIR Team
2019-09-04 12:26:11 -07:00
committed by A. Unique TensorFlower
parent c6f8adad8e
commit b5652720c1
2 changed files with 20 additions and 12 deletions

View File

@@ -124,14 +124,14 @@ Type LLVMTypeConverter::convertFunctionType(FunctionType type) {
// Convert a MemRef to an LLVM type. If the memref is statically-shaped, then
// we return a pointer to the converted element type. Otherwise we return an
// LLVM stucture type, where the first element of the structure type is a
// LLVM structure type, where the first element of the structure type is a
// pointer to the elemental type of the MemRef and the following N elements are
// values of the Index type, one for each of N dynamic dimensions of the MemRef.
Type LLVMTypeConverter::convertMemRefType(MemRefType type) {
LLVM::LLVMType elementType = unwrap(convertType(type.getElementType()));
if (!elementType)
return {};
auto ptrType = elementType.getPointerTo();
auto ptrType = elementType.getPointerTo(type.getMemorySpace());
// Extra value for the memory space.
unsigned numDynamicSizes = type.getNumDynamicDims();
@@ -189,7 +189,7 @@ static Type getMemRefElementPtrType(MemRefType t, LLVMTypeConverter &lowering) {
auto converted = lowering.convertType(elementType);
if (!converted)
return {};
return converted.cast<LLVM::LLVMType>().getPointerTo();
return converted.cast<LLVM::LLVMType>().getPointerTo(t.getMemorySpace());
}
LLVMOpLowering::LLVMOpLowering(StringRef rootOpName, MLIRContext *context,
@@ -514,14 +514,11 @@ struct ConstLLVMOpLowering
using Super::Super;
};
// Check if the MemRefType `type` is supported by the lowering. We currently do
// not support memrefs with affine maps and non-default memory spaces.
// Check if the MemRefType `type` is supported by the lowering. We currently
// only support memrefs with identity maps.
static bool isSupportedMemRefType(MemRefType type) {
if (!type.getAffineMaps().empty())
return false;
if (type.getMemorySpace() != 0)
return false;
return true;
return llvm::all_of(type.getAffineMaps(),
[](AffineMap map) { return map.isIdentity(); });
}
// An `alloc` is converted into a definition of a memref descriptor value and
@@ -598,8 +595,8 @@ struct AllocOpLowering : public LLVMLegalizationPattern<AllocOp> {
cumulativeSize)
.getResult(0);
auto structElementType = lowering.convertType(elementType);
auto elementPtrType =
structElementType.cast<LLVM::LLVMType>().getPointerTo();
auto elementPtrType = structElementType.cast<LLVM::LLVMType>().getPointerTo(
type.getMemorySpace());
allocated = rewriter.create<LLVM::BitcastOp>(op->getLoc(), elementPtrType,
ArrayRef<Value *>(allocated));

View File

@@ -0,0 +1,11 @@
// RUN: mlir-opt %s -lower-to-llvm | FileCheck %s
// CHECK-LABEL: func @address_space(%{{.*}}: !llvm<"float addrspace(7)*">)
func @address_space(%arg0 : memref<32xf32, (d0) -> (d0), 7>) {
%0 = alloc() : memref<32xf32, (d0) -> (d0), 5>
%1 = constant 7 : index
// CHECK: llvm.load %{{.*}} : !llvm<"float addrspace(5)*">
%2 = load %0[%1] : memref<32xf32, (d0) -> (d0), 5>
std.return
}