[mlir][sparse] move emitter ownership into environment

[aartbik]

Summary

last bits and pieces of the environment refactoring

[aartbik]

clang format

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

public:
 
   //
   // General methods.
   //
 
   linalg::GenericOp op() const { return linalgOp; }
   const SparsificationOptions &options() const { return sparseOptions; }
   Merger &merger() { return latticeMerger; }
-  LoopEmitter *emitter() { return loopEmitter; }
+  LoopEmitter &emitter() { return loopEmitter; }
 
-  void startEmit(OpOperand *so, unsigned lv, LoopEmitter *le);
+  void startEmit(OpOperand *so, unsigned lv);
 
   /// Generates loop boundary statements (entering/exiting loops). The function
   /// passes and updates the passed-in parameters.
   Optional<Operation *> genLoopBoundary(
       function_ref<Optional<Operation *>(MutableArrayRef<Value> parameters)>
           callback);
 
   //
public:
   DimLevelType dlt(unsigned b) const {
     return latticeMerger.getDimLevelType(b);
   }
 
   //
   // Topological delegate and sort methods.
   //
 
-  // TODO: get rid of this one!
-  std::vector<unsigned> &topSortRef() { return topSort; }
-
   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);
   }
 
private:
   linalg::GenericOp linalgOp;
 
   // Sparsification options.
   SparsificationOptions sparseOptions;
 
   // Merger helper class.
   Merger latticeMerger;
 
-  // Loop emitter helper class (keep reference in scope!).
-  // TODO: move emitter constructor up in time?
-  LoopEmitter *loopEmitter;
+  // Loop emitter helper class.
+  LoopEmitter loopEmitter;
 
   // Topological sort.
   std::vector<unsigned> topSort;
 
   // Sparse tensor as output. Implemented either through direct injective
   // insertion in lexicographic index order or through access pattern
   // expansion in the innermost loop nest (`expValues` through `expCount`).
   OpOperand *sparseOut;

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

 //===----------------------------------------------------------------------===//
 // 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(nullptr),
