//===- SparseTensorCodegen.cpp - Sparse tensor primitives conversion ------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// A pass that converts sparse tensor types and primitives to actual compiler
// visible buffers and actual compiler IR that implements these primitives on
// the selected sparse tensor storage schemes. This pass provides an alternative
// to the SparseTensorConversion pass, eliminating the dependence on a runtime
// support library, and providing much more opportunities for subsequent
// compiler optimization of the generated code.
//
//===----------------------------------------------------------------------===//

#include "CodegenUtils.h"

#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/DialectConversion.h"

using namespace mlir;
using namespace mlir::sparse_tensor;

namespace {

static constexpr uint64_t DimSizesIdx = 0;
static constexpr uint64_t MemSizesIdx = 1;
static constexpr uint64_t FieldsIdx = 2;

//===----------------------------------------------------------------------===//
// Helper methods.
//===----------------------------------------------------------------------===//

/// Returns the "tuple" value of the adapted tensor.
static UnrealizedConversionCastOp getTuple(Value tensor) {
  return llvm::cast<UnrealizedConversionCastOp>(tensor.getDefiningOp());
}

/// Packs the given values as a "tuple" value.
static Value genTuple(OpBuilder &builder, Location loc, Type tp,
                      ValueRange values) {
  return builder.create<UnrealizedConversionCastOp>(loc, TypeRange(tp), values)
      .getResult(0);
}

/// Flatten a list of operands that may contain sparse tensors.
static void flattenOperands(ValueRange operands,
                            SmallVectorImpl<Value> &flattened) {
  // In case of
  // sparse_tensor, c, sparse_tensor
  // ==>
  // memref ..., c, memref ...
  for (auto operand : operands) {
    if (auto tuple = getTuple(operand);
        tuple && getSparseTensorEncoding(tuple->getResultTypes()[0]))
      // An unrealized_conversion_cast will be inserted by type converter to
      // inter-mix the gap between 1:N conversion between sparse tensors and
      // fields. In this case, take the operands in the cast and replace the
      // sparse tensor output with the flattened type array.
      flattened.append(tuple.getOperands().begin(), tuple.getOperands().end());
    else
      flattened.push_back(operand);
  }
}

/// Adds index conversions where needed.
static Value toType(OpBuilder &builder, Location loc, Value value, Type tp) {
  if (value.getType() != tp)
    return builder.create<arith::IndexCastOp>(loc, tp, value);
  return value;
}

/// Generates a load with proper index typing.
static Value genLoad(OpBuilder &builder, Location loc, Value mem, Value idx) {
  idx = toType(builder, loc, idx, builder.getIndexType());
  return builder.create<memref::LoadOp>(loc, mem, idx);
}

/// Generates a store with proper index typing and (for indices) proper value.
static void genStore(OpBuilder &builder, Location loc, Value val, Value mem,
                     Value idx) {
  idx = toType(builder, loc, idx, builder.getIndexType());
  val = toType(builder, loc, val,
               mem.getType().cast<ShapedType>().getElementType());
  builder.create<memref::StoreOp>(loc, val, mem, idx);
}

/// Creates a straightforward counting for-loop.
static scf::ForOp createFor(OpBuilder &builder, Location loc, Value upper,
                            SmallVectorImpl<Value> &fields,
                            Value lower = Value()) {
  Type indexType = builder.getIndexType();
  if (!lower)
    lower = constantZero(builder, loc, indexType);
  Value one = constantOne(builder, loc, indexType);
  scf::ForOp forOp = builder.create<scf::ForOp>(loc, lower, upper, one, fields);
  for (unsigned i = 0, e = fields.size(); i < e; i++)
    fields[i] = forOp.getRegionIterArg(i);
  builder.setInsertionPointToStart(forOp.getBody());
  return forOp;
}

/// Gets the dimension size for the given sparse tensor at the given
/// original dimension 'dim'. Returns None if no sparse encoding is
/// attached to the given tensor type.
static Optional<Value> sizeFromTensorAtDim(OpBuilder &builder, Location loc,
                                           RankedTensorType tensorTp,
                                           Value adaptedValue, unsigned dim) {
  auto enc = getSparseTensorEncoding(tensorTp);
  if (!enc)
    return llvm::None;

  // Access into static dimension can query original type directly.
  // Note that this is typically already done by DimOp's folding.
  auto shape = tensorTp.getShape();
  if (!ShapedType::isDynamic(shape[dim]))
    return constantIndex(builder, loc, shape[dim]);

  // Any other query can consult the dimSizes array at field DimSizesIdx,
  // accounting for the reordering applied to the sparse storage.
  auto tuple = getTuple(adaptedValue);
  Value idx = constantIndex(builder, loc, toStoredDim(tensorTp, dim));
  return builder
      .create<memref::LoadOp>(loc, tuple.getInputs()[DimSizesIdx], idx)
      .getResult();
}

// Gets the dimension size at the given stored dimension 'd', either as a
// constant for a static size, or otherwise dynamically through memSizes.
Value sizeAtStoredDim(OpBuilder &builder, Location loc, RankedTensorType rtp,
                      SmallVectorImpl<Value> &fields, unsigned d) {
  unsigned dim = toOrigDim(rtp, d);
  auto shape = rtp.getShape();
  if (!ShapedType::isDynamic(shape[dim]))
    return constantIndex(builder, loc, shape[dim]);
  return genLoad(builder, loc, fields[DimSizesIdx],
                 constantIndex(builder, loc, d));
}

/// Translates field index to memSizes index.
static unsigned getMemSizesIndex(unsigned field) {
  assert(FieldsIdx <= field);
  return field - FieldsIdx;
}

/// Creates a pushback op for given field and updates the fields array
/// accordingly. This operation also updates the memSizes contents.
static void createPushback(OpBuilder &builder, Location loc,
                           SmallVectorImpl<Value> &fields, unsigned field,
                           Value value, Value repeat = Value()) {
  assert(FieldsIdx <= field && field < fields.size());
  Type etp = fields[field].getType().cast<ShapedType>().getElementType();
  fields[field] = builder.create<PushBackOp>(
      loc, fields[field].getType(), fields[MemSizesIdx], fields[field],
      toType(builder, loc, value, etp), APInt(64, getMemSizesIndex(field)),
      repeat);
}

/// Returns field index of sparse tensor type for pointers/indices, when set.
static unsigned getFieldIndex(Type type, unsigned ptrDim, unsigned idxDim) {
  assert(getSparseTensorEncoding(type));
  RankedTensorType rType = type.cast<RankedTensorType>();
  unsigned field = FieldsIdx; // start past header
  for (unsigned r = 0, rank = rType.getShape().size(); r < rank; r++) {
    if (isCompressedDim(rType, r)) {
      if (r == ptrDim)
        return field;
      field++;
      if (r == idxDim)
        return field;
      field++;
    } else if (isSingletonDim(rType, r)) {
      if (r == idxDim)
        return field;
      field++;
    } else {
      assert(isDenseDim(rType, r)); // no fields
    }
  }
  assert(ptrDim == -1u && idxDim == -1u);
  return field + 1; // return values field index
}

/// Maps a sparse tensor type to the appropriate compounded buffers.
static Optional<LogicalResult>
convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
  auto enc = getSparseTensorEncoding(type);
  if (!enc)
    return llvm::None;
  // Construct the basic types.
  auto *context = type.getContext();
  unsigned idxWidth = enc.getIndexBitWidth();
  unsigned ptrWidth = enc.getPointerBitWidth();
  RankedTensorType rType = type.cast<RankedTensorType>();
  Type indexType = IndexType::get(context);
  Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
  Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
  Type eltType = rType.getElementType();
  //
  // Sparse tensor storage scheme for rank-dimensional tensor is organized
  // as a single compound type with the following fields. Note that every
  // memref with ? size actually behaves as a "vector", i.e. the stored
  // size is the capacity and the used size resides in the memSizes array.
  //
  // struct {
  //   memref<rank x index> dimSizes     ; size in each dimension
  //   memref<n x index> memSizes        ; sizes of ptrs/inds/values
  //   ; per-dimension d:
  //   ;  if dense:
  //        <nothing>
  //   ;  if compresed:
  //        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
  //        memref<? x idx>  indices-d   ; indices for sparse dim d
  //   ;  if singleton:
  //        memref<? x idx>  indices-d   ; indices for singleton dim d
  //   memref<? x eltType> values        ; values
  // };
  //
  unsigned rank = rType.getShape().size();
  unsigned lastField = getFieldIndex(type, -1u, -1u);
  // The dimSizes array and memSizes array.
  fields.push_back(MemRefType::get({rank}, indexType));
  fields.push_back(MemRefType::get({getMemSizesIndex(lastField)}, indexType));
  // Per-dimension storage.
  for (unsigned r = 0; r < rank; r++) {
    // Dimension level types apply in order to the reordered dimension.
    // As a result, the compound type can be constructed directly in the given
    // order. Clients of this type know what field is what from the sparse
    // tensor type.
    if (isCompressedDim(rType, r)) {
      fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, ptrType));
      fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
    } else if (isSingletonDim(rType, r)) {
      fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
    } else {
      assert(isDenseDim(rType, r)); // no fields
    }
  }
  // The values array.
  fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));
  assert(fields.size() == lastField);
  return success();
}

