//===- Sparsification.cpp - Implementation of sparsification --------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements converting sparse tensor types to actual sparse code.
//
//===----------------------------------------------------------------------===//

#include "CodegenUtils.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.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/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Dialect/SparseTensor/Utils/Merger.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TensorEncoding.h"
#include "llvm/ADT/SmallBitVector.h"

using namespace mlir;
using namespace mlir::sparse_tensor;

//===----------------------------------------------------------------------===//
// Declarations of data structures.
//===----------------------------------------------------------------------===//

namespace {

// Iteration graph sorting.
enum SortMask {
  kSparseOnly = 0x0,
  kIncludeDense = 0x1,
  kIncludeUndef = 0x2,
  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,
          std::vector<unsigned> &ts)
      : options(o), loopEmitter(tensors, /*hasOutput=*/true,
                                /*isSparseOut=*/op != nullptr),
        sparseOut(op), outerParNest(nest), topSort(ts) {
    if (op)
      insChain = op->get();
  }
  /// Sparsification options.
  SparsificationOptions options;
  /// Loop emitter helper class.
  SparseTensorLoopEmitter loopEmitter;
  /// Current reduction, updated during code generation. When indices of a
  /// reduction are exhausted, all inner loops can use a scalarized reduction.
  unsigned redExp = -1u;
  Value redVal;
  Reduction redKind = kNoReduc;
  unsigned redCustom = -1u;
  // Sparse tensor as output. Implemented either through direct injective
  // insertion in lexicographic index order or through access pattern expansion
  // in the innermost loop nest (`expValues` through `expCount`).
  OpOperand *sparseOut;
  unsigned outerParNest;
  Value insChain; // bookkeeping for insertion chain
  Value expValues;
  Value expFilled;
  Value expAdded;
  Value expCount;
  // Topsort (reference should remain in scope).
  std::vector<unsigned> &topSort;

  Value getLoopIdxValue(size_t loopIdx) const {
    for (unsigned lv = 0; lv < topSort.size(); lv++)
      if (topSort[lv] == loopIdx)
        return loopEmitter.getLoopIV(lv);

    llvm_unreachable("invalid loop index");
  }
};

} // namespace

//===----------------------------------------------------------------------===//
// Sparse compiler analysis methods.
//===----------------------------------------------------------------------===//

/// Helper method to construct a permuted dimension ordering
/// that adheres to the given topological sort.
static AffineMap permute(MLIRContext *context, AffineMap m,
                         std::vector<unsigned> &topSort) {
  unsigned sz = topSort.size();
  assert(m.getNumResults() == sz && "TopoSort/AffineMap size mismatch");
  // Construct the inverse of `m`; to avoid the asymptotic complexity
  // of calling `m.getPermutedPosition` repeatedly.
  SmallVector<unsigned, 4> inv(sz);
  for (unsigned i = 0; i < sz; i++)
    inv[i] = m.getDimPosition(i);
  // Construct the permutation.
  SmallVector<unsigned, 4> perm(sz);
  for (unsigned i = 0; i < sz; i++)
    perm[i] = inv[topSort[i]];
  return AffineMap::getPermutationMap(perm, context);
}

/// Helper method to inspect affine expressions. Rejects cases where the
/// same index is used more than once. Also rejects compound affine
/// expressions in sparse dimensions.
static bool findAffine(Merger &merger, unsigned tensor, unsigned dim,
                       AffineExpr a, DimLevelType dlt,
                       bool setLvlFormat = true) {
  switch (a.getKind()) {
  case AffineExprKind::DimId: {
    unsigned idx = a.cast<AffineDimExpr>().getPosition();
    if (!isUndefDLT(merger.getDimLevelType(tensor, idx)))
      return false; // used more than once

    if (setLvlFormat)
      merger.setDimAndDimLevelType(tensor, idx, dim, dlt);
    return true;
  }
  case AffineExprKind::Add:
  case AffineExprKind::Mul: {
    if (!isDenseDLT(dlt))
      return false; // compound only in dense dim
    auto binOp = a.cast<AffineBinaryOpExpr>();
    // We do not set dim level format for affine expresssion like d0 + d1 on
    // both loop index at d0 and d1,
    return findAffine(merger, tensor, dim, binOp.getLHS(), dlt, false) &&
           findAffine(merger, tensor, dim, binOp.getRHS(), dlt, false);
  }
  case AffineExprKind::Constant:
    return isDenseDLT(dlt); // const only in dense dim
  default:
    return false;
  }
}

/// Helper method to inspect sparse encodings in the tensor types.
/// Fills the per-dimension sparsity information for all tensors.
/// Returns true if the sparse annotations and affine subscript
/// expressions of all tensors are admissible. Returns false if
/// no annotations are found or inadmissible constructs occur.
static bool findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
  bool annotated = false;
  for (OpOperand &t : op->getOpOperands()) {
    auto map = op.getMatchingIndexingMap(&t);
    auto enc = getSparseTensorEncoding(t.get().getType());
    if (enc)
      annotated = true;
    assert(map.getNumResults() == op.getRank(&t));
    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
      unsigned tensor = t.getOperandNumber();
      AffineExpr a = map.getResult(toOrigDim(enc, d));
      if (!findAffine(merger, tensor, d, a, getDimLevelType(enc, d)))
        return false; // inadmissible affine expression
    }
  }
  return annotated;
}

/// A helper to compute a topological sort. O(n^2) time complexity
/// as we use adj matrix for the graph.
/// The sorted result will put the first Reduction iterator to the
/// latest possible index.
static bool topSortOptimal(unsigned n, ArrayRef<StringRef> iteratorTypes,
                           std::vector<unsigned> &topSort,
                           std::vector<unsigned> &inDegree,
                           std::vector<std::vector<bool>> &adjM) {
  std::vector<unsigned> redIt; // reduce iterator with 0 degree
  std::vector<unsigned> parIt; // parallel iterator with 0 degree
  for (unsigned i = 0; i < n; i++) {
    if (inDegree[i] == 0) {
      if (linalg::isReductionIterator(iteratorTypes[i]))
        redIt.push_back(i);
      else
        parIt.push_back(i);
    }
  }

  while (!redIt.empty() || !parIt.empty()) {
    // We always choose parallel iterator if there is any.
    auto &it = !parIt.empty() ? parIt : redIt;
    auto src = it.back();
    topSort.push_back(src);
    it.pop_back();
    // Update in-degree, and push 0-degree node into worklist.
    for (unsigned dst = 0; dst < n; dst++)
      if (adjM[src][dst] && --inDegree[dst] == 0) {
        if (linalg::isReductionIterator(iteratorTypes[dst]))
          redIt.push_back(dst);
        else
          parIt.push_back(dst);
      }
  }
  return topSort.size() == n;
}