+      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, LoopEmitter *le) {
-  assert(sparseOut == nullptr && loopEmitter == nullptr &&
-         insChain == nullptr && "must only start emitting once");
+void CodegenEnv::startEmit(OpOperand *so, unsigned lv) {
+  assert(sparseOut == nullptr && insChain == nullptr &&
+         "must only start emitting once");
   sparseOut = so;
   outerParNest = lv;
-  loopEmitter = le;
   if (sparseOut) {
     insChain = sparseOut->get();
     latticeMerger.setHasSparseOut(true);
   }
+  // Initialize loop emitter.
+  SmallVector<Value> tensors;
+  for (OpOperand &t : linalgOp->getOpOperands())
+    tensors.push_back(t.get());
+  loopEmitter.initialize(tensors,
+                         StringAttr::get(linalgOp.getContext(),
+                                         linalg::GenericOp::getOperationName()),
+                         /*hasOutput=*/true,
+                         /*isSparseOut=*/sparseOut != nullptr, topSort);
 }
 
 Optional<Operation *> CodegenEnv::genLoopBoundary(
     function_ref<Optional<Operation *>(MutableArrayRef<Value> parameters)>
         callback) {
   SmallVector<Value> params;
   if (isReduc())
     params.push_back(redVal);

 // 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 {
-  return getTopSortSlice(0, loopEmitter->getCurrentDepth());
+  return getTopSortSlice(0, loopEmitter.getCurrentDepth());
 }
 
 Value CodegenEnv::getLoopIdxValue(size_t loopIdx) const {
   for (unsigned lv = 0, lve = topSort.size(); lv < lve; lv++)
     if (topSort[lv] == loopIdx)
-      return loopEmitter->getLoopIV(lv);
+      return loopEmitter.getLoopIV(lv);
   llvm_unreachable("invalid loop index");
 }
 
 //===----------------------------------------------------------------------===//
 // Code generation environment sparse tensor output and expansion methods
 //===----------------------------------------------------------------------===//
 
 void CodegenEnv::updateInsertionChain(Value chain) {

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

 //===----------------------------------------------------------------------===//
 
 /// 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();
   assert(op.getNumOperands() == op.getNumDpsInputs() + 1);
 
-  env.emitter()->initializeLoopEmit(
+  env.emitter().initializeLoopEmit(
       builder, loc,
       /// Generates buffer for the output tensor.
       /// Note that all sparse kernels assume that when all elements are written
       /// to (viz. x(i) = y(i) * z(i)), the output buffer is already initialized
       /// to all zeroes and only nonzeroes values are computed and written out.
       /// For updates (viz. x(i) += y(i) * z(i)), only nonzeroes values are used
       /// for the updates and no assumption on the original contents of the
       /// output buffer is necessary.

 static Value genSubscript(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
                           SmallVectorImpl<Value> &args) {
   linalg::GenericOp op = env.op();
   unsigned tensor = t->getOperandNumber();
   auto map = op.getMatchingIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   unsigned rank = map.getNumResults();
   if (enc) {
-    Value pidx = env.emitter()->getPidxs()[tensor].back();
+    Value pidx = env.emitter().getPidxs()[tensor].back();
     assert(pidx);
     args.push_back(pidx); // position index
   } else {
     for (unsigned d = 0; d < rank; d++) {
       AffineExpr a = map.getResult(d);
-      args.push_back(env.emitter()->genAffine(builder, a, op.getLoc()));
+      args.push_back(env.emitter().genAffine(builder, a, op.getLoc()));
     }
   }
-  return env.emitter()->getValBuffer()[tensor];
+  return env.emitter().getValBuffer()[tensor];
 }
 
 /// Generates insertion code to implement dynamic tensor load.
 static Value genInsertionLoad(CodegenEnv &env, OpBuilder &builder,
                               OpOperand *t) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   // Direct lexicographic index order, tensor loads as zero.
static void genInsertionStore(CodegenEnv &env, OpBuilder &builder, OpOperand *t,
                               Value rhs) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   // Direct insertion in lexicographic index order.
   if (!env.isExpand()) {
     unsigned rank = op.getRank(t);
     SmallVector<Value> indices;
     for (unsigned i = 0; i < rank; i++) {
-      assert(env.emitter()->getLoopIV(i));
-      indices.push_back(env.emitter()->getLoopIV(i));
+      assert(env.emitter().getLoopIV(i));
+      indices.push_back(env.emitter().getLoopIV(i));
     }
     Value chain = env.getInsertionChain();
     env.updateInsertionChain(
         builder.create<InsertOp>(loc, rhs, chain, indices));
     return;
   }
   // Generates insertion code along expanded access pattern.
   //   if (!expFilled[i]) then
if (atStart) {
     Type t4 = builder.getIndexType();
     auto r = builder.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
     assert(r.getNumResults() == 4);
     env.startExpand(r.getResult(0), r.getResult(1), r.getResult(2),
                     r.getResult(3));
   } else {
     SmallVector<Value> indices;
     for (unsigned i = 0; i < at; i++)
-      indices.push_back(env.emitter()->getLoopIV(i));
+      indices.push_back(env.emitter().getLoopIV(i));
     Value values = env.getExpandValues();
     Value filled = env.getExpandFilled();
     Value added = env.getExpandAdded();
     Value count = env.getExpandCount();
     Value chain = env.getInsertionChain();
     Value compress = builder.create<CompressOp>(loc, values, filled, added,
                                                 count, chain, indices);
     env.updateInsertionChain(compress);
if (env.merger().isFilterLoop(idx)) {
       // tids/dims must only have one value because filter loops only
       // corresponding to the one and only sparse tensor level.
       assert(isSparse && tids.size() == 1 && dims.size() == 1);
       OpOperand *t = &op->getOpOperand(tid);
       auto enc = getSparseTensorEncoding(t->get().getType());
       // Retrieves the affine expression for the filter loop.
       AffineExpr a =
           op.getMatchingIndexingMap(t).getResult(toOrigDim(enc, dim));
-      return env.emitter()->enterFilterLoopOverTensorAtDim(builder, loc, tid,
-                                                           dim, a, reduc);
+      return env.emitter().enterFilterLoopOverTensorAtDim(builder, loc, tid,
     }
-    return env.emitter()->enterLoopOverTensorAtDim(builder, loc, tids, dims,
-                                                   reduc, isParallel);
+    return env.emitter().enterLoopOverTensorAtDim(builder, loc, tids, dims,
   });
   assert(loop);
   return loop;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
 static Operation *genWhile(CodegenEnv &env, OpBuilder &builder, unsigned idx,
                            bool needsUniv, ArrayRef<size_t> tids,
                            ArrayRef<size_t> dims) {
   Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     // Construct the while-loop with a parameter for each
     // index.
-    return env.emitter()->enterCoIterationOverTensorsAtDims(
+    return env.emitter().enterCoIterationOverTensorsAtDims(
         builder, env.op().getLoc(), tids, dims, needsUniv, reduc);
   });
   assert(loop);
   return loop;
 }
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, unsigned idx,
   for (unsigned b = 0, be = conditions.size(); b < be; b++) {
     if (!conditions[b])
       continue;
     unsigned tensor = env.merger().tensor(b);
     assert(idx == env.merger().index(b));
     Value clause;
     if (isCompressedDLT(env.dlt(b)) || isSingletonDLT(env.dlt(b))) {
       auto dim = *env.merger().getDimNum(tensor, idx);
-      Value op1 = env.emitter()->getCoord()[tensor][dim];
+      Value op1 = env.emitter().getCoord()[tensor][dim];
       Value op2 = env.getLoopIdxValue(idx);
       clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, op1,
                                              op2);
     } else {
       assert(isDenseDLT(env.merger().getDimLevelType(b)) ||
              isUndefDLT(env.merger().getDimLevelType(b)));
       clause = constantI1(builder, loc, true);
     }
env.merger().foreachTidDimPairInBits(
           needsUniv = true;
         } else {
           // sparse/singleton dim levels.
           tids.push_back(tid);
           dims.push_back(*dim);
         }
       });
 
-  env.emitter()->enterNewLoopSeq(builder, env.op().getLoc(), tids, dims);
+  env.emitter().enterNewLoopSeq(builder, env.op().getLoc(), tids, dims);
 
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     unsigned lsize = env.set(lts).size();
     for (unsigned i = 1; i < lsize; i++) {
       unsigned li = env.set(lts)[i];
       if (!env.merger().hasAnySparse(env.lat(li).simple))
static void genConstantDenseAddressFromLevel(CodegenEnv &env,
   OpOperand *input = op.getDpsInputOperands()[tid];
   ArrayRef<AffineExpr> affines = op.getMatchingIndexingMap(input).getResults();
   auto enc = getSparseTensorEncoding(input->get().getType());
   if (enc) {
     for (unsigned i = lvl, e = affines.size(); i < e; i++) {
       AffineExpr affine = affines[toOrigDim(enc, i)];
       if (isDenseDLT(getDimLevelType(enc, i)) &&
           affine.isa<AffineConstantExpr>())
-        env.emitter()->genDenseAffineAddressAtCurLevel(
+        env.emitter().genDenseAffineAddressAtCurLevel(
             builder, op.getLoc(), input->getOperandNumber(), i, affine);
       else
         return; // break on first non-dense non-constant level
     }
   }
 }
 
 static void genInitConstantDenseAddress(CodegenEnv &env,
static Operation *startLoop(CodegenEnv &env, OpBuilder &builder, unsigned at,
   SmallVector<size_t> affineTids, affineDims;
   SmallVector<AffineExpr> affines;
   bool isFor = translateBitsToTidDimPairs(
       env, li, env.topSortAt(at), tids, dims, affineTids, affineDims, affines);
 
   // Emit the for/while-loop control.
   Operation *loop = genLoop(env, builder, at, needsUniv, tids, dims, isFor);
   for (auto [tid, dim, exp] : llvm::zip(affineTids, affineDims, affines)) {
-    env.emitter()->genDenseAffineAddressAtCurLevel(builder, env.op().getLoc(),
-                                                   tid, dim, exp);
+    env.emitter().genDenseAffineAddressAtCurLevel(builder, env.op().getLoc(),
   }
 
   // Until now, we have entered every <tid, dim> pair in {cond, extra,
   // affine}Tids/Dims. The addresses of the upcoming levels which are dependent
   // on constant affines expression may now be determined.
   auto allTids = llvm::concat<size_t>(tids, affineTids);
   auto allDims = llvm::concat<size_t>(dims, affineDims);
   for (auto [tid, dim] : llvm::zip(allTids, allDims)) {
static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop,
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
     finalizeWhileOp(env, rewriter, idx, needsUniv, env.lat(li).bits, whileOp);
   } else {
     needsUniv = false;
   }
 
   env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
-    env.emitter()->exitCurrentLoop(rewriter, env.op().getLoc(), reduc);
+    env.emitter().exitCurrentLoop(rewriter, env.op().getLoc(), reduc);
     return std::nullopt;
   });
 
   return needsUniv;
 }
 
 /// Ends a loop sequence at given level.
 static void endLoopSeq(CodegenEnv &env, OpBuilder &builder, unsigned exp,
                        unsigned at, unsigned idx, unsigned ldx) {
   assert(env.getLoopIdxValue(idx) == nullptr);
-  env.emitter()->exitCurrentLoopSeq();
+  env.emitter().exitCurrentLoopSeq();
   // Unmark bookkeeping of invariants and loop index.
   genInvariants(env, builder, exp, ldx, /*atStart=*/false);
   // Finalize access pattern expansion for sparse tensor output.
   genExpand(env, builder, at, /*atStart=*/false);
 }
 
 /// Recursively generates code while computing iteration lattices in order
 /// to manage the complexity of implementing co-iteration over unions
if (getSparseTensorEncoding(resType)) {
     if (Value chain = env.getInsertionChain()) {
       hasInserts = true;
       tensor = chain;
     }
     rewriter.replaceOpWithNewOp<LoadOp>(op, resType, tensor, hasInserts);
   } else {
     // To rematerialize an non-annotated tensor, simply load it
     // from the bufferized value.
-    Value val = env.emitter()->getValBuffer().back(); // value array
+    Value val = env.emitter().getValBuffer().back(); // value array
     rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler rewriting methods.
 //===----------------------------------------------------------------------===//
 
for (auto mask : allMask)
         }
         // else try a set of less strict constraints.
       }
     if (hasCycle)
       return resolveCycle(env, rewriter); // one last shot
     if (!isAdmissible)
       return failure(); // inadmissible expression, reject
 
-    // Updates environment with a loop emitter.
-    // TODO: refactor so that emitter can be constructed earlier
-    //       and updating is made easy, i.e. remove this whole block?
-    SmallVector<Value> tensors;
-    for (OpOperand &t : op->getOpOperands())
-      tensors.push_back(t.get());
-    LoopEmitter lpe(
-        tensors,
-        StringAttr::get(op.getContext(), linalg::GenericOp::getOperationName()),
-        /*hasOutput=*/true, /*isSparseOut=*/sparseOut != nullptr,
-        env.topSortRef());
-    env.startEmit(sparseOut, outerParNest, &lpe);
-
     // Recursively generates code if admissible.
+    env.startEmit(sparseOut, outerParNest);
     genBuffers(env, rewriter);
     genInitConstantDenseAddress(env, rewriter);
     genStmt(env, rewriter, exp, 0);
     genResult(env, rewriter);
     return success();
   }
 
 private: