[mlir][sparse] Improve the rewriting for dense-to-annotated-dense convert operator.

[bixia]

Summary

Previously, we use InsertOp to write the destination tensor. We now directly
store to the destination tensor values buffer. For the case where the source is
a dense tensor and the destination has an identity ordering, we use memref
CopyOp to write the destination.

[Peiming]

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

805

I really feel that we should not allow all dense sparse encoding right now... we should canonicalize it to non-encoded tensor instead. This brings too much unnecessary complexity here. @aartbik WDYT?

820

There if/else here and below are too hard to follow, if we are going to allow all dense encoding, we should still probably split this function up.

e.g., genDirectConvert and genIndirectConvert or some other more meaningful methods.

[bixia]

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

805

We currently distinguish tensors through these methods:
(1) sparse tensor pass won't handle, not encoded
(2) sparse tensor pass handle, but essentially not sparse, annotated all dense
(3) sparse tensor pass handle and real sparse

[bixia]

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

820

Is genDirectConvert for conversion without a tmpCOO and genIndirectConvert for conversion with a tmpCOO?
Or is genDirectConvert for conversion to annotated all dense?

[Peiming]

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

805

yeah, I understand.

But all dense encoded tensor is essentially the same as non-encoded tensor, right? Seems that we are treating allDense as a special case everywhere, which could be avoided.

820

maybe with/without tmpCOO? WDYT? If you want, you can further split it up.

[aartbik]

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

805

The all-dense encoding is really just something that is helpful to debug the sparse codegen. I don't expect clients to ever favor it over truly dense tensors. As such, is the amount of special cases worth it? Can you say a bit more on what it improves?

[wrengr]

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

805

For source programs, I think it'd be good to canonicalize all-dense "sparse"-tensors to standard dense-tensors. However, internally, there are a few places where we currently rely on the fact that all-dense encodings are in fact allowed. One example is in the runtime pass's dense2sparse conversionOp (or was it in the sparse2dense?). So if we do want to go the canonicalization route, then it's important to distinguish (1) "all-dense tensors", i.e., tensor values of type RTT with all-dense STEA, (2) the type "RTT with all-dense STEA", (3) the attribute "all-dense STEA". Since just because we canonicalize the first of those three doesn't mean we necessarily canonicalize or invalidate either of the other two. And if we also canonicalize the second one, that still doesn't necessarily mean canonicalizing nor invalidating the third.

Also N.B., in D141975 we have code for explicitly going the other direction: allowing us to treat RTT-without-encoding in the same way that we treat RTT-with-STEA. I did that in order to remove a bunch of other special cases in the code. So it's also important to see whether that's sufficient to remove a bunch of the special cases you have in mind, and to ensure that any canonicalization effort doesn't end up introducing special cases somewhere else instead.

[bixia]

Split out the handling of annotated all dense result.

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

805

I edit the PR description to add more detail on what the PR does.

820

Split out the dense case.

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

LogicalResult matchAndRewrite(ConvertOp op,
     if (!encSrc && encDst)
       return dense2SparseRewrite(op, rewriter);
 
     // Dense-to-dense convert is a nop and handled by canonicalization.
     return failure();
   }
 
 private:
