[Mlir-commits] [mlir] [mlir][sparse] remove (some) deprecated dim/lvl methods (PR #71125)

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/71125

This removes the most obvious ones. The others are still TBD.

>From b68be85b4b2df0c5d1d5af57a5956e132245fc80 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Thu, 2 Nov 2023 16:02:44 -0700
Subject: [PATCH] [mlir][sparse] remove (some) deprecated dim/lvl methods

This removes the most obvious ones. The others are still TBD.
---
 .../mlir/Dialect/SparseTensor/IR/SparseTensor.h  |  2 --
 .../SparseTensor/IR/SparseTensorDialect.cpp      | 16 +---------------
 .../Transforms/SparseTensorCodegen.cpp           |  2 +-
 .../Transforms/SparseTensorRewriting.cpp         |  3 ++-
 4 files changed, 4 insertions(+), 19 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
index 241d90a87165928..2fb61746d048999 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
@@ -162,12 +162,10 @@ bool isBlockSparsity(AffineMap dimToLvl);
 /// [deprecated] Convenience method to translate the given level to the
 /// corresponding dimension. Requires: `0 <= l < lvlRank`.
 Dimension toOrigDim(SparseTensorEncodingAttr enc, Level l);
-Dimension toOrigDim(RankedTensorType type, Level l);
 
 /// [deprecated] Convenience method to translate the given dimension to
 /// the corresponding level. Requires: `0 <= d < dimRank`.
 Level toStoredDim(SparseTensorEncodingAttr enc, Dimension d);
-Level toStoredDim(RankedTensorType type, Dimension d);
 
 } // namespace sparse_tensor
 } // namespace mlir
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index b3a155686824555..aedac67865aacec 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -409,6 +409,7 @@ SparseTensorEncodingAttr::tranlateShape(ArrayRef<int64_t> srcShape,
 
   if (isPermutation()) {
     for (unsigned r = 0; r < rank; r++) {
+      // FIXME: `toOrigDim` and `toStoredDim` are deprecated.
       unsigned trans = dir == CrdTransDirectionKind::dim2lvl
                            ? toOrigDim(*this, r)
                            : toStoredDim(*this, r);
@@ -953,21 +954,6 @@ Level mlir::sparse_tensor::toStoredDim(SparseTensorEncodingAttr enc,
   return d;
 }
 
-// TODO: Remove this definition once all use-sites have been fixed to
-// properly handle non-permutations.
-Dimension mlir::sparse_tensor::toOrigDim(RankedTensorType type, Level l) {
-  const auto enc = getSparseTensorEncoding(type);
-  assert(!enc || l < enc.getLvlRank());
-  return toOrigDim(enc, l);
-}
-
-// TODO: Remove this definition once all use-sites have been fixed to
-// properly handle non-permutations.
-Level mlir::sparse_tensor::toStoredDim(RankedTensorType type, Dimension d) {
-  assert(d < static_cast<Dimension>(type.getRank()));
-  return toStoredDim(getSparseTensorEncoding(type), d);
-}
-
 /// We normalized sparse tensor encoding attribute by always using
 /// ordered/unique DLT such that "compressed_nu_no" and "compressed_nu" (as well
 /// as other variants) lead to the same storage specifier type, and stripping
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 8c6312150f4c832..bc696267d5279bf 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -248,7 +248,7 @@ static void createAllocFields(OpBuilder &builder, Location loc,
   for (Level lvlRank = stt.getLvlRank(), l = 0; l < lvlRank; l++) {
     // Fills dim sizes array.
     // FIXME: `toOrigDim` is deprecated.
-    desc.setLvlSize(builder, loc, l, dimSizes[toOrigDim(stt, l)]);
+    desc.setLvlSize(builder, loc, l, dimSizes[toOrigDim(stt.getEncoding(), l)]);
     // Pushes a leading zero to positions memref.
     if (stt.isCompressedLvl(l))
       createPushback(builder, loc, desc, SparseTensorFieldKind::PosMemRef, l,
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 7770bd857e88093..c541a5d27e93860 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -872,8 +872,9 @@ struct SparseTensorDimOpRewriter : public OpRewritePattern<tensor::DimOp> {
       return failure();
 
     if (stt.isPermutation()) {
+      // FIXME: `toStoredDim` is deprecated
       rewriter.replaceOpWithNewOp<LvlOp>(op, op.getSource(),
-                                         toStoredDim(stt, *dim));
+                                         toStoredDim(stt.getEncoding(), *dim));
       return success();
     }