/// Helper method to add all constraints from the indices in one affine
/// expression before all indices in the other affine expression. For
/// example i0+i1 < i2+i3+1 yields i0<i2, i0<i3, i1<i2, and i1<i3.
static void addAffineOrderings(std::vector<std::vector<bool>> &adjM,
                               std::vector<unsigned> &inDegree, AffineExpr a,
                               AffineExpr b, unsigned fidx) {
  switch (a.getKind()) {
  case AffineExprKind::DimId: {
    unsigned idx = a.cast<AffineDimExpr>().getPosition();
    if (b)
      addAffineOrderings(adjM, inDegree, b, AffineExpr(), idx);
    else if (!adjM[fidx][idx]) {
      adjM[fidx][idx] = true;
      inDegree[idx]++;
    }
    break;
  }
  case AffineExprKind::Add:
  case AffineExprKind::Mul: {
    auto binOp = a.cast<AffineBinaryOpExpr>();
    addAffineOrderings(adjM, inDegree, binOp.getLHS(), b, fidx);
    addAffineOrderings(adjM, inDegree, binOp.getRHS(), b, fidx);
    break;
  }
  default:
    break;
  }
}

/// Computes a topologically sorted iteration graph for the linalg operation.
/// Ensures all tensors are visited in natural index order. This is essential
/// for sparse storage formats since these only support access along fixed
/// dimensions. Even for dense storage formats, however, the natural index
/// order yields innermost unit-stride access with better spatial locality.
static bool computeIterationGraph(Merger &merger, linalg::GenericOp op,
                                  std::vector<unsigned> &topSort, unsigned mask,
                                  OpOperand *skip = nullptr) {
  // Set up an n x n from/to adjacency matrix of the iteration graph
  // for the implicit loop indices i_0 .. i_n-1.
  unsigned n = op.getNumLoops();
  std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
  std::vector<unsigned> inDegree(n, 0); // in-degree of each node.
  auto iteratorTypes = op.getIteratorTypesArray();
  // Iterate over the indexing maps of every tensor in the tensor expression.
  for (OpOperand &t : op->getOpOperands()) {
    // Skip tensor during cycle resolution.
    if (&t == skip)
      continue;
    // Get map and encoding.
    auto map = op.getMatchingIndexingMap(&t);
    auto enc = getSparseTensorEncoding(t.get().getType());
    assert(map.getNumDims() == n);
    // Skip dense tensor constraints when not requested.
    if (!(mask & SortMask::kIncludeDense) && !enc)
      continue;
    // Each tensor expression and optional dimension ordering (row-major
    // by default) puts an ordering constraint on the loop indices. For
    // example, the tensor expresion A_ijk forces the ordering i < j < k
    // on the loop indices if no explicit dimension ordering is given.
    for (unsigned d = 1, rank = map.getNumResults(); d < rank; d++) {
      AffineExpr f = map.getResult(toOrigDim(enc, d - 1));
      AffineExpr t = map.getResult(toOrigDim(enc, d));
      addAffineOrderings(adjM, inDegree, f, t, 0);
    }
    // Push unrelated loops into sparse iteration space, so these
    // will be skipped more often.
    if (mask & SortMask::kIncludeUndef) {
      unsigned tensor = t.getOperandNumber();
      for (unsigned i = 0; i < n; i++)
        if (isCompressedDLT(merger.getDimLevelType(tensor, i)) ||
            isSingletonDLT(merger.getDimLevelType(tensor, i))) {
          for (unsigned j = 0; j < n; j++)
            if (isUndefDLT(merger.getDimLevelType(tensor, j))) {
              adjM[i][j] = true;
              inDegree[j]++;
            }
        } else {
          assert(isDenseDLT(merger.getDimLevelType(tensor, i)) ||
                 isUndefDLT(merger.getDimLevelType(tensor, i)));
        }
    }
  }
  // Topologically sort the iteration graph to determine loop order.
  // Report failure for a cyclic iteration graph.
  topSort.clear();
  topSort.reserve(n);
  return topSortOptimal(n, iteratorTypes, topSort, inDegree, adjM);
}

/// Returns true if tensor materializes uninitialized into the computation.
static bool isMaterializing(Value val) {
  return val.getDefiningOp<tensor::EmptyOp>() ||
         val.getDefiningOp<bufferization::AllocTensorOp>();
}

/// Returns true when the tensor expression is admissible for codegen.
/// Since all sparse input tensors are admissible, we just need to check
/// whether the out tensor in the tensor expression codegen is admissible.
/// Sets `sparseOut` to the tensor and `outerParNest` to the outer injective
/// nesting depth when a "truly dynamic" sparse tensor output occurs.
static bool isAdmissibleTensorExp(Merger &merger, linalg::GenericOp op,
                                  std::vector<unsigned> &topSort, unsigned exp,
                                  OpOperand **sparseOut,
                                  unsigned &outerParNest) {
  OpOperand *lhs = op.getDpsInitOperand(0);
  unsigned tensor = lhs->getOperandNumber();
  auto enc = getSparseTensorEncoding(lhs->get().getType());
  // An non-annotated output tensor is assumed dense, and becomes a random
  // access n-dim memref. Admissible since insertions cannot occur.
  if (!enc)
    return true;
  // An all-dense annotated "sparse" output tensor becomes a linearized random
  // access 1-dim memref. Also admissible since insertions cannot occur.
  bool allDense = true;
  auto iteratorTypes = op.getIteratorTypesArray();
  unsigned numLoops = iteratorTypes.size();
  for (unsigned i = 0; i < numLoops; i++)
    if (isCompressedDLT(merger.getDimLevelType(tensor, i)) ||
        isSingletonDLT(merger.getDimLevelType(tensor, i))) {
      allDense = false;
      break;
    } else {
      assert(isDenseDLT(merger.getDimLevelType(tensor, i)) ||
             isUndefDLT(merger.getDimLevelType(tensor, i)));
    }
  if (allDense)
    return true;
  // A tensor expression with a sparse output tensor that changes its values
  // but not its nonzero structure, an operation called "simply dynamic" in
  // [Bik96,Ch9], is also admissible without special codegen.
  if (merger.isSingleCondition(tensor, exp))
    return true;
  // Accept "truly dynamic" if the output tensor materializes uninitialized
  // into the computation and insertions occur in lexicographic index order.
  if (isMaterializing(lhs->get())) {
    unsigned nest = 0;
    for (unsigned i = 0; i < numLoops; i++) {
      if (linalg::isReductionIterator(iteratorTypes[topSort[i]]))
        break; // terminate at first reduction
      nest++;
    }
    // Determine admissible dynamic insertion situations:
    // (1) fully injective, since there are no reductions,
    // (2) admissible 1-d expansion in innermost dimension.
    if (nest >= op.getRank(lhs) - 1) {
      *sparseOut = lhs;
      outerParNest = nest;
      return true;
    }
  }
  return false;
}

