StandardTypes.h

Go to the documentation of this file.

//===- StandardTypes.h - MLIR Standard Type Classes -------------*- C++ -*-===//
//
// Part of the MLIR 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_IR_STANDARDTYPES_H
#define MLIR_IR_STANDARDTYPES_H

#include "mlir/IR/Types.h"

namespace llvm {
struct fltSemantics;
} // namespace llvm

namespace mlir {
class AffineExpr;
class AffineMap;
class FloatType;
class IndexType;
class IntegerType;
class Location;
class MLIRContext;

namespace detail {

struct IntegerTypeStorage;
struct ShapedTypeStorage;
struct VectorTypeStorage;
struct RankedTensorTypeStorage;
struct UnrankedTensorTypeStorage;
struct MemRefTypeStorage;
struct UnrankedMemRefTypeStorage;
struct ComplexTypeStorage;
struct TupleTypeStorage;

} // namespace detail

namespace StandardTypes {
enum Kind {
  // Floating point.
  BF16 = Type::Kind::FIRST_STANDARD_TYPE,
  F16,
  F32,
  F64,
  FIRST_FLOATING_POINT_TYPE = BF16,
  LAST_FLOATING_POINT_TYPE = F64,

  // Target pointer sized integer, used (e.g.) in affine mappings.
  Index,

  // Derived types.
  Integer,
  Vector,
  RankedTensor,
  UnrankedTensor,
  MemRef,
  UnrankedMemRef,
  Complex,
  Tuple,
  None,
};

} // namespace StandardTypes

class IndexType : public Type::TypeBase<IndexType, Type> {
public:
  using Base::Base;

  static IndexType get(MLIRContext *context);

  static bool kindof(unsigned kind) { return kind == StandardTypes::Index; }
};

class IntegerType
    : public Type::TypeBase<IntegerType, Type, detail::IntegerTypeStorage> {
public:
  using Base::Base;

  static IntegerType get(unsigned width, MLIRContext *context);

  static IntegerType getChecked(unsigned width, MLIRContext *context,
                                Location location);

  static LogicalResult verifyConstructionInvariants(Optional<Location> loc,
                                                    MLIRContext *context,
                                                    unsigned width);

  unsigned getWidth() const;

  static bool kindof(unsigned kind) { return kind == StandardTypes::Integer; }

  static constexpr unsigned kMaxWidth = 4096;
};

class FloatType : public Type::TypeBase<FloatType, Type> {
public:
  using Base::Base;

  static FloatType get(StandardTypes::Kind kind, MLIRContext *context);

  // Convenience factories.
  static FloatType getBF16(MLIRContext *ctx) {
    return get(StandardTypes::BF16, ctx);
  }
  static FloatType getF16(MLIRContext *ctx) {
    return get(StandardTypes::F16, ctx);
  }
  static FloatType getF32(MLIRContext *ctx) {
    return get(StandardTypes::F32, ctx);
  }
  static FloatType getF64(MLIRContext *ctx) {
    return get(StandardTypes::F64, ctx);
  }

  static bool kindof(unsigned kind) {
    return kind >= StandardTypes::FIRST_FLOATING_POINT_TYPE &&
           kind <= StandardTypes::LAST_FLOATING_POINT_TYPE;
  }

  unsigned getWidth();

  const llvm::fltSemantics &getFloatSemantics();
};

class ComplexType
    : public Type::TypeBase<ComplexType, Type, detail::ComplexTypeStorage> {
public:
  using Base::Base;

  static ComplexType get(Type elementType);

  static ComplexType getChecked(Type elementType, Location location);

  static LogicalResult verifyConstructionInvariants(Optional<Location> loc,
                                                    MLIRContext *context,
                                                    Type elementType);

  Type getElementType();

  static bool kindof(unsigned kind) { return kind == StandardTypes::Complex; }
};

class ShapedType : public Type {
public:
  using ImplType = detail::ShapedTypeStorage;
  using Type::Type;

