Skip to content

Commit 63bdcaf

Browse files
wrengrwrengr
committed
[mlir][sparse] Moving delete coo into codegen instead of runtime library
Prior to this change there were a number of places where the allocation and deallocation of SparseTensorCOO objects were not cleanly paired, leading to inconsistencies regarding whether each function released its tensor/coo arguments or not, as well as making it easy to run afoul of memory leaks, use-after-free, or double-free errors. This change cleans up the codegen vs runtime boundary to resolve those issues. Now, the only time the runtime library frees an object is either (a) because it's a function explicitly designed to do so, or (b) because the allocated object is entirely local to the function and would be a memory leak if not released. Thus, now the codegen takes complete responsibility for releasing any objects it caused to be allocated. Reviewed By: aartbik Differential Revision: https://reviews.llvm.org/D122435
1 parent ec0b332 commit 63bdcaf

File tree

3 files changed

+95
-44
lines changed

3 files changed

+95
-44
lines changed

‎mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,14 @@ static Value genIndexAndValueForDense(ConversionPatternRewriter &rewriter,
253253
return val;
254254
}
255255

256+
/// Generates a call to release/delete a `SparseTensorCOO`.
257+
static void genDelCOOCall(OpBuilder &builder, Operation *op, Type elemTp,
258+
Value coo) {
259+
SmallString<21> name{"delSparseTensorCOO", primaryTypeFunctionSuffix(elemTp)};
260+
TypeRange noTp;
261+
createFuncCall(builder, op, name, noTp, coo, EmitCInterface::Off);
262+
}
263+
256264
/// Generates a call that adds one element to a coordinate scheme.
257265
/// In particular, this generates code like the following:
258266
/// val = a[i1,..,ik];
@@ -501,7 +509,9 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
501509
params[4] = constantIndexTypeEncoding(rewriter, loc, encDst);
502510
params[6] = constantAction(rewriter, loc, Action::kFromCOO);
503511
params[7] = coo;
504-
rewriter.replaceOp(op, genNewCall(rewriter, op, params));
512+
Value dst = genNewCall(rewriter, op, params);
513+
genDelCOOCall(rewriter, op, stp.getElementType(), coo);
514+
rewriter.replaceOp(op, dst);
505515
return success();
506516
}
507517
if (!encDst && encSrc) {
@@ -545,6 +555,7 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
545555
insertScalarIntoDenseTensor(rewriter, loc, elemPtr, dst, rank, ind);
546556
rewriter.create<scf::YieldOp>(loc);
547557
rewriter.setInsertionPointAfter(whileOp);
558+
genDelCOOCall(rewriter, op, elemTp, iter);
548559
rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, dst);
549560
return success();
550561
}
@@ -584,7 +595,7 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
584595
SmallVector<Value, 8> params;
585596
sizesFromSrc(rewriter, sizes, loc, src);
586597
newParams(rewriter, params, op, stp, encDst, Action::kEmptyCOO, sizes);
587-
Value ptr = genNewCall(rewriter, op, params);
598+
Value coo = genNewCall(rewriter, op, params);
588599
Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
589600
Value perm = params[2];
590601
SmallVector<Value> lo;
@@ -620,13 +631,15 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
620631
ivs, rank);
621632
else
622633
val = genIndexAndValueForDense(rewriter, loc, src, ind, ivs);
623-
genAddEltCall(rewriter, op, eltType, ptr, val, ind, perm);
634+
genAddEltCall(rewriter, op, eltType, coo, val, ind, perm);
624635
return {};
625636
});
626637
// Final call to construct sparse tensor storage.
627638
params[6] = constantAction(rewriter, loc, Action::kFromCOO);
628-
params[7] = ptr;
629-
rewriter.replaceOp(op, genNewCall(rewriter, op, params));
639+
params[7] = coo;
640+
Value dst = genNewCall(rewriter, op, params);
641+
genDelCOOCall(rewriter, op, eltType, coo);
642+
rewriter.replaceOp(op, dst);
630643
return success();
631644
}
632645
};
@@ -822,8 +835,9 @@ class SparseTensorOutConverter : public OpConversionPattern<OutOp> {
822835
Type eltType = srcType.getElementType();
823836
SmallString<18> name{"outSparseTensor", primaryTypeFunctionSuffix(eltType)};
824837
TypeRange noTp;
825-
replaceOpWithFuncCall(rewriter, op, name, noTp, params,
826-
EmitCInterface::Off);
838+
createFuncCall(rewriter, op, name, noTp, params, EmitCInterface::Off);
839+
genDelCOOCall(rewriter, op, eltType, coo);
840+
rewriter.eraseOp(op);
827841
return success();
828842
}
829843
};

