mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp

//===- LoopEmitter.cpp ----------------------------------------------------===//
//
// 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 "LoopEmitter.h"
#include "CodegenUtils.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

using namespace mlir;
using namespace mlir::sparse_tensor;

//===----------------------------------------------------------------------===//
// File local shorthand macros
//===----------------------------------------------------------------------===//

#define CMPI(p, l, r)                                                          \
  (arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::p, (l), (r))      \
       .getResult())

#define C_IDX(v) (constantIndex(builder, loc, (v)))
#define YIELD(vs) (scf::YieldOp::create(builder, loc, (vs)))
#define ADDI(lhs, rhs) (arith::AddIOp::create(builder, loc, (lhs), (rhs)))
#define ANDI(lhs, rhs) (arith::AndIOp::create(builder, loc, (lhs), (rhs)))
#define SUBI(lhs, rhs) (arith::SubIOp::create(builder, loc, (lhs), (rhs)))
#define MULI(lhs, rhs) (arith::MulIOp::create(builder, loc, (lhs), (rhs)))
#define REMUI(lhs, rhs) (arith::RemUIOp::create(builder, loc, (lhs), (rhs)))
#define DIVUI(lhs, rhs) (arith::DivUIOp::create(builder, loc, (lhs), (rhs)))
#define SELECT(c, l, r) (arith::SelectOp::create(builder, loc, (c), (l), (r)))

//===----------------------------------------------------------------------===//
// Debugging utils
//===----------------------------------------------------------------------===//

#ifndef NDEBUG
[[maybe_unused]] static void dumpIndexMemRef(OpBuilder &builder, Location loc,
                                             Value memref) {
  memref = memref::CastOp::create(
      builder, loc, UnrankedMemRefType::get(builder.getIndexType(), 0), memref);
  createFuncCall(builder, loc, "printMemrefInd", TypeRange{},
                 ValueRange{memref}, EmitCInterface::On);
}
#endif

//===----------------------------------------------------------------------===//
// File local helper functions.
//===----------------------------------------------------------------------===//

// For index reduction loops, since the tensor are sliced into non-continuous
// fragments, we need a triple [pLo, pHi, pPtr], in which the pair (pLo, pHi)
// specifies the range of the fragment, and pPtr specifies the index of the
// corresponding fragment in the child level (i.e., a pointer to the sliced
// position array).
static Value genSliceOffset(OpBuilder &builder, Location loc, Value tensor,
                            Level lvl) {
  auto enc = getSparseTensorEncoding(tensor.getType());
  return createOrFoldSliceOffsetOp(builder, loc, tensor, toDim(enc, lvl));
}

static Value genSliceStride(OpBuilder &builder, Location loc, Value tensor,
                            Level lvl) {
  auto enc = getSparseTensorEncoding(tensor.getType());
  return createOrFoldSliceStrideOp(builder, loc, tensor, toDim(enc, lvl));
}

static bool isIntOrFPZero(Attribute attr) {
  if (auto f = llvm::dyn_cast<FloatAttr>(attr); f && f.getValue().isZero())
    return true;
  if (auto i = llvm::dyn_cast<IntegerAttr>(attr); i && i.getValue().isZero())
    return true;
  return false;
}

static Value unFoldOpIntResult(OpBuilder &builder, Location loc,
                               OpFoldResult ofr) {
  if (std::optional<int64_t> i = getConstantIntValue(ofr); i.has_value())
    return constantIndex(builder, loc, *i);
  return cast<Value>(ofr);
}

static Value tryFoldTensors(Value t) {
  // TODO: this should be done through a folding pass after switching to
  // `sparse_tensor.iterate`-based sparsification.
  auto stt = tryGetSparseTensorType(t);
  auto padOp = t.getDefiningOp<tensor::PadOp>();
  if (padOp && stt.has_value() && stt->hasEncoding() &&
      padOp.getSourceType().getEncoding() == stt->getEncoding() &&
      stt->getEncoding().isIdentity()) {
    // Try fusing padOp with zeros.
    Attribute padCst;
    if (matchPattern(padOp.getBody()->getTerminator(),
                     m_Op<tensor::YieldOp>(m_Constant(&padCst))) &&
        isIntOrFPZero(padCst)) {
      return padOp.getSource();
    }
  }
  return t;
}

//===----------------------------------------------------------------------===//
// Sparse tensor loop emitter class implementations
//===----------------------------------------------------------------------===//

LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                         bool isSparseOut, unsigned numLoops,
                         DependentLvlGetter dimGetter,
                         SparseEmitStrategy emitStrategy) {
  initialize(tensors, loopTag, hasOutput, isSparseOut, numLoops, dimGetter);
}

