[mlir][sparse] Add new utility class to help generates loop structures over sparse tensors; Implement foreach operator.

[Peiming]

Update comments

[Peiming]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_codegen_foreach.mlir

110

Note how this sequence is printed in col-major order

[Peiming]

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

373

Should we always lowering foreach? or only when RT is disabled? I think the operator is general enough to be used for conversion as well.

[Peiming]

try fixing windows fail...

[Peiming]

code cleanup.

[Peiming]

Update comments.

[Peiming]

rebase + remove dup utils

[bixia]

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

346–362

Can you use RewriterBase::MergeBlocks to avoid manually remap like this?

[bixia]

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

307–311

The style guide actually discourage the use of "auto" unless the type is obvious from the context (line 309) or it makes the code safer.

[aartbik]

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

26

structure
over sparse tensors

Note that this also shows the danger of repeating comments. In general, google style says that if something is documented in the header, it does not need to be repeated in cc file. I would just pick a shorter "header" comment here to separate the parts

65

period

138

note that with an eye on more dim level types, I like to make these asserts very defensive.
So
assert(isCompressed || isDenseDim()) would capture future omissions

240

you can remove this one

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

29

yeah, as you suggested yourself, you can merge these with the sparsetensor.h ones

35

oh nice, I just did this, and renamed it to toOrigDim

49

loop structure

to (co-iterate) over sparse tensors?

52

relies -> rely

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

334

if we don't do this, do we break debugging path?

If it is not expensive, shall we always do it, just to avoid inconsistent paths under prod/debug?

373

Yeah, if it is there, we lower (right now when RT is enabled, we won't see it).
You can add an assert if you are worried?

[bixia]

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

299

the foreach

[Peiming]

address comments

[Peiming]

fix comments

[Peiming]

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

334

No, it affects nothing in this case.

I can remove it.

373

Okay, I will then make it accessible universally (I just found that my integrate test did not disableRT either.

[Peiming]

remove spaces

[Peiming]

rebase + remove dup utils.

[Peiming]

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

346–362

Thanks! It does make the code simpler!

[Peiming]

fix errors.

[Peiming]

move assert into if.

[aartbik]

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

30–31

Emitter (one m in comment)

37

Add a doc on what constructor expects (the list of tensors)

45

period at end?

53

period at end?

57

I just talked with the IREE team and we may find a better separation of codegen concerns by migrating the vectorization and parallelization to a downstream pass. If it simplifies things for you know, I would omit the parallel/vector capabilities from these utils for now

73

period at end?

119

sparis -> sparsification

125

for each tensor

129

A lot of these are taken from the codegen structured (and sometimes renamed, which makes "recognition" a bit harder for me)
But can you at least copy the comments from the codegen struct here, so that each section has a bit of a explanation

E.g.

/ Sparse iteration information (by tensor and index). These arrays
/ are updated to remain current within the current loop.

was useful to explain the vector of vector

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

314

I like how this reads now!

334

int64_t when dealing with rank?

345

same

[Peiming]

update comments.

[Peiming]

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

129

One thing that need to be mentioned is that all the variables here now are indexed by tensor+dim instead of tensor+loop idx.

[Peiming]

Address comments

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

57

Sounds good!

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

314

;-)

[Peiming]

rebase against main.

[aartbik]

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

74

I like the idea of encforcing proper API usage through asserts and other means a lot.
Can we add some more? E.g. here you can check that nothing has been called yet, so you don't try to use this method again later.

Same for methods below.

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

34

Please document the class also. In general, showing "usage" is very illustrative. So you can give a very short pseudo-code example of starting entering exiting and leaving a loop using the API of the class.

55

typo: generated

80

This is not really a getter. Just give it its own doc stating that it will return the coordinate lists in the passed array

85

this is a lot of code for a header, and it distracts from reading the API

why not move this block into the cpp file?

109

I am not sure I am a big fan of using two completely different data structures for debug/prod, since that has the risk of introducing bugs on one path and not the other. Can't we just simply always keep all the bookkeeping with proper detection of wrong usage?

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

333

Actually, I would remove this comment completely. The API of the emitter requires every loop to be opened/closed properly (and asserting on this behavior is a good thing). There is really no good benefit telling the reader here that a "shortcut" usage is possible ;-)

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_codegen_foreach.mlir

2

