[mlir][sparse] minor merger code cleanup

[aartbik]

Summary

Moved larger constructor from header to CPP file.
Used toMLIRString() for DimLvlType debug.
Minor layout changes.

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

 /// A class to handle all iteration lattice operations. This class abstracts
 /// away from some implementation details of storing iteration lattices and
 /// tensor expressions. This allows for fine-tuning performance characteristics
 /// independently from the basic algorithm if bottlenecks are identified.
 class Merger {
 public:
   /// Constructs a merger for the given number of tensors, native loops, and
   /// filter loops. The user supplies the number of tensors involved in the
-  /// kernel, with the last tensor in this set denoting the output tensor. The
-  /// merger adds an additional synthetic tensor at the end of this set to
-  /// represent all invariant expressions in the kernel.
-  /// In addition to natives
-  /// loops (which are specified by the GenericOp), extra filter loops are
-  /// needed in order to handle affine expressions on sparse dimensions.
-  /// E.g., (d0, d1, d2) => (d0 + d1, d2), a naive implementation of the filter
-  /// loop could be generated as:
+  /// kernel, with the last tensor in this set denoting the output tensor.
+  /// The merger adds an additional synthetic tensor at the end of this set
+  /// to represent all invariant expressions in the kernel.
+  ///
+  /// In addition to natives loops (which are specified by the GenericOp),
+  /// extra filter loops are needed in order to handle affine expressions on
+  /// sparse dimensions. E.g., (d0, d1, d2) => (d0 + d1, d2), a naive
+  /// implementation of the filter loop could be generated as
+  ///
   /// for (coord : sparse_dim[0])
   ///   if (coord == d0 + d1) {
   ///      generated_code;
   ///   }
   /// }
-  /// to filter out coordinates that are not equal to the affine expression
-  /// result.
+  ///
+  /// to filter out coordinates that are not equal to the affine expression.
+  ///
   /// TODO: we want to make the filter loop more efficient in the future, e.g.,
   /// by avoiding scanning the full stored index sparse (keeping the last
   /// position in ordered list) or even apply binary search to find the index.
-  Merger(unsigned t, unsigned l, unsigned fl)
-      : outTensor(t - 1), syntheticTensor(t), numTensors(t + 1),
-        numNativeLoops(l), numLoops(l + fl), hasSparseOut(false),
-        dimTypes(numTensors,
-                 std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
-        loopIdxToDim(numTensors,
-                     std::vector<Optional<unsigned>>(numLoops, std::nullopt)),
-        dimToLoopIdx(numTensors,
-                     std::vector<Optional<unsigned>>(numLoops, std::nullopt)) {}
+  ///
+  Merger(unsigned t, unsigned l, unsigned fl);
 
   /// Adds a tensor expression. Returns its index.
   unsigned addExp(Kind k, unsigned e0, unsigned e1 = -1u, Value v = Value(),
                   Operation *op = nullptr);
   unsigned addExp(Kind k, unsigned e, Value v, Operation *op = nullptr) {
     return addExp(k, e, -1u, v, op);
   }
   unsigned addExp(Kind k, Value v, Operation *op = nullptr) {
private:
 
   /// Merger data structures.
   const unsigned outTensor;
   const unsigned syntheticTensor;
   const unsigned numTensors;
   const unsigned numNativeLoops;
   const unsigned numLoops;
   bool hasSparseOut;
+
   // Map that converts pair<tensor id, loop id> to the corresponding dimension
   // level type.
   std::vector<std::vector<DimLevelType>> dimTypes;
+
   // Map that converts pair<tensor id, loop id> to the corresponding
   // dimension.
   std::vector<std::vector<Optional<unsigned>>> loopIdxToDim;
+
   // Map that converts pair<tensor id, dim> to the corresponding loop id.
   std::vector<std::vector<Optional<unsigned>>> dimToLoopIdx;
+
   llvm::SmallVector<TensorExp> tensorExps;
   llvm::SmallVector<LatPoint> latPoints;
   llvm::SmallVector<SmallVector<unsigned>> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

 
 LatPoint::LatPoint(unsigned n, unsigned e, unsigned b)
     : bits(n, false), exp(e) {
   bits.set(b);
 }
 
 LatPoint::LatPoint(const BitVector &b, unsigned e) : bits(b), exp(e) {}
 
+Merger::Merger(unsigned t, unsigned l, unsigned fl)
+    : outTensor(t - 1), syntheticTensor(t), numTensors(t + 1),
+      numNativeLoops(l), numLoops(l + fl), hasSparseOut(false),
+      dimTypes(numTensors,
+               std::vector<DimLevelType>(numLoops, DimLevelType::Undef)),
+      loopIdxToDim(numTensors,
+                   std::vector<Optional<unsigned>>(numLoops, std::nullopt)),
+      dimToLoopIdx(numTensors,
+                   std::vector<Optional<unsigned>>(numLoops, std::nullopt)) {}
+
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 unsigned Merger::addExp(Kind k, unsigned e0, unsigned e1, Value v,
                         Operation *op) {
   unsigned e = tensorExps.size();
   tensorExps.push_back(TensorExp(k, e0, e1, v, op));

 }
 
 void Merger::dumpBits(const BitVector &bits) const {
   for (unsigned b = 0, be = bits.size(); b < be; b++) {
     if (bits[b]) {
       unsigned t = tensor(b);
       unsigned i = index(b);
       DimLevelType dlt = dimTypes[t][i];
-      llvm::dbgs() << " i_" << t << "_" << i << "_";
-      if (isDenseDLT(dlt))
-        llvm::dbgs() << "D";
-      else if (isCompressedDLT(dlt))
-        llvm::dbgs() << "C";
-      else if (isSingletonDLT(dlt))
-        llvm::dbgs() << "S";
-      else if (isUndefDLT(dlt))
-        llvm::dbgs() << "U";
-      llvm::dbgs() << "[O=" << isOrderedDLT(dlt) << ",U=" << isUniqueDLT(dlt)
-                   << "]";
+      llvm::dbgs() << " i_" << t << "_" << i << "_" << toMLIRString(dlt);
     }
   }
 }
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.