//===----------------------------------------------------------------------===//
// Sparse compiler synthesis methods (reductions).
//===----------------------------------------------------------------------===//

/// Maps operation to reduction.
static Reduction getReduction(Kind kind) {
  switch (kind) {
  case Kind::kAddF:
  case Kind::kAddC:
  case Kind::kAddI:
  case Kind::kSubF:
  case Kind::kSubC:
  case Kind::kSubI:
    return kSum;
  case Kind::kMulF:
  case Kind::kMulC:
  case Kind::kMulI:
    return kProduct;
  case Kind::kAndI:
    return kAnd;
  case Kind::kOrI:
    return kOr;
  case Kind::kXorI:
    return kXor;
  case Kind::kReduce:
    return kCustom;
  default:
    llvm_unreachable("unexpected reduction operator");
  }
}

/// Updates scalarized reduction value.
static void updateReduc(Merger &merger, CodeGen &codegen, Value reduc) {
  assert(codegen.redKind != kNoReduc);
  codegen.redVal = merger.exp(codegen.redExp).val = reduc;
}

/// Extracts identity from custom reduce.
static Value getCustomRedId(Operation *op) {
  return dyn_cast<sparse_tensor::ReduceOp>(op).getIdentity();
}

//===----------------------------------------------------------------------===//
// Sparse compiler synthesis methods (statements and expressions).
//===----------------------------------------------------------------------===//

/// Generates loop boundary statements (entering/exiting loops). The function
/// passes and updates the reduction value.
static Optional<Operation *> genLoopBoundary(
    CodeGen &codegen, Merger &merger,
    function_ref<Optional<Operation *>(MutableArrayRef<Value> reduc)>
        callback) {
  SmallVector<Value, 4> reduc;
  if (codegen.redVal)
    reduc.push_back(codegen.redVal);
  if (codegen.expValues)
    reduc.push_back(codegen.expCount);
  if (codegen.insChain)
    reduc.push_back(codegen.insChain);

  auto r = callback(reduc);

  // Callback should do in-place update on reduction value vector.
  unsigned i = 0;
  if (codegen.redVal)
    updateReduc(merger, codegen, reduc[i++]);
  if (codegen.expValues)
    codegen.expCount = reduc[i++];
  if (codegen.insChain)
    codegen.insChain = reduc[i];

  return r;
}

/// Local bufferization of all dense and sparse data structures.
static void genBuffers(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                       linalg::GenericOp op) {
  Location loc = op.getLoc();
  assert(op.getNumOperands() == op.getNumDpsInputs() + 1);

  codegen.loopEmitter.initializeLoopEmit(
      builder, loc,
      /// Generates buffer for the output tensor.
      /// Note that all sparse kernels assume that when all elements are written
      /// to (viz. x(i) = y(i) * z(i)), the output buffer is already initialized
      /// to all zeroes and only nonzeroes values are computed and written out.
      /// For updates (viz. x(i) += y(i) * z(i)), only nonzeroes values are used
      /// for the updates and no assumption on the original contents of the
      /// output buffer is necessary.
      [&op](OpBuilder &builder, Location loc, Value memref,
            Value tensor) -> Value {
        // Must not be a sparse tensor.
        assert(!getSparseTensorEncoding(tensor.getType()));
        OpOperand *lhs = op.getDpsInitOperand(0);
        // Two output tensors references should pointed to the same object.
        assert(lhs->get() == tensor);
        bool isInit = op.isInitTensor(lhs);
        // An output tensor can simply materialize from the buffer of the tensor
        // that appears in the outs() clause. For updates, this has the
        // advantage that only the nonzero value are involved in the
        // computation, keeping the operation O(nnz). In all other cases, we are
        // forced to zero out the buffer to enforce the assumption above, which
        // may negatively impact running complexity (viz. O(n^2 + nnz) vs.
        // O(nnz) for matrices).
        // TODO: use better analysis to avoid zeroing out the buffer?
        Value init = memref;
        if (!isInit) {
          Value zero = constantZero(builder, loc,
                                    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);
}

/// Generates subscript for load/store on a dense or sparse tensor.
static Value genSubscript(CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, OpOperand *t,
                          SmallVector<Value, 4> &args) {
  unsigned tensor = t->getOperandNumber();
  auto map = op.getMatchingIndexingMap(t);
  auto enc = getSparseTensorEncoding(t->get().getType());
  unsigned rank = map.getNumResults();
  if (enc) {
    // Note that currently, all sparse subscripts are simple.
    // TODO: accept affine too?
    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()));
    }
  }
  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) {
  Location loc = op.getLoc();
  // Direct lexicographic index order, tensor loads as zero.
  if (!codegen.expValues) {
    Type tp = getElementTypeOrSelf(t->get().getType());
    return constantZero(builder, loc, tp);
  }
  // Load from expanded access pattern.
  Value index = genIndex(codegen, op, t);
  return builder.create<memref::LoadOp>(loc, codegen.expValues, index);
}

/// Generates insertion code to implement dynamic tensor load for reduction.
static Value genInsertionLoadReduce(Merger &merger, CodeGen &codegen,
                                    OpBuilder &builder, linalg::GenericOp op,
                                    OpOperand *t) {
  Location loc = op.getLoc();
  Value identity = getCustomRedId(merger.exp(codegen.redCustom).op);
  // Direct lexicographic index order, tensor loads as identity.
  if (!codegen.expValues) {
    return identity;
  }
  // Load from expanded access pattern if filled, identity otherwise.
  Value index = genIndex(codegen, op, t);
  Value isFilled =
      builder.create<memref::LoadOp>(loc, codegen.expFilled, index);
  Value valAtIndex =
      builder.create<memref::LoadOp>(loc, codegen.expValues, index);
  return builder.create<arith::SelectOp>(loc, isFilled, valAtIndex, identity);
}