This seems a copy and paste error (kDirect only impacts sparse2sparse)

16

Can we use the common names for this, DCSC etc.

23

Use common comment format with capital/period at end

25

such a nifty test for this construct!

[Peiming]

address comments.

[Peiming]

update names

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

74

see line.90.

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

55

Really need to configure a spell checker on my editor ;-(

109

Make sense, yet another pre-mature optimization ;-(

[Peiming]

rebase

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

 
 namespace mlir {
 class Location;
 class Type;
 class Value;
 
 namespace sparse_tensor {
 
 //===----------------------------------------------------------------------===//
-// ExecutionEngine/SparseTensorUtils helper functions.
+// SparseTensorLoopEmmiter class, manages sparse tensors and helps to generate
+// loop structure to (co-iterate) sparse tensors.
 //===----------------------------------------------------------------------===//
 
+// TODO: Sparsification should also rely on this class to generate loops.
+class SparseTensorLoopEmitter {
+public:
+  explicit SparseTensorLoopEmitter(ValueRange tensors);
+
+  ///
+  /// Core functions.
+  ///
+
+  /// Starts a loop emitting session:
+  /// 1. Generates all the buffers needed to iterate tensors.
+  /// 2. Generates the lo/hi bounds to iterate tensors[0]
+  void initializeLoopEmit(OpBuilder &builder, Location loc);
+
+  // TODO: Gets rid of `dim` in the argument list? Track the dimension we
+  // are currently at internally. Then it would be enterNextDimForTensor.
+
+  /// Emits loop over tensor[dim], it assumes that loops between
+  /// tensor[0...dim - 1] have already been genereted.
+  /// It also prepares to enter tensor[dim + 1]
+  Operation *enterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
+                                      size_t tid, size_t dim,
+                                      ArrayRef<Value> reduc = {},
+                                      bool isParallel = false,
+                                      bool isVector = false);
+
+  /// Emits a coiteration loop over a set of tensors.
+  // TODO: not yet implemented
+  void enterCoiterationOverTensorsAtDims(OpBuilder &builder, Location loc,
+                                         ArrayRef<size_t> ts,
+                                         ArrayRef<size_t> ds);
+
+  /// Emits extra locals, since the locals might not be in simplified lattices
+  /// point used to generate the loops, but are still required to generates
+  /// expressions.
+  Value emitExtraLocalsForTensorsAtDims(OpBuilder &builder, Location loc,
+                                        size_t tid, size_t dim);
+
+  ///
+  /// Getters
+  ///
+  Value getTensorValueBuffer(size_t tid) { return valBuffer[tid]; }
+  Value getLastLevelTensorPointerIndex(size_t tid) {
+    return pidxs[tid].back();
+  };
+  void getCoordinateArray(SmallVectorImpl<Value> &coords) {
+    for (auto &l : loopStack)
+      coords.push_back(l.idx);
+  }
+
+  void exitCurrentLoop() {
+#ifndef NDEBUG
+    // Clean up the values, it would help use to discover potential bug at a
+    // earlier stage (instead of silently using a wrong value).
+    LoopLevelInfo &loopInfo = loopStack.back();
+    assert(loopInfo.tensors.size() == loopInfo.dims.size());
+    for (auto info : llvm::zip(loopInfo.tensors, loopInfo.dims)) {
+      auto tid = std::get<0>(info);
+      auto dim = std::get<1>(info);
+      assert(pidxs[tid][dim] && coord[tid][dim] && highs[tid][dim]);
+      // Reset to null.
+      pidxs[tid][dim] = Value();
+      coord[tid][dim] = Value();
+      if (!isDenseDim(dims[tid][dim]))
+        // Dense dimension, high is fixed.
+        highs[tid][dim] = Value();
+    }
+#endif
+    loopStack.pop_back();
+  }
+
+private:
+  struct LoopLevelInfo {
+#ifndef NDEBUG
+    LoopLevelInfo(ArrayRef<size_t> ts, ArrayRef<size_t> ds, Value idx)
+        : tensors(ts), dims(ds), idx(idx) {}
+    llvm::SmallVector<size_t, 4> tensors;
+    llvm::SmallVector<size_t, 4> dims;
+#else
+    LoopLevelInfo(Value idx) : idx(idx) {}
+#endif
+    Value idx;
+  };
+
+  /// Return false if tid[dim] is a dense dimension that does not need to be
+  /// prepared (to be used by sparisfication for needUniv).
+  bool prepareLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid,
+                                  size_t dim);
+
+  /// Input (TODO: and output) tensors.
+  std::vector<Value> tensors;
+  /// The dim type array for each tensors.
+  std::vector<std::vector<SparseTensorEncodingAttr::DimLevelType>> dims;
+  std::vector<std::vector<Value>> pidxs;
+  std::vector<std::vector<Value>> coord;
+  std::vector<std::vector<Value>> highs;
+  std::vector<std::vector<Value>> sizes;
+  std::vector<std::vector<Value>> ptrBuffer; // to_pointers
+  std::vector<std::vector<Value>> idxBuffer; // to_indices
+  std::vector<Value> valBuffer;              // to_value
+
+  std::vector<LoopLevelInfo> loopStack;
+  std::vector<size_t> curLv;
+};
+
+//===----------------------------------------------------------------------===//
+// ExecutionEngine/SparseTensorUtils helper functions.
+//===----------------------------------------------------------------------===//
+//
 /// Converts an overhead storage bitwidth to its internal type-encoding.
 OverheadType overheadTypeEncoding(unsigned width);
 
 /// Converts an overhead storage type to its internal type-encoding.
 OverheadType overheadTypeEncoding(Type tp);
 
 /// Converts the internal type-encoding for overhead storage to an mlir::Type.
 Type getOverheadType(Builder &builder, OverheadType ot);

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

 //===- CodegenUtils.cpp - Utilities for generating MLIR -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
+/// Generates a pointer/index load from the sparse storage scheme. Narrower
+/// data types need to be zero extended before casting the value into the
+/// index type used for looping and indexing.
+static Value genIndexLoad(OpBuilder &builder, Location loc, Value ptr,
+                          Value s) {
+  // For the scalar case, we simply zero extend narrower indices into 64-bit
+  // values before casting to index without a performance penalty. Here too,
+  // however, indices that already are 64-bit, in theory, cannot express the
+  // full range as explained above.
+  Value load = builder.create<memref::LoadOp>(loc, ptr, s);
+  if (!load.getType().isa<IndexType>()) {
+    if (load.getType().getIntOrFloatBitWidth() < 64)
+      load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
+    load =
+        builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
+  }
+  return load;
+}
+
+//===----------------------------------------------------------------------===//
+// Sparse tensor loop emitter class implementations
+//===----------------------------------------------------------------------===//
+
+SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors)
+    : tensors(tensors.begin(), tensors.end()), dims(tensors.size()),
+      pidxs(tensors.size()), coord(tensors.size()), highs(tensors.size()),
+      sizes(tensors.size()), ptrBuffer(tensors.size()),
+      idxBuffer(tensors.size()), valBuffer(tensors.size()), loopStack(),
+      curLv(tensors.size(), 0) {
+  for (size_t i = 0, e = tensors.size(); i < e; i++) {
+    auto t = tensors[i];
+    auto rtp = t.getType().cast<RankedTensorType>();
+    auto rank = static_cast<size_t>(rtp.getRank());
+    auto enc = getSparseTensorEncoding(rtp);
+    if (enc)
+      for (auto dimTp : enc.getDimLevelType())
+        dims[i].push_back(dimTp);
+    else
+      dims[i].assign(rank, SparseTensorEncodingAttr::DimLevelType::Dense);
+
+    // Initialize using empty value.
+    pidxs[i].assign(rank, Value());
+    coord[i].assign(rank, Value());
+    highs[i].assign(rank, Value());
+    sizes[i].assign(rank, Value());
+    ptrBuffer[i].assign(rank, Value());
+    idxBuffer[i].assign(rank, Value());
+  }
+}
+
+void SparseTensorLoopEmitter::initializeLoopEmit(OpBuilder &builder,
+                                                 Location loc) {
+  // For every tensor, find lower and upper bound on dimensions, set the
+  // same bounds on loop indices, and obtain dense or sparse buffer(s).
+  // TODO: Provides ability to generate loop on output buffer (with undef
+  // dim level in Merger in GenericOp Sparsification).
+  for (size_t t = 0, e = tensors.size(); t < e; t++) {
+    auto tensor = tensors[t];
+    auto rtp = tensor.getType().cast<RankedTensorType>();
+    auto rank = rtp.getRank();
+    auto shape = rtp.getShape();
+    auto enc = getSparseTensorEncoding(rtp);
+    auto dynShape = {ShapedType::kDynamicSize};
+    // Scan all dimensions of current tensor.
+    for (int64_t d = 0; d < rank; d++) {
+      // Handle sparse storage schemes.
+      if (isCompressedDim(dims[t][d])) {
+        auto ptrTp =
+            MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
+        auto indTp =
+            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
+        auto dim = builder.getIndexAttr(d);
+        // Generate sparse primitives to obtains pointer and indices.
+        ptrBuffer[t][d] = builder.create<ToPointersOp>(loc, ptrTp, tensor, dim);
+        idxBuffer[t][d] = builder.create<ToIndicesOp>(loc, indTp, tensor, dim);
+      } else if (isSingletonDim(dims[t][d])) {
+        llvm_unreachable("TODO: not implemented yet");
+      }
+
+      // Find upper bound in current dimension.
+      unsigned p = toOrigDim(enc, d);
+      Value up = mlir::linalg::createOrFoldDimOp(builder, loc, tensor, p);
+      sizes[t][d] = highs[t][d] = up;
+    }
+    // Perform the required bufferization. Dense inputs materialize
+    // from the input tensors. Dense outputs need special handling.
+    // Sparse inputs use sparse primitives to obtain the values.
+    Type elementType = rtp.getElementType();
+
+    if (!enc) {
+      // Non-annotated dense tensors.
+      auto denseTp = MemRefType::get(shape, elementType);
+      // This is not the output tensor
+      valBuffer[t] =
+          builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
+    } else {
+      // Annotated sparse tensors.
+      auto dynShape = {ShapedType::kDynamicSize};
+      auto sparseTp = MemRefType::get(dynShape, elementType);
+      valBuffer[t] = builder.create<ToValuesOp>(loc, sparseTp, tensor);
+    }
+    // Prepare to enter the first dim for all (input) tensors
+    prepareLoopOverTensorAtDim(builder, loc, t, 0);
+  }
+}
+
+Operation *SparseTensorLoopEmitter::enterLoopOverTensorAtDim(
+    OpBuilder &builder, Location loc, size_t tid, size_t dim,
+    ArrayRef<Value> reduc, bool isParallel, bool isVector) {
+  assert(!isVector && "unimplemented!");
+  assert(dims[tid].size() > dim);
+  Value step = constantIndex(builder, loc, 1);
+  bool isCompressed = isCompressedDim(dims[tid][dim]);
+  assert(isDenseDim(dims[tid][dim]) || isCompressedDim(dims[tid][dim]));
+
+  Value lo = isCompressed ? pidxs[tid][dim] : constantIndex(builder, loc, 0);
+  Value hi = highs[tid][dim];
+  Operation *loop = nullptr;
+  Value iv; // induction variable
+  if (isParallel) {
+    // Can not be parallel + vector or parallel + reduction at the same time.
+    // TODO: support reduction on parallel for.
+    assert(!isVector && reduc.empty());
+    scf::ParallelOp parOp = builder.create<scf::ParallelOp>(loc, lo, hi, step);
+    builder.setInsertionPointToStart(parOp.getBody());
+    iv = parOp.getInductionVars()[0];
+    loop = parOp;
+  } else {
+    // TODO: different steps value for vector
+    // TODO: support reduction.
+    scf::ForOp forOp;
+    if (!reduc.empty())
+      llvm_unreachable("TODO: not implemented yet");
+    else
+      forOp = builder.create<scf::ForOp>(loc, lo, hi, step, reduc);
+    builder.setInsertionPointToStart(forOp.getBody());
+    iv = forOp.getInductionVar();
+    loop = forOp;
+  }
+
+  assert(iv);
+  if (isCompressed) {
+    pidxs[tid][dim] = iv;
+    // Generating a load on the indices array yields the coordinate.
+    Value ptr = idxBuffer[tid][dim];
+    // TODO: generates load for vector value.
+    coord[tid][dim] = genIndexLoad(builder, loc, ptr, iv);
+  } else {
+    // Dense tensor, the coordinates is the inducation variable.
+    coord[tid][dim] = iv;
+    // generate pidx for dense dim (pidx = i * sz + j)
+    // TODO: handle vector loop.
+    Value p = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
+    Value mul = builder.create<arith::MulIOp>(loc, sizes[tid][dim], p);
+    Value add = builder.create<arith::AddIOp>(loc, mul, iv);
+    pidxs[tid][dim] = add;
+  }
+
+  // Prepares for next dim if this is not currently the innermost dimension.
+  if (dim != dims[tid].size() - 1)
+    prepareLoopOverTensorAtDim(builder, loc, tid, dim + 1);
+
+#ifndef NDEBUG
+  loopStack.push_back(LoopLevelInfo({tid}, {dim}, coord[tid][dim]));
+#else
+  loopStack.push_back(coord[tid][dim]);
+#endif
+  return loop;
+}
+
+void SparseTensorLoopEmitter::enterCoiterationOverTensorsAtDims(
+    OpBuilder &builder, Location loc, ArrayRef<size_t> ts,
+    ArrayRef<size_t> ds) {
+  llvm_unreachable("TODO: unimplemented");
+}
+
+bool SparseTensorLoopEmitter::prepareLoopOverTensorAtDim(OpBuilder &builder,
+                                                         Location loc,
+                                                         size_t tid,
+                                                         size_t dim) {
+  // TODO: generate loop iteration on output tensor based on the shape
+  // instead of pointer/indices arrays.
+  assert(dims[tid].size() > dim);
+
+  if (isDenseDim(dims[tid][dim]))
+    return false;
+
+  // Either the first dimension, or the previous dimension has been set.
+  assert(dim == 0 || pidxs[tid][dim - 1]);
+  if (isCompressedDim(dims[tid][dim])) {
+    Value ptr = ptrBuffer[tid][dim];
+    Value c1 = constantIndex(builder, loc, 1);
+    Value pLo = dim == 0 ? constantIndex(builder, loc, 0) : pidxs[tid][dim - 1];
+    Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
+
+    pidxs[tid][dim] = genIndexLoad(builder, loc, ptr, pLo);
+    highs[tid][dim] = genIndexLoad(builder, loc, ptr, pHi);
+
+    return true;
+  }
+
+  if (isSingletonDim(dims[tid][dim]))
+    llvm_unreachable("TODO: not implemented yet");
+
+  llvm_unreachable("Unrecognizable dimesion type!");
+}
+
+Value SparseTensorLoopEmitter::emitExtraLocalsForTensorsAtDims(
+    OpBuilder &builder, Location loc, size_t tid, size_t dim) {
+  llvm_unreachable("TODO: not implemented yet");
+}
+
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
 //===----------------------------------------------------------------------===//
 
 OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) {
   switch (width) {
   case 64:
     return OverheadType::kU64;
   case 32:
     return OverheadType::kU32;
   case 16:
     return OverheadType::kU16;
   case 8:

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
LogicalResult matchAndRewrite(ReshapeOp op,
     if (encSrc) {
       RankedTensorType rtp =
           op.getSrc().getType().template cast<RankedTensorType>();
       auto denseTp =
           RankedTensorType::get(rtp.getShape(), rtp.getElementType());
       auto convert = rewriter.create<ConvertOp>(loc, denseTp, op.getSrc());
       op->setOperand(0, convert);
       return success();
-    } else if (encDst) {
+    }
+    if (encDst) {
       RankedTensorType rtp =
           op.getResult().getType().template cast<RankedTensorType>();
       auto denseTp =
           RankedTensorType::get(rtp.getShape(), rtp.getElementType());
       auto reshape = rewriter.create<ReshapeOp>(loc, denseTp, op.getSrc(),
                                                 op.getReassociation());
       Value convert = rewriter.create<ConvertOp>(loc, rtp, reshape);
       rewriter.replaceOp(op, convert);
       return success();
     }
     return failure();
   }
 };
 
+/// Sparse rewriting rule for the foreach operator.
+struct ForeachRewriter : public OpRewritePattern<ForeachOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForeachOp op,
+                                PatternRewriter &rewriter) const override {
+
+    auto loc = op.getLoc();
+    Value input = op.getTensor();
+    auto rtp = input.getType().cast<RankedTensorType>();
+    int64_t rank = rtp.getRank();
+    auto enc = getSparseTensorEncoding(rtp);
+
+    // 1. Generates loop for the sparse input.
+    SparseTensorLoopEmitter loopEmitter(ValueRange{input});
+    loopEmitter.initializeLoopEmit(rewriter, loc);
+    for (int64_t i = 0; i < rank; i++)
+      loopEmitter.enterLoopOverTensorAtDim(rewriter, loc, 0, i);
+
+    Value vals = loopEmitter.getTensorValueBuffer(0);
+    Value idx = loopEmitter.getLastLevelTensorPointerIndex(0);
+    Value val = rewriter.create<memref::LoadOp>(op.getLoc(), vals, idx);
+
+    SmallVector<Value, 4> coords;
+    coords.reserve(rank);
+    loopEmitter.getCoordinateArray(coords);
+
+    // NOTE: we do not actually need to exit the loop since the loop
+    // structure is simple. i.e., no complex composition as following:
+    // for {
+    //   for {}
+    //   for {}
+    // }
+    // We do it here solely for debuging purposes.
+    for (int i = 0; i < rank; i++)
+      loopEmitter.exitCurrentLoop();
+
+    // 2. Inline the block in the foreach operator.
+    Block::iterator inlinePos = rewriter.getInsertionPoint();
+    Block *srcBlock = op.getBody();
+    // Remove sparse_tensor.yield.
+    rewriter.eraseOp(srcBlock->getTerminator());
+
+    SmallVector<Value, 4> args;
+    // Remap coordinates.
+    for (int i = 0; i < rank; i++) {
+      Value actual = coords[toOrigDim(enc, i)];
+      args.push_back(actual);
+    }
+    // Remap value.
+    args.push_back(val);
+
+    // Inline body.
+    rewriter.mergeBlockBefore(srcBlock, &*inlinePos, args);
+    // delete the foreach operator.
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 
 void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns,
-                                         bool /*enableRT*/) {
+                                         bool enableRT) {
   patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd,
                ReshapeRewriter<tensor::ExpandShapeOp>,
-               ReshapeRewriter<tensor::CollapseShapeOp>>(patterns.getContext());
+               ReshapeRewriter<tensor::CollapseShapeOp>, ForeachRewriter>(
+      patterns.getContext());
   // TODO: If RT not enabled, rewrite concatenate ops, etc here.
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_codegen_foreach.mlir

+// Force this file to use the kDirect method for sparse2sparse.
+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#Tensor1 = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "dense" ]
+}>
+
+#Tensor2 = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ]
+}>
+
+#Tensor3 = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+module {
+
+  // uses foreach operator to print coords and values
+  func.func @foreach_print_1(%arg0: tensor<2x2xf64, #Tensor1>) {
+    sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #Tensor1> do {
+      ^bb0(%1: index, %2: index, %v: f64) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %v: f64
+     }
+     return
+  }
+
+  func.func @foreach_print_2(%arg0: tensor<2x2xf64, #Tensor2>) {
+    sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #Tensor2> do {
+      ^bb0(%1: index, %2: index, %v: f64) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %v: f64
+     }
+     return
+  }
+
+  func.func @foreach_print_3(%arg0: tensor<2x2xf64, #Tensor3>) {
+    sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #Tensor3> do {
+      ^bb0(%1: index, %2: index, %v: f64) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %v: f64
+     }
+     return
+  }
+
+  //
+  // Main driver.
+  //
+  func.func @entry() {
+    //
+    // Initialize a 3-dim dense tensor.
+    //
+    %src = arith.constant dense<
+       [[  1.0,  2.0],
+        [  5.0,  6.0]]
+    > : tensor<2x2xf64>
+
+    //
+    // Convert dense tensor directly to various sparse tensors.
+    //
+    %s1 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Tensor1>
+    %s2 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Tensor2>
+    %s3 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Tensor3>
+    // CHECK: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
+    call @foreach_print_1(%s1) : (tensor<2x2xf64, #Tensor1>) -> ()
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
+    call @foreach_print_2(%s2) : (tensor<2x2xf64, #Tensor2>) -> ()
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
+    call @foreach_print_3(%s3) : (tensor<2x2xf64, #Tensor3>) -> ()
+    
+    bufferization.dealloc_tensor %s1 : tensor<2x2xf64, #Tensor1>
+    bufferization.dealloc_tensor %s2 : tensor<2x2xf64, #Tensor2>
+    bufferization.dealloc_tensor %s3 : tensor<2x2xf64, #Tensor3>
+
+    return
+  }
+}