[mlir][sparse] add print methods to Merger (for debugging)

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

 #define MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_
 
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/BitVector.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
+/// Tensor expression kind.
 enum class Kind { kTensor, kInvariant, kMulF, kMulI, kAddF, kAddI };
+
+/// Dimension level type for a tensor (undef means index does not appear).
 enum class Dim { kSparse, kDense, kSingle, kUndef };
 
 /// Tensor expression. Represents a MLIR expression in tensor index notation.
 /// For tensors, e0 denotes the tensor index. For invariants, the IR value is
 /// stored directly. For binary operations, e0 and e1 denote the index of the
 /// children tensor expressions.
 struct TensorExp {
   TensorExp(Kind k, unsigned x, unsigned y, Value v)
       : kind(k), e0(x), e1(y), val(v) {
     assert((kind == Kind::kTensor && e0 != -1u && e1 == -1u && !val) ||
            (kind == Kind::kInvariant && e0 == -1u && e1 == -1u && val) ||
            (kind >= Kind::kMulF && e0 != -1u && e1 != -1u && !val));
   }
+
+  /// Tensor expression kind.
   Kind kind;
+
   /// Indices of children expression(s).
   unsigned e0;
   unsigned e1;
+
   /// Direct link to IR for an invariant. During code generation,
   /// field is used to cache "hoisted" loop invariant tensor loads.
   Value val;
 };
 
 /// Lattice point. Each lattice point consists of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding
 /// tensor expression.
 struct LatPoint {
   LatPoint(unsigned n, unsigned e, unsigned b) : bits(n, false), exp(e) {
     bits.set(b);
   }
   LatPoint(const llvm::BitVector &b, unsigned e) : bits(b), exp(e) {}
+
   /// Conjunction of tensor loop indices as bitvector. This represents
   /// all indices involved in the tensor expression
   llvm::BitVector bits;
+
   /// Simplified conjunction of tensor loop indices as bitvector. This
   /// represents a simplified condition under which this tensor expression
   /// must execute. Pre-computed during codegen to avoid repeated eval.
   llvm::BitVector simple;
+
   /// Index of the tensor expresssion.
   unsigned exp;
 };
 
 /// 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.
public:
   /// Setter
   void setDim(unsigned t, unsigned i, Dim d) { dims[t][i] = d; }
 
   /// Getters.
   TensorExp &exp(unsigned e) { return tensorExps[e]; }
   LatPoint &lat(unsigned l) { return latPoints[l]; }
   SmallVector<unsigned, 16> &set(unsigned s) { return latSets[s]; }
 
+#ifndef NDEBUG
+  /// Print methods (for debugging).
+  void dumpExp(unsigned e) const;
+  void dumpLat(unsigned p) const;
+  void dumpSet(unsigned s) const;
+  void dumpBits(const llvm::BitVector &bits) const;
+#endif
+
 private:
   const unsigned outTensor;
   const unsigned numTensors;
   const unsigned numLoops;
 
   std::vector<std::vector<Dim>> dims;
   llvm::SmallVector<TensorExp, 32> tensorExps;
   llvm::SmallVector<LatPoint, 16> latPoints;
   llvm::SmallVector<SmallVector<unsigned, 16>, 8> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

 //===- Merger.cpp - Implementation of iteration lattices ------------------===//
 //
 // 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 "mlir/Dialect/SparseTensor/Utils/Merger.h"
 
