[mlir][sparse] move tensor expression builder into Merger utility

[aartbik]

Summary

Rationale:
Follow-up on migrating lattice and tensor expression related methods into the new utility.
This also prepares the next step of generalizing the op kinds that are handled.

[aartbik]

edit

[aartbik]

fixed lint issue

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

 //===- Merger.h - Utilities for defining lattices ---------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines utilities for dealing with iteration lattices.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_
 #define MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_
 
+#include "mlir/Dialect/Linalg/IR/LinalgOps.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 };
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
 
+  /// 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);
+
+  /// Builds a tensor expression from the given Linalg operation.
+  /// Returns index of the root expression on success.
+  Optional<unsigned> buildTensorExpFromLinalg(linalg::GenericOp op);
+
 private:
+  /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
+  Optional<unsigned> buildTensorExp(linalg::GenericOp op, Value val);
+
   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 computeIterationGraph(Merger &merger, linalg::GenericOp op,
   for (unsigned i = 0; i < n; i++)
     if (visit[i] == 0)
       if (!topSortDFS(i, visit, topSort, adjM))
         return false; // cycle!
   std::reverse(std::begin(topSort), std::end(topSort));
   return true;
 }
 
-/// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
-/// This simplifies constructing (sub)expressions during iteration lattice
-/// building (compared to using the SSA representation everywhere).
-static Optional<unsigned> buildTensorExp(Merger &merger, linalg::GenericOp op,
-                                         Value val) {
-  if (auto arg = val.dyn_cast<BlockArgument>()) {
-    unsigned argN = arg.getArgNumber();
-    // Any argument of the generic op that is not marked as a scalar
-    // argument is considered a tensor, indexed by the implicit loop
-    // bounds. This includes rank-0 tensor arguments.
-    if (arg.getOwner()->getParentOp() == op) {
-      OpOperand *t = op.getInputAndOutputOperands()[argN];
-      if (!op.isScalar(t))
-        return merger.addExp(Kind::kTensor, argN);
-      val = t->get(); // get scalar value
-    }
-    // Any other argument (marked as scalar argument for the generic op
-    // or belonging to an enveloping op) is considered invariant.
-    return merger.addExp(Kind::kInvariant, val);
-  }
-  Operation *def = val.getDefiningOp();
-  if (def->getBlock() != &op.region().front()) {
-    // Something defined outside is invariant.
-    return merger.addExp(Kind::kInvariant, val);
-  } else if (def->getNumOperands() == 2) {
-    // Construct binary operations if subexpressions could be built.
-    auto x = buildTensorExp(merger, op, def->getOperand(0));
-    auto y = buildTensorExp(merger, op, def->getOperand(1));
-    if (x.hasValue() && y.hasValue()) {
-      unsigned e0 = x.getValue();
-      unsigned e1 = y.getValue();
-      if (isa<MulFOp>(def))
-        return merger.addExp(Kind::kMulF, e0, e1);
-      if (isa<MulIOp>(def))
-        return merger.addExp(Kind::kMulI, e0, e1);
-      if (isa<AddFOp>(def))
-        return merger.addExp(Kind::kAddF, e0, e1);
-      if (isa<AddIOp>(def))
-        return merger.addExp(Kind::kAddI, e0, e1);
-    }
-  }
-  // Cannot build (yet).
-  return None;
-}
-
 /// Returns true if given tensor co-iterates with conjunction only.
 /// For the output tensor, this defines a "simply dynamic" operation.
 /// For instance: A(I) = A(I) * B(I) * C(I)
 static unsigned isConjunction(Merger &merger, unsigned tensor, unsigned exp) {
   switch (merger.exp(exp).kind) {
   case Kind::kTensor:
     return merger.exp(exp).e0 == tensor;
   case Kind::kMulF:
LogicalResult matchAndRewrite(linalg::GenericOp op,
     // tensors are visited in natural index order. Fails on cycles.
     // This assumes that higher-level passes have already put the
     // tensors in each tensor expression in a feasible order.
     std::vector<unsigned> topSort;
     if (!computeIterationGraph(merger, op, topSort, /*sparseOnly=*/false) &&
         !computeIterationGraph(merger, op, topSort, /*sparseOnly=*/true))
       return failure();
 
-    // Finds the terminating yield statement and builds the tensor
-    // expression for the Linalg operation in SSA form.
-    Operation *yield = op.region().front().getTerminator();
-    Optional<unsigned> exp = buildTensorExp(merger, op, yield->getOperand(0));
+    // Builds the tensor expression for the Linalg operation in SSA form.
+    Optional<unsigned> exp = merger.buildTensorExpFromLinalg(op);
     if (!exp.hasValue())
-      return failure(); // build failure
+      return failure();
 
-    // Reject an inadmissable tensor expression.
+    // Rejects an inadmissable tensor expression.
     if (!isAdmissableTensorExp(merger, op, exp.getValue()))
       return failure();
 
     // Recursively generates code.
     CodeGen codegen(options, numTensors, numLoops);
     if (!genBuffers(merger, codegen, rewriter, op))
       return failure(); // could not bufferize
     genStmt(merger, codegen, rewriter, op, topSort, exp.getValue(), 0);

mlir/lib/Dialect/SparseTensor/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorUtils
   Merger.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRLinalg
 )

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 {
 
+//
+// Lattice methods.
+//
+
 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;
 }
 
 unsigned Merger::addLat(unsigned t, unsigned i, unsigned e) {
   assert(t < numTensors && i < numLoops);
unsigned Merger::optimizeSet(unsigned s0) {
   unsigned s = addSet();
   assert(latSets[s0].size() != 0);
   unsigned p0 = latSets[s0][0];
   for (unsigned p1 : latSets[s0]) {
     bool add = true;
     if (p0 != p1) {
       // Is this a straightforward copy?
       unsigned e = latPoints[p1].exp;
-      if (exp(e).kind == Kind::kTensor && exp(e).e0 == outTensor)
+      if (tensorExps[e].kind == Kind::kTensor && tensorExps[e].e0 == outTensor)
         continue;
       // Conjunction already covered?
       for (unsigned p2 : latSets[s]) {
         assert(!latGT(p1, p2)); // Lj => Li would be bad
         if (onlyDenseDiff(p2, p1)) {
           add = false;
           break;
         }

 
 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) {
   case Kind::kTensor:
+    if (tensorExps[e].e0 == syntheticTensor)
+      llvm::dbgs() << "synthetic_";
+    else if (tensorExps[e].e0 == outTensor)
+      llvm::dbgs() << "output_";
     llvm::dbgs() << "tensor_" << tensorExps[e].e0;
     break;
   case Kind::kInvariant:
     llvm::dbgs() << "invariant";
     break;
   default:
   case Kind::kMulI:
     llvm::dbgs() << "(";
if (bits[b]) {
         break;
       }
     }
   }
 }
 
 #endif // NDEBUG
 
+//
+// Builder methods.
+//
+
+unsigned Merger::buildLattices(unsigned e, unsigned idx) {
+  Kind kind = tensorExps[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 ? tensorExps[e].e0 : syntheticTensor;
+    latSets[s].push_back(addLat(t, idx, e));
+    return s;
+  }
+  unsigned s0 = buildLattices(tensorExps[e].e0, idx);
+  unsigned s1 = buildLattices(tensorExps[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");
+}
+
+Optional<unsigned> Merger::buildTensorExpFromLinalg(linalg::GenericOp op) {
+  Operation *yield = op.region().front().getTerminator();
+  return buildTensorExp(op, yield->getOperand(0));
+}
+
+Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value val) {
+  if (auto arg = val.dyn_cast<BlockArgument>()) {
+    unsigned argN = arg.getArgNumber();
+    // Any argument of the generic op that is not marked as a scalar
+    // argument is considered a tensor, indexed by the implicit loop
+    // bounds. This includes rank-0 tensor arguments.
+    if (arg.getOwner()->getParentOp() == op) {
+      OpOperand *t = op.getInputAndOutputOperands()[argN];
+      if (!op.isScalar(t))
+        return addExp(Kind::kTensor, argN);
+      val = t->get(); // get scalar value
+    }
+    // Any other argument (marked as scalar argument for the generic op
+    // or belonging to an enveloping op) is considered invariant.
+    return addExp(Kind::kInvariant, val);
+  }
+  // Something defined outside is invariant.
+  Operation *def = val.getDefiningOp();
+  if (def->getBlock() != &op.region().front())
+    return addExp(Kind::kInvariant, val);
+  // Construct binary operations if subexpressions could be built.
+  if (def->getNumOperands() == 2) {
+    auto x = buildTensorExp(op, def->getOperand(0));
+    auto y = buildTensorExp(op, def->getOperand(1));
+    if (x.hasValue() && y.hasValue()) {
+      unsigned e0 = x.getValue();
+      unsigned e1 = y.getValue();
+      if (isa<MulFOp>(def))
+        return addExp(Kind::kMulF, e0, e1);
+      if (isa<MulIOp>(def))
+        return addExp(Kind::kMulI, e0, e1);
+      if (isa<AddFOp>(def))
+        return addExp(Kind::kAddF, e0, e1);
+      if (isa<AddIOp>(def))
+        return addExp(Kind::kAddI, e0, e1);
+    }
+  }
+  // Cannot build.
+  return None;
+}
+
 } // namespace sparse_tensor
 } // namespace mlir