/// Generates insertion code to implement dynamic tensor store.
static void genInsertionStore(CodeGen &codegen, OpBuilder &builder,
                              linalg::GenericOp op, OpOperand *t, Value rhs) {
  Location loc = op.getLoc();
  // Direct insertion in lexicographic index order.
  if (!codegen.expValues) {
    unsigned rank = op.getRank(t);
    SmallVector<Value, 4> indices;
    for (unsigned i = 0; i < rank; i++) {
      assert(codegen.loopEmitter.getLoopIV(i));
      indices.push_back(codegen.loopEmitter.getLoopIV(i));
    }
    codegen.insChain =
        builder.create<InsertOp>(loc, rhs, codegen.insChain, indices);
    return;
  }
  // Generates insertion code along expanded access pattern.
  //   if (!expFilled[i]) then
  //     expFilled[i] = true
  //     expAdded[inserts++] = i
  //   endif
  //   values[i] = rhs
  Value index = genIndex(codegen, op, t);
  Value fval = constantI1(builder, loc, false);
  Value tval = constantI1(builder, loc, true);
  // If statement.
  Value filled = builder.create<memref::LoadOp>(loc, codegen.expFilled, index);
  Value cond = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                             filled, fval);
  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, builder.getIndexType(), cond,
                                             /*else=*/true);
  // True branch.
  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
  builder.create<memref::StoreOp>(loc, tval, codegen.expFilled, index);
  builder.create<memref::StoreOp>(loc, index, codegen.expAdded,
                                  codegen.expCount);
  Value one = constantIndex(builder, loc, 1);
  Value add = builder.create<arith::AddIOp>(loc, codegen.expCount, one);
  builder.create<scf::YieldOp>(loc, add);
  // False branch.
  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
  builder.create<scf::YieldOp>(loc, codegen.expCount);
  builder.setInsertionPointAfter(ifOp);
  // Value assignment.
  codegen.expCount = ifOp.getResult(0);
  builder.create<memref::StoreOp>(loc, rhs, codegen.expValues, index);
}

/// Generates a load on a dense or sparse tensor.
static Value genTensorLoad(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, unsigned exp) {
  // Test if the load was hoisted to a higher loop nest.
  Value val = merger.exp(exp).val;
  if (val)
    return val;

  // Load during insertion.
  OpOperand &t = op->getOpOperand(merger.exp(exp).tensor);
  if (&t == codegen.sparseOut) {
    if (codegen.redCustom != -1u)
      return genInsertionLoadReduce(merger, codegen, builder, op, &t);
    return genInsertionLoad(codegen, builder, op, &t);
  }
  // Actual load.
  SmallVector<Value, 4> args;
  Value ptr = genSubscript(codegen, builder, op, &t, args);
  return builder.create<memref::LoadOp>(op.getLoc(), ptr, args);
}

/// Generates a store on a dense or sparse tensor.
static void genTensorStore(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, unsigned exp, Value rhs) {
  Location loc = op.getLoc();
  // Test if this is a scalarized reduction.
  if (codegen.redVal) {
    updateReduc(merger, codegen, rhs);
    return;
  }
  // Store during insertion.
  OpOperand *t = op.getDpsInitOperand(0);
  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);
      SmallVector<Type, 1> types;
      types.push_back(codegen.insChain.getType());
      scf::IfOp ifOp =
          builder.create<scf::IfOp>(loc, types, 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);
      // Yield original insertion chain along false branch.
      builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
      builder.create<scf::YieldOp>(loc, insChain);
      // Done with if statement.
      codegen.insChain = ifOp->getResult(0);
      builder.setInsertionPointAfter(ifOp);
    } else {
      genInsertionStore(codegen, builder, op, t, rhs);
    }
    return;
  }
  // Actual store.
  SmallVector<Value, 4> args;
  Value ptr = genSubscript(codegen, builder, op, t, args);
  builder.create<memref::StoreOp>(loc, rhs, ptr, args);
}

/// Generates an invariant value.
inline static Value genInvariantValue(Merger &merger, CodeGen &codegen,
                                      OpBuilder &builder, unsigned exp) {
  return merger.exp(exp).val;
}

/// Generates an index value.
inline static Value genIndexValue(CodeGen &codegen, OpBuilder &builder,
                                  unsigned idx) {
  return codegen.getLoopIdxValue(idx);
}

/// Semi-ring branches are simply inlined by the sparse compiler. Prior
/// analysis has verified that all computations are "local" to the inlined
/// branch or otherwise invariantly defined outside the loop nest, with the
/// exception of index computations, which need to be relinked to actual
/// inlined cloned code.
static Value relinkBranch(CodeGen &codegen, RewriterBase &rewriter,
                          Block *block, Value e, unsigned ldx) {
  if (Operation *def = e.getDefiningOp()) {
    if (auto indexOp = dyn_cast<linalg::IndexOp>(def))
      return genIndexValue(codegen, rewriter, indexOp.getDim());
    if (def->getBlock() == block) {
      for (unsigned i = 0, n = def->getNumOperands(); i < n; i++)
        def->setOperand(
            i, relinkBranch(codegen, rewriter, block, def->getOperand(i), ldx));
    }
  }
  return e;
}

/// Recursively generates tensor expression.
static Value genExp(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                    linalg::GenericOp op, unsigned exp, unsigned ldx) {
  Location loc = op.getLoc();
  if (exp == -1u)
    return Value();
  if (merger.exp(exp).kind == Kind::kTensor)
    return genTensorLoad(merger, codegen, rewriter, op, exp);
  if (merger.exp(exp).kind == Kind::kInvariant)
    return genInvariantValue(merger, codegen, rewriter, exp);
  if (merger.exp(exp).kind == Kind::kIndex)
    return genIndexValue(codegen, rewriter, merger.exp(exp).index);

  if (merger.exp(exp).kind == Kind::kReduce) {
    // Make custom reduction identity accessible for expanded access pattern.
    assert(codegen.redCustom == -1u);
    codegen.redCustom = exp;
  }

  Value v0 =
      genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e0, ldx);
  Value v1 =
      genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e1, ldx);
  Value ee = merger.buildExp(rewriter, loc, exp, v0, v1);
  if (ee && (merger.exp(exp).kind == Kind::kUnary ||
             merger.exp(exp).kind == Kind::kBinary ||
             merger.exp(exp).kind == Kind::kBinaryBranch ||
             merger.exp(exp).kind == Kind::kReduce ||
             merger.exp(exp).kind == Kind::kSelect))
    ee = relinkBranch(codegen, rewriter, ee.getParentBlock(), ee, ldx);

  if (merger.exp(exp).kind == kSelect) {
    assert(!merger.exp(exp).val);
    merger.exp(exp).val = v0; // Preserve value for later use.
  }

  if (merger.exp(exp).kind == Kind::kReduce) {
    assert(codegen.redCustom != -1u);
    codegen.redCustom = -1u;
  }

  return ee;
}

