Implement the conversion from sparse constant to sparse tensors.

The sparse constant provides a constant tensor in coordinate format. We first split the sparse constant into a constant tensor for indices and a constant tensor for values. We then generate a loop to fill a sparse tensor in coordinate format using the tensors for the indices and the values. Finally, we convert the sparse tensor in coordinate format to the destination sparse tensor format.

Add tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D110373

Author: bixia1

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

@@ -195,16 +195,39 @@ static Value genIsNonzero(ConversionPatternRewriter &rewriter, Location loc,
   llvm_unreachable("Unknown element type");
 }
 
+/// Generates the code to read the value from tensor[ivs], and conditionally
+/// stores the indices ivs to the memory in ind. The generated code looks like
+/// the following and the insertion point after this routine is inside the
+/// if-then branch behind the assignment to ind. This is to ensure that the
+/// addEltX call generated after is inside the if-then branch.
+///    if (tensor[ivs]!=0) {
+///      ind = ivs
+static Value genIndexAndValueForDense(ConversionPatternRewriter &rewriter,
+                                      Operation *op, Type eltType, Value tensor,
+                                      Value ind, ValueRange ivs) {
+  Location loc = op->getLoc();
+  Value val = rewriter.create<tensor::ExtractOp>(loc, tensor, ivs);
+  Value cond = genIsNonzero(rewriter, loc, eltType, val);
+  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, cond, /*else*/ false);
+  rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
+  unsigned i = 0;
+  for (auto iv : ivs) {
+    Value idx = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(i++));
+    rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
+  }
+  return val;
+}
+
 /// Generates a call that adds one element to a coordinate scheme.
 /// In particular, this generates code like the following:
 ///   val = a[i1,..,ik];
 ///   if val != 0
 ///     t->add(val, [i1,..,ik], [p1,..,pk]);
 static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
-                          Value ptr, Value tensor, Value ind, Value perm,
-                          ValueRange ivs) {
+                          Type eltType, Value ptr, Value val, Value ind,
+                          Value perm) {
+  Location loc = op->getLoc();
   StringRef name;
-  Type eltType = tensor.getType().cast<ShapedType>().getElementType();
   if (eltType.isF64())
     name = "addEltF64";
   else if (eltType.isF32())
@@ -219,16 +242,6 @@ static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
     name = "addEltI8";
   else
     llvm_unreachable("Unknown element type");
-  Location loc = op->getLoc();
-  Value val = rewriter.create<tensor::ExtractOp>(loc, tensor, ivs);
-  Value cond = genIsNonzero(rewriter, loc, eltType, val);
-  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, cond, /*else*/ false);
-  rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
-  unsigned i = 0;
-  for (auto iv : ivs) {
-    Value idx = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(i++));
-    rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
-  }
   SmallVector<Value, 8> params;
   params.push_back(ptr);
   params.push_back(val);
@@ -240,6 +253,41 @@ static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
       params);
 }
 
+/// If the tensor is a sparse constant, generates and returns the pair of
+/// the constants for the indices and the values.
+static Optional<std::pair<Value, Value>>
+genSplitSparseConstant(ConversionPatternRewriter &rewriter, ConvertOp op,
+                       Value tensor) {
+  if (auto constOp = tensor.getDefiningOp<ConstantOp>()) {
+    if (auto attr = constOp.value().dyn_cast<SparseElementsAttr>()) {
+      Location loc = op->getLoc();
+      DenseElementsAttr indicesAttr = attr.getIndices();
+      Value indices = rewriter.create<ConstantOp>(loc, indicesAttr);
+      DenseElementsAttr valuesAttr = attr.getValues();
+      Value values = rewriter.create<ConstantOp>(loc, valuesAttr);
+      return std::make_pair(indices, values);
+    }
+  }
+  return {};
+}
+
+/// Generates the code to copy the index at indices[ivs] to ind, and return
+/// the value at value[ivs].
+static Value genIndexAndValueForSparse(ConversionPatternRewriter &rewriter,
+                                       Operation *op, Value indices,
+                                       Value values, Value ind, ValueRange ivs,
+                                       unsigned rank) {
+  Location loc = op->getLoc();
+  for (unsigned i = 0; i < rank; i++) {
+    Value idx = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(i));
+    Value val = rewriter.create<tensor::ExtractOp>(loc, indices,
+                                                   ValueRange{ivs[0], idx});
+    val = rewriter.create<IndexCastOp>(loc, val, rewriter.getIndexType());
+    rewriter.create<memref::StoreOp>(loc, val, ind, idx);
+  }
+  return rewriter.create<tensor::ExtractOp>(loc, values, ivs[0]);
+}
+
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
@@ -330,15 +378,26 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
       // TODO: sparse => dense
       return failure();
     }
-    // This is a dense => sparse conversion, which is handled as follows:
+    // This is a dense => sparse conversion or a sparse constant in COO =>
+    // sparse conversion, which is handled as follows:
     //   t = newSparseCOO()
