CodegenUtils.cpp | Explore in Territory

//===- CodegenUtils.cpp - Utilities for generating MLIR -------------------===//
//
// 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 "CodegenUtils.h"
#include "SparseTensorDescriptor.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include <optional>

usingnamespacemlir;
usingnamespacemlir::sparse_tensor;

//===----------------------------------------------------------------------===//
// ExecutionEngine/SparseTensorUtils helper functions.
//===----------------------------------------------------------------------===//

OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) { … }

OverheadType mlir::sparse_tensor::overheadTypeEncoding(Type tp) { … }

Type mlir::sparse_tensor::getOverheadType(Builder &builder, OverheadType ot) { … }

OverheadType
mlir::sparse_tensor::posTypeEncoding(SparseTensorEncodingAttr enc) { … }

OverheadType
mlir::sparse_tensor::crdTypeEncoding(SparseTensorEncodingAttr enc) { … }

// TODO: we ought to add some `static_assert` tests to ensure that the
// `STEA::get{Pos,Crd}Type` methods agree with `getOverheadType(builder,
// {pos,crd}OverheadTypeEncoding(enc))`

// TODO: Adjust the naming convention for the constructors of
// `OverheadType` so we can use the `MLIR_SPARSETENSOR_FOREVERY_O` x-macro
// here instead of `MLIR_SPARSETENSOR_FOREVERY_FIXED_O`; to further reduce
// the possibility of typo bugs or things getting out of sync.
StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(OverheadType ot) { … }

StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(Type tp) { … }

PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) { … }

StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(PrimaryType pt) { … }

StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(Type elemTp) { … }

//===----------------------------------------------------------------------===//
// Misc code generators.
//===----------------------------------------------------------------------===//

Value sparse_tensor::genCast(OpBuilder &builder, Location loc, Value value,
                             Type dstTp) { … }

Value sparse_tensor::genScalarToTensor(OpBuilder &builder, Location loc,
                                       Value elem, Type dstTp) { … }

Value sparse_tensor::genIndexLoad(OpBuilder &builder, Location loc, Value mem,
                                  ValueRange s) { … }

mlir::TypedAttr mlir::sparse_tensor::getOneAttr(Builder &builder, Type tp) { … }

Value mlir::sparse_tensor::genIsNonzero(OpBuilder &builder, mlir::Location loc,
                                        Value v) { … }

void mlir::sparse_tensor::genReshapeDstShape(
    OpBuilder &builder, Location loc, SmallVectorImpl<Value> &dstShape,
    ArrayRef<Value> srcShape, ArrayRef<Size> staticDstShape,
    ArrayRef<ReassociationIndices> reassociation) { … }

void mlir::sparse_tensor::reshapeCvs(
    OpBuilder &builder, Location loc,
    ArrayRef<ReassociationIndices> reassociation, // NOLINT
    ValueRange srcSizes, ValueRange srcCvs,       // NOLINT
    ValueRange dstSizes, SmallVectorImpl<Value> &dstCvs) { … }

FlatSymbolRefAttr mlir::sparse_tensor::getFunc(ModuleOp module, StringRef name,
                                               TypeRange resultType,
                                               ValueRange operands,
                                               EmitCInterface emitCInterface) { … }

func::CallOp mlir::sparse_tensor::createFuncCall(
    OpBuilder &builder, Location loc, StringRef name, TypeRange resultType,
    ValueRange operands, EmitCInterface emitCInterface) { … }

Type mlir::sparse_tensor::getOpaquePointerType(MLIRContext *ctx) { … }

Type mlir::sparse_tensor::getOpaquePointerType(Builder &builder) { … }

Value mlir::sparse_tensor::genAlloca(OpBuilder &builder, Location loc,
                                     unsigned sz, Type tp, bool staticShape) { … }

Value mlir::sparse_tensor::genAlloca(OpBuilder &builder, Location loc, Value sz,
                                     Type tp) { … }

Value mlir::sparse_tensor::genAllocaScalar(OpBuilder &builder, Location loc,
                                           Type tp) { … }

Value mlir::sparse_tensor::allocaBuffer(OpBuilder &builder, Location loc,
                                        ValueRange values) { … }

Value mlir::sparse_tensor::allocDenseTensor(OpBuilder &builder, Location loc,
                                            RankedTensorType tensorTp,
                                            ValueRange sizes) { … }

void mlir::sparse_tensor::deallocDenseTensor(OpBuilder &builder, Location loc,
                                             Value buffer) { … }

void mlir::sparse_tensor::sizesFromSrc(OpBuilder &builder,
                                       SmallVectorImpl<Value> &sizes,
                                       Location loc, Value src) { … }

Operation *mlir::sparse_tensor::getTop(Operation *op) { … }

void sparse_tensor::foreachInSparseConstant(
    OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order,
    function_ref<void(ArrayRef<Value>, Value)> callback) { … }

SmallVector<Value> sparse_tensor::loadAll(OpBuilder &builder, Location loc,
                                          size_t size, Value mem,
                                          size_t offsetIdx, Value offsetVal) { … }

void sparse_tensor::storeAll(OpBuilder &builder, Location loc, Value mem,
                             ValueRange vs, size_t offsetIdx, Value offsetVal) { … }

TypedValue<BaseMemRefType>
sparse_tensor::genToMemref(OpBuilder &builder, Location loc, Value tensor) { … }

Value sparse_tensor::genValMemSize(OpBuilder &builder, Location loc,
                                   Value tensor) { … }

Value sparse_tensor::createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc,
                                               Value tensor, Dimension dim) { … }

Value sparse_tensor::createOrFoldSliceStrideOp(OpBuilder &builder, Location loc,
                                               Value tensor, Dimension dim) { … }

Value sparse_tensor::genReader(OpBuilder &builder, Location loc,
                               SparseTensorType stt, Value tensor,
                               /*out*/ SmallVectorImpl<Value> &dimSizesValues,
                               /*out*/ Value &dimSizesBuffer) { … }

Value sparse_tensor::genMapBuffers(
    OpBuilder &builder, Location loc, SparseTensorType stt,
    ArrayRef<Value> dimSizesValues, Value dimSizesBuffer,
    /*out*/ SmallVectorImpl<Value> &lvlSizesValues,
    /*out*/ Value &dim2lvlBuffer,
    /*out*/ Value &lvl2dimBuffer) { … }
llvm/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
