[mlir][sparse] add initialize() API to LoopEmitter to support post-constructor initialization

[Peiming]

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

219–231

I know that this is not recommended, but it cause name collision between parameter and fields.

I can change the parameter names instead of using this if you prefer.

[aartbik]

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

586

I would move the comment from the constructor here, and document it as a setup, and then document the constructor as doing a direct setup. But just a suggestion, your call in the end.

[Peiming]

address comments.

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

 
 class SparseTensorLoopEmitter {
 public:
   /// Optional callback function to setup dense output tensors when
   /// initializing the loop emitter (e.g., to fill a dense output with zeros).
   using OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc,
                                            Value memref, Value tensor)>;
 
-  /// Constructor: take an array of tensors inputs, on which the generated
-  /// loops will iterate on. The index of the tensor in the array is also the
-  /// tensor id (tid) used in related functions.
-  /// If isSparseOut is set, loop emitter assume that the sparse output tensor
-  /// is empty, and will always generate loops on it based on the dim sizes.
-  /// An optional array could be provided (by sparsification) to indicate the
-  /// loop id sequence that will be generated. It is used to establish the
-  /// mapping between affineDimExpr to the corresponding loop index in the
-  /// loop stack that are maintained by the loop emitter.
+  SparseTensorLoopEmitter() = default;
+
+  /// Takes an array of tensors inputs, on which the generated loops will
+  /// iterate on. The index of the tensor in the array is also the tensor id
+  /// (tid) used in related functions. If isSparseOut is set, loop emitter
+  /// assume that the sparse output tensor is empty, and will always generate
+  /// loops on it based on the dim sizes. An optional array could be provided
+  /// (by sparsification) to indicate the loop id sequence that will be
+  /// generated. It is used to establish the mapping between affineDimExpr to
+  /// the corresponding loop index in the loop stack that are maintained by the
+  /// loop emitter.
+  void initialize(ValueRange tensors, StringAttr loopTag = nullptr,
+                  bool hasOutput = false, bool isSparseOut = false,
+                  ArrayRef<unsigned> topSort = {});
+
   explicit SparseTensorLoopEmitter(ValueRange tensors,
                                    StringAttr loopTag = nullptr,
                                    bool hasOutput = false,
                                    bool isSparseOut = false,
                                    ArrayRef<unsigned> topSort = {});
 
   /// Starts a loop emitting session by generating all the buffers needed to
   /// iterate tensors.
   void initializeLoopEmit(OpBuilder &builder, Location loc,
                           OutputUpdater updater = nullptr);
 
   /// Generates a list of operations to compute the affine expression.
   Value genAffine(OpBuilder &builder, AffineExpr a, Location loc);
 
   /// Enters a new loop sequence, the loops within the same sequence starts

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

 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
                                                  StringAttr loopTag,
                                                  bool hasOutput,
                                                  bool isSparseOut,
-                                                 ArrayRef<unsigned> topSort)
-    : loopTag(loopTag), hasOutput(hasOutput), isSparseOut(isSparseOut),
-      tensors(tensors.begin(), tensors.end()), dimTypes(tensors.size()),
-      pidxs(tensors.size()), coord(tensors.size()), highs(tensors.size()),
-      ptrBuffer(tensors.size()), idxBuffer(tensors.size()),
-      valBuffer(tensors.size()), loopStack(),
-      sparsiferLoopLvlMap(topSort.size(), 0) {
+                                                 ArrayRef<unsigned> topSort) {
+  initialize(tensors, loopTag, hasOutput, isSparseOut, topSort);
+}
+
+void SparseTensorLoopEmitter::initialize(ValueRange tensors, StringAttr loopTag,
+                                         bool hasOutput, bool isSparseOut,
+                                         ArrayRef<unsigned> topSort) {
+  // First initializes fields.
+  this->loopTag = loopTag;
+  this->hasOutput = hasOutput;
+  this->isSparseOut = isSparseOut;
+  this->tensors.assign(tensors.begin(), tensors.end());
+  this->dimTypes.assign(tensors.size(), std::vector<DimLevelType>());
+  this->pidxs.assign(tensors.size(), std::vector<Value>());
+  this->coord.assign(tensors.size(), std::vector<Value>());
+  this->highs.assign(tensors.size(), std::vector<Value>());
+  this->ptrBuffer.assign(tensors.size(), std::vector<Value>());
+  this->idxBuffer.assign(tensors.size(), std::vector<Value>());
+  this->valBuffer.assign(tensors.size(), nullptr);
+  this->loopStack.reserve(topSort.size());
+  this->sparsiferLoopLvlMap.assign(topSort.size(), 0);
+
   for (size_t tid = 0, e = tensors.size(); tid < e; tid++) {
     auto t = tensors[tid];
     // a scalar or 0-dimension tensors
     if (isZeroRankedTensorOrScalar(t.getType()))
       continue;
     auto rtp = t.getType().cast<RankedTensorType>();
     auto rank = static_cast<size_t>(rtp.getRank());
     auto enc = getSparseTensorEncoding(rtp);
     // We always treat sparse output tensor as dense so that we always iterate
     // it based on dim size.
     if (enc && !(isOutputTensor(tid) && isSparseOut))
       for (auto dimTp : enc.getDimLevelType())
         dimTypes[tid].push_back(dimTp);
     else
       dimTypes[tid].assign(rank, DimLevelType::Dense);
 
     // Initialize using empty value.
     pidxs[tid].assign(rank, Value());
     coord[tid].assign(rank, Value());
     highs[tid].assign(rank, Value());
     ptrBuffer[tid].assign(rank, Value());
     idxBuffer[tid].assign(rank, Value());
   }
 
+  // FIXME: This map should be maintained outside loop emitter.
   for (unsigned i = 0, e = topSort.size(); i < e; i++) {
     // This is an inverse map of the topologically sorted loop index from
     // sparsifier. This is needed to map the AffineDimExpr back to the loopStack
     // index used in loop emitter.
     sparsiferLoopLvlMap[topSort[i]] = i;
   }
 }