From c6b17735dcbe66a45df5083edc68a3e3d3c5af4f Mon Sep 17 00:00:00 2001 From: Yinying Li Date: Mon, 9 Oct 2023 22:07:08 +0000 Subject: [PATCH 1/3] [mlir][sparse] Fix errors in doc and tests --- .../SparseTensor/IR/SparseTensorAttrDefs.td | 14 ++++----- .../SparseTensor/roundtrip_encoding.mlir | 30 +++++++++---------- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td index cacc8176c6782..0d7340f663ac4 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td @@ -133,9 +133,9 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding", level-expressions collectively define an affine map from dimension-coordinates to level-coordinates. The dimension-expressions collectively define the inverse map, which only needs to be provided for elaborate cases where it cannot be inferred - automatically. Within the sparse storage format, we refer to indices that are - stored explicitly as **coordinates** and offsets into the storage format as - **positions**. + automatically. Each dimension could also have an optional `SparseTensorDimSliceAttr`. + Within the sparse storage format, we refer to indices that are stored explicitly + as **coordinates** and offsets into the storage format as **positions**. The supported level-formats are the following: @@ -176,9 +176,6 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding", coordinate over all levels). The choices are `8`, `16`, `32`, `64`, or, the default, `0` to indicate a native bitwidth. - - An optional array of `SparseTensorDimSliceAttr`, which specifies - how the sparse tensor is partitioned on each dimension. - Examples: ```mlir @@ -228,7 +225,8 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding", // Same block sparse row storage (2x3 blocks) but this time // also with a redundant reverse mapping, which can be inferred. #BSR_explicit = #sparse_tensor.encoding<{ - map = ( i = ib * 2 + ii, + map = {ib, jb, ii, jj} + ( i = ib * 2 + ii, j = jb * 3 + jj) -> ( ib = i floordiv 2 : dense, jb = j floordiv 3 : compressed, @@ -265,7 +263,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding", j : #sparse_tensor) -> (i : dense, j : compressed) }> - ... tensor ... + ... tensor ... ``` }]; diff --git a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir index c4ef50bee01ea..ae3805d8b7741 100644 --- a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir +++ b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir @@ -84,18 +84,18 @@ func.func private @sparse_sorted_coo(tensor<10x10xf64, #SortedCOO>) // ----- -#BCSR = #sparse_tensor.encoding<{ +#BSR = #sparse_tensor.encoding<{ map = ( i, j ) -> - ( i floordiv 2 : compressed, + ( i floordiv 2 : dense, j floordiv 3 : compressed, i mod 2 : dense, j mod 3 : dense ) }> -// CHECK-LABEL: func private @sparse_bcsr( -// CHECK-SAME: tensor<10x60xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : compressed, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>> -func.func private @sparse_bcsr(tensor<10x60xf64, #BCSR>) +// CHECK-LABEL: func private @sparse_bsr( +// CHECK-SAME: tensor<10x60xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>> +func.func private @sparse_bsr(tensor<10x60xf64, #BSR>) // ----- @@ -143,39 +143,39 @@ func.func private @sparse_2_out_of_4(tensor) // ----- -#BCSR = #sparse_tensor.encoding<{ +#BSR = #sparse_tensor.encoding<{ map = ( i, j ) -> - ( i floordiv 2 : compressed, + ( i floordiv 2 : dense, j floordiv 3 : compressed, i mod 2 : dense, j mod 3 : dense ) }> -// CHECK-LABEL: func private @BCSR( -// CHECK-SAME: tensor (d0 floordiv 2 : compressed, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>> -func.func private @BCSR(%arg0: tensor) { +// CHECK-LABEL: func private @BSR( +// CHECK-SAME: tensor (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>> +func.func private @BSR(%arg0: tensor) { return } // ----- -#BCSR_explicit = #sparse_tensor.encoding<{ +#BSR_explicit = #sparse_tensor.encoding<{ map = {il, jl, ii, jj} ( i = il * 2 + ii, j = jl * 3 + jj ) -> - ( il = i floordiv 2 : compressed, + ( il = i floordiv 2 : dense, jl = j floordiv 3 : compressed, ii = i mod 2 : dense, jj = j mod 3 : dense ) }> -// CHECK-LABEL: func private @BCSR_explicit( -// CHECK-SAME: tensor (d0 floordiv 2 : compressed, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>> -func.func private @BCSR_explicit(%arg0: tensor) { +// CHECK-LABEL: func private @BSR_explicit( +// CHECK-SAME: tensor (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>> +func.func private @BSR_explicit(%arg0: tensor) { return } From 26a1f9ff44250a618f55e0ae12d033aa1c4d777c Mon Sep 17 00:00:00 2001 From: Yinying Li Date: Mon, 9 Oct 2023 22:13:07 +0000 Subject: [PATCH 2/3] format --- .../mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td index 0d7340f663ac4..23b565857c84d 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td @@ -225,7 +225,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding", // Same block sparse row storage (2x3 blocks) but this time // also with a redundant reverse mapping, which can be inferred. #BSR_explicit = #sparse_tensor.encoding<{ - map = {ib, jb, ii, jj} + map = { ib, jb, ii, jj } ( i = ib * 2 + ii, j = jb * 3 + jj) -> ( ib = i floordiv 2 : dense, From ad847d7d44cd20b5105f397c4330b8e99639d239 Mon Sep 17 00:00:00 2001 From: Yinying Li Date: Mon, 9 Oct 2023 23:49:20 +0000 Subject: [PATCH 3/3] address comment --- .../mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td index 23b565857c84d..afd978c1c57eb 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td @@ -133,7 +133,9 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding", level-expressions collectively define an affine map from dimension-coordinates to level-coordinates. The dimension-expressions collectively define the inverse map, which only needs to be provided for elaborate cases where it cannot be inferred - automatically. Each dimension could also have an optional `SparseTensorDimSliceAttr`. + automatically. + + Each dimension could also have an optional `SparseTensorDimSliceAttr`. Within the sparse storage format, we refer to indices that are stored explicitly as **coordinates** and offsets into the storage format as **positions**.