+  // Handles sparse constant or dense tensor to annotate all dense sparse tensor
+  // conversion as follows:
+  // If the source is a dense tensor:
+  //   if (destination has an identity ordering)
+  //     copy src to dst_tensor_values
+  //   else for i1 in dim1
+  //    ..
+  //     for ik in dimk
+  //       val = src[i1,..,ik]
+  //       if val != 0
+  //         dst_tensor_values[i1,..,ik] = val
+  //
+  // If the source is a sparse constant in COO format:
+  //   for i in range(NNZ)
+  //     val = values[i]
+  //     [i1,..,ik] = indices[i]
+  //     dst_tensor_values[i1,..,ik] = val
+  LogicalResult dense2AnnotateAllDenseRewrite(ConvertOp op,
+                                              PatternRewriter &rewriter) const {
+    Location loc = op.getLoc();
+    Value src = op.getSource();
+    auto dstTp = getRankedTensorType(op);
+    SmallVector<Value> sizes;
+    sizesFromSrc(rewriter, sizes, loc, src);
+    SmallVector<Value> dynSizes;
+    getDynamicSizes(dstTp, sizes, dynSizes);
+
+    bool fromSparseConst = false;
+    if (auto constOp = op.getSource().getDefiningOp<arith::ConstantOp>()) {
+      if (constOp.getValue().dyn_cast<SparseElementsAttr>()) {
+        fromSparseConst = true;
+      }
+    }
+
+    SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
+    assert(encDst.isAllDense());
+
+    Value tensor =
+        rewriter.create<AllocTensorOp>(loc, dstTp, dynSizes).getResult();
+    int64_t rank = dstTp.getRank();
+    // Create a view of the values buffer to match the unannotated dense
+    // tensor.
+    Value valuesBuffer = genToValues(rewriter, loc, tensor);
+    Value idxBuffer = genAlloca(rewriter, loc, rank, rewriter.getIndexType(),
+                                /*staticShape=*/true);
+    Value dst = reshapeValuesToLevels(rewriter, loc, encDst, sizes,
+                                      valuesBuffer, idxBuffer);
+
+    if (!fromSparseConst && encDst.hasIdDimOrdering()) {
+      auto srcTp = MemRefType::get(
+          op.getSource().getType().cast<RankedTensorType>().getShape(),
+          dstTp.getElementType());
+      src = rewriter.create<bufferization::ToMemrefOp>(loc, srcTp, src);
+      rewriter.create<memref::CopyOp>(loc, src, dst);
+      dst = rewriter.create<ConvertOp>(loc, dstTp, tensor).getResult();
+      rewriter.replaceOp(op, dst);
+      return success();
+    }
+
+    rewriter.create<ForeachOp>(
+        loc, src, std::nullopt,
+        [&](OpBuilder &builder, Location loc, ValueRange indices, Value v,
+            ValueRange) {
+          uint64_t rank = dstTp.getRank();
+          SmallVector<Value> indicesArray(rank, Value());
+          for (uint64_t i = 0; i < rank; i++)
+            indicesArray[toStoredDim(encDst, i)] = indices[i];
+          if (fromSparseConst) {
+            builder.create<memref::StoreOp>(loc, v, dst, indicesArray);
+          } else {
+            Value cond = genIsNonzero(builder, loc, v);
+            auto ifOp = builder.create<scf::IfOp>(loc, TypeRange(), cond,
+                                                  /*else=*/false);
+            builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+            builder.create<memref::StoreOp>(loc, v, dst, indicesArray);
+            builder.setInsertionPointAfter(ifOp);
+          }
+          builder.create<sparse_tensor::YieldOp>(loc);
+        });
+    rewriter.setInsertionPointAfter(op);
+    dst = rewriter.create<ConvertOp>(loc, dstTp, tensor).getResult();
+    rewriter.replaceOp(op, dst);
+
+    return success();
+  }
+
   // Handles sparse constant to sparse tensor or dense tensor to sparse tensor
   // conversion as follows:
   //   t = new sparse COO tensor
   //   fill t using src
   //   dst = convert t
   //
   // To fill the COO tensor from a dense tensor:
   //   for i1 in dim1
