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

//===- CodegenEnv.cpp -  Code generation environment class ----------------===//
//
// 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 "CodegenEnv.h"

#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include <optional>

usingnamespacemlir;
usingnamespacemlir::sparse_tensor;

//===----------------------------------------------------------------------===//
// Code generation environment helper functions
//===----------------------------------------------------------------------===//

/// Returns true if tensor materializes uninitialized into the computation.
static bool isMaterializing(Value val) { … }

/// Sorts the dependent loops such that it is ordered in the same sequence in
/// which loops will be generated.
static void sortDependentLoops(std::vector<LoopCoeffPair> &target) { … }
//===----------------------------------------------------------------------===//
// Code generation environment constructor and general methods
//===----------------------------------------------------------------------===//

CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                       unsigned numTensors, unsigned numLoops, unsigned maxRank)
    : … { … }

LogicalResult CodegenEnv::initTensorExp() { … }

void CodegenEnv::startEmit(SparseEmitStrategy emitStrategy) { … }

std::optional<Operation *> CodegenEnv::genLoopBoundary(
    function_ref<std::optional<Operation *>(MutableArrayRef<Value> parameters)>
        callback) { … }

//===----------------------------------------------------------------------===//
// Code generation environment verify functions.
//===----------------------------------------------------------------------===//

bool CodegenEnv::isAdmissibleTensorExp(ExprId exp) { … }

//===----------------------------------------------------------------------===//
// Code generation environment topological sort methods
//===----------------------------------------------------------------------===//

Value CodegenEnv::getLoopVar(LoopId i) const { … }

//===----------------------------------------------------------------------===//
// Code generation environment sparse tensor output and expansion methods
//===----------------------------------------------------------------------===//

void CodegenEnv::updateInsertionChain(Value chain) { … }

bool CodegenEnv::atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const { … }

void CodegenEnv::startExpand(Value values, Value filled, Value added,
                             Value count) { … }

void CodegenEnv::updateExpandCount(Value count) { … }

void CodegenEnv::endExpand() { … }

//===----------------------------------------------------------------------===//
// Code generation environment reduction methods
//===----------------------------------------------------------------------===//

void CodegenEnv::startReduc(ExprId exp, Value val) { … }

void CodegenEnv::updateReduc(Value val) { … }

Value CodegenEnv::endReduc() { … }

void CodegenEnv::startValidLexInsert(Value val) { … }

void CodegenEnv::updateValidLexInsert(Value val) { … }

void CodegenEnv::endValidLexInsert() { … }

void CodegenEnv::startCustomReduc(ExprId exp) { … }

Value CodegenEnv::getCustomRedId() const { … }

void CodegenEnv::endCustomReduc() { … }