[mlir][sparse] Return actual identity map instead of null map (#70365)

Changes:

1. For both dimToLvl and lvlToDim, always returns the actual map instead
of AffineMap() for identity map.
2. Updated custom builder for encoding to have default values.
3. Non-inferable lvlToDim will still return AffineMap() during
inference, so it will be caught by verifier.

Author: yinying-lisa-li

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

@@ -303,10 +303,13 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
 
   let builders = [
     AttrBuilder<(ins "ArrayRef<::mlir::sparse_tensor::DimLevelType>":$lvlTypes,
-                     "AffineMap":$dimToLvl,
-                     "AffineMap":$lvlToDim,
-                     "unsigned":$posWidth,
-                     "unsigned":$crdWidth), [{
+                     CArg<"AffineMap", "{}">:$dimToLvl,
+                     CArg<"AffineMap", "{}">:$lvlToDim,
+                     CArg<"unsigned", "0">:$posWidth,
+                     CArg<"unsigned", "0">:$crdWidth), [{
+      if (!dimToLvl) {
+        dimToLvl = ::mlir::AffineMap::getMultiDimIdentityMap(lvlTypes.size(), $_ctxt);
+      }
       if (!lvlToDim) {
         lvlToDim = ::mlir::sparse_tensor::inferLvlToDim(dimToLvl, $_ctxt);
       }

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

@@ -313,8 +313,7 @@ AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
   lvlAffines.reserve(getLvlRank());
   for (const auto &lvlSpec : lvlSpecs)
     lvlAffines.push_back(lvlSpec.getExpr().getAffineExpr());
-  auto map =  AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
-  if (map.isIdentity()) return AffineMap();
+  auto map = AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
   return map;
 }
 
@@ -328,7 +327,9 @@ AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
     }
   }
   auto map = AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
-  if (dimAffines.empty() || map.isIdentity())
+  // If no lvlToDim map was passed in, returns a null AffineMap and infers it
+  // in SparseTensorEncodingAttr::parse.
+  if (dimAffines.empty())
     return AffineMap();
   return map;
 }

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

@@ -291,7 +291,7 @@ SparseTensorEncodingAttr
 SparseTensorEncodingAttr::withDimToLvl(AffineMap dimToLvl) const {
   assert(getImpl() && "Uninitialized SparseTensorEncodingAttr");
   return SparseTensorEncodingAttr::get(getContext(), getLvlTypes(), dimToLvl,
-                                       getLvlToDim(), getPosWidth(),
+                                       AffineMap(), getPosWidth(),
                                        getCrdWidth());
 }
 

mlir/test/python/dialects/sparse_tensor/dialect.py

@@ -30,9 +30,9 @@ def testEncodingAttr1D():
 
         # CHECK: lvl_types: [<DimLevelType.compressed: 8>]
         print(f"lvl_types: {casted.lvl_types}")
-        # CHECK: dim_to_lvl: None
+        # CHECK: dim_to_lvl: (d0) -> (d0)
         print(f"dim_to_lvl: {casted.dim_to_lvl}")
-        # CHECK: lvl_to_dim: None
+        # CHECK: lvl_to_dim: (d0) -> (d0)
         print(f"lvl_to_dim: {casted.lvl_to_dim}")
         # CHECK: pos_width: 16
         print(f"pos_width: {casted.pos_width}")