[mlir][sparse] Update verifier for block sparsity and singleton (#69389)

Updates:
1. Verification of block sparsity.
2. Verification of singleton level type can only follow compressed or
loose_compressed levels. And all level types after singleton should be
singleton.
3. Added getBlockSize function.
4. Added an invalid encoding test for an incorrect lvlToDim map that
user provides.

Author: yinying-lisa-li

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

@@ -173,6 +173,21 @@ AffineMap inferLvlToDim(AffineMap dimToLvl, MLIRContext *context);
 /// Asserts on failure (so only use when known to succeed).
 AffineMap inverseBlockSparsity(AffineMap dimToLvl, MLIRContext *context);
 
+/// Given the dimToLvl map, returns the block sizes in a vector.
+/// For instance, a 2x3 block will return [2, 3]. Unblocked dimension i
+/// will return 0, and i floordiv 1, i mod 1 will return 1. Therefore,
+/// the example below will return [0, 1].
+/// map = ( i, j ) ->
+///       ( i : dense,
+///         j floordiv 1 : compressed,
+///         j mod 1      : dense
+///       )
+/// Only valid block sparsity will be accepted.
+SmallVector<unsigned> getBlockSize(AffineMap dimToLvl);
+
+/// Given the dimToLvl map, returns if it's block sparsity.
+bool isBlockSparsity(AffineMap dimToLvl);
+
 //
 // Reordering.
 //

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp

@@ -356,10 +356,15 @@ AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
 AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
   SmallVector<AffineExpr> dimAffines;
   dimAffines.reserve(getDimRank());
-  for (const auto &dimSpec : dimSpecs)
-    dimAffines.push_back(dimSpec.getExpr().getAffineExpr());
+  for (const auto &dimSpec : dimSpecs) {
+    auto expr = dimSpec.getExpr().getAffineExpr();
+    if (expr) {
+      dimAffines.push_back(expr);
+    }
+  }
   auto map = AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
-  if (map.isIdentity()) return AffineMap();
+  if (dimAffines.empty() || map.isIdentity())
+    return AffineMap();
   return map;
 }
 

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

@@ -455,6 +455,7 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
   SmallVector<DimLevelType> lvlTypes;
   SmallVector<SparseTensorDimSliceAttr> dimSlices;
   AffineMap dimToLvl = {};
+  AffineMap lvlToDim = {};
   unsigned posWidth = 0;
   unsigned crdWidth = 0;
   StringRef attrName;
@@ -568,6 +569,7 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
 
       ERROR_IF(dimToLvl, "Cannot mix `dimToLvl` with `map`")
       dimToLvl = dlm.getDimToLvlMap(parser.getContext());
+      lvlToDim = dlm.getLvlToDimMap(parser.getContext());
       break;
     }
     } // switch
@@ -582,8 +584,9 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
 #undef RETURN_ON_FAIL
 
   // Construct struct-like storage for attribute.
-  // TODO: Fetch lvlToDim if user provides one
-  AffineMap lvlToDim = inferLvlToDim(dimToLvl, parser.getContext());
+  if (!lvlToDim || lvlToDim.isEmpty()) {
+    lvlToDim = inferLvlToDim(dimToLvl, parser.getContext());
+  }
   return parser.getChecked<SparseTensorEncodingAttr>(
       parser.getContext(), lvlTypes, dimToLvl, lvlToDim, posWidth, crdWidth,
       dimSlices);
@@ -663,6 +666,17 @@ SparseTensorEncodingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
     return emitError() << "unexpected position bitwidth: " << posWidth;
   if (!acceptBitWidth(crdWidth))
     return emitError() << "unexpected coordinate bitwidth: " << crdWidth;
+  if (auto it = std::find_if(lvlTypes.begin(), lvlTypes.end(), isSingletonDLT);
+      it != std::end(lvlTypes)) {
+    if (it == lvlTypes.begin() ||
+        (!isCompressedDLT(*(it - 1)) && !isLooseCompressedDLT(*(it - 1))))
+      return emitError() << "expected compressed or loose_compressed level "
+                            "before singleton level";
+    if (!std::all_of(it, lvlTypes.end(),
+                     [](DimLevelType i) { return isSingletonDLT(i); }))
+      return emitError() << "expected all singleton lvlTypes "
+                            "following a singleton level";
+  }
   // Before we can check that the level-rank is consistent/coherent
   // across all fields, we need to define it.  The source-of-truth for
   // the `getLvlRank` method is the length of the level-types array,
@@ -678,19 +692,12 @@ SparseTensorEncodingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
       return emitError()
              << "level-rank mismatch between dimToLvl and lvlTypes: "
              << dimToLvl.getNumResults() << " != " << lvlRank;
