[Mlir-commits] [mlir] 1e6ef0c - [mlir][sparse] refine trait of sparse_tensor.convert

[Aart Bik]

Author: Aart Bik
Date: 2021-10-26T14:36:49-07:00
New Revision: 1e6ef0cfb0964f17066b01402b7c6fbf1a2967fe

URL: https://github.com/llvm/llvm-project/commit/1e6ef0cfb0964f17066b01402b7c6fbf1a2967fe
DIFF: https://github.com/llvm/llvm-project/commit/1e6ef0cfb0964f17066b01402b7c6fbf1a2967fe.diff

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

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.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D112576

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
    mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
    mlir/test/Dialect/SparseTensor/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index d1724b4c6f5cc..c4682a16a0eff 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -74,15 +74,17 @@ def SparseTensor_InitOp : SparseTensor_Op<"init", []>,
 }
 
 def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
-  [NoSideEffect, SameOperandsAndResultType]>,
+  [NoSideEffect, SameOperandsAndResultElementType]>,
     Arguments<(ins AnyTensor:$source)>,
     Results<(outs AnyTensor:$dest)> {
   string summary = "Converts between 
diff erent 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 
diff erent. 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 
diff erent.
+    The name `convert` was preferred over `cast`, since the operation may incur
+    a non-trivial cost.
 
     When converting between two 
diff erent sparse tensor types, only explicitly
     stored values are moved from one underlying sparse storage format to

diff  --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 3b123545baf1c..1ee943cd7adbd 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -237,7 +237,8 @@ static LogicalResult verify(InitOp op) {
 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

diff  --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir
index 89553595c8142..399cf3fba8d50 100644
--- a/mlir/test/Dialect/SparseTensor/invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid.mlir
@@ -160,9 +160,19 @@ func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
 
 // -----
 
+#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+func @sparse_convert_rank_mismatch(%arg0: tensor<10x10xf64, #DCSR>) -> tensor<?xf64> {
+  // expected-error at +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 at +1 {{unexpected conversion mismatch in dimension 1}}
   %0 = sparse_tensor.convert %arg0 : tensor<10x?xf32> to tensor<10x10xf32, #CSR>
   return %0 : tensor<10x10xf32, #CSR>