[mlir][SparseTensor] Fix type conversion rule

[matthias-springer]

A type conversion rule cannot make any assumptions about the number of pre-existing types in the results vector.

This commit fixes a failed assertion in a SparseTensor type conversion rule. This is only reproducible when type conversion caching is deactivated. There's no way to do this at the moment. This commit is in preparation of adding context-aware type conversions, which will deactivate type caching in such cases.

Depends on #140347.

[llvmbot]

@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

A type conversion rule cannot make any assumptions about the number of pre-existing types in the results vector.

This commit fixes a failed assertion in a SparseTensor type conversion rule. This is only reproducible when type conversion caching is deactivated. There's no way to do this at the moment. This commit is in preparation of adding context-aware type conversions, which will deactivate type caching in such cases.

Depends on #140347.

---

Full diff: https://github.com/llvm/llvm-project/pull/140350.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp (+5-4)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp
index 8bbb2cac5efdf..79602a22dc1fe 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp
@@ -38,12 +38,13 @@ convertSparseTensorType(RankedTensorType rtp, SmallVectorImpl<Type> &fields) {
   if (!stt.hasEncoding())
     return std::nullopt;
 
+  unsigned numFields = fields.size();
   foreachFieldAndTypeInSparseTensor(
       stt,
-      [&fields](Type fieldType, FieldIndex fieldIdx,
-                SparseTensorFieldKind /*fieldKind*/, Level /*lvl*/,
-                LevelType /*lt*/) -> bool {
-        assert(fieldIdx == fields.size());
+      [&](Type fieldType, FieldIndex fieldIdx,
+          SparseTensorFieldKind /*fieldKind*/, Level /*lvl*/,
+          LevelType /*lt*/) -> bool {
+        assert(numFields + fieldIdx == fields.size());
         fields.push_back(fieldType);
         return true;
       });

[PeimingLiu]

Thx！