[mlir][sparse] refactoring isAdmissibleTensorExp into codegen

[Peiming]

Summary

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**.

[Peiming]

fix style.

[Peiming]

cleanup.

[aartbik]

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

91

These section names correspond 1:1 with the comments in the header file.
You now introduce the "verify functions" section, but there is none in the header

So I would org them consistently

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

45

empty line after the general methods section

54

I would group initTensorExp/getTensorExp together

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

1444

is this a functional change?
(as part of the refactoring?)

[Peiming]

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

1444

Maybe? but not really? it was in isAdmissibleTensorExp that was called in the loop that iterates over sort masks, but this value is loop invariant, so i think it might be better to move it out to reject the kernel earlier. (See L568 in the original code)

[Peiming]

address comments.

[aartbik]

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

91

The order does not match the header. In the header we have this order

topo / verify / output

but in the CC you have

verify / topo / output

Probably easiest to just fix this in header?

131

here and below, no longer needed to call op()

136

here and below, no longer need to call merger()

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

110

/// style more consistent with rest

112

empty line before second comment/method

112

You need to switch the order here in the header

admissibleTensorExp appears before admissibleTopoOrder in CC file

(also, that is the logical order of calling, so I like to see that reflected here)

112

also, add

// Sets outerParNest on success with sparse output

113

Setting sparseOutput is also rather essential in making sure we can split it into two admissible methods, i.e. first admissable tensor exp then admissible topoorder. I had to think about that for a bit!

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

1444

Fair enough but, AFAIK, the "NFC" tag is really reserved for textual refactoring only, not for something that actually changes behavior (even if the overall behavior stays the same) ;-)

1458–1460

Keep a comment here on what it does

[Peiming]

address comments.

[aartbik]

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

1458–1460

Typo: Constructs

1459

add: returns failure if the tree cannot be build, or the tensor expression is inadmissible

[Peiming]

fix typo.

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

public:
   /// Constructs a code generation environment which can be
   /// passed around during sparsification for bookkeeping
   /// together with some consistency asserts.
   CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
              unsigned numTensors, unsigned numLoops, unsigned numFilterLoops);
 
   //
   // 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.
   std::optional<Operation *>
   genLoopBoundary(function_ref<
                   std::optional<Operation *>(MutableArrayRef<Value> parameters)>
                       callback);
 
   //
   // Merger delegates.
   //
 
   TensorExp &exp(unsigned e) { return latticeMerger.exp(e); }
   LatPoint &lat(unsigned l) { return latticeMerger.lat(l); }
   SmallVector<unsigned> &set(unsigned s) { return latticeMerger.set(s); }
   DimLevelType dlt(unsigned t, unsigned i) const {
     return latticeMerger.getDimLevelType(t, i);
   }
   DimLevelType dlt(unsigned b) const {
     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.
   //
 
   size_t topSortSize() const { return topSort.size(); }
   unsigned topSortAt(unsigned i) const { return topSort.at(i); }
   void topSortPushBack(unsigned i) { topSort.push_back(i); }
   void topSortClear(unsigned capacity = 0) {
     topSort.clear();
     topSort.reserve(capacity);
   }
 
   ArrayRef<unsigned> getTopSortSlice(size_t n, size_t m) const;
   ArrayRef<unsigned> getLoopCurStack() const;
   Value getLoopIdxValue(size_t loopIdx) const;
 
   //
   // Sparse tensor output and expansion methods.
   //
 
   bool hasSparseOutput() const { return sparseOut != nullptr; }
   bool isSparseOutput(OpOperand *o) const { return sparseOut == o; }
 
   Value getInsertionChain() const { return insChain; }
   void updateInsertionChain(Value chain);
 
   bool atExpandLevel(OpOperand *o, unsigned rank, unsigned lv) const;
   void startExpand(Value values, Value filled, Value added, Value count);
   bool isExpand() const { return expValues != nullptr; }
   void updateExpandCount(Value count);
   Value getExpandValues() const { return expValues; }
   Value getExpandFilled() const { return expFilled; }
private:
   Value expCount;
 
   // Bookkeeping for reductions (up-to-date value of the reduction, and indices
   // into the merger's expression tree. When the indices of a tensor reduction
   // expression are exhausted, all inner loops can use a scalarized reduction.
   Value redVal;
   unsigned redExp;
   unsigned redCustom;
+
+  // The root tensor expression of the kernel.
+  unsigned tensorExp;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENENV_H_

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

 //===- CodegenEnv.cpp -  Code generation environment class ----------------===//
 //
 // 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 "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
 //===----------------------------------------------------------------------===//
 
 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                        unsigned numTensors, unsigned numLoops,
                        unsigned numFilterLoops)
     : linalgOp(linop), sparseOptions(opts),
       latticeMerger(numTensors, numLoops, numFilterLoops), loopEmitter(),
       topSort(), sparseOut(nullptr), outerParNest(-1u), insChain(), expValues(),
       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);
   }
   // Initialize loop emitter.
   SmallVector<Value> tensors;
   for (OpOperand &t : linalgOp->getOpOperands())
     tensors.push_back(t.get());
