[mlir][sparse] Only try to compute a better iteraton graph when needed

[Peiming]

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

362

Or should I return an enum? it might look awkward (but type converter used the same thing for Optional<LogicalResult>)

[Peiming]

add missing period...

[Peiming]

more missing periods..

[Peiming]

only resolve cycle when there is one

[aartbik]

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

362

I would be okay with this solution, since it is used elsewhere too (like your codegen type converter :-), but it seems the hasCycle is all we need

1860

This part is good!

[Peiming]

remove obscue code.

[aartbik]

Good work, thanks Peiming!

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

return val.getDefiningOp<linalg::InitTensorOp>() ||
          val.getDefiningOp<bufferization::AllocTensorOp>();
 }
 
 /// Returns true when the tensor expression is admissable for codegen.
 /// Since all sparse input tensors are admissable, we just need to check
 /// whether the out tensor in the tensor expression codegen is admissable.
 /// Sets `sparseOut` to the tensor and `outerParNest` to the outer injective
 /// nesting depth when a "truly dynamic" sparse tensor output occurs.
-static bool isAdmissableTensorExp(Merger &merger, linalg::GenericOp op,
-                                  std::vector<unsigned> &topSort, unsigned exp,
-                                  OpOperand **sparseOut,
+/// Return true on accept.
+/// Return false on reject.
+/// Return llvm::None on try a looser order.
+static Optional<bool> isAdmissableTensorExp(Merger &merger,
+                                            linalg::GenericOp op,
+                                            std::vector<unsigned> &topSort,
+                                            unsigned exp, OpOperand **sparseOut,
                                             unsigned &outerParNest) {
   OpOperand *lhs = op.getOutputOperand(0);
   unsigned tensor = lhs->getOperandNumber();
   auto enc = getSparseTensorEncoding(lhs->get().getType());
   // An non-annotated output tensor is assumed dense, and becomes a random
   // access n-dim memref. Admissable since insertions cannot occur.
   if (!enc)
     return true;
   // An all-dense annotated "sparse" output tensor becomes a linearized random
if (isMaterializing(lhs->get())) {
     // Determine admissable dynamic insertion situations:
     // (1) fully injective, since there are no reductions,
     // (2) admissable 1-d expansion in innermost dimension.
     if (nest >= op.getRank(lhs) - 1) {
       *sparseOut = lhs;
       outerParNest = nest;
       return true;
     }
+    // The current loop sequences is too strict
+    return llvm::None;
   }
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (reductions).
 //===----------------------------------------------------------------------===//
 
LogicalResult matchAndRewrite(linalg::GenericOp op,
     OpOperand *sparseOut = nullptr;
     unsigned outerParNest = 0;
     // Computes a topologically sorted iteration graph to ensure tensors
     // are visited in natural index order. Gradually relaxes the considered
     // constraints until an acyclic iteration graph results, such that sparse
     // code generation can proceed. As a last resort, an attempt is made
     // to resolve cycles by inserting a conversion.
     std::vector<unsigned> topSort;
-    // Whether the current GenericOp is admissible
-    bool isAdmissible = false;
+    // Whether the current GenericOp is admissible.
+    Optional<bool> isAdmissible = llvm::None;
     // An const list of all masks that we used for interation graph
     // computation. Must be ordered from strict -> loose.
     const auto allMask = {SortMask::kIncludeAll, SortMask::kIncludeUndef,
                           SortMask::kIncludeDense, SortMask::kSparseOnly};
     for (auto mask : allMask) {
-      if (computeIterationGraph(merger, op, topSort, mask) &&
-          isAdmissableTensorExp(merger, op, topSort, exp, &sparseOut,
-                                outerParNest)) {
-        // This is an admissible GenericOp.
-        isAdmissible = true;
+      if (computeIterationGraph(merger, op, topSort, mask)) {
+        isAdmissible = isAdmissableTensorExp(merger, op, topSort, exp,
+                                             &sparseOut, outerParNest);
+        if (!isAdmissible)
+          // Due to cycle, try looser constraints.
+          continue;
+        // Either admissible or inadmissable due to other reasons.
         break;
       }
       // else try a less strict constraints.
     }
 
-    if (!isAdmissible)
+    if (!isAdmissible) // None
       // Give it one last shot to resolve the cycle.
       return resolveCycle(merger, rewriter, op);
+    if (!isAdmissible.value())
+      // Inadmissible due to other reason, simply reject.
+      return failure();
 
     // Recursively generates code if admissible.
     merger.setHasSparseOut(sparseOut != nullptr);
     CodeGen codegen(options, numTensors, numLoops, sparseOut, outerParNest);
     genBuffers(merger, codegen, rewriter, op);
     genStmt(merger, codegen, rewriter, op, topSort, exp, 0);
     genResult(merger, codegen, rewriter, op);
     return success();