[Mlir-commits] [mlir] [mlir][sparse] Order the inferred lvlToDim results by dimension (PR #227674)

[Samarth Narang]

https://github.com/snarang181 created https://github.com/llvm/llvm-project/pull/227674

inverseBlockSparsity appended the expressions of the plain dimensions first and those of the blocked dimensions second, so the result order followed the kind of the expression instead of the dimension it belongs to. For a map with a plain dimension after a blocked one, such as
(i, j) -> (j : dense, i floordiv 2 : compressed, i mod 2 : singleton), the inferred lvlToDim swapped the two dimensions. As a result, lvl_to_dim coordinate translation returns them in the wrong order, and the correct hand-written lvlToDim was rejected by the verifier.

This PR fills the result vector by dimension position instead. Dimensions that do not appear in the map still get no result, as before, since the function is also applied to indexing maps.

Fixes #225249

>From 441da4b6facc59dbcd64807336087904b47433c7 Mon Sep 17 00:00:00 2001
From: Samarth Narang <sn2910 at nyu.edu>
Date: Wed, 30 Sep 2026 11:54:26 +0000
Subject: [PATCH] [mlir][sparse] Order the inferred lvlToDim results by
 dimension

inverseBlockSparsity appended the expressions of the plain dimensions first
and those of the blocked dimensions second, so the result order followed the
kind of the expression instead of the dimension it belongs to. For a map with
a plain dimension after a blocked one, such as
(i, j) -> (j : dense, i floordiv 2 : compressed, i mod 2 : singleton), the
inferred lvlToDim swapped the two dimensions: lvl_to_dim coordinate
translation returned them in the wrong order, and the correct hand-written
lvlToDim was rejected by the verifier as not being the inverse.

Fill the result vector by dimension position instead. Dimensions that do not
appear in the map still get no result, as before, since the function is also
applied to indexing maps.

Fixes #225249
---
 .../SparseTensor/IR/SparseTensorDialect.cpp     | 13 +++++++++----
 .../SparseTensor/roundtrip_encoding.mlir        | 13 +++++++++++++
 .../Dialect/SparseTensor/sparse_foreach.mlir    | 17 +++++++++++++++++
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 781189210ecd39..0ebaf73f06938f 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -1092,9 +1092,10 @@ AffineMap mlir::sparse_tensor::inferLvlToDim(AffineMap dimToLvl,
 
 AffineMap mlir::sparse_tensor::inverseBlockSparsity(AffineMap dimToLvl,
                                                     MLIRContext *context) {
-  SmallVector<AffineExpr> lvlExprs;
+  // The results of lvlToDim follow the dimension order, so the vector is
+  // filled by dimension position rather than by expression kind.
+  SmallVector<AffineExpr> lvlExprs(dimToLvl.getNumDims());
   auto numLvls = dimToLvl.getNumResults();
-  lvlExprs.reserve(numLvls);
   // lvlExprComponents stores information of the floordiv and mod operations
   // applied to the same dimension, so as to build the lvlToDim map.
   std::map<unsigned, SmallVector<AffineExpr, 3>> lvlExprComponents;
@@ -1124,7 +1125,8 @@ AffineMap mlir::sparse_tensor::inverseBlockSparsity(AffineMap dimToLvl,
         assert(false && "expected floordiv or mod");
       }
     } else {
-      lvlExprs.push_back(getAffineDimExpr(i, context));
+      auto pos = cast<AffineDimExpr>(result).getPosition();
+      lvlExprs[pos] = getAffineDimExpr(i, context);
     }
   }
   // Build lvlExprs from lvlExprComponents.
@@ -1138,8 +1140,11 @@ AffineMap mlir::sparse_tensor::inverseBlockSparsity(AffineMap dimToLvl,
         AffineExprKind::Mul, components.second[0], components.second[1]);
     auto addOp =
         getAffineBinaryOpExpr(AffineExprKind::Add, mulOp, components.second[2]);
-    lvlExprs.push_back(addOp);
+    lvlExprs[components.first] = addOp;
   }
+  // Dimensions that do not appear in dimToLvl (this is also applied to
+  // indexing maps) get no result.
+  llvm::erase_if(lvlExprs, [](AffineExpr expr) { return !expr; });
   return dimToLvl.get(dimToLvl.getNumResults(), 0, lvlExprs, context);
 }
 
diff --git a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
index 44710cad246c6a..bb7f8eb6687b81 100644
--- a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
@@ -378,3 +378,16 @@ func.func private @NV_24(%arg0: tensor<?x?x?xf64, #NV_24>) {
 func.func private @NOutOfM(%arg0: tensor<?x?x?xf64, #NOutOfM>) {
   return
 }
+
+// -----
+
+// The explicit lvlToDim of a blocked dimension that follows a plain one is the
+// inverse of dimToLvl and must be accepted.
+#MixedBlock = #sparse_tensor.encoding<{
+  map = {jl, il, ii} (i = il * 2 + ii, j = jl) -> (jl = j : dense, il = i floordiv 2 : compressed, ii = i mod 2 : singleton)
+}>
+
+// CHECK: #[[$MIXED:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d1 : dense, d0 floordiv 2 : compressed, d0 mod 2 : singleton) }>
+// CHECK-LABEL: func private @sparse_mixed_block(
+// CHECK-SAME: tensor<?x?xf64, #[[$MIXED]]>)
+func.func private @sparse_mixed_block(tensor<?x?xf64, #MixedBlock>)
diff --git a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
index 84d7e5af351eb8..e0bc3558f1bf0c 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
@@ -183,3 +183,20 @@ func.func @symbolic_crd_translate(%i: index, %j: index)
       as #Symbolic : index, index, index
   return %l0, %l1, %l2 : index, index, index
 }
+
+// -----
+
+// A plain dimension that follows a blocked one: the inferred lvlToDim must
+// return the dimensions in dimension order.
+#MixedBlock = #sparse_tensor.encoding<{
+  map = (i, j) -> (j : dense, i floordiv 2 : compressed, i mod 2 : singleton)
+}>
+
+// CHECK-LABEL: func.func @sparse_crd_translate_mixed_block(
+// CHECK-SAME:    %[[L0:.*]]: index, %[[L1:.*]]: index, %[[L2:.*]]: index)
+// CHECK:         %[[D0:.*]] = affine.apply {{.*}}[%[[L1]], %[[L2]]]
+// CHECK:         return %[[D0]], %[[L0]] : index, index
+func.func @sparse_crd_translate_mixed_block(%l0: index, %l1: index, %l2: index) -> (index, index) {
+  %d0, %d1 = sparse_tensor.crd_translate lvl_to_dim [%l0, %l1, %l2] as #MixedBlock : index, index
+  return %d0, %d1 : index, index
+}