[mlir][sparse] Adding new `Merger::addLat` overload

[wrengr]

rebase

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

ExprId addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op = nullptr) {
     return addExp(k, e, kInvalidId, v, op);
   }
   ExprId addExp(TensorExp::Kind k, Value v, Operation *op = nullptr) {
     return addExp(k, kInvalidId, kInvalidId, v, op);
   }
 
   /// Constructs a new iteration lattice point, and returns its identifier.
   LatPointId addLat(TensorId t, LoopId i, ExprId e);
+  LatPointId addLat(const BitVector &bits, ExprId e);
 
   /// Constructs a new (initially empty) set, and returns its identifier.
   LatSetId addSet();
 
   /// Computes a single conjunction of two lattice points by taking the "union"
   /// of `LoopId` (effectively constructing a larger "intersection" of those
   /// loops) with a newly constructed tensor (sub)expression of given kind.
   /// Returns the identifier of the new lattice point.

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

 
 LatPointId Merger::addLat(TensorId t, LoopId i, ExprId e) {
   assert(t < numTensors && i < numLoops);
   const LatPointId p = latPoints.size();
   latPoints.emplace_back(numTensors, numLoops, t, i, e);
   return p;
 }
 
+LatPointId Merger::addLat(const BitVector &bits, ExprId e) {
+  assert(bits.size() == numLoops * numTensors);
+  const LatPointId p = latPoints.size();
+  latPoints.emplace_back(bits, e);
+  return p;
+}
+
 LatSetId Merger::addSet() {
   const LatSetId s = latSets.size();
   latSets.emplace_back();
   return s;
 }
 
 LatPointId Merger::conjLat(TensorExp::Kind kind, LatPointId p0, LatPointId p1,
                            Operation *op) {

 }
 
 LatSetId Merger::mapSet(TensorExp::Kind kind, LatSetId s0, Value v,
                         Operation *op) {
   assert(TensorExp::Kind::kAbsF <= kind && kind <= TensorExp::Kind::kSelect);
   const LatSetId s = addSet();
   for (const LatPointId p : latSets[s0]) {
     const ExprId e = addExp(kind, latPoints[p].exp, v, op);
-    latPoints.emplace_back(latPoints[p].bits, e);
-    latSets[s].push_back(latPoints.size() - 1);
+    latSets[s].push_back(addLat(latPoints[p].bits, e));
   }
   return s;
 }
 
 LatSetId Merger::optimizeSet(LatSetId s0) {
   const LatSetId s = addSet();
   assert(!latSets[s0].empty());
   const LatPointId p0 = latSets[s0][0];