[mlir][sparse] avoid slice rewriting when conditions are not met

Reviewed By: Peiming

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

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 4e1e66d..c1cb092 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1165,11 +1165,9 @@ public:
     MLIRContext *ctx = op.getContext();
     auto srcEnc = getSparseTensorEncoding(op.getSourceType());
     auto dstEnc = getSparseTensorEncoding(op.getResult().getType());
-    if (!srcEnc && !dstEnc)
-      return failure();
-
     // TODO: We should check these in ExtractSliceOp::verify.
-    assert(srcEnc && dstEnc && dstEnc.isSlice());
+    if (!srcEnc || !dstEnc || !dstEnc.isSlice())
+      return failure();
     assert(srcEnc.getDimLevelType() == dstEnc.getDimLevelType());
     assert(srcEnc.getDimOrdering() == dstEnc.getDimOrdering());
     assert(srcEnc.getHigherOrdering() == dstEnc.getHigherOrdering());