void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
                             bool isSparseOut, unsigned numLoops,
                             DependentLvlGetter dimGetter,
                             SparseEmitStrategy emitStrategy) {
  // First initialize the top-level type of the fields.
  this->loopTag = loopTag;
  this->hasOutput = hasOutput;
  this->isSparseOut = isSparseOut;
  this->emitStrategy = emitStrategy;

  const unsigned numManifestTensors = ts.size();
  const unsigned synTensorId = numManifestTensors;
  const unsigned numTensors = numManifestTensors + 1;
  // tensors array (len == numManifestTensor).
  this->tensors.assign(ts.begin(), ts.end());
  // Arrays with len == numTensor.
  this->valBuffer.assign(numTensors, nullptr);
  this->lvls.resize(numTensors);
  this->iters.resize(numTensors);
  this->spIterVals.resize(numTensors);

  // These zeros will be overwritten below, but we need to initialize
  // them to something since we'll need random-access assignment.
  this->loopStack.reserve(numLoops);
  this->loopSeqStack.reserve(numLoops);

  // Index-reduction related fields.
  this->dependentLvlMap.assign(
      numTensors, std::vector<std::vector<std::pair<TensorLevel, unsigned>>>());
  this->sliceMeta.assign(
      numTensors, std::vector<std::vector<std::pair<Value, unsigned>>>());
  this->levelReducedDep.assign(numTensors, std::vector<unsigned>());

  // Initialize nested types of `TensorId`-indexed fields.
  for (TensorId tid = 0; tid < numTensors; tid++) {
    Level lvlRank;
    if (tid == synTensorId) {
      // Synthetic tensor (conceptually) is an all-dense tensor with rank equal
      // to the total number of loops (each level can potentially be mapped to
      // one of the loop being generated).
      lvlRank = numLoops;
    } else {
      const Value t = tensors[tid];
      // a scalar or 0-dimension tensors
      if (isZeroRankedTensorOrScalar(t.getType()))
        continue;

      auto rtp = getRankedTensorType(t);
      const SparseTensorType stt(rtp);
      lvlRank = stt.getLvlRank();
    }

    lvls[tid].resize(lvlRank);
    iters[tid].resize(lvlRank);
    spIterVals[tid].resize(lvlRank);
    loopHighs.assign(numLoops, nullptr);

    // Slice-driven loops related initialization.
    levelReducedDep[tid].assign(lvlRank, 0);
    dependentLvlMap[tid].assign(
        lvlRank, std::vector<std::pair<TensorLevel, unsigned>>());
    sliceMeta[tid].assign(lvlRank, std::vector<std::pair<Value, unsigned>>());
    if (dimGetter && !isSynTensor(tid)) {
      for (Level l = 0; l < lvlRank; l++) {
        std::vector<std::pair<LoopId, unsigned>> deps = dimGetter(tid, l);
        // Sort the loop by order.
        llvm::sort(deps, llvm::less_first());

        dependentLvlMap[tid][l] = std::move(deps);
        unsigned depends = dependentLvlMap[tid][l].size();
        if (depends == 0)
          continue;
        sliceMeta[tid][l].reserve(depends);
      }
    }
  }
}

std::unique_ptr<SparseIterator>
LoopEmitter::makeLevelIterator(OpBuilder &builder, Location loc, TensorId t,
                               Level l) {
  Value tensor = tensors[t];
  auto stt = getSparseTensorType(tensor);
  auto it = makeSimpleIterator(*lvls[t][l], emitStrategy);

  Value folded = tryFoldTensors(tensor);
  if (folded != tensor) {
    auto padOp = tensor.getDefiningOp<tensor::PadOp>();
    assert(padOp);
    if (padOp.getPaddedDims().test(l)) {
      Value low = unFoldOpIntResult(builder, loc, padOp.getMixedLowPad()[l]);
      Value high = unFoldOpIntResult(builder, loc, padOp.getMixedHighPad()[l]);
      auto padIt = makePaddedIterator(std::move(it), low, high, emitStrategy);
      return padIt;
    }
  }

  if (stt.hasEncoding() && stt.getEncoding().isSlice()) {
    Value offset = genSliceOffset(builder, loc, tensor, l);
    Value stride = genSliceStride(builder, loc, tensor, l);
    auto slicedIt = makeSlicedLevelIterator(
        std::move(it), offset, stride, lvls[t][l]->getSize(), emitStrategy);
    return slicedIt;
  }

  return it;
}

