diff options
author | Aart Bik <ajcbik@google.com> | 2024-04-10 09:42:12 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-04-10 09:42:12 -0700 |
commit | f388a3a446ef2566d73b6a73ba300738f8c2c002 (patch) | |
tree | 0a3e2fa54ea65c07a8a547618729cf32b07cebfe | |
parent | 3d468566eb395995ac54fcf90d3afb9b9f822eb3 (diff) | |
download | llvm-f388a3a446ef2566d73b6a73ba300738f8c2c002.zip llvm-f388a3a446ef2566d73b6a73ba300738f8c2c002.tar.gz llvm-f388a3a446ef2566d73b6a73ba300738f8c2c002.tar.bz2 |
[mlir][sparse] update doc and examples of the [dis]assemble operations (#88213)
The doc and examples of the [dis]assemble operations did not reflect all
the recent changes on order of the operands. Also clarified some of the
text.
4 files changed, 61 insertions, 58 deletions
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td index 5df8a17..0cfc64f 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td @@ -58,37 +58,35 @@ def SparseTensor_AssembleOp : SparseTensor_Op<"assemble", [Pure]>, Arguments<(ins Variadic<TensorOf<[AnySignlessIntegerOrIndex]>>:$levels, TensorOf<[AnyType]>:$values)>, Results<(outs AnySparseTensor: $result)> { - let summary = "Returns a sparse tensor assembled from the given values and levels"; + let summary = "Returns a sparse tensor assembled from the given levels and values"; let description = [{ - Assembles the values and per-level coordinate or postion arrays into a sparse tensor. - The order and types of provided levels must be consistent with the actual storage - layout of the returned sparse tensor described below. + Assembles the per-level position and coordinate arrays together with + the values arrays into a sparse tensor. The order and types of the + provided levels must be consistent with the actual storage layout of + the returned sparse tensor described below. - - `values : tensor<? x V>` - supplies the value for each stored element in the sparse tensor. - `levels: [tensor<? x iType>, ...]` - each supplies the sparse tensor coordinates scheme in the sparse tensor for - the corresponding level as specifed by `sparse_tensor::StorageLayout`. - - This operation can be used to assemble a sparse tensor from external - sources; e.g., when passing two numpy arrays from Python. - - Disclaimer: This is the user's responsibility to provide input that can be - correctly interpreted by the sparsifier, which does not perform - any sanity test during runtime to verify data integrity. + supplies the sparse tensor position and coordinate arrays + of the sparse tensor for the corresponding level as specifed by + `sparse_tensor::StorageLayout`. + - `values : tensor<? x V>` + supplies the values array for the stored elements in the sparse tensor. - TODO: The returned tensor is allowed (in principle) to have non-identity - dimOrdering/higherOrdering mappings. However, the current implementation - does not yet support them. + This operation can be used to assemble a sparse tensor from an + external source; e.g., by passing numpy arrays from Python. It + is the user's responsibility to provide input that can be correctly + interpreted by the sparsifier, which does not perform any sanity + test to verify data integrity. Example: ```mlir - %values = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64> - %coordinates = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex> - %st = sparse_tensor.assemble %values, %coordinates - : tensor<3xf64>, tensor<3x2xindex> to tensor<3x4xf64, #COO> + %pos = arith.constant dense<[0, 3]> : tensor<2xindex> + %index = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex> + %values = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64> + %s = sparse_tensor.assemble (%pos, %index), %values + : (tensor<2xindex>, tensor<3x2xindex>), tensor<3xf64> to tensor<3x4xf64, #COO> // yields COO format |1.1, 0.0, 0.0, 0.0| // of 3x4 matrix |0.0, 0.0, 2.2, 3.3| // |0.0, 0.0, 0.0, 0.0| @@ -96,8 +94,8 @@ def SparseTensor_AssembleOp : SparseTensor_Op<"assemble", [Pure]>, }]; let assemblyFormat = - "` ` `(` $levels `)` `,` $values attr-dict" - " `:` `(` type($levels) `)` `,` type($values) `to` type($result)"; + "` ` `(` $levels `)` `,` $values attr-dict `:`" + " `(` type($levels) `)` `,` type($values) `to` type($result)"; let hasVerifier = 1; } @@ -110,21 +108,20 @@ def SparseTensor_DisassembleOp : SparseTensor_Op<"disassemble", [Pure, SameVaria TensorOf<[AnyType]>:$ret_values, Variadic<AnyIndexingScalarLike>:$lvl_lens, AnyIndexingScalarLike:$val_len)> { - let summary = "Returns the (values, coordinates) pair disassembled from the input tensor"; + let summary = "Copies the levels and values of the given sparse tensor"; let description = [{ The disassemble operation is the inverse of `sparse_tensor::assemble`. - It returns the values and per-level position and coordinate array to the - user from the sparse tensor along with the actual length of the memory used - in each returned buffer. This operation can be used for returning an - disassembled MLIR sparse tensor to frontend; e.g., returning two numpy arrays - to Python. - - Disclaimer: This is the user's responsibility to allocate large enough buffers - to hold the sparse tensor. The sparsifier simply copies each fields - of the sparse tensor into the user-supplied buffer without bound checking. + It copies the per-level position and coordinate arrays together with + the values array of the given sparse tensor into the user-supplied buffers + along with the actual length of the memory used in each returned buffer. - TODO: the current implementation does not yet support non-identity mappings. + This operation can be used for returning a disassembled MLIR sparse tensor; + e.g., copying the sparse tensor contents into pre-allocated numpy arrays + back to Python. It is the user's responsibility to allocate large enough + buffers of the appropriate types to hold the sparse tensor contents. + The sparsifier simply copies all fields of the sparse tensor into the + user-supplied buffers without any sanity test to verify data integrity. Example: @@ -132,26 +129,26 @@ def SparseTensor_DisassembleOp : SparseTensor_Op<"disassemble", [Pure, SameVaria // input COO format |1.1, 0.0, 0.0, 0.0| // of 3x4 matrix |0.0, 0.0, 2.2, 3.3| // |0.0, 0.0, 0.0, 0.0| - %v, %p, %c, %v_len, %p_len, %c_len = - sparse_tensor.disassemble %sp : tensor<3x4xf64, #COO> - out_lvls(%op, %oi) : tensor<2xindex>, tensor<3x2xindex>, - out_vals(%od) : tensor<3xf64> -> - tensor<3xf64>, (tensor<2xindex>, tensor<3x2xindex>), index, (index, index) - // %v = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64> + %p, %c, %v, %p_len, %c_len, %v_len = + sparse_tensor.disassemble %s : tensor<3x4xf64, #COO> + out_lvls(%op, %oi : tensor<2xindex>, tensor<3x2xindex>) + out_vals(%od : tensor<3xf64>) -> + (tensor<2xindex>, tensor<3x2xindex>), tensor<3xf64>, (index, index), index // %p = arith.constant dense<[ 0, 3 ]> : tensor<2xindex> // %c = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex> - // %v_len = 3 + // %v = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64> // %p_len = 2 // %c_len = 6 (3x2) + // %v_len = 3 ``` }]; let assemblyFormat = - "$tensor `:` type($tensor) " + "$tensor attr-dict `:` type($tensor)" "`out_lvls` `(` $out_levels `:` type($out_levels) `)` " - "`out_vals` `(` $out_values `:` type($out_values) `)` attr-dict" - "`->` `(` type($ret_levels) `)` `,` type($ret_values) `,` " - "`(` type($lvl_lens) `)` `,` type($val_len)"; + "`out_vals` `(` $out_values `:` type($out_values) `)` `->`" + "`(` type($ret_levels) `)` `,` type($ret_values) `,` " + "`(` type($lvl_lens) `)` `,` type($val_len)"; let hasVerifier = 1; } diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir index 18851f2..7f5c051 100644 --- a/mlir/test/Dialect/SparseTensor/invalid.mlir +++ b/mlir/test/Dialect/SparseTensor/invalid.mlir @@ -60,7 +60,7 @@ func.func @invalid_pack_mis_position(%values: tensor<6xf64>, %coordinates: tenso func.func @invalid_unpack_type(%sp: tensor<100xf32, #SparseVector>, %values: tensor<6xf64>, %pos: tensor<2xi32>, %coordinates: tensor<6x1xi32>) { // expected-error@+1 {{input/output element-types don't match}} - %rv, %rp, %rc, %vl, %pl, %cl = sparse_tensor.disassemble %sp : tensor<100xf32, #SparseVector> + %rp, %rc, %rv, %pl, %cl, %vl = sparse_tensor.disassemble %sp : tensor<100xf32, #SparseVector> out_lvls(%pos, %coordinates : tensor<2xi32>, tensor<6x1xi32>) out_vals(%values : tensor<6xf64>) -> (tensor<2xi32>, tensor<6x1xi32>), tensor<6xf64>, (index, index), index @@ -73,7 +73,7 @@ func.func @invalid_unpack_type(%sp: tensor<100xf32, #SparseVector>, %values: ten func.func @invalid_unpack_type(%sp: tensor<100x2xf64, #SparseVector>, %values: tensor<6xf64>, %pos: tensor<2xi32>, %coordinates: tensor<6x3xi32>) { // expected-error@+1 {{input/output trailing COO level-ranks don't match}} - %rv, %rp, %rc, %vl, %pl, %cl = sparse_tensor.disassemble %sp : tensor<100x2xf64, #SparseVector> + %rp, %rc, %rv, %pl, %cl, %vl = sparse_tensor.disassemble %sp : tensor<100x2xf64, #SparseVector> out_lvls(%pos, %coordinates : tensor<2xi32>, tensor<6x3xi32> ) out_vals(%values : tensor<6xf64>) -> (tensor<2xi32>, tensor<6x3xi32>), tensor<6xf64>, (index, index), index @@ -86,7 +86,7 @@ func.func @invalid_unpack_type(%sp: tensor<100x2xf64, #SparseVector>, %values: t func.func @invalid_unpack_mis_position(%sp: tensor<2x100xf64, #CSR>, %values: tensor<6xf64>, %coordinates: tensor<6xi32>) { // expected-error@+1 {{inconsistent number of fields between input/output}} - %rv, %rc, %vl, %pl = sparse_tensor.disassemble %sp : tensor<2x100xf64, #CSR> + %rc, %rv, %cl, %vl = sparse_tensor.disassemble %sp : tensor<2x100xf64, #CSR> out_lvls(%coordinates : tensor<6xi32>) out_vals(%values : tensor<6xf64>) -> (tensor<6xi32>), tensor<6xf64>, (index), index diff --git a/mlir/test/Dialect/SparseTensor/roundtrip.mlir b/mlir/test/Dialect/SparseTensor/roundtrip.mlir index a47a3d5..12f69c1 100644 --- a/mlir/test/Dialect/SparseTensor/roundtrip.mlir +++ b/mlir/test/Dialect/SparseTensor/roundtrip.mlir @@ -33,21 +33,21 @@ func.func @sparse_pack(%pos: tensor<2xi32>, %index: tensor<6x1xi32>, %data: tens #SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed), crdWidth=32}> // CHECK-LABEL: func @sparse_unpack( // CHECK-SAME: %[[T:.*]]: tensor<100xf64, # -// CHECK-SAME: %[[OD:.*]]: tensor<6xf64> -// CHECK-SAME: %[[OP:.*]]: tensor<2xindex> -// CHECK-SAME: %[[OI:.*]]: tensor<6x1xi32> +// CHECK-SAME: %[[OP:.*]]: tensor<2xindex>, +// CHECK-SAME: %[[OI:.*]]: tensor<6x1xi32>, +// CHECK-SAME: %[[OD:.*]]: tensor<6xf64>) // CHECK: %[[P:.*]]:2, %[[D:.*]], %[[PL:.*]]:2, %[[DL:.*]] = sparse_tensor.disassemble %[[T]] // CHECK: return %[[P]]#0, %[[P]]#1, %[[D]] func.func @sparse_unpack(%sp : tensor<100xf64, #SparseVector>, - %od : tensor<6xf64>, %op : tensor<2xindex>, - %oi : tensor<6x1xi32>) + %oi : tensor<6x1xi32>, + %od : tensor<6xf64>) -> (tensor<2xindex>, tensor<6x1xi32>, tensor<6xf64>) { - %rp, %ri, %rd, %vl, %pl, %cl = sparse_tensor.disassemble %sp : tensor<100xf64, #SparseVector> + %rp, %ri, %d, %rpl, %ril, %dl = sparse_tensor.disassemble %sp : tensor<100xf64, #SparseVector> out_lvls(%op, %oi : tensor<2xindex>, tensor<6x1xi32>) out_vals(%od : tensor<6xf64>) -> (tensor<2xindex>, tensor<6x1xi32>), tensor<6xf64>, (index, index), index - return %rp, %ri, %rd : tensor<2xindex>, tensor<6x1xi32>, tensor<6xf64> + return %rp, %ri, %d : tensor<2xindex>, tensor<6x1xi32>, tensor<6xf64> } // ----- diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir index 7ecccad..5415625 100644 --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir @@ -231,7 +231,7 @@ module { %od = tensor.empty() : tensor<3xf64> %op = tensor.empty() : tensor<2xi32> %oi = tensor.empty() : tensor<3x2xi32> - %p, %i, %d, %dl, %pl, %il = sparse_tensor.disassemble %s5 : tensor<10x10xf64, #SortedCOOI32> + %p, %i, %d, %pl, %il, %dl = sparse_tensor.disassemble %s5 : tensor<10x10xf64, #SortedCOOI32> out_lvls(%op, %oi : tensor<2xi32>, tensor<3x2xi32>) out_vals(%od : tensor<3xf64>) -> (tensor<2xi32>, tensor<3x2xi32>), tensor<3xf64>, (i32, i64), index @@ -244,10 +244,13 @@ module { %vi = vector.transfer_read %i[%c0, %c0], %i0 : tensor<3x2xi32>, vector<3x2xi32> vector.print %vi : vector<3x2xi32> + // CHECK-NEXT: 3 + vector.print %dl : index + %d_csr = tensor.empty() : tensor<4xf64> %p_csr = tensor.empty() : tensor<3xi32> %i_csr = tensor.empty() : tensor<3xi32> - %rp_csr, %ri_csr, %rd_csr, %ld_csr, %lp_csr, %li_csr = sparse_tensor.disassemble %csr : tensor<2x2xf64, #CSR> + %rp_csr, %ri_csr, %rd_csr, %lp_csr, %li_csr, %ld_csr = sparse_tensor.disassemble %csr : tensor<2x2xf64, #CSR> out_lvls(%p_csr, %i_csr : tensor<3xi32>, tensor<3xi32>) out_vals(%d_csr : tensor<4xf64>) -> (tensor<3xi32>, tensor<3xi32>), tensor<4xf64>, (i32, i64), index @@ -256,6 +259,9 @@ module { %vd_csr = vector.transfer_read %rd_csr[%c0], %f0 : tensor<4xf64>, vector<3xf64> vector.print %vd_csr : vector<3xf64> + // CHECK-NEXT: 3 + vector.print %ld_csr : index + %bod = tensor.empty() : tensor<6xf64> %bop = tensor.empty() : tensor<4xindex> %boi = tensor.empty() : tensor<6x2xindex> |