[Mlir-commits] [mlir] [MLIR][DataLayout] Add support for scalable vectors (PR #89349)
Christian Ulmann
llvmlistbot at llvm.org
Fri Apr 19 02:22:49 PDT 2024
https://github.com/Dinistro updated https://github.com/llvm/llvm-project/pull/89349
>From c241a68274702ca908e20ab77a723545f922b376 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Fri, 19 Apr 2024 06:22:21 +0000
Subject: [PATCH 1/3] [MLIR][DataLayout] Add support for scalable vectors
This commit extends the data layout to support scalable vectors. For
scalable vectors, the `TypeSize`'s scalable field is set accordingly,
and the alignment information remains the same as for normal vectors.
This behavior is in sync with what LLVM's data layout queries are
producing.
---
mlir/lib/Interfaces/DataLayoutInterfaces.cpp | 16 ++++++++++------
.../Interfaces/DataLayoutInterfaces/query.mlir | 12 ++++++++++++
.../lib/Dialect/DLTI/TestDataLayoutQuery.cpp | 11 +++++++++--
3 files changed, 31 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index e93a9efbb76c17..ecf83c020799ee 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -75,10 +75,12 @@ mlir::detail::getDefaultTypeSizeInBits(Type type, const DataLayout &dataLayout,
// there is no bit-packing at the moment element sizes are taken in bytes and
// multiplied with 8 bits.
// TODO: make this extensible.
- if (auto vecType = dyn_cast<VectorType>(type))
- return vecType.getNumElements() / vecType.getShape().back() *
- llvm::PowerOf2Ceil(vecType.getShape().back()) *
- dataLayout.getTypeSize(vecType.getElementType()) * 8;
+ if (auto vecType = dyn_cast<VectorType>(type)) {
+ uint64_t baseSize = vecType.getNumElements() / vecType.getShape().back() *
+ llvm::PowerOf2Ceil(vecType.getShape().back()) *
+ dataLayout.getTypeSize(vecType.getElementType()) * 8;
+ return llvm::TypeSize::get(baseSize, vecType.isScalable());
+ }
if (auto typeInterface = dyn_cast<DataLayoutTypeInterface>(type))
return typeInterface.getTypeSizeInBits(dataLayout, params);
@@ -138,9 +140,11 @@ getFloatTypeABIAlignment(FloatType fltType, const DataLayout &dataLayout,
uint64_t mlir::detail::getDefaultABIAlignment(
Type type, const DataLayout &dataLayout,
ArrayRef<DataLayoutEntryInterface> params) {
- // Natural alignment is the closest power-of-two number above.
+ // Natural alignment is the closest power-of-two number above. For scalable
+ // vectors, aligning them to the same as the base vector is sufficient.
+ // This should be consitent with LLVM.
if (isa<VectorType>(type))
- return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type));
+ return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type).getKnownMinValue());
if (auto fltType = dyn_cast<FloatType>(type))
return getFloatTypeABIAlignment(fltType, dataLayout, params);
diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
index d3bc91339d164b..11c06c38452c85 100644
--- a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
+++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
@@ -32,6 +32,18 @@ func.func @no_layout_builtin() {
// CHECK: preferred = 8
// CHECK: size = 8
"test.data_layout_query"() : () -> index
+ // CHECK: alignment = 16
+ // CHECK: bitsize = 128
+ // CHECK: index = 0
+ // CHECK: preferred = 16
+ // CHECK: size = 16
+ "test.data_layout_query"() : () -> vector<4xi32>
+ // CHECK: alignment = 16
+ // CHECK: bitsize = "scalable"
+ // CHECK: index = 0
+ // CHECK: preferred = 16
+ // CHECK: size = "scalable"
+ "test.data_layout_query"() : () -> vector<[4]xi32>
return
}
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index a4464bba7e8584..b5eb2d99e71140 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -46,9 +46,16 @@ struct TestDataLayoutQuery
Attribute programMemorySpace = layout.getProgramMemorySpace();
Attribute globalMemorySpace = layout.getGlobalMemorySpace();
uint64_t stackAlignment = layout.getStackAlignment();
+
+ auto convertTypeSizeToAttr = [&](llvm::TypeSize typeSize) -> Attribute {
+ if (typeSize.isScalable())
+ return builder.getStringAttr("scalable");
+ return builder.getIndexAttr(typeSize);
+ };
+
op->setAttrs(
- {builder.getNamedAttr("size", builder.getIndexAttr(size)),
- builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
+ {builder.getNamedAttr("size", convertTypeSizeToAttr(size)),
+ builder.getNamedAttr("bitsize", convertTypeSizeToAttr(bitsize)),
builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
builder.getNamedAttr("preferred", builder.getIndexAttr(preferred)),
builder.getNamedAttr("index", builder.getIndexAttr(index)),
>From 2859a05780ee029593c7de2d8caf3b6caa2a1ddb Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Fri, 19 Apr 2024 06:33:06 +0000
Subject: [PATCH 2/3] fix comment
---
mlir/lib/Interfaces/DataLayoutInterfaces.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index ecf83c020799ee..15cfb3dbaf7451 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -142,7 +142,6 @@ uint64_t mlir::detail::getDefaultABIAlignment(
ArrayRef<DataLayoutEntryInterface> params) {
// Natural alignment is the closest power-of-two number above. For scalable
// vectors, aligning them to the same as the base vector is sufficient.
- // This should be consitent with LLVM.
if (isa<VectorType>(type))
return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type).getKnownMinValue());
>From 5eb2bc22e74ad53eee3c49ec9e1556ce5055ab9a Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Fri, 19 Apr 2024 09:22:34 +0000
Subject: [PATCH 3/3] address review comments
---
mlir/test/Interfaces/DataLayoutInterfaces/query.mlir | 4 ++--
mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp | 12 +++++++++---
2 files changed, 11 insertions(+), 5 deletions(-)
diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
index 11c06c38452c85..5df32555000ad0 100644
--- a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
+++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
@@ -39,10 +39,10 @@ func.func @no_layout_builtin() {
// CHECK: size = 16
"test.data_layout_query"() : () -> vector<4xi32>
// CHECK: alignment = 16
- // CHECK: bitsize = "scalable"
+ // CHECK: bitsize = {minimal_size = 128 : index, scalable}
// CHECK: index = 0
// CHECK: preferred = 16
- // CHECK: size = "scalable"
+ // CHECK: size = {minimal_size = 16 : index, scalable}
"test.data_layout_query"() : () -> vector<[4]xi32>
return
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index b5eb2d99e71140..84f45b31603192 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -48,9 +48,15 @@ struct TestDataLayoutQuery
uint64_t stackAlignment = layout.getStackAlignment();
auto convertTypeSizeToAttr = [&](llvm::TypeSize typeSize) -> Attribute {
- if (typeSize.isScalable())
- return builder.getStringAttr("scalable");
- return builder.getIndexAttr(typeSize);
+ if (!typeSize.isScalable())
+ return builder.getIndexAttr(typeSize);
+
+ return builder.getDictionaryAttr({
+ builder.getNamedAttr("scalable", builder.getUnitAttr()),
+ builder.getNamedAttr(
+ "minimal_size",
+ builder.getIndexAttr(typeSize.getKnownMinValue())),
+ });
};
op->setAttrs(
More information about the Mlir-commits
mailing list