void LoopEmitter::initializeLoopEmit(
    OpBuilder &builder, Location loc, LoopEmitter::OutputUpdater updater,
    LoopEmitter::SynTensorBoundSetter synSetter) {

  // For every manifest tensor, set up the values buffer.
  for (TensorId t = 0, numTensors = getNumManifestTensors(); t < numTensors;
       t++) {
    // TODO: this should be done through a folding pass after switching to
    // `sparse_tensor.iterate`-based sparsification.
    const Value tensor = tryFoldTensors(tensors[t]);
    const auto rtp = dyn_cast<RankedTensorType>(tensor.getType());
    // Skips only scalar, zero ranked tensor still need to be bufferized and
    // (probably) filled with zeros by users.
    if (!rtp)
      continue;

    auto stt = getSparseTensorType(tensor);
    const auto shape = rtp.getShape();

    // Perform the required bufferization. Dense inputs materialize from the
    // input tensors. Sparse inputs use sparse primitives to obtain the values.
    // Delegates extra output initialization to clients.
    bool isOutput = isOutputTensor(t);
    Type elementType = stt.getElementType();
    if (!stt.hasEncoding()) {
      // Non-annotated dense tensors.
      BaseMemRefType denseTp = MemRefType::get(shape, elementType);

      // TODO: if we unconditionally use fully dynamic layout here, it breaks
      // some vectorization passes which requires static stride = 1.
      // Is it possible to call vectorization pass after bufferization?
      if (llvm::isa_and_nonnull<tensor::ExtractSliceOp>(tensor.getDefiningOp()))
        denseTp = bufferization::getMemRefTypeWithFullyDynamicLayout(rtp);

      Value denseVal =
          bufferization::ToBufferOp::create(builder, loc, denseTp, tensor);
      // Dense outputs need special handling.
      if (isOutput && updater)
        denseVal = updater(builder, loc, denseVal, tensor);

      valBuffer[t] = denseVal;
    } else {
      // Annotated sparse tensors.
      // We also need the value buffer for all-dense annotated "sparse"
      // tensors.
      valBuffer[t] = ToValuesOp::create(builder, loc, tensor);
    }
  }

  // The sparse iterator values will only be available after the loop is
  // constructed.
  if (emitStrategy == SparseEmitStrategy::kSparseIterator)
    return;

  // For every synthetic tensor, set the high bound by calling the callback.
  if (synSetter) {
    TensorId synId = getSynTensorId();
    for (unsigned i = 0, e = loopHighs.size(); i < e; i++) {
      Value sz = loopHighs[i] = synSetter(builder, loc, i);
      auto [stl, it] = makeSynLevelAndIterator(sz, synId, i, emitStrategy);
      lvls[synId][i] = std::move(stl);
      iters[synId][i].emplace_back(std::move(it));
    }
  }

  // For every manifest tensor:
  // * For every level:
  //   * get the positions and coordinates buffers
  //   * get/compute the level-size, which is also used as the upper-bound
  //     on positions.
  for (TensorId t = 0, numTensors = getNumManifestTensors(); t < numTensors;
       t++) {
    // TODO: this should be done through a folding pass after switching to
    // `sparse_tensor.iterate`-based sparsification.
    const Value tensor = tryFoldTensors(tensors[t]);
    const auto rtp = dyn_cast<RankedTensorType>(tensor.getType());
    if (!rtp)
      // Skips only scalar, zero ranked tensor still need to be bufferized and
      // (probably) filled with zeros by users.
      continue;

    auto stt = getSparseTensorType(tensor);
    const Level lvlRank = stt.getLvlRank();

    // Scan all levels of current tensor.
    for (Level l = 0; l < lvlRank; l++) {
      // Find upper bound in current dimension.
      lvls[t][l] = makeSparseTensorLevel(builder, loc, tensor, t, l);
      if (!dependentLvlMap[t][l].empty())
        continue;

      auto it = makeLevelIterator(builder, loc, t, l);
      iters[t][l].emplace_back(std::move(it));
    }
    // NOTE: we can also prepare for 0 lvl here in advance, this will hoist
    // some loop preparation from tensor iteration, but will also (undesirably)
    // hoist the code ouside if-conditions.
  }
  // TODO: avoid treating subsection iterator as a special case.
  initSubSectIterator(builder, loc);
}

void LoopEmitter::initSubSectIterator(OpBuilder &builder, Location loc) {
  Value c0 = C_IDX(0);
  for (TensorId t = 0, e = tensors.size(); t < e; t++) {
    auto rtp = dyn_cast<RankedTensorType>(tensors[t].getType());
    if (!rtp)
      continue;

    Level lvlRank = SparseTensorType(rtp).getLvlRank();

    // Compute the dependency reduction order.
    auto remDepStack = dependentLvlMap;
    std::vector<std::tuple<LoopId, TensorId, Level>> depRedOrder;
    for (Level lvl = 0; lvl < lvlRank; lvl++) {
      // Reverse queue into a stack.
      std::reverse(remDepStack[t][lvl].begin(), remDepStack[t][lvl].end());
      for (auto [loop, coeff] : dependentLvlMap[t][lvl])
        depRedOrder.emplace_back(loop, t, lvl);
    }

    if (depRedOrder.empty())
      continue;

    llvm::sort(depRedOrder, llvm::less_first());

    SmallVector<SparseIterator *> lastIter(tensors.size(), nullptr);
    for (auto [loop, t, lvl] : depRedOrder) {
      std::pair<LoopId, unsigned> curDep = remDepStack[t][lvl].back();
      assert(curDep.first == loop);
      remDepStack[t][lvl].pop_back();

      auto lvlIt = makeLevelIterator(builder, loc, t, lvl);
      const SparseIterator *parent = lastIter[t];
      if (!parent && lvl > 0) {
        if (dependentLvlMap[t][lvl - 1].empty()) {
          parent = iters[t][lvl - 1].back().get();
        }
      }

      std::unique_ptr<SparseIterator> it;
      if (!remDepStack[t][lvl].empty()) {
        // Compute the subsection size.
        Value size = c0;
        for (auto [loop, stride] : remDepStack[t][lvl]) {
          Value idxMax = SUBI(loopHighs[loop], C_IDX(1));
          size = ADDI(size, ADDI(MULI(idxMax, C_IDX(stride)), C_IDX(1)));
        }
        it = makeNonEmptySubSectIterator(builder, loc, parent, loopHighs[loop],
                                         std::move(lvlIt), size, curDep.second,
                                         emitStrategy);
      } else {
        const SparseIterator &subSectIter = *iters[t][lvl].back();
        it = makeTraverseSubSectIterator(builder, loc, subSectIter, *parent,
                                         std::move(lvlIt), loopHighs[loop],
                                         curDep.second, emitStrategy);
      }
      lastIter[t] = it.get();
      iters[t][lvl].emplace_back(std::move(it));
    }
  }
}

