[Mlir-commits] [mlir] [mlir][sparse] first proof-of-concept non-permutation rewriter (PR #70863)

[Aart Bik]

================
@@ -18,10 +20,135 @@ using namespace mlir::sparse_tensor;
 
 namespace {
 
-// TODO:
-//   (1) insert the zero-cost sparse_tensor.reinterpret_map ops
-//   (2) rewrite linalg.generic ops traits on level crds
-//   (3) compute topsort, and resolve cyles with sparse_tensor.convert ops
+//===----------------------------------------------------------------------===//
+// Helper methods.
+//===----------------------------------------------------------------------===//
+
+// Translates a "simple" map according to an identify lvl-map.
+static AffineMap translateMap(OpBuilder &builder, SparseTensorType stt,
+                              AffineMap map) {
+  unsigned lvlRank = stt.getLvlRank();
+  AffineMap lvl2dim = stt.getLvlToDim();
+  assert(lvl2dim.getNumInputs() == lvlRank);
+  SmallVector<AffineExpr> exps;
+  for (unsigned i = 0, n = map.getNumResults(); i < n; i++) {
+    unsigned pos = map.getResult(i).cast<AffineDimExpr>().getPosition();
+    exps.push_back(lvl2dim.getResult(pos));
+  }
+  return AffineMap::get(lvlRank, 0, exps, builder.getContext());
+}
+
+// Generates a "de"mapping reinterpretation of the map.
+static Value genDemap(OpBuilder &builder, SparseTensorEncodingAttr enc,
+                      Value val) {
+  unsigned lvlRank = enc.getLvlTypes().size();
+  AffineMap idMap =
+      AffineMap::getMultiDimIdentityMap(lvlRank, builder.getContext());
+  auto newEnc = SparseTensorEncodingAttr::get(
+      builder.getContext(), enc.getLvlTypes(), idMap, idMap, enc.getPosWidth(),
+      enc.getCrdWidth());
----------------
aartbik wrote:

Yes, I noticed your other helpers for this too. But let me make this change at least.

https://github.com/llvm/llvm-project/pull/70863