Skip to content

Commit 2a07f0f

Browse files
[mlir][sparse] Migrate more tests to use new syntax (#66443)
**Dense** `lvlTypes = [ "dense", "dense" ]` to `map = (d0, d1) -> (d0 : dense, d1 : dense)` `lvlTypes = [ "dense", "dense" ], dimToLvl = affine_map<(i,j) -> (j,i)>` to `map = (d0, d1) -> (d1 : dense, d0 : dense)` **DCSR** `lvlTypes = [ "compressed", "compressed" ]` to `map = (d0, d1) -> (d0 : compressed, d1 : compressed)` **DCSC** `lvlTypes = [ "compressed", "compressed" ], dimToLvl = affine_map<(i,j) -> (j,i)>` to `map = (d0, d1) -> (d1 : compressed, d0 : compressed)` **Block Row** `lvlTypes = [ "compressed", "dense" ]` to `map = (d0, d1) -> (d0 : compressed, d1 : dense)` **Block Column** `lvlTypes = [ "compressed", "dense" ], dimToLvl = affine_map<(i,j) -> (j,i)>` to `map = (d0, d1) -> (d1 : compressed, d0 : dense)` This is an ongoing effort: #66146, #66309
1 parent 4db6803 commit 2a07f0f

File tree

69 files changed

+128
-165
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

69 files changed

+128
-165
lines changed

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -206,8 +206,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
206206

207207
// Doubly compressed sparse column storage with specific bitwidths.
208208
#DCSC = #sparse_tensor.encoding<{
209-
lvlTypes = [ "compressed", "compressed" ],
210-
dimToLvl = affine_map<(i, j) -> (j, i)>,
209+
map = (d0, d1) -> (d1 : compressed, d0 : compressed),
211210
posWidth = 32,
212211
crdWidth = 8
213212
}>

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -351,8 +351,7 @@ static bool findDepIdxSet(Merger &merger, TensorId tensor, Level lvl,
351351
/// Get the total number of compound affine expressions in the
352352
/// `getMatchingIndexingMap` for the given tensor. For the following inputs:
353353
///
354-
/// map = (d0, d1, d2) => (d0 + d1, d2)
355-
/// lvlTypes = ["compressed", "compressed"]
354+
/// map = (d0, d1, d2) => (d0 + d1 : compressed, d2 : compressed)
356355
///
357356
/// Returns 1 (because the first level is compressed and its corresponding
358357
/// indexing-expression is `d0 + d1`)

mlir/test/Dialect/Bufferization/invalid.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ func.func @escape_attr_non_bufferizable(%m0: memref<?xf32>) {
5858

5959
// -----
6060

61-
#DCSR = #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed" ] }>
61+
#DCSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed, d1 : compressed) }>
6262

6363
func.func @sparse_alloc_direct_return() -> tensor<20x40xf32, #DCSR> {
6464
// expected-error @+1{{sparse tensor allocation should not escape function}}
@@ -68,7 +68,7 @@ func.func @sparse_alloc_direct_return() -> tensor<20x40xf32, #DCSR> {
6868

6969
// -----
7070

71-
#DCSR = #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed" ] }>
71+
#DCSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed, d1 : compressed) }>
7272

7373
func.func private @foo(tensor<20x40xf32, #DCSR>) -> ()
7474

mlir/test/Dialect/SparseTensor/codegen.mlir

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,13 @@
99
}>
1010

1111
#Dense2D = #sparse_tensor.encoding<{
12-
lvlTypes = [ "dense", "dense" ],
12+
map = (d0, d1) -> (d0 : dense, d1 : dense),
1313
crdWidth = 64,
1414
posWidth = 32
1515
}>
1616

1717
#Row = #sparse_tensor.encoding<{
18-
lvlTypes = [ "compressed", "dense" ],
18+
map = (d0, d1) -> (d0 : compressed, d1 : dense),
1919
crdWidth = 64,
2020
posWidth = 32
2121
}>
@@ -35,7 +35,7 @@
3535
}>
3636

3737
#DCSR = #sparse_tensor.encoding<{
38-
lvlTypes = [ "compressed", "compressed" ],
38+
map = (d0, d1) -> (d0 : compressed, d1 : compressed),
3939
crdWidth = 64,
4040
posWidth = 32
4141
}>