void LoopEmitter::categorizeIterators(
    ArrayRef<TensorLevel> tidLvls, SmallVectorImpl<SparseIterator *> &raIters,
    SmallVectorImpl<SparseIterator *> &spIters) {
  // Finds out the tensor level that we should use to generate loops. Amongs all
  // the tensor levels, there is at most one sparse tensor level.
  for (auto [t, l] : unpackTensorLevelRange(tidLvls)) {
    SparseIterator *it = &getCurIterator(t, l);
    if (it->randomAccessible())
      raIters.push_back(it);
    else
      spIters.push_back(it);
  }

  llvm::stable_sort(spIters, [](auto lhs, auto rhs) {
    // AffineUnRed > Affine > Slice > Trivial
    return static_cast<uint8_t>(lhs->kind) > static_cast<uint8_t>(rhs->kind);
  });
}

void LoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
                                  ArrayRef<TensorLevel> tidLvls) {
  // TODO: sort
  assert(loopSeqStack.size() == loopStack.size());

  if (emitStrategy != SparseEmitStrategy::kSparseIterator) {
    // Prepares for all the tensors used in the current loop sequence.
    for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
      levelReducedDep[tid][lvl]++;
      prepareLoopOverTensorAtLvl(builder, loc, tid, lvl);
    }
  }

  // Universal Index starts from 0.
  loopSeqStack.emplace_back(C_IDX(0), tidLvls.vec());
}

void LoopEmitter::exitCurrentLoopSeq(OpBuilder &builder, Location loc) {
  assert(loopSeqStack.size() == loopStack.size() + 1);

  // Depending on whether the slice is resolved or not at current loop sequence,
  // end them in different ways.
  for (auto [tid, lvl] : unpackTensorLevelRange(loopSeqStack.back().second))
    levelReducedDep[tid][lvl]--;

  loopSeqStack.pop_back();
}

Value LoopEmitter::genAffine(OpBuilder &builder, Location loc, AffineExpr a) {
  switch (a.getKind()) {
  case AffineExprKind::DimId: {
    // FIXME: since the one callsite in Sparsification passes in a
    // level-expression, the `getPosition` must in fact be a `Dimension`.
    // However, elsewhere we have been lead to expect that `loopIdToOrd`
    // should be indexed by `LoopId`...
    const auto loopId = cast<AffineDimExpr>(a).getPosition();
    return loopStack[loopId].iv;
  }
  case AffineExprKind::Add: {
    auto binOp = cast<AffineBinaryOpExpr>(a);
    return ADDI(genAffine(builder, loc, binOp.getLHS()),
                genAffine(builder, loc, binOp.getRHS()));
  }
  case AffineExprKind::Mul: {
    auto binOp = cast<AffineBinaryOpExpr>(a);
    return MULI(genAffine(builder, loc, binOp.getLHS()),
                genAffine(builder, loc, binOp.getRHS()));
  }
  case AffineExprKind::Constant: {
    int64_t c = cast<AffineConstantExpr>(a).getValue();
    return C_IDX(c);
  }
  default:
    llvm_unreachable("unexpected affine subscript");
  }
}

std::pair<Operation *, Value> LoopEmitter::emitForLoopOverTensorAtLvl(
    OpBuilder &builder, Location loc, SparseIterator &iter,
    MutableArrayRef<Value> reduc, bool isParallel) {

  // TODO: support dynamic slices.
  // Uses the first dimension here to build the loop bound (which is also the
  // biggest range).

  Value step = C_IDX(1);
  auto [lo, hi] = iter.genForCond(builder, loc);
  Operation *loop = nullptr;
  Value iv;
  if (isParallel) {
    scf::ParallelOp parOp =
        scf::ParallelOp::create(builder, loc, lo, hi, step, reduc);
    builder.setInsertionPointToStart(parOp.getBody());
    assert(parOp.getNumReductions() == reduc.size());
    iv = parOp.getInductionVars()[0];

    // In-place update on the reduction variable vector.
    // Note that the init vals is not the actual reduction variables but instead
    // used as a "special handle" to (temporarily) represent them. The
    // expression on init vals will be moved into scf.reduce and replaced with
    // the block arguments when exiting the loop (see exitForLoop). This is
    // needed as we can not build the actual reduction block and get the actual
    // reduction variable before users fill parallel loop body.
    for (int i = 0, e = reduc.size(); i < e; i++)
      reduc[i] = parOp.getInitVals()[i];
    loop = parOp;
  } else {
    scf::ForOp forOp = scf::ForOp::create(builder, loc, lo, hi, step, reduc);
    builder.setInsertionPointToStart(forOp.getBody());
    iv = forOp.getInductionVar();

    // In-place update on the reduction variable vector.
    assert(forOp.getNumRegionIterArgs() == reduc.size());
    for (int i = 0, e = reduc.size(); i < e; i++)
      reduc[i] = forOp.getRegionIterArg(i);
    loop = forOp;
  }
  assert(loop && iv);

  Value crd = iv;
  if (!iter.randomAccessible()) {
    iter.linkNewScope(iv);
    crd = iter.deref(builder, loc);
  } else {
    iter.locate(builder, loc, iv);
  }

  return {loop, crd};
}

