[mlir][sparse] annotate loops that are generated by loop emitter.

[Peiming]

use operation name as the tag.

[aartbik]

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

499

maybe a doc comment?

[Peiming]

address comment.

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

public:
   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.
-  explicit SparseTensorLoopEmitter(ValueRange tensors, bool hasOutput = false,
+  explicit SparseTensorLoopEmitter(ValueRange tensors,
+                                   StringAttr loopTag = nullptr,
+                                   bool hasOutput = false,
                                    bool isSparseOut = false);
 
   /// Starts a loop emitting session by generating all the buffers needed to
   /// iterate tensors.
   void initializeLoopEmit(OpBuilder &builder, Location loc,
                           OutputUpdater updater = nullptr);
 
   /// Enters a new loop sequence, the loops within the same sequence starts from
public:
   const std::vector<std::vector<Value>> &getPtrBuffer() const {
     return ptrBuffer;
   };
   const std::vector<std::vector<Value>> &getIdxBuffer() const {
     return idxBuffer;
   };
   const std::vector<Value> &getValBuffer() const { return valBuffer; };
 
+  constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName() {
+    return llvm::StringLiteral("Emitted from");
+  }
+
 private:
   struct LoopLevelInfo {
     LoopLevelInfo(ArrayRef<size_t> tids, ArrayRef<size_t> dims, Operation *loop,
-                  Value iv)
-        : tids(tids), dims(dims), loop(loop), iv(iv) {}
+                  Value iv, StringAttr loopTag)
+        : tids(tids), dims(dims), loop(loop), iv(iv) {
+      // Attached a special tag to loop emitter generated loop.
+      if (loopTag)
+        loop->setAttr(SparseTensorLoopEmitter::getLoopEmitterLoopAttrName(),
+                      loopTag);
+    }
     // TODO: maybe use a vector<pair> for tid and dim?
     // The set of tensors that the loop is operating on
     const llvm::SmallVector<size_t> tids;
     // The corresponding dims for the tensors
     const llvm::SmallVector<size_t> dims;
     const Operation *loop; // the loop operation
     const Value iv;        // the induction variable for the loop
   };
private:
   /// Return %ret to user, while %val is provided by users (`reduc`).
   void exitForLoop(RewriterBase &rewriter, Location loc,
                    MutableArrayRef<Value> reduc);
 
   /// Exits a while loop, returns the reduction results.
   void exitCoIterationLoop(OpBuilder &builder, Location loc,
                            MutableArrayRef<Value> reduc);
 
-  // Whether the loop emitter needs to treat the last tensor as the output
-  // tensor.
+  /// A optional string attribute that should be attached to the loop generated
+  /// by loop emitter, it might help following passes to identify loops that
+  /// operates on sparse tensors more easily.
+  StringAttr loopTag;
+  /// Whether the loop emitter needs to treat the last tensor as the output
+  /// tensor.
   bool hasOutput;
   bool isSparseOut;
   /// Input and (optional) output tensors.
   std::vector<Value> tensors;
   /// The dim type array for each tensor.
   std::vector<std::vector<DimLevelType>> dimTypes;
   /// Sparse iteration information (by tensor and dim). These arrays
   /// are updated to remain current within the current loop.

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

static Value genIndexAndValueForDense(OpBuilder &builder, Location loc,
   return val;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
+                                                 StringAttr loopTag,
                                                  bool hasOutput,
                                                  bool isSparseOut)
-    : hasOutput(hasOutput), isSparseOut(isSparseOut),
+    : 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() {
   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()))
if (isSparseInput) {
     auto enc = getSparseTensorEncoding(tensors[tid].getType());
     if (enc && !isSparseOutput(tid))
       pidxs[tid][dim] = genAddress(builder, loc, tid, dim, iv);
   }
 
   // NOTE: we can also prepares for next dim here in advance
   // Push the loop into stack
   loopStack.emplace_back(ArrayRef<size_t>(tid), ArrayRef<size_t>(dim), loop,
-                         coord[tid][dim]);
+                         coord[tid][dim], loopTag);
   // Emit extra locals.
   emitExtraLocalsForTensorsAtDenseDims(builder, loc, extraTids, extraDims);
 
   return loop;
 }
 
 Operation *SparseTensorLoopEmitter::enterCoIterationOverTensorsAtDims(
     OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
if (isDenseDLT(dimTypes[tid][dim])) {
       // generate pidx for dense dim (pidx = i * sz + j)
       auto enc = getSparseTensorEncoding(tensors[tid].getType());
       if (enc && !isSparseOutput(tid))
         pidxs[tid][dim] = genAddress(builder, loc, tid, dim, min);
     }
     // NOTE: we can also prepares for next dim here in advance
   }
   // Sets up the loop stack.
