[mlir][sparse] support singleton in loop emitter.

[Peiming]

fix memory leak

[Peiming]

fix warning

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 //===----------------------------------------------------------------------===//
 
 // TODO: Sparsification should also rely on this class to generate loops.
 class SparseTensorLoopEmitter {
 public:
   /// Constructor: take an array of tensors inputs, on which the generated loops
   /// will iterate on. The index of the tensor in the array is also the
   /// tensor id (tid) used in related functions.
-  explicit SparseTensorLoopEmitter(ValueRange tensors);
+  explicit SparseTensorLoopEmitter(ValueRange tensors,
+                                   bool isLastOutput = false);
 
   ///
   /// Core functions.
   ///
 
   /// Starts a loop emitting session:
   /// 1. Generates all the buffers needed to iterate tensors.
   /// 2. Generates the lo/hi bounds to iterate tensors[0].
private:
   std::vector<std::vector<Value>> highs;
   /// Universal dense indices and upper bounds (by index). The sizes array is
   /// set once with the inferred dimension sizes.
   std::vector<std::vector<Value>> sizes;
   std::vector<std::vector<Value>> ptrBuffer; // to_pointers
   std::vector<std::vector<Value>> idxBuffer; // to_indices
   std::vector<Value> valBuffer;              // to_value
 
+  bool isLastOutput; // Is the last tensor output tensor
   std::vector<LoopLevelInfo> loopStack;
   // TODO: not yet used, it should track the current level for each tensor
   // to help eliminate `dim` paramters from above APIs.
   std::vector<size_t> curLv;
 };
 
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.

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

static Value genIndexLoad(OpBuilder &builder, Location loc, Value ptr,
   }
   return load;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
-SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors)
+SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
+                                                 bool isLastOutput)
     : tensors(tensors.begin(), tensors.end()), dims(tensors.size()),
       pidxs(tensors.size()), coord(tensors.size()), highs(tensors.size()),
       sizes(tensors.size()), ptrBuffer(tensors.size()),