std::pair<Operation *, Value> LoopEmitter::emitWhileLoopOverTensorsAtLvls(
    OpBuilder &builder, Location loc, ArrayRef<SparseIterator *> spIters,
    MutableArrayRef<Value> reduc, bool needsUniv) {
  return genCoIteration(builder, loc, spIters, reduc,
                        needsUniv ? loopSeqStack.back().first : nullptr);
}

bool LoopEmitter::shouldIteratedByForLoop(ArrayRef<SparseIterator *> spIters) {
  // If we need to co-iterate over two sparse tensors, we need a while loop
  if (spIters.size() > 1)
    return false;

  if (spIters.size() == 1)
    return spIters.front()->iteratableByFor();

  return true;
}

Region *LoopEmitter::enterCurrentCoIterationCase(OpBuilder &builder,
                                                 Location loc,
                                                 I64BitSet caseBit,
                                                 unsigned caseIdx,
                                                 MutableArrayRef<Value> reduc) {
  auto coIterOp = cast<CoIterateOp>(loopStack.back().loop);
  SmallVector<Attribute> cases(coIterOp.getCases().getAsRange<Attribute>());
  cases[caseIdx] = builder.getI64IntegerAttr(caseBit);

  coIterOp.setCasesAttr(builder.getArrayAttr(cases));
  Region &caseRegion = coIterOp.getRegion(caseIdx);
  assert(caseRegion.getBlocks().empty() &&
         "re-initialize the same coiteration case region.");

  // Each block starts with by a list of user-provided iteration arguments.
  TypeRange iterArgsTps = coIterOp.getInitArgs().getTypes();
  // Followed by a list of used coordinates of index type.
  SmallVector<Type> blockArgTps(coIterOp.getCrdUsedLvls().count(),
                                builder.getIndexType());

  blockArgTps.append(iterArgsTps.begin(), iterArgsTps.end());
  // Ends with a set of iterators that defines the actually iteration space.
  for (auto i : caseBit.bits()) {
    blockArgTps.push_back(
        cast<IterSpaceType>(coIterOp.getIterSpaces()[i].getType())
            .getIteratorType());
  }
  SmallVector<Location> locs(blockArgTps.size(), loc);
  caseRegion.emplaceBlock().addArguments(blockArgTps, locs);

  // Entering the new region scope, updating the SSA chain.
  builder.setInsertionPointToStart(&caseRegion.front());
  // Update the coordinates.
  loopStack.back().iv = coIterOp.getCrds(caseIdx).front();
  // Updates loop iteration arguments.
  ValueRange iterArgs = coIterOp.getRegionIterArgs(caseIdx);
  llvm::copy(iterArgs, reduc.begin());
  // Updates sparse iterator values.
  ValueRange iters = coIterOp.getRegionIterators(caseIdx);
  ArrayRef<TensorLevel> tidLvls = loopStack.back().tidLvls;
  for (auto [i, tl] : llvm::enumerate(unpackTensorLevelRange(tidLvls))) {
    if (caseBit[i]) {
      spIterVals[tl.first][tl.second] = iters.front();
      iters = iters.drop_front();
    } else {
      spIterVals[tl.first][tl.second] = nullptr;
    }
  }
  // Must have consumed all iterator SSA values.
  assert(iters.empty());
  return &caseRegion;
}

Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
    OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls,
    unsigned numCases, MutableArrayRef<Value> reduc, bool tryParallel,
    bool needsUniv) {
  // TODO: Argument `numCases` only used when generating iterator-based sparse
  // loops. Simplify the code upon feature complete.
  // TODO: handle coiteration with sparse iterator.
  if (emitStrategy == SparseEmitStrategy::kSparseIterator) {
    if (tidLvls.size() == 1) {
      auto [tid, lvl] = unpackTensorLevel(tidLvls.front());
      Value t = tensors[tid];

      // Extract and iterate over the iteration space.
      ExtractIterSpaceOp extractSpaceOp =
          lvl == 0 ? ExtractIterSpaceOp::create(builder, loc, t)
                   : ExtractIterSpaceOp::create(builder, loc, t,
                                                spIterVals[tid][lvl - 1], lvl);

      IterateOp iterOp = IterateOp::create(
          builder, loc, extractSpaceOp.getExtractedSpace(), reduc);
      spIterVals[tid][lvl] = iterOp.getIterator();

      // Update the reduction varaibles.
      llvm::copy(iterOp.getRegionIterArgs(), reduc.begin());
      // Set the insertion point to loop body.
      builder.setInsertionPointToStart(iterOp.getBody());
      loopStack.emplace_back(tidLvls, iterOp, builder.getInsertionBlock(),
                             iterOp.getCrds().front(), loopTag);
      return iterOp;
    }

    // CoIteration Loops.
    SmallVector<Value> spaces;
    for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
      Value t = tensors[tid];
      ExtractIterSpaceOp extractSpaceOp =
          lvl == 0 ? ExtractIterSpaceOp::create(builder, loc, t)
                   : ExtractIterSpaceOp::create(builder, loc, t,
                                                spIterVals[tid][lvl - 1], lvl);
      spaces.push_back(extractSpaceOp.getExtractedSpace());
    }
    auto coIterOp = CoIterateOp::create(builder, loc, spaces, reduc, numCases);
    // The CoIterationOp does not have insertion block nor induction variable.
    // TODO: the `struct LoopInfo` should be simplied after full migration.
    loopStack.emplace_back(tidLvls, coIterOp, /*insertion block*/ nullptr,
                           /*induction variable*/ nullptr, loopTag);
    return coIterOp;
  }

  // TODO: support multiple return on parallel for?
  tryParallel = tryParallel && reduc.size() <= 1;

  SmallVector<SparseIterator *> raIters;
  SmallVector<SparseIterator *> spIters;
  categorizeIterators(tidLvls, raIters, spIters);

  // Only when there is at least one sparse conditions, do we really need the
  // universal index.
  // TODO: Maybe we should instead requires merger to pass in a valid value at
  // the first place instead of adjusting it in LoopEmitter?
  needsUniv = !spIters.empty() && needsUniv;
  // The TensorLevel used for loop conditions.
  // If there is any sparse level, we need to use the sparse condition.
  // If all levels are dense, we can pick arbitrary one (dense slice-driven loop
  // can be generated using a simple ForOp as well).
  Operation *l = nullptr;
  Value iv = nullptr;
  SmallVector<TensorLevel> tls;

  // Generates loops differently depending on whether we need a slice-driven
  // loop or a simple level traversal loop.
  if (shouldIteratedByForLoop(spIters) && !needsUniv) {
    assert(spIters.size() <= 1);
    SparseIterator &it = spIters.empty() ? *raIters.front() : *spIters.front();
    std::tie(l, iv) =
        emitForLoopOverTensorAtLvl(builder, loc, it, reduc, tryParallel);
    tls.push_back(makeTensorLevel(it.tid, it.lvl));
  } else {
    for (auto *it : spIters) {
      tls.push_back(makeTensorLevel(it->tid, it->lvl));
    }

    if (needsUniv)
      for (auto *it : raIters)
        tls.push_back(makeTensorLevel(it->tid, it->lvl));

    std::tie(l, iv) =
        emitWhileLoopOverTensorsAtLvls(builder, loc, spIters, reduc, needsUniv);
  }

  // Enter dense tensor levels.
  for (SparseIterator *it : raIters)
    it->locate(builder, loc, iv);

  // NOTE: we can also prepare for next dim here in advance
  // Pushes the loop into stack.
  loopStack.emplace_back(tls, l, builder.getInsertionBlock(), iv, loopTag);
  return l;
}

