[Mlir-commits] [mlir] [mlir][sparse] cleanup sparse runtime library (PR #82807)
Aart Bik
llvmlistbot at llvm.org
Fri Feb 23 10:26:59 PST 2024
https://github.com/aartbik updated https://github.com/llvm/llvm-project/pull/82807
>From 9579fcfe665e11841deddb8b3e720460e0b2dab9 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Fri, 23 Feb 2024 10:14:26 -0800
Subject: [PATCH 1/2] [mlir][sparse] cleanup sparse runtime library
remove some obsoleted APIs from the library that have
been fully replaced with actual direct IR codegen
---
.../mlir/Dialect/SparseTensor/IR/Enums.h | 9 ++--
.../ExecutionEngine/SparseTensor/Storage.h | 42 ++--------------
.../ExecutionEngine/SparseTensorRuntime.h | 24 +--------
.../ExecutionEngine/SparseTensor/Storage.cpp | 7 ---
.../ExecutionEngine/SparseTensorRuntime.cpp | 50 ++-----------------
5 files changed, 11 insertions(+), 121 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
index a00c9c31256c96..1c81d80ea7ec4e 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
@@ -145,12 +145,9 @@ constexpr bool isComplexPrimaryType(PrimaryType valTy) {
/// The actions performed by @newSparseTensor.
enum class Action : uint32_t {
kEmpty = 0,
- kEmptyForward = 1,
- kFromCOO = 2,
- kFromReader = 4,
- kToCOO = 5,
- kPack = 7,
- kSortCOOInPlace = 8,
+ kFromReader = 1,
+ kPack = 2,
+ kSortCOOInPlace = 3,
};
/// This enum defines all supported storage format without the level properties.
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
index eff1aca42d3e8f..c3c1c9c94f9a3c 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
@@ -149,13 +149,6 @@ class SparseTensorStorageBase {
MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETVALUES)
#undef DECL_GETVALUES
- /// Element-wise forwarding insertions. The first argument is the
- /// dimension-coordinates for the value being inserted.
-#define DECL_FORWARDINGINSERT(VNAME, V) \
- virtual void forwardingInsert(const uint64_t *, V);
- MLIR_SPARSETENSOR_FOREVERY_V(DECL_FORWARDINGINSERT)
-#undef DECL_FORWARDINGINSERT
-
/// Element-wise insertion in lexicographic coordinate order. The first
/// argument is the level-coordinates for the value being inserted.
#define DECL_LEXINSERT(VNAME, V) virtual void lexInsert(const uint64_t *, V);
@@ -171,9 +164,6 @@ class SparseTensorStorageBase {
MLIR_SPARSETENSOR_FOREVERY_V(DECL_EXPINSERT)
#undef DECL_EXPINSERT
- /// Finalizes forwarding insertions.
- virtual void endForwardingInsert() = 0;
-
/// Finalizes lexicographic insertions.
virtual void endLexInsert() = 0;
@@ -248,7 +238,7 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
static SparseTensorStorage<P, C, V> *
newEmpty(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
const uint64_t *lvlSizes, const LevelType *lvlTypes,
- const uint64_t *dim2lvl, const uint64_t *lvl2dim, bool forwarding);
+ const uint64_t *dim2lvl, const uint64_t *lvl2dim);
/// Allocates a new sparse tensor and initializes it from the given COO.
static SparseTensorStorage<P, C, V> *
@@ -284,13 +274,6 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
*out = &values;
}
- /// Partially specialize forwarding insertions based on template types.
- void forwardingInsert(const uint64_t *dimCoords, V val) final {
- assert(dimCoords && coo);
- map.pushforward(dimCoords, lvlCursor.data());
- coo->add(lvlCursor, val);
- }
-
/// Partially specialize lexicographical insertions based on template types.
void lexInsert(const uint64_t *lvlCoords, V val) final {
assert(lvlCoords);
@@ -345,21 +328,6 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
}
}
- /// Finalizes forwarding insertions.
- void endForwardingInsert() final {
- // Ensure COO is sorted.
- assert(coo);
- coo->sort();
- // Now actually insert the `elements`.
- const auto &elements = coo->getElements();
- const uint64_t nse = elements.size();
- assert(values.size() == 0);
- values.reserve(nse);
- fromCOO(elements, 0, nse, 0);
- delete coo;
- coo = nullptr;
- }
-
/// Finalizes lexicographic insertions.
void endLexInsert() final {
if (!allDense) {
@@ -653,13 +621,9 @@ template <typename P, typename C, typename V>
SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newEmpty(
uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
const uint64_t *lvlSizes, const LevelType *lvlTypes,
- const uint64_t *dim2lvl, const uint64_t *lvl2dim, bool forwarding) {
- SparseTensorCOO<V> *lvlCOO = nullptr;
- if (forwarding)
- lvlCOO = new SparseTensorCOO<V>(lvlRank, lvlSizes);
+ const uint64_t *dim2lvl, const uint64_t *lvl2dim) {
return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
- lvlTypes, dim2lvl, lvl2dim, lvlCOO,
- !forwarding);
+ lvlTypes, dim2lvl, lvl2dim, nullptr, true);
}
template <typename P, typename C, typename V>
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h b/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
index 8b0829aab0d8d0..d916186c835c2e 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
@@ -38,15 +38,12 @@ extern "C" {
/// This is the "swiss army knife" method for materializing sparse
/// tensors into the computation. The types of the `ptr` argument and
/// the result depend on the action, as explained in the following table,
-/// where "STS" means a sparse-tensor-storage object and "COO" means
-/// a coordinate-scheme object.
+/// where "STS" means a sparse-tensor-storage object.
///
/// Action: `ptr`: Returns:
+/// ---------------------------------------------------------------------------
/// kEmpty - STS, empty
-/// kEmptyForward - STS, empty, with forwarding COO
-/// kFromCOO COO STS, copied from the COO source
/// kFromReader reader STS, input from reader
-/// kToCOO STS COO, copied from the STS source
/// kPack buffers STS, from level buffers
/// kSortCOOInPlace STS STS, sorted in place
MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_newSparseTensor( // NOLINT
@@ -80,14 +77,6 @@ MLIR_SPARSETENSOR_FOREVERY_O(DECL_SPARSEPOSITIONS)
MLIR_SPARSETENSOR_FOREVERY_O(DECL_SPARSECOORDINATES)
#undef DECL_SPARSECOORDINATES
-/// Tensor-storage method for a dim to lvl forwarding insertion.
-#define DECL_FORWARDINGINSERT(VNAME, V) \
- MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_forwardingInsert##VNAME( \
- void *tensor, StridedMemRefType<V, 0> *vref, \
- StridedMemRefType<index_type, 1> *dimCoordsRef); \
- MLIR_SPARSETENSOR_FOREVERY_V(DECL_FORWARDINGINSERT)
-#undef DECL_FORWARDINGINSERT
-
/// Tensor-storage method to insert elements in lexicographical
/// level-coordinate order.
#define DECL_LEXINSERT(VNAME, V) \
@@ -160,21 +149,12 @@ MLIR_CRUNNERUTILS_EXPORT index_type sparseLvlSize(void *tensor, index_type l);
/// Tensor-storage method to get the size of the given dimension.
MLIR_CRUNNERUTILS_EXPORT index_type sparseDimSize(void *tensor, index_type d);
-/// Tensor-storage method to finalize forwarding insertions.
-MLIR_CRUNNERUTILS_EXPORT void endForwardingInsert(void *tensor);
-
/// Tensor-storage method to finalize lexicographic insertions.
MLIR_CRUNNERUTILS_EXPORT void endLexInsert(void *tensor);
/// Releases the memory for the tensor-storage object.
MLIR_CRUNNERUTILS_EXPORT void delSparseTensor(void *tensor);
-/// Releases the memory for the coordinate-scheme object.
-#define DECL_DELCOO(VNAME, V) \
- MLIR_CRUNNERUTILS_EXPORT void delSparseTensorCOO##VNAME(void *coo);
-MLIR_SPARSETENSOR_FOREVERY_V(DECL_DELCOO)
-#undef DECL_DELCOO
-
/// Helper function to read a sparse tensor filename from the environment,
/// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(index_type id);
diff --git a/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp b/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp
index 9e8b240899d808..bbe10b0dcdd465 100644
--- a/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp
@@ -74,13 +74,6 @@ MLIR_SPARSETENSOR_FOREVERY_FIXED_O(IMPL_GETCOORDINATES)
MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETVALUES)
#undef IMPL_GETVALUES
-#define IMPL_FORWARDINGINSERT(VNAME, V) \
- void SparseTensorStorageBase::forwardingInsert(const uint64_t *, V) { \
- FATAL_PIV("forwardingInsert" #VNAME); \
- }
-MLIR_SPARSETENSOR_FOREVERY_V(IMPL_FORWARDINGINSERT)
-#undef IMPL_FORWARDINGINSERT
-
#define IMPL_LEXINSERT(VNAME, V) \
void SparseTensorStorageBase::lexInsert(const uint64_t *, V) { \
FATAL_PIV("lexInsert" #VNAME); \
diff --git a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
index a5e75a77b4e47f..0bc90b281b7ced 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
@@ -117,20 +117,7 @@ extern "C" {
switch (action) { \
case Action::kEmpty: { \
return SparseTensorStorage<P, C, V>::newEmpty( \
- dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \
- false); \
- } \
- case Action::kEmptyForward: { \
- return SparseTensorStorage<P, C, V>::newEmpty( \
- dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \
- true); \
- } \
- case Action::kFromCOO: { \
- assert(ptr && "Received nullptr for SparseTensorCOO object"); \
- auto &coo = *static_cast<SparseTensorCOO<V> *>(ptr); \
- return SparseTensorStorage<P, C, V>::newFromCOO( \
- dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim, \
- coo); \
+ dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim); \
} \
case Action::kFromReader: { \
assert(ptr && "Received nullptr for SparseTensorReader object"); \
@@ -138,11 +125,6 @@ extern "C" {
return static_cast<void *>(reader.readSparseTensor<P, C, V>( \
lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim)); \
} \
- case Action::kToCOO: { \
- assert(ptr && "Received nullptr for SparseTensorStorage object"); \
- auto &tensor = *static_cast<SparseTensorStorage<P, C, V> *>(ptr); \
- return tensor.toCOO(); \
- } \
case Action::kPack: { \
assert(ptr && "Received nullptr for SparseTensorStorage object"); \
intptr_t *buffers = static_cast<intptr_t *>(ptr); \
@@ -341,21 +323,6 @@ MLIR_SPARSETENSOR_FOREVERY_O(IMPL_SPARSECOORDINATES)
#undef IMPL_SPARSECOORDINATES
#undef IMPL_GETOVERHEAD
-#define IMPL_FORWARDINGINSERT(VNAME, V) \
- void _mlir_ciface_forwardingInsert##VNAME( \
- void *t, StridedMemRefType<V, 0> *vref, \
- StridedMemRefType<index_type, 1> *dimCoordsRef) { \
- assert(t &&vref); \
- ASSERT_NO_STRIDE(dimCoordsRef); \
- const index_type *dimCoords = MEMREF_GET_PAYLOAD(dimCoordsRef); \
- assert(dimCoords); \
- const V *value = MEMREF_GET_PAYLOAD(vref); \
- static_cast<SparseTensorStorageBase *>(t)->forwardingInsert(dimCoords, \
- *value); \
- }
-MLIR_SPARSETENSOR_FOREVERY_V(IMPL_FORWARDINGINSERT)
-#undef IMPL_FORWARDINGINSERT
-
#define IMPL_LEXINSERT(VNAME, V) \
void _mlir_ciface_lexInsert##VNAME( \
void *t, StridedMemRefType<index_type, 1> *lvlCoordsRef, \
@@ -427,8 +394,8 @@ void _mlir_ciface_getSparseTensorReaderDimSizes(
const uint64_t cSize = MEMREF_GET_USIZE(cref); \
const uint64_t vSize = MEMREF_GET_USIZE(vref); \
ASSERT_USIZE_EQ(lvl2dimRef, dimRank); \
- assert(cSize >= lvlRank * vSize); \
- assert(vSize >= reader.getNSE() && "Not enough space in buffers"); \
+ assert(cSize >= lvlRank * reader.getNSE()); \
+ assert(vSize >= reader.getNSE()); \
(void)dimRank; \
(void)cSize; \
(void)vSize; \
@@ -488,10 +455,6 @@ index_type sparseDimSize(void *tensor, index_type d) {
return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
}
-void endForwardingInsert(void *tensor) {
- return static_cast<SparseTensorStorageBase *>(tensor)->endForwardingInsert();
-}
-
void endLexInsert(void *tensor) {
return static_cast<SparseTensorStorageBase *>(tensor)->endLexInsert();
}
@@ -500,13 +463,6 @@ void delSparseTensor(void *tensor) {
delete static_cast<SparseTensorStorageBase *>(tensor);
}
-#define IMPL_DELCOO(VNAME, V) \
- void delSparseTensorCOO##VNAME(void *coo) { \
- delete static_cast<SparseTensorCOO<V> *>(coo); \
- }
-MLIR_SPARSETENSOR_FOREVERY_V(IMPL_DELCOO)
-#undef IMPL_DELCOO
-
char *getTensorFilename(index_type id) {
constexpr size_t bufSize = 80;
char var[bufSize];
>From 7a3cb2c767382397fbf07bd236b497baee2bdea7 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Fri, 23 Feb 2024 10:26:31 -0800
Subject: [PATCH 2/2] clang-format
---
mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
index c3c1c9c94f9a3c..fe0e08b5c84036 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
@@ -623,7 +623,8 @@ SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newEmpty(
const uint64_t *lvlSizes, const LevelType *lvlTypes,
const uint64_t *dim2lvl, const uint64_t *lvl2dim) {
return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
- lvlTypes, dim2lvl, lvl2dim, nullptr, true);
+ lvlTypes, dim2lvl, lvl2dim, nullptr,
+ true);
}
template <typename P, typename C, typename V>
More information about the Mlir-commits
mailing list