[mlir][sparse] Adding `SparseTensorType::{operator==, hasSameDimToLvlMap}`

[wrengr]

Summary

Depends On D143800

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h

/*implicit*/ operator T() const {
     return rtp;
   }
 
   /// Explicitly convert to `RankedTensorType`.  This method is
   /// a convenience for resolving overload-ambiguity issues with
   /// implicit conversion.
   RankedTensorType getRankedTensorType() const { return rtp; }
 
+  bool operator==(const SparseTensorType &other) const {
+    // All other fields are derived from `rtp` and therefore don't need
+    // to be checked.
+    return rtp == other.rtp;
+  }
+
+  bool operator!=(const SparseTensorType &other) const {
+    return !(*this == other);
+  }
+
   MLIRContext *getContext() const { return rtp.getContext(); }
 
   Type getElementType() const { return rtp.getElementType(); }
 
   /// Returns the encoding (or the null-attribute for dense-tensors).
   SparseTensorEncodingAttr getEncoding() const { return enc; }
 
   /// Returns true for tensors which have an encoding, and false for
public:
   /// (This is always true for dense-tensors.)
   bool isAllOrdered() const { return enc.isAllOrdered(); }
 
   /// Returns true if the dimToLvl mapping is the identity.
   /// (This is always true for dense-tensors.)
   bool isIdentity() const { return !dim2lvl; }
 
   /// Returns the dimToLvl mapping (or the null-map for the identity).
+  /// If you intend to compare the results of this method for equality,
+  /// see `hasSameDimToLvlMap` instead.
   AffineMap getDimToLvlMap() const { return dim2lvl; }
 
   /// Returns the dimToLvl mapping, where the identity map is expanded out
   /// into a full `AffineMap`.  This method is provided as a convenience,
   /// but for most purposes other methods (`isIdentity`, `getDimToLvlMap`,
   /// etc) will be more helpful.
   AffineMap getExpandedDimToLvlMap() const {
     return dim2lvl
                ? dim2lvl
                : AffineMap::getMultiDimIdentityMap(getDimRank(), getContext());
   }
 
+  /// Returns true iff the two types have the same mapping.  This method
+  /// takes care to handle identity maps properly, so it should be preferred
+  /// over using `getDimToLvlMap` followed by `AffineMap::operator==`.
+  bool hasSameDimToLvlMap(const SparseTensorType &other) const {
+    // If the maps are the identity, then we need to check the rank
+    // to be sure they're the same size identity.  (And since identity
+    // means dimRank==lvlRank, we use lvlRank as a minor optimization.)
+    return isIdentity() ? (other.isIdentity() && lvlRank == other.lvlRank)
+                        : (dim2lvl == other.dim2lvl);
+  }
+
   /// Returns the dimension-rank.
   Dimension getDimRank() const { return rtp.getRank(); }
 
   /// Returns the level-rank.
   Level getLvlRank() const { return lvlRank; }
 
   /// Returns the dimension-shape.
   ArrayRef<DynSize> getDimShape() const { return rtp.getShape(); }

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 
 // Helper to detect a sparse tensor type operand.
 static bool isSparseTensor(OpOperand *op) {
   auto enc = getSparseTensorEncoding(op->get().getType());
   return enc &&
          llvm::is_contained(enc.getDimLevelType(), DimLevelType::Compressed);
 }
 
-static bool hasSameDimOrdering(RankedTensorType rtp1, RankedTensorType rtp2) {
-  assert(rtp1.getRank() == rtp2.getRank());
-  return SparseTensorType(rtp1).getDimToLvlMap() ==
-         SparseTensorType(rtp2).getDimToLvlMap();
-}
-
 // Helper method to find zero/uninitialized allocation.
 static bool isAlloc(OpOperand *op, bool isZero) {
   Value val = op->get();
   // Check allocation, with zero alloc when required.
   if (auto alloc = val.getDefiningOp<AllocTensorOp>()) {
     Value copy = alloc.getCopy();
     if (isZero)
       return copy && isZeroValue(copy);
LogicalResult sparse2SparseRewrite(ConvertOp op,
     sizesForTensor(rewriter, srcSizes, loc, srcRTT, src);
     Value tmpCoo = Value();
     Value nnz = rewriter.create<NumberOfEntriesOp>(loc, src);
     // We need a tmp COO buffer if and only if
     // 1. the src tensor is not a COO and
     // 2. the src tensor is not ordered in the same way as the target
     // tensor (e.g., src tensor is not ordered or src tensor haves a different
     // dimOrdering).
-    if (!isUniqueCOOType(srcRTT) && !(SparseTensorType(srcRTT).isAllOrdered() &&
-                                      hasSameDimOrdering(srcRTT, dstTp))) {
+    if (const SparseTensorType srcTp(srcRTT);
+        !isUniqueCOOType(srcRTT) &&
+        !(srcTp.isAllOrdered() && srcTp.hasSameDimToLvlMap(dstTp))) {
       // Construct a COO tensor from the src tensor.
       // TODO: there may be cases for which more efficiently without
       // going through an intermediate COO, such as cases that only change
       // the overhead types.
       SmallVector<Value> dynSrcSizes;
       getDynamicSizes(srcRTT, srcSizes, dynSrcSizes);
       srcRTT =
           getUnorderedCOOFromTypeWithOrdering(srcRTT, dstTp.getDimToLvlMap());
LogicalResult sparse2SparseRewrite(ConvertOp op,
     // the guarantee from the sparse tensor encoding).
     if (!srcTp.isAllOrdered()) {
       // Retrieve the values-array.
       Value y = genToValues(rewriter, loc, src);
       const auto encSrc = srcTp.getEncoding();
       // Sort the COO tensor so that its elements are ordered via increasing
       // indices for the storage ordering of the dst tensor. Use SortCoo if the
       // COO tensor has the same dim ordering as the dst tensor.
-      if (dimRank > 1 && hasSameDimOrdering(srcTp, dstTp)) {
+      if (dimRank > 1 && srcTp.hasSameDimToLvlMap(dstTp)) {
         MemRefType indTp =
             get1DMemRefType(getIndexOverheadType(rewriter, encSrc),
                             /*withLayout=*/false);
         Value xs = rewriter.create<ToIndicesBufferOp>(loc, indTp, src);
         rewriter.create<SortCooOp>(
             loc, nnz, xs, ValueRange{y}, rewriter.getIndexAttr(dimRank),
             rewriter.getIndexAttr(0), SparseTensorSortKind::HybridQuickSort);
       } else {