void LoopEmitter::locateLvlAtAffineAddress(OpBuilder &builder, Location loc,
                                           TensorLevel tidLvl,
                                           AffineExpr lvlExpr) {
  auto [tid, lvl] = unpackTensorLevel(tidLvl);

  const SparseIterator *parent =
      lvl == 0 ? nullptr : iters[tid][lvl - 1].back().get();
  auto &it = getCurIterator(tid, lvl);
  it.genInit(builder, loc, parent);

  assert(it.kind == IterKind::kTrivial && it.randomAccessible());
  Value lvlCrd = genAffine(builder, loc, lvlExpr);
  it.locate(builder, loc, lvlCrd);
}

void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                             TensorId tid, Level lvl) {
  // if this is the first level, there is no parent iterator for the current
  // iterator.
  // If the current iterator is a subsection-based iterator, the parent iterator
  // is memorized by the iterator.
  bool hasParent = lvl == 0 || !dependentLvlMap[tid][lvl].empty();

  const SparseIterator *parent =
      hasParent ? nullptr : iters[tid][lvl - 1].back().get();
  auto &it = getCurIterator(tid, lvl);
  it.genInit(builder, loc, parent);

  // Locates the randon accessible iterator to 0.
  if (it.randomAccessible())
    it.locate(builder, loc, C_IDX(0));
}

