|
23 | 23 | using namespace mlir;
|
24 | 24 | using namespace mlir::sparse_tensor;
|
25 | 25 |
|
26 |
| -/// If the tensor is a sparse constant, generates and returns the pair of |
27 |
| -/// the constants for the coordinates and the values. |
28 |
| -static std::optional<std::pair<Value, Value>> |
29 |
| -genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor) { |
30 |
| - if (auto constOp = tensor.getDefiningOp<arith::ConstantOp>()) { |
31 |
| - if (auto a = dyn_cast<SparseElementsAttr>(constOp.getValue())) { |
32 |
| - auto coordinates = builder.create<arith::ConstantOp>(loc, a.getIndices()); |
33 |
| - auto values = builder.create<arith::ConstantOp>(loc, a.getValues()); |
34 |
| - return std::make_pair(coordinates, values); |
35 |
| - } |
36 |
| - } |
37 |
| - return {}; |
38 |
| -} |
39 |
| - |
40 |
| -/// Reads `coordinates[k][0..rank-1]` and `value[k]`, appending the |
41 |
| -/// former onto `cvs` and returning the latter. |
42 |
| -// FIXME: Change the `rank` argument to `Dimension dimRank` or `Level lvlRank`, |
43 |
| -// to clarify its intended meaning. |
44 |
| -static Value genCoordsAndValueForSparse(OpBuilder &builder, Location loc, |
45 |
| - Value coordinates, Value values, |
46 |
| - SmallVectorImpl<Value> &cvs, Value k, |
47 |
| - unsigned rank) { |
48 |
| - for (unsigned d = 0; d < rank; d++) { |
49 |
| - Value dim = constantIndex(builder, loc, d); |
50 |
| - Value crd = |
51 |
| - builder.create<tensor::ExtractOp>(loc, coordinates, ValueRange{k, dim}); |
52 |
| - crd = builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), crd); |
53 |
| - // builder.create<memref::StoreOp>(loc, crd, cvs, dim); |
54 |
| - cvs.push_back(crd); |
55 |
| - } |
56 |
| - return builder.create<tensor::ExtractOp>(loc, values, k); |
57 |
| -} |
58 |
| - |
59 |
| -/// Generates code to read the value from `tensor[ivs]`, and open |
60 |
| -/// a conditional for whether the value is non-zero. The generated code |
61 |
| -/// looks like the following and the insertion point after this routine |
62 |
| -/// is inside the then-branch. |
63 |
| -/// if (tensor[ivs] != 0) |
64 |
| -/// insert_point |
65 |
| -static Value genCoordsAndValueForDense(OpBuilder &builder, Location loc, |
66 |
| - Value tensor, |
67 |
| - SmallVectorImpl<Value> &cvs, |
68 |
| - ValueRange ivs) { |
69 |
| - Value val = genValueForDense(builder, loc, tensor, ivs); |
70 |
| - cvs.append(ivs.begin(), ivs.end()); |
71 |
| - return val; |
72 |
| -} |
73 |
| - |
74 | 26 | //===----------------------------------------------------------------------===//
|
75 | 27 | // ExecutionEngine/SparseTensorUtils helper functions.
|
76 | 28 | //===----------------------------------------------------------------------===//
|
@@ -450,65 +402,6 @@ void mlir::sparse_tensor::deallocDenseTensor(OpBuilder &builder, Location loc,
|
450 | 402 | builder.create<memref::DeallocOp>(loc, buffer);
|
451 | 403 | }
|
452 | 404 |
|
453 |
| -Value mlir::sparse_tensor::genValueForDense(OpBuilder &builder, Location loc, |
454 |
| - Value tensor, ValueRange ivs) { |
455 |
| - Value val = builder.create<tensor::ExtractOp>(loc, tensor, ivs); |
456 |
| - Value cond = genIsNonzero(builder, loc, val); |
457 |
| - scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else*/ false); |
458 |
| - builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); |
459 |
| - return val; |
460 |
| -} |
461 |
| - |
462 |
| -// FIXME: |
463 |
| -// 1. Dense tensors loop should be generated by loop emitter. |
464 |
| -// 2. Support reduction variables to propagate SSA chains properly. |
465 |
| -// 3. Change the `rank` argument to `Dimension dimRank` or `Level lvlRank`, |
466 |
| -// to clarify its meaning. |
467 |
| -void mlir::sparse_tensor::genDenseTensorOrSparseConstantIterLoop( |
468 |
| - OpBuilder &builder, Location loc, Value src, unsigned rank, |
469 |
| - function_ref<void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder) { |
470 |
| - // `cvs` is actually the flattened coordinates array for all elements, |
471 |
| - // not just for one element (since we do not `SmallVector::clear` after |
472 |
| - // each iteration of the body of the loopnest. |
473 |
| - SmallVector<Value> cvs; |
474 |
| - SmallVector<Value> lo; |
475 |
| - SmallVector<Value> hi; |
476 |
| - SmallVector<Value> st; |
477 |
| - const Value zero = constantIndex(builder, loc, 0); |
478 |
| - const Value one = constantIndex(builder, loc, 1); |
479 |
| - const auto splitSrc = genSplitSparseConstant(builder, loc, src); |
480 |
| - if (splitSrc.has_value()) { |
481 |
| - const Value srcCoordinates = splitSrc->first; |
482 |
| - const Value srcValues = splitSrc->second; |
483 |
| - lo.push_back(zero); |
484 |
| - hi.push_back(linalg::createOrFoldDimOp(builder, loc, srcValues, 0)); |
485 |
| - st.push_back(one); |
486 |
| - scf::buildLoopNest(builder, loc, lo, hi, st, {}, |
487 |
| - [&](OpBuilder &builder, Location loc, ValueRange ivs, |
488 |
| - ValueRange /*args*/) -> scf::ValueVector { |
489 |
| - Value val = genCoordsAndValueForSparse( |
490 |
| - builder, loc, srcCoordinates, srcValues, cvs, |
491 |
| - ivs[0], rank); |
492 |
| - bodyBuilder(builder, loc, val, cvs); |
493 |
| - return {}; |
494 |
| - }); |
495 |
| - } else { |
496 |
| - for (unsigned i = 0; i < rank; i++) { |
497 |
| - lo.push_back(zero); |
498 |
| - hi.push_back(linalg::createOrFoldDimOp(builder, loc, src, i)); |
499 |
| - st.push_back(one); |
500 |
| - } |
501 |
| - scf::buildLoopNest(builder, loc, lo, hi, st, {}, |
502 |
| - [&](OpBuilder &builder, Location loc, ValueRange ivs, |
503 |
| - ValueRange /*args*/) -> scf::ValueVector { |
504 |
| - Value val = genCoordsAndValueForDense(builder, loc, |
505 |
| - src, cvs, ivs); |
506 |
| - bodyBuilder(builder, loc, val, cvs); |
507 |
| - return {}; |
508 |
| - }); |
509 |
| - } |
510 |
| -} |
511 |
| - |
512 | 405 | void mlir::sparse_tensor::sizesFromSrc(OpBuilder &builder,
|
513 | 406 | SmallVectorImpl<Value> &sizes,
|
514 | 407 | Location loc, Value src) {
|
|
0 commit comments