SparseTensorIterator.cpp | Explore in Territory

//===- SparseTensorIterator.cpp -------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "SparseTensorIterator.h"
#include "CodegenUtils.h"

#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

usingnamespacemlir;
usingnamespacemlir::sparse_tensor;
ValuePair;
ValueTuple;

//===----------------------------------------------------------------------===//
// File local helper functions/macros.
//===----------------------------------------------------------------------===//
#define CMPI …

#define C_FALSE …
#define C_TRUE …
#define C_IDX …
#define YIELD …
#define ADDI …
#define ORI(lhs, rhs) …
#define ANDI …
#define SUBI …
#define MULI …
#define MINUI(lhs, rhs) …
#define REMUI(lhs, rhs) …
#define DIVUI(lhs, rhs) …
#define SELECT …

//===----------------------------------------------------------------------===//
// SparseTensorLevel derived classes.
//===----------------------------------------------------------------------===//

namespace {

template <bool hasPosBuffer>
class SparseLevel : public SparseTensorLevel { … };

class DenseLevel : public SparseTensorLevel { … };

class BatchLevel : public SparseTensorLevel { … };

class CompressedLevel : public SparseLevel</*hasPosBuf=*/true> { … }; // namespace

class LooseCompressedLevel : public SparseLevel</*hasPosBuf=*/true> { … }; // namespace

class SingletonLevel : public SparseLevel</*hasPosBuf=*/false> { … };

class NOutOfMLevel : public SparseLevel</*hasPosBuf=*/false> { … };

} // namespace

//===----------------------------------------------------------------------===//
// File local helpers
//===----------------------------------------------------------------------===//

static scf::ValueVector genWhenInBound(
    OpBuilder &b, Location l, SparseIterator &it, ValueRange elseRet,
    llvm::function_ref<scf::ValueVector(OpBuilder &, Location, Value)>
        builder) { … }

/// Generates code to compute the *absolute* offset of the slice based on the
/// provide minimum coordinates in the slice.
/// E.g., when reducing d0 + d1 + d2, we need two slices to fully reduced the
/// expression, i,e, s1 = slice(T, d0), s2 = slice(s1, d1). The *absolute*
/// offset is the offset computed relative to the initial tensors T.
///
/// When isNonEmpty == true, the computed offset is meaningless and should not
/// be used during runtime, the method generates code to return 0 currently in
/// that case.
///
/// offset = minCrd >= size ? minCrd - size + 1 : 0;
static Value offsetFromMinCrd(OpBuilder &b, Location l, Value minCrd,
                              Value size) { … }

//===----------------------------------------------------------------------===//
// SparseIterator derived classes.
//===----------------------------------------------------------------------===//

namespace {

// The iterator that traverses a concrete sparse tensor levels. High-level
// abstract iterators wrap it to achieve more complex goals (such as collapsing
// several levels). It also holds the common storage to hold the mlir::Values
// for itself as well as for wrappers.
class ConcreteIterator : public SparseIterator { … };

class TrivialIterator : public ConcreteIterator { … };

class DedupIterator : public ConcreteIterator { … };

// A util base-iterator that delegates all methods to the wrapped iterator.
class SimpleWrapIterator : public SparseIterator { … };

//
// A filter iterator wrapped from another iterator. The filter iterator update
// the wrapped iterator *in-place*.
//
class FilterIterator : public SimpleWrapIterator { … };

//
// A pad iterator wrapped from another iterator. The pad iterator updates
// the wrapped iterator *in-place*.
//
class PadIterator : public SimpleWrapIterator { … };

class NonEmptySubSectIterator : public SparseIterator { … };

class SubSectIterator;

// A wrapper that helps generating code to traverse a subsection, used
// by both `NonEmptySubSectIterator`and `SubSectIterator`.
struct SubSectIterHelper { … };

class SubSectIterator : public SparseIterator { … };

} // namespace

//===----------------------------------------------------------------------===//
// SparseIterator derived classes implementation.
//===----------------------------------------------------------------------===//

void SparseIterator::genInit(OpBuilder &b, Location l,
                             const SparseIterator *p) { … }

Value SparseIterator::genNotEnd(OpBuilder &b, Location l) { … }

void SparseIterator::locate(OpBuilder &b, Location l, Value crd) { … }

Value SparseIterator::deref(OpBuilder &b, Location l) { … }

ValueRange SparseIterator::forward(OpBuilder &b, Location l) { … }

ValueRange SparseIterator::forwardIf(OpBuilder &b, Location l, Value cond) { … }

Value DedupIterator::genSegmentHigh(OpBuilder &b, Location l, Value pos) { … }

Value FilterIterator::genCrdNotLegitPredicate(OpBuilder &b, Location l,
                                              Value wrapCrd) { … }

