| //===-- ShapeUtils.h - TOSA shape support declarations ----------*- 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 |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Class declarations for shape utilities meant to assist shape propagation. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H |
| #define MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H |
| |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "mlir/IR/Types.h" |
| #include "mlir/Interfaces/InferTypeOpInterface.h" |
| #include "llvm/ADT/Sequence.h" |
| #include "llvm/ADT/SmallVector.h" |
| |
| namespace mlir { |
| namespace tosa { |
| /// Statically known information for a particular Value. |
| /// |
| /// This struct currently tracks only information relevant for tensor/array-like |
| /// shaped types. It is fine to associate a `ValueKnowledge` with a non-shaped |
| /// type as long as it is in the default "no knowledge" state returned by |
| /// `getPessimisticValueState`. The important invariant is that we cannot |
| /// claim to know something about a value which is false. |
| /// |
| /// This class could also be called "dataflow facts", "lattice value", etc. |
| struct ValueKnowledge { |
| ValueKnowledge() = delete; |
| ValueKnowledge(bool hasRank, llvm::ArrayRef<int64_t> newSizes, Type dtype) |
| : hasError(false), hasRank(hasRank), dtype(dtype) { |
| sizes.reserve(newSizes.size()); |
| for (auto size : newSizes) |
| sizes.push_back(size); |
| } |
| |
| operator bool() const { return !hasError; } |
| |
| // Get the static knowledge intrinsic to `type`. |
| static ValueKnowledge getKnowledgeFromType(Type type) { |
| ValueKnowledge result = getPessimisticValueState(); |
| if (auto shapedType = type.dyn_cast<ShapedType>()) { |
| if (shapedType.hasRank()) { |
| result.hasRank = true; |
| result.sizes.reserve(shapedType.getRank()); |
| for (auto dim : shapedType.getShape()) |
| result.sizes.push_back(dim); |
| } |
| result.dtype = shapedType.getElementType(); |
| } |
| return result; |
| } |
| |
| // Return a pessimistic/conservative value state without assuming any knowlege |
| // about the IR. |
| static ValueKnowledge getPessimisticValueState() { |
| return ValueKnowledge(false, {}, Type()); |
| } |
| |
| ShapedTypeComponents getShapedTypeComponents() const { |
| return hasRank ? ShapedTypeComponents(sizes) : ShapedTypeComponents(); |
| } |
| |
| Type getType() const { |
| if (hasRank) |
| return RankedTensorType::get(llvm::makeArrayRef(sizes), dtype); |
| return UnrankedTensorType::get(dtype); |
| } |
| |
| bool operator==(const ValueKnowledge &rhs) const { |
| return hasRank == rhs.hasRank && sizes == rhs.sizes && dtype == rhs.dtype; |
| } |
| |
| // Given two pieces of static knowledge, calculate conservatively the |
| // information we can be sure about. |
| static ValueKnowledge join(const ValueKnowledge &lhs, |
| const ValueKnowledge &rhs) { |
| // Mental model: All conditions are checking how to change from the safe "no |
| // knowledge" default-initialized state to a state with more knowledge |
| // consistent with lhs and rhs. |
| ValueKnowledge result = getPessimisticValueState(); |
| result.hasError = true; |
| |
| if (!lhs || !rhs || lhs.dtype != rhs.dtype) |
| return result; |
| |
| result.hasError = false; |
| result.dtype = lhs.dtype; |
| |
| if (!lhs.hasRank && !rhs.hasRank) |
| return result; |
| |
| if (!rhs.hasRank) { |
| result.hasRank = true; |
| result.sizes = lhs.sizes; |
| return result; |
| } |
| |
| if (!lhs.hasRank) { |
| result.hasRank = true; |
| result.sizes = rhs.sizes; |
| return result; |
| } |
| |
| if (lhs.sizes.size() != rhs.sizes.size()) |
| return result; |
| |
| result.hasRank = true; |
| result.sizes.resize(lhs.sizes.size(), ShapedType::kDynamicSize); |
| for (auto i : llvm::seq<unsigned>(0, result.sizes.size())) { |
| int64_t lhsSize = lhs.sizes[i]; |
| int64_t rhsSize = rhs.sizes[i]; |
| int64_t &resultSize = result.sizes[i]; |
| if (lhsSize == ShapedType::kDynamicSize) { |
| resultSize = rhsSize; |
| } else if (rhsSize == ShapedType::kDynamicSize) { |
| resultSize = lhsSize; |
| } else if (lhsSize == rhsSize) { |
| resultSize = lhsSize; |
| } else { |
| result.hasError = true; |
| } |
| } |
| |
| return result; |
| } |
| |
| // Given to types, generate a new ValueKnowledge that meets to cover both |
| // cases. E.g. if the rank of the LHS and RHS differ, the resulting tensor |
| // has unknown rank. |
| static ValueKnowledge meet(const ValueKnowledge &lhs, |
| const ValueKnowledge &rhs) { |
| ValueKnowledge result = getPessimisticValueState(); |
| result.hasError = true; |
| |
| if (!rhs || !rhs || lhs.dtype != rhs.dtype) |
| return result; |
| |
| result.hasError = false; |
| result.dtype = lhs.dtype; |
| |
| if (!lhs.hasRank || !rhs.hasRank) { |
| result.hasRank = false; |
| return result; |
| } |
| |
| if (lhs.sizes.size() != rhs.sizes.size()) { |
| result.hasRank = false; |
| return result; |
| } |
| |
| result.hasRank = true; |
| result.sizes.resize(lhs.sizes.size(), ShapedType::kDynamicSize); |
| for (int i = 0, e = lhs.sizes.size(); i < e; i++) { |
| if (lhs.sizes[i] == rhs.sizes[i]) { |
| result.sizes[i] = lhs.sizes[i]; |
| } |
| } |
| |
| return result; |
| } |
| |
| // Whether the value information has an error. |
| bool hasError; |
| // Whether the value has known rank. |
| bool hasRank; |
| // If `hasRank`, the sizes along each rank. Unknown sizes are represented as |
| // `ShapedType::kDynamicSize`. |
| llvm::SmallVector<int64_t> sizes; |
| // The dtype of a tensor. |
| // This is equal to nullptr if we don't know that it is a specific concrete |
| // type. |
| Type dtype; |
| }; |
| } // namespace tosa |
| } // namespace mlir |
| |
| #endif // MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H |