[clang] [mlir][sparse] introduce MapRef, unify conversion/codegen for reader (PR #68360)
https://github.com/aartbik closed https://github.com/llvm/llvm-project/pull/68360 ___ cfe-commits mailing list cfe-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
[clang] [mlir][sparse] introduce MapRef, unify conversion/codegen for reader (PR #68360)
https://github.com/yinying-lisa-li approved this pull request. https://github.com/llvm/llvm-project/pull/68360 ___ cfe-commits mailing list cfe-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
[clang] [mlir][sparse] introduce MapRef, unify conversion/codegen for reader (PR #68360)
@@ -729,3 +729,92 @@ Value sparse_tensor::createOrFoldSliceStrideOp(OpBuilder
, Location loc,
return constantIndex(builder, loc, *stride);
return builder.create(loc, tensor, APInt(64, dim));
}
+
+void sparse_tensor::fillDimShape(OpBuilder , Location loc,
+ SparseTensorType stt,
+ SmallVectorImpl ) {
+ out.clear();
+ out.reserve(stt.getDimRank());
+ for (const DynSize sh : stt.getDimShape()) {
+const auto s = ShapedType::isDynamic(sh) ? 0 : sh;
+out.push_back(constantIndex(builder, loc, s));
+ }
+}
+
+Value sparse_tensor::genReader(OpBuilder , Location loc,
+ SparseTensorType stt, Value tensor,
+ /*out*/ SmallVectorImpl ,
+ /*out*/ Value ) {
+ // Construct the dimShapes buffer. The buffer contains the static size
+ // per dimension, or otherwise a zero for a dynamic size.
+ fillDimShape(builder, loc, stt, dimShapesValues);
+ Value dimShapesBuffer = allocaBuffer(builder, loc, dimShapesValues);
+ // Create the `CheckedSparseTensorReader`. This reader performs a
+ // consistency check on the static sizes, but accepts any size
+ // of each dimension with a dynamic size.
+ Type opaqueTp = getOpaquePointerType(builder);
+ Type eltTp = stt.getElementType();
+ Value valTp = constantPrimaryTypeEncoding(builder, loc, eltTp);
+ Value reader =
+ createFuncCall(builder, loc, "createCheckedSparseTensorReader", opaqueTp,
+ {tensor, dimShapesBuffer, valTp}, EmitCInterface::On)
+ .getResult(0);
+ // For static shapes, the shape buffer can be used right away. For dynamic
+ // shapes, use the information from the reader to construct a buffer that
+ // supplies the actual size for each dynamic dimension.
+ dimSizesBuffer = dimShapesBuffer;
+ if (stt.hasDynamicDimShape()) {
+Type indexTp = builder.getIndexType();
+auto memTp = MemRefType::get({ShapedType::kDynamic}, indexTp);
+dimSizesBuffer =
+createFuncCall(builder, loc, "getSparseTensorReaderDimSizes", memTp,
+ reader, EmitCInterface::On)
+.getResult(0);
+ }
+ return reader;
+}
+
+Value sparse_tensor::genReaderBuffers(OpBuilder , Location loc,
+ SparseTensorType stt,
+ ArrayRef dimShapesValues,
+ Value dimSizesBuffer,
+ /*out*/ Value ,
+ /*out*/ Value ) {
+ const Dimension dimRank = stt.getDimRank();
+ const Level lvlRank = stt.getLvlRank();
+ // For an identify mapping, the dim2lvl and lvl2dim mappings are
yinying-lisa-li wrote:
identity?
https://github.com/llvm/llvm-project/pull/68360
___
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
[clang] [mlir][sparse] introduce MapRef, unify conversion/codegen for reader (PR #68360)
https://github.com/PeimingLiu approved this pull request. https://github.com/llvm/llvm-project/pull/68360 ___ cfe-commits mailing list cfe-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
[clang] [mlir][sparse] introduce MapRef, unify conversion/codegen for reader (PR #68360)
https://github.com/aartbik updated
https://github.com/llvm/llvm-project/pull/68360
>From 6094912685a0cfa5c13e023e8ec97238a84fca2f Mon Sep 17 00:00:00 2001
From: Aart Bik
Date: Thu, 5 Oct 2023 13:22:28 -0700
Subject: [PATCH 1/4] [mlir][sparse] introduce MapRef, unify conversion/codegen
for reader
This revision introduces a MapRef, which will support a future
generalization beyond permutations (e.g. block sparsity). This
revision also unifies the conversion/codegen paths for the
sparse_tensor.new operation from file (eg. the readers). Note
that more unification is planned as well as general affine
dim2lvl and lvl2dim (all marked with TODOs).
---
.../mlir/ExecutionEngine/SparseTensor/File.h | 156 ++--
.../ExecutionEngine/SparseTensor/MapRef.h | 96 ++
.../ExecutionEngine/SparseTensor/Storage.h| 108 +--
.../ExecutionEngine/SparseTensorRuntime.h | 8 -
.../SparseTensor/Transforms/CodegenUtils.cpp | 89 +
.../SparseTensor/Transforms/CodegenUtils.h| 18 ++
.../Transforms/SparseTensorCodegen.cpp| 73 ++--
.../Transforms/SparseTensorConversion.cpp | 111 ++-
.../SparseTensor/CMakeLists.txt | 1 +
.../ExecutionEngine/SparseTensor/MapRef.cpp | 52 ++
.../ExecutionEngine/SparseTensorRuntime.cpp | 60 +++---
mlir/test/Dialect/SparseTensor/codegen.mlir | 172 +-
.../test/Dialect/SparseTensor/conversion.mlir | 18 +-
13 files changed, 475 insertions(+), 487 deletions(-)
create mode 100644 mlir/include/mlir/ExecutionEngine/SparseTensor/MapRef.h
create mode 100644 mlir/lib/ExecutionEngine/SparseTensor/MapRef.cpp
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
index 78c1a0544e3a521..9157bfa7e773239 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
@@ -20,6 +20,7 @@
#ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_FILE_H
#define MLIR_EXECUTIONENGINE_SPARSETENSOR_FILE_H
+#include "mlir/ExecutionEngine/SparseTensor/MapRef.h"
#include "mlir/ExecutionEngine/SparseTensor/Storage.h"
#include
@@ -75,6 +76,10 @@ inline V readValue(char **linePtr, bool isPattern) {
} // namespace detail
+//===--===//
+//
+// Reader class.
+//
//===--===//
/// This class abstracts over the information stored in file headers,
@@ -132,6 +137,7 @@ class SparseTensorReader final {
/// Reads and parses the file's header.
void readHeader();
+ /// Returns the stored value kind.
ValueKind getValueKind() const { return valueKind_; }
/// Checks if a header has been successfully read.
@@ -185,58 +191,37 @@ class SparseTensorReader final {
/// valid after parsing the header.
void assertMatchesShape(uint64_t rank, const uint64_t *shape) const;
- /// Reads a sparse tensor element from the next line in the input file and
- /// returns the value of the element. Stores the coordinates of the element
- /// to the `dimCoords` array.
- template
- V readElement(uint64_t dimRank, uint64_t *dimCoords) {
-assert(dimRank == getRank() && "rank mismatch");
-char *linePtr = readCoords(dimCoords);
-return detail::readValue(, isPattern());
- }
-
- /// Allocates a new COO object for `lvlSizes`, initializes it by reading
- /// all the elements from the file and applying `dim2lvl` to their
- /// dim-coordinates, and then closes the file. Templated on V only.
- template
- SparseTensorCOO *readCOO(uint64_t lvlRank, const uint64_t *lvlSizes,
- const uint64_t *dim2lvl);
-
/// Allocates a new sparse-tensor storage object with the given encoding,
/// initializes it by reading all the elements from the file, and then
/// closes the file. Templated on P, I, and V.
template
SparseTensorStorage *
readSparseTensor(uint64_t lvlRank, const uint64_t *lvlSizes,
- const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
- const uint64_t *dim2lvl) {
-auto *lvlCOO = readCOO(lvlRank, lvlSizes, dim2lvl);
+ const DimLevelType *lvlTypes, const uint64_t *dim2lvl,
+ const uint64_t *lvl2dim) {
+const uint64_t dimRank = getRank();
+MapRef map(dimRank, lvlRank, dim2lvl, lvl2dim);
+auto *coo = readCOO(map, lvlSizes);
auto *tensor = SparseTensorStorage::newFromCOO(
-getRank(), getDimSizes(), lvlRank, lvlTypes, lvl2dim, *lvlCOO);
-delete lvlCOO;
+dimRank, getDimSizes(), lvlRank, lvlTypes, lvl2dim, *coo);
+delete coo;
return tensor;
}
/// Reads the COO tensor from the file, stores the coordinates and values to
/// the given buffers, returns a boolean value to indicate whether the COO
/// elements are sorted.
- /// Precondition: the buffers should
