@@ -58,46 +58,44 @@ def SparseTensor_AssembleOp : SparseTensor_Op<"assemble", [Pure]>,
5858 Arguments<(ins Variadic<TensorOf<[AnySignlessIntegerOrIndex]>>:$levels,
5959 TensorOf<[AnyType]>:$values)>,
6060 Results<(outs AnySparseTensor: $result)> {
61- let summary = "Returns a sparse tensor assembled from the given values and levels ";
61+ let summary = "Returns a sparse tensor assembled from the given levels and values ";
6262
6363 let description = [{
64- Assembles the values and per-level coordinate or postion arrays into a sparse tensor.
65- The order and types of provided levels must be consistent with the actual storage
66- layout of the returned sparse tensor described below.
64+ Assembles the per-level position and coordinate arrays together with
65+ the values arrays into a sparse tensor. The order and types of the
66+ provided levels must be consistent with the actual storage layout of
67+ the returned sparse tensor described below.
6768
68- - `values : tensor<? x V>`
69- supplies the value for each stored element in the sparse tensor.
7069 - `levels: [tensor<? x iType>, ...]`
71- each supplies the sparse tensor coordinates scheme in the sparse tensor for
72- the corresponding level as specifed by `sparse_tensor::StorageLayout`.
73-
74- This operation can be used to assemble a sparse tensor from external
75- sources; e.g., when passing two numpy arrays from Python.
76-
77- Disclaimer: This is the user's responsibility to provide input that can be
78- correctly interpreted by the sparsifier, which does not perform
79- any sanity test during runtime to verify data integrity.
70+ supplies the sparse tensor position and coordinate arrays
71+ of the sparse tensor for the corresponding level as specifed by
72+ `sparse_tensor::StorageLayout`.
73+ - `values : tensor<? x V>`
74+ supplies the values array for the stored elements in the sparse tensor.
8075
81- TODO: The returned tensor is allowed (in principle) to have non-identity
82- dimOrdering/higherOrdering mappings. However, the current implementation
83- does not yet support them.
76+ This operation can be used to assemble a sparse tensor from an
77+ external source; e.g., by passing numpy arrays from Python. It
78+ is the user's responsibility to provide input that can be correctly
79+ interpreted by the sparsifier, which does not perform any sanity
80+ test to verify data integrity.
8481
8582 Example:
8683
8784 ```mlir
88- %values = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64>
89- %coordinates = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
90- %st = sparse_tensor.assemble %values, %coordinates
91- : tensor<3xf64>, tensor<3x2xindex> to tensor<3x4xf64, #COO>
85+ %pos = arith.constant dense<[0, 3]> : tensor<2xindex>
86+ %index = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
87+ %values = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64>
88+ %s = sparse_tensor.assemble (%pos, %index), %values
89+ : (tensor<2xindex>, tensor<3x2xindex>), tensor<3xf64> to tensor<3x4xf64, #COO>
9290 // yields COO format |1.1, 0.0, 0.0, 0.0|
9391 // of 3x4 matrix |0.0, 0.0, 2.2, 3.3|
9492 // |0.0, 0.0, 0.0, 0.0|
9593 ```
9694 }];
9795
9896 let assemblyFormat =
99- "` ` `(` $levels `)` `,` $values attr-dict"
100- " `:` `(` type($levels) `)` `,` type($values) `to` type($result)";
97+ "` ` `(` $levels `)` `,` $values attr-dict `:` "
98+ " `(` type($levels) `)` `,` type($values) `to` type($result)";
10199
102100 let hasVerifier = 1;
103101}
@@ -110,48 +108,47 @@ def SparseTensor_DisassembleOp : SparseTensor_Op<"disassemble", [Pure, SameVaria
110108 TensorOf<[AnyType]>:$ret_values,
111109 Variadic<AnyIndexingScalarLike>:$lvl_lens,
112110 AnyIndexingScalarLike:$val_len)> {
113- let summary = "Returns the (values, coordinates) pair disassembled from the input tensor";
111+ let summary = "Copies the levels and values of the given sparse tensor";
114112
115113 let description = [{
116114 The disassemble operation is the inverse of `sparse_tensor::assemble`.
117- It returns the values and per-level position and coordinate array to the
118- user from the sparse tensor along with the actual length of the memory used
119- in each returned buffer. This operation can be used for returning an
120- disassembled MLIR sparse tensor to frontend; e.g., returning two numpy arrays
121- to Python.
122-
123- Disclaimer: This is the user's responsibility to allocate large enough buffers
124- to hold the sparse tensor. The sparsifier simply copies each fields
125- of the sparse tensor into the user-supplied buffer without bound checking.
115+ It copies the per-level position and coordinate arrays together with
116+ the values array of the given sparse tensor into the user-supplied buffers
117+ along with the actual length of the memory used in each returned buffer.
126118
127- TODO: the current implementation does not yet support non-identity mappings.
119+ This operation can be used for returning a disassembled MLIR sparse tensor;
120+ e.g., copying the sparse tensor contents into pre-allocated numpy arrays
121+ back to Python. It is the user's responsibility to allocate large enough
122+ buffers of the appropriate types to hold the sparse tensor contents.
123+ The sparsifier simply copies all fields of the sparse tensor into the
124+ user-supplied buffers without any sanity test to verify data integrity.
128125
129126 Example:
130127
131128 ```mlir
132129 // input COO format |1.1, 0.0, 0.0, 0.0|
133130 // of 3x4 matrix |0.0, 0.0, 2.2, 3.3|
134131 // |0.0, 0.0, 0.0, 0.0|
135- %v, %p, %c, %v_len, %p_len, %c_len =
136- sparse_tensor.disassemble %sp : tensor<3x4xf64, #COO>
137- out_lvls(%op, %oi) : tensor<2xindex>, tensor<3x2xindex>,
138- out_vals(%od) : tensor<3xf64> ->
139- tensor<3xf64>, (tensor<2xindex>, tensor<3x2xindex>), index, (index, index)
140- // %v = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64>
132+ %p, %c, %v, %p_len, %c_len, %v_len =
133+ sparse_tensor.disassemble %s : tensor<3x4xf64, #COO>
134+ out_lvls(%op, %oi : tensor<2xindex>, tensor<3x2xindex>)
135+ out_vals(%od : tensor<3xf64>) ->
136+ (tensor<2xindex>, tensor<3x2xindex>), tensor<3xf64>, (index, index), index
141137 // %p = arith.constant dense<[ 0, 3 ]> : tensor<2xindex>
142138 // %c = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
143- // %v_len = 3
139+ // %v = arith.constant dense<[ 1.1, 2.2, 3.3 ]> : tensor<3xf64>
144140 // %p_len = 2
145141 // %c_len = 6 (3x2)
142+ // %v_len = 3
146143 ```
147144 }];
148145
149146 let assemblyFormat =
150- "$tensor `:` type($tensor) "
147+ "$tensor attr-dict `:` type($tensor)"
151148 "`out_lvls` `(` $out_levels `:` type($out_levels) `)` "
152- "`out_vals` `(` $out_values `:` type($out_values) `)` attr-dict "
153- "`->` ` (` type($ret_levels) `)` `,` type($ret_values) `,` "
154- "`(` type($lvl_lens) `)` `,` type($val_len)";
149+ "`out_vals` `(` $out_values `:` type($out_values) `)` `->` "
150+ "`(` type($ret_levels) `)` `,` type($ret_values) `,` "
151+ "`(` type($lvl_lens) `)` `,` type($val_len)";
155152
156153 let hasVerifier = 1;
157154}
0 commit comments