Skip to content

Commit 4daf86e

Browse files
authored
[mlir][sparse] refactoring sparse runtime lib into less paths (#85332)
Two constructors could be easily refactored into one after a lot of previous deprecated code has been removed.
1 parent e4b7724 commit 4daf86e

File tree

3 files changed

+34
-58
lines changed

3 files changed

+34
-58
lines changed

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ class SparseTensorReader final {
206206
auto *lvlCOO = readCOO<V>(map, lvlSizes);
207207
auto *tensor = SparseTensorStorage<P, I, V>::newFromCOO(
208208
dimRank, getDimSizes(), lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim,
209-
*lvlCOO);
209+
lvlCOO);
210210
delete lvlCOO;
211211
return tensor;
212212
}

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

Lines changed: 32 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -201,33 +201,18 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
201201

202202
public:
203203
/// Constructs a sparse tensor with the given encoding, and allocates
204-
/// overhead storage according to some simple heuristics. When the
205-
/// `bool` argument is true and `lvlTypes` are all dense, then this
206-
/// ctor will also initialize the values array with zeros. That
207-
/// argument should be true when an empty tensor is intended; whereas
208-
/// it should usually be false when the ctor will be followed up by
209-
/// some other form of initialization.
204+
/// overhead storage according to some simple heuristics. When lvlCOO
205+
/// is set, the sparse tensor initializes with the contents from that
206+
/// data structure. Otherwise, an empty sparse tensor results.
210207
SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
211208
uint64_t lvlRank, const uint64_t *lvlSizes,
212209
const LevelType *lvlTypes, const uint64_t *dim2lvl,
213-
const uint64_t *lvl2dim, SparseTensorCOO<V> *lvlCOO,
214-
bool initializeValuesIfAllDense);
210+
const uint64_t *lvl2dim, SparseTensorCOO<V> *lvlCOO);
215211

216212
/// Constructs a sparse tensor with the given encoding, and initializes
217-
/// the contents from the COO. This ctor performs the same heuristic
218-
/// overhead-storage allocation as the ctor above.
219-
SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
220-
uint64_t lvlRank, const uint64_t *lvlSizes,
221-
const LevelType *lvlTypes, const uint64_t *dim2lvl,
222-
const uint64_t *lvl2dim, SparseTensorCOO<V> &lvlCOO);
223-
224-
/// Constructs a sparse tensor with the given encoding, and initializes
225-
/// the contents from the level buffers. This ctor allocates exactly
226-
/// the required amount of overhead storage, not using any heuristics.
227-
/// It assumes that the data provided by `lvlBufs` can be directly used to
228-
/// interpret the result sparse tensor and performs *NO* integrity test on the
229-
/// input data. It also assume that the trailing COO coordinate buffer is
230-
/// passed in as a single AoS memory.
213+
/// the contents from the level buffers. The constructor assumes that the
214+
/// data provided by `lvlBufs` can be directly used to interpret the result
215+
/// sparse tensor and performs no integrity test on the input data.
231216
SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
232217
uint64_t lvlRank, const uint64_t *lvlSizes,
233218
const LevelType *lvlTypes, const uint64_t *dim2lvl,
@@ -244,16 +229,14 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
244229
newFromCOO(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
245230
const uint64_t *lvlSizes, const LevelType *lvlTypes,
246231
const uint64_t *dim2lvl, const uint64_t *lvl2dim,
247-
SparseTensorCOO<V> &lvlCOO);
232+
SparseTensorCOO<V> *lvlCOO);
248233

249-
/// Allocates a new sparse tensor and initialize it with the data stored level
250-
/// buffers directly.
234+
/// Allocates a new sparse tensor and initialize it from the given buffers.
251235
static SparseTensorStorage<P, C, V> *
252-
packFromLvlBuffers(uint64_t dimRank, const uint64_t *dimSizes,
253-
uint64_t lvlRank, const uint64_t *lvlSizes,
254-
const LevelType *lvlTypes, const uint64_t *dim2lvl,
255-
const uint64_t *lvl2dim, uint64_t srcRank,
256-
const intptr_t *buffers);
236+
newFromBuffers(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
237+
const uint64_t *lvlSizes, const LevelType *lvlTypes,
238+
const uint64_t *dim2lvl, const uint64_t *lvl2dim,
239+
uint64_t srcRank, const intptr_t *buffers);
257240

258241
~SparseTensorStorage() final = default;
259242

