[Mlir-commits] [mlir] b165650 - [mlir][sparse] Return actual identity map instead of null map (#70365)

[llvmlistbot at llvm.org]

Author: Yinying Li
Date: 2023-10-26T18:30:34-04:00
New Revision: b165650aee851b15c4f87f0d510b37cb1bf6b814

URL: https://github.com/llvm/llvm-project/commit/b165650aee851b15c4f87f0d510b37cb1bf6b814
DIFF: https://github.com/llvm/llvm-project/commit/b165650aee851b15c4f87f0d510b37cb1bf6b814.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
    mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
    mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
    mlir/test/python/dialects/sparse_tensor/dialect.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index 6c3cfc2c49b2e91..7e2ad11752b34d8 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/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);
       }

diff  --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
index 851867926fe679e..6a81a11a932f94a 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
+++ b/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;
 }

diff  --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index be44a0d31c92a7d..4f1c446faec3714 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/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());
 }
 

diff  --git a/mlir/test/python/dialects/sparse_tensor/dialect.py b/mlir/test/python/dialects/sparse_tensor/dialect.py
index 240db6ebd1d1eb3..fe7b41e536e2763 100644
--- a/mlir/test/python/dialects/sparse_tensor/dialect.py
+++ b/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}")