-  loopStack.emplace_back(tids, dims, whileOp, min);
+  loopStack.emplace_back(tids, dims, whileOp, min, loopTag);
   assert(loopStack.size() == loopSeqStack.size());
 
   // Emits extra locals
   emitExtraLocalsForTensorsAtDenseDims(builder, loc, extraTids, extraDims);
 
   // Updates reduction variables
   assert(after->getNumArguments() == o + reduc.size() + (needsUniv ? 1 : 0));
   // In-place update on reduction variable.

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

if (auto constOp = input.getDefiningOp<arith::ConstantOp>()) {
         return success();
       }
     }
 
     // Otherwise, use loop emitter to generate loops.
     auto enc = getSparseTensorEncoding(rtp);
 
     // 1. Generates loop for the sparse input.
-    SparseTensorLoopEmitter loopEmitter(ValueRange{input});
+    SparseTensorLoopEmitter loopEmitter(
+        ValueRange{input},
+        StringAttr::get(getContext(), ForeachOp::getOperationName()));
     loopEmitter.initializeLoopEmit(rewriter, loc);
     for (int64_t i = 0; i < rank; i++) {
       // TODO: provide utility function for loop sequences that only contains
       // one for loop?
       loopEmitter.enterNewLoopSeq(rewriter, loc, 0, static_cast<size_t>(i));
       // Note that reduc will be taken care of by loop emitter and get updated
       // in place.
       loopEmitter.enterLoopOverTensorAtDim(rewriter, loc, 0, i, reduc);

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

enum SortMask {
   kIncludeAll = 0x3
 };
 
 // Reduction kinds.
 enum Reduction { kNoReduc, kSum, kProduct, kAnd, kOr, kXor, kCustom };
 
 // Code generation.
 struct CodeGen {
-  CodeGen(SparsificationOptions o, ValueRange tensors, unsigned numTensors,
-          unsigned numLoops, OpOperand *op, unsigned nest,
+  CodeGen(SparsificationOptions o, MLIRContext *context, ValueRange tensors,
+          unsigned numTensors, unsigned numLoops, OpOperand *op, unsigned nest,
           std::vector<unsigned> &ts)
-      : options(o), loopEmitter(tensors, /*hasOutput=*/true,
-                                /*isSparseOut=*/op != nullptr),
+      : options(o),
+        loopEmitter(
+            tensors,
+            StringAttr::get(context, linalg::GenericOp::getOperationName()),
+            /*hasOutput=*/true,
         sparseOut(op), outerParNest(nest), topSort(ts) {
     if (op)
       insChain = op->get();
   }
   /// Sparsification options.
   SparsificationOptions options;
   /// Loop emitter helper class.
   SparseTensorLoopEmitter loopEmitter;
if (t == codegen.sparseOut) {
     if (!rhs) {
       // Only unary and binary are allowed to return uninitialized rhs
       // to indicate missing output.
       assert(merger.exp(exp).kind == kUnary || merger.exp(exp).kind == kBinary);
     } else if (merger.exp(exp).kind == kSelect) {
       // Select operation insertion.
       Value insChain = codegen.insChain;
       assert(insChain);
-      scf::IfOp ifOp = builder.create<scf::IfOp>(
-          loc, insChain.getType(), rhs, /*else=*/true);
+      scf::IfOp ifOp = builder.create<scf::IfOp>(loc, insChain.getType(), rhs,
+                                                 /*else=*/true);
       builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
       // Existing value was preserved to be used here.
       assert(merger.exp(exp).val);
       Value v0 = merger.exp(exp).val;
       genInsertionStore(codegen, builder, op, t, v0);
       merger.exp(exp).val = Value();
       // Yield modified insertion chain along true branch.
       builder.create<scf::YieldOp>(op.getLoc(), codegen.insChain);
LogicalResult matchAndRewrite(linalg::GenericOp op,
 
     merger.setHasSparseOut(sparseOut != nullptr);
 
     SmallVector<Value> tensors;
     for (OpOperand &t : op->getOpOperands())
       tensors.push_back(t.get());
 
     // Recursively generates code if admissible.
-    CodeGen codegen(options, tensors, numTensors, numLoops, sparseOut,
-                    outerParNest, topSort);
+    CodeGen codegen(options, op.getContext(), tensors, numTensors, numLoops,
+                    sparseOut, outerParNest, topSort);
     genBuffers(merger, codegen, rewriter, op);
     genStmt(merger, codegen, rewriter, op, exp, 0);
     genResult(merger, codegen, rewriter, op);
     return success();
   }
 
 private:
   // Last resort cycle resolution.