/// Generates code that allocates a sparse storage scheme for given rank.
static void allocSchemeForRank(OpBuilder &builder, Location loc,
                               RankedTensorType rtp,
                               SmallVectorImpl<Value> &fields, unsigned field,
                               unsigned r0) {
  unsigned rank = rtp.getShape().size();
  Value linear = constantIndex(builder, loc, 1);
  for (unsigned r = r0; r < rank; r++) {
    if (isCompressedDim(rtp, r)) {
      // Append linear x pointers, initialized to zero. Since each compressed
      // dimension initially already has a single zero entry, this maintains
      // the desired "linear + 1" length property at all times.
      unsigned ptrWidth = getSparseTensorEncoding(rtp).getPointerBitWidth();
      Type indexType = builder.getIndexType();
      Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
      Value ptrZero = constantZero(builder, loc, ptrType);
      createPushback(builder, loc, fields, field, ptrZero, linear);
      return;
    } else if (isSingletonDim(rtp, r)) {
      return; // nothing to do
    } else {
      // Keep compounding the size, but nothing needs to be initialized
      // at this level. We will eventually reach a compressed level or
      // otherwise the values array for the from-here "all-dense" case.
      assert(isDenseDim(rtp, r));
      Value size = sizeAtStoredDim(builder, loc, rtp, fields, r);
      linear = builder.create<arith::MulIOp>(loc, linear, size);
    }
  }
  // Reached values array so prepare for an insertion.
  Value valZero = constantZero(builder, loc, rtp.getElementType());
  createPushback(builder, loc, fields, field, valZero, linear);
  assert(fields.size() == ++field);
}

