diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td index cacc8176c6782..afd978c1c57eb 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td @@ -133,9 +133,11 @@ 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 +178,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 +227,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 +265,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 }