[mlir][sparse] refactoring: move genAffine to loopEmitter

[Peiming]

Summary

This is the beginning patch of a sequence of dependent patches that in together provide the affine expression on matched indexing mapping for sparse tensors.

This patch itself simply move genAffine into loop emitter to be prepared for upcoming patches.

D138169 provides support for affine expression on dense dimensions only (except for constant affine expression)
D138170 provides support for constant affine expressions on dense dimensions
D138171 provides merger support for affine expression on sparse dimension (without codegen)
D138172 provides codegen support (by generating a "filter" loop) for affine expression on sparse dimensions.
D138173 fixes a crash on resolveCycle when dealing with affine expressions.

[Peiming]

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

558

Should I make it a std::function in this patch? I think block sparsity might also need a (different) mapping?

It is not a complex change anyway.

[Peiming]

rebase

[aartbik]

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

351

update doc to match all parameters

558

I would do that change when needed? but up to you

[Peiming]

address comments

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

 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.
   explicit SparseTensorLoopEmitter(ValueRange tensors,
                                    StringAttr loopTag = nullptr,
                                    bool hasOutput = false,
-                                   bool isSparseOut = 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 from
   /// the break points of previous loop instead of starting over from 0.
   /// e.g.,
   /// {
   ///   // loop sequence start.
   ///   p0 = while(xxx)
   ///     ...
   ///     break p0
private:
 
   // Loop Stack, stores the information of all the nested loops that are alive.
   std::vector<LoopLevelInfo> loopStack;
 
   // Loop Sequence Stack, stores the unversial index for the current loop
   // sequence.
   std::vector<Value> loopSeqStack;
 
+  // Maps AffineDimExpr to the index of the loop in loopStack.
+  // TODO: We should probably use a callback function here to make it more
+  // general.
+  std::vector<unsigned> sparsiferLoopLvlMap;
+
   // TODO: not yet used, it should track the current level for each tensor
   // to help eliminate `dim` paramters from above APIs.
   // std::vector<size_t> curLv;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

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

 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
                                                  StringAttr loopTag,
                                                  bool hasOutput,
-                                                 bool isSparseOut)
+                                                 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() {
+      valBuffer(tensors.size()), loopStack(),
+      sparsiferLoopLvlMap(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());
   }
+
+  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;
+  }
 }
 
 void SparseTensorLoopEmitter::initializeLoopEmit(
     OpBuilder &builder, Location loc,
     SparseTensorLoopEmitter::OutputUpdater updater) {
   // For every tensor, find lower and upper bound on dimensions, set the
   // same bounds on loop indices, and obtain dense or sparse buffer(s).
   for (size_t t = 0, e = tensors.size(); t < e; t++) {
void SparseTensorLoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
   assert(loopSeqStack.size() == loopStack.size());
   // Universal index starts from 0
   loopSeqStack.emplace_back(constantIndex(builder, loc, 0));
   // Prepares for all the tensors used in the current loop sequence.
   for (auto [tid, dim] : llvm::zip(tids, dims))
     prepareLoopOverTensorAtDim(builder, loc, tid, dim);
 }
 
+Value SparseTensorLoopEmitter::genAffine(OpBuilder &builder, AffineExpr a,
+                                         Location loc) {
+  switch (a.getKind()) {
+  case AffineExprKind::DimId: {
+    unsigned idx = a.cast<AffineDimExpr>().getPosition();
+    return loopStack[sparsiferLoopLvlMap[idx]].iv;
+  }
+  case AffineExprKind::Add: {
+    auto binOp = a.cast<AffineBinaryOpExpr>();
+    return builder.create<arith::AddIOp>(
+        loc, genAffine(builder, binOp.getLHS(), loc),
+        genAffine(builder, binOp.getRHS(), loc));
+  }
+  case AffineExprKind::Mul: {
+    auto binOp = a.cast<AffineBinaryOpExpr>();
+    return builder.create<arith::MulIOp>(
+        loc, genAffine(builder, binOp.getLHS(), loc),
+        genAffine(builder, binOp.getRHS(), loc));
+  }
+  case AffineExprKind::Constant: {
+    int64_t c = a.cast<AffineConstantExpr>().getValue();
+    return constantIndex(builder, loc, c);
+  }
+  default:
+    llvm_unreachable("unexpected affine subscript");
+  }
+}
+
 Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
     OpBuilder &builder, Location loc, size_t tid, size_t dim,
     MutableArrayRef<Value> reduc, bool isParallel, ArrayRef<size_t> extraTids,
     ArrayRef<size_t> extraDims) {
 
   assert(dimTypes[tid].size() > dim);
   // We can not re-enter the same level.
   assert(!coord[tid][dim]);

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

struct CodeGen {
   CodeGen(SparsificationOptions o, MLIRContext *context, ValueRange tensors,
           unsigned numTensors, unsigned numLoops, OpOperand *op, unsigned nest,
           std::vector<unsigned> &ts)
       : options(o),
         loopEmitter(
             tensors,
             StringAttr::get(context, linalg::GenericOp::getOperationName()),
             /*hasOutput=*/true,
-            /*isSparseOut=*/op != nullptr),
+            /*isSparseOut=*/op != nullptr, ts),
         sparseOut(op), outerParNest(nest), topSort(ts) {
     if (op)
       insChain = op->get();
   }
   /// Sparsification options.
   SparsificationOptions options;
   /// Loop emitter helper class.
   SparseTensorLoopEmitter loopEmitter;
codegen.loopEmitter.initializeLoopEmit(
                                     getElementTypeOrSelf(tensor.getType()));
           builder.create<linalg::FillOp>(loc, ValueRange{zero},
                                          ValueRange{init});
         }
         return init;
       });
 }
 