/// Creates allocation operation.
static Value createAllocation(OpBuilder &builder, Location loc, Type type,
                              Value sz, bool enableInit) {
  auto memType = MemRefType::get({ShapedType::kDynamicSize}, type);
  Value buffer = builder.create<memref::AllocOp>(loc, memType, sz);
  if (enableInit) {
    Value fillValue =
        builder.create<arith::ConstantOp>(loc, type, builder.getZeroAttr(type));
    builder.create<linalg::FillOp>(loc, fillValue, buffer);
  }
  return buffer;
}

/// Creates allocation for each field in sparse tensor type. Note that
/// for all dynamic memrefs, the memory size is really the capacity of
/// the "vector", while the actual size resides in the sizes array.
///
/// TODO: for efficiency, we will need heuristis to make educated guesses
///       on the required capacities (see heuristic variable).
///
static void createAllocFields(OpBuilder &builder, Location loc, Type type,
                              ValueRange dynSizes, bool enableInit,
                              SmallVectorImpl<Value> &fields) {
  auto enc = getSparseTensorEncoding(type);
  assert(enc);
  // Construct the basic types.
  unsigned idxWidth = enc.getIndexBitWidth();
  unsigned ptrWidth = enc.getPointerBitWidth();
  RankedTensorType rtp = type.cast<RankedTensorType>();
  Type indexType = builder.getIndexType();
  Type idxType = idxWidth ? builder.getIntegerType(idxWidth) : indexType;
  Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
  Type eltType = rtp.getElementType();
  auto shape = rtp.getShape();
  unsigned rank = shape.size();
  Value heuristic = constantIndex(builder, loc, 16);
  // Build original sizes.
  SmallVector<Value, 8> sizes;
  for (unsigned r = 0, o = 0; r < rank; r++) {
    if (ShapedType::isDynamic(shape[r]))
      sizes.push_back(dynSizes[o++]);
    else
      sizes.push_back(constantIndex(builder, loc, shape[r]));
  }
  // The dimSizes array and memSizes array.
  unsigned lastField = getFieldIndex(type, -1u, -1u);
  Value dimSizes =
      builder.create<memref::AllocOp>(loc, MemRefType::get({rank}, indexType));
  Value memSizes = builder.create<memref::AllocOp>(
      loc, MemRefType::get({getMemSizesIndex(lastField)}, indexType));
  fields.push_back(dimSizes);
  fields.push_back(memSizes);
  // Per-dimension storage.
  for (unsigned r = 0; r < rank; r++) {
    if (isCompressedDim(rtp, r)) {
      fields.push_back(
          createAllocation(builder, loc, ptrType, heuristic, enableInit));
      fields.push_back(
          createAllocation(builder, loc, idxType, heuristic, enableInit));
    } else if (isSingletonDim(rtp, r)) {
      fields.push_back(
          createAllocation(builder, loc, idxType, heuristic, enableInit));
    } else {
      assert(isDenseDim(rtp, r)); // no fields
    }
  }
  // The values array.
  fields.push_back(
      createAllocation(builder, loc, eltType, heuristic, enableInit));
  assert(fields.size() == lastField);
  // Initialize the storage scheme to an empty tensor. Initialized memSizes
  // to all zeros, sets the dimSizes to known values and gives all pointer
  // fields an initial zero entry, so that it is easier to maintain the
  // "linear + 1" length property.
  builder.create<linalg::FillOp>(
      loc, ValueRange{constantZero(builder, loc, indexType)},
      ValueRange{memSizes}); // zero memSizes
  Value ptrZero = constantZero(builder, loc, ptrType);
  for (unsigned r = 0, field = FieldsIdx; r < rank; r++) {
    unsigned ro = toOrigDim(rtp, r);
    genStore(builder, loc, sizes[ro], dimSizes, constantIndex(builder, loc, r));
    if (isCompressedDim(rtp, r)) {
      createPushback(builder, loc, fields, field, ptrZero);
      field += 2;
    } else if (isSingletonDim(rtp, r)) {
      field += 1;
    }
  }
  allocSchemeForRank(builder, loc, rtp, fields, FieldsIdx, /*rank=*/0);
}

