Implement the conversion from sparse constant to sparse tensors.

[bixia]

Summary

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.

[bixia]

Rebased and resolved conflict.

[aartbik]

Oh, this is solid work Bixia! Thank you, I love how we no longer have to iterate over a dense iteration space to initialize from a sparse constant!

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

259

You probably saw the clang-tidy message already, but our style prefers camel case for locals, e.g. indicesAttr or something like that

261

same

274

unsigned

276

unsigned

374

coo -> COO

418

since you use this here and below in the loop gen at L437 can we give this a meaningful name

isCOOConstant = iValues.hasValue();
if (isCOOConstant)

...

[aartbik]

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

207

this method leaves the insertion point in the if val != 0
for subsequent code addition

Probably something to document in the comment at L198
(ie. both the generation of the if as well as the post-condition on the insertion point)

391

maybe use the informal

val = values[i]
[i1,..ik] = indices[i]
t->add(val, [i1,..ik], [p1,..,pk])

notation to make the fact that L391 assign a vector more clear

[bixia]

Addressed review comments.

[aartbik]

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

418

I don't see this change yet?

[bixia]

Used bool isCOOConstant.

[bixia]

Added [mlir][sparse] to PR subject.

[aartbik]

Ship it! Great work.

[bixia]

resolved conflict.

[bixia]

Resolved conflicts.

[bixia]

Fixed PR description.

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

static Value genIsNonzero(ConversionPatternRewriter &rewriter, Location loc,
   Value zero = rewriter.create<ConstantOp>(loc, rewriter.getZeroAttr(t));
   if (t.isa<FloatType>())
     return rewriter.create<CmpFOp>(loc, CmpFPredicate::UNE, v, zero);
   if (t.isIntOrIndex())
     return rewriter.create<CmpIOp>(loc, CmpIPredicate::ne, v, zero);
   llvm_unreachable("Unknown element type");
 }
 
+/// Generates the code to read the value of the tensor[ivs], and store the
+/// indices ivs to the memory in ind. Returns the value.
+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())
     name = "addEltF32";
   else if (eltType.isInteger(64))
     name = "addEltI64";
   else if (eltType.isInteger(32))
     name = "addEltI32";
   else if (eltType.isInteger(16))
     name = "addEltI16";
   else if (eltType.isInteger(8))
     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);
   params.push_back(ind);
   params.push_back(perm);
   Type pTp = LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
   rewriter.create<CallOp>(
       loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true),
       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 indices_attr = attr.getIndices();
+      Value indices = rewriter.create<ConstantOp>(loc, indices_attr);
+      DenseElementsAttr values_attr = attr.getValues();
+      Value values = rewriter.create<ConstantOp>(loc, values_attr);
+      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,
+                                       int rank) {
+  Location loc = op->getLoc();
+  for (int 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.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<ReturnOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
if (encDst && encSrc) {
       Value coo = genNewCall(rewriter, op, encDst, 3, perm, operands[0]);
       rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, coo));
       return success();
     }
     if (!encDst || encSrc) {
       // 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]
+    //     ind = indices[i]
+    //     t->add(val, ind, [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.
     // Also note that the code below only generates the "new" ops and
     // the loop-nest per se; whereas the entire body of the innermost
     // loop is generated by genAddElt().
     Location loc = op->getLoc();
     ShapedType shape = resType.cast<ShapedType>();
     auto memTp =
         MemRefType::get({ShapedType::kDynamicSize}, rewriter.getIndexType());
     Value perm;
     Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
-    Value tensor = operands[0];
     Value arg = rewriter.create<ConstantOp>(
         loc, rewriter.getIndexAttr(shape.getRank()));
     Value ind = rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
     SmallVector<Value> lo;
     SmallVector<Value> hi;
     SmallVector<Value> st;
     Value zero = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(0));
     Value one = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(1));
+    Value tensor = operands[0];
+    auto indices_values = genSplitSparseConstant(rewriter, op, tensor);
+    Value indices;
+    Value values;
+    if (indices_values.hasValue()) {
+      indices = indices_values->first;
+      values = indices_values->second;
+      lo.push_back(zero);
+      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();
+    int 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 (indices_values.hasValue())
+                           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));
     return success();
   }
 };
 
 /// Sparse conversion rule for pointer accesses.

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

void runOnOperation() override {
     });
     target.addDynamicallyLegalOp<ReturnOp>(
         [&](ReturnOp op) { return converter.isLegal(op.getOperandTypes()); });
     target.addDynamicallyLegalOp<tensor::DimOp>([&](tensor::DimOp op) {
       return converter.isLegal(op.getOperandTypes());
     });
     // 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.
     populateFuncOpTypeConversionPattern(patterns, converter);
     populateCallOpTypeConversionPattern(patterns, converter);
     populateSparseTensorConversionPatterns(converter, patterns);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))

mlir/test/Dialect/SparseTensor/conversion.mlir

 //       CHECK: }
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix> {
   %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to 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
 //   CHECK-DAG: %[[C1:.*]] = constant 1 : index
 //   CHECK-DAG: %[[C2:.*]] = constant 2 : index
 //   CHECK-DAG: %[[U:.*]] = constant dense<[0, 1, 1]> : tensor<3xi8>
 //   CHECK-DAG: %[[V:.*]] = constant dense<0> : tensor<3xi64>
 //   CHECK-DAG: %[[W:.*]] = constant dense<[1, 2, 0]> : tensor<3xi64>

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

+// 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
+  }
+}
+