[Mlir-commits] [mlir] [mlir][sparse] Fix bug in new syntax parser (PR #66024)

[llvmlistbot at llvm.org]

https://github.com/yinying-lisa-li created https://github.com/llvm/llvm-project/pull/66024:

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.

>From b4162b41fe7c50e051b96c7f3147f4823c909ce6 Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Mon, 11 Sep 2023 21:58:57 +0000
Subject: [PATCH] [mlir][sparse] Fix bug in new syntax parser

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.
---
 mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

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 {