llvm/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp

//===- IterationGraphSorter.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 "IterationGraphSorter.h"

#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/BuiltinTypes.h"

usingnamespacemlir;
usingnamespacemlir::sparse_tensor;

namespace {

/// A helper class that visits an affine expression and tries to find
/// an AffineDimExpr to which the corresponding iterator from a GenericOp
/// matches the desired iterator type. If there is no matched iterator
/// type, the method returns the first DimExpr in the expression.
class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> { … };

/// Flattens an affine expression into a list of AffineDimExprs.
struct AffineDimCollector : public AffineExprVisitor<AffineDimCollector> { … };

} // namespace

inline static bool includesAny(SortMask mask1, SortMask mask2) { … }

inline static bool includesDenseInput(SortMask mask) { … }

inline static bool includesDenseOutput(SortMask mask) { … }

AffineMap IterationGraphSorter::topoSort() { … }

IterationGraphSorter
IterationGraphSorter::fromGenericOp(linalg::GenericOp genericOp) { … }

IterationGraphSorter::IterationGraphSorter(
    SmallVector<Value> &&ins, SmallVector<AffineMap> &&loop2InsLvl, Value out,
    AffineMap loop2OutLvl, SmallVector<utils::IteratorType> &&iterTypes)
    : … { … }

AffineMap IterationGraphSorter::sort(SortMask mask, Value ignored) { … }

void IterationGraphSorter::addConstraints(Value t, AffineMap loop2LvlMap) { … }