[Mlir-commits] [mlir] [mlir][vector] Make CompressstoreOp` + ExpandloadOp support scalable vectors (PR #210288)
Andrzej Warzyński
llvmlistbot at llvm.org
Fri Jul 17 02:57:16 PDT 2026
https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/210288
Extends `vector.compressstore` + `vector.expandload` to support scalable
vectors and updates relevant tests.
>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] [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)
More information about the Mlir-commits
mailing list