text: //===-- 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 "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());
    }
  
    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
  