[mlir][sparse] fix various bug to support sparse pooling

[Peiming]

code cleanup

[Peiming]

code cleanup.

[Peiming]

fix leak

[Peiming]

rebase.

[aartbik]

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

1368

just wondering if we can have a few more asserts here to guard against unexpected case
e.g. when there is no yield, assert on iter args

but just something to ponder on, not needed now

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

15

do we actually vectorize?

if not, then we should probably not introduce the extra runs
but I hope we indeed are ;-) but please check

[Peiming]

address comments.

[Peiming]

remove useless test run

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

15

It vectorized other places that we are not particularly interested in this test, I removed them.

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

if (isDenseDLT(lvlTypes[tid][lvl])) {
       }
     }
   }
 }
 
 void LoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
                               MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
-  rewriter.setInsertionPointToEnd(loopInfo.userCodeBlock);
-  if (!loopInfo.userCodeBlock->empty() &&
-      llvm::isa<scf::ForOp>(loopInfo.loop) &&
-      llvm::isa<scf::YieldOp>(&loopInfo.userCodeBlock->back())) {
-    // scf::For inserts an implicit yield op when there is no loop iter args. In
-    // this case, we need to insert the code before the yield.
-    assert(reduc.empty());
-    rewriter.setInsertionPoint(&loopInfo.userCodeBlock->back());
-  }
   for (auto [tid, lvl, reduced] : loopInfo.sliceDrivenInfo) {
     SliceInfo &info = sliceStack[tid].back();
     assert(isDenseDLT(lvlTypes[tid][lvl]));
     assert(*info.slicedOnLvl == lvl && !reduced);
     (void)reduced;
     // Resets slices pointers as the resolved slices are invalidated after we
     // moves forward to the next slice.
     invalidateSliceIterIdx(rewriter, loc, tid, lvl);
if (!isDenseDLT(lvlTypes[tid][lvl]))
       highs[tid][lvl] = Value();
   }
 }
 
 void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
                                 MutableArrayRef<Value> reduc) {
   const LoopInfo &loopInfo = loopStack.back();
   auto whileOp = llvm::cast<scf::WhileOp>(loopInfo.loop);
-  builder.setInsertionPointToEnd(loopInfo.userCodeBlock);
   Value iv = loopInfo.iv;
 
   // Finalize the induction. Note that the induction could be performed
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
if (operands.size() + delta < whileOp.getNumResults()) {
     assert(operands.size() + delta + 1 == whileOp.getNumResults());
     // The last one is the universial index.
     operands.push_back(ADDI(iv, one));
     // update the loop starting point of current loop sequence
     loopSeqStack.back().first = whileOp->getResult(o++);
   }
 
   assert(o == operands.size() + delta);
+  if (!operands.empty())
     YIELD(operands);
+
   builder.setInsertionPointAfter(whileOp);
 }
 
 void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                   MutableArrayRef<Value> reduc) {
   // Clean up the values, it would help use to discover potential bug at a
   // earlier stage (instead of silently using a wrong value).
   const LoopInfo &loopInfo = loopStack.back();
-  SmallVector<Value> red;
+
+  // Sets the insertion point to the right position.
+  rewriter.setInsertionPointToEnd(loopInfo.userCodeBlock);
+  if (!loopInfo.userCodeBlock->empty() &&
+      llvm::isa<scf::YieldOp>(&loopInfo.userCodeBlock->back())) {
+    // scf::While/For inserts an implicit yield op when there is no loop
+    // iter args. In this case, we need to insert the code before the yield.
+    assert(loopInfo.userCodeBlock->back().getNumResults() == 0);
+    rewriter.setInsertionPoint(&loopInfo.userCodeBlock->back());
+  }
+
   if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
     exitWhileLoop(rewriter, loc, reduc);
   } else {
     exitForLoop(rewriter, loc, reduc);
   }
 
   assert(loopStack.size() == loopSeqStack.size());
   loopStack.pop_back();