/// Helper method that generates block specific to compressed case:
///
///  plo = pointers[d][pos[d-1]]
///  phi = pointers[d][pos[d-1]+1]
///  msz = indices[d].size()
///  if (plo < phi) {
///    present = indices[d][phi-1] == i[d]
///  } else { // first insertion
///    present = false
///    pointers[d][pos[d-1]] = msz
///  }
///  if (present) { // index already present
///    next = phi-1
///  } else {
///    indices[d].push_back(i[d])
///    pointers[d][pos[d-1]+1] = msz+1
///    next = msz
///    <prepare dimension d + 1>
///  }
///  pos[d] = next
static Value genCompressed(OpBuilder &builder, Location loc,
                           RankedTensorType rtp, SmallVectorImpl<Value> &fields,
                           SmallVectorImpl<Value> &indices, Value value,
                           Value pos, unsigned field, unsigned d) {
  unsigned rank = rtp.getShape().size();
  SmallVector<Type, 4> types;
  Type indexType = builder.getIndexType();
  Type boolType = builder.getIntegerType(1);
  Value one = constantIndex(builder, loc, 1);
  Value pp1 = builder.create<arith::AddIOp>(loc, pos, one);
  Value plo = genLoad(builder, loc, fields[field], pos);
  Value phi = genLoad(builder, loc, fields[field], pp1);
  Value psz = constantIndex(builder, loc, getMemSizesIndex(field + 1));
  Value msz = genLoad(builder, loc, fields[MemSizesIdx], psz);
  Value phim1 = builder.create<arith::SubIOp>(
      loc, toType(builder, loc, phi, indexType), one);
  // Conditional expression.
  Value lt =
      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, plo, phi);
  types.push_back(boolType);
  scf::IfOp ifOp1 = builder.create<scf::IfOp>(loc, types, lt, /*else*/ true);
  types.pop_back();
  builder.setInsertionPointToStart(&ifOp1.getThenRegion().front());
  Value crd = genLoad(builder, loc, fields[field + 1], phim1);
  Value eq = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                           toType(builder, loc, crd, indexType),
                                           indices[d]);
  builder.create<scf::YieldOp>(loc, eq);
  builder.setInsertionPointToStart(&ifOp1.getElseRegion().front());
  if (d > 0)
    genStore(builder, loc, msz, fields[field], pos);
  builder.create<scf::YieldOp>(loc, constantI1(builder, loc, false));
  builder.setInsertionPointAfter(ifOp1);
  Value p = ifOp1.getResult(0);
  // If present construct. Note that for a non-unique dimension level, we simply
  // set the condition to false and rely on CSE/DCE to clean up the IR.
  //
  // TODO: generate less temporary IR?
  //
  for (unsigned i = 0, e = fields.size(); i < e; i++)
    types.push_back(fields[i].getType());
  types.push_back(indexType);
  if (!isUniqueDim(rtp, d))
    p = constantI1(builder, loc, false);
  scf::IfOp ifOp2 = builder.create<scf::IfOp>(loc, types, p, /*else*/ true);
  // If present (fields unaffected, update next to phim1).
  builder.setInsertionPointToStart(&ifOp2.getThenRegion().front());
  fields.push_back(phim1);
  builder.create<scf::YieldOp>(loc, fields);
  fields.pop_back();
  // If !present (changes fields, update next).
  builder.setInsertionPointToStart(&ifOp2.getElseRegion().front());
  Value mszp1 = builder.create<arith::AddIOp>(loc, msz, one);
  genStore(builder, loc, mszp1, fields[field], pp1);
  createPushback(builder, loc, fields, field + 1, indices[d]);
  // Prepare the next dimension "as needed".
  if ((d + 1) < rank)
    allocSchemeForRank(builder, loc, rtp, fields, field + 2, d + 1);
  fields.push_back(msz);
  builder.create<scf::YieldOp>(loc, fields);
  fields.pop_back();
  // Update fields and return next pos.
  builder.setInsertionPointAfter(ifOp2);
  unsigned o = 0;
  for (unsigned i = 0, e = fields.size(); i < e; i++)
    fields[i] = ifOp2.getResult(o++);
  return ifOp2.getResult(o);
}

