SparseTensorRewriting.cpp File Reference

#include "Utils/CodegenUtils.h"
#include "Utils/LoopEmitter.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.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/SCF/IR/SCF.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Support/LLVM.h"

Go to the source code of this file.

Functions
static bool	isSparseTensor (Value v)
static bool	isSparseTensor (OpOperand *op)
static bool	isMaterializing (OpOperand *op, bool isZero)
static bool	isSampling (GenericOp op)
static bool	isMulChain (Value val, Value x)
static bool	isSumOfMul (GenericOp op)
static bool	isZeroYield (GenericOp op)
static void	sizesForTensor (OpBuilder &builder, SmallVectorImpl< Value > &sizes, Location loc, ShapedType stp, Value tensor)
	Populates given sizes array from type (for static sizes) and from the tensor (for dynamic sizes).
static RankedTensorType	getBufferType (const SparseTensorType &stt, bool needTmpCOO)
static void	getDynamicSizes (RankedTensorType tp, ValueRange sizes, SmallVectorImpl< Value > &dynSizes)
	Collects the dynamic dimension sizes for tp with the assumption that sizes are the dimension sizes for the type.
static LogicalResult	genForeachOnSparseConstant (ForeachOp op, RewriterBase &rewriter, SparseElementsAttr attr)
static void	concatSizesFromInputs (OpBuilder &builder, SmallVectorImpl< Value > &sizes, Location loc, ShapedType dstTp, ValueRange srcs, unsigned dim)
	Populates the given sizes array for concatenation from types (for static sizes) and from the source tensors (for dynamic sizes).

Function Documentation

concatSizesFromInputs()

void concatSizesFromInputs ( OpBuilder & builder, SmallVectorImpl< Value > & sizes, Location loc, ShapedType dstTp, ValueRange srcs, unsigned dim )

static

Populates the given sizes array for concatenation from types (for static sizes) and from the source tensors (for dynamic sizes).

Definition at line 180 of file SparseTensorRewriting.cpp.

References mlir::sparse_tensor::constantIndex(), mlir::linalg::createOrFoldDimOp(), and mlir::sparse_tensor::sizesFromSrc().

genForeachOnSparseConstant()

LogicalResult genForeachOnSparseConstant ( ForeachOp op, RewriterBase & rewriter, SparseElementsAttr attr )

static

Definition at line 149 of file SparseTensorRewriting.cpp.

References mlir::OpBuilder::clone(), mlir::RewriterBase::eraseOp(), mlir::sparse_tensor::foreachInSparseConstant(), mlir::Operation::getOperands(), mlir::RewriterBase::inlineBlockBefore(), mlir::RewriterBase::replaceOp(), and success().

getBufferType()

RankedTensorType getBufferType ( const SparseTensorType & stt, bool needTmpCOO )

static

Definition at line 132 of file SparseTensorRewriting.cpp.

References mlir::sparse_tensor::SparseTensorType::getCOOType(), and mlir::sparse_tensor::SparseTensorType::getRankedTensorType().

getDynamicSizes()

void getDynamicSizes ( RankedTensorType tp, ValueRange sizes, SmallVectorImpl< Value > & dynSizes )

static

Collects the dynamic dimension sizes for tp with the assumption that sizes are the dimension sizes for the type.

Stores the dynamic dimension sizes to dynSizes.

Definition at line 141 of file SparseTensorRewriting.cpp.

isMaterializing()

bool isMaterializing ( OpOperand * op, bool isZero )

static

Definition at line 51 of file SparseTensorRewriting.cpp.

References copy(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Value::getDefiningOp(), and mlir::isZeroIntegerOrFloat().

Referenced by mlir::sparse_tensor::CodegenEnv::isAdmissibleTensorExp().

isMulChain()

bool isMulChain ( Value val, Value x )

static

Definition at line 83 of file SparseTensorRewriting.cpp.

References mlir::Value::getDefiningOp(), and isMulChain().

Referenced by isMulChain(), and isSumOfMul().

isSampling()

bool isSampling ( GenericOp op )

static

Definition at line 68 of file SparseTensorRewriting.cpp.

isSparseTensor() [1/2]

bool isSparseTensor ( OpOperand * op )

static

Definition at line 48 of file SparseTensorRewriting.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), and isSparseTensor().

isSparseTensor() [2/2]

bool isSparseTensor ( Value v )

static

Definition at line 43 of file SparseTensorRewriting.cpp.

References mlir::sparse_tensor::getSparseTensorEncoding(), and mlir::Value::getType().

Referenced by isSparseTensor().

isSumOfMul()

bool isSumOfMul ( GenericOp op )

static

Definition at line 95 of file SparseTensorRewriting.cpp.

References isMulChain().

isZeroYield()

bool isZeroYield ( GenericOp op )

static

Definition at line 108 of file SparseTensorRewriting.cpp.

References mlir::isZeroIntegerOrFloat().

sizesForTensor()

void sizesForTensor ( OpBuilder & builder, SmallVectorImpl< Value > & sizes, Location loc, ShapedType stp, Value tensor )

static

Populates given sizes array from type (for static sizes) and from the tensor (for dynamic sizes).

Definition at line 120 of file SparseTensorRewriting.cpp.

References mlir::sparse_tensor::constantIndex().