+    //   ...code to fill the COO tensor t...
+    //   s = newSparseTensor(t)
+    //
+    // To fill the COO tensor from a dense tensor:
     //   for i1 in dim1
     //    ..
     //     for ik in dimk
     //       val = a[i1,..,ik]
     //       if val != 0
     //         t->add(val, [i1,..,ik], [p1,..,pk])
-    //   s = newSparseTensor(t)
+    //
+    // To fill the COO tensor from a sparse constant in COO format:
+    //   for i in range(NNZ)
+    //     val = values[i]
+    //     [i1,..,ik] = indices[i]
+    //     t->add(val, [i1,..,ik], [p1,..,pk])
+    //
     // Note that the dense tensor traversal code is actually implemented
     // using MLIR IR to avoid having to expose too much low-level
     // memref traversal details to the runtime support library.
@@ -351,7 +410,6 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
         MemRefType::get({ShapedType::kDynamicSize}, rewriter.getIndexType());
     Value perm;
     Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
-    Value tensor = adaptor.getOperands()[0];
     Value arg = rewriter.create<ConstantOp>(
         loc, rewriter.getIndexAttr(shape.getRank()));
     Value ind = rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
@@ -360,16 +418,38 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
     SmallVector<Value> st;
     Value zero = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(0));
     Value one = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(1));
-    for (unsigned i = 0, rank = shape.getRank(); i < rank; i++) {
+    Value tensor = adaptor.getOperands()[0];
+    auto indicesValues = genSplitSparseConstant(rewriter, op, tensor);
+    bool isCOOConstant = indicesValues.hasValue();
+    Value indices;
+    Value values;
+    if (isCOOConstant) {
+      indices = indicesValues->first;
+      values = indicesValues->second;
       lo.push_back(zero);
-      hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, tensor, i));
+      hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, values, 0));
       st.push_back(one);
+    } else {
+      for (unsigned i = 0, rank = shape.getRank(); i < rank; i++) {
+        lo.push_back(zero);
+        hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, tensor, i));
+        st.push_back(one);
+      }
     }
+    Type eltType = shape.getElementType();
+    unsigned rank = shape.getRank();
     scf::buildLoopNest(rewriter, op.getLoc(), lo, hi, st, {},
                        [&](OpBuilder &builder, Location loc, ValueRange ivs,
                            ValueRange args) -> scf::ValueVector {
-                         genAddEltCall(rewriter, op, ptr, tensor, ind, perm,
-                                       ivs);
+                         Value val;
+                         if (isCOOConstant)
+                           val = genIndexAndValueForSparse(
+                               rewriter, op, indices, values, ind, ivs, rank);
+                         else
+                           val = genIndexAndValueForDense(rewriter, op, eltType,
+                                                          tensor, ind, ivs);
+                         genAddEltCall(rewriter, op, eltType, ptr, val, ind,
+                                       perm);
                          return {};
                        });
     rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, ptr));

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

@@ -114,8 +114,8 @@ struct SparseTensorConversionPass
     });
     // The following operations and dialects may be introduced by the
     // rewriting rules, and are therefore marked as legal.
-    target.addLegalOp<ConstantOp, tensor::CastOp, tensor::ExtractOp, CmpFOp,
-                      CmpIOp>();
+    target.addLegalOp<ConstantOp, IndexCastOp, tensor::CastOp,
+                      tensor::ExtractOp, CmpFOp, CmpIOp>();
     target.addLegalDialect<scf::SCFDialect, LLVM::LLVMDialect,
                            memref::MemRefDialect>();
     // Populate with rules and apply rewriting rules.

mlir/test/Dialect/SparseTensor/conversion.mlir

@@ -182,6 +182,45 @@ func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix
   return %0 : tensor<2x4xf64, #SparseMatrix>
 }
 