/// Generates code along an insertion path without the need for a "cursor".
/// This current insertion strategy comes at the expense of some testing
/// overhead for each insertion. The strategy will be optimized later for
/// common insertion patterns. The current insertion strategy also assumes
/// insertions occur in "a reasonable order" that enables building the
/// storage scheme in an appending/inserting kind of fashion (i.e. no
/// in-between insertions that need data movement). The implementation
/// relies on CSE/DCE to clean up all bookkeeping that is not needed.
///
/// TODO: better unord/not-unique; also generalize, optimize, specialize!
///
static void genInsert(OpBuilder &builder, Location loc, RankedTensorType rtp,
                      SmallVectorImpl<Value> &fields,
                      SmallVectorImpl<Value> &indices, Value value) {
  unsigned rank = rtp.getShape().size();
  assert(rank == indices.size());
  unsigned field = FieldsIdx; // start past header
  Value pos = constantZero(builder, loc, builder.getIndexType());
  // Generate code for every dimension.
  for (unsigned d = 0; d < rank; d++) {
    if (isCompressedDim(rtp, d)) {
      // Create:
      //   if (!present) {
      //     indices[d].push_back(i[d])
      //     <update pointers and prepare dimension d + 1>
      //   }
      //   pos[d] = indices.size() - 1
      //   <insert @ pos[d] at next dimension d + 1>
      pos = genCompressed(builder, loc, rtp, fields, indices, value, pos, field,
                          d);
      field += 2;
    } else if (isSingletonDim(rtp, d)) {
      // Create:
      //   indices[d].push_back(i[d])
      //   pos[d] = pos[d-1]
      //   <insert @ pos[d] at next dimension d + 1>
      createPushback(builder, loc, fields, field, indices[d]);
      field += 1;
    } else {
      assert(isDenseDim(rtp, d));
      // Construct the new position as:
      //   pos[d] = size * pos[d-1] + i[d]
      //   <insert @ pos[d] at next dimension d + 1>
      Value size = sizeAtStoredDim(builder, loc, rtp, fields, d);
      Value mult = builder.create<arith::MulIOp>(loc, size, pos);
      pos = builder.create<arith::AddIOp>(loc, mult, indices[d]);
    }
  }
  // Reached the actual value append/insert.
  if (!isDenseDim(rtp, rank - 1))
    createPushback(builder, loc, fields, field++, value);
  else
    genStore(builder, loc, value, fields[field++], pos);
  assert(fields.size() == field);
}

/// Generations insertion finalization code.
static void genEndInsert(OpBuilder &builder, Location loc, RankedTensorType rtp,
                         SmallVectorImpl<Value> &fields) {
  unsigned rank = rtp.getShape().size();
  unsigned field = FieldsIdx; // start past header
  for (unsigned d = 0; d < rank; d++) {
    if (isCompressedDim(rtp, d)) {
      // Compressed dimensions need a pointer cleanup for all entries
      // that were not visited during the insertion pass.
      //
      // TODO: avoid cleanup and keep compressed scheme consistent at all times?
      //
      if (d > 0) {
        unsigned ptrWidth = getSparseTensorEncoding(rtp).getPointerBitWidth();
        Type indexType = builder.getIndexType();
        Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
        Value mz = constantIndex(builder, loc, getMemSizesIndex(field));
        Value hi = genLoad(builder, loc, fields[MemSizesIdx], mz);
        Value zero = constantIndex(builder, loc, 0);
        Value one = constantIndex(builder, loc, 1);
        SmallVector<Value, 1> inits;
        inits.push_back(genLoad(builder, loc, fields[field], zero));
        scf::ForOp loop = createFor(builder, loc, hi, inits, one);
        Value i = loop.getInductionVar();
        Value oldv = loop.getRegionIterArg(0);
        Value newv = genLoad(builder, loc, fields[field], i);
        Value ptrZero = constantZero(builder, loc, ptrType);
        Value cond = builder.create<arith::CmpIOp>(
            loc, arith::CmpIPredicate::eq, newv, ptrZero);
        scf::IfOp ifOp = builder.create<scf::IfOp>(loc, TypeRange(ptrType),
                                                   cond, /*else*/ true);
        builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
        genStore(builder, loc, oldv, fields[field], i);
        builder.create<scf::YieldOp>(loc, oldv);
        builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
        builder.create<scf::YieldOp>(loc, newv);
        builder.setInsertionPointAfter(ifOp);
        builder.create<scf::YieldOp>(loc, ifOp.getResult(0));
        builder.setInsertionPointAfter(loop);
      }
      field += 2;
    } else if (isSingletonDim(rtp, d)) {
      field++;
    } else {
      assert(isDenseDim(rtp, d));
    }
  }
  assert(fields.size() == ++field);
}

//===----------------------------------------------------------------------===//
// Codegen rules.
//===----------------------------------------------------------------------===//

/// Sparse tensor storage conversion rule for returns.
class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SmallVector<Value, 8> flattened;
    flattenOperands(adaptor.getOperands(), flattened);
    // Create a return with the flattened value extracted from sparse tensors.
    rewriter.replaceOpWithNewOp<func::ReturnOp>(op, flattened);
    return success();
  }
};

