llvm/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp

#include "Utils/CodegenUtils.h"
#include "Utils/LoopEmitter.h"
#include "Utils/SparseTensorIterator.h"

#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Transforms/OneToNTypeConversion.h"

usingnamespacemlir;
usingnamespacemlir::sparse_tensor;

static void convertLevelType(SparseTensorEncodingAttr enc, Level lvl,
                             SmallVectorImpl<Type> &fields) { … }

static std::optional<LogicalResult>
convertIterSpaceType(IterSpaceType itSp, SmallVectorImpl<Type> &fields) { … }

static std::optional<LogicalResult>
convertIteratorType(IteratorType itTp, SmallVectorImpl<Type> &fields) { … }

static ValueRange
genCoIterateBranchNest(PatternRewriter &rewriter, Location loc, CoIterateOp op,
                       Value loopCrd,
                       ArrayRef<std::unique_ptr<SparseIterator>> iters,
                       ArrayRef<Region *> subCases, ArrayRef<Value> userReduc) { … }

static ValueRange genLoopWithIterator(
    PatternRewriter &rewriter, Location loc, SparseIterator *it,
    ValueRange reduc,
    function_ref<SmallVector<Value>(PatternRewriter &rewriter, Location loc,
                                    Region &loopBody, SparseIterator *it,
                                    ValueRange reduc)>
        bodyBuilder) { … }

namespace {

/// Sparse codegen rule for number of entries operator.
class ExtractIterSpaceConverter
    : public OneToNOpConversionPattern<ExtractIterSpaceOp> { … };

/// Sparse codegen rule for number of entries operator.
class ExtractValOpConverter : public OneToNOpConversionPattern<ExtractValOp> { … };

class SparseIterateOpConverter : public OneToNOpConversionPattern<IterateOp> { … };

class SparseCoIterateOpConverter
    : public OneToNOpConversionPattern<CoIterateOp> { … };

} // namespace

mlir::SparseIterationTypeConverter::SparseIterationTypeConverter() { … }

void mlir::populateLowerSparseIterationToSCFPatterns(
    const TypeConverter &converter, RewritePatternSet &patterns) { … }