[mlir][sparse] replace magic constant with symbol

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D137177
This commit is contained in:
Aart Bik
2022-11-01 10:34:13 -07:00
parent 5baa4b8e11
commit 10db57b7ea

View File

@@ -31,9 +31,7 @@ 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 DimSizesIdx = 0;
static constexpr uint64_t MemSizesIdx = 2;
static constexpr uint64_t FieldsIdx = 3;
@@ -88,11 +86,12 @@ static Optional<Value> sizeFromTensorAtDim(OpBuilder &rewriter, Location loc,
if (!ShapedType::isDynamic(shape[dim]))
return constantIndex(rewriter, loc, shape[dim]);
// Any other query can consult the dimSizes array at field 0 using,
// Any other query can consult the dimSizes array at field DimSizesIdx,
// accounting for the reordering applied to the sparse storage.
auto tuple = getTuple(adaptedValue);
Value idx = constantIndex(rewriter, loc, toStoredDim(tensorTp, dim));
return rewriter.create<memref::LoadOp>(loc, tuple.getInputs().front(), idx)
return rewriter
.create<memref::LoadOp>(loc, tuple.getInputs()[DimSizesIdx], idx)
.getResult();
}