[mlir][sparse] refactoring loop emitter into its own files.

Reviewed By: aartbik

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

Author: Peiming Liu

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

@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRSparseTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   CodegenEnv.cpp
   CodegenUtils.cpp
+  LoopEmitter.cpp
   SparseBufferRewriting.cpp
   SparseStorageSpecifierToLLVM.cpp
   SparseTensorCodegen.cpp

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

@@ -24,8 +24,7 @@ CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
       expFilled(), expAdded(), expCount(), redVal(), redExp(-1u),
       redCustom(-1u) {}
 
-void CodegenEnv::startEmit(OpOperand *so, unsigned lv,
-                           SparseTensorLoopEmitter *le) {
+void CodegenEnv::startEmit(OpOperand *so, unsigned lv, LoopEmitter *le) {
   assert(sparseOut == nullptr && loopEmitter == nullptr &&
          insChain == nullptr && "must only start emitting once");
   sparseOut = so;

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

@@ -14,6 +14,7 @@
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENENV_H_
 
 #include "CodegenUtils.h"
+#include "LoopEmitter.h"
 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
@@ -45,9 +46,9 @@ class CodegenEnv {
   linalg::GenericOp op() const { return linalgOp; }
   const SparsificationOptions &options() const { return sparseOptions; }
   Merger &merger() { return latticeMerger; }
-  SparseTensorLoopEmitter *emitter() { return loopEmitter; }
+  LoopEmitter *emitter() { return loopEmitter; }
 
-  void startEmit(OpOperand *so, unsigned lv, SparseTensorLoopEmitter *le);
+  void startEmit(OpOperand *so, unsigned lv, LoopEmitter *le);
 
   /// Generates loop boundary statements (entering/exiting loops). The function
   /// passes and updates the passed-in parameters.
@@ -135,7 +136,7 @@ class CodegenEnv {
 
   // Loop emitter helper class (keep reference in scope!).
   // TODO: move emitter constructor up in time?
-  SparseTensorLoopEmitter *loopEmitter;
+  LoopEmitter *loopEmitter;
 
   // Topological sort.
   std::vector<unsigned> topSort;