scf::WhileOp whileOp = builder.create<scf::WhileOp>(
       [this, c1, tid, lvl, sliceHi, genYield,
        bodyBuilder](OpBuilder &builder, Location loc, ValueRange args) {
         Value iv = args[0];
         Value coord =
             genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], iv);
         Value cont = CMPI(ult, coord, sliceHi);
         TypeRange types = args.drop_front(2).getTypes();
 
-        auto ifOp = builder.create<scf::IfOp>(loc, types, cont, true);
+        auto ifOp = builder.create<scf::IfOp>(loc, types, cont,
+                                              /*withElseBlock=*/!types.empty());
         {
           // 2 reduction variable maintained by us.
           SmallVector<Value> ifRet = args.drop_front(2);
           assert(ifRet.size() == args.size() - 2);
 
           OpBuilder::InsertionGuard guard(builder);
           // If coord >= sliceHi.
+          if (!ifRet.empty()) {
             builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
             YIELD(ifRet);
+          }
 
           // If coord < sliceHi.
           builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
           // Delegates to users' callback.
           bodyBuilder(builder, loc, iv, ifRet);
           if (genYield) {
             builder.setInsertionPointToEnd(&ifOp.getThenRegion().front());
             YIELD(ifRet);

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

static void addSliceBasedConstraints(CodegenEnv &env, OpOperand &t,
   // and d_x > d_y && {d0, d1, d3} - d_x > {d4, d5, d6} - d_y
   const Level lvlRank = map.getNumResults();
   assert(!enc || lvlRank == enc.getLvlRank());
   for (Level lvl = 1; lvl < lvlRank; lvl++) {
     // FIXME: `toOrigDim` is deprecated.
     const AffineExpr fa = map.getResult(toOrigDim(enc, lvl - 1));
     const AffineExpr ta = map.getResult(toOrigDim(enc, lvl));
 
-    if (auto fdim = fa.dyn_cast<AffineDimExpr>()) {
+    if (fa.isa<AffineDimExpr>() || ta.isa<AffineDimExpr>()) {
+      AffineDimCollector fCollector;
+      fCollector.walkPostOrder(fa);
+
       AffineDimCollector tCollector;
       tCollector.walkPostOrder(ta);
-
-      const LoopId f = env.makeLoopId(fdim.getPosition());
+      for (auto fd : fCollector.dims) {
         for (auto td : tCollector.dims) {
+          const LoopId f = env.makeLoopId(fd.getPosition());
           const LoopId t = env.makeLoopId(td.getPosition());
           addIterOrdering(f, t, adjM, inDegree);
         }
+      }
       continue;
     }
 
     // This is a heuristic, we pick an abitrary reduction loop from lhs and
     // rhs and use them as d_x and d_y.
     finder.walkPostOrder(fa);
     const AffineDimExpr fexp = finder.getDimExpr();
     const LoopId fldx = env.makeLoopId(fexp.getPosition());

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

+// DEFINE: %{option} = "enable-runtime-library=false enable-index-reduction=true"
+// DEFINE: %{compile} = mlir-opt %s --sparse-compiler=%{option}
+// DEFINE: %{run} = mlir-cpu-runner \
+// DEFINE:  -e entry -entry-point-result=void  \
+// DEFINE:  -shared-libs=%mlir_c_runner_utils | \
+// DEFINE: FileCheck %s
+//
+// RUN: %{compile} | %{run}
+//
+// Do the same run, but now with direct IR generation.
+// REDEFINE: %{option} = "enable-runtime-library=true enable-buffer-initialization=true enable-index-reduction=true"
+// RUN: %{compile} | %{run}
+
+#CCCC = #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "compressed", "compressed" ], posWidth = 32, crdWidth = 32 }>
+
+func.func @pooling_nhwc_sum_CCCC(%input: tensor<1x4x4x1xf32, #CCCC>, %filter: tensor<2x2xf32>) -> tensor<1x3x3x1xf32, #CCCC> {
+  %init = bufferization.alloc_tensor() : tensor<1x3x3x1xf32, #CCCC>
+  %0 = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>,
+                                strides = dense<1> : tensor<2xi64>}
+     ins (%input, %filter: tensor<1x4x4x1xf32, #CCCC>, tensor<2x2xf32>)
+    outs (%init: tensor<1x3x3x1xf32, #CCCC>) -> tensor<1x3x3x1xf32, #CCCC>
+  return %0 : tensor<1x3x3x1xf32, #CCCC>
+}
+
+func.func @pooling_nhwc_sum(%input: tensor<1x4x4x1xf32>, %filter: tensor<2x2xf32>) -> tensor<1x3x3x1xf32> {
+  %init = arith.constant dense<[[ [[0.0], [0.0], [0.0]],
+                                  [[0.0], [0.0], [0.0]],
+                                  [[0.0], [0.0], [0.0]] ]]> : tensor<1x3x3x1xf32>
+  %0 = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>,
+                                strides = dense<1> : tensor<2xi64>}
+     ins (%input, %filter: tensor<1x4x4x1xf32>, tensor<2x2xf32>)
+    outs (%init: tensor<1x3x3x1xf32>) -> tensor<1x3x3x1xf32>
+  return %0 : tensor<1x3x3x1xf32>
+}
+
+
+func.func @entry() {
+  %c0 = arith.constant 0 : index
+  %zero = arith.constant 0.00000e+00 : f32
+
+  %filter = arith.constant dense<
+     [[  1.0,  0.0],
+      [  0.0,  1.0]]
+  > : tensor<2x2xf32>
+
+  %in_dense = arith.constant dense<
+     [[[[1.0],  [2.0],  [1.0],  [2.0]],
+       [[1.0],  [2.0],  [1.0],  [2.0]],
+       [[1.0],  [2.0],  [1.0],  [2.0]],
+       [[1.0],  [2.0],  [1.0],  [2.0]]]]
+  > : tensor<1x4x4x1xf32>
+
+  %in_CCCC = sparse_tensor.convert %in_dense : tensor<1x4x4x1xf32> to tensor<1x4x4x1xf32, #CCCC>
+
+  %dense_ret = call @pooling_nhwc_sum(%in_dense, %filter) : (tensor<1x4x4x1xf32>, tensor<2x2xf32>) -> tensor<1x3x3x1xf32>
+  %CCCC_ret = call @pooling_nhwc_sum_CCCC(%in_CCCC, %filter) : (tensor<1x4x4x1xf32, #CCCC>, tensor<2x2xf32>) -> tensor<1x3x3x1xf32, #CCCC>
+
+  // CHECK: ( ( ( ( 6 ), ( 6 ), ( 6 ) ), ( ( 6 ), ( 6 ), ( 6 ) ), ( ( 6 ), ( 6 ), ( 6 ) ) ) )
+  %dense_v = vector.transfer_read %dense_ret[%c0, %c0, %c0, %c0], %zero
+      : tensor<1x3x3x1xf32>, vector<1x3x3x1xf32>
+  vector.print %dense_v : vector<1x3x3x1xf32>
+
+  //
+  // Sparse pooling should have the same output.
+  //
+
+  // CHECK-NEXT: ( ( ( ( 6 ), ( 6 ), ( 6 ) ), ( ( 6 ), ( 6 ), ( 6 ) ), ( ( 6 ), ( 6 ), ( 6 ) ) ) )
+  %s1 = sparse_tensor.convert %CCCC_ret : tensor<1x3x3x1xf32, #CCCC> to tensor<1x3x3x1xf32>
+  %v1 = vector.transfer_read %s1[%c0, %c0, %c0, %c0], %zero
+      : tensor<1x3x3x1xf32>, vector<1x3x3x1xf32>
+  vector.print %v1 : vector<1x3x3x1xf32>
+
+  // Releases resources.
+  bufferization.dealloc_tensor %in_CCCC : tensor<1x4x4x1xf32, #CCCC>
+  bufferization.dealloc_tensor %CCCC_ret : tensor<1x3x3x1xf32, #CCCC>
+  bufferization.dealloc_tensor %dense_ret : tensor<1x3x3x1xf32>
+  return
+}