[mlir][sparse] Fixes bug in VarSet ctor

[wrengr]

Summary

Previously, the commented out code in the DimLvlMap ctor would result in VarSet::add raising an OOB error; which should be impossible because the ctor asserted DimLvlMap::isWF which ensures that all variables occuring in the map are within bounds for the ranks.

The root cause of that bug was the VarSet ctor using SmallBitVector::reserve which does not actually change the size of the bitvectors (hence the subsequent OOB). This is corrected by using any of SmallBitVector::resize, the move-ctor, or the copy-ctor. Since the default-initialized bitvectors being modified/overwritten have size zero, there shouldn't be any significant performance difference between these three implementations.

[wrengr]

Removed fixme comment in VarSet::contains(Var). The potential logic-bug mentioned there is the same one in the VarSet ctor that this CL fixes. (I.e., if we were to assert against OOB in that method, then everything would still pass not that we've corrected the VarSet ctor)

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

 //===----------------------------------------------------------------------===//
 // `DimLvlMap` implementation.
 //===----------------------------------------------------------------------===//
 
 DimLvlMap::DimLvlMap(unsigned symRank, ArrayRef<DimSpec> dimSpecs,
                      ArrayRef<LvlSpec> lvlSpecs)
     : symRank(symRank), dimSpecs(dimSpecs), lvlSpecs(lvlSpecs) {
   // First, check integrity of the variable-binding structure.
+  // NOTE: This establishes the invariant that calls to `VarSet::add`
+  // below cannot cause OOB errors.
   assert(isWF());
 
   // TODO: Second, we need to infer/validate the `lvlToDim` mapping.
   // Along the way we should set every `DimSpec::elideExpr` according
   // to whether the given expression is inferable or not.  Notably, this
   // needs to happen before the code for setting every `LvlSpec::elideVar`,
   // since if the LvlVar is only used in elided DimExpr, then the
   // LvlVar should also be elided.
+  // NOTE: Whenever we set a new DimExpr, we must make sure to validate it
+  // against our ranks, to restore the invariant established by `isWF` above.
+  // TODO(wrengr): We might should adjust the `DimLvlExpr` ctor to take a
+  // `Ranks` argument and perform the validation then.
 
   // Third, we set every `LvlSpec::elideVar` according to whether that
   // LvlVar occurs in a non-elided DimExpr (TODO: or CountingExpr).
+  // NOTE: The invariant established by `isWF` ensures that the following
+  // calls to `VarSet::add` cannot raise OOB errors.
   VarSet usedVars(getRanks());
-  // NOTE TO Wren: bypassed for now
-  // for (const auto &dimSpec : dimSpecs)
-  //  if (!dimSpec.canElideExpr())
-  //    usedVars.add(dimSpec.getExpr());
+  for (const auto &dimSpec : dimSpecs)
+    if (!dimSpec.canElideExpr())
+      usedVars.add(dimSpec.getExpr());
   for (auto &lvlSpec : this->lvlSpecs)
     lvlSpec.setElideVar(!usedVars.contains(lvlSpec.getBoundVar()));
 }
 
 bool DimLvlMap::isWF() const {
   const auto ranks = getRanks();
   unsigned dimNum = 0;
   for (const auto &dimSpec : dimSpecs)

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

 //===----------------------------------------------------------------------===//
 // `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].reserve(ranks.getRank(vk));
+    impl[vk] = llvm::SmallBitVector(ranks.getRank(vk));
 }
 
 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()];
   const auto num = var.getNum();
-  // FIXME(wrengr): If we `assert(num < bits.size())` then
-  // "roundtrip_encoding.mlir" will fail.  So we need to figure out
-  // where exactly the OOB `var` is coming from, to determine whether
-  // that's a logic bug or not.
   return num < bits.size() && bits[num];
 }
 
 bool VarSet::occursIn(VarSet const &other) const {
   for (const auto vk : everyVarKind)
     if (impl[vk].anyCommon(other.impl[vk]))
       return true;
   return false;