[mlir][sparse] improved tensor type lowering

[aartbik]

Summary

Also includes a first codegen example (although full support need tuple access)

[aartbik]

typo

[aartbik]

minor edit

[bixia]

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

66

I don't quite understand the implication of this optimization yet: Either casting a sparse tensor created from a dynamic shape to its "true static type" is not a noop, or we will have problems in reading the sparse tensor using its "true static type"?

66–67

Need to indent from the previous line.

[aartbik]

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

66

For a static type, codegen never needs to query this, since all information is implicit in the context (I showed the dimop lowering as the first example of that).

[aartbik]

fixed indentation

[aartbik]

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

66

The cast operation back and forth will indeed need to add/remove the sizes back as metadata.

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

 using namespace mlir::sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
+/// Reorders stored dimension to logical dimension.
+static unsigned reorder(const SparseTensorEncodingAttr &enc, unsigned d) {
+  auto order = enc.getDimOrdering();
+  if (order) {
+    assert(order.isPermutation());
+    return order.getDimPosition(d);
+  }
+  return d;
+}
+
 /// 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();
+  ArrayRef<int64_t> shape = rType.getShape();
   //
   // 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
+  //   ; if dynamic shape:
+  //     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
+  //   memref<? x eltType> values        ; values
   // };
   //
-  // TODO: fill in the ? when statically known
-  //
-  // TODO: emit dimSizes when not needed (e.g. all-dense)
-  //
+  int64_t linear = 1;
+  bool allDense = true;
   unsigned rank = rType.getShape().size();
   SmallVector<Type, 8> fields;
+  // The dimSizes array.
+  if (!rType.hasStaticShape())
     fields.push_back(MemRefType::get({rank}, indexType));
+  // Per-dimension storage.
   for (unsigned r = 0; r < rank; r++) {
+    // Get the original dimension (ro) for the current stored dimension (r).
+    unsigned ro = reorder(enc, 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:
+      // Linearize the size of consecutive dense dimensions.
+      if (ShapedType::isDynamic(shape[ro]) || ShapedType::isDynamic(linear))
+        linear = ShapedType::kDynamicSize;
+      else
+        linear *= shape[ro];
       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));
+      allDense = false;
+      linear = 1;
       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));
+      allDense = false;
+      linear = 1;
       break;
     }
   }
-  fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));
+  // The values array.
+  int64_t nnz =
+      (rType.hasStaticShape() && allDense) ? linear : ShapedType::kDynamicSize;
+  fields.push_back(MemRefType::get({nnz}, 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.
+// Codegen rules.
 //===----------------------------------------------------------------------===//
 
-/// Sparse conversion rule for returns.
+/// Sparse codegen rule for returns.
 class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<func::ReturnOp>(op, adaptor.getOperands());
     return success();
   }
 };
 
+/// Sparse codegen rule for dimension accesses.
+class SparseDimOpConverter : public OpConversionPattern<tensor::DimOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+    Type type = op.getSource().getType();
+    // Only rewrite annotated DimOp with constant index.
+    auto enc = getSparseTensorEncoding(type);
+    if (!enc)
+      return failure();
+    Optional<int64_t> index = op.getConstantIndex();
+    if (!index)
+      return failure();
+    // Access into static shape can query original type directly.
+    // Note that this is typically already done by DimOp's folding.
+    RankedTensorType rType = type.cast<RankedTensorType>();
+    if (rType.hasStaticShape()) {
+      rewriter.replaceOp(
+          op, constantIndex(rewriter, loc, rType.getShape()[*index]));
+      return success();
+    }
+    // Any other query can consult the dimSize array.
+    // TODO: this needs tuple access
+    return failure();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into an actual buffer.
 //===----------------------------------------------------------------------===//
 
 mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
   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());
+  patterns.add<SparseReturnConverter, SparseDimOpConverter>(
+      typeConverter, patterns.getContext());
 }

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

struct SparseTensorCodegenPass
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeToBufferConverter converter;
     ConversionTarget target(*ctx);
     // Everything in the sparse dialect must go!
     target.addIllegalDialect<SparseTensorDialect>();
-    // All dynamic rules below accept new function, call, return, and various
-    // tensor and bufferization operations as legal output of the rewriting.
+    // All dynamic rules below accept new function, call, return.
     target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
       return converter.isSignatureLegal(op.getFunctionType());
     });
     target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
       return converter.isSignatureLegal(op.getCalleeType());
     });
     target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
       return converter.isLegal(op.getOperandTypes());
     });
+    // Legal dialects may occur in generated code.
+    target.addLegalDialect<arith::ArithmeticDialect,
+                           bufferization::BufferizationDialect,
+                           memref::MemRefDialect, scf::SCFDialect>();
     // Populate with rules and apply rewriting rules.
     populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                    converter);
     populateCallOpTypeConversionPattern(patterns, converter);
     scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                          target);
     populateSparseTensorCodegenPatterns(converter, patterns);
     if (failed(applyPartialConversion(getOperation(), target,

mlir/test/Dialect/SparseTensor/codegen.mlir

 // RUN: mlir-opt %s --sparse-tensor-codegen  --canonicalize --cse | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed" ],
   indexBitWidth = 64,
   pointerBitWidth = 32
 }>
 
-#Dense = #sparse_tensor.encoding<{
+#Dense2D = #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
 }>
 
+#Dense3D = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "dense", "dense" ],
+  indexBitWidth = 64,
+  pointerBitWidth = 32,
+  dimOrdering = affine_map<(i,j,k) -> (k, i,j)>
+}>
+
 // CHECK-LABEL: func @sparse_nop(
 //  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-LABEL: func @sparse_dense_2d(
 //  CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xf64>>)
-func.func @sparse_dense(%arg0: tensor<?x?xf64, #Dense>) {
+func.func @sparse_dense_2d(%arg0: tensor<?x?xf64, #Dense2D>) {
   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
 }
+
+//
+// Just a linearized array in the end. Dim op is statically known.
+//
+// CHECK-LABEL: func @sparse_dense_3d(
+//  CHECK-SAME: %[[A:.*]]: tuple<memref<6000xf64>>) -> index
+//       CHECK: %[[C:.*]] = arith.constant 20 : index
+//       CHECK: return %[[C]] : index
+func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index {
+  %c = arith.constant 1 : index
+  %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D>
+  return %0 : index
+}