@@ -563,23 +546,24 @@ SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newEmpty(
563546
uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
564547
const uint64_t *lvlSizes, const LevelType *lvlTypes,
565548
const uint64_t *dim2lvl, const uint64_t *lvl2dim) {
549+
SparseTensorCOO<V> *noLvlCOO = nullptr;
566550
return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
567-
lvlTypes, dim2lvl, lvl2dim, nullptr,
568-
true);
551+
lvlTypes, dim2lvl, lvl2dim, noLvlCOO);
569552
}
570553

571554
template <typename P, typename C, typename V>
572555
SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newFromCOO(
573556
uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
574557
const uint64_t *lvlSizes, const LevelType *lvlTypes,
575558
const uint64_t *dim2lvl, const uint64_t *lvl2dim,
576-
SparseTensorCOO<V> &lvlCOO) {
559+
SparseTensorCOO<V> *lvlCOO) {
560+
assert(lvlCOO);
577561
return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
578562
lvlTypes, dim2lvl, lvl2dim, lvlCOO);
579563
}
580564

581565
template <typename P, typename C, typename V>
582-
SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::packFromLvlBuffers(
566+
SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newFromBuffers(
583567
uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
584568
const uint64_t *lvlSizes, const LevelType *lvlTypes,
585569
const uint64_t *dim2lvl, const uint64_t *lvl2dim, uint64_t srcRank,
@@ -599,10 +583,9 @@ SparseTensorStorage<P, C, V>::SparseTensorStorage(
599583
uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
600584
const uint64_t *lvlSizes, const LevelType *lvlTypes,
601585
const uint64_t *dim2lvl, const uint64_t *lvl2dim,
602-
SparseTensorCOO<V> *lvlCOO, bool initializeValuesIfAllDense)
586+
SparseTensorCOO<V> *lvlCOO)
603587
: SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
604588
dim2lvl, lvl2dim) {
605-
assert(!lvlCOO || lvlRank == lvlCOO->getRank());
606589
// Provide hints on capacity of positions and coordinates.
607590
// TODO: needs much fine-tuning based on actual sparsity; currently
608591
// we reserve position/coordinate space based on all previous dense
@@ -633,27 +616,20 @@ SparseTensorStorage<P, C, V>::SparseTensorStorage(
633616
sz = detail::checkedMul(sz, lvlSizes[l]);
634617
}
635618
}
636-
if (allDense && initializeValuesIfAllDense)
619+
if (lvlCOO) {
620+
/* New from COO: ensure it is sorted. */
621+
assert(lvlCOO->getRank() == lvlRank);
622+
lvlCOO->sort();
623+
// Now actually insert the `elements`.
624+
const auto &elements = lvlCOO->getElements();
625+
const uint64_t nse = elements.size();
626+
assert(values.size() == 0);
627+
values.reserve(nse);
628+
fromCOO(elements, 0, nse, 0);
629+
} else if (allDense) {
630+
/* New empty (all dense) */
637631
values.resize(sz, 0);
638-
}
639-
640-
template <typename P, typename C, typename V>
641-
SparseTensorStorage<P, C, V>::SparseTensorStorage( // NOLINT
642-
uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
643-
const uint64_t *lvlSizes, const LevelType *lvlTypes,
644-
const uint64_t *dim2lvl, const uint64_t *lvl2dim,
645-
SparseTensorCOO<V> &lvlCOO)
646-
: SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
647-
dim2lvl, lvl2dim, nullptr, false) {
648-
// Ensure lvlCOO is sorted.
649-
assert(lvlRank == lvlCOO.getRank());
650-
lvlCOO.sort();
651-
// Now actually insert the `elements`.
652-
const auto &elements = lvlCOO.getElements();
653-
const uint64_t nse = elements.size();
654-
assert(values.size() == 0);
655-
values.reserve(nse);
656-
fromCOO(elements, 0, nse, 0);
632+
}
657633
}
658634

659635
template <typename P, typename C, typename V>

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ extern "C" {
127127
case Action::kPack: { \
128128
assert(ptr && "Received nullptr for SparseTensorStorage object"); \
129129
intptr_t *buffers = static_cast<intptr_t *>(ptr); \
130-
return SparseTensorStorage<P, C, V>::packFromLvlBuffers( \
130+
return SparseTensorStorage<P, C, V>::newFromBuffers( \
131131
dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \
132132
dimRank, buffers); \
133133
} \

0 commit comments

Comments
 (0)