[MLIR][Sparse] Move `buildLattices` into Merger

[gussmith23]

Summary

This allows us to use buildLattices in the Merger unittests.

[aartbik]

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

152

Add: Returns index of the root expression.

since it is now public API

153

Since we are pulling this out, perhaps we should use the Linalg in the name, something like "buildLatticesFromLinalgOp". That way, we can later add other "builders"

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

140–141

we don't generally repeat the comment from header in cc file

(style in a nutshell: header comments explain API, CC comments explain implementation)

[aartbik]

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

153

I was actually confusing my feedback with buildTensorExp(), which perhaps should be moved into the merger as well.

For this case, it only uses the linalg op for getNumInputsAndOutputs(), which the merger already knows about. So I think you can simplify the API and not rely on linalg at this point

[gussmith23]

Address Aart's comments

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

153

Removed the dependence on the GenericOp!

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

140–141

Whoops, I forgot to delete this! Thanks for catching.

[aartbik]

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

147

I was thinking, since we have this in its own self-contained file and class now,
shall we add a new constant

const unsigned syntheticTensor;

and initialize that to t in the constructor? This will read a bit nice than all the numTensors - 1 an outPut tensor + 1 we have at other places ;-)

[gussmith23]

Address Aart's comment -- add syntheticTensor

[gussmith23]

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

147

Added! I looked for other places to use syntheticTensor but this seems to be the only place we need to use it.

[gussmith23]

Rebase

[aartbik]

Looking good! Ship it!

[gussmith23]

Rebase

[gussmith23]

Rebase

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

 class Merger {
 public:
   /// Constructs a merger for the given number of tensors and 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.
   Merger(unsigned t, unsigned l)
-      : outTensor(t - 1), numTensors(t + 1), numLoops(l),
+      : outTensor(t - 1), syntheticTensor(t), numTensors(t + 1), numLoops(l),
         dims(t + 1, std::vector<Dim>(l, Dim::kUndef)) {}
 
   /// Adds a tensor expression. Returns its index.
   unsigned addExp(Kind k, unsigned e0, unsigned e1 = -1u, Value v = Value());
   unsigned addExp(Kind k, Value v) { return addExp(k, -1u, -1u, v); }
 
   /// Adds an iteration lattice point. Returns its index.
   unsigned addLat(unsigned t, unsigned i, unsigned e);
public:
   bool isDim(unsigned t, unsigned i, Dim d) const {
     assert(t < numTensors && i < numLoops);
     return dims[t][i] == d;
   }
 
   /// Returns true if any set bit corresponds to queried dim.
   bool hasAnyDimOf(const llvm::BitVector &bits, Dim d) const;
 
+  /// Builds the iteration lattices in a bottom-up traversal given the remaining
+  /// tensor (sub)expression and the next loop index in the iteration graph.
+  /// Returns index of the root expression.
+  unsigned buildLattices(unsigned exp, unsigned idx);
+
   /// 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 syntheticTensor;
   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/Transforms/Sparsification.cpp

static bool isAdmissableTensorExp(Merger &merger, linalg::GenericOp op,
   // [Bik96,Ch9], is also admissable without special codegen.
   if (isConjunction(merger, tensor, exp))
     return true;
   // Reject for now since this requires changes to the nonzero structure.
   // TODO: implement "workspaces" [Kjolstad2019]
   return false;
 }
 
-/// Builds the iteration lattices in a bottom-up traversal given the remaining
-/// tensor (sub)expression and the next loop index in the iteration graph.
-static unsigned buildLattices(Merger &merger, linalg::GenericOp op,
-                              unsigned exp, unsigned idx) {
-  Kind kind = merger.exp(exp).kind;
-  if (kind == Kind::kTensor || kind == Kind::kInvariant) {
-    // Either the index is really used in the tensor expression, or it is
-    // set to the undefined index in that dimension. An invariant expression
-    // is set to a synthetic tensor with undefined indices only.
-    unsigned s = merger.addSet();
-    unsigned t = kind == Kind::kTensor ? merger.exp(exp).e0
-                                       : op.getNumInputsAndOutputs();
-    merger.set(s).push_back(merger.addLat(t, idx, exp));
-    return s;
-  }
-  unsigned s0 = buildLattices(merger, op, merger.exp(exp).e0, idx);
-  unsigned s1 = buildLattices(merger, op, merger.exp(exp).e1, idx);
-  switch (kind) {
-  case Kind::kTensor:
-  case Kind::kInvariant:
-    llvm_unreachable("handled above");
-  case Kind::kMulF:
-  case Kind::kMulI:
-    return merger.takeConj(kind, s0, s1);
-  case Kind::kAddF:
-  case Kind::kAddI:
-    return merger.takeDisj(kind, s0, s1);
-  }
-  llvm_unreachable("unexpected expression kind");
-}
-
 /// Maps sparse integer option to actual integral storage type.
 static Type genIntType(PatternRewriter &rewriter, unsigned width) {
   if (width == 0)
     return rewriter.getIndexType();
   return rewriter.getIntegerType(width);
 }
 
 /// Detects in-place annotation on tensor argument.
static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
   assert(codegen.curVecLength == 1);
 
   // Construct iteration lattices for current loop index, with L0 at top.
   // Then emit initialization code for the loop sequence at this level.
   // We maintain the universal dense index if dense indices are still
   // in play for a non-singleton loop sequence.
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
-  unsigned lts = merger.optimizeSet(buildLattices(merger, op, exp, idx));
+  unsigned lts = merger.optimizeSet(merger.buildLattices(exp, idx));
   unsigned lsize = merger.set(lts).size();
   assert(lsize != 0);
   unsigned l0 = merger.set(lts)[0];
   unsigned ldx = at == 0 ? -1u : topSort[at - 1];
   genInvariants(merger, codegen, rewriter, op, exp, ldx, /*hoist=*/true);
   bool needsUniv = false;
   if (genInit(merger, codegen, rewriter, op, topSort, at,
               merger.lat(l0).bits)) {

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

 
 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;
 }
 
+unsigned Merger::buildLattices(unsigned e, unsigned idx) {
+  Kind kind = exp(e).kind;
+  if (kind == Kind::kTensor || kind == Kind::kInvariant) {
+    // Either the index is really used in the tensor expression, or it is
+    // set to the undefined index in that dimension. An invariant expression
+    // is set to a synthetic tensor with undefined indices only.
+    unsigned s = addSet();
+    unsigned t = kind == Kind::kTensor ? exp(e).e0 : syntheticTensor;
+    set(s).push_back(addLat(t, idx, e));
+    return s;
+  }
+  unsigned s0 = buildLattices(exp(e).e0, idx);
+  unsigned s1 = buildLattices(exp(e).e1, idx);
+  switch (kind) {
+  case Kind::kTensor:
+  case Kind::kInvariant:
+    llvm_unreachable("handled above");
+  case Kind::kMulF:
+  case Kind::kMulI:
+    return takeConj(kind, s0, s1);
+  case Kind::kAddF:
+  case Kind::kAddI:
+    return takeDisj(kind, s0, s1);
+  }
+  llvm_unreachable("unexpected expression kind");
+}
+
 #ifndef NDEBUG
 
 //
 // Print methods (for debugging).
 //
 
 void Merger::dumpExp(unsigned e) const {
   switch (tensorExps[e].kind) {