/// Determines if affine expression is invariant.
static bool isInvariantAffine(const CodeGen &codegen, AffineExpr a,
                              unsigned ldx, bool &atLevel) {
  switch (a.getKind()) {
  case AffineExprKind::DimId: {
    unsigned idx = a.cast<AffineDimExpr>().getPosition();
    if (idx == ldx)
      atLevel = true;
    return codegen.getLoopIdxValue(idx) != nullptr; // no longer in play?
  }
  case AffineExprKind::Add:
  case AffineExprKind::Mul: {
    auto binOp = a.cast<AffineBinaryOpExpr>();
    return isInvariantAffine(codegen, binOp.getLHS(), ldx, atLevel) &&
           isInvariantAffine(codegen, binOp.getRHS(), ldx, atLevel);
  }
  default:
    return true;
  }
}

/// Hoists loop invariant tensor loads for which indices have been exhausted.
static void genInvariants(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, unsigned exp, unsigned ldx,
                          bool atStart, unsigned last = -1u) {
  if (exp == -1u)
    return;
  if (merger.exp(exp).kind == Kind::kTensor) {
    // Inspect tensor indices.
    bool atLevel = ldx == -1u;
    OpOperand &t = op->getOpOperand(merger.exp(exp).tensor);
    auto map = op.getMatchingIndexingMap(&t);
    auto enc = getSparseTensorEncoding(t.get().getType());
    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
      AffineExpr a = map.getResult(toOrigDim(enc, d));
      if (!isInvariantAffine(codegen, a, ldx, atLevel))
        return; // still in play
    }
    // All exhausted at this level (atLevel denotes exactly at this level).
    if (!atLevel)
      return;
    OpOperand *lhs = op.getDpsInitOperand(0);
    if (lhs == &t) {
      // Start or end a scalarized reduction
      if (atStart) {
        Kind kind = merger.exp(last).kind;
        Value load = kind == Kind::kReduce
                         ? getCustomRedId(merger.exp(last).op)
                         : genTensorLoad(merger, codegen, builder, op, exp);
        codegen.redKind = getReduction(kind);
        codegen.redExp = exp;
        updateReduc(merger, codegen, load);
      } else {
        Value redVal = codegen.redVal;
        updateReduc(merger, codegen, Value());
        codegen.redExp = -1u;
        codegen.redKind = kNoReduc;
        genTensorStore(merger, codegen, builder, op, exp, redVal);
      }
    } else {
      // Start or end loop invariant hoisting of a tensor load.
      merger.exp(exp).val =
          atStart ? genTensorLoad(merger, codegen, builder, op, exp) : Value();
    }
  } else if (merger.exp(exp).kind != Kind::kInvariant &&
             merger.exp(exp).kind != Kind::kIndex) {
    // Traverse into the binary operations. Note that we only hoist
    // tensor loads, since subsequent MLIR/LLVM passes know how to
    // deal with all other kinds of derived loop invariants.
    unsigned e0 = merger.exp(exp).children.e0;
    unsigned e1 = merger.exp(exp).children.e1;
    genInvariants(merger, codegen, builder, op, e0, ldx, atStart, exp);
    genInvariants(merger, codegen, builder, op, e1, ldx, atStart, exp);
  }
}

/// Generates an expanded access pattern in innermost dimension.
static void genExpansion(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                         linalg::GenericOp op, unsigned at, bool atStart) {
  OpOperand *lhs = codegen.sparseOut;
  if (!lhs || codegen.outerParNest != op.getRank(lhs) - 1 ||
      at != codegen.outerParNest)
    return; // not needed at this level
  assert(codegen.redVal == nullptr);
  // Generate start or end of an expanded access pattern. Note that because
  // an expension does not rely on the ongoing contents of the sparse storage
  // scheme, we can use the original tensor as incoming SSA value (which
  // simplifies codegen a bit). If expansion on the actual contents is ever
  // needed, we will need to use the SSA value in the insertion chain instead.
  Value tensor = lhs->get();
  Location loc = op.getLoc();
  if (atStart) {
    auto dynShape = {ShapedType::kDynamicSize};
    Type etp = tensor.getType().cast<ShapedType>().getElementType();
    Type t1 = MemRefType::get(dynShape, etp);
    Type t2 = MemRefType::get(dynShape, builder.getI1Type());
    Type t3 = MemRefType::get(dynShape, builder.getIndexType());
    Type t4 = builder.getIndexType();
    auto res =
        builder.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
    assert(res.getNumResults() == 4);
    assert(!codegen.expValues);
    codegen.expValues = res.getResult(0);
    codegen.expFilled = res.getResult(1);
    codegen.expAdded = res.getResult(2);
    codegen.expCount = res.getResult(3);
  } else {
    assert(codegen.expValues);
    SmallVector<Value, 4> indices;
    for (unsigned i = 0; i < at; i++) {
      assert(codegen.loopEmitter.getLoopIV(i));
      indices.push_back(codegen.loopEmitter.getLoopIV(i));
    }
    codegen.insChain = builder.create<CompressOp>(
        loc, codegen.expValues, codegen.expFilled, codegen.expAdded,
        codegen.expCount, codegen.insChain, indices);
    codegen.expValues = codegen.expFilled = codegen.expAdded =
        codegen.expCount = Value();
  }
}

/// Returns parallelization strategy. Any implicit loop in the Linalg
/// operation that is marked "parallel" is a candidate. Whether it is actually
/// converted to a parallel operation depends on the requested strategy.
static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isSparse) {
  // Reject parallelization of sparse output.
  if (codegen.sparseOut)
    return false;
  // Parallel loops on tensor expansion can cause data races.
  if (codegen.expCount)
    return false;
  // Inspect strategy.
  switch (codegen.options.parallelizationStrategy) {
  case SparseParallelizationStrategy::kNone:
    return false;
  case SparseParallelizationStrategy::kDenseOuterLoop:
    return isOuter && !isSparse;
  case SparseParallelizationStrategy::kAnyStorageOuterLoop:
    return isOuter;
  case SparseParallelizationStrategy::kDenseAnyLoop:
    return !isSparse;
  case SparseParallelizationStrategy::kAnyStorageAnyLoop:
    return true;
  }
  llvm_unreachable("unexpected parallelization strategy");
}

