327 lines
12 KiB
C++
327 lines
12 KiB
C++
//===- BuiltinAttributeInterfaces.h - Builtin Attr Interfaces ---*- 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_IR_BUILTINATTRIBUTEINTERFACES_H
|
|
#define MLIR_IR_BUILTINATTRIBUTEINTERFACES_H
|
|
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Attributes.h"
|
|
#include "mlir/IR/BuiltinTypeInterfaces.h"
|
|
#include "mlir/IR/Types.h"
|
|
#include "llvm/Support/raw_ostream.h"
|
|
#include <complex>
|
|
#include <optional>
|
|
|
|
namespace mlir {
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ElementsAttr
|
|
//===----------------------------------------------------------------------===//
|
|
namespace detail {
|
|
/// This class provides support for indexing into the element range of an
|
|
/// ElementsAttr. It is used to opaquely wrap either a contiguous range, via
|
|
/// `ElementsAttrIndexer::contiguous`, or a non-contiguous range, via
|
|
/// `ElementsAttrIndexer::nonContiguous`, A contiguous range is an array-like
|
|
/// range, where all of the elements are layed out sequentially in memory. A
|
|
/// non-contiguous range implies no contiguity, and elements may even be
|
|
/// materialized when indexing, such as the case for a mapped_range.
|
|
struct ElementsAttrIndexer {
|
|
public:
|
|
ElementsAttrIndexer()
|
|
: ElementsAttrIndexer(/*isContiguous=*/true, /*isSplat=*/true) {}
|
|
ElementsAttrIndexer(ElementsAttrIndexer &&rhs)
|
|
: isContiguous(rhs.isContiguous), isSplat(rhs.isSplat) {
|
|
if (isContiguous)
|
|
conState = rhs.conState;
|
|
else
|
|
new (&nonConState) NonContiguousState(std::move(rhs.nonConState));
|
|
}
|
|
ElementsAttrIndexer(const ElementsAttrIndexer &rhs)
|
|
: isContiguous(rhs.isContiguous), isSplat(rhs.isSplat) {
|
|
if (isContiguous)
|
|
conState = rhs.conState;
|
|
else
|
|
new (&nonConState) NonContiguousState(rhs.nonConState);
|
|
}
|
|
~ElementsAttrIndexer() {
|
|
if (!isContiguous)
|
|
nonConState.~NonContiguousState();
|
|
}
|
|
|
|
/// Construct an indexer for a non-contiguous range starting at the given
|
|
/// iterator. A non-contiguous range implies no contiguity, and elements may
|
|
/// even be materialized when indexing, such as the case for a mapped_range.
|
|
template <typename IteratorT>
|
|
static ElementsAttrIndexer nonContiguous(bool isSplat, IteratorT &&iterator) {
|
|
ElementsAttrIndexer indexer(/*isContiguous=*/false, isSplat);
|
|
new (&indexer.nonConState)
|
|
NonContiguousState(std::forward<IteratorT>(iterator));
|
|
return indexer;
|
|
}
|
|
|
|
// Construct an indexer for a contiguous range starting at the given element
|
|
// pointer. A contiguous range is an array-like range, where all of the
|
|
// elements are layed out sequentially in memory.
|
|
template <typename T>
|
|
static ElementsAttrIndexer contiguous(bool isSplat, const T *firstEltPtr) {
|
|
ElementsAttrIndexer indexer(/*isContiguous=*/true, isSplat);
|
|
new (&indexer.conState) ContiguousState(firstEltPtr);
|
|
return indexer;
|
|
}
|
|
|
|
/// Access the element at the given index.
|
|
template <typename T>
|
|
T at(uint64_t index) const {
|
|
if (isSplat)
|
|
index = 0;
|
|
return isContiguous ? conState.at<T>(index) : nonConState.at<T>(index);
|
|
}
|
|
|
|
private:
|
|
ElementsAttrIndexer(bool isContiguous, bool isSplat)
|
|
: isContiguous(isContiguous), isSplat(isSplat), conState(nullptr) {}
|
|
|
|
/// This class contains all of the state necessary to index a contiguous
|
|
/// range.
|
|
class ContiguousState {
|
|
public:
|
|
ContiguousState(const void *firstEltPtr) : firstEltPtr(firstEltPtr) {}
|
|
|
|
/// Access the element at the given index.
|
|
template <typename T>
|
|
const T &at(uint64_t index) const {
|
|
return *(reinterpret_cast<const T *>(firstEltPtr) + index);
|
|
}
|
|
|
|
private:
|
|
const void *firstEltPtr;
|
|
};
|
|
|
|
/// This class contains all of the state necessary to index a non-contiguous
|
|
/// range.
|
|
class NonContiguousState {
|
|
private:
|
|
/// This class is used to represent the abstract base of an opaque iterator.
|
|
/// This allows for all iterator and element types to be completely
|
|
/// type-erased.
|
|
struct OpaqueIteratorBase {
|
|
virtual ~OpaqueIteratorBase() = default;
|
|
virtual std::unique_ptr<OpaqueIteratorBase> clone() const = 0;
|
|
};
|
|
/// This class is used to represent the abstract base of an opaque iterator
|
|
/// that iterates over elements of type `T`. This allows for all iterator
|
|
/// types to be completely type-erased.
|
|
template <typename T>
|
|
struct OpaqueIteratorValueBase : public OpaqueIteratorBase {
|
|
virtual T at(uint64_t index) = 0;
|
|
};
|
|
/// This class is used to represent an opaque handle to an iterator of type
|
|
/// `IteratorT` that iterates over elements of type `T`.
|
|
template <typename IteratorT, typename T>
|
|
struct OpaqueIterator : public OpaqueIteratorValueBase<T> {
|
|
template <typename ItTy, typename FuncTy, typename FuncReturnTy>
|
|
static void isMappedIteratorTestFn(
|
|
llvm::mapped_iterator<ItTy, FuncTy, FuncReturnTy>) {}
|
|
template <typename U, typename... Args>
|
|
using is_mapped_iterator =
|
|
decltype(isMappedIteratorTestFn(std::declval<U>()));
|
|
template <typename U>
|
|
using detect_is_mapped_iterator =
|
|
llvm::is_detected<is_mapped_iterator, U>;
|
|
|
|
/// Access the element within the iterator at the given index.
|
|
template <typename ItT>
|
|
static std::enable_if_t<!detect_is_mapped_iterator<ItT>::value, T>
|
|
atImpl(ItT &&it, uint64_t index) {
|
|
return *std::next(it, index);
|
|
}
|
|
template <typename ItT>
|
|
static std::enable_if_t<detect_is_mapped_iterator<ItT>::value, T>
|
|
atImpl(ItT &&it, uint64_t index) {
|
|
// Special case mapped_iterator to avoid copying the function.
|
|
return it.getFunction()(*std::next(it.getCurrent(), index));
|
|
}
|
|
|
|
public:
|
|
template <typename U>
|
|
OpaqueIterator(U &&iterator) : iterator(std::forward<U>(iterator)) {}
|
|
std::unique_ptr<OpaqueIteratorBase> clone() const final {
|
|
return std::make_unique<OpaqueIterator<IteratorT, T>>(iterator);
|
|
}
|
|
|
|
/// Access the element at the given index.
|
|
T at(uint64_t index) final { return atImpl(iterator, index); }
|
|
|
|
private:
|
|
IteratorT iterator;
|
|
};
|
|
|
|
public:
|
|
/// Construct the state with the given iterator type.
|
|
template <typename IteratorT, typename T = typename llvm::remove_cvref_t<
|
|
decltype(*std::declval<IteratorT>())>>
|
|
NonContiguousState(IteratorT iterator)
|
|
: iterator(std::make_unique<OpaqueIterator<IteratorT, T>>(iterator)) {}
|
|
NonContiguousState(const NonContiguousState &other)
|
|
: iterator(other.iterator->clone()) {}
|
|
NonContiguousState(NonContiguousState &&other) = default;
|
|
|
|
/// Access the element at the given index.
|
|
template <typename T>
|
|
T at(uint64_t index) const {
|
|
auto *valueIt = static_cast<OpaqueIteratorValueBase<T> *>(iterator.get());
|
|
return valueIt->at(index);
|
|
}
|
|
|
|
/// The opaque iterator state.
|
|
std::unique_ptr<OpaqueIteratorBase> iterator;
|
|
};
|
|
|
|
/// A boolean indicating if this range is contiguous or not.
|
|
bool isContiguous;
|
|
/// A boolean indicating if this range is a splat.
|
|
bool isSplat;
|
|
/// The underlying range state.
|
|
union {
|
|
ContiguousState conState;
|
|
NonContiguousState nonConState;
|
|
};
|
|
};
|
|
|
|
/// This class implements a generic iterator for ElementsAttr.
|
|
template <typename T>
|
|
class ElementsAttrIterator
|
|
: public llvm::iterator_facade_base<ElementsAttrIterator<T>,
|
|
std::random_access_iterator_tag, T,
|
|
std::ptrdiff_t, T, T> {
|
|
public:
|
|
ElementsAttrIterator(ElementsAttrIndexer indexer, size_t dataIndex)
|
|
: indexer(std::move(indexer)), index(dataIndex) {}
|
|
|
|
// Boilerplate iterator methods.
|
|
ptrdiff_t operator-(const ElementsAttrIterator &rhs) const {
|
|
return index - rhs.index;
|
|
}
|
|
bool operator==(const ElementsAttrIterator &rhs) const {
|
|
return index == rhs.index;
|
|
}
|
|
bool operator<(const ElementsAttrIterator &rhs) const {
|
|
return index < rhs.index;
|
|
}
|
|
ElementsAttrIterator &operator+=(ptrdiff_t offset) {
|
|
index += offset;
|
|
return *this;
|
|
}
|
|
ElementsAttrIterator &operator-=(ptrdiff_t offset) {
|
|
index -= offset;
|
|
return *this;
|
|
}
|
|
|
|
/// Return the value at the current iterator position.
|
|
T operator*() const { return indexer.at<T>(index); }
|
|
|
|
private:
|
|
ElementsAttrIndexer indexer;
|
|
ptrdiff_t index;
|
|
};
|
|
|
|
/// This class provides iterator utilities for an ElementsAttr range.
|
|
template <typename IteratorT>
|
|
class ElementsAttrRange : public llvm::iterator_range<IteratorT> {
|
|
public:
|
|
using reference = typename IteratorT::reference;
|
|
|
|
ElementsAttrRange(ShapedType shapeType,
|
|
const llvm::iterator_range<IteratorT> &range)
|
|
: llvm::iterator_range<IteratorT>(range), shapeType(shapeType) {}
|
|
ElementsAttrRange(ShapedType shapeType, IteratorT beginIt, IteratorT endIt)
|
|
: ElementsAttrRange(shapeType, llvm::make_range(beginIt, endIt)) {}
|
|
|
|
/// Return the value at the given index.
|
|
reference operator[](ArrayRef<uint64_t> index) const;
|
|
reference operator[](uint64_t index) const {
|
|
return *std::next(this->begin(), index);
|
|
}
|
|
|
|
/// Return the size of this range.
|
|
size_t size() const { return llvm::size(*this); }
|
|
|
|
private:
|
|
/// The shaped type of the parent ElementsAttr.
|
|
ShapedType shapeType;
|
|
};
|
|
|
|
} // namespace detail
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// MemRefLayoutAttrInterface
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace detail {
|
|
|
|
// Verify the affine map 'm' can be used as a layout specification
|
|
// for memref with 'shape'.
|
|
LogicalResult
|
|
verifyAffineMapAsLayout(AffineMap m, ArrayRef<int64_t> shape,
|
|
function_ref<InFlightDiagnostic()> emitError);
|
|
|
|
// Return the strides and offsets that can be inferred from the given affine
|
|
// layout map given the map and a memref shape.
|
|
LogicalResult getAffineMapStridesAndOffset(AffineMap map,
|
|
ArrayRef<int64_t> shape,
|
|
SmallVectorImpl<int64_t> &strides,
|
|
int64_t &offset);
|
|
} // namespace detail
|
|
|
|
} // namespace mlir
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Tablegen Interface Declarations
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/IR/BuiltinAttributeInterfaces.h.inc"
|
|
#include "mlir/IR/OpAsmAttrInterface.h.inc"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ElementsAttr
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace mlir {
|
|
namespace detail {
|
|
/// Return the value at the given index.
|
|
template <typename IteratorT>
|
|
auto ElementsAttrRange<IteratorT>::operator[](ArrayRef<uint64_t> index) const
|
|
-> reference {
|
|
// Skip to the element corresponding to the flattened index.
|
|
return (*this)[ElementsAttr::getFlattenedIndex(shapeType, index)];
|
|
}
|
|
} // namespace detail
|
|
|
|
/// Return the elements of this attribute as a value of type 'T'.
|
|
template <typename T>
|
|
auto ElementsAttr::value_begin() const -> DefaultValueCheckT<T, iterator<T>> {
|
|
if (std::optional<iterator<T>> iterator = try_value_begin<T>())
|
|
return std::move(*iterator);
|
|
llvm::errs()
|
|
<< "ElementsAttr does not provide iteration facilities for type `"
|
|
<< llvm::getTypeName<T>() << "`, see attribute: " << *this << "\n";
|
|
llvm_unreachable("invalid `T` for ElementsAttr::getValues");
|
|
}
|
|
template <typename T>
|
|
auto ElementsAttr::try_value_begin() const
|
|
-> DefaultValueCheckT<T, std::optional<iterator<T>>> {
|
|
FailureOr<detail::ElementsAttrIndexer> indexer =
|
|
getValuesImpl(TypeID::get<T>());
|
|
if (failed(indexer))
|
|
return std::nullopt;
|
|
return iterator<T>(std::move(*indexer), 0);
|
|
}
|
|
} // namespace mlir.
|
|
|
|
#endif // MLIR_IR_BUILTINATTRIBUTEINTERFACES_H
|