[mlir][vector] Make CompressstoreOp + ExpandloadOp support scalable vectors (#210288)
Extends `vector.compressstore` + `vector.expandload` to support scalable vectors and updates relevant tests. An e2e test for `vector.compressstore` is added. For `vector.expandload`, we need to wait for QEMU support: https://github.com/llvm/llvm-project/issues/210942.
This commit is contained in:
Andrzej Warzyński
GitHub
parent
bb22aa8127
commit
8b690a0854
8 changed files with 371 additions and 33 deletions
|
|
@ -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)> {
|
||||
|
|
@ -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
|
||||
|
|
@ -2381,7 +2379,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)> {
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -6655,6 +6655,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();
|
||||
|
|
@ -6709,6 +6712,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();
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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>) -> ()
|
||||
|
||||
// -----
|
||||
|
|
|
|||
|
|
@ -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@+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@+1 {{'vector.expandload' op expected result shape to match mask shape}}
|
||||
|
|
@ -1742,18 +1734,17 @@ func.func @expand_non_power_of_2_alignment(%base: memref<?xf32>, %mask: vector<1
|
|||
|
||||
// -----
|
||||
|
||||
func.func @compress_base_type_mismatch(%base: memref<?xf64>, %mask: vector<16xi1>, %value: vector<16xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
// expected-error@+1 {{'vector.compressstore' op base element type ('f64') does not match valueToStore element type ('f32')}}
|
||||
vector.compressstore %base[%c0], %mask, %value : memref<?xf64>, vector<16xi1>, vector<16xf32>
|
||||
func.func @expand_scalable_dims_mismatch(%base: memref<?xf32>, %mask: vector<16xi1>, %pass_thru: vector<[16]xf32>, %c0: index) {
|
||||
// expected-error@+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_scalable(%base: memref<?xf32>, %mask: vector<[16]xi1>, %value: 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@+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>
|
||||
// expected-error@+1 {{'vector.compressstore' op base element type ('f64') does not match valueToStore element type ('f32')}}
|
||||
vector.compressstore %base[%c0], %mask, %value : memref<?xf64>, vector<16xi1>, vector<16xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
|
@ -1796,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@+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@+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} :
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -732,6 +732,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>
|
||||
|
|
@ -778,6 +791,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>
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
214
mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
Normal file
214
mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/compress.mlir
Normal file
|
|
@ -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> contains 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 -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
|
||||
|
||||
// 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)
|
||||
Loading…
Reference in a new issue