[mlir][sparse] refactoring isAdmissibleTensorExp into codegen

This patch moves some utils into CodegenEnv class, it should make the code easier to follow and it eliminates several indirect value assignment that use `ptr**`.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D142040

Author: Peiming Liu

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp

@@ -7,11 +7,25 @@
 //===----------------------------------------------------------------------===//
 
 #include "CodegenEnv.h"
+
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
+//===----------------------------------------------------------------------===//
+// Code generation environment helper functions
+//===----------------------------------------------------------------------===//
+
+/// Returns true if tensor materializes uninitialized into the computation.
+static bool isMaterializing(Value val) {
+  return val.getDefiningOp<tensor::EmptyOp>() ||
+         val.getDefiningOp<bufferization::AllocTensorOp>();
+}
+
 //===----------------------------------------------------------------------===//
 // Code generation environment constructor and general methods
 //===----------------------------------------------------------------------===//
@@ -25,11 +39,18 @@ CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
       expFilled(), expAdded(), expCount(), redVal(), redExp(-1u),
       redCustom(-1u) {}
 
-void CodegenEnv::startEmit(OpOperand *so, unsigned lv) {
-  assert(sparseOut == nullptr && insChain == nullptr &&
-         "must only start emitting once");
-  sparseOut = so;
-  outerParNest = lv;
+LogicalResult CodegenEnv::initTensorExp() {
+  // Builds the tensor expression for the Linalg operation in SSA form.
+  std::optional<unsigned> optExp = latticeMerger.buildTensorExpFromLinalg(op());
+  if (!optExp || !isAdmissibleTensorExp(*optExp))
+    return failure();
+
+  tensorExp = *optExp;
+  return success();
+}
+
+void CodegenEnv::startEmit() {
+  assert(insChain == nullptr && "must only start emitting once");
   if (sparseOut) {
     insChain = sparseOut->get();
     latticeMerger.setHasSparseOut(true);
@@ -66,6 +87,71 @@ std::optional<Operation *> CodegenEnv::genLoopBoundary(
   return r;
 }
 
+//===----------------------------------------------------------------------===//
+// Code generation environment verify functions.
+//===----------------------------------------------------------------------===//
+
+bool CodegenEnv::isAdmissibleTensorExp(unsigned exp) {
+  // We reject any expression that makes a reduction from `-outTensor`, as those
+  // expressions create a dependency between the current iteration (i) and the
+  // previous iteration (i-1). It would require iterating over the whole
+  // coordinate space, which prevent exploiting sparsity for faster code.
+  for (utils::IteratorType it : linalgOp.getIteratorTypesArray()) {
+    if (it == utils::IteratorType::reduction) {
+      if (latticeMerger.hasNegateOnOut(exp))
+        return false;
+      break;
+    }
+  }
+
+  OpOperand *lhs = linalgOp.getDpsInitOperand(0);
+  unsigned tensor = lhs->getOperandNumber();
+  auto enc = getSparseTensorEncoding(lhs->get().getType());
+  // An non-annotated output tensor is assumed dense, and becomes a random
+  // access n-dim memref. Admissible since insertions cannot occur.
+  if (!enc || enc.isAllDense())
+    return true;
+
+  // A tensor expression with a sparse output tensor that changes its values
+  // but not its nonzero structure, an operation called "simply dynamic" in
+  // [Bik96,Ch9], is also admissible without special env.
+  if (latticeMerger.isSingleCondition(tensor, exp))
+    return true;
+
+  // Accept "truly dynamic" if the output tensor materializes uninitialized
+  // into the computation and insertions occur in lexicographic index order.
+  sparseOut = lhs;
+  return isMaterializing(lhs->get());
+}
+
+bool CodegenEnv::isAdmissibleTopoOrder() {
+  if (!hasSparseOutput())
+    return true;
+
+  OpOperand *lhs = linalgOp.getDpsInitOperand(0);
+  // Accept "truly dynamic" if the output tensor materializes uninitialized
+  // into the computation and insertions occur in lexicographic index order.
+  unsigned nest = 0;
+  auto iteratorTypes = linalgOp.getIteratorTypesArray();
+  for (unsigned i = 0, e = latticeMerger.getNumLoops(); i < e; i++) {
+    if (!latticeMerger.isFilterLoop(topSortAt(i))) {
+      // We only count non-filter loops as filter loops should be considered
+      // as a special type of parallel loops.
+      if (linalg::isReductionIterator(iteratorTypes[topSortAt(i)]))
+        break; // terminate at first reduction
+      nest++;
+    }
+  }
+  // Determine admissible dynamic insertion situations:
+  // (1) fully injective, since there are no reductions,
+  // (2) admissible 1-d expansion in innermost dimension.
+  if (nest >= linalgOp.getRank(lhs) - 1) {
+    outerParNest = nest;
+    return true;
+  }
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // Code generation environment topological sort methods
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.h

@@ -44,12 +44,15 @@ class CodegenEnv {
   // General methods.
   //
 
+  LogicalResult initTensorExp();
+  unsigned getTensorExp() const { return tensorExp; }
+
   linalg::GenericOp op() const { return linalgOp; }
   const SparsificationOptions &options() const { return sparseOptions; }
   Merger &merger() { return latticeMerger; }
   LoopEmitter &emitter() { return loopEmitter; }
 
-  void startEmit(OpOperand *so, unsigned lv);
+  void startEmit();
 
   /// Generates loop boundary statements (entering/exiting loops). The function
   /// passes and updates the passed-in parameters.
@@ -72,6 +75,18 @@ class CodegenEnv {
     return latticeMerger.getDimLevelType(b);
   }
 
+  //
+  // Code generation environment verify functions.
+  //
+
+  /// Whether the tensor expression is admissible for codegen.
+  /// It also sets the sparseOut if the output tensor is sparse.
+  bool isAdmissibleTensorExp(unsigned exp);
+
+  /// Whether the iteration graph is sorted in admissible topoOrder.
+  /// Sets outerParNest on success with sparse output
+  bool isAdmissibleTopoOrder();
+
   //
   // Topological delegate and sort methods.
   //
@@ -156,6 +171,9 @@ class CodegenEnv {
   Value redVal;
   unsigned redExp;
   unsigned redCustom;
+
+  // The root tensor expression of the kernel.
+  unsigned tensorExp;
 };
 
 } // namespace sparse_tensor