/// Generates a for-loop on a single index.
static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                         linalg::GenericOp op, bool isOuter, bool isInner,
                         unsigned idx, size_t tid, size_t dim,
                         ArrayRef<size_t> extraTids,
                         ArrayRef<size_t> extraDims) {
  Location loc = op.getLoc();
  auto iteratorTypes = op.getIteratorTypesArray();
  bool isSparse = isCompressedDLT(merger.getDimLevelType(tid, idx)) ||
                  isSingletonDLT(merger.getDimLevelType(tid, idx));
  bool isParallel = isParallelFor(codegen, isOuter, isSparse);

  Operation *loop =
      genLoopBoundary(codegen, merger, [&](MutableArrayRef<Value> reduc) {
        return codegen.loopEmitter.enterLoopOverTensorAtDim(
            builder, loc, tid, dim, reduc, isParallel, extraTids, extraDims);
      }).value();
  assert(loop);
  return loop;
}

/// Emit a while-loop for co-iteration over multiple indices.
static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, unsigned idx, bool needsUniv,
                           ArrayRef<size_t> condTids, ArrayRef<size_t> condDims,
                           ArrayRef<size_t> extraTids,
                           ArrayRef<size_t> extraDims) {

  Operation *loop =
      genLoopBoundary(codegen, merger, [&](MutableArrayRef<Value> reduc) {
        // Construct the while-loop with a parameter for each index.
        return codegen.loopEmitter.enterCoIterationOverTensorsAtDims(
            builder, op.getLoc(), condTids, condDims, needsUniv, reduc,
            extraTids, extraDims);
      }).value();
  assert(loop);
  return loop;
}

/// Generates a for-loop or a while-loop, depending on whether it implements
/// singleton iteration or co-iteration over the given conjunction.
static Operation *genLoop(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, unsigned at, bool needsUniv,
                          ArrayRef<size_t> condTids, ArrayRef<size_t> condDims,
                          ArrayRef<size_t> extraTids,
                          ArrayRef<size_t> extraDims) {
  assert(condTids.size() == condDims.size());
  assert(extraTids.size() == extraDims.size());
  unsigned idx = codegen.topSort[at];
  if (condTids.size() == 1) {
    bool isOuter = at == 0;
    bool isInner = at == codegen.topSort.size() - 1;
    return genFor(merger, codegen, builder, op, isOuter, isInner, idx,
                  condTids.front(), condDims.front(), extraTids, extraDims);
  }
  return genWhile(merger, codegen, builder, op, idx, needsUniv, condTids,
                  condDims, extraTids, extraDims);
}

/// Generates the induction structure for a while-loop.
static void finalizeWhileOp(Merger &merger, CodeGen &codegen,
                            OpBuilder &builder, linalg::GenericOp op,
                            unsigned idx, bool needsUniv, BitVector &induction,
                            scf::WhileOp whileOp) {
  Location loc = op.getLoc();
  // Finalize each else branch of all if statements.
  if (codegen.redVal || codegen.expValues || codegen.insChain) {
    while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
               builder.getInsertionBlock()->getParentOp())) {
      unsigned y = 0;
      SmallVector<Value, 4> yields;
      if (codegen.redVal) {
        yields.push_back(codegen.redVal);
        updateReduc(merger, codegen, ifOp.getResult(y++));
      }
      if (codegen.expValues) {
        yields.push_back(codegen.expCount);
        codegen.expCount = ifOp->getResult(y++);
      }
      if (codegen.insChain) {
        yields.push_back(codegen.insChain);
        codegen.insChain = ifOp->getResult(y++);
      }
      assert(y == yields.size());
      builder.create<scf::YieldOp>(loc, yields);
      builder.setInsertionPointAfter(ifOp);
    }
  }
  builder.setInsertionPointToEnd(&whileOp.getAfter().front());
}

/// Generates a single if-statement within a while-loop.
static scf::IfOp genIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                       linalg::GenericOp op, unsigned idx,
                       BitVector &conditions) {
  Location loc = op.getLoc();
  SmallVector<Type, 4> types;
  Value cond;
  for (unsigned b = 0, be = conditions.size(); b < be; b++) {
    if (!conditions[b])
      continue;
    unsigned tensor = merger.tensor(b);
    assert(idx == merger.index(b));
    Value clause;
    if (isCompressedDLT(merger.getDimLevelType(b)) ||
        isSingletonDLT(merger.getDimLevelType(b))) {
      auto dim = merger.getDimNum(tensor, idx).value();
      Value op1 = codegen.loopEmitter.getCoord()[tensor][dim];
      Value op2 = codegen.getLoopIdxValue(idx);
      clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, op1,
                                             op2);
    } else {
      assert(isDenseDLT(merger.getDimLevelType(b)) ||
             isUndefDLT(merger.getDimLevelType(b)));
      clause = constantI1(builder, loc, true);
    }
    cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
  }
  if (codegen.redVal)
    types.push_back(codegen.redVal.getType());
  if (codegen.expValues)
    types.push_back(builder.getIndexType());
  if (codegen.insChain)
    types.push_back(codegen.insChain.getType());
  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, cond, /*else=*/true);
  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
  return ifOp;
}

/// Generates end of true branch of if-statement within a while-loop.
static void endIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                  linalg::GenericOp op, scf::IfOp ifOp, Operation *loop,
                  Value redInput, Value cntInput, Value insInput) {
  SmallVector<Value, 4> operands;
  if (codegen.redVal) {
    operands.push_back(codegen.redVal);
    updateReduc(merger, codegen, redInput);
  }
  if (codegen.expValues) {
    operands.push_back(codegen.expCount);
    codegen.expCount = cntInput;
  }
  if (codegen.insChain) {
    operands.push_back(codegen.insChain);
    codegen.insChain = insInput;
  }
  if (!operands.empty())
    builder.create<scf::YieldOp>(op.getLoc(), operands);
  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
}

//===----------------------------------------------------------------------===//
// Sparse compiler synthesis methods (loop sequence).
//===----------------------------------------------------------------------===//

/// Starts a loop sequence at given level. Returns true if
/// the universal loop index must be maintained at this level.
static bool startLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                         linalg::GenericOp op, unsigned exp, unsigned at,
                         unsigned idx, unsigned ldx, unsigned lts) {
  assert(!codegen.getLoopIdxValue(idx));
  // Emit invariants at this loop sequence level.
  genInvariants(merger, codegen, builder, op, exp, ldx, /*atStart=*/true);
  // Emit access pattern expansion for sparse tensor output.
  genExpansion(merger, codegen, builder, op, at, /*atStart=*/true);
  // Emit further intitialization at this loop sequence level.
  unsigned l0 = merger.set(lts)[0];
  bool needsUniv = false;

  SmallVector<size_t> tids;
  SmallVector<size_t> dims;
  merger.foreachTidDimPairInBits(
      merger.lat(l0).bits,
      [&](unsigned b, unsigned tid, Optional<unsigned> dim, DimLevelType dlt) {
        assert(merger.index(b) == idx);
        if (isDenseDLT(dlt) || isUndefDLT(dlt)) {
          needsUniv = true;
        } else {
          // sparse/singleton dim levels.
          tids.push_back(tid);
          dims.push_back(dim.value());
        }
      });

  codegen.loopEmitter.enterNewLoopSeq(builder, op.getLoc(), tids, dims);

  // Maintain the universal index only if it is actually
  // consumed by a subsequent lattice point.
  if (needsUniv) {
    unsigned lsize = merger.set(lts).size();
    for (unsigned i = 1; i < lsize; i++) {
      unsigned li = merger.set(lts)[i];
      if (!merger.hasAnySparse(merger.lat(li).simple))
        return true;
    }
  }
  return false;
}

