diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp @@ -567,6 +567,18 @@ Value newBuffer = rewriter.create(loc, bufferType, buffer, capacity); + // Initialize the buffer to support debugging. + const char *envStr = std::getenv("SPARSE_TENSOR_INITIALIZE_BUFFER"); + if (envStr && strcmp(envStr, "1") == 0) { + Value fillSize = rewriter.create(loc, capacity, newSize); + Value fillValue = rewriter.create( + loc, value.getType(), rewriter.getZeroAttr(value.getType())); + Value subBuffer = rewriter.create( + loc, newBuffer, /*offset=*/ValueRange{newSize}, + /*size=*/ValueRange{fillSize}, + /*step=*/ValueRange{constantIndex(rewriter, loc, 1)}); + rewriter.create(loc, fillValue, subBuffer); + } rewriter.create(loc, newBuffer); // False branch.