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

[Peiming]

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

365

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

365

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

1853

This part is good!

[Peiming]

remove obscue code.

[aartbik]

Good work, thanks Peiming!

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

 /// 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,
                                   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
   // access 1-dim memref. Also admissable since insertions cannot occur.
   bool allDense = true;
   auto iteratorTypes = op.iterator_types().getValue();
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
+    // Whether the current GenericOp is admissible.
     bool isAdmissible = false;
+    bool hasCycle = true;
     // 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,
+    for (auto mask : allMask)
+      if (computeIterationGraph(merger, op, topSort, mask)) {
+        hasCycle = false;
+        if (isAdmissableTensorExp(merger, op, topSort, exp, &sparseOut,
                                   outerParNest)) {
-        // This is an admissible GenericOp.
           isAdmissible = true;
           break;
         }
-      // else try a less strict constraints.
+        // else try a set of less strict constraints.
       }
 
-    if (!isAdmissible)
+    if (hasCycle)
       // Give it one last shot to resolve the cycle.
       return resolveCycle(merger, rewriter, op);
+    if (!isAdmissible)
+      // Inadmissible expression, 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();