mlir/test/Dialect/SparseTensor/dense.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
// latter class is linearized into one-dimensional buffers that are backed
88
// by the runtime support library.
99

10-
#DenseMatrix = #sparse_tensor.encoding<{ lvlTypes = [ "dense", "dense" ] }>
10+
#DenseMatrix = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : dense) }>
1111

1212
#trait_2d = {
1313
indexing_maps = [

mlir/test/Dialect/SparseTensor/invalid.mlir

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,7 @@ func.func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
371371

372372
// -----
373373

374-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
374+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
375375

376376
func.func @sparse_convert_rank_mismatch(%arg0: tensor<10x10xf64, #DCSR>) -> tensor<?xf64> {
377377
// expected-error@+1 {{unexpected conversion mismatch in rank}}
@@ -714,7 +714,7 @@ func.func @invalid_concat_size_mismatch(%arg0: tensor<2x4xf64, #DC>,
714714

715715
// -----
716716

717-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
717+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
718718
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
719719
// expected-error@+1 {{Unmatched number of arguments in the block}}
720720
sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
@@ -725,7 +725,7 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
725725

726726
// -----
727727

728-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
728+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
729729
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
730730
// expected-error@+1 {{Expecting Index type for argument at index 1}}
731731
sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
@@ -736,7 +736,7 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
736736

737737
// -----
738738

739-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
739+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
740740
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
741741
// expected-error@+1 {{Unmatched element type between input tensor and block argument}}
742742
sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
@@ -747,7 +747,7 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
747747

748748
// -----
749749

750-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
750+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
751751
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
752752
// expected-error@+1 {{Unmatched element type between input tensor and block argument}}
753753
sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
@@ -758,7 +758,7 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
758758

759759
// -----
760760

761-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
761+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
762762
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
763763
// expected-error@+1 {{Mismatch in number of init arguments and results}}
764764
sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 do {
@@ -769,7 +769,7 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> (
769769

770770
// -----
771771

772-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
772+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
773773
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
774774
// expected-error@+1 {{Mismatch in types of init arguments and results}}
775775
%1 = sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 -> i32 do {
@@ -780,7 +780,7 @@ func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> (
780780

781781
// -----
782782

783-
#DCSR = #sparse_tensor.encoding<{lvlTypes = ["compressed", "compressed"]}>
783+
#DCSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
784784
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
785785
// expected-error@+1 {{Mismatch in types of yield values and results}}
786786
%1 = sparse_tensor.foreach in %arg0 init(%arg1) : tensor<2x4xf64, #DCSR>, f32 -> f32 do {

mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22
// RUN: mlir-opt %s -test-tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" | FileCheck %s --check-prefix=CHECK-FUNC
33

44
#DCSR = #sparse_tensor.encoding<{
5-
lvlTypes = [ "compressed", "compressed" ],
6-
dimToLvl = affine_map<(i,j) -> (i,j)>
5+
map = (d0, d1) -> (d0 : compressed, d1 : compressed)
76
}>
87

98
// CHECK-LABEL: func @bufferization_alloc_tensor

mlir/test/Dialect/SparseTensor/one_trip.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// RUN: mlir-opt %s -sparsification -cse | FileCheck %s
22

33
#Dense = #sparse_tensor.encoding<{
4-
lvlTypes = [ "dense" , "dense" ]
4+
map = (d0, d1) -> (d0 : dense, d1 : dense)
55
}>
66

77
#trait_scale = {

mlir/test/Dialect/SparseTensor/post_rewriting.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
}>
66

77
#SparseMatrix = #sparse_tensor.encoding<{
8-
lvlTypes = ["compressed", "compressed"]
8+
map = (d0, d1) -> (d0 : compressed, d1 : compressed)
99
}>
1010

1111
// CHECK-LABEL: func.func @expand_dense(

mlir/test/Dialect/SparseTensor/pre_rewriting.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
}>
1010

1111
#DCSR = #sparse_tensor.encoding<{
12-
lvlTypes = ["compressed", "compressed"]
12+
map = (d0, d1) -> (d0 : compressed, d1 : compressed)
1313
}>
1414

1515
#Slice = #sparse_tensor.encoding<{

0 commit comments

Comments
 (0)