[mlir][sparse] build proper insertion chain

The alloc->insert/compress->load chain needs to be
properly represented with an SSA chain now in loops
and if statements to properly reflect the modifying
behavior (runtime support lib is forgiving on breaking
this, but the new codegen is not).

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D136966
This commit is contained in:
Aart Bik
2022-10-28 11:14:58 -07:00
parent eb536613a2
commit 5661647e85
13 changed files with 467 additions and 388 deletions

View File

@@ -31,6 +31,9 @@ using namespace mlir::sparse_tensor;
namespace {
// TODO: start using these when insertions are implemented
// static constexpr uint64_t DimSizesIdx = 0;
// static constexpr uint64_t DimCursorIdx = 1;
static constexpr uint64_t MemSizesIdx = 2;
static constexpr uint64_t FieldsIdx = 3;
@@ -632,13 +635,7 @@ public:
filled, index);
rewriter.create<scf::YieldOp>(loc, fields);
// Deallocate the buffers on exit of the full loop nest.
Operation *parent = op;
for (; isa<scf::ForOp>(parent->getParentOp()) ||
isa<scf::WhileOp>(parent->getParentOp()) ||
isa<scf::ParallelOp>(parent->getParentOp()) ||
isa<scf::IfOp>(parent->getParentOp());
parent = parent->getParentOp())
;
Operation *parent = getTop(op);
rewriter.setInsertionPointAfter(parent);
rewriter.create<memref::DeallocOp>(loc, values);
rewriter.create<memref::DeallocOp>(loc, filled);