LogicalResult dense2SparseRewrite(ConvertOp op,
     bool fromSparseConst = false;
     if (auto constOp = op.getSource().getDefiningOp<arith::ConstantOp>()) {
       if (constOp.getValue().dyn_cast<SparseElementsAttr>()) {
         fromSparseConst = true;
       }
     }
 
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
+    if (encDst.isAllDense())
+      return dense2AnnotateAllDenseRewrite(op, rewriter);
+
     // We don't need a temporary COO tensor if the destination has an identity
     // ordering. Otherwise, we use the destination ordering for the temporary
     // COO tensor.
     // TODO: enhance foreachOp to take ordering to remove the need of a
     // temporary COO tensor here.
     RankedTensorType bufferTp = encDst.hasIdDimOrdering()
                                     ? dstTp
                                     : getUnorderedCOOFromTypeWithOrdering(
                                           dstTp, encDst.getDimOrdering());
     auto buffer =
         rewriter.create<AllocTensorOp>(loc, bufferTp, dynSizes).getResult();
     auto foreachOp = rewriter.create<ForeachOp>(
         loc, src, buffer,
         [&](OpBuilder &builder, Location loc, ValueRange indices, Value v,
             ValueRange reduc) {
           Value input = reduc.front();
           uint64_t rank = dstTp.getRank();
           SmallVector<Value> indicesArray(rank, Value());
           for (uint64_t i = 0; i < rank; i++)
             indicesArray[toStoredDim(encDst, i)] = indices[i];
           if (fromSparseConst) {
             input = builder.create<InsertOp>(loc, v, input, indicesArray);
           } else {
             Value cond = genIsNonzero(builder, loc, v);
             auto ifOp = builder.create<scf::IfOp>(
                 loc, TypeRange(input.getType()), cond, /*else*/ true);

mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir

 // RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
 // RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-foreach=false" \
 // RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
+#DD = #sparse_tensor.encoding<{
+  dimLevelType = ["dense", "dense"]
+}>
+
 #CSC = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   dimOrdering = affine_map<(i, j) -> (j, i)>
 }>
 
 #SparseTensor = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed", "compressed"],
   dimOrdering = affine_map<(i,j,k) -> (k,i,j)>

 //       CHECK-RWT:     %[[R:.*]] = sparse_tensor.load %[[T2]] hasInserts
 //       CHECK-RWT:     return %[[R]]
 //       CHECK-RWT:   }
 func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {
   %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #CSR>
   return %0 : tensor<2x4xf64, #CSR>
 }
 
+// CHECK-RWT-LABEL:   func.func @sparse_convert_2d_dd(
+//  CHECK-RWT-SAME:     %[[T0:.*]]: tensor<2x4xf64>) -> tensor<2x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> {
+//   CHECK-RWT-DAG:     %[[TMP_c2:.*]] = arith.constant 2 : index
+//   CHECK-RWT-DAG:     %[[TMP_c4:.*]] = arith.constant 4 : index
+//   CHECK-RWT-DAG:     %[[TMP_c0:.*]] = arith.constant 0 : index
+//   CHECK-RWT-DAG:     %[[TMP_c1:.*]] = arith.constant 1 : index
+//       CHECK-RWT:     %[[T1:.*]] = bufferization.alloc_tensor()
+//       CHECK-RWT:     %[[VAL_0:.*]] = sparse_tensor.values %[[T1]]
+//       CHECK-RWT:     %[[DIM_0:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK-RWT:     memref.store %[[TMP_c2]], %[[DIM_0]][%[[TMP_c0]]]
+//       CHECK-RWT:     memref.store %[[TMP_c4]], %[[DIM_0]][%[[TMP_c1]]]
+//       CHECK-RWT:     %[[VAL_1:.*]] = memref.reshape %[[VAL_0]](%[[DIM_0]])
+//       CHECK-RWT:     %[[VAL_2:.*]] = bufferization.to_memref %[[T0]]
+//       CHECK-RWT:     memref.copy %[[VAL_2]], %[[VAL_1]]
+//       CHECK-RWT:     %[[R:.*]] = sparse_tensor.convert %[[T1]]
+//       CHECK-RWT:     return %[[R]]
+//       CHECK-RWT:   }
+func.func @sparse_convert_2d_dd(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #DD> {
+  %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #DD>
+  return %0 : tensor<2x4xf64, #DD>
+}
+
 // CHECK-LABEL: func @sparse_constant() -> !llvm.ptr<i8> {
 //   CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32
 //   CHECK-DAG: %[[FromCOO:.*]] = arith.constant 2 : i32
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
 //   CHECK-DAG: %[[LvlTypes:.*]] = memref.alloca() : memref<2xi8>
 //   CHECK-DAG: %[[DimSizes:.*]] = memref.alloca() : memref<2xindex>

 func.func @sparse_constant() -> tensor<8x7xf32, #CSR>{
   // Initialize a tensor.
   %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
   // Convert the tensor to a sparse tensor.
   %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSR>
   return %1 : tensor<8x7xf32, #CSR>
 }
 
+// CHECK-RWT-LABEL:   func.func @sparse_constant_dd() -> tensor<8x7xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> {
+//   CHECK-RWT-DAG:     %[[TMP_c8:.*]] = arith.constant 8 : index
+//   CHECK-RWT-DAG:     %[[TMP_c7:.*]] = arith.constant 7 : index
+//   CHECK-RWT-DAG:     %[[TMP_c0:.*]] = arith.constant 0 : index
+//   CHECK-RWT-DAG:     %[[TMP_c1:.*]] = arith.constant 1 : index
+//       CHECK-RWT:     %[[F0:.*]] = arith.constant sparse<{{\[\[}}0, 0], [1, 6]], [1.000000e+00, 5.000000e+00]> : tensor<8x7xf32>
+//       CHECK-RWT:     %[[T0:.*]] = bufferization.alloc_tensor()
+//       CHECK-RWT:     %[[VAL_0:.*]] = sparse_tensor.values %[[T0]]
+//       CHECK-RWT:     %[[DIM_0:.*]] = memref.alloca()
+//       CHECK-RWT:     memref.store %[[TMP_c8]], %[[DIM_0]][%[[TMP_c0]]]
+//       CHECK-RWT:     memref.store %[[TMP_c7]], %[[DIM_0]][%[[TMP_c1]]]
+//       CHECK-RWT:     %[[VAL_1:.*]] = memref.reshape %[[VAL_0]](%[[DIM_0]])
+//       CHECK-RWT:     sparse_tensor.foreach in %[[F0]]
+//       CHECK-RWT:     ^bb0(%[[L0I0:.*]]: index, %[[L0I1:.*]]: index, %[[L0V:.*]]: f32):
+//       CHECK-RWT:       memref.store %[[L0V]], %[[VAL_1]]{{\[}}%[[L0I0]], %[[L0I1]]]
+//       CHECK-RWT:     }
+//       CHECK-RWT:     %[[R:.*]] = sparse_tensor.convert %[[T0]]
+//       CHECK-RWT:     return %[[R]]
+//       CHECK-RWT:   }
+func.func @sparse_constant_dd() -> tensor<8x7xf32, #DD>{
+  // Initialize a tensor.
+  %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
+  // Convert the tensor to a sparse tensor.
+  %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #DD>
+  return %1 : tensor<8x7xf32, #DD>
+}
+
 // CHECK-RWT-LABEL:   func.func @sparse_constant_csc() -> tensor<8x7xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> {
 //       CHECK-RWT:     %[[F0:.*]] = arith.constant sparse<{{\[\[}}0, 0], [1, 6]], [1.000000e+00, 5.000000e+00]> : tensor<8x7xf32>
 //       CHECK-RWT:     %[[T0:.*]] = bufferization.alloc_tensor()
 //       CHECK-RWT:     %[[T1:.*]] = sparse_tensor.foreach in %[[F0]] init(%[[T0]])
 //       CHECK-RWT:     ^bb0(%[[L0I0:.*]]: index, %[[L0I1:.*]]: index, %[[L0V:.*]]: f32, %[[L0T:.*]]: tensor
 //       CHECK-RWT:       %[[L0T2:.*]] = sparse_tensor.insert %[[L0V]] into %[[L0T]]{{\[}}%[[L0I1]], %[[L0I0]]]
 //       CHECK-RWT:       sparse_tensor.yield %[[L0T2]]
 //       CHECK-RWT:     }