/// Returns the list of entries.
DataLayoutEntryListRef getEntries() const;
+ /// Returns the alloca memory space identifier.
+ StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
+
/// Parses an instance of this attribute.
static DataLayoutSpecAttr parse(AsmParser &parser);
}];
let extraClassDeclaration = [{
+ // Top level attribute name.
constexpr const static ::llvm::StringLiteral
kDataLayoutAttrName = "dlti.dl_spec";
+ // Constants used in entries.
constexpr const static ::llvm::StringLiteral
kDataLayoutEndiannessKey = "dlti.endianness";
constexpr const static ::llvm::StringLiteral
kDataLayoutEndiannessLittle = "little";
+
+ constexpr const static ::llvm::StringLiteral
+ kDataLayoutAllocaMemorySpaceKey = "dlti.alloca_memory_space";
}];
let useDefaultAttributePrinterParser = 1;
getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout,
ArrayRef<DataLayoutEntryInterface> params);
+/// Default handler for alloca memory space request. Dispatches to the
+/// DataLayoutInterface if specified, otherwise returns the default.
+Attribute getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry);
+
/// Given a list of data layout entries, returns a new list containing the
/// entries with keys having the given type ID, i.e. belonging to the same type
/// class.
/// Returns the preferred of the given type in the current scope.
unsigned getTypePreferredAlignment(Type t) const;
+ /// Returns the memory space used for AllocaOps.
+ Attribute getAllocaMemorySpace() const;
+
private:
/// Combined layout spec at the given scope.
const DataLayoutSpecInterface originalLayout;
mutable DenseMap<Type, unsigned> bitsizes;
mutable DenseMap<Type, unsigned> abiAlignments;
mutable DenseMap<Type, unsigned> preferredAlignments;
+
+ /// Cache for alloca memory space.
+ mutable std::optional<Attribute> allocaMemorySpace;
};
} // namespace mlir
/*methodName=*/"getEntries",
/*args=*/(ins)
>,
+ InterfaceMethod<
+ /*description=*/"Returns the alloca memory space identifier.",
+ /*retTy=*/"::mlir::StringAttr",
+ /*methodName=*/"getAllocaMemorySpaceIdentifier",
+ /*args=*/(ins "::mlir::MLIRContext *":$context)
+ >,
// Implementations may override this if they have an efficient lookup
// mechanism.
InterfaceMethod<
params);
}]
>,
+ StaticInterfaceMethod<
+ /*description=*/"Returns the memory space used by the ABI computed "
+ "using the relevant entries. The data layout object "
+ "can be used for recursive queries.",
+ /*retTy=*/"::mlir::Attribute",
+ /*methodName=*/"getAllocaMemorySpace",
+ /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::getDefaultAllocaMemorySpace(entry);
+ }]
+ >,
];
let verify = [{ return ::mlir::detail::verifyDataLayoutOp($_op); }];
//===----------------------------------------------------------------------===//
//
constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword;
+constexpr const StringLiteral
+ mlir::DLTIDialect::kDataLayoutAllocaMemorySpaceKey;
namespace mlir {
namespace impl {
return getImpl()->entries;
}
+StringAttr
+DataLayoutSpecAttr::getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
+ return Builder(context).getStringAttr(
+ DLTIDialect::kDataLayoutAllocaMemorySpaceKey);
+}
+
/// Parses an attribute with syntax
/// attr ::= `#target.` `dl_spec` `<` attr-list? `>`
/// attr-list ::= attr
<< DLTIDialect::kDataLayoutEndiannessBig << "' or '"
<< DLTIDialect::kDataLayoutEndiannessLittle << "'";
}
+ if (entryName == DLTIDialect::kDataLayoutAllocaMemorySpaceKey)
+ return success();
return emitError(loc) << "unknown data layout entry name: " << entryName;
}
};
reportMissingDataLayout(type);
}
+// Returns the memory space used for allocal operations if specified in the
+// given entry. If the entry is empty the default memory space represented by
+// an empty attribute is returned.
+Attribute
+mlir::detail::getDefaultAllocaMemorySpace(DataLayoutEntryInterface entry) {
+ if (entry == DataLayoutEntryInterface()) {
+ return Attribute();
+ }
+
+ return entry.getValue();
+}
+
DataLayoutEntryList
mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries,
TypeID typeID) {
mlir::DataLayout::DataLayout() : DataLayout(ModuleOp()) {}
mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
- : originalLayout(getCombinedDataLayout(op)), scope(op) {
+ : originalLayout(getCombinedDataLayout(op)), scope(op),
+ allocaMemorySpace(std::nullopt) {
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
checkMissingLayout(originalLayout, op);
collectParentLayouts(op, layoutStack);
}
mlir::DataLayout::DataLayout(ModuleOp op)
- : originalLayout(getCombinedDataLayout(op)), scope(op) {
+ : originalLayout(getCombinedDataLayout(op)), scope(op),
+ allocaMemorySpace(std::nullopt) {
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
checkMissingLayout(originalLayout, op);
collectParentLayouts(op, layoutStack);
});
}
+mlir::Attribute mlir::DataLayout::getAllocaMemorySpace() const {
+ checkValid();
+ MLIRContext *context = scope->getContext();
+ if (allocaMemorySpace)
+ return *allocaMemorySpace;
+ DataLayoutEntryInterface entry;
+ if (originalLayout)
+ entry = originalLayout.getSpecForIdentifier(
+ originalLayout.getAllocaMemorySpaceIdentifier(context));
+ if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
+ allocaMemorySpace = iface.getAllocaMemorySpace(entry);
+ else
+ allocaMemorySpace = detail::getDefaultAllocaMemorySpace(entry);
+ return *allocaMemorySpace;
+}
+
//===----------------------------------------------------------------------===//
// DataLayoutSpecInterface
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/DLTI/DLTI.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
#include "mlir/Tools/mlir-translate/Translation.h"
// Remaining unhandled default layout defaults
// e (little endian if not set)
// p[n]:64:64:64 (non zero address spaces have 64-bit properties)
+ // Alloca address space defaults to 0.
std::string append =
"p:64:64:64-S0-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f16:16:16-f64:"
- "64:64-f128:128:128-v64:64:64-v128:128:128-a:0:64";
+ "64:64-f128:128:128-v64:64:64-v128:128:128-a:0:64-A0";
if (layoutstr.empty())
layoutstr = append;
else
StringAttr::get(context, DLTIDialect::kDataLayoutEndiannessKey),
value);
entries.emplace_back(entry);
+ } else if (symbol == 'A') {
+ unsigned addressSpace;
+ if (parameter.getAsInteger(/*Radix=*/10, addressSpace))
+ return nullptr;
+ // Skip storing if generic address space is defined.
+ if (addressSpace != 0) {
+ auto entry = DataLayoutEntryAttr::get(
+ StringAttr::get(context,
+ DLTIDialect::kDataLayoutAllocaMemorySpaceKey),
+ mlir::Builder(context).getUI32IntegerAttr(addressSpace));
+ entries.emplace_back(entry);
+ }
}
}
layoutStream.flush();
continue;
}
+ if (key.getValue() == DLTIDialect::kDataLayoutAllocaMemorySpaceKey) {
+ auto value = entry.getValue().cast<IntegerAttr>();
+ if (value != 0) {
+ // Only emit non-default address space.
+ layoutStream << "A" << value;
+ layoutStream.flush();
+ }
+ continue;
+ }
emitError(*loc) << "unsupported data layout key " << key;
return failure();
}
// CHECK: @no_spec
func.func @no_spec() {
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<i8>
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<i32>
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<bf16>
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<!llvm.ptr<i8>>
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<i8, 3>
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 0
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<i8, 5>
// CHECK: alignment = 8
- // CHECK: bitsize = 64
+ // CHECK: alloca_memory_space = 0
+ // CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<5>
module attributes { dlti.dl_spec = #dlti.dl_spec<
#dlti.dl_entry<!llvm.ptr<i8>, dense<[32, 32, 64]> : vector<3xi32>>,
#dlti.dl_entry<!llvm.ptr<i8, 5>, dense<[64, 64, 64]> : vector<3xi32>>,
- #dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64]> : vector<3xi32>>
+ #dlti.dl_entry<!llvm.ptr<4>, dense<[32, 64, 64]> : vector<3xi32>>,
+ #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>
>} {
// CHECK: @spec
func.func @spec() {
// CHECK: alignment = 4
+ // CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<i8>
// CHECK: alignment = 4
+ // CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<i32>
// CHECK: alignment = 4
+ // CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<bf16>
// CHECK: alignment = 4
+ // CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<!llvm.ptr<i8>>
// CHECK: alignment = 4
+ // CHECK: alloca_memory_space = 5
// CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<i8, 3>
// CHECK: alignment = 8
+ // CHECK: alloca_memory_space = 5
// CHECK: bitsize = 64
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> !llvm.ptr<i8, 5>
// CHECK: alignment = 4
- // CHECK: bitsize = 32
+ // CHECK: alloca_memory_space = 5
+ // CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<3>
// CHECK: alignment = 8
- // CHECK: bitsize = 32
+ // CHECK: alloca_memory_space = 5
+ // CHECK: bitsize = 32
// CHECK: preferred = 8
// CHECK: size = 4
"test.data_layout_query"() : () -> !llvm.ptr<4>
unsigned bitsize = layout.getTypeSizeInBits(op.getType());
unsigned alignment = layout.getTypeABIAlignment(op.getType());
unsigned preferred = layout.getTypePreferredAlignment(op.getType());
+ Attribute allocaMemorySpace = layout.getAllocaMemorySpace();
op->setAttrs(
{builder.getNamedAttr("size", builder.getIndexAttr(size)),
builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
- builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))});
+ builder.getNamedAttr("preferred", builder.getIndexAttr(preferred)),
+ builder.getNamedAttr("alloca_memory_space",
+ allocaMemorySpace == Attribute()
+ ? builder.getUI32IntegerAttr(0)
+ : allocaMemorySpace)
+
+ });
});
}
};
namespace {
constexpr static llvm::StringLiteral kAttrName = "dltest.layout";
+constexpr static llvm::StringLiteral kAllocaKeyName =
+ "dltest.alloca_memory_space";
/// Trivial array storage for the custom data layout spec attribute, just a list
/// of entries.
}
DataLayoutEntryListRef getEntries() const { return getImpl()->entries; }
LogicalResult verifySpec(Location loc) { return success(); }
+ StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const {
+ return Builder(context).getStringAttr(kAllocaKeyName);
+ }
};
/// A type subject to data layout that exits the program if it is queried more
executed = true;
return 4;
}
+
+ Attribute getAllocaMemorySpace(DataLayoutEntryInterface entry) {
+ static bool executed = false;
+ if (executed)
+ llvm::report_fatal_error("repeated call");
+
+ executed = true;
+ return Attribute();
+ }
};
/// A types that is not subject to data layout.
EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u);
EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u);
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u);
+
+ EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
}
TEST(DataLayout, NullSpec) {
auto op =
cast<DataLayoutOpInterface>(module->getBody()->getOperations().front());
DataLayout layout(op);
+
EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u);
EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u);
EXPECT_EQ(layout.getTypeSizeInBits(IntegerType::get(&ctx, 42)), 8u * 42u);
EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u);
EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u);
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
+
+ EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
}
TEST(DataLayout, EmptySpec) {
EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u);
EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u);
EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u);
+
+ EXPECT_EQ(layout.getAllocaMemorySpace(), Attribute());
}
TEST(DataLayout, SpecWithEntries) {
const char *ir = R"MLIR(
"dltest.op_with_layout"() { dltest.layout = #dltest.spec<
#dlti.dl_entry<i42, 5>,
- #dlti.dl_entry<i16, 6>
+ #dlti.dl_entry<i16, 6>,
+ #dlti.dl_entry<"dltest.alloca_memory_space", 5 : i32>
> } : () -> ()
)MLIR";
EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u);
EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u);
EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u);
+
+ EXPECT_EQ(layout.getAllocaMemorySpace(), Builder(&ctx).getI32IntegerAttr(5));
}
TEST(DataLayout, Caching) {