[mlir][sparse] Adding `Ranks::operator==` and `VarSet::getRank` etc

[wrengr]

Summary

Depends On D156001

[wrengr]

rebasing. And adding a missing newline between function defns

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.h

constexpr Ranks(unsigned symRank, unsigned dimRank, unsigned lvlRank)
     impl[to_index(VarKind::Symbol)] = symRank;
     impl[to_index(VarKind::Dimension)] = dimRank;
     impl[to_index(VarKind::Level)] = lvlRank;
   }
   Ranks(VarKindArray<unsigned> const &ranks)
       : Ranks(ranks[VarKind::Symbol], ranks[VarKind::Dimension],
               ranks[VarKind::Level]) {}
 
+  bool operator==(Ranks const &other) const;
+  bool operator!=(Ranks const &other) const { return !(*this == other); }
+
   constexpr unsigned getRank(VarKind vk) const { return impl[to_index(vk)]; }
   constexpr unsigned getSymRank() const { return getRank(VarKind::Symbol); }
   constexpr unsigned getDimRank() const { return getRank(VarKind::Dimension); }
   constexpr unsigned getLvlRank() const { return getRank(VarKind::Level); }
 
   [[nodiscard]] constexpr bool isValid(Var var) const {
     return var.getNum() < getRank(var.getKind());
   }
class VarSet final {
   // bitvector (akin to how `mlir::presburger::PresburgerSpace` does it);
   // however, doing so would greatly complicate the implementation of the
   // `occursIn(VarSet)` method.
   VarKindArray<llvm::SmallBitVector> impl;
 
 public:
   explicit VarSet(Ranks const &ranks);
 
+  unsigned getRank(VarKind vk) const { return impl[vk].size(); }
+  unsigned getSymRank() const { return getRank(VarKind::Symbol); }
+  unsigned getDimRank() const { return getRank(VarKind::Dimension); }
+  unsigned getLvlRank() const { return getRank(VarKind::Level); }
+  Ranks getRanks() const {
+    return Ranks(getSymRank(), getDimRank(), getLvlRank());
+  }
+
   bool contains(Var var) const;
   bool occursIn(VarSet const &vars) const;
   bool occursIn(DimLvlExpr expr) const;
 
   void add(Var var);
   void add(VarSet const &vars);
   void add(DimLvlExpr expr);
 };

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.cpp

 //===- Var.cpp ------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Var.h"
 #include "DimLvlMap.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 using namespace mlir::sparse_tensor::ir_detail;
 
 //===----------------------------------------------------------------------===//
+// `VarKind` helpers.
+//===----------------------------------------------------------------------===//
+
+/// For use in foreach loops.
+static constexpr const VarKind everyVarKind[] = {
+    VarKind::Dimension, VarKind::Symbol, VarKind::Level};
+
+//===----------------------------------------------------------------------===//
 // `Var` implementation.
 //===----------------------------------------------------------------------===//
 
 void Var::print(AsmPrinter &printer) const { print(printer.getStream()); }
 
 void Var::print(llvm::raw_ostream &os) const {
   os << toChar(getKind()) << getNum();
 }
 
 void Var::dump() const {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
 //===----------------------------------------------------------------------===//
 // `Ranks` implementation.
 //===----------------------------------------------------------------------===//
 
+bool Ranks::operator==(Ranks const &other) const {
+  for (const auto vk : everyVarKind)
+    if (getRank(vk) != other.getRank(vk))
+      return false;
+  return true;
+}
+
 bool Ranks::isValid(DimLvlExpr expr) const {
   assert(expr);
   // Compute the maximum identifiers for symbol-vars and dim/lvl-vars
   // (each `DimLvlExpr` only allows one kind of non-symbol variable).
   int64_t maxSym = -1, maxVar = -1;
   mlir::getMaxDimAndSymbol<ArrayRef<AffineExpr>>({{expr.getAffineExpr()}},
                                                  maxVar, maxSym);
   // TODO(wrengr): We may want to add a call to `LLVM_DEBUG` like
   // `willBeValidAffineMap` does.  And/or should return `InFlightDiagnostic`
   // instead of bool.
   return maxSym < getSymRank() && maxVar < getRank(expr.getAllowedVarKind());
 }
 
 //===----------------------------------------------------------------------===//
 // `VarSet` implementation.
 //===----------------------------------------------------------------------===//
 
-static constexpr const VarKind everyVarKind[] = {
-    VarKind::Dimension, VarKind::Symbol, VarKind::Level};
-
 VarSet::VarSet(Ranks const &ranks) {
   // NOTE: We must not use `reserve` here, since that doesn't change
   // the `size` of the bitvectors and therefore will result in unexpected
   // OOB errors.  Either `resize` or copy/move-ctor work; we opt for the
   // move-ctor since it should be (marginally) more efficient.
   for (const auto vk : everyVarKind)
     impl[vk] = llvm::SmallBitVector(ranks.getRank(vk));
+  assert(getRanks() == ranks);
 }
 
 bool VarSet::contains(Var var) const {
   // NOTE: We make sure to return false on OOB, for consistency with
   // the `anyCommon` implementation of `VarSet::occursIn(VarSet)`.
   // However beware that, as always with silencing OOB, this can hide
   // bugs in client code.
   const llvm::SmallBitVector &bits = impl[var.getKind()];