[mlir][sparse] support coiteration with fused reshape tensor

[Peiming]

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

567

Surprised to know that structured bindings cannot be captured by lambda expressions until C++20.
https://en.cppreference.com/w/cpp/language/structured_binding

[aartbik]

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

136

Should this be // Generate
or are you focusing on that it yields....?

567

but lvl is okay? strange

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

193

Add a /// doc

[Peiming]

address comments.

[Peiming]

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

567

lvl is not captured..

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

struct LoopLevelInfo {
     Block *const userCodeBlock; // the block holding users' generated code.
     const Value iv;             // the induction variable for the loop
   };
 
   /// Linearizes address for dense dimension (i.e., p = (i * d0) + j).
   Value genAddress(OpBuilder &builder, Location loc, size_t tid, size_t dim,
                    Value iv);
 
+  /// Generates instructions to compute the coordinate of tesnors[tid] on `l`
+  /// under the current loop context.
+  Value genSparseCoord(OpBuilder &builder, Location loc, size_t tid, size_t l);
+
   bool isOutputTensor(size_t tid) {
     return hasOutput && tid == tensors.size() - 1;
   }
 
   bool isSparseOutput(size_t tid) { return isOutputTensor(tid) && isSparseOut; }
 
   /// Setups [lo, hi] for iterating tensor[dim], it assumes that tensor[0
   /// ...dims-1] has already been setup.

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

Value LoopEmitter::genAddress(OpBuilder &builder, Location loc, size_t tid,
   if (isSparseSlices[tid]) {
     auto enc = getSparseTensorEncoding(tensors[tid].getType());
     iv = toSliceCoord(builder, loc, iv, enc, dim);
   }
   Value add = builder.create<arith::AddIOp>(loc, mul, iv);
   return add;
 }
 