/// Sparse tensor storage conversion rule for calls.
class SparseCallConverter : public OpConversionPattern<func::CallOp> {
public:
  // The default CallOp converter can not handle 1:N type conversion.
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(func::CallOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    // In case of:
    //  sparse_tensor, f, sparse_tensor = call @foo(...)
    // ==>
    //  memref..., f, memref = call @foo(...) replace with
    //  cast(memref...)->sparse_tensor, f, cast(memref...)->sparse_tensor
    SmallVector<Type, 8> finalRetTy;
    if (failed(typeConverter->convertTypes(op.getResultTypes(), finalRetTy)))
      return failure();

    // (1) Genereates new call with flattened return value.
    SmallVector<Value, 8> flattened;
    flattenOperands(adaptor.getOperands(), flattened);
    auto newCall = rewriter.create<func::CallOp>(loc, op.getCallee(),
                                                 finalRetTy, flattened);
    // (2) Create cast operation for sparse tensor returns.
    SmallVector<Value, 4> castedRet;
    // Tracks the offset of current return value (of the orignal call)
    // relative to the new call (after sparse tensor flattening);
    unsigned retOffset = 0;
    // Temporal buffer to hold the flattened list of type for
    // a sparse tensor.
    SmallVector<Type, 8> sparseFlat;
    for (auto ret : op.getResults()) {
      assert(retOffset < newCall.getNumResults());
      auto retType = ret.getType();
      if (failed(typeConverter->convertType(retType, sparseFlat)))
        // This should never happen.
        llvm_unreachable("Failed to convert type in sparse tensor codegen");

      // Converted types can not be empty when the type conversion succeed.
      assert(!sparseFlat.empty());
      if (sparseFlat.size() > 1) {
        auto flatSize = sparseFlat.size();
        ValueRange fields(iterator_range<ResultRange::iterator>(
            newCall.result_begin() + retOffset,
            newCall.result_begin() + retOffset + flatSize));
        castedRet.push_back(genTuple(rewriter, loc, retType, fields));
        retOffset += flatSize;
      } else {
        // If this is an 1:1 conversion, no need for casting.
        castedRet.push_back(newCall.getResult(retOffset));
        retOffset++;
      }
      sparseFlat.clear();
    }

    assert(castedRet.size() == op.getNumResults());
    rewriter.replaceOp(op, castedRet);
    return success();
  }
};

/// Sparse codegen rule for dimension accesses.
class SparseDimOpConverter : public OpConversionPattern<tensor::DimOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Optional<int64_t> index = op.getConstantIndex();
    if (!index)
      return failure();
    auto sz =
        sizeFromTensorAtDim(rewriter, op.getLoc(),
                            op.getSource().getType().cast<RankedTensorType>(),
                            adaptor.getSource(), *index);
    if (!sz)
      return failure();

    rewriter.replaceOp(op, *sz);
    return success();
  }
};

/// Sparse codegen rule for trivial tensor casts.
class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(tensor::CastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only rewrite identically annotated source/dest.
    auto encDst = getSparseTensorEncoding(op.getType());
    auto encSrc = getSparseTensorEncoding(op.getSource().getType());
    if (!encDst || encDst != encSrc)
      return failure();
    rewriter.replaceOp(op, adaptor.getOperands());
    return success();
  }
};

/// Sparse codgen rule for the alloc operator.
class SparseTensorAllocConverter
    : public OpConversionPattern<bufferization::AllocTensorOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  SparseTensorAllocConverter(TypeConverter &typeConverter, MLIRContext *context,
                             bool enableInit)
      : OpConversionPattern(typeConverter, context),
        enableBufferInitialization(enableInit) {}
  LogicalResult
  matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    RankedTensorType resType = op.getType();
    auto enc = getSparseTensorEncoding(resType);
    if (!enc)
      return failure();
    if (op.getCopy())
      return rewriter.notifyMatchFailure(op, "tensor copy not implemented");

    // Construct allocation for each field.
    Location loc = op.getLoc();
    SmallVector<Value, 8> fields;
    createAllocFields(rewriter, loc, resType, adaptor.getOperands(),
                      enableBufferInitialization, fields);
    // Replace operation with resulting memrefs.
    rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields));
    return success();
  }

private:
  bool enableBufferInitialization;
};

/// Sparse codegen rule for the dealloc operator.
class SparseTensorDeallocConverter
    : public OpConversionPattern<bufferization::DeallocTensorOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto enc = getSparseTensorEncoding(op.getTensor().getType());
    if (!enc)
      return failure();

    // Replace the sparse tensor deallocation with field deallocations.
    Location loc = op.getLoc();
    auto tuple = getTuple(adaptor.getTensor());
    for (auto input : tuple.getInputs())
      // Deallocate every buffer used to store the sparse tensor handler.
      rewriter.create<memref::DeallocOp>(loc, input);

    rewriter.eraseOp(op);
    return success();
  }
};

/// Sparse codegen rule for tensor rematerialization.
class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(LoadOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    RankedTensorType srcType =
        op.getTensor().getType().cast<RankedTensorType>();
    auto tuple = getTuple(adaptor.getTensor());
    // Prepare fields.
    SmallVector<Value, 8> fields(tuple.getInputs());
    // Generate optional insertion finalization code.
    if (op.getHasInserts())
      genEndInsert(rewriter, op.getLoc(), srcType, fields);
    // Replace operation with resulting memrefs.
    rewriter.replaceOp(op, genTuple(rewriter, op.getLoc(), srcType, fields));
    return success();
  }
};