+#include "mlir/IR/Operation.h"
+#include "llvm/Support/Debug.h"
+
 namespace mlir {
 namespace sparse_tensor {
 
 unsigned Merger::addExp(Kind k, unsigned e0, unsigned e1, Value v) {
   unsigned e = tensorExps.size();
   tensorExps.push_back(TensorExp(k, e0, e1, v));
   return e;
 }

 
 bool Merger::hasAnyDimOf(const llvm::BitVector &bits, Dim d) const {
   for (unsigned b = 0, be = bits.size(); b < be; b++)
     if (bits[b] && isDim(b, d))
       return true;
   return false;
 }
 
+#ifndef NDEBUG
+
+//
+// Print methods (for debugging).
+//
+
+void Merger::dumpExp(unsigned e) const {
+  switch (tensorExps[e].kind) {
+  case Kind::kTensor:
+    llvm::dbgs() << "tensor_" << tensorExps[e].e0;
+    break;
+  case Kind::kInvariant:
+    llvm::dbgs() << "invariant";
+    break;
+  default:
+  case Kind::kMulI:
+    llvm::dbgs() << "(";
+    dumpExp(tensorExps[e].e0);
+    llvm::dbgs() << " * ";
+    dumpExp(tensorExps[e].e1);
+    llvm::dbgs() << ")";
+    break;
+  case Kind::kAddF:
+  case Kind::kAddI:
+    llvm::dbgs() << "(";
+    dumpExp(tensorExps[e].e0);
+    llvm::dbgs() << " + ";
+    dumpExp(tensorExps[e].e1);
+    llvm::dbgs() << ")";
+    break;
+  }
+}
+
+void Merger::dumpLat(unsigned p) const {
+  llvm::dbgs() << "lat(";
+  dumpBits(latPoints[p].bits);
+  llvm::dbgs() << " :";
+  dumpBits(latPoints[p].simple);
+  llvm::dbgs() << " / ";
+  dumpExp(latPoints[p].exp);
+  llvm::dbgs() << " )\n";
+}
+
+void Merger::dumpSet(unsigned s) const {
+  llvm::dbgs() << "{ #" << latSets[s].size() << "\n";
+  for (unsigned p : latSets[s]) {
+    llvm::dbgs() << "  ";
+    dumpLat(p);
+  }
+  llvm::dbgs() << "}\n";
+}
+
+void Merger::dumpBits(const llvm::BitVector &bits) const {
+  for (unsigned b = 0, be = bits.size(); b < be; b++) {
+    if (bits[b]) {
+      unsigned t = tensor(b);
+      unsigned i = index(b);
+      llvm::dbgs() << " i_" << t << "_" << i << "_";
+      switch (dims[t][i]) {
+      case Dim::kSparse:
+        llvm::dbgs() << "S";
+        break;
+      case Dim::kDense:
+        llvm::dbgs() << "D";
+        break;
+      case Dim::kSingle:
+        llvm::dbgs() << "T";
+        break;
+      case Dim::kUndef:
+        llvm::dbgs() << "U";
+        break;
+      }
+    }
+  }
+}
+
+#endif // NDEBUG
+
 } // namespace sparse_tensor
 } // namespace mlir

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/SparseTensor/Transforms/Passes.td",
     deps = [":PassBaseTdFiles"],
 )
 
 cc_library(
     name = "SparseTensor",
-    srcs = glob(["lib/Dialect/SparseTensor/IR/*.cpp"]),
+    srcs = glob([
+        "lib/Dialect/SparseTensor/IR/*.cpp",
+    ]),
     hdrs = ["include/mlir/Dialect/SparseTensor/IR/SparseTensor.h"],
     includes = ["include"],
     deps = [
         ":IR",
         ":SideEffectInterfaces",
         ":SparseTensorAttrDefsIncGen",
         ":SparseTensorOpsIncGen",
+        ":SparseTensorUtils",
+        ":StandardOps",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
+    name = "SparseTensorUtils",
+    srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
+    hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),
+    includes = ["include"],
+    deps = [
+        ":IR",
+        ":SideEffectInterfaces",
+        ":SparseTensorAttrDefsIncGen",
+        ":SparseTensorOpsIncGen",
         ":StandardOps",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorTransforms",
     srcs = glob(["lib/Dialect/SparseTensor/Transforms/*.cpp"]),
deps = [
         ":StandardOpsTransforms",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )
 
-cc_library(
-    name = "SparseTensorUtils",
-    srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
-    hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),
-    includes = ["include"],
-    deps = [
-        ":IR",
-        "//llvm:Support",
-    ],
-)
-
 td_library(
     name = "StdOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/StandardOps/IR/Ops.td",
         "include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td",
     ],
     includes = ["include"],
     deps = [