  // TODO(ntv): merge these two special values in a single one used everywhere.
  // Unfortunately, uses of `-1` have crept deep into the codebase now and are
  // hard to track.
  static constexpr int64_t kDynamicSize = -1;
  static constexpr int64_t kDynamicStrideOrOffset =
      std::numeric_limits<int64_t>::min();

  Type getElementType() const;

  unsigned getElementTypeBitWidth() const;

  int64_t getNumElements() const;

  int64_t getRank() const;

  bool hasRank() const;

  ArrayRef<int64_t> getShape() const;

  bool hasStaticShape() const;

  bool hasStaticShape(ArrayRef<int64_t> shape) const;

  int64_t getNumDynamicDims() const;

  int64_t getDimSize(int64_t i) const;

  unsigned getDynamicDimIndex(unsigned index) const;

  int64_t getSizeInBits() const;

  static bool classof(Type type) {
    return type.getKind() == StandardTypes::Vector ||
           type.getKind() == StandardTypes::RankedTensor ||
           type.getKind() == StandardTypes::UnrankedTensor ||
           type.getKind() == StandardTypes::UnrankedMemRef ||
           type.getKind() == StandardTypes::MemRef;
  }

  static constexpr bool isDynamic(int64_t dSize) { return dSize < 0; }
  static constexpr bool isDynamicStrideOrOffset(int64_t dStrideOrOffset) {
    return dStrideOrOffset == kDynamicStrideOrOffset;
  }
};

class VectorType
    : public Type::TypeBase<VectorType, ShapedType, detail::VectorTypeStorage> {
public:
  using Base::Base;

  static VectorType get(ArrayRef<int64_t> shape, Type elementType);

  static VectorType getChecked(ArrayRef<int64_t> shape, Type elementType,
                               Location location);

  static LogicalResult verifyConstructionInvariants(Optional<Location> loc,
                                                    MLIRContext *context,
                                                    ArrayRef<int64_t> shape,
                                                    Type elementType);

  static bool isValidElementType(Type t) { return t.isIntOrFloat(); }

  ArrayRef<int64_t> getShape() const;

  static bool kindof(unsigned kind) { return kind == StandardTypes::Vector; }
};

class TensorType : public ShapedType {
public:
  using ShapedType::ShapedType;

  static bool isValidElementType(Type type) {
    // Note: Non standard/builtin types are allowed to exist within tensor
    // types. Dialects are expected to verify that tensor types have a valid
    // element type within that dialect.
    return type.isIntOrFloat() || type.isa<ComplexType>() ||
           type.isa<VectorType>() || type.isa<OpaqueType>() ||
           (type.getKind() > Type::Kind::LAST_STANDARD_TYPE);
  }

  static bool classof(Type type) {
    return type.getKind() == StandardTypes::RankedTensor ||
           type.getKind() == StandardTypes::UnrankedTensor;
  }
};

class RankedTensorType
    : public Type::TypeBase<RankedTensorType, TensorType,
                            detail::RankedTensorTypeStorage> {
public:
  using Base::Base;

  static RankedTensorType get(ArrayRef<int64_t> shape, Type elementType);

  static RankedTensorType getChecked(ArrayRef<int64_t> shape, Type elementType,
                                     Location location);

  static LogicalResult verifyConstructionInvariants(Optional<Location> loc,
                                                    MLIRContext *context,
                                                    ArrayRef<int64_t> shape,
                                                    Type elementType);

  ArrayRef<int64_t> getShape() const;

  static bool kindof(unsigned kind) {
    return kind == StandardTypes::RankedTensor;
  }
};

class UnrankedTensorType
    : public Type::TypeBase<UnrankedTensorType, TensorType,
                            detail::UnrankedTensorTypeStorage> {
public:
  using Base::Base;

  static UnrankedTensorType get(Type elementType);

  static UnrankedTensorType getChecked(Type elementType, Location location);

  static LogicalResult verifyConstructionInvariants(Optional<Location> loc,
                                                    MLIRContext *context,
                                                    Type elementType);

  ArrayRef<int64_t> getShape() const { return llvm::None; }

  static bool kindof(unsigned kind) {
    return kind == StandardTypes::UnrankedTensor;
  }
};

class BaseMemRefType : public ShapedType {
public:
  using ShapedType::ShapedType;

