[Mlir-commits] [mlir] [mlir][ABI] Fix scalable vector mapping in ABITypeMapper (PR #206617)

[Fedor Nikolaev]

https://github.com/felichita updated https://github.com/llvm/llvm-project/pull/206617

>From 63c23741f0c5cc9ac418dc5a484dbbae9b0f3337 Mon Sep 17 00:00:00 2001
From: Fedor Nikolaev <fridrixnm at gmail.com>
Date: Tue, 30 Jun 2026 01:40:44 +0200
Subject: [PATCH] [mlir][ABI] Fix scalable vector mapping in ABITypeMapper

mapVectorType always used ElementCount::getFixed, even for scalable
vectors. Use the correct scalable flag based on the VectorType.

ElementCount can only express a single scalable dimension, so vectors
with more than one scalable dim (e.g. vector<[2]x[4]xf32>) cannot be
represented correctly; return nullptr in that case instead of
silently computing a wrong size.
---
 mlir/lib/ABI/ABITypeMapper.cpp           | 10 +++++++-
 mlir/unittests/ABI/ABITypeMapperTest.cpp | 30 ++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/ABI/ABITypeMapper.cpp b/mlir/lib/ABI/ABITypeMapper.cpp
index 83dc6990ec5bb..109b70c93952f 100644
--- a/mlir/lib/ABI/ABITypeMapper.cpp
+++ b/mlir/lib/ABI/ABITypeMapper.cpp
@@ -70,12 +70,20 @@ const llvm::abi::Type *ABITypeMapper::mapVectorType(mlir::VectorType type) {
   if (!elementTy)
     return nullptr;
 
+  // ABI lowering cannot represent vectors with more than one scalable
+  // dimension (e.g. vector<[2]x[4]xf32>) since ElementCount only has a
+  // single scalable bit and would silently produce vscale*N instead of
+  // the correct vscale^2*N.
+  if (llvm::count(type.getScalableDims(), true) > 1)
+    return nullptr;
+
   auto shape = type.getShape();
   uint64_t totalElements = 1;
   for (int64_t dim : shape)
     totalElements *= dim;
 
-  llvm::ElementCount ec = llvm::ElementCount::getFixed(totalElements);
+  bool isScalable = llvm::is_contained(type.getScalableDims(), true);
+  llvm::ElementCount ec = llvm::ElementCount::get(totalElements, isScalable);
   uint64_t abiAlign = dl.getTypeABIAlignment(type);
   return builder.getVectorType(elementTy, ec, llvm::Align(abiAlign));
 }
diff --git a/mlir/unittests/ABI/ABITypeMapperTest.cpp b/mlir/unittests/ABI/ABITypeMapperTest.cpp
index 4a7989298a149..573ab9df9896f 100644
--- a/mlir/unittests/ABI/ABITypeMapperTest.cpp
+++ b/mlir/unittests/ABI/ABITypeMapperTest.cpp
@@ -170,4 +170,34 @@ TEST_F(ABITypeMapperTest, MapUnsignedI32) {
   EXPECT_FALSE(intTy->isSigned());
 }
 
+TEST_F(ABITypeMapperTest, MapScalableVectorOf4xF32) {
+  DataLayout dl(module);
+  ABITypeMapper mapper(dl);
+
+  auto f32 = Float32Type::get(&ctx);
+  auto vec = VectorType::get({4}, f32, /* scalableDims=*/true);
+  const llvm::abi::Type *result = mapper.map(vec);
+
+  ASSERT_NE(result, nullptr);
+  EXPECT_TRUE(result->isVector());
+
+  auto *vecTy = llvm::cast<llvm::abi::VectorType>(result);
+  EXPECT_TRUE(vecTy->getNumElements().isScalable());
+  EXPECT_EQ(vecTy->getNumElements().getKnownMinValue(), 4u);
+  EXPECT_TRUE(vecTy->getElementType()->isFloat());
+}
+
+TEST_F(ABITypeMapperTest, MapMultiScalableVectorReturnsNull) {
+  DataLayout dl(module);
+  ABITypeMapper mapper(dl);
+
+  auto f32 = Float32Type::get(&ctx);
+  auto vec = VectorType::get({2, 4}, f32, /*scalableDims=*/{true, true});
+  const llvm::abi::Type *result = mapper.map(vec);
+
+  // (vscale*2)*(vscale*4) = vscale^2*8,
+  // ElementCount -> must return nullptr
+  EXPECT_EQ(result, nullptr);
+}
+
 } // namespace