static void translateBitsToTidDimPairs(Merger &merger, CodeGen &codegen,
                                       unsigned li, unsigned idx,
                                       SmallVectorImpl<size_t> &condTids,
                                       SmallVectorImpl<size_t> &condDims,
                                       SmallVectorImpl<size_t> &extraTids,
                                       SmallVectorImpl<size_t> &extraDims) {
  const BitVector &all = merger.lat(li).bits;
  const BitVector &simple = merger.lat(li).simple;

  // Converts bits to array + dim pair
  merger.foreachTidDimPairInBits(all, [&, idx](unsigned b, unsigned tid,
                                               Optional<unsigned> dim,
                                               DimLevelType dlt) {
    if (simple.test(b)) {
      if (isUndefDLT(dlt)) {
        // An undefined dlt in the lattices, we probably mean to iterate based
        // on the dim of output tensor.
        // E.g., this could be a synthetic tensor (for invariants and sparse
        // output tensor).
        // out[i][j] = invariant; or a broadcast
        // out[i][j] = in[i] (j is undef for input)
        tid = merger.getOutTensorID();
        dim = merger.getDimNum(tid, idx);
        // Skips invalid dim (e.g., when this is a zero ranked tensor).
        if (!dim)
          return;
      }
      condTids.push_back(tid);
      condDims.push_back(dim.value());
    } else if (isDenseDLT(dlt)) {
      // TODO: get rid of extraTids and extraDims.
      extraTids.push_back(tid);
      extraDims.push_back(dim.value());
    }
  });

  if (isDenseDLT(merger.getDimLevelType(merger.getOutTensorID(), idx))) {
    // Note that we generate dense indices of the output tensor
    // unconditionally, since they may not appear in the lattice, but may be
    // needed for linearized codegen.
    // Only dense dimensions should be optimized from conditions.
    auto dim = merger.getDimNum(merger.getOutTensorID(), idx).value();
    extraTids.push_back(merger.getOutTensorID());
    extraDims.push_back(dim);
  }
}

/// Starts a single loop in current sequence.
static Operation *startLoop(Merger &merger, CodeGen &codegen,
                            OpBuilder &builder, linalg::GenericOp op,
                            unsigned at, unsigned li, bool needsUniv) {
  // The set of tensors + dims to generate loops on
  SmallVector<size_t, 4> condTids, condDims;
  // The set of (dense) tensors that is optimized from condition, yet still
  // need extra locals to iterate on them.
  SmallVector<size_t, 4> extraTids, extraDims;

  translateBitsToTidDimPairs(merger, codegen, li, codegen.topSort[at], condTids,
                             condDims, extraTids, extraDims);
  // Emit the for/while-loop control.
  Operation *loop = genLoop(merger, codegen, builder, op, at, needsUniv,
                            condTids, condDims, extraTids, extraDims);
  return loop;
}

/// Ends a single loop in current sequence. Returns new values for needsUniv.
static bool endLoop(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                    linalg::GenericOp op, Operation *loop, unsigned idx,
                    unsigned li, bool needsUniv) {
  // End a while-loop.
  if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
    finalizeWhileOp(merger, codegen, rewriter, op, idx, needsUniv,
                    merger.lat(li).bits, whileOp);
  } else {
    needsUniv = false;
  }

  genLoopBoundary(codegen, merger, [&](MutableArrayRef<Value> reduc) {
    codegen.loopEmitter.exitCurrentLoop(rewriter, op.getLoc(), reduc);
    return llvm::None;
  });

  return needsUniv;
}

/// Ends a loop sequence at given level.
static void endLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                       linalg::GenericOp op, unsigned exp, unsigned at,
                       unsigned idx, unsigned ldx) {
  assert(codegen.getLoopIdxValue(idx) == nullptr);
  codegen.loopEmitter.exitCurrentLoopSeq();
  // Unmark bookkeeping of invariants and loop index.
  genInvariants(merger, codegen, builder, op, exp, ldx, /*atStart=*/false);
  // Finalize access pattern expansion for sparse tensor output.
  genExpansion(merger, codegen, builder, op, at, /*atStart=*/false);
}

/// Recursively generates code while computing iteration lattices in order
/// to manage the complexity of implementing co-iteration over unions
/// and intersections of sparse iterations spaces.
static void genStmt(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                    linalg::GenericOp op, unsigned exp, unsigned at) {
  // At each leaf, assign remaining tensor (sub)expression to output tensor.
  if (at == codegen.topSort.size()) {
    unsigned ldx = codegen.topSort[at - 1];
    Value rhs = genExp(merger, codegen, rewriter, op, exp, ldx);
    genTensorStore(merger, codegen, rewriter, op, exp, rhs);
    return;
  }

  // Construct iteration lattices for current loop index, with L0 at top.
  unsigned idx = codegen.topSort[at];
  unsigned ldx = at == 0 ? -1u : codegen.topSort[at - 1];
  unsigned lts = merger.optimizeSet(merger.buildLattices(exp, idx));

  // TODO: sort
  // TODO: dedup

  // Start a loop sequence.
  bool needsUniv =
      startLoopSeq(merger, codegen, rewriter, op, exp, at, idx, ldx, lts);

  // Emit a loop for every lattice point L0 >= Li in this loop sequence.
  unsigned lsize = merger.set(lts).size();
  for (unsigned i = 0; i < lsize; i++) {
    // Start a loop.
    unsigned li = merger.set(lts)[i];
    Operation *loop =
        startLoop(merger, codegen, rewriter, op, at, li, needsUniv);

    // Visit all lattices points with Li >= Lj to generate the
    // loop-body, possibly with if statements for coiteration.
    Value redInput = codegen.redVal;
    Value cntInput = codegen.expCount;
    Value insInput = codegen.insChain;
    bool isWhile = dyn_cast<scf::WhileOp>(loop) != nullptr;
    for (unsigned j = 0; j < lsize; j++) {
      unsigned lj = merger.set(lts)[j];
      unsigned ej = merger.lat(lj).exp;
      if (li == lj || merger.latGT(li, lj)) {
        // Recurse into body of each branch.
        if (isWhile) {
          scf::IfOp ifOp =
              genIf(merger, codegen, rewriter, op, idx, merger.lat(lj).simple);
          genStmt(merger, codegen, rewriter, op, ej, at + 1);
          endIf(merger, codegen, rewriter, op, ifOp, loop, redInput, cntInput,
                insInput);
        } else {
          genStmt(merger, codegen, rewriter, op, ej, at + 1);
        }
      }
    }

    // End a loop.
    needsUniv =
        endLoop(merger, codegen, rewriter, op, loop, idx, li, needsUniv);
  }

  // End a loop sequence.
  endLoopSeq(merger, codegen, rewriter, op, exp, at, idx, ldx);
}