std::optional<Operation *> CodegenEnv::genLoopBoundary(
   if (isExpand())
     updateExpandCount(params[i++]);
   if (insChain != nullptr)
     updateInsertionChain(params[i]);
   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
 //===----------------------------------------------------------------------===//
 
 ArrayRef<unsigned> CodegenEnv::getTopSortSlice(size_t n, size_t m) const {
   return ArrayRef<unsigned>(topSort).slice(n, m);
 }
 
 ArrayRef<unsigned> CodegenEnv::getLoopCurStack() const {

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

 static unsigned getNumCompoundAffineOnSparseDims(linalg::GenericOp op) {
   unsigned num = 0;
   for (OpOperand &t : op->getOpOperands())
     num += getNumCompoundAffineOnSparseDims(op.getMatchingIndexingMap(&t),
                                             t.get());
   return num;
 }
 
+static bool hasCompoundAffineOnSparseOut(linalg::GenericOp op) {
+  OpOperand *out = op.getDpsInitOperand(0);
+  auto enc = getSparseTensorEncoding(out->get().getType());
+  if (!enc || enc.isAllDense())
+    return false;
+
+  return getNumCompoundAffineOnSparseDims(op.getMatchingIndexingMap(out),
+                                          out->get());
+}
+
 /// Helper method to inspect sparse encodings in the tensor types.
 /// Fills the per-dimension sparsity information for all tensors.
 /// Returns true if the sparse annotations and affine subscript
 /// expressions of all tensors are admissible. Returns false if
 /// no annotations are found or inadmissible constructs occur.
 static bool findSparseAnnotations(CodegenEnv &env) {
   bool annotated = false;
   unsigned filterLdx = env.merger().getFilterLoopStartingIdx();
for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
         // above.
         addAffineOrderings(adjM, inDegree, fa, ta, fldx, tldx);
       }
     }
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     if (mask & SortMask::kIncludeUndef) {
       unsigned tensor = t.getOperandNumber();
-      for (unsigned i = 0; i < n; i++)
+      for (unsigned i = 0; i < n; i++) {
         if (isCompressedDLT(env.dlt(tensor, i)) ||
             isSingletonDLT(env.dlt(tensor, i))) {
           for (unsigned j = 0; j < n; j++)
             if (isUndefDLT(env.dlt(tensor, j))) {
               adjM[i][j] = true;
               inDegree[j]++;
             }
         } else {
           assert(isDenseDLT(env.dlt(tensor, i)) ||
                  isUndefDLT(env.dlt(tensor, i)));
         }
       }
     }
+  }
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
   env.topSortClear(n);
   return topSortOptimal(env, n, iteratorTypes, inDegree, adjM);
 }
 