‎mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -400,7 +400,6 @@ class SparseTensorStorage : public SparseTensorStorageBase {
400400
assert(shape[r] == 0 || shape[r] == tensor->getSizes()[perm[r]]);
401401
n = new SparseTensorStorage<P, I, V>(tensor->getSizes(), perm, sparsity,
402402
tensor);
403-
delete tensor;
404403
} else {
405404
std::vector<uint64_t> permsz(rank);
406405
for (uint64_t r = 0; r < rank; r++) {
@@ -748,7 +747,6 @@ static void outSparseTensor(void *tensor, void *dest, bool sort) {
748747
file.flush();
749748
file.close();
750749
assert(file.good());
751-
delete coo;
752750
}
753751

754752
/// Initializes sparse tensor from an external COO-flavored format.
@@ -780,17 +778,19 @@ toMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape, V *values,
780778
#endif
781779

782780
// Convert external format to internal COO.
783-
auto *tensor = SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm, nse);
781+
auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm, nse);
784782
std::vector<uint64_t> idx(rank);
785783
for (uint64_t i = 0, base = 0; i < nse; i++) {
786784
for (uint64_t r = 0; r < rank; r++)
787785
idx[perm[r]] = indices[base + r];
788-
tensor->add(idx, values[i]);
786+
coo->add(idx, values[i]);
789787
base += rank;
790788
}
791789
// Return sparse tensor storage format as opaque pointer.
792-
return SparseTensorStorage<uint64_t, uint64_t, V>::newSparseTensor(
793-
rank, shape, perm, sparsity, tensor);
790+
auto *tensor = SparseTensorStorage<uint64_t, uint64_t, V>::newSparseTensor(
791+
rank, shape, perm, sparsity, coo);
792+
delete coo;
793+
return tensor;
794794
}
795795

796796
/// Converts a sparse tensor to an external COO-flavored format.
@@ -847,28 +847,31 @@ extern "C" {
847847

848848
#define CASE(p, i, v, P, I, V) \
849849
if (ptrTp == (p) && indTp == (i) && valTp == (v)) { \
850-
SparseTensorCOO<V> *tensor = nullptr; \
850+
SparseTensorCOO<V> *coo = nullptr; \
851851
if (action <= Action::kFromCOO) { \
852852
if (action == Action::kFromFile) { \
853853
char *filename = static_cast<char *>(ptr); \
854-
tensor = openSparseTensorCOO<V>(filename, rank, shape, perm); \
854+
coo = openSparseTensorCOO<V>(filename, rank, shape, perm); \
855855
} else if (action == Action::kFromCOO) { \
856-
tensor = static_cast<SparseTensorCOO<V> *>(ptr); \
856+
coo = static_cast<SparseTensorCOO<V> *>(ptr); \
857857
} else { \
858858
assert(action == Action::kEmpty); \
859859
} \
860-
return SparseTensorStorage<P, I, V>::newSparseTensor(rank, shape, perm, \
861-
sparsity, tensor); \
860+
auto *tensor = SparseTensorStorage<P, I, V>::newSparseTensor( \
861+
rank, shape, perm, sparsity, coo); \
862+
if (action == Action::kFromFile) \
863+
delete coo; \
864+
return tensor; \
862865
} \
863866
if (action == Action::kEmptyCOO) \
864867
return SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm); \
865-
tensor = static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm); \
868+
coo = static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm); \
866869
if (action == Action::kToIterator) { \
867-
tensor->startIterator(); \
870+
coo->startIterator(); \
868871
} else { \
869872
assert(action == Action::kToCOO); \
870873
} \
871-
return tensor; \
874+
return coo; \
872875
}
873876

874877
#define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
@@ -924,10 +927,8 @@ extern "C" {
924927
const uint64_t isize = iref->sizes[0]; \
925928
auto iter = static_cast<SparseTensorCOO<V> *>(tensor); \
926929
const Element<V> *elem = iter->getNext(); \
927-
if (elem == nullptr) { \
928-
delete iter; \
930+
if (elem == nullptr) \
929931
return false; \
930-
} \
931932
for (uint64_t r = 0; r < isize; r++) \
932933
indx[r] = elem->indices[r]; \
933934
*value = elem->value; \
@@ -1208,6 +1209,19 @@ void delSparseTensor(void *tensor) {
12081209
delete static_cast<SparseTensorStorageBase *>(tensor);
12091210
}
12101211

1212+
/// Releases sparse tensor coordinate scheme.
1213+
#define IMPL_DELCOO(VNAME, V) \
1214+
void delSparseTensorCOO##VNAME(void *coo) { \
1215+
delete static_cast<SparseTensorCOO<V> *>(coo); \
1216+
}
1217+
IMPL_DELCOO(F64, double)
1218+
IMPL_DELCOO(F32, float)
1219+
IMPL_DELCOO(I64, int64_t)
1220+
IMPL_DELCOO(I32, int32_t)
1221+
IMPL_DELCOO(I16, int16_t)
1222+
IMPL_DELCOO(I8, int8_t)
1223+
#undef IMPL_DELCOO
1224+
12111225
/// Initializes sparse tensor from a COO-flavored format expressed using C-style
12121226
/// data structures. The expected parameters are:
12131227
///

0 commit comments

Comments
 (0)