[mlir][ABI] Fix scalable vector mapping in ABITypeMapper (#206617)

mapVectorType in ABITypeMapper incorrectly used ElementCount::getFixed
for all vector types, including scalable ones (e.g. vector<[4]xf32>).
This results in scalable vectors being mapped as fixed-size vectors,
which produces incorrect ABI lowering for SVE and RISC-V V targets.

mlir/lib/ABI/ABITypeMapper.cpp

@@ -70,12 +70,19 @@ const llvm::abi::Type *ABITypeMapper::mapVectorType(mlir::VectorType type) {
   if (!elementTy)
     return nullptr;
 
+  // ElementCount supports at most one scalable dimension, hence
+  // vectors with more than one scalable dimension (e.g. vector<[2]x[4]xf32>)
+  // cannot be represented and mapVectorType must return nullptr in this case.
+  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);
+  llvm::ElementCount ec =
+      llvm::ElementCount::get(totalElements, type.isScalable());
   uint64_t abiAlign = dl.getTypeABIAlignment(type);
   return builder.getVectorType(elementTy, ec, llvm::Align(abiAlign));
 }

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, MapWithTwoScalableDimsReturnsNull) {
+  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);
+
+  // LLVM supports at most 1 scalable dimension, hence this case should
+  // return nullptr.
+  EXPECT_EQ(result, nullptr);
+}
+
 } // namespace