/// Converts the result computed by the sparse kernel into the required form.
static void genResult(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                      linalg::GenericOp op) {
  OpOperand *lhs = op.getDpsInitOperand(0);
  Value tensor = lhs->get();
  Type resType = tensor.getType();
  if (getSparseTensorEncoding(resType)) {
    // The sparse tensor rematerializes from the original sparse tensor's
    // underlying sparse storage format. For an insertion chain, the
    // tensor materializes from the chain with 'hasInserts' enabled.
    bool hasInserts = codegen.sparseOut == lhs;
    if (hasInserts)
      tensor = codegen.insChain;
    rewriter.replaceOpWithNewOp<LoadOp>(op, resType, tensor, hasInserts);
  } else {
    // To rematerialize an non-annotated tensor, simply load it
    // from the bufferized value.
    Value val = codegen.loopEmitter.getValBuffer().back(); // value array
    rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
  }
}

//===----------------------------------------------------------------------===//
// Sparse compiler rewriting methods.
//===----------------------------------------------------------------------===//

namespace {

/// Sparse rewriting rule for generic Lingalg operation.
struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
public:
  GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
      : OpRewritePattern<linalg::GenericOp>(context), options(o) {}

  LogicalResult matchAndRewrite(linalg::GenericOp op,
                                PatternRewriter &rewriter) const override {
    // Detects sparse annotations and translate the per-dimension sparsity
    // information for all tensors to loop indices in the kernel.
    if (op.getNumDpsInits() != 1)
      return failure();
    unsigned numTensors = op->getNumOperands();
    unsigned numLoops = op.getNumLoops();
    Merger merger(numTensors, numLoops);
    if (!findSparseAnnotations(merger, op))
      return failure();

    // Builds the tensor expression for the Linalg operation in SSA form.
    Optional<unsigned> optExp = merger.buildTensorExpFromLinalg(op);
    if (!optExp.has_value())
      return failure();

    unsigned exp = optExp.value();
    OpOperand *sparseOut = nullptr;
    unsigned outerParNest = 0;
    // Computes a topologically sorted iteration graph to ensure tensors
    // are visited in natural index order. Gradually relaxes the considered
    // constraints until an acyclic iteration graph results, such that sparse
    // code generation can proceed. As a last resort, an attempt is made
    // to resolve cycles by inserting a conversion.
    std::vector<unsigned> topSort;
    // Whether the current GenericOp is admissible.
    bool isAdmissible = false;
    bool hasCycle = true;
    // An const list of all masks that we used for interation graph
    // computation. Must be ordered from strict -> loose.
    const auto allMask = {SortMask::kIncludeAll, SortMask::kIncludeUndef,
                          SortMask::kIncludeDense, SortMask::kSparseOnly};
    for (auto mask : allMask)
      if (computeIterationGraph(merger, op, topSort, mask)) {
        hasCycle = false;
        if (isAdmissibleTensorExp(merger, op, topSort, exp, &sparseOut,
                                  outerParNest)) {
          isAdmissible = true;
          break;
        }
        // else try a set of less strict constraints.
      }

    if (hasCycle)
      // Give it one last shot to resolve the cycle.
      return resolveCycle(merger, rewriter, op);
    if (!isAdmissible)
      // Inadmissible expression, reject.
      return failure();

    merger.setHasSparseOut(sparseOut != nullptr);

    SmallVector<Value, 4> 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);
    genBuffers(merger, codegen, rewriter, op);
    genStmt(merger, codegen, rewriter, op, exp, 0);
    genResult(merger, codegen, rewriter, op);
    return success();
  }

private:
  // Last resort cycle resolution.
  LogicalResult resolveCycle(Merger &merger, PatternRewriter &rewriter,
                             linalg::GenericOp op) const {
    // Compute topological sort while leaving out every
    // sparse input tensor in succession until an acylic
    // iteration graph results.
    std::vector<unsigned> topSort;
    for (OpOperand *t : op.getDpsInputOperands()) {
      unsigned tensor = t->getOperandNumber();
      Value tval = t->get();
      auto srcEnc = getSparseTensorEncoding(tval.getType());
      if (!srcEnc ||
          !computeIterationGraph(merger, op, topSort, SortMask::kSparseOnly, t))
        continue;
      // Found an input tensor that resolves the cycle by inserting a
      // conversion into a sparse tensor that adheres to the iteration
      // graph order. Also releases the temporary sparse tensor.
      //
      // TODO: investigate fusing the conversion with computation,
      //       especially if it is a direct yield!
      //
      auto srcTp = tval.getType().cast<RankedTensorType>();
      auto dstEnc = SparseTensorEncodingAttr::get(
          op->getContext(), srcEnc.getDimLevelType(),
          permute(getContext(), op.getMatchingIndexingMap(t),
                  topSort), // new order
          srcEnc.getHigherOrdering(), srcEnc.getPointerBitWidth(),
          srcEnc.getIndexBitWidth());
      auto dstTp = RankedTensorType::get(srcTp.getShape(),
                                         srcTp.getElementType(), dstEnc);
      auto convert = rewriter.create<ConvertOp>(tval.getLoc(), dstTp, tval);
      op->setOperand(tensor, convert);
      rewriter.setInsertionPointAfter(op);
      rewriter.create<bufferization::DeallocTensorOp>(tval.getLoc(), convert);
      return success();
    }
    // Cannot be resolved with a single conversion.
    // TODO: convert more than one?
    return failure();
  }

  /// Options to control sparse code generation.
  SparsificationOptions options;
};

} // namespace

/// Populates the given patterns list with rewriting rules required for
/// the sparsification of linear algebra operations.
void mlir::populateSparsificationPatterns(
    RewritePatternSet &patterns, const SparsificationOptions &options) {
  patterns.add<GenericOpSparsifier>(patterns.getContext(), options);
}