[mlir][sparse] sparse storage scheme type conversion

[aartbik]

Summary

This builds a compound type for the buffers required for the sparse storage scheme defined by the MLIR sparse tensor types. The use of a tuple allows for a simple 1:1 type conversion. A subsequent pass can expand this tuple into its component with an isolated 1:N type conversion.

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

63

Good to see it here!

74

So the permutation is not handled here? I assume that it will be considered later when trying to mapping back from dim => tuple index?

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

74

That is right. It may be surprising, but the dimension level types are interpreted *after* applying the reordering (see my comment in this posting: https://discourse.llvm.org/t/sparse-tensors-in-mlir/3389/12). So we can construct the basic type without considering the ordering. However, when filling in the sizes *and* when querying the type, reordering needs to be included. I think I will make some top level utilities for this, so we don't make mistakes in the codegen parts that we will add below.

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

 using namespace mlir::sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
-/// Maps each sparse tensor type to the appropriate buffer.
-static Optional<Type> convertSparseTensorTypes(Type type) {
-  if (getSparseTensorEncoding(type) != nullptr) {
-    // TODO: this is just a dummy rule to get the ball rolling....
-    RankedTensorType rTp = type.cast<RankedTensorType>();
-    return MemRefType::get({ShapedType::kDynamicSize}, rTp.getElementType());
-  }
+/// Maps a sparse tensor type to the appropriate compounded buffers.
+static Optional<Type> convertSparseTensorType(Type type) {
+  auto enc = getSparseTensorEncoding(type);
+  if (!enc)
     return llvm::None;
+  // Construct the basic types.
+  auto context = type.getContext();
+  unsigned idxWidth = enc.getIndexBitWidth();
+  unsigned ptrWidth = enc.getPointerBitWidth();
+  RankedTensorType rType = type.cast<RankedTensorType>();
+  Type indexType = IndexType::get(context);
+  Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
+  Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
+  Type eltType = rType.getElementType();
+  //
+  // Sparse tensor storage for rank-dimensional tensor is organized as a
+  // single compound type with the following fields:
+  //
+  // struct {
+  //   memref<rank x index> dimSize  ; size in each dimension
+  //   ; per-dimension d:
+  //   ;  if dense:
+  //        <nothing>
+  //   ;  if compresed:
+  //        memref<? x idx>  indices-d   ; indices for sparse dim d
+  //        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
+  //   ;  if singleton:
+  //        memref<? x idx>  indices-d   ; indices for singleton dim d
+  //   memref<? x eltType> values    ; values
+  // };
+  //
+  // TODO: fill in the ? when statically known
+  //
+  // TODO: emit dimSizes when not needed (e.g. all-dense)
+  //
+  unsigned rank = rType.getShape().size();
+  SmallVector<Type, 8> fields;
+  fields.push_back(MemRefType::get({rank}, indexType));
+  for (unsigned r = 0; r < rank; r++) {
+    // Dimension level types apply in order to the reordered dimension.
+    // As a result, the compound type can be constructed directly in the given
+    // order. Clients of this type know what field is what from the sparse
+    // tensor type.
+    switch (enc.getDimLevelType()[r]) {
+    case SparseTensorEncodingAttr::DimLevelType::Dense:
+      break;
+    case SparseTensorEncodingAttr::DimLevelType::Compressed:
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
+      fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
+      fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, ptrType));
+      break;
+    case SparseTensorEncodingAttr::DimLevelType::Singleton:
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
+      fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
+      break;
+    }
+  }
+  fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));
+  // Sparse tensor storage (temporarily) lives in a tuple. This allows a
+  // simple 1:1 type conversion during codegen. A subsequent pass uses
+  // a 1:N type conversion to expand the tuple into its fields.
+  return TupleType::get(context, fields);
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {

 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into an actual buffer.
 //===----------------------------------------------------------------------===//
 
 mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
-  addConversion(convertSparseTensorTypes);
+  addConversion(convertSparseTensorType);
 }
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter>(typeConverter, patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

 // RUN: mlir-opt %s --sparse-tensor-codegen  --canonicalize --cse | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{
-  dimLevelType = ["compressed"]
+  dimLevelType = [ "compressed" ],
+  indexBitWidth = 64,
+  pointerBitWidth = 32
 }>
 
-// TODO: just a dummy memref rewriting to get the ball rolling....
+#Dense = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "dense" ],
+  indexBitWidth = 64,
+  pointerBitWidth = 32
+}>
+
+#Row = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "dense" ],
+  indexBitWidth = 64,
+  pointerBitWidth = 32
+}>
+
+#CSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  indexBitWidth = 64,
+  pointerBitWidth = 32
+}>
+
+#DCSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ],
+  indexBitWidth = 64,
+  pointerBitWidth = 32
+}>
 
 // CHECK-LABEL: func @sparse_nop(
-//  CHECK-SAME: %[[A:.*]]: memref<?xf64>) -> memref<?xf64> {
-//       CHECK: return %[[A]] : memref<?xf64>
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<1xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>) -> tuple<memref<1xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>
+//       CHECK: return %[[A]] : tuple<memref<1xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>
 func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
   return %arg0 : tensor<?xf64, #SparseVector>
 }
+
+// CHECK-LABEL: func @sparse_dense(
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xf64>>)
+func.func @sparse_dense(%arg0: tensor<?x?xf64, #Dense>) {
+  return
+}
+
+// CHECK-LABEL: func @sparse_row(
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>)
+func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) {
+  return
+}
+
+// CHECK-LABEL: func @sparse_csr(
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>)
+func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) {
+  return
+}
+
+// CHECK-LABEL: func @sparse_dcsr(
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xf64>>)
+func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) {
+  return
+}