SparseTensor.h

Go to the documentation of this file.

//===- SparseTensor.h - Sparse tensor dialect -------------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
#define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
 
#include "mlir/Dialect/SparseTensor/IR/Enums.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/TensorEncoding.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
 
// We must include Enums.h.inc before AttrDefs.h.inc due to dependency between
// StorageSpecifierKindAttr and StorageSpeciferKind Enum.
 
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.h.inc"
 
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"
 
#define GET_TYPEDEF_CLASSES
#include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.h.inc"
 
#define GET_OP_CLASSES
#include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
 
#include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
 
namespace mlir {
namespace sparse_tensor {
 
/// Convenience method to abbreviate casting `getType()`.
template <typename T>
inline RankedTensorType getRankedTensorType(T t) {
  return t.getType().template cast<RankedTensorType>();
}
 
/// Convenience method to abbreviate casting `getType()`.
template <typename T>
inline MemRefType getMemRefType(T t) {
  return t.getType().template cast<MemRefType>();
}
 
/// Convenience method to get a sparse encoding attribute from a type.
/// Returns null-attribute for any type without an encoding.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
/// Returns true iff the given type is a type for a COO tensor with the last
/// dimension level type being unique.
bool isUniqueCOOType(RankedTensorType tp);
 
/// Returns the starting dimension for a trailing COO region that spans across
/// at least two dimensions. If no such COO region is found, returns the rank
/// of the tensor.
unsigned getCOOStart(SparseTensorEncodingAttr enc);
 
//
// Dimension level types.
//
 
// MSVC does not allow this function to be constexpr, because
// `SparseTensorEncodingAttr::operator bool` isn't declared constexpr.
// And therefore all functions calling it cannot be constexpr either.
// TODO: since Clang does allow these to be constexpr, perhaps we should
// define a macro to abstract over `inline` vs `constexpr` annotations.
inline DimLevelType getDimLevelType(SparseTensorEncodingAttr enc, uint64_t d) {
  if (enc) {
    auto types = enc.getDimLevelType();
    assert(d < types.size() && "Dimension out of bounds");
    return types[d];
  }
  return DimLevelType::Dense; // unannotated tensor is dense
}
 
inline DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {
  return getDimLevelType(getSparseTensorEncoding(type), d);
}
 
/// Convenience function to test for dense dimension (0 <= d < rank).
inline bool isDenseDim(RankedTensorType type, uint64_t d) {
  return isDenseDLT(getDimLevelType(type, d));
}
 
/// Convenience function to test for compressed dimension (0 <= d < rank).
inline bool isCompressedDim(RankedTensorType type, uint64_t d) {
  return isCompressedDLT(getDimLevelType(type, d));
}
 
/// Convenience function to test for singleton dimension (0 <= d < rank).
inline bool isSingletonDim(RankedTensorType type, uint64_t d) {
  return isSingletonDLT(getDimLevelType(type, d));
}
 
/// Convenience function to test for dense dimension (0 <= d < rank).
inline bool isDenseDim(SparseTensorEncodingAttr enc, uint64_t d) {
  return isDenseDLT(getDimLevelType(enc, d));
}
 
/// Convenience function to test for compressed dimension (0 <= d < rank).
inline bool isCompressedDim(SparseTensorEncodingAttr enc, uint64_t d) {
  return isCompressedDLT(getDimLevelType(enc, d));
}
 
/// Convenience function to test for singleton dimension (0 <= d < rank).
inline bool isSingletonDim(SparseTensorEncodingAttr enc, uint64_t d) {
  return isSingletonDLT(getDimLevelType(enc, d));
}
 
//
// Dimension level properties.
//
 
/// Convenience function to test for ordered property in the
/// given dimension (0 <= d < rank).
inline bool isOrderedDim(RankedTensorType type, uint64_t d) {
  return isOrderedDLT(getDimLevelType(type, d));
}
 
/// Convenience function to test for unique property in the
/// given dimension (0 <= d < rank).
inline bool isUniqueDim(RankedTensorType type, uint64_t d) {
  return isUniqueDLT(getDimLevelType(type, d));
}
 
//
// Reordering.
//
 
uint64_t toOrigDim(SparseTensorEncodingAttr enc, uint64_t d);
uint64_t toStoredDim(SparseTensorEncodingAttr enc, uint64_t d);
 
/// Convenience method to translate the given stored dimension
/// to the original dimension (0 <= d < rank).
uint64_t toOrigDim(RankedTensorType type, uint64_t d);
 
/// Convenience method to translate the given original dimension
/// to the stored dimension (0 <= d < rank).
uint64_t toStoredDim(RankedTensorType type, uint64_t d);
 
} // namespace sparse_tensor
} // namespace mlir
 
#endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_