diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h --- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h @@ -331,7 +331,9 @@ /// 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 @@ -413,11 +415,20 @@ }; const std::vector &getValBuffer() const { return valBuffer; }; + constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName() { + return llvm::StringLiteral("Emitted from"); + } + private: struct LoopLevelInfo { LoopLevelInfo(ArrayRef tids, ArrayRef 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 for tid and dim? // The set of tensors that the loop is operating on const llvm::SmallVector tids; @@ -485,8 +496,12 @@ void exitCoIterationLoop(OpBuilder &builder, Location loc, MutableArrayRef 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. diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp @@ -95,9 +95,10 @@ //===----------------------------------------------------------------------===// 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()), @@ -284,7 +285,7 @@ // NOTE: we can also prepares for next dim here in advance // Push the loop into stack loopStack.emplace_back(ArrayRef(tid), ArrayRef(dim), loop, - coord[tid][dim]); + coord[tid][dim], loopTag); // Emit extra locals. emitExtraLocalsForTensorsAtDenseDims(builder, loc, extraTids, extraDims); @@ -386,7 +387,7 @@ // 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 diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp @@ -789,7 +789,9 @@ 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 diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp @@ -53,11 +53,15 @@ // 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 &ts) - : options(o), loopEmitter(tensors, /*hasOutput=*/true, - /*isSparseOut=*/op != nullptr), + : options(o), + loopEmitter( + tensors, + StringAttr::get(context, linalg::GenericOp::getOperationName()), + /*hasOutput=*/true, + /*isSparseOut=*/op != nullptr), sparseOut(op), outerParNest(nest), topSort(ts) { if (op) insChain = op->get(); @@ -670,8 +674,8 @@ // Select operation insertion. Value insChain = codegen.insChain; assert(insChain); - scf::IfOp ifOp = builder.create( - loc, insChain.getType(), rhs, /*else=*/true); + scf::IfOp ifOp = builder.create(loc, insChain.getType(), rhs, + /*else=*/true); builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); // Existing value was preserved to be used here. assert(merger.exp(exp).val); @@ -1372,8 +1376,8 @@ 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);