+Value LoopEmitter::genSparseCoord(OpBuilder &builder, Location loc, size_t tid,
+                                  size_t l) {
+  Value c = constantIndex(builder, loc, 0);
+  auto reass = getCollapseReassociation(tid, l);
+  for (unsigned i = 0; i < reass.size(); i++) {
+    auto lvl = reass[i];
+    // A load on the indices array yields the coordinate.
+    Value ptr = idxBuffer[tid][lvl];
+    Value off = genIndexLoad(builder, loc, ptr, pidxs[tid][l]);
+    // Linearized the coordinates within the same collapse reassociation.
+    c = builder.create<arith::AddIOp>(loc, c, off);
+    if (i != reass.size() - 1) {
+      c = builder.create<arith::MulIOp>(loc, c,
+                                        this->lvlSizes[tid][reass[i + 1]]);
+    }
+  }
+  return c;
+}
+
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                          bool isSparseOut, ArrayRef<unsigned> topSort) {
   initialize(tensors, loopTag, hasOutput, isSparseOut, topSort);
 }
 
 void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
                              bool isSparseOut, ArrayRef<unsigned> topSort) {
   // First initializes fields.
for (int i = 0, e = reduc.size(); i < e; i++)
       reduc[i] = forOp.getRegionIterArg(i);
     loop = forOp;
   }
   assert(loop && iv);
 
   Value c;
   if (isSparseInput) {
     assert(reass.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
-    c = constantIndex(builder, loc, 0);
-    for (unsigned i = 0; i < reass.size(); i++) {
-      auto lvl = reass[i];
-      // For COO, the pidxs are always the same across consecutive levels.
-      pidxs[tid][lvl] = iv;
-      // Generating a load on the indices array yields the coordinate.
-      Value ptr = idxBuffer[tid][lvl];
-      Value off = genIndexLoad(builder, loc, ptr, iv);
-      c = builder.create<arith::AddIOp>(loc, c, off);
-      if (i != reass.size() - 1) {
-        c = builder.create<arith::MulIOp>(loc, c,
-                                          this->lvlSizes[tid][reass[i + 1]]);
-      }
-    }
+    // For COO, the position is the same across consecutive levels.
+    llvm::for_each(reass, [this, tid, iv](int lvl) { pidxs[tid][lvl] = iv; });
+    c = genSparseCoord(builder, loc, tid, dim);
   } else {
     // Dense tensor, the coordinates is the inducation variable.
     c = iv;
   }
 
   if (isSparseSlices[tid] && isSparseInput) {
     // For sparse level slices, we need to filter out invalid coordinates that
     // are not included in the slice.
Operation *LoopEmitter::enterCoIterationOverTensorsAtDims(
   Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
   Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
-  for (auto [tid, dim] : llvm::zip(tids, dims)) {
-    if (isCompressedDLT(dimTypes[tid][dim]) ||
-        isSingletonDLT(dimTypes[tid][dim])) {
+  for (auto [t, lvl] : llvm::zip(tids, dims)) {
+    unsigned tid = t; // Why `t` can not be captured by lambda?
+    if (isCompressedDLT(dimTypes[tid][lvl]) ||
+        isSingletonDLT(dimTypes[tid][lvl])) {
       Value op1 = before->getArgument(o);
-      Value op2 = highs[tid][dim];
+      Value op2 = highs[tid][lvl];
       Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
                                                 op1, op2);
       cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
-      // Update
-      pidxs[tid][dim] = after->getArgument(o++);
+      // Update positions
+      Value pos = after->getArgument(o++);
+      auto reass = getCollapseReassociation(tid, lvl);
+      assert(reass.size() == 1 || isUniqueCOOType(tensors[tid].getType()));
+      // For COO, the position is the same across consecutive levels.
+      llvm::for_each(reass,
+                     [this, tid, pos](int lvl) { pidxs[tid][lvl] = pos; });
     }
   }
   builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
 
   // Generates while body.
   builder.setInsertionPointToStart(&whileOp.getAfter().front());
 
   SmallVector<std::pair<Value, unsigned>> slicesPreds;
   unsigned i = 0;
   for (auto [tid, dim] : llvm::zip(tids, dims)) {
     // Prepares for next level.
     if (isCompressedDLT(dimTypes[tid][dim]) ||
         isSingletonDLT(dimTypes[tid][dim])) {
-      Value ptr = idxBuffer[tid][dim];
-      Value s = pidxs[tid][dim];
-      Value load = genIndexLoad(builder, loc, ptr, s);
-      coord[tid][dim] = load;
+      coord[tid][dim] = genSparseCoord(builder, loc, tid, dim);
       if (isSparseSlices[tid]) {
+        Value load =
+            genIndexLoad(builder, loc, idxBuffer[tid][dim], pidxs[tid][dim]);
         auto enc = getSparseTensorEncoding(tensors[tid].getType());
         auto [trans, pred] =
             genSliceLegitPredicate(builder, loc, load, enc, dim);
         slicesPreds.emplace_back(pred, i);
         // Updates to the relative coordinate to the slice.
         coord[tid][dim] = trans;
       }
       i++;

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

 
 #COO_2D = #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ], pointerBitWidth = 32, indexBitWidth = 32 }>
 #COO_3D = #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton-nu", "singleton" ], pointerBitWidth = 32, indexBitWidth = 32 }>
 
 module {
   func.func private @printMemref3dF32(%ptr : tensor<?x?x?xf32>) attributes { llvm.emit_c_interface }
   func.func private @printMemref2dF32(%ptr : tensor<?x?xf32>) attributes { llvm.emit_c_interface }
 
-  func.func @test_sparse(%arg0: tensor<5x6xf32>, %arg1: tensor<6x2x3xf32, #COO_3D>) -> tensor<?x?x?xf32> {
+  func.func @test_sparse_rhs(%arg0: tensor<5x6xf32>, %arg1: tensor<6x2x3xf32, #COO_3D>) -> tensor<?x?x?xf32> {
     %collapsed = tensor.collapse_shape %arg1 [[0], [1, 2]] : tensor<6x2x3xf32, #COO_3D> into tensor<6x6xf32, #COO_2D>
     %0 = tensor.empty() : tensor<5x6xf32>
     %cst = arith.constant 0.000000e+00 : f32
     %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
     %2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32>, tensor<6x6xf32, #COO_2D>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
     %expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
     %ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>
     return %ret1 : tensor<?x?x?xf32>
   }
 
+  func.func @test_sparse_all(%arg0: tensor<5x6xf32, #COO_2D>, %arg1: tensor<6x2x3xf32, #COO_3D>) -> tensor<?x?x?xf32> {
+    %collapsed = tensor.collapse_shape %arg1 [[0], [1, 2]] : tensor<6x2x3xf32, #COO_3D> into tensor<6x6xf32, #COO_2D>
+    %0 = tensor.empty() : tensor<5x6xf32>
+    %cst = arith.constant 0.000000e+00 : f32
+    %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
+    %2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32, #COO_2D>, tensor<6x6xf32, #COO_2D>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
+    %expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
+    %ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>
+    return %ret1 : tensor<?x?x?xf32>
+  }
+
   func.func @test_dense(%arg0: tensor<5x6xf32>, %arg1: tensor<6x2x3xf32>) -> tensor<?x?x?xf32> {
     %collapsed = tensor.collapse_shape %arg1 [[0], [1, 2]] : tensor<6x2x3xf32> into tensor<6x6xf32>
     %0 = tensor.empty() : tensor<5x6xf32>
     %cst = arith.constant 0.000000e+00 : f32
     %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<5x6xf32>) -> tensor<5x6xf32>
     %2 = linalg.matmul ins(%arg0, %collapsed : tensor<5x6xf32>, tensor<6x6xf32>) outs(%1 : tensor<5x6xf32>) -> tensor<5x6xf32>
     %expanded = tensor.expand_shape %2 [[0], [1, 2]] : tensor<5x6xf32> into tensor<5x2x3xf32>
     %ret1 = tensor.cast %expanded : tensor<5x2x3xf32> to tensor<?x?x?xf32>
     return %ret1 :  tensor<?x?x?xf32>
   }
 
 
   func.func @entry() {
     // Setup two sparse vectors.
     %d1 = arith.constant sparse<
         [ [0, 0], [1, 1], [2, 2], [2, 3], [4, 5] ],
           [1.0,      2.0,    3.0,    4.0,   5.0]
     > : tensor<5x6xf32>
 
     %d2 = arith.constant sparse<
       [ [0, 0, 0], [1, 1, 1], [2, 1, 1] ],
         [     6.0,       7.0,      8.0]
     > : tensor<6x2x3xf32>
 
+    %s1 = sparse_tensor.convert %d1 : tensor<5x6xf32> to tensor<5x6xf32, #COO_2D>
+    %s2 = sparse_tensor.convert %d2 : tensor<6x2x3xf32> to tensor<6x2x3xf32, #COO_3D>
+
     //      CHECK: Memref base@ = {{.*}} rank = 3 offset = 0 sizes = [5, 2, 3] strides = [6, 3, 1] data =
     // CHECK-NEXT:[
     // CHECK-SAME: [
     // CHECK-SAME:  [6,    0,    0],
     // CHECK-NEXT:  [0,    0,    0]],
     // CHECK-NEXT: [
     // CHECK-SAME:  [0,    0,    0],
     // CHECK-NEXT:  [0,    14,    0]],
func.func @entry() {
     // CHECK-SAME:  [0,    0,    0],
     // CHECK-NEXT:  [0,    24,    0]],
     // CHECK-NEXT: [
     // CHECK-SAME:  [0,    0,    0],
     // CHECK-NEXT:  [0,    0,    0]],
     // CHECK-NEXT: [
     // CHECK-SAME:  [0,    0,    0],
     // CHECK-NEXT:  [0,    0,    0]]]
-    %s2 = sparse_tensor.convert %d2 : tensor<6x2x3xf32> to tensor<6x2x3xf32, #COO_3D>
-    %so1 = call @test_sparse(%d1, %s2): (tensor<5x6xf32>, tensor<6x2x3xf32, #COO_3D>) -> tensor<?x?x?xf32>
+    %so1 = call @test_sparse_rhs(%d1, %s2): (tensor<5x6xf32>, tensor<6x2x3xf32, #COO_3D>) -> tensor<?x?x?xf32>
     call @printMemref3dF32(%so1) : (tensor<?x?x?xf32>) -> ()
 
+    // Same results.
+    // CHECK-NEXT: Memref base@ = {{.*}} rank = 3 offset = 0 sizes = [5, 2, 3] strides = [6, 3, 1] data =
+    // CHECK-NEXT:[
+    // CHECK-SAME: [
+    // CHECK-SAME:  [6,    0,    0],
+    // CHECK-NEXT:  [0,    0,    0]],
+    // CHECK-NEXT: [
+    // CHECK-SAME:  [0,    0,    0],
+    // CHECK-NEXT:  [0,    14,    0]],
+    // CHECK-NEXT: [
+    // CHECK-SAME:  [0,    0,    0],
+    // CHECK-NEXT:  [0,    24,    0]],
+    // CHECK-NEXT: [
+    // CHECK-SAME:  [0,    0,    0],
+    // CHECK-NEXT:  [0,    0,    0]],
+    // CHECK-NEXT: [
+    // CHECK-SAME:  [0,    0,    0],
+    // CHECK-NEXT:  [0,    0,    0]]]
+    %so2 = call @test_sparse_all(%s1, %s2): (tensor<5x6xf32, #COO_2D>, tensor<6x2x3xf32, #COO_3D>) -> tensor<?x?x?xf32>
+    call @printMemref3dF32(%so2) : (tensor<?x?x?xf32>) -> ()
+
+    bufferization.dealloc_tensor %s1 : tensor<5x6xf32, #COO_2D>
     bufferization.dealloc_tensor %s2 : tensor<6x2x3xf32, #COO_3D>
     bufferization.dealloc_tensor %do1 : tensor<?x?x?xf32>
     bufferization.dealloc_tensor %so1 : tensor<?x?x?xf32>
+    bufferization.dealloc_tensor %so2 : tensor<?x?x?xf32>
     return
   }
 }