[Mlir-commits] [mlir] [mlir] Fix alignment for predicate (i1) vectors (PR #175961)

[Andrzej Warzyński]

https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/175961

Legal scalable predicate vectors (legal in the LLVM sense), e.g.
vector<[16]xi1> (or <vscale x 16 x i1>, using LLVM syntax) ought to have
alignment 2 rather than 16, see e.g. [1].

MLIR currently computes the vector “size in bits” as:

```cpp
  vecType.getNumElements()
    * dataLayout.getTypeSize(vecType.getElementType()) * 8
```

but `getTypeSize()` returns a size in *bytes* (rounded up from bits), so for
`i1` it returns 1. Multiplying by 8 converts that storage byte back to 8 bits
per element, which overestimates predicate vector sizes.

Instead, use:

```cpp
  vecType.getNumElements()
    * dataLayout.getTypeSizeInBits(vecType.getElementType())
```

For `vector<[16]xi1>` this changes:
  * [before]: `16 * (1 byte * 8) = 128 bits`
  * [after]:  `16 * 1 bit        = 16 bits`

This is a very small update that, based on the available tests, only
affects types like `vector<[16]xi1>`. It aligns MLIR with LLVM, making
sure that the corresponding alignment is 2 rather that 16. For context,
LLVM computes the alignment in this case via `getTypeStoreSize`, which
for `16 x i1` returns 2 bytes. Perhaps MLIR should follow similar path
in the future.

[1] https://developer.arm.com/documentation/ddi0602/2025-12/SVE-Instructions/LDR--predicate---Load-predicate-register-?lang=en


>From fc98105ef968dd8e3c1b371f72cd073504b45c13 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Wed, 14 Jan 2026 13:40:20 +0000
Subject: [PATCH] [mlir] Fix alignment for predicate (i1) vectors
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Legal scalable predicate vectors (legal in the LLVM sense), e.g.
vector<[16]xi1> (or <vscale x 16 x i1>, using LLVM syntax) ought to have
alignment 2 rather than 16, see e.g. [1].

MLIR currently computes the vector “size in bits” as:

```cpp
  vecType.getNumElements()
    * dataLayout.getTypeSize(vecType.getElementType()) * 8
```

but `getTypeSize()` returns a size in *bytes* (rounded up from bits), so for
`i1` it returns 1. Multiplying by 8 converts that storage byte back to 8 bits
per element, which overestimates predicate vector sizes.

Instead, use:

```cpp
  vecType.getNumElements()
    * dataLayout.getTypeSizeInBits(vecType.getElementType())
```

For `vector<[16]xi1>` this changes:
  * [before]: `16 * (1 byte * 8) = 128 bits`
  * [after]:  `16 * 1 bit        = 16 bits`

This is a very small update that, based on the available tests, only
affects types like `vector<[16]xi1>`. It aligns MLIR with LLVM, making
sure that the corresponding alignment is 2 rather that 16. For context,
LLVM computes the alignment in this case via `getTypeStoreSize`, which
for `16 x i1` returns 2 bytes. Perhaps MLIR should follow similar path
in the future.

[1] https://developer.arm.com/documentation/ddi0602/2025-12/SVE-Instructions/LDR--predicate---Load-predicate-register-?lang=en
---
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp         | 2 +-
 mlir/test/Interfaces/DataLayoutInterfaces/query.mlir | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index 782384999c70c..a6922ee5f4b5b 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -78,7 +78,7 @@ mlir::detail::getDefaultTypeSizeInBits(Type type, const DataLayout &dataLayout,
   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;
+                        dataLayout.getTypeSizeInBits(vecType.getElementType());
     return llvm::TypeSize::get(baseSize, vecType.isScalable());
   }
 
diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
index 5df32555000ad..97ef8b2a8ae1c 100644
--- a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
+++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir
@@ -44,6 +44,12 @@ func.func @no_layout_builtin() {
   // CHECK: preferred = 16
   // CHECK: size = {minimal_size = 16 : index, scalable}
   "test.data_layout_query"() : () -> vector<[4]xi32>
+  // CHECK: alignment = 2
+  // CHECK: bitsize = {minimal_size = 16 : index, scalable}
+  // CHECK: index = 0
+  // CHECK: preferred = 2
+  // CHECK: size = {minimal_size = 2 : index, scalable}
+  "test.data_layout_query"() : () -> vector<[16]xi1>
   return
 
 }