Value FilterIterator::genShouldFilter(OpBuilder &b, Location l) { … }

Value FilterIterator::genNotEndImpl(OpBuilder &b, Location l) { … }

ValueRange FilterIterator::forwardImpl(OpBuilder &b, Location l) { … }

SubSectIterHelper::SubSectIterHelper(const NonEmptySubSectIterator &subSect)
    : … { … }

SubSectIterHelper::SubSectIterHelper(const SubSectIterator &iter)
    : … { … }

void SubSectIterHelper::deserializeFromTupleId(OpBuilder &b, Location l,
                                               Value tupleId) { … }

void SubSectIterHelper::locate(OpBuilder &b, Location l, Value crd) { … }

Value SubSectIterHelper::genNotEnd(OpBuilder &b, Location l) { … }

Value SubSectIterHelper::deref(OpBuilder &b, Location l) { … }

ValueRange SubSectIterHelper::forward(OpBuilder &b, Location l) { … }

ValueRange NonEmptySubSectIterator::inflateSubSectTree(
    OpBuilder &b, Location l, ValueRange reduc, TraverseBuilder builder) const { … }

void TrivialIterator::genInitImpl(OpBuilder &b, Location l,
                                  const SparseIterator *parent) { … }

void NonEmptySubSectIterator::genInitImpl(OpBuilder &b, Location l,
                                          const SparseIterator *) { … }

ValueRange NonEmptySubSectIterator::forwardImpl(OpBuilder &b, Location l) { … }

//===----------------------------------------------------------------------===//
// SparseIterationSpace Implementation
//===----------------------------------------------------------------------===//

mlir::sparse_tensor::SparseIterationSpace::SparseIterationSpace(
    Location l, OpBuilder &b, Value t, unsigned tid,
    std::pair<Level, Level> lvlRange, ValueRange parentPos)
    : … { … }

SparseIterationSpace mlir::sparse_tensor::SparseIterationSpace::fromValues(
    IterSpaceType dstTp, ValueRange values, unsigned int tid) { … }

std::unique_ptr<SparseIterator>
SparseIterationSpace::extractIterator(OpBuilder &b, Location l) const { … }

//===----------------------------------------------------------------------===//
// SparseIterator factory functions.
//===----------------------------------------------------------------------===//

/// Helper function to create a TensorLevel object from given `tensor`.
std::unique_ptr<SparseTensorLevel>
sparse_tensor::makeSparseTensorLevel(LevelType lt, Value sz, ValueRange b,
                                     unsigned t, Level l) { … }

std::unique_ptr<SparseTensorLevel>
sparse_tensor::makeSparseTensorLevel(OpBuilder &b, Location l, Value t,
                                     unsigned tid, Level lvl) { … }

std::pair<std::unique_ptr<SparseTensorLevel>, std::unique_ptr<SparseIterator>>
sparse_tensor::makeSynLevelAndIterator(Value sz, unsigned tid, unsigned lvl,
                                       SparseEmitStrategy strategy) { … }

std::unique_ptr<SparseIterator>
sparse_tensor::makeSimpleIterator(OpBuilder &b, Location l,
                                  const SparseIterationSpace &iterSpace) { … }

std::unique_ptr<SparseIterator>
sparse_tensor::makeSimpleIterator(const SparseTensorLevel &stl,
                                  SparseEmitStrategy strategy) { … }

std::unique_ptr<SparseIterator>
sparse_tensor::makeSlicedLevelIterator(std::unique_ptr<SparseIterator> &&sit,
                                       Value offset, Value stride, Value size,
                                       SparseEmitStrategy strategy) { … }

std::unique_ptr<SparseIterator>
sparse_tensor::makePaddedIterator(std::unique_ptr<SparseIterator> &&sit,
                                  Value padLow, Value padHigh,
                                  SparseEmitStrategy strategy) { … }

static const SparseIterator *tryUnwrapFilter(const SparseIterator *it) { … }

std::unique_ptr<SparseIterator> sparse_tensor::makeNonEmptySubSectIterator(
    OpBuilder &b, Location l, const SparseIterator *parent, Value loopBound,
    std::unique_ptr<SparseIterator> &&delegate, Value size, unsigned stride,
    SparseEmitStrategy strategy) { … }

std::unique_ptr<SparseIterator> sparse_tensor::makeTraverseSubSectIterator(
    OpBuilder &b, Location l, const SparseIterator &subSectIter,
    const SparseIterator &parent, std::unique_ptr<SparseIterator> &&wrap,
    Value loopBound, unsigned stride, SparseEmitStrategy strategy) { … }

#undef CMPI
#undef C_IDX
#undef YIELD
#undef ADDI
#undef ANDI
#undef SUBI
#undef MULI
#undef SELECT
llvm/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
