[mlir][sparse] refine trait of sparse_tensor.convert

[aartbik]

Summary

Rationale:
The currently used trait was demanding that all types are the same
which is not true (since the sparse part may change and the dim sizes
may be relaxed). This revision uses the correct trait and makes the
rank match test explicit in the verify method.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

%0 = linalg.matmul
       ins(%a, %b: tensor<?x?xf32>, tensor<?x?xf32>)
       outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
     ```
   }];
   let assemblyFormat = "`[` $sizes `]` attr-dict `:` type($result)";
 }
 
 def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
-  [NoSideEffect, SameOperandsAndResultType]>,
+  [NoSideEffect, SameOperandsAndResultElementType]>,
     Arguments<(ins AnyTensor:$source)>,
     Results<(outs AnyTensor:$dest)> {
   string summary = "Converts between different tensor types";
   string description = [{
     Converts one sparse or dense tensor type to another tensor type. The rank
-    and dimensions of the source and destination types must match, but the sparse
-    encoding of these types can obviously be different. The name `convert` was
-    preferred over `cast`, since the operation may incur a non-trivial cost.
+    of the source and destination types must match exactly, and the dimension
+    sizes must either match exactly or relax from a static to a dynamic size.
+    The sparse encoding of the two types can obviously be completely different.
+    The name `convert` was preferred over `cast`, since the operation may incur
+    a non-trivial cost.
 
     When converting between two different sparse tensor types, only explicitly
     stored values are moved from one underlying sparse storage format to
     the other. When converting from an unannotated dense tensor type to a
     sparse tensor type, an explicit test for nonzero values is used. When
     converting to an unannotated dense tensor type, implicit zeroes in the
     sparse storage format are made explicit. Note that the conversions can have
     non-trivial costs associated with them, since they may involve elaborate

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

if (!constantOp ||
              << shape[i];
   }
   return success();
 }
 
 static LogicalResult verify(ConvertOp op) {
   if (auto tp1 = op.source().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = op.dest().getType().dyn_cast<RankedTensorType>()) {
-      assert(tp1.getRank() == tp2.getRank());
+      if (tp1.getRank() != tp2.getRank())
+        return op.emitError("unexpected conversion mismatch in rank");
       auto shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       // Accept size matches between the source and the destination type
       // (e.g. 10 vs. 10, 10 vs. ?, or ? vs. ?), but reject direct mismatches or
       // matches that would need a runtime assert (e.g. 10 vs. 20 or ? vs. 10).
       for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++) {
         if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamicSize)
           return op.emitError("unexpected conversion mismatch in dimension ")

mlir/test/Dialect/SparseTensor/invalid.mlir

 func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
   // expected-error@+1 {{unexpected type in convert}}
   %0 = sparse_tensor.convert %arg0 : tensor<*xf32> to tensor<10xf32>
   return %0 : tensor<10xf32>
 }
 
 // -----
 
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+func @sparse_convert_rank_mismatch(%arg0: tensor<10x10xf64, #DCSR>) -> tensor<?xf64> {
+  // expected-error@+1 {{unexpected conversion mismatch in rank}}
+  %0 = sparse_tensor.convert %arg0 : tensor<10x10xf64, #DCSR> to tensor<?xf64>
+  return %0 : tensor<?xf64>
+}
+
+// -----
+
 #CSR = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
 
-func @sparse_convert_mismatch(%arg0: tensor<10x?xf32>) -> tensor<10x10xf32, #CSR> {
+func @sparse_convert_dim_mismatch(%arg0: tensor<10x?xf32>) -> tensor<10x10xf32, #CSR> {
   // expected-error@+1 {{unexpected conversion mismatch in dimension 1}}
   %0 = sparse_tensor.convert %arg0 : tensor<10x?xf32> to tensor<10x10xf32, #CSR>
   return %0 : tensor<10x10xf32, #CSR>
 }