-/// Generates an affine expression.
-//
-// TODO: generalize for sparse tensor subscripts
-//
-static Value genAffine(CodeGen &codegen, OpBuilder &builder, AffineExpr a,
-                       Location loc) {
-  switch (a.getKind()) {
-  case AffineExprKind::DimId: {
-    unsigned idx = a.cast<AffineDimExpr>().getPosition();
-    return codegen.getLoopIdxValue(idx); // universal dense index
-  }
-  case AffineExprKind::Add: {
-    auto binOp = a.cast<AffineBinaryOpExpr>();
-    return builder.create<arith::AddIOp>(
-        loc, genAffine(codegen, builder, binOp.getLHS(), loc),
-        genAffine(codegen, builder, binOp.getRHS(), loc));
-  }
-  case AffineExprKind::Mul: {
-    auto binOp = a.cast<AffineBinaryOpExpr>();
-    return builder.create<arith::MulIOp>(
-        loc, genAffine(codegen, builder, binOp.getLHS(), loc),
-        genAffine(codegen, builder, binOp.getRHS(), loc));
-  }
-  case AffineExprKind::Constant: {
-    int64_t c = a.cast<AffineConstantExpr>().getValue();
-    return constantIndex(builder, loc, c);
-  }
-  default:
-    llvm_unreachable("unexpected affine subscript");
-  }
-}
-
 /// Generates index for load/store on sparse tensor.
 static Value genIndex(CodeGen &codegen, linalg::GenericOp op, OpOperand *t) {
   auto map = op.getMatchingIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   AffineExpr a = map.getResult(toOrigDim(enc, map.getNumResults() - 1));
   assert(a.getKind() == AffineExprKind::DimId);
   unsigned idx = a.cast<AffineDimExpr>().getPosition();
   return codegen.getLoopIdxValue(idx);
if (enc) {
     assert(map.getResult(toOrigDim(enc, rank - 1)).getKind() ==
            AffineExprKind::DimId);
     Value pidx = codegen.loopEmitter.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(genAffine(codegen, builder, a, op.getLoc()));
+      args.push_back(codegen.loopEmitter.genAffine(builder, a, op.getLoc()));
     }
   }
   return codegen.loopEmitter.getValBuffer()[tensor];
 }
 
 /// Generates insertion code to implement dynamic tensor load.
 static Value genInsertionLoad(CodeGen &codegen, OpBuilder &builder,
                               linalg::GenericOp op, OpOperand *t) {