[mlir][sparse] Fix bug in new syntax parser

[yinying-lisa-li]

Currently, dimlvlmap with identity affine map will be treated as empty affine map. But the new syntax would treat it as an actual identity affine map such as {d0} -> {d0}. This mismatch could raise an error when we are comparing sparse encodings.

[llvmbot]

@llvm/pr-subscribers-mlir-sparse

Changes

Currently, dimlvlmap with identity affine map will be treated as empty affine map. But the new syntax would treat it as an actual identity affine map such as {d0} -> {d0}. This mismatch could raise error in reshape rewriting.

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

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp (+6-2)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
index 81302f200f686bb..05fce96043826f1 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
@@ -348,7 +348,9 @@ AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
   lvlAffines.reserve(getLvlRank());
   for (const auto &lvlSpec : lvlSpecs)
     lvlAffines.push_back(lvlSpec.getExpr().getAffineExpr());
-  return AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+  auto map =  AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+  if (map.isIdentity()) return AffineMap();
+  return map;
 }
 
 AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
@@ -356,7 +358,9 @@ AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
   dimAffines.reserve(getDimRank());
   for (const auto &dimSpec : dimSpecs)
     dimAffines.push_back(dimSpec.getExpr().getAffineExpr());
-  return AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+  auto map = AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+  if (map.isIdentity()) return AffineMap();
+  return map;
 }
 
 void DimLvlMap::dump() const {

[llvmbot]

@llvm/pr-subscribers-mlir-core

Changes

Currently, dimlvlmap with identity affine map will be treated as empty affine map. But the new syntax would treat it as an actual identity affine map such as {d0} -> {d0}. This mismatch could raise error in reshape rewriting.

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

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp (+6-2)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
index 81302f200f686bb..05fce96043826f1 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
@@ -348,7 +348,9 @@ AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
   lvlAffines.reserve(getLvlRank());
   for (const auto &lvlSpec : lvlSpecs)
     lvlAffines.push_back(lvlSpec.getExpr().getAffineExpr());
-  return AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+  auto map =  AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+  if (map.isIdentity()) return AffineMap();
+  return map;
 }
 
 AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
@@ -356,7 +358,9 @@ AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
   dimAffines.reserve(getDimRank());
   for (const auto &dimSpec : dimSpecs)
     dimAffines.push_back(dimSpec.getExpr().getAffineExpr());
-  return AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+  auto map = AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+  if (map.isIdentity()) return AffineMap();
+  return map;
 }
 
 void DimLvlMap::dump() const {

[llvmbot]

@llvm/pr-subscribers-mlir

Changes

Currently, dimlvlmap with identity affine map will be treated as empty affine map. But the new syntax would treat it as an actual identity affine map such as {d0} -> {d0}. This mismatch could raise error in reshape rewriting.

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

1 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp (+6-2)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
index 81302f200f686bb..05fce96043826f1 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp
@@ -348,7 +348,9 @@ AffineMap DimLvlMap::getDimToLvlMap(MLIRContext *context) const {
   lvlAffines.reserve(getLvlRank());
   for (const auto &lvlSpec : lvlSpecs)
     lvlAffines.push_back(lvlSpec.getExpr().getAffineExpr());
-  return AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+  auto map =  AffineMap::get(getDimRank(), getSymRank(), lvlAffines, context);
+  if (map.isIdentity()) return AffineMap();
+  return map;
 }
 
 AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
@@ -356,7 +358,9 @@ AffineMap DimLvlMap::getLvlToDimMap(MLIRContext *context) const {
   dimAffines.reserve(getDimRank());
   for (const auto &dimSpec : dimSpecs)
     dimAffines.push_back(dimSpec.getExpr().getAffineExpr());
-  return AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+  auto map = AffineMap::get(getLvlRank(), getSymRank(), dimAffines, context);
+  if (map.isIdentity()) return AffineMap();
+  return map;
 }
 
 void DimLvlMap::dump() const {

[aartbik]

Nice job Yinying!

Only change I would request is in the description, since the bug will show up anywhere we do a comparison on the encoding (the reshaping test was just one that exposed it). But good to go once you make that tiny change.