[Mlir-commits] [mlir] [mlir][sparse] first proof-of-concept non-permutation rewriter (PR #70863)
Aart Bik
llvmlistbot at llvm.org
Tue Oct 31 14:59:53 PDT 2023
================
@@ -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
More information about the Mlir-commits
mailing list