+#CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
+
+// CHECK-LABEL:   func @entry() -> !llvm.ptr<i8> {
+// CHECK:           %[[C1:.*]] = constant 1 : i32
+// CHECK:           %[[Offset:.*]] = constant dense<[0, 1]> : tensor<2xi64>
+// CHECK:           %[[Dims:.*]] = constant dense<[8, 7]> : tensor<2xi64>
+// CHECK:           %[[Base:.*]] = constant dense<[0, 1]> : tensor<2xi8>
+// CHECK:           %[[I2:.*]] = constant 2 : index
+// CHECK:           %[[SparseV:.*]] = constant dense<[1.000000e+00, 5.000000e+00]> : tensor<2xf32>
+// CHECK:           %[[SparseI:.*]] = constant dense<{{\[\[}}0, 0], [1, 6]]> : tensor<2x2xi64>
+// CHECK:           %[[I1:.*]] = constant 1 : index
+// CHECK:           %[[I0:.*]] = constant 0 : index
+// CHECK:           %[[C2:.*]] = constant 2 : i32
+// CHECK:           %[[BaseD:.*]] = tensor.cast %[[Base]] : tensor<2xi8> to tensor<?xi8>
+// CHECK:           %[[DimsD:.*]] = tensor.cast %[[Dims]] : tensor<2xi64> to tensor<?xi64>
+// CHECK:           %[[OffsetD:.*]] = tensor.cast %[[Offset]] : tensor<2xi64> to tensor<?xi64>
+// CHECK:           %[[TCOO:.*]] = call @newSparseTensor(%[[BaseD]], %[[DimsD]], %[[OffsetD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[C2]], %{{.}})
+// CHECK:           %[[Index:.*]] = memref.alloca() : memref<2xindex>
+// CHECK:           %[[IndexD:.*]] = memref.cast %[[Index]] : memref<2xindex> to memref<?xindex>
+// CHECK:           scf.for %[[IV:.*]] = %[[I0]] to %[[I2]] step %[[I1]] {
+// CHECK:             %[[VAL0:.*]] = tensor.extract %[[SparseI]]{{\[}}%[[IV]], %[[I0]]] : tensor<2x2xi64>
+// CHECK:             %[[VAL1:.*]] = index_cast %[[VAL0]] : i64 to index
+// CHECK:             memref.store %[[VAL1]], %[[Index]]{{\[}}%[[I0]]] : memref<2xindex>
+// CHECK:             %[[VAL2:.*]] = tensor.extract %[[SparseI]]{{\[}}%[[IV]], %[[I1]]] : tensor<2x2xi64>
+// CHECK:             %[[VAL3:.*]] = index_cast %[[VAL2]] : i64 to index
+// CHECK:             memref.store %[[VAL3]], %[[Index]]{{\[}}%[[I1]]] : memref<2xindex>
+// CHECK:             %[[VAL4:.*]] = tensor.extract %[[SparseV]]{{\[}}%[[IV]]] : tensor<2xf32>
+// CHECK:             call @addEltF32(%[[TCOO]], %[[VAL4]], %[[IndexD]], %[[OffsetD]])
+// CHECK:           }
+// CHECK:           %[[T:.*]] = call @newSparseTensor(%[[BaseD]], %[[DimsD]], %[[OffsetD]], %{{.*}}, %{{.*}}, %[[C1]], %{{.*}})
+// CHECK:           return %[[T]] : !llvm.ptr<i8>
+func @entry() -> tensor<8x7xf32, #CSR>{
+  // Initialize a tensor.
+  %0 = 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-LABEL: func @sparse_convert_3d(
 //  CHECK-SAME: %[[A:.*]]: tensor<?x?x?xf64>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[C0:.*]] = constant 0 : index

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse-constant_to_sparse_tensor.mlir

@@ -0,0 +1,51 @@
+// RUN: mlir-opt %s \
+// RUN:   --sparsification --sparse-tensor-conversion \
+// RUN:   --convert-vector-to-scf --convert-scf-to-std \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize  \
+// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm --reconcile-unrealized-casts | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#Tensor1  = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed"]
+}>
+
+//
+// Integration tests for conversions from sparse constants to sparse tensors.
+//
+module {
+  func @entry() {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %d0 = constant 0.0 : f64
+
+    // A tensor in COO format.
+    %ti = constant sparse<[[0, 0], [0, 7], [1, 2], [4, 2], [5, 3], [6, 4], [6, 6], [9, 7]],
+                          [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]> : tensor<10x8xf64>
+
+    // Convert the tensor in COO format to a sparse tensor with annotation #Tensor1.
+    %ts = sparse_tensor.convert %ti : tensor<10x8xf64> to tensor<10x8xf64, #Tensor1>
+
+    // CHECK: ( 0, 1, 4, 5, 6, 9 )
+    %i0 = sparse_tensor.indices %ts, %c0 : tensor<10x8xf64, #Tensor1> to memref<?xindex>
+    %i0r = vector.transfer_read %i0[%c0], %c0: memref<?xindex>, vector<6xindex>
+    vector.print %i0r : vector<6xindex>
+
+    // CHECK: ( 0, 7, 2, 2, 3, 4, 6, 7 )
+    %i1 = sparse_tensor.indices %ts, %c1 : tensor<10x8xf64, #Tensor1> to memref<?xindex>
+    %i1r = vector.transfer_read %i1[%c0], %c0: memref<?xindex>, vector<8xindex>
+    vector.print %i1r : vector<8xindex>
+
+    // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8 )
+    %v = sparse_tensor.values %ts : tensor<10x8xf64, #Tensor1> to memref<?xf64>
+    %vr = vector.transfer_read %v[%c0], %d0: memref<?xf64>, vector<8xf64>
+    vector.print %vr : vector<8xf64>
+
+    return
+  }
+}
+