-/// Returns true if tensor materializes uninitialized into the computation.
-static bool isMaterializing(Value val) {
-  return val.getDefiningOp<tensor::EmptyOp>() ||
-         val.getDefiningOp<bufferization::AllocTensorOp>();
-}
-
-/// Returns true when the tensor expression is admissible for codegen.
-/// Since all sparse input tensors are admissible, we just need to check
-/// whether the out tensor in the tensor expression codegen is admissible.
-/// Sets `sparseOut` to the tensor and `outerParNest` to the outer injective
-/// nesting depth when a "truly dynamic" sparse tensor output occurs.
-static bool isAdmissibleTensorExp(CodegenEnv &env, unsigned exp,
-                                  OpOperand **sparseOut,
-                                  unsigned *outerParNest) {
-  // 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 : env.op().getIteratorTypesArray()) {
-    if (it == utils::IteratorType::reduction) {
-      if (env.merger().hasNegateOnOut(exp))
-        return false;
-      break;
-    }
-  }
-
-  OpOperand *lhs = env.op().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)
-    return true;
-  // An all-dense annotated "sparse" output tensor becomes a linearized random
-  // access 1-dim memref. Also admissible since insertions cannot occur.
-  bool allDense = true;
-  unsigned numLoops =
-      env.merger().getNumLoops(); // numNativeLoops + numFilterLoops
-  for (unsigned i = 0; i < env.merger().getNumLoops(); i++)
-    if (isCompressedDLT(env.dlt(tensor, i)) ||
-        isSingletonDLT(env.dlt(tensor, i))) {
-      allDense = false;
-      break;
-    } else {
-      assert(isDenseDLT(env.dlt(tensor, i)) || isUndefDLT(env.dlt(tensor, i)));
-    }
-  if (allDense)
-    return true;
-
-  // TODO: support compound affine expression on sparse output.
-  if (getNumCompoundAffineOnSparseDims(env.op().getMatchingIndexingMap(lhs),
-                                       lhs->get()) != 0)
-    return false;
-
-  // 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 (env.merger().isSingleCondition(tensor, exp))
-    return true;
-
-  // Accept "truly dynamic" if the output tensor materializes uninitialized
-  // into the computation and insertions occur in lexicographic index order.
-  if (isMaterializing(lhs->get())) {
-    auto iteratorTypes = env.op().getIteratorTypesArray();
-    unsigned nest = 0;
-    for (unsigned i = 0; i < numLoops; i++) {
-      if (!env.merger().isFilterLoop(env.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[env.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 >= env.op().getRank(lhs) - 1) {
-      *sparseOut = lhs;
-      *outerParNest = nest;
-      return true;
-    }
-  }
-  return false;
-}
-
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
 /// Local bufferization of all dense and sparse data structures.
 static void genBuffers(CodegenEnv &env, OpBuilder &builder) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
static bool translateBitsToTidDimPairs(
     CodegenEnv &env, unsigned li, unsigned idx, SmallVectorImpl<size_t> &tids,
     SmallVectorImpl<size_t> &dims, SmallVectorImpl<size_t> &affineTids,
     SmallVectorImpl<size_t> &affineDims, SmallVectorImpl<AffineExpr> &exps) {
   const BitVector &all = env.lat(li).bits;
   const BitVector &simple = env.lat(li).simple;
 
   unsigned numloopCond = 0;
   // Converts bits to array + dim pair
-  env.merger().foreachTidDimPairInBits(all, [&,
-                                             idx](unsigned b, unsigned tid,
-                                                  std::optional<unsigned> dim,
-                                                  DimLevelType dlt) {
+  env.merger().foreachTidDimPairInBits(
+      all, [&, idx](unsigned b, unsigned tid, std::optional<unsigned> dim,
         if (simple.test(b)) {
           if (isUndefDLT(dlt)) {
-        // An undefined dlt in the lattices, we probably mean to iterate based
-        // on the dim of output tensor.
+            // An undefined dlt in the lattices, we probably mean to iterate
+            // based on the dim of output tensor.
             // E.g., this could be a synthetic tensor (for invariants and sparse
             // output tensor).
             // out[i][j] = invariant; or a broadcast
             // out[i][j] = in[i] (j is undef for input)
             tid = env.merger().getOutTensorID();
             dim = env.merger().getDimNum(tid, idx);
             // Skips invalid dim (e.g., when this is a zero ranked tensor).
             if (!dim)
               return;
           }
           tids.push_back(tid);
           dims.push_back(*dim);
           numloopCond++;
         } else if (isDenseDLT(dlt)) {
           tids.push_back(tid);
           dims.push_back(*dim);
         } else {
           assert(isUndefDLT(dlt));
           linalg::GenericOp op = env.op();
           if (tid >= op.getNumDpsInputs())
             // We only handle affine expression on input tensors (for now).
             return;
           OpOperand *operand = &op->getOpOperand(tid);
           auto enc = getSparseTensorEncoding(operand->get().getType());
           // Non-annotated dense tensors requires no special handling.
           if (!enc)
             return;
 
           ArrayRef<AffineExpr> affines =
               op.getMatchingIndexingMap(operand).getResults();
           assert(affines.size() == enc.getDimLevelType().size());
           for (unsigned i = 0, e = affines.size(); i < e; i++) {
             AffineExpr exp = affines[toOrigDim(enc, i)];
             // Skip simple affine expression and non dense dimensions (which has
             // it own filter loop).
-        if (exp.isa<AffineDimExpr>() || !isDenseDLT(getDimLevelType(enc, i)))
+            if (exp.isa<AffineDimExpr>() ||
+                !isDenseDLT(getDimLevelType(enc, i)))
               continue;
 
             // Constant affine expression are handled in genLoop
             if (!exp.isa<AffineConstantExpr>()) {
               bool atLevel = false;
               if (isInvariantAffine(env, exp, idx, atLevel) && atLevel) {
                 // If the compound affine is invariant and we are right at the
-            // level. We need to generate the address according to the affine
-            // expression. This is also the best place we can do it to avoid
-            // putting it inside inner loops.
+                // level. We need to generate the address according to the
+                // affine expression. This is also the best place we can do it
+                // to avoid putting it inside inner loops.
                 // NOTE: It assumes that the levels of the input tensor are
                 // initialized in order (and it is also currently guaranteed by
-            // computeIterationGraph), another more admissible approach might be
-            // accepting out-of-order access between consecutive dense levels.
+                // computeIterationGraph), another more admissible approach
+                // might be accepting out-of-order access between consecutive
+                // dense levels.
                 affineTids.push_back(tid);
                 affineDims.push_back(i);
                 exps.push_back(exp);
               }
             }
           }
         }
       });
 
   if (isDenseDLT(env.dlt(env.merger().getOutTensorID(), idx))) {
     // Note that we generate dense indices of the output tensor
     // unconditionally, since they may not appear in the lattice, but may be
     // needed for linearized env.
     auto dim = *env.merger().getDimNum(env.merger().getOutTensorID(), idx);
     tids.push_back(env.merger().getOutTensorID());
     dims.push_back(dim);

 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
 public:
   GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
       : OpRewritePattern<linalg::GenericOp>(context), options(o) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
-    // Only accept single output operations.
-    if (op.getNumDpsInits() != 1)
+    // Only accept single output operations without affine index on sparse
+    // output.
+    if (op.getNumDpsInits() != 1 || hasCompoundAffineOnSparseOut(op))
       return failure();
 
     // Sets up a code generation environment.
     unsigned numTensors = op->getNumOperands();
     unsigned numLoops = op.getNumLoops();
     unsigned numFilterLoops = getNumCompoundAffineOnSparseDims(op);
     CodegenEnv env(op, options, numTensors, numLoops, numFilterLoops);
 
     // Detects sparse annotations and translates the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     if (!findSparseAnnotations(env))
       return failure();
 
-    // Builds the tensor expression for the Linalg operation in SSA form.
-    std::optional<unsigned> optExp = env.merger().buildTensorExpFromLinalg(op);
-    if (!optExp)
+    // Constructs the tensor expressions tree from `op`, returns failure if the
+    // tree can not be built or the tensor expression is inadmissible.
+    if (failed(env.initTensorExp()))
       return failure();
-    unsigned exp = *optExp;
 
     // Computes a topologically sorted iteration graph to ensure tensors
     // are visited in natural index order. Gradually relaxes the considered
     // constraints until an acyclic iteration graph results, such that sparse
     // code generation can proceed. As a last resort, an attempt is made
     // to resolve cycles by inserting a conversion.
     bool isAdmissible = false;
     bool hasCycle = true;
-    OpOperand *sparseOut = nullptr;
-    unsigned outerParNest = -1u;
+
     // An const list of all masks that we used for interation graph
     // computation. Must be ordered from more strict to less strict.
     const auto allMask = {SortMask::kIncludeAll, SortMask::kIncludeUndef,
                           SortMask::kIncludeDense, SortMask::kSparseOnly};
-    for (auto mask : allMask)
+    for (auto mask : allMask) {
       if (computeIterationGraph(env, mask)) {
         hasCycle = false;
-        if (isAdmissibleTensorExp(env, exp, &sparseOut, &outerParNest)) {
+        if (env.isAdmissibleTopoOrder()) {
           isAdmissible = true;
           break;
         }
         // else try a set of less strict constraints.
       }
+    }
     if (hasCycle)
       return resolveCycle(env, rewriter); // one last shot
     if (!isAdmissible)
       return failure(); // inadmissible expression, reject
 
     // Recursively generates code if admissible.
-    env.startEmit(sparseOut, outerParNest);
+    env.startEmit();
     genBuffers(env, rewriter);
     genInitConstantDenseAddress(env, rewriter);
-    genStmt(env, rewriter, exp, 0);
+    genStmt(env, rewriter, env.getTensorExp(), 0);
     genResult(env, rewriter);
     return success();
   }
 
 private:
   // Last resort cycle resolution.
   LogicalResult resolveCycle(CodegenEnv &env, PatternRewriter &rewriter) const {
     // Compute topological sort while leaving out every