void LoopEmitter::exitForLoop(RewriterBase &rewriter, Location loc,
                              MutableArrayRef<Value> reduc) {
  const LoopInfo &loopInfo = loopStack.back();
  if (emitStrategy == SparseEmitStrategy::kSparseIterator) {
    auto iterateOp = llvm::cast<IterateOp>(loopInfo.loop);
    assert(reduc.size() == iterateOp.getNumResults());
    sparse_tensor::YieldOp::create(rewriter, loc, reduc);
    // Exit the loop.
    rewriter.setInsertionPointAfter(iterateOp);
    // In-place update reduction variables.
    llvm::copy(iterateOp.getResults(), reduc.begin());
    return;
  }
  if (auto forOp = llvm::dyn_cast<scf::ForOp>(loopInfo.loop)) {
    if (!reduc.empty()) {
      assert(reduc.size() == forOp.getNumResults());
      scf::YieldOp::create(rewriter, loc, reduc);
    }
    // Exit the loop.
    rewriter.setInsertionPointAfter(forOp);
    // In-place update reduction variables.
    llvm::copy(forOp.getResults(), reduc.begin());
  } else {
    auto parOp = llvm::cast<scf::ParallelOp>(loopInfo.loop);
    if (!reduc.empty()) {
      assert(reduc.size() == parOp.getInitVals().size() && reduc.size() == 1);
      Operation *redExp = reduc.front().getDefiningOp();
      // Reduction expression should have no use.
      assert(redExp->getUses().empty());
      // This must be a binary operation.
      // NOTE: This is users' responsibility to ensure the operation are
      // commutative.
      assert(redExp->getNumOperands() == 2 && redExp->getNumResults() == 1);

      Value redVal = parOp.getInitVals().front();
      Value curVal;
      if (redExp->getOperand(0) == redVal)
        curVal = redExp->getOperand(1);
      else if (redExp->getOperand(1) == redVal)
        curVal = redExp->getOperand(0);
      // One of the operands must be the init value (which is also the
      // previous reduction value).
      assert(curVal);
#ifndef NDEBUG
      // The reduction expression should be the only user of the reduction val
      // inside the parallel for.
      unsigned numUsers = 0;
      for (Operation *op : redVal.getUsers()) {
        if (op->getParentOp() == parOp)
          numUsers++;
      }
      assert(numUsers == 1);
#endif // NDEBUG

      rewriter.setInsertionPointAfter(redExp);
      auto redOp = scf::ReduceOp::create(rewriter, loc, curVal);
      // Attach to the reduction op.
      Block *redBlock = &redOp.getReductions().front().front();
      rewriter.setInsertionPointToEnd(redBlock);
      Operation *newRed = rewriter.clone(*redExp);
      // Replaces arguments of the reduction expression by using the block
      // arguments from scf.reduce.
      rewriter.modifyOpInPlace(
          newRed, [&]() { newRed->setOperands(redBlock->getArguments()); });
      // Erases the out-dated reduction expression.
      rewriter.eraseOp(redExp);
      rewriter.setInsertionPointToEnd(redBlock);
      scf::ReduceReturnOp::create(rewriter, loc, newRed->getResult(0));
    }
    rewriter.setInsertionPointAfter(parOp);
    // In-place update reduction variables.
    for (unsigned i = 0, e = parOp.getResults().size(); i < e; i++)
      reduc[i] = parOp.getResult(i);
  }
}

void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
                                MutableArrayRef<Value> reduc) {
  const LoopInfo &loopInfo = loopStack.back();
  auto whileOp = llvm::cast<scf::WhileOp>(loopInfo.loop);
  Value iv = loopInfo.iv;
  Value one = C_IDX(1);

  // Finalize the induction. Note that the induction could be performed
  // in the individual if-branches to avoid re-evaluating the conditions.
  // However, that would result in a rather elaborate forest of yield
  // instructions during code generation. Moreover, performing the induction
  // after the if-statements more closely resembles code generated by TACO.
  SmallVector<Value> operands;
  ValueRange whileRes = whileOp.getResults();

  for (auto [tid, lvl] : unpackTensorLevelRange(loopInfo.tidLvls)) {
    SparseIterator &it = getCurIterator(tid, lvl);
    if (!it.randomAccessible()) {
      // Forward the sparse iterator.
      Value cmp = CMPI(eq, it.getCrd(), iv);
      it.forwardIf(builder, loc, cmp);
      operands.append(it.getCursor().begin(), it.getCursor().end());
      // const Value newPos = whileOp->getResult(o++);
      // Following loops continue iteration from the break point of the
      // current while loop.
      whileRes = it.linkNewScope(whileRes);
    }
    // Note: random-accessible (dense) iterators are located after the while
    // loop exits; see below.
  }

  // Reduction value from users.
  for (auto &i : reduc) {
    operands.push_back(i);
    // Update user reduction variables.
    i = whileRes.front();
    whileRes = whileRes.drop_front();
  }

  // An (optional) universal index.
  if (operands.size() < whileOp.getNumResults()) {
    assert(operands.size() + 1 == whileOp.getNumResults());
    // The last one is the universial index.
    operands.push_back(ADDI(iv, one));
    // update the loop starting point of current loop sequence
    loopSeqStack.back().first = whileOp->getResults().back();
  }

  if (!operands.empty())
    YIELD(operands);

  builder.setInsertionPointAfter(whileOp);

  // Locate random-accessible (dense) iterators at the final universal index.
  // This must happen after the while loop exits so that the while results are
  // in scope. Doing it inside the after block (using while results there) would
  // cause domination violations when the iterator is later dereferenced.
  for (auto [tid, lvl] : unpackTensorLevelRange(loopInfo.tidLvls)) {
    SparseIterator &it = getCurIterator(tid, lvl);
    if (it.randomAccessible()) {
      Value uniIdx = whileOp.getResults().back();
      it.locate(builder, loc, uniIdx);
    }
  }
}