  static bool classof(Type type) {
    return type.getKind() == StandardTypes::MemRef ||
           type.getKind() == StandardTypes::UnrankedMemRef;
  }
};

class MemRefType : public Type::TypeBase<MemRefType, BaseMemRefType,
                                         detail::MemRefTypeStorage> {
public:
  using Base::Base;

  static MemRefType get(ArrayRef<int64_t> shape, Type elementType,
                        ArrayRef<AffineMap> affineMapComposition = {},
                        unsigned memorySpace = 0);

  static MemRefType getChecked(ArrayRef<int64_t> shape, Type elementType,
                               ArrayRef<AffineMap> affineMapComposition,
                               unsigned memorySpace, Location location);

  ArrayRef<int64_t> getShape() const;

  ArrayRef<AffineMap> getAffineMaps() const;

  unsigned getMemorySpace() const;

  // TODO(ntv): merge these two special values in a single one used everywhere.
  // Unfortunately, uses of `-1` have crept deep into the codebase now and are
  // hard to track.
  static constexpr int64_t kDynamicSize = -1;
  static int64_t getDynamicStrideOrOffset() {
    return ShapedType::kDynamicStrideOrOffset;
  }

  static bool kindof(unsigned kind) { return kind == StandardTypes::MemRef; }

private:
  static MemRefType getImpl(ArrayRef<int64_t> shape, Type elementType,
                            ArrayRef<AffineMap> affineMapComposition,
                            unsigned memorySpace, Optional<Location> location);
  using Base::getImpl;
};

class UnrankedMemRefType
    : public Type::TypeBase<UnrankedMemRefType, BaseMemRefType,
                            detail::UnrankedMemRefTypeStorage> {
public:
  using Base::Base;

  static UnrankedMemRefType get(Type elementType, unsigned memorySpace);

  static UnrankedMemRefType getChecked(Type elementType, unsigned memorySpace,
                                       Location location);

  static LogicalResult verifyConstructionInvariants(Optional<Location> loc,
                                                    MLIRContext *context,
                                                    Type elementType,
                                                    unsigned memorySpace);

  ArrayRef<int64_t> getShape() const { return llvm::None; }

  unsigned getMemorySpace() const;
  static bool kindof(unsigned kind) {
    return kind == StandardTypes::UnrankedMemRef;
  }
};

class TupleType
    : public Type::TypeBase<TupleType, Type, detail::TupleTypeStorage> {
public:
  using Base::Base;

  static TupleType get(ArrayRef<Type> elementTypes, MLIRContext *context);

  static TupleType get(MLIRContext *context) { return get({}, context); }

  ArrayRef<Type> getTypes() const;

  void getFlattenedTypes(SmallVectorImpl<Type> &types);

  size_t size() const;

  using iterator = ArrayRef<Type>::iterator;
  iterator begin() const { return getTypes().begin(); }
  iterator end() const { return getTypes().end(); }

  Type getType(size_t index) const {
    assert(index < size() && "invalid index for tuple type");
    return getTypes()[index];
  }

  static bool kindof(unsigned kind) { return kind == StandardTypes::Tuple; }
};

class NoneType : public Type::TypeBase<NoneType, Type> {
public:
  using Base::Base;

  static NoneType get(MLIRContext *context);

  static bool kindof(unsigned kind) { return kind == StandardTypes::None; }
};

LogicalResult getStridesAndOffset(MemRefType t,
                                  SmallVectorImpl<int64_t> &strides,
                                  int64_t &offset);
LogicalResult getStridesAndOffset(MemRefType t,
                                  SmallVectorImpl<AffineExpr> &strides,
                                  AffineExpr &offset);

AffineMap makeStridedLinearLayoutMap(ArrayRef<int64_t> strides, int64_t offset,
                                     MLIRContext *context);

MemRefType canonicalizeStridedLayout(MemRefType t);

bool isStrided(MemRefType t);

} // end namespace mlir

#endif // MLIR_IR_STANDARDTYPES_H