[mlir][sparse] move loop boundary method to codegenenv

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

public:
 
   linalg::GenericOp op() const { return linalgOp; }
   const SparsificationOptions &options() const { return sparseOptions; }
   Merger &merger() { return latticeMerger; }
   SparseTensorLoopEmitter *emitter() { return loopEmitter; }
 
   void startEmit(OpOperand *so, unsigned lv, SparseTensorLoopEmitter *le);
 
+  /// 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);
+
   //
   // 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 {

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

void CodegenEnv::startEmit(OpOperand *so, unsigned lv,
   outerParNest = lv;
   loopEmitter = le;
   if (sparseOut) {
     insChain = sparseOut->get();
     latticeMerger.setHasSparseOut(true);
   }
 }
 
+Optional<Operation *> CodegenEnv::genLoopBoundary(
+    function_ref<Optional<Operation *>(MutableArrayRef<Value> parameters)>
+        callback) {
+  SmallVector<Value> params;
+  if (isReduc())
+    params.push_back(redVal);
+  if (isExpand())
+    params.push_back(expCount);
+  if (insChain != nullptr)
+    params.push_back(insChain);
+  auto r = callback(params); // may update parameters
+  unsigned i = 0;
+  if (isReduc())
+    updateReduc(params[i++]);
+  if (isExpand())
+    updateExpandCount(params[i++]);
+  if (insChain != nullptr)
+    updateInsertionChain(params[i]);
+  return r;
+}
+
 //===----------------------------------------------------------------------===//
 // Code generation environment topological sort methods
 //===----------------------------------------------------------------------===//
 
 ArrayRef<unsigned> CodegenEnv::getTopSortSlice(size_t n, size_t m) const {
   return ArrayRef<unsigned>(topSort).slice(n, m);
 }
 

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

static bool isAdmissibleTensorExp(CodegenEnv &env, unsigned exp,
   }
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
-/// Generates loop boundary statements (entering/exiting loops). The function
-/// passes and updates the reduction value.
-static Optional<Operation *> genLoopBoundary(
-    CodegenEnv &env,
-    function_ref<Optional<Operation *>(MutableArrayRef<Value> reduc)>
-        callback) {
-  SmallVector<Value> reduc;
-  if (env.isReduc())
-    reduc.push_back(env.getReduc());
-  if (env.isExpand())
-    reduc.push_back(env.getExpandCount());
-  if (env.getInsertionChain())
-    reduc.push_back(env.getInsertionChain());
-
-  auto r = callback(reduc);
-
-  // Callback should do in-place update on reduction value vector.
-  unsigned i = 0;
-  if (env.isReduc())
-    env.updateReduc(reduc[i++]);
-  if (env.isExpand())
-    env.updateExpandCount(reduc[i++]);
-  if (env.getInsertionChain())
-    env.updateInsertionChain(reduc[i]);
-
-  return r;
-}
-
 /// 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(
       builder, loc,
static Operation *genFor(CodegenEnv &env, OpBuilder &builder, bool isOuter,
                          ArrayRef<size_t> extraDims) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   auto iteratorTypes = op.getIteratorTypesArray();
   bool isSparse =
       isCompressedDLT(env.dlt(tid, idx)) || isSingletonDLT(env.dlt(tid, idx));
   bool isParallel = isParallelFor(env, isOuter, isSparse);
 
-  Operation *loop = *genLoopBoundary(env, [&](MutableArrayRef<Value> reduc) {
+  Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     if (env.merger().isFilterLoop(idx)) {
       // extraTids/extraDims must be empty because filter loops only
       // corresponding to the one and only sparse tensor level.
       assert(isSparse && extraTids.empty() && extraDims.empty());
       OpOperand *t = &op->getOpOperand(tid);
       auto enc = getSparseTensorEncoding(t->get().getType());
       // Retrieves the affine expression for the filter loop.
       AffineExpr a =

 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
 static Operation *genWhile(CodegenEnv &env, OpBuilder &builder, unsigned idx,
                            bool needsUniv, ArrayRef<size_t> condTids,
                            ArrayRef<size_t> condDims,
                            ArrayRef<size_t> extraTids,
                            ArrayRef<size_t> extraDims) {
-  Operation *loop = *genLoopBoundary(env, [&](MutableArrayRef<Value> reduc) {
+  Operation *loop = *env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     // Construct the while-loop with a parameter for each
     // index.
     return env.emitter()->enterCoIterationOverTensorsAtDims(
         builder, env.op().getLoc(), condTids, condDims, needsUniv, reduc,
         extraTids, extraDims);
   });
   assert(loop);
   return loop;
static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop,
                     unsigned idx, unsigned li, bool needsUniv) {
   // 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;
   }
 
-  genLoopBoundary(env, [&](MutableArrayRef<Value> reduc) {
+  env.genLoopBoundary([&](MutableArrayRef<Value> reduc) {
     env.emitter()->exitCurrentLoop(rewriter, env.op().getLoc(), reduc);
     return std::nullopt;
   });
 
   return needsUniv;
 }
 
 /// Ends a loop sequence at given level.