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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

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

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


3 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+9-4) 
- (modified) mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir (+13) 
- (modified) mlir/test/Dialect/SparseTensor/sparse_foreach.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 781189210ecd3..0ebaf73f06938 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 44710cad246c6..bb7f8eb6687b8 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 84d7e5af351eb..e0bc3558f1bf0 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
+}

``````````

</details>


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