[mlir][sparse] remove a few rewriting failures

[aartbik]

Summary

Rationale:
Make sure preconditions are tested already during verfication.
Currently, the only way a sparse rewriting rule can fail is if
(1) the linalg op does not have sparse annotations, or
(2) a yet to be handled operation is encounted inside the op

[penpornk]

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

837

Question please. How do you choose between return failure(); and assert()?
The old code return failure(); for op.getNumOutputs() != 1 before.

862

Do we remove this because it's already checked elsewhere?

[aartbik]

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

837

The return failure() denotes to the calling method that the rewriting rule has failed (did not apply), but it is not an error condition. I have replaced that with an assert since we now verify that much earlier, and we should never encounter that condition anymore

862

Yes, just like above. We verify that earlier now, so the rewriting rule always applies to something with sparse annotations that passed the verifier.

[nicolasvasilache]

LGTM!

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 template <>
 LogicalResult AnnotationsVerifier<GenericOp>::verify(GenericOp op) {
   ArrayAttr sparseAttr = op.sparseAttr();
   if (!sparseAttr)
     return success();
   // Verify consistency of sparse annotations.
   if (!op.hasTensorSemantics())
     return op.emitOpError("expected sparse annotations on tensors only");
+  if (op.getNumOutputs() != 1)
+    return op.emitOpError("expected single output tensor");
   unsigned numTensors = op.getNumInputsAndOutputs();
   if (sparseAttr.size() != numTensors)
     return op.emitOpError("expected one sparse annotation for each tensor");
   for (unsigned t = 0; t < numTensors; t++) {
     auto dimAttr = sparseAttr[t].dyn_cast_or_null<ArrayAttr>();
     if (!dimAttr)
       return op.emitOpError("expected sparse annotation array for tensor ")
              << t;

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

 namespace {
 
 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
   using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
-    unsigned numTensors = op.getNumInputsAndOutputs();
-    unsigned numLoops = op.iterator_types().getValue().size();
-    Merger merger(numTensors, numLoops);
-
     // Detects sparse annotations and translate the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     if (!op.hasSparseSemantics())
       return failure();
+    assert(op.getNumOutputs() == 1);
+    unsigned numTensors = op.getNumInputsAndOutputs();
+    unsigned numLoops = op.iterator_types().getValue().size();
+    Merger merger(numTensors, numLoops);
     findSparseAnnotations(op, merger.sparse());
 
-    // Accept only single, dense result.
-    if (op.getNumOutputs() != 1 ||
-        std::any_of(merger.sparse().back().begin(),
-                    merger.sparse().back().end(), [](bool b) { return b; }))
-      return failure();
-
     // Computes a topologically sorted iteration graph to ensure
     // tensors are visited in natural index order. Fails on cycles.
     // This assumes that higher-level passes have already put the
     // tensors in each tensor expression in a feasible order.
     // TODO: try again without *dense* constraints on failure or
     //       even try to insert sparse reorderings to resolve cycles
     std::vector<unsigned> topSort;
     if (!computeIterationGraph(op, topSort))
       return failure();
 
     // Finds the terminating yield statement and builds the tensor
     // expression for the Linalg operation in SSA form.
-    auto &region = op.region();
-    if (!llvm::hasSingleElement(region))
-      return failure(); // single block only
-    Operation *yield = region.front().getTerminator();
+    Operation *yield = op.region().front().getTerminator();
     Optional<unsigned> exp = buildTensorExp(merger, op, yield->getOperand(0));
     if (!exp.hasValue())
       return failure(); // build failure
 
     // Recursively generates code.
     CodeGen codegen(numTensors, numLoops);
     genBuffers(merger, codegen, rewriter, op);
     genStmt(merger, codegen, rewriter, op, topSort, exp.getValue(), 0);

mlir/test/Dialect/Linalg/sparse_invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 #trait_memref = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [
     [ "S" ],  // a
     [ "D" ]   // x
   ],
   iterator_types = ["parallel"]
 }
 
 func @invalid_memref(%arga: memref<32xf32>, %argb: f32) -> tensor<32xf32> {
   // expected-error@+1 {{'linalg.generic' op expected sparse annotations on tensors only}}
   %0 = linalg.generic #trait_memref
     ins(%arga: memref<32xf32>) {
       ^bb(%a: f32):
         %0 = addf %a, %argb  : f32
         linalg.yield %0 : f32
   } -> tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
 // -----
 
+#trait_two_out = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>,  // x (out)
+    affine_map<(i) -> (i)>   // y (out)
+  ],
+  sparse = [
+    [ "S" ],  // a
+    [ "D" ],  // x
+    [ "D" ]   // y
+  ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_two_out(%arga: tensor<32xf32>) -> tensor<32xf32> {
+  // expected-error@+1 {{'linalg.generic' op expected single output tensor}}
+  %0, %1 = linalg.generic #trait_two_out
+    ins(%arga: tensor<32xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %a : f32
+        linalg.yield %a, %0 : f32, f32
+  } -> tensor<32xf32>, tensor<32xf32>
+  return %1 : tensor<32xf32>
+}
+
+// -----
+
+#trait_two_blocks = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  sparse = [
+    [ "S" ],  // a
+    [ "D" ]   // x
+  ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_two_blocks(%arga: tensor<32xf32>) -> tensor<32xf32> {
+  // expected-error@+1 {{'linalg.generic' op expects region #0 to have 0 or 1 blocks}}
+  %0 = linalg.generic #trait_two_blocks
+    ins(%arga: tensor<32xf32>) {
+      ^bb1(%a: f32):
+        %0 = addf %a, %a : f32
+      ^bb2:
+        linalg.yield %0 : f32
+  } -> tensor<32xf32>
+  return %0 : tensor<32xf32>
+}
+
+// -----
+
+#trait_no_block = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>  // a
+  ],
+  sparse = [
+    [ "S" ]  // a
+  ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_no_block(%arga: tensor<32xf32>) {
+  // expected-error@+1 {{'linalg.generic' op expected region with 1 block}}
+  linalg.generic #trait_no_block
+    ins(%arga: tensor<32xf32>) {
+    }
+  return
+}
+
+// -----
+
 #trait_too_many = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [
     [ "S" ],  // a
     [ "S" ],  // b
     [ "D" ]   // x
   ],
   iterator_types = ["parallel"]
 }
 
 func @invalid_too_many(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> {
   // expected-error@+1 {{'linalg.generic' op expected one sparse annotation for each tensor}}
   %0 = linalg.generic #trait_too_many
     ins(%arga: tensor<32xf32>) {
       ^bb(%a: f32):
         %0 = addf %a, %argb  : f32
         linalg.yield %0 : f32
   } -> tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
 // -----
 
 #trait_no_array = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [ 1, 2 ],
   iterator_types = ["parallel"]
 }
 
 func @invalid_no_array(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> {
   // expected-error@+1 {{'linalg.generic' op expected sparse annotation array for tensor 0}}
   %0 = linalg.generic #trait_no_array
     ins(%arga: tensor<32xf32>) {
       ^bb(%a: f32):
         %0 = addf %a, %argb  : f32
         linalg.yield %0 : f32
   } -> tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
 // -----
 
 #trait_wrong_rank = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [
     [ "S" ],
     [ "D", "D" ]
   ],
   iterator_types = ["parallel"]
 }
 
 func @invalid_wrong_rank(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> {
   // expected-error@+1 {{'linalg.generic' op expected sparse annotation with rank 1 for tensor 1}}
   %0 = linalg.generic #trait_wrong_rank
     ins(%arga: tensor<32xf32>) {
       ^bb(%a: f32):
         %0 = addf %a, %argb  : f32
         linalg.yield %0 : f32
   } -> tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
 // -----
 
 #trait_no_string = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // a
     affine_map<(i,j) -> (i,j)>   // x (out)
   ],
   sparse = [
     [ "S", 1 ],
     [ "D", "D" ]
   ],
   iterator_types = ["parallel","parallel"]
 }
 
 func @invalid_no_string(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> {
   // expected-error@+1 {{'linalg.generic' op expected sparse annotation at position 1 for tensor 0}}
   %0 = linalg.generic #trait_no_string
     ins(%arga: tensor<32x16xf32>) {
       ^bb(%a: f32):
         %0 = addf %a, %argb  : f32
         linalg.yield %0 : f32
   } -> tensor<32x16xf32>
   return %0 : tensor<32x16xf32>
 }