-    // TODO:  The following is attempting to match the old error-conditions
-    // from prior to merging dimOrdering and higherOrdering into dimToLvl.
-    // That is, we currently require `dimToLvl` to be either a permutation
-    // (as when higherOrdering is the identity) or expansive (as per the
-    // constraints on higherOrdering).  However, those constraints do
-    // not match the intended semantics of `dimToLvl`.  As we improve the
-    // compiler to actually handle non-permutations, we need to update these
-    // checks to match what is actually supported.  In particular, this is
-    // where we'll have to check that when `lvlToDim` is provided then it
-    // is indeed an inverse of `dimToLvl`, and when it isn't provided then
-    // it can be automatically inferred.
-    if (dimRank == lvlRank && !dimToLvl.isPermutation())
-      return emitError() << "expected a permutation affine map for dimToLvl";
+    auto inferRes = inferLvlToDim(dimToLvl, dimToLvl.getContext());
+    // Symbols can't be inferred but are acceptable.
+    if (!inferRes && dimToLvl.getNumSymbols() == 0)
+      return emitError() << "failed to infer lvlToDim from dimToLvl";
+    if (lvlToDim && (inferRes != lvlToDim))
+      return emitError() << "expected lvlToDim to be an inverse of dimToLvl";
     if (dimRank > lvlRank)
       return emitError() << "unexpected dimToLvl mapping from " << dimRank
                          << " to " << lvlRank;
@@ -758,8 +765,7 @@ AffineMap mlir::sparse_tensor::inferLvlToDim(AffineMap dimToLvl,
     lvlToDim = AffineMap();
   } else if (map.isPermutation()) {
     lvlToDim = inversePermutation(map);
-  } else {
-    // TODO: check if it's block sparsity
+  } else if (isBlockSparsity(map)) {
     lvlToDim = inverseBlockSparsity(map, context);
   }
   return lvlToDim;
@@ -818,6 +824,53 @@ AffineMap mlir::sparse_tensor::inverseBlockSparsity(AffineMap dimToLvl,
   return dimToLvl.get(dimToLvl.getNumResults(), 0, lvlExprs, context);
 }
 
+SmallVector<unsigned> mlir::sparse_tensor::getBlockSize(AffineMap dimToLvl) {
+  assert(isBlockSparsity(dimToLvl) &&
+         "expected dimToLvl to be block sparsity for calling getBlockSize");
+  SmallVector<unsigned> blockSize;
+  for (auto result : dimToLvl.getResults()) {
+    if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
+      if (result.getKind() == AffineExprKind::Mod) {
+        blockSize.push_back(
+            binOp.getRHS().dyn_cast<AffineConstantExpr>().getValue());
+      }
+    } else {
+      blockSize.push_back(0);
+    }
+  }
+  return blockSize;
+}
+
+bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
+  if (!dimToLvl)
+    return false;
+  std::map<unsigned, int64_t> coeffientMap;
+  for (auto result : dimToLvl.getResults()) {
+    if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
+      auto pos = binOp.getLHS().dyn_cast<AffineDimExpr>().getPosition();
+      if (result.getKind() == AffineExprKind::FloorDiv) {
+        // Expect only one floordiv for each dimension.
+        if (coeffientMap.find(pos) != coeffientMap.end())
+          return false;
+        coeffientMap[pos] =
+            binOp.getRHS().dyn_cast<AffineConstantExpr>().getValue();
+      } else if (result.getKind() == AffineExprKind::Mod) {
+        // Expect floordiv before mod.
+        if (coeffientMap.find(pos) == coeffientMap.end())
+          return false;
+        // Expect mod to have the same coefficient as floordiv.
+        if (binOp.getRHS().dyn_cast<AffineConstantExpr>().getValue() !=
+            coeffientMap[pos]) {
+          return false;
+        }
+      } else {
+        return false;
+      }
+    }
+  }
+  return !coeffientMap.empty();
+}
+
 bool mlir::sparse_tensor::isCOOType(SparseTensorEncodingAttr enc,
                                     Level startLvl, bool isUnique) {
   if (!enc ||

mlir/test/Dialect/SparseTensor/invalid_encoding.mlir

@@ -60,7 +60,7 @@ func.func private @tensor_sizes_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
 // -----
 
-// expected-error@+1 {{unexpected dimToLvl mapping from 2 to 1}}
+// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
 #a = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense)}>
 func.func private @tensor_sizes_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
@@ -119,7 +119,7 @@ func.func private @tensor_dimtolvl_mismatch(%arg0: tensor<8xi32, #a>) -> ()
 
 // -----
 
-// expected-error@+1 {{expected a permutation affine map for dimToLvl}}
+// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
 #a = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d0 : compressed)}>
 func.func private @tensor_no_permutation(%arg0: tensor<16x32xf32, #a>) -> ()
 
@@ -251,3 +251,22 @@ func.func private @too_few_lvl_decl(%arg0: tensor<?x?xf64, #TooFewLvlDecl>) {
 func.func private @wrong_order_lvl_decl(%arg0: tensor<?x?xf64, #WrongOrderLvlDecl>) {
   return
 }
+
+// -----
+
+// expected-error@+1 {{expected lvlToDim to be an inverse of dimToLvl}}
+#BSR_explicit = #sparse_tensor.encoding<{
+  map =
+  {il, jl, ii, jj}
+  ( i = il * 3 + ii,
+    j = jl * 2 + jj
+  ) ->
+  ( il = i floordiv 2 : dense,
+    jl = j floordiv 3 : compressed,
+    ii = i mod 2      : dense,
+    jj = j mod 3      : dense
+  )
+}>
+func.func private @BSR_explicit(%arg0: tensor<?x?xf64, #BSR_explicit>) {
+  return
+}