/// Sparse codegen rule for the expand op.
class SparseExpandConverter : public OpConversionPattern<ExpandOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    RankedTensorType srcType =
        op.getTensor().getType().cast<RankedTensorType>();
    Type eltType = srcType.getElementType();
    Type boolType = rewriter.getIntegerType(1);
    Type idxType = rewriter.getIndexType();
    // All initialization should be done on entry of the loop nest.
    rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
    // Determine the size for access expansion (always the innermost stored
    // dimension size, translated back to original dimension). Note that we
    // recursively rewrite the new DimOp on the **original** tensor.
    unsigned innerDim = toOrigDim(srcType, srcType.getRank() - 1);
    auto sz = sizeFromTensorAtDim(rewriter, loc, srcType, adaptor.getTensor(),
                                  innerDim);
    assert(sz); // This for sure is a sparse tensor
    // Generate a memref for `sz` elements of type `t`.
    auto genAlloc = [&](Type t) {
      auto memTp = MemRefType::get({ShapedType::kDynamicSize}, t);
      return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{*sz});
    };
    // Allocate temporary buffers for values/filled-switch and added.
    // We do not use stack buffers for this, since the expanded size may
    // be rather large (as it envelops a single expanded dense dimension).
    Value values = genAlloc(eltType);
    Value filled = genAlloc(boolType);
    Value added = genAlloc(idxType);
    Value zero = constantZero(rewriter, loc, idxType);
    // Reset the values/filled-switch to all-zero/false. Note that this
    // introduces an O(N) operation into the computation, but this reset
    // operation is amortized over the innermost loops for the access
    // pattern expansion. As noted in the operation doc, we would like
    // to amortize this setup cost even between kernels.
    rewriter.create<linalg::FillOp>(
        loc, ValueRange{constantZero(rewriter, loc, eltType)},
        ValueRange{values});
    rewriter.create<linalg::FillOp>(
        loc, ValueRange{constantZero(rewriter, loc, boolType)},
        ValueRange{filled});
    // Replace expansion op with these buffers and initial index.
    assert(op.getNumResults() == 4);
    rewriter.replaceOp(op, {values, filled, added, zero});
    return success();
  }
};

/// Sparse codegen rule for the compress operator.
class SparseCompressConverter : public OpConversionPattern<CompressOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(CompressOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op->getLoc();
    RankedTensorType dstType =
        op.getTensor().getType().cast<RankedTensorType>();
    Type eltType = dstType.getElementType();
    auto tuple = getTuple(adaptor.getTensor());
    Value values = adaptor.getValues();
    Value filled = adaptor.getFilled();
    Value added = adaptor.getAdded();
    Value count = adaptor.getCount();
    // Prepare fields and indices.
    SmallVector<Value, 8> fields(tuple.getInputs());
    SmallVector<Value, 8> indices(adaptor.getIndices());
    // If the innermost dimension is ordered, we need to sort the indices
    // in the "added" array prior to applying the compression.
    unsigned rank = dstType.getShape().size();
    if (isOrderedDim(dstType, rank - 1))
      rewriter.create<SortOp>(loc, count, ValueRange{added}, ValueRange{});
    // While performing the insertions, we also need to reset the elements
    // of the values/filled-switch by only iterating over the set elements,
    // to ensure that the runtime complexity remains proportional to the
    // sparsity of the expanded access pattern.
    //
    // Generate
    //    out_memrefs = for (i = 0; i < count; i++)(in_memrefs) {
    //      index = added[i];
    //      value = values[index];
    //      insert({prev_indices, index}, value);
    //      new_memrefs = insert(in_memrefs, {prev_indices, index}, value);
    //      values[index] = 0;
    //      filled[index] = false;
    //      yield new_memrefs
    //    }
    scf::ForOp loop = createFor(rewriter, loc, count, fields);
    Value i = loop.getInductionVar();
    Value index = genLoad(rewriter, loc, added, i);
    Value value = genLoad(rewriter, loc, values, index);
    indices.push_back(index);
    // TODO: faster for subsequent insertions?
    genInsert(rewriter, loc, dstType, fields, indices, value);
    genStore(rewriter, loc, constantZero(rewriter, loc, eltType), values,
             index);
    genStore(rewriter, loc, constantI1(rewriter, loc, false), filled, index);
    rewriter.create<scf::YieldOp>(loc, fields);
    rewriter.setInsertionPointAfter(loop);
    Value result = genTuple(rewriter, loc, dstType, loop->getResults());
    // Deallocate the buffers on exit of the full loop nest.
    Operation *parent = getTop(op);
    rewriter.setInsertionPointAfter(parent);
    rewriter.create<memref::DeallocOp>(loc, values);
    rewriter.create<memref::DeallocOp>(loc, filled);
    rewriter.create<memref::DeallocOp>(loc, added);
    // Replace operation with resulting memrefs.
    rewriter.replaceOp(op, result);
    return success();
  }
};

/// Sparse codegen rule for the insert operator.
class SparseInsertConverter : public OpConversionPattern<InsertOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(InsertOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    RankedTensorType dstType =
        op.getTensor().getType().cast<RankedTensorType>();
    auto tuple = getTuple(adaptor.getTensor());
    // Prepare fields and indices.
    SmallVector<Value, 8> fields(tuple.getInputs());
    SmallVector<Value, 8> indices(adaptor.getIndices());
    // Generate insertion.
    Value value = adaptor.getValue();
    genInsert(rewriter, op->getLoc(), dstType, fields, indices, value);
    // Replace operation with resulting memrefs.
    rewriter.replaceOp(op, genTuple(rewriter, op.getLoc(), dstType, fields));
    return success();
  }
};

