[Mlir-commits] [mlir] [mlir][sparse] Populate lvlToDim (PR #68937)

[Peiming Liu]

================
@@ -749,6 +748,71 @@ mlir::sparse_tensor::getSparseTensorEncoding(Type type) {
   return nullptr;
 }
 
+AffineMap mlir::sparse_tensor::inferLvlToDim(AffineMap dimToLvl,
+                                             MLIRContext *context) {
+  auto map = static_cast<AffineMap>(dimToLvl);
+  AffineMap lvlToDim;
+  // TODO: support ELL instead of returning an empty lvlToDim.
+  if (!map || map.getNumSymbols() != 0) {
+    lvlToDim = AffineMap();
+  } else if (map.isPermutation()) {
+    lvlToDim = inversePermutation(map);
+  } else {
+    lvlToDim = inverseBlockSparsity(map, context);
+  }
+  return lvlToDim;
+}
+
+AffineMap mlir::sparse_tensor::inverseBlockSparsity(AffineMap dimToLvl,
+                                                    MLIRContext *context) {
+  SmallVector<AffineExpr> lvlExprs;
+  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.
+  // Map key is the position of the dimension in dimToLvl.
+  // Map value is a SmallVector that contains lvl var for floordiv, multiplier,
+  // lvl var for mod in dimToLvl.
+  // For example, for il = i floordiv 2 and ii = i mod 2, the SmalleVector
+  // would be [il, 2, ii]. It could be used to build the AffineExpr
+  // i = il * 2 + ii in lvlToDim.
+  std::map<unsigned, SmallVector<AffineExpr, 3>> lvlExprComponents;
----------------
PeimingLiu wrote:

This can be a vector right? Seems like you always index the map by dim.

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