-      idxBuffer(tensors.size()), valBuffer(tensors.size()), loopStack(),
-      curLv(tensors.size(), 0) {
+      idxBuffer(tensors.size()), valBuffer(tensors.size()),
+      isLastOutput(isLastOutput), loopStack(), curLv(tensors.size(), 0) {
   for (size_t i = 0, e = tensors.size(); i < e; i++) {
     auto t = tensors[i];
-    auto rtp = t.getType().cast<RankedTensorType>();
+    auto rtp = t.getType().dyn_cast<RankedTensorType>();
+    if (!rtp) // a scalar (0-dimension tensors)
+      continue;
+
     auto rank = static_cast<size_t>(rtp.getRank());
     auto enc = getSparseTensorEncoding(rtp);
     if (enc)
       for (auto dimTp : enc.getDimLevelType())
         dims[i].push_back(dimTp);
     else
       dims[i].assign(rank, SparseTensorEncodingAttr::DimLevelType::Dense);
 
for (int64_t d = 0; d < rank; d++) {
             MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         // Generate sparse primitives to obtains pointer and indices.
         ptrBuffer[t][d] = builder.create<ToPointersOp>(loc, ptrTp, tensor, dim);
         idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
       } else if (isSingletonDim(dims[t][d])) {
-        llvm_unreachable("TODO: not implemented yet");
+        // Singleton dimension, fetch indices.
+        auto indTp =
+            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
+        auto dim = builder.getIndexAttr(d);
+        idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
+      } else {
+        // Dense dimension, nothing to fetch.
+        assert(isDenseDim(dims[t][d]));
       }
 
       // Find upper bound in current dimension.
       unsigned p = toOrigDim(enc, d);
       Value up = mlir::linalg::createOrFoldDimOp(builder, loc, tensor, p);
       sizes[t][d] = highs[t][d] = up;
     }
     // Perform the required bufferization. Dense inputs materialize
     // from the input tensors. Dense outputs need special handling.
     // Sparse inputs use sparse primitives to obtain the values.
     Type elementType = rtp.getElementType();
 
     if (!enc) {
       // Non-annotated dense tensors.
       auto denseTp = MemRefType::get(shape, elementType);
-      // This is not the output tensor
+      if (isLastOutput && t == tensors.size() - 1)
+        llvm_unreachable("TODO: not yet handled");
+      else
         valBuffer[t] =
             builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
     } else {
       // Annotated sparse tensors.
       auto dynShape = {ShapedType::kDynamicSize};
       auto sparseTp = MemRefType::get(dynShape, elementType);
       valBuffer[t] = builder.create<ToValuesOp>(loc, sparseTp, tensor);
     }
     // Prepare to enter the first dim for all (input) tensors
     prepareLoopOverTensorAtDim(builder, loc, t, 0);
   }
 }
 
 Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
     OpBuilder &builder, Location loc, size_t tid, size_t dim,
     ArrayRef<Value> reduc) {
   assert(dims[tid].size() > dim);
   // We can not re-enter the same level.
   assert(!coord[tid][dim]);
   Value step = constantIndex(builder, loc, 1);
-  bool isCompressed = isCompressedDim(dims[tid][dim]);
-  assert(isDenseDim(dims[tid][dim]) || isCompressedDim(dims[tid][dim]));
+  auto dimType = dims[tid][dim];
+  bool isSparse = isCompressedDim(dimType) || isSingletonDim(dimType);
+  assert(isDenseDim(dimType) || isCompressedDim(dimType) ||
+         isSingletonDim(dimType));
 
-  Value lo = isCompressed ? pidxs[tid][dim] : constantIndex(builder, loc, 0);
+  Value lo = isSparse ? pidxs[tid][dim] : constantIndex(builder, loc, 0);
   Value hi = highs[tid][dim];
 
   // TODO: support reduction.
   if (!reduc.empty())
     llvm_unreachable("TODO: not implemented yet");
 
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
   builder.setInsertionPointToStart(forOp.getBody());
   Value iv = forOp.getInductionVar();
   Operation *loop = forOp;
 
   assert(iv);
-  if (isCompressed) {
+  if (isSparse) {
     pidxs[tid][dim] = iv;
     // Generating a load on the indices array yields the coordinate.
     Value ptr = idxBuffer[tid][dim];
     // TODO: generates load for vector value.
     coord[tid][dim] = genIndexLoad(builder, loc, ptr, iv);
   } else {
     // Dense tensor, the coordinates is the inducation variable.
     coord[tid][dim] = iv;

 
 bool SparseTensorLoopEmitter::prepareLoopOverTensorAtDim(OpBuilder &builder,
                                                          Location loc,
                                                          size_t tid,
                                                          size_t dim) {
   // TODO: generate loop iteration on output tensor based on the shape
   // instead of pointer/indices arrays.
   assert(dims[tid].size() > dim);
+  auto dimType = dims[tid][dim];
 
-  if (isDenseDim(dims[tid][dim]))
+  if (isDenseDim(dimType))
     return false;
 
   // Either the first dimension, or the previous dimension has been set.
   assert(dim == 0 || pidxs[tid][dim - 1]);
-  if (isCompressedDim(dims[tid][dim])) {
-    Value ptr = ptrBuffer[tid][dim];
-    Value c1 = constantIndex(builder, loc, 1);
-    Value pLo = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
+  Value c0 = constantIndex(builder, loc, 0);
+  if (isCompressedDim(dimType)) {
+    Value ptr = ptrBuffer[tid][dim];
+
+    Value pLo = dim == 0 ? c0 : pidxs[tid][dim - 1];
     Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
 
     pidxs[tid][dim] = genIndexLoad(builder, loc, ptr, pLo);
     highs[tid][dim] = genIndexLoad(builder, loc, ptr, pHi);
-
     return true;
   }
+  if (isSingletonDim(dimType)) {
+    Value pLo = dim == 0 ? c0 : pidxs[tid][dim - 1];
+    Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
 
-  if (isSingletonDim(dims[tid][dim]))
-    llvm_unreachable("TODO: not implemented yet");
+    pidxs[tid][dim] = pLo;
+    highs[tid][dim] = pHi;
+    return true;
+  }
 
   llvm_unreachable("Unrecognizable dimesion type!");
 }
 
 Value SparseTensorLoopEmitter::emitExtraLocalsForTensorsAtDims(
     OpBuilder &builder, Location loc, size_t tid, size_t dim) {
   llvm_unreachable("TODO: not implemented yet");
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_codegen_foreach.mlir

#CSR = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ]
 }>
 
 #DCSC = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (j,i)>
 }>
 
+#SortedCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ]
+}>
+
+#SortedCOOPerm = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
 module {
 
   /// uses foreach operator to print coords and values.
   func.func @foreach_print_1(%arg0: tensor<2x2xf64, #Row>) {
     sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #Row> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #DCSC> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
      }
      return
   }
 
+  func.func @foreach_print_4(%arg0: tensor<2x2xf64, #SortedCOO>) {
+    sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #SortedCOO> do {
+      ^bb0(%1: index, %2: index, %v: f64) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %v: f64
+     }
+     return
+  }
+
+  func.func @foreach_print_5(%arg0: tensor<2x2xf64, #SortedCOOPerm>) {
+    sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #SortedCOOPerm> do {
+      ^bb0(%1: index, %2: index, %v: f64) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %v: f64
+     }
+     return
+  }
+
   //
   // Main driver.
   //
   func.func @entry() {
     //
     // Initialize a 3-dim dense tensor.
     //
     %src = arith.constant dense<
        [[  1.0,  2.0],
         [  5.0,  6.0]]
     > : tensor<2x2xf64>
 
     //
     // Convert dense tensor directly to various sparse tensors.
     //
     %s1 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Row>
     %s2 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #CSR>
     %s3 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #DCSC>
+    %s4 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOO>
+    %s5 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOOPerm>
     // CHECK: 0
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
func.func @entry() {
     // CHECK-NEXT: 5
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1
     // CHECK-NEXT: 1
     // CHECK-NEXT: 6
     call @foreach_print_3(%s3) : (tensor<2x2xf64, #DCSC>) -> ()
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
+    call @foreach_print_4(%s4) : (tensor<2x2xf64, #SortedCOO>) -> ()
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
+    call @foreach_print_5(%s5) : (tensor<2x2xf64, #SortedCOOPerm>) -> ()
     
     bufferization.dealloc_tensor %s1 : tensor<2x2xf64, #Row>
     bufferization.dealloc_tensor %s2 : tensor<2x2xf64, #CSR>
     bufferization.dealloc_tensor %s3 : tensor<2x2xf64, #DCSC>
+    bufferization.dealloc_tensor %s4 : tensor<2x2xf64, #SortedCOO>
+    bufferization.dealloc_tensor %s5 : tensor<2x2xf64, #SortedCOOPerm>
 
     return
   }
 }