[Mlir-commits] [mlir] [mlir][vector] Make CompressstoreOp + ExpandloadOp support scalable vectors (PR #210288)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Jul 21 03:37:45 PDT 2026
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/210288
>From d9c4ebfc9363b1c9fd3af3f7dbb04815dbf9d6a1 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 17 Jul 2026 09:54:59 +0000
Subject: [PATCH 1/2] [mlir][vector] Make CompressstoreOp` + ExpandloadOp
support scalable vectors
Extends `vector.compressstore` + `vector.expandload` to support scalable
vectors and updates relevant tests.
---
.../mlir/Dialect/Vector/IR/VectorOps.td | 4 +-
.../vector-to-llvm-interface.mlir | 53 ++++-
mlir/test/Dialect/Vector/invalid.mlir | 16 --
mlir/test/Dialect/Vector/ops.mlir | 20 ++
.../vector-dropleadunitdim-transforms.mlir | 24 ++
.../Dialect/Vector/vector-mem-transforms.mlir | 50 ++++
.../Dialect/Vector/CPU/ArmSVE/compress.mlir | 214 ++++++++++++++++++
7 files changed, 351 insertions(+), 30 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 6e0134fc0cdc6..ebfb379684dd3 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2289,7 +2289,7 @@ def Vector_ExpandLoadOp :
]>,
Arguments<(ins Arg<AnyMemRef, "", [MemRead]>:$base,
Variadic<Index>:$indices,
- FixedVectorOfNonZeroRankOf<[I1]>:$mask,
+ VectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$pass_thru,
OptionalAttr<IntValidAlignment<I64Attr>>: $alignment)>,
Results<(outs AnyVectorOfNonZeroRank:$result)> {
@@ -2381,7 +2381,7 @@ def Vector_CompressStoreOp :
]>,
Arguments<(ins Arg<AnyMemRef, "", [MemWrite]>:$base,
Variadic<Index>:$indices,
- FixedVectorOfNonZeroRankOf<[I1]>:$mask,
+ VectorOfNonZeroRankOf<[I1]>:$mask,
AnyVectorOfNonZeroRank:$valueToStore,
OptionalAttr<IntValidAlignment<I64Attr>>: $alignment)> {
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
index cc35c6600d12c..4cbea04e8076e 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
@@ -1944,13 +1944,13 @@ func.func @negative_scatter_on_strided_memref(%arg0: memref<?xf32, strided<[2],
// vector.expandload
//===----------------------------------------------------------------------===//
-func.func @expand_load_op(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2: vector<11xf32>) -> vector<11xf32> {
+func.func @expandload(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2: vector<11xf32>) -> vector<11xf32> {
%c0 = arith.constant 0: index
%0 = vector.expandload %arg0[%c0], %arg1, %arg2 : memref<?xf32>, vector<11xi1>, vector<11xf32> into vector<11xf32>
return %0 : vector<11xf32>
}
-// CHECK-LABEL: func @expand_load_op
+// CHECK-LABEL: func @expandload
// CHECK: %[[CO:.*]] = arith.constant 0 : index
// CHECK: %[[C:.*]] = builtin.unrealized_conversion_cast %[[CO]] : index to i64
// CHECK: %[[P:.*]] = llvm.getelementptr %{{.*}}[%[[C]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
@@ -1959,21 +1959,36 @@ func.func @expand_load_op(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2: vec
// -----
-func.func @expand_load_op_index(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>) -> vector<11xindex> {
+func.func @expandload_scalable(%arg0: memref<?xf32>, %arg1: vector<[11]xi1>, %arg2: vector<[11]xf32>) -> vector<[11]xf32> {
+ %c0 = arith.constant 0: index
+ %0 = vector.expandload %arg0[%c0], %arg1, %arg2 : memref<?xf32>, vector<[11]xi1>, vector<[11]xf32> into vector<[11]xf32>
+ return %0 : vector<[11]xf32>
+}
+
+// CHECK-LABEL: func @expandload_scalable
+// CHECK: %[[CO:.*]] = arith.constant 0 : index
+// CHECK: %[[C:.*]] = builtin.unrealized_conversion_cast %[[CO]] : index to i64
+// CHECK: %[[P:.*]] = llvm.getelementptr %{{.*}}[%[[C]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+// CHECK: %[[E:.*]] = "llvm.intr.masked.expandload"(%[[P]], %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<[11]xi1>, vector<[11]xf32>) -> vector<[11]xf32>
+// CHECK: return %[[E]] : vector<[11]xf32>
+
+// -----
+
+func.func @expandload_index(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>) -> vector<11xindex> {
%c0 = arith.constant 0: index
%0 = vector.expandload %arg0[%c0], %arg1, %arg2 : memref<?xindex>, vector<11xi1>, vector<11xindex> into vector<11xindex>
return %0 : vector<11xindex>
}
-// CHECK-LABEL: func @expand_load_op_index
+// CHECK-LABEL: func @expandload_index
// CHECK: %{{.*}} = "llvm.intr.masked.expandload"(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<11xi1>, vector<11xi64>) -> vector<11xi64>
// -----
-func.func @expand_load_op_with_alignment(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>, %c0: index) -> vector<11xindex> {
+func.func @expandload_with_alignment(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>, %c0: index) -> vector<11xindex> {
%0 = vector.expandload %arg0[%c0], %arg1, %arg2 { alignment = 8 } : memref<?xindex>, vector<11xi1>, vector<11xindex> into vector<11xindex>
return %0 : vector<11xindex>
}
-// CHECK-LABEL: func @expand_load_op_with_alignment
+// CHECK-LABEL: func @expandload_with_alignment
// CHECK: %{{.*}} = "llvm.intr.masked.expandload"(%{{.*}}, %{{.*}}, %{{.*}}) <{arg_attrs = [{llvm.align = 8 : i64}, {}, {}]}> : (!llvm.ptr, vector<11xi1>, vector<11xi64>) -> vector<11xi64>
// -----
@@ -1982,13 +1997,13 @@ func.func @expand_load_op_with_alignment(%arg0: memref<?xindex>, %arg1: vector<1
// vector.compressstore
//===----------------------------------------------------------------------===//
-func.func @compress_store_op(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2: vector<11xf32>) {
+func.func @compressstore(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2: vector<11xf32>) {
%c0 = arith.constant 0: index
vector.compressstore %arg0[%c0], %arg1, %arg2 : memref<?xf32>, vector<11xi1>, vector<11xf32>
return
}
-// CHECK-LABEL: func @compress_store_op
+// CHECK-LABEL: func @compressstore
// CHECK: %[[CO:.*]] = arith.constant 0 : index
// CHECK: %[[C:.*]] = builtin.unrealized_conversion_cast %[[CO]] : index to i64
// CHECK: %[[P:.*]] = llvm.getelementptr %{{.*}}[%[[C]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
@@ -1996,21 +2011,35 @@ func.func @compress_store_op(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2:
// -----
-func.func @compress_store_op_index(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>) {
+func.func @compressstore_scalable(%arg0: memref<?xf32>, %arg1: vector<[11]xi1>, %arg2: vector<[11]xf32>) {
+ %c0 = arith.constant 0: index
+ vector.compressstore %arg0[%c0], %arg1, %arg2 : memref<?xf32>, vector<[11]xi1>, vector<[11]xf32>
+ return
+}
+
+// CHECK-LABEL: func @compressstore_scalable
+// CHECK: %[[CO:.*]] = arith.constant 0 : index
+// CHECK: %[[C:.*]] = builtin.unrealized_conversion_cast %[[CO]] : index to i64
+// CHECK: %[[P:.*]] = llvm.getelementptr %{{.*}}[%[[C]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+// CHECK: "llvm.intr.masked.compressstore"(%{{.*}}, %[[P]], %{{.*}}) : (vector<[11]xf32>, !llvm.ptr, vector<[11]xi1>) -> ()
+
+// -----
+
+func.func @compressstore_index(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>) {
%c0 = arith.constant 0: index
vector.compressstore %arg0[%c0], %arg1, %arg2 : memref<?xindex>, vector<11xi1>, vector<11xindex>
return
}
-// CHECK-LABEL: func @compress_store_op_index
+// CHECK-LABEL: func @compressstore_index
// CHECK: "llvm.intr.masked.compressstore"(%{{.*}}, %{{.*}}, %{{.*}}) : (vector<11xi64>, !llvm.ptr, vector<11xi1>) -> ()
// -----
-func.func @compress_store_op_with_alignment(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>, %c0: index) {
+func.func @compressstore_with_alignment(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %arg2: vector<11xindex>, %c0: index) {
vector.compressstore %arg0[%c0], %arg1, %arg2 { alignment = 8 } : memref<?xindex>, vector<11xi1>, vector<11xindex>
return
}
-// CHECK-LABEL: func @compress_store_op_with_alignment
+// CHECK-LABEL: func @compressstore_with_alignment
// CHECK: "llvm.intr.masked.compressstore"(%{{.*}}, %{{.*}}, %{{.*}}) <{arg_attrs = [{}, {llvm.align = 8 : i64}, {}]}> : (vector<11xi64>, !llvm.ptr, vector<11xi1>) -> ()
// -----
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index aaa55cface958..d4f305f20b595 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1688,14 +1688,6 @@ func.func @expand_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>,
// -----
-func.func @expand_base_scalable(%base: memref<?xf32>, %mask: vector<[16]xi1>, %pass_thru: vector<[16]xf32>) {
- %c0 = arith.constant 0 : index
- // expected-error at +1 {{'vector.expandload' op operand #2 must be fixed-length vector of 1-bit signless integer values, but got 'vector<[16]xi1>}}
- %0 = vector.expandload %base[%c0], %mask, %pass_thru : memref<?xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
-}
-
-// -----
-
func.func @expand_dim_mask_mismatch(%base: memref<?xf32>, %mask: vector<17xi1>, %pass_thru: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error at +1 {{'vector.expandload' op expected result shape to match mask shape}}
@@ -1750,14 +1742,6 @@ func.func @compress_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1
// -----
-func.func @compress_scalable(%base: memref<?xf32>, %mask: vector<[16]xi1>, %value: vector<[16]xf32>) {
- %c0 = arith.constant 0 : index
- // expected-error at +1 {{'vector.compressstore' op operand #2 must be fixed-length vector of 1-bit signless integer values, but got 'vector<[16]xi1>}}
- vector.compressstore %base[%c0], %mask, %value : memref<?xf32>, vector<[16]xi1>, vector<[16]xf32>
-}
-
-// -----
-
func.func @compress_dim_mask_mismatch(%base: memref<?xf32>, %mask: vector<17xi1>, %value: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error at +1 {{'vector.compressstore' op expected valueToStore shape to match mask shape}}
diff --git a/mlir/test/Dialect/Vector/ops.mlir b/mlir/test/Dialect/Vector/ops.mlir
index e84bd3f1dce17..1b4109a3368ea 100644
--- a/mlir/test/Dialect/Vector/ops.mlir
+++ b/mlir/test/Dialect/Vector/ops.mlir
@@ -885,6 +885,16 @@ func.func @expand_and_compress(%base: memref<?xf32>, %mask: vector<16xi1>, %pass
return
}
+// CHECK-LABEL: @expand_and_compress_scalable
+func.func @expand_and_compress_scalable(%base: memref<?xf32>, %mask: vector<[16]xi1>, %pass_thru: vector<[16]xf32>) {
+ %c0 = arith.constant 0 : index
+ // CHECK: %[[X:.*]] = vector.expandload %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}} : memref<?xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
+ %0 = vector.expandload %base[%c0], %mask, %pass_thru : memref<?xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
+ // CHECK: vector.compressstore %{{.*}}[%{{.*}}], %{{.*}}, %[[X]] : memref<?xf32>, vector<[16]xi1>, vector<[16]xf32>
+ vector.compressstore %base[%c0], %mask, %0 : memref<?xf32>, vector<[16]xi1>, vector<[16]xf32>
+ return
+}
+
// CHECK-LABEL: @expand_and_compress2d
func.func @expand_and_compress2d(%base: memref<?x?xf32>, %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
%c0 = arith.constant 0 : index
@@ -895,6 +905,16 @@ func.func @expand_and_compress2d(%base: memref<?x?xf32>, %mask: vector<16xi1>, %
return
}
+// CHECK-LABEL: @expand_and_compress2d_scalable
+func.func @expand_and_compress2d_scalable(%base: memref<?x?xf32>, %mask: vector<[16]xi1>, %pass_thru: vector<[16]xf32>) {
+ %c0 = arith.constant 0 : index
+ // CHECK: %[[X:.*]] = vector.expandload %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}, %{{.*}} : memref<?x?xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
+ %0 = vector.expandload %base[%c0, %c0], %mask, %pass_thru : memref<?x?xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
+ // CHECK: vector.compressstore %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}, %[[X]] : memref<?x?xf32>, vector<[16]xi1>, vector<[16]xf32>
+ vector.compressstore %base[%c0, %c0], %mask, %0 : memref<?x?xf32>, vector<[16]xi1>, vector<[16]xf32>
+ return
+}
+
// CHECK-LABEL: @multi_reduction
func.func @multi_reduction(%0: vector<4x8x16x32xf32>, %acc0: vector<4x16xf32>,
%acc1: f32) -> f32 {
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index bf01c8a8589d9..ff978665125ed 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -733,6 +733,19 @@ func.func @cast_away_expandload_leading_one_dims(%base: memref<16xf32>, %i: inde
// -----
+// CHECK-LABEL: func.func @cast_away_expandload_leading_one_dims_scalable
+// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<[4]xi1> from vector<1x[4]xi1>
+// CHECK: %[[P:.+]] = vector.extract %{{.*}}[0] : vector<[4]xf32> from vector<1x[4]xf32>
+// CHECK: %[[L:.+]] = vector.expandload %{{.*}}[%{{.*}}], %[[M]], %[[P]] : memref<16xf32>, vector<[4]xi1>, vector<[4]xf32> into vector<[4]xf32>
+// CHECK: %[[B:.+]] = vector.broadcast %[[L]] : vector<[4]xf32> to vector<1x[4]xf32>
+// CHECK: return %[[B]] : vector<1x[4]xf32>
+func.func @cast_away_expandload_leading_one_dims_scalable(%base: memref<16xf32>, %i: index, %mask: vector<1x[4]xi1>, %pass: vector<1x[4]xf32>) -> vector<1x[4]xf32> {
+ %0 = vector.expandload %base[%i], %mask, %pass : memref<16xf32>, vector<1x[4]xi1>, vector<1x[4]xf32> into vector<1x[4]xf32>
+ return %0 : vector<1x[4]xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @cast_away_gather_leading_one_dims
// CHECK: %[[I:.+]] = vector.extract %{{.*}}[0] : vector<4xi32> from vector<1x4xi32>
// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
@@ -779,6 +792,17 @@ func.func @cast_away_compressstore_leading_one_dims(%base: memref<16xf32>, %i: i
// -----
+// CHECK-LABEL: func.func @cast_away_compressstore_leading_one_dims_scalable
+// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<[4]xi1> from vector<1x[4]xi1>
+// CHECK: %[[V:.+]] = vector.extract %{{.*}}[0] : vector<[4]xf32> from vector<1x[4]xf32>
+// CHECK: vector.compressstore %{{.*}}[%{{.*}}], %[[M]], %[[V]] : memref<16xf32>, vector<[4]xi1>, vector<[4]xf32>
+func.func @cast_away_compressstore_leading_one_dims_scalable(%base: memref<16xf32>, %i: index, %mask: vector<1x[4]xi1>, %val: vector<1x[4]xf32>) {
+ vector.compressstore %base[%i], %mask, %val : memref<16xf32>, vector<1x[4]xi1>, vector<1x[4]xf32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: func.func @cast_away_scatter_leading_one_dims
// CHECK: %[[I:.+]] = vector.extract %{{.*}}[0] : vector<4xi32> from vector<1x4xi32>
// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
diff --git a/mlir/test/Dialect/Vector/vector-mem-transforms.mlir b/mlir/test/Dialect/Vector/vector-mem-transforms.mlir
index 2004a47851e2e..cd7ea84cdb2eb 100644
--- a/mlir/test/Dialect/Vector/vector-mem-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-mem-transforms.mlir
@@ -175,6 +175,20 @@ func.func @fold_expandload_all_true(%base: memref<16xf32>, %pass_thru: vector<16
return %ld : vector<16xf32>
}
+// CHECK-LABEL: func @fold_expandload_all_true_scalable(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[PASS_THRU:.*]]: vector<[16]xf32>) -> vector<[16]xf32> {
+// CHECK-DAG: %[[C:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %[[T:.*]] = vector.load %[[BASE]][%[[C]]] : memref<16xf32>, vector<[16]xf32>
+// CHECK-NEXT: return %[[T]] : vector<[16]xf32>
+func.func @fold_expandload_all_true_scalable(%base: memref<16xf32>, %pass_thru: vector<[16]xf32>) -> vector<[16]xf32> {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [16] : vector<[16]xi1>
+ %ld = vector.expandload %base[%c0], %mask, %pass_thru
+ : memref<16xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
+ return %ld : vector<[16]xf32>
+}
+
// CHECK-LABEL: func @fold_expandload_all_false(
// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
// CHECK-SAME: %[[PASS_THRU:.*]]: vector<16xf32>) -> vector<16xf32> {
@@ -187,6 +201,18 @@ func.func @fold_expandload_all_false(%base: memref<16xf32>, %pass_thru: vector<1
return %ld : vector<16xf32>
}
+// CHECK-LABEL: func @fold_expandload_all_false_scalable(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[PASS_THRU:.*]]: vector<[16]xf32>) -> vector<[16]xf32> {
+// CHECK-NEXT: return %[[PASS_THRU]] : vector<[16]xf32>
+func.func @fold_expandload_all_false_scalable(%base: memref<16xf32>, %pass_thru: vector<[16]xf32>) -> vector<[16]xf32> {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [0] : vector<[16]xi1>
+ %ld = vector.expandload %base[%c0], %mask, %pass_thru
+ : memref<16xf32>, vector<[16]xi1>, vector<[16]xf32> into vector<[16]xf32>
+ return %ld : vector<[16]xf32>
+}
+
//-----------------------------------------------------------------------------
// [Pattern: CompressStoreFolder]
//-----------------------------------------------------------------------------
@@ -204,6 +230,19 @@ func.func @fold_compressstore_all_true(%base: memref<16xf32>, %value: vector<16x
return
}
+// CHECK-LABEL: func @fold_compressstore_all_true_scalable(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[VALUE:.*]]: vector<[16]xf32>) {
+// CHECK-NEXT: %[[C:.*]] = arith.constant 0 : index
+// CHECK-NEXT: vector.store %[[VALUE]], %[[BASE]][%[[C]]] : memref<16xf32>, vector<[16]xf32>
+// CHECK-NEXT: return
+func.func @fold_compressstore_all_true_scalable(%base: memref<16xf32>, %value: vector<[16]xf32>) {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [16] : vector<[16]xi1>
+ vector.compressstore %base[%c0], %mask, %value : memref<16xf32>, vector<[16]xi1>, vector<[16]xf32>
+ return
+}
+
// CHECK-LABEL: func @fold_compressstore_all_false(
// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
// CHECK-SAME: %[[VALUE:.*]]: vector<16xf32>) {
@@ -214,3 +253,14 @@ func.func @fold_compressstore_all_false(%base: memref<16xf32>, %value: vector<16
vector.compressstore %base[%c0], %mask, %value : memref<16xf32>, vector<16xi1>, vector<16xf32>
return
}
+
+// CHECK-LABEL: func @fold_compressstore_all_false_scalable(
+// CHECK-SAME: %[[BASE:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[VALUE:.*]]: vector<[16]xf32>) {
+// CHECK-NEXT: return
+func.func @fold_compressstore_all_false_scalable(%base: memref<16xf32>, %value: vector<[16]xf32>) {
+ %c0 = arith.constant 0 : index
+ %mask = vector.constant_mask [0] : vector<[16]xi1>
+ vector.compressstore %base[%c0], %mask, %value : memref<16xf32>, vector<[16]xi1>, vector<[16]xf32>
+ return
+}
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
new file mode 100644
index 0000000000000..d8752cfb7fc92
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
@@ -0,0 +1,214 @@
+// REQUIRES: arm-emulator
+
+/// End-to-end test for vector.compressstore for SVE
+
+// In order to demonstrate the impact of using scalable vectors, vscale is set
+// to 2 so that vector<[16]xi32> constains 32 rather than 16 elements at
+// run-time
+//
+// Note that you can also tweak the size of vscale by passing this flag to
+// QEMU:
+// * -cpu max,sve-max-vq=[1-16]
+// (select the value between 1 and 16).
+
+// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm
+// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+sve"\
+// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils,%native_mlir_arm_runner_utils
+
+// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// @compress_16
+//
+// The number of inserted elements is 16 x vscale. Insertion index is
+// hard-coded to 0
+//===----------------------------------------------------------------------===//
+func.func @compress_16(%base: memref<?xi32>,
+ %mask: vector<[16]xi1>, %value: vector<[16]xi32>) {
+ %c0 = arith.constant 0: index
+ vector.compressstore %base[%c0], %mask, %value
+ : memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @compress_16_at_8
+//
+// Same as @compress_16, but the insertion index is hard-coded to 8 instead of 0
+//===----------------------------------------------------------------------===//
+func.func @compress_16_at_8(%base: memref<?xi32>,
+ %mask: vector<[16]xi1>, %value: vector<[16]xi32>) {
+ %c8 = arith.constant 8: index
+ vector.compressstore %base[%c8], %mask, %value
+ : memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @print1DMemRef
+//
+// TODO: Move to an utility file
+//===----------------------------------------------------------------------===//
+func.func @print1DMemRef(%ptr: memref<?xi32>) -> () {
+ %cast = memref.cast %ptr: memref<?xi32> to memref<*xi32>
+
+ call @printMemrefI32(%cast): (memref<*xi32>) -> ()
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @reset_mem_i32
+//
+// Resets the input memory to 0.
+
+// TODO: Create a run-time utility funcion.
+//===----------------------------------------------------------------------===//
+func.func @reset_mem_i32(%ptr: memref<?xi32>, %size: index) {
+ %c0_idx = arith.constant 0: index
+ %c0 = arith.constant 0: i32
+ %step = arith.constant 1: index
+
+ scf.for %i = %c0_idx to %size step %step {
+ memref.store %c0, %ptr[%i] : memref<?xi32>
+ }
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @main
+//
+// The main entry point - sets the value of vscale.
+//===----------------------------------------------------------------------===//
+func.func @main() {
+ // Set vscale to 2 (vector width = 256). This will have identical effect to:
+ // * qemu-aarch64 -cpu max,sve-max-vq=2 (...)
+ %c256 = arith.constant 256 : i32
+ func.call @setArmVLBits(%c256) : (i32) -> ()
+
+ // Run the tests.
+ func.call @test() : () -> ()
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @test
+//
+// Set-up and run tests.
+//===----------------------------------------------------------------------===//
+func.func @test() {
+ //
+ // Shared constants.
+ //
+ %vs = vector.vscale
+
+ //
+ // Set up memory.
+ //
+ %c16 = arith.constant 16: index
+ %vs_16 = arith.muli %vs, %c16 : index
+ %A = memref.alloc(%vs_16) : memref<?xi32>
+ call @reset_mem_i32(%A, %vs_16) : (memref<?xi32>, index) -> ()
+
+ //
+ // Set the input vector.
+ //
+ %value = vector.step : vector<[16]xi32>
+ vector.print %value : vector<[16]xi32>
+
+ //
+ // Set up masks.
+ //
+
+ %f = arith.constant 0: i1
+ %t = arith.constant 1: i1
+
+ // %none = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ..., 0]
+ %none = vector.constant_mask [0] : vector<[16]xi1>
+
+ // %all = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1]
+ %all = vector.constant_mask [16] : vector<[16]xi1>
+ vector.print %all : vector<[16]xi1>
+
+ // %first_vscale_4 = [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, ..., 0]
+ %c4 = arith.constant 4 : index
+ %vs_4 = arith.muli %vs, %c4 : index
+ %first_vscale_4 = vector.create_mask %vs_4 : vector<[16]xi1>
+ vector.print %first_vscale_4 : vector<[16]xi1>
+
+ // %odd = [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, ..., 0]
+ %0 = vector.insert %t, %none[1] : i1 into vector<[16]xi1>
+ %1 = vector.insert %t, %0[3] : i1 into vector<[16]xi1>
+ %2 = vector.insert %t, %1[5] : i1 into vector<[16]xi1>
+ %3 = vector.insert %t, %2[7] : i1 into vector<[16]xi1>
+ %4 = vector.insert %t, %3[9] : i1 into vector<[16]xi1>
+ %5 = vector.insert %t, %4[11] : i1 into vector<[16]xi1>
+ %odd = vector.insert %t, %5[13] : i1 into vector<[16]xi1>
+
+ // %even = [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, ..., 0]
+ %6 = vector.insert %t, %none[0] : i1 into vector<[16]xi1>
+ %7 = vector.insert %t, %6[2] : i1 into vector<[16]xi1>
+ %8 = vector.insert %t, %7[4] : i1 into vector<[16]xi1>
+ %9 = vector.insert %t, %8[6] : i1 into vector<[16]xi1>
+ %10 = vector.insert %t, %9[8] : i1 into vector<[16]xi1>
+ %11 = vector.insert %t, %10[10] : i1 into vector<[16]xi1>
+ %even = vector.insert %t, %11[12] : i1 into vector<[16]xi1>
+
+
+ //
+ // Tests.
+ //
+
+ call @compress_16(%A, %none, %value)
+ : (memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>) -> ()
+ call @print1DMemRef(%A) : (memref<?xi32>) -> ()
+ // CHECK: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+ /// (...)
+ // CHECK-SAME: 0, 0]
+
+ call @compress_16(%A, %first_vscale_4, %value)
+ : (memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>) -> ()
+ call @print1DMemRef(%A) : (memref<?xi32>) -> ()
+ // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 0, 0, 0, 0, 0, 0, 0,
+ /// (...)
+ // CHECK-SAME: 0, 0]
+
+ call @compress_16(%A, %all, %value)
+ : (memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>) -> ()
+ call @print1DMemRef(%A) : (memref<?xi32>) -> ()
+ // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
+ /// (...)
+ // CHECK-SAME: 30, 31]
+
+ call @reset_mem_i32(%A, %vs_16) : (memref<?xi32>, index) -> ()
+ call @compress_16_at_8(%A, %first_vscale_4, %value)
+ : (memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>) -> ()
+ call @print1DMemRef(%A) : (memref<?xi32>) -> ()
+ // CHECK: [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 0, 0,
+ /// (...)
+ // CHECK-SAME: 0, 0]
+
+ call @reset_mem_i32(%A, %vs_16) : (memref<?xi32>, index) -> ()
+ call @compress_16(%A, %odd, %value)
+ : (memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>) -> ()
+ call @print1DMemRef(%A) : (memref<?xi32>) -> ()
+ // CHECK: [1, 3, 5, 7, 9, 11, 13, 0, 0, 0, 0,
+ /// (...)
+ // CHECK-SAME: 0, 0]
+
+ call @reset_mem_i32(%A, %vs_16) : (memref<?xi32>, index) -> ()
+ call @compress_16(%A, %even, %value)
+ : (memref<?xi32>, vector<[16]xi1>, vector<[16]xi32>) -> ()
+ call @print1DMemRef(%A) : (memref<?xi32>) -> ()
+ // CHECK: [0, 2, 4, 6, 8, 10, 12, 0, 0, 0, 0,
+ /// (...)
+ // CHECK-SAME: 0, 0]
+
+ memref.dealloc %A : memref<?xi32>
+ return
+}
+
+func.func private @printMemrefI32(%ptr : memref<*xi32>)
+func.func private @setArmVLBits(%bits : i32)
>From ced689efcae970392f8c8b1da697dc8ae351d352 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 21 Jul 2026 09:26:10 +0000
Subject: [PATCH 2/2] Address PR comments
---
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 4 ----
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 6 ++++++
mlir/test/Dialect/Vector/invalid.mlir | 15 +++++++++++++++
.../Dialect/Vector/CPU/ArmSVE/compress.mlir | 4 ++--
4 files changed, 23 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index ebfb379684dd3..e2f8f1a8817a5 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2328,8 +2328,6 @@ def Vector_ExpandLoadOp :
memory at an address aligned to this boundary. Violating this requirement
triggers immediate undefined behavior.
- Note, at the moment this Op is only available for fixed-width vectors.
-
Examples:
```mlir
@@ -2419,8 +2417,6 @@ def Vector_CompressStoreOp :
memory at an address aligned to this boundary. Violating this requirement
triggers immediate undefined behavior.
- Note, at the moment this Op is only available for fixed-width vectors.
-
Examples:
```mlir
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 94f627d12752e..e2b77c51e9217 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6653,6 +6653,9 @@ LogicalResult ExpandLoadOp::verify() {
return emitOpError("requires ") << memType.getRank() << " indices";
if (resVType.getShape() != maskVType.getShape())
return emitOpError("expected result shape to match mask shape");
+ if (resVType.getScalableDims() != maskVType.getScalableDims())
+ return emitOpError(
+ "expected result scalable dims to match mask scalable dims");
if (resVType != passVType)
return emitOpError("expected pass_thru of same type as result type");
return success();
@@ -6707,6 +6710,9 @@ LogicalResult CompressStoreOp::verify() {
return emitOpError("requires ") << memType.getRank() << " indices";
if (valueVType.getShape() != maskVType.getShape())
return emitOpError("expected valueToStore shape to match mask shape");
+ if (valueVType.getScalableDims() != maskVType.getScalableDims())
+ return emitOpError(
+ "expected valueToStore scalable dims to match mask scalable dims");
return success();
}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index d4f305f20b595..0097db39f6ed9 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1734,6 +1734,13 @@ func.func @expand_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector<1
// -----
+func.func @expand_scalable_dims_mismatch(%base: memref<?xf32>, %mask: vector<16xi1>, %pass_thru: vector<[16]xf32>, %c0: index) {
+ // expected-error at +1 {{expected result scalable dims to match mask scalable dims}}
+ %0 = vector.expandload %base[%c0], %mask, %pass_thru : memref<?xf32>, vector<16xi1>, vector<[16]xf32> into vector<[16]xf32>
+}
+
+// -----
+
func.func @compress_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>, %value: vector<16xf32>) {
%c0 = arith.constant 0 : index
// expected-error at +1 {{'vector.compressstore' op base element type ('f64') does not match valueToStore element type ('f32')}}
@@ -1780,6 +1787,14 @@ func.func @compress_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector
// -----
+func.func @compress_scalable_dims_mismatch(%base: memref<?xf32>, %mask: vector<16xi1>, %value: vector<[16]xf32>) {
+ %c0 = arith.constant 0 : index
+ // expected-error at +1 {{expected valueToStore scalable dims to match mask scalable dims}}
+ vector.compressstore %base[%c0], %mask, %value : memref<?xf32>, vector<16xi1>, vector<[16]xf32>
+}
+
+// -----
+
func.func @scan_reduction_dim_constraint(%arg0: vector<2x3xi32>, %arg1: vector<3xi32>) -> vector<3xi32> {
// expected-error at +1 {{'vector.scan' op reduction dimension 5 has to be less than 2}}
%0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = true, reduction_dim = 5} :
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
index d8752cfb7fc92..eddc7c664d3f1 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
@@ -3,7 +3,7 @@
/// End-to-end test for vector.compressstore for SVE
// In order to demonstrate the impact of using scalable vectors, vscale is set
-// to 2 so that vector<[16]xi32> constains 32 rather than 16 elements at
+// to 2 so that vector<[16]xi32> contains 32 rather than 16 elements at
// run-time
//
// Note that you can also tweak the size of vscale by passing this flag to
@@ -11,7 +11,7 @@
// * -cpu max,sve-max-vq=[1-16]
// (select the value between 1 and 16).
-// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm
+// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm -o %t
// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+sve"\
// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils,%native_mlir_arm_runner_utils
More information about the Mlir-commits
mailing list