/// Base class for getter-like operations, e.g., to_indices, to_pointers.
template <typename SourceOp, typename Base>
class SparseGetterOpConverter : public OpConversionPattern<SourceOp> {
public:
  using OpAdaptor = typename SourceOp::Adaptor;
  using OpConversionPattern<SourceOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Replace the requested pointer access with corresponding field.
    // The cast_op is inserted by type converter to intermix 1:N type
    // conversion.
    auto tuple = getTuple(adaptor.getTensor());
    unsigned idx = Base::getIndexForOp(tuple, op);
    auto fields = tuple.getInputs();
    assert(idx < fields.size());
    rewriter.replaceOp(op, fields[idx]);
    return success();
  }
};

/// Sparse codegen rule for pointer accesses.
class SparseToPointersConverter
    : public SparseGetterOpConverter<ToPointersOp, SparseToPointersConverter> {
public:
  using SparseGetterOpConverter::SparseGetterOpConverter;
  // Callback for SparseGetterOpConverter.
  static unsigned getIndexForOp(UnrealizedConversionCastOp /*tuple*/,
                                ToPointersOp op) {
    uint64_t dim = op.getDimension().getZExtValue();
    return getFieldIndex(op.getTensor().getType(), /*ptrDim=*/dim, -1u);
  }
};

/// Sparse codegen rule for index accesses.
class SparseToIndicesConverter
    : public SparseGetterOpConverter<ToIndicesOp, SparseToIndicesConverter> {
public:
  using SparseGetterOpConverter::SparseGetterOpConverter;
  // Callback for SparseGetterOpConverter.
  static unsigned getIndexForOp(UnrealizedConversionCastOp /*tuple*/,
                                ToIndicesOp op) {
    uint64_t dim = op.getDimension().getZExtValue();
    return getFieldIndex(op.getTensor().getType(), -1u, /*idxDim=*/dim);
  }
};

/// Sparse codegen rule for value accesses.
class SparseToValuesConverter
    : public SparseGetterOpConverter<ToValuesOp, SparseToValuesConverter> {
public:
  using SparseGetterOpConverter::SparseGetterOpConverter;
  // Callback for SparseGetterOpConverter.
  static unsigned getIndexForOp(UnrealizedConversionCastOp tuple,
                                ToValuesOp /*op*/) {
    // The last field holds the value buffer.
    return tuple.getInputs().size() - 1;
  }
};

/// Sparse codegen rule for the convert operator.
class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType());
    SparseTensorEncodingAttr encSrc =
        getSparseTensorEncoding(op.getSource().getType());
    if (encDst != encSrc) {
      // This should be handled by rewriting before codegen.
      return failure();
    }
    rewriter.replaceOp(op, adaptor.getSource());
    return success();
  }
};

/// Sparse codegen rule for number of entries operator.
class SparseNumberOfEntriesConverter
    : public OpConversionPattern<NumberOfEntriesOp> {
public:
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Query memSizes for the actually stored values size.
    auto tuple = getTuple(adaptor.getTensor());
    auto fields = tuple.getInputs();
    unsigned lastField = fields.size() - 1;
    Value field =
        constantIndex(rewriter, op.getLoc(), getMemSizesIndex(lastField));
    rewriter.replaceOpWithNewOp<memref::LoadOp>(op, fields[MemSizesIdx], field);
    return success();
  }
};

} // namespace

//===----------------------------------------------------------------------===//
// Sparse tensor type conversion into an actual buffer.
//===----------------------------------------------------------------------===//

mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
  addConversion([](Type type) { return type; });
  addConversion(convertSparseTensorType);

  // Required by scf.for 1:N type conversion.
  addSourceMaterialization([](OpBuilder &builder, RankedTensorType tp,
                              ValueRange inputs,
                              Location loc) -> Optional<Value> {
    if (!getSparseTensorEncoding(tp))
      // Not a sparse tensor.
      return llvm::None;
    // Sparse compiler knows how to cancel out these casts.
    return genTuple(builder, loc, tp, inputs);
  });
}

//===----------------------------------------------------------------------===//
// Public method for populating conversion rules.
//===----------------------------------------------------------------------===//

/// Populates the given patterns list with conversion rules required for
/// the sparsification of linear algebra operations.
void mlir::populateSparseTensorCodegenPatterns(
    TypeConverter &typeConverter, RewritePatternSet &patterns,
    bool enableBufferInitialization) {
  patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
               SparseCastConverter, SparseTensorAllocConverter,
               SparseTensorDeallocConverter, SparseTensorLoadConverter,
               SparseExpandConverter, SparseCompressConverter,
               SparseInsertConverter, SparseToPointersConverter,
               SparseToIndicesConverter, SparseToValuesConverter,
               SparseConvertConverter, SparseNumberOfEntriesConverter>(
      typeConverter, patterns.getContext());
  patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                           enableBufferInitialization);
}