void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                  MutableArrayRef<Value> reduc) {
  // Clean up the values, it would help use to discover potential bug at a
  // earlier stage (instead of silently using a wrong value).
  const LoopInfo &loopInfo = loopStack.back();
  if (emitStrategy == SparseEmitStrategy::kSparseIterator) {
    Operation *p = loopInfo.loop;
    if (isa<IterateOp>(p))
      sparse_tensor::YieldOp::create(rewriter, loc, reduc);

    // Exit the loop.
    rewriter.setInsertionPointAfter(p);
    // In-place update reduction variables.
    llvm::copy(p->getResults(), reduc.begin());
    loopStack.pop_back();
    return;
  }

  // Sets the insertion point to the right position.
  rewriter.setInsertionPointToEnd(loopInfo.userCodeBlock);
  if (!loopInfo.userCodeBlock->empty() &&
      llvm::isa<scf::YieldOp>(&loopInfo.userCodeBlock->back())) {
    // scf::While/For inserts an implicit yield op when there is no loop
    // iter args. In this case, we need to insert the code before the yield.
    assert(loopInfo.userCodeBlock->back().getNumResults() == 0);
    rewriter.setInsertionPoint(&loopInfo.userCodeBlock->back());
  }

  if (llvm::isa<scf::WhileOp>(loopInfo.loop)) {
    exitWhileLoop(rewriter, loc, reduc);
  } else {
    exitForLoop(rewriter, loc, reduc);
  }

  assert(loopStack.size() == loopSeqStack.size());
  loopStack.pop_back();
}

//===----------------------------------------------------------------------===//
// Loop generation utils
//===----------------------------------------------------------------------===//

std::pair<Operation *, Value> sparse_tensor::genCoIteration(
    OpBuilder &builder, Location loc, ArrayRef<SparseIterator *> spIters,
    MutableArrayRef<Value> reduc, Value uniIdx, bool userReducFirst) {
  // NOTE: the slice driven tensor-related reduction variable must
  // appear before normal tensors.

  // The set of induction variables for the while loop.
  SmallVector<Value> ivs;

  // TODO: remove the flag after full migration. Currently
  // `sparse_tensor.coiterate` operation (must) put user provided reduction
  // values at the front of the block list, while direct sparsification to scf
  // loops put them at the end.
  if (userReducFirst)
    ivs.append(reduc.begin(), reduc.end());

  // Construct the while-loop with a parameter for each coordinate.
  for (SparseIterator *it : spIters) {
    ValueRange itVals = it->getCursor();
    ivs.append(itVals.begin(), itVals.end());
  }

  if (!userReducFirst)
    ivs.append(reduc.begin(), reduc.end());

  // Update universal index.
  if (uniIdx)
    ivs.push_back(uniIdx);

  // Ensures all operands are valid.
  assert(!llvm::is_contained(ivs, nullptr));
  TypeRange types = ValueRange(ivs).getTypes();
  auto whileOp = scf::WhileOp::create(builder, loc, types, ivs);

  SmallVector<Location> locs(types.size(), loc);
  Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
  Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);

  // Generates loop conditions.
  builder.setInsertionPointToStart(before);
  ValueRange bArgs = before->getArguments();
  Value whileCond = nullptr; // bool values for loop condition.

  for (SparseIterator *it : spIters) {
    auto [cond, remArgs] = it->genWhileCond(builder, loc, bArgs);
    whileCond = !whileCond ? cond : ANDI(whileCond, cond);
    bArgs = remArgs;
  }
  // The remaining block arguments are user-provided reduction values and an
  // optional universal index. Make sure their sizes match.
  assert(bArgs.size() == reduc.size() + (uniIdx ? 1 : 0));
  scf::ConditionOp::create(builder, loc, whileCond, before->getArguments());

  // Generates loop body.
  builder.setInsertionPointToStart(after);
  ValueRange aArgs = after->getArguments();

  for (SparseIterator *it : spIters) {
    aArgs = it->linkNewScope(aArgs);
    // Dereference the iterator to cache the coordinate.
    it->deref(builder, loc);
  }

  // In-place update on reduction variable.
  for (unsigned i = 0, e = reduc.size(); i < e; i++)
    reduc[i] = aArgs[i];

  Value min;
  // Finds the minimum coordinate
  if (!uniIdx) {
    for (SparseIterator *it : spIters) {
      if (min) {
        Value cmp = CMPI(ult, it->getCrd(), min);
        min = SELECT(cmp, it->getCrd(), min);
      } else {
        min = it->getCrd();
      }
    }
  } else {
    // Otherwise, universal index is the minimal pos.
    min = whileOp.getAfterArguments().back();
  }

  return {whileOp, min};
}

#undef CMPI
#undef C_IDX
#undef YIELD
#undef ADDI
#undef ANDI
#undef SUBI
#undef MULI
#undef REMUI
#undef DIVUI
#undef SELECT