//===- TosaProfileCompliance.h - Tosa Profile-based Compliance Validation -===//
//
// 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_TOSA_TRANSFORMS_TOSAPROFILECOMPILANCE_H
#define MLIR_DIALECT_TOSA_TRANSFORMS_TOSAPROFILECOMPILANCE_H

#include <unordered_map>

#include "mlir/Dialect/Tosa/IR/TargetEnv.h"
#include "mlir/Dialect/Tosa/Transforms/Passes.h"

#include "mlir/Support/TypeID.h"

using namespace mlir;
using namespace mlir::tosa;

//===----------------------------------------------------------------------===//
// Type Compilance Definition
//===----------------------------------------------------------------------===//

struct TypeInfo {
  TypeInfo(mlir::TypeID typeID, uint32_t bitWidth)
      : typeID(typeID), bitWidth(bitWidth), valueTypeID(mlir::TypeID()),
        scaleTypeID(mlir::TypeID()), blockShape(std::nullopt) {}

  TypeInfo(mlir::TypeID typeID, uint32_t bitWidth, mlir::TypeID valueTypeID,
           mlir::TypeID scaleTypeID, tosa::BlockShape blockShape)
      : typeID(typeID), bitWidth(bitWidth), valueTypeID(valueTypeID),
        scaleTypeID(scaleTypeID), blockShape(blockShape) {}

  mlir::TypeID typeID;
  uint32_t bitWidth;
  mlir::TypeID valueTypeID;
  mlir::TypeID scaleTypeID;
  std::optional<tosa::BlockShape> blockShape;
};

enum CheckCondition {
  invalid,
  // Valid when any of the profile (extension) requirement is meet.
  anyOf,
  // Valid when all of the profile (extension) requirement are meet.
  allOf
};

using VersionedTypeInfo =
    std::pair<SmallVector<TypeInfo>, SpecificationVersion>;

template <typename T>
struct OpComplianceInfo {
  // Certain operations require multiple modes enabled.
  // e.g. cast bf16 to fp8e4m3 requires EXT-BF16 and EXT-FP8E4M3.
  SmallVector<T> mode;
  SmallVector<VersionedTypeInfo> operandTypeInfoSet;
  CheckCondition condition = CheckCondition::anyOf;
};

using OperationProfileComplianceMap =
    std::unordered_map<std::string, SmallVector<OpComplianceInfo<Profile>>>;
using OperationExtensionComplianceMap =
    std::unordered_map<std::string, SmallVector<OpComplianceInfo<Extension>>>;

//===----------------------------------------------------------------------===//
// Tosa Profile And Extension Information Depot
//===----------------------------------------------------------------------===//

class ProfileInfoDepot {
public:
  ProfileInfoDepot(Operation *op) {
    if (failed(populatationDispatch(op)))
      op->emitOpError() << "fail to populate the profile info\n";
  }

  void addType(Type t) { tyInfo.push_back(convertTypeToInfo(t)); }
  void addValue(Value v) { tyInfo.push_back(convertValueToInfo(v)); }
  SmallVector<TypeInfo> getInfo() { return tyInfo; }

private:
  TypeInfo convertTypeToInfo(Type type) {
    if (auto blockScaledTy = dyn_cast<tosa::BlockScaledType>(type)) {
      Type valueTy = blockScaledTy.getValueType();
      Type scaleTy = blockScaledTy.getScaleType();
      return {type.getTypeID(), tosa::getBitWidth(valueTy), valueTy.getTypeID(),
              scaleTy.getTypeID(), blockScaledTy.getBlockShape()};
    }
    return {type.getTypeID(), tosa::getBitWidth(type)};
  }

  TypeInfo convertValueToInfo(Value value) {
    return convertTypeToInfo(getElementTypeOrSelf(value.getType()));
  }

  LogicalResult populatationDispatch(Operation *op);

  // Add input operands and output results to the profile type info list
  LogicalResult populateProfileInfo(ValueRange operands, ValueRange results);

  // Base
  template <typename T>
  LogicalResult populateProfileInfo(T op) {
    return op->emitOpError()
           << "profile requirement for this op has not been defined";
  }
  // For conv2d, conv3d, transpose_conv2d, and depthwise_conv2d.
  template <typename T>
  LogicalResult populateProfileInfoConv(T op);

  // For reshape, slice, tile, and transpose.
  template <typename T>
  LogicalResult populateProfileInfoDataLayout(T op);

private:
  SmallVector<TypeInfo> tyInfo;
};

//===----------------------------------------------------------------------===//
// Tosa Profile And Extension Compliance Checker
//===----------------------------------------------------------------------===//

class TosaProfileCompliance {
public:
  explicit TosaProfileCompliance();

  // Accessor of the compliance info map.
  template <typename T>
  std::unordered_map<std::string, SmallVector<OpComplianceInfo<T>>>
  getProfileComplianceMap() {
    // Only profile and extension compliance info are provided.
    return {};
  }

  // Verify if the operation is allowed to be executed in the given target
  // environment.
  LogicalResult checkProfile(Operation *op, const tosa::TargetEnv &targetEnv);
  LogicalResult checkExtension(Operation *op, const tosa::TargetEnv &targetEnv);
  LogicalResult checkInvalid(Operation *op);

  template <typename T>
  LogicalResult checkProfileOrExtension(
      Operation *op, const tosa::TargetEnv &targetEnv,
      const SmallVector<ArrayRef<T>> &specDefinedProfileSet);

  bool isSameTypeInfo(TypeInfo a, TypeInfo b) {
    return a.typeID == b.typeID && a.bitWidth == b.bitWidth &&
           a.valueTypeID == b.valueTypeID && a.scaleTypeID == b.scaleTypeID &&
           a.blockShape == b.blockShape;
  }

  // Find the required profiles or extensions from the compliance info according
  // to the operand type combination.
  template <typename T>
  SmallVector<OpComplianceInfo<T>>
  findMatchedEntries(Operation *op, SmallVector<OpComplianceInfo<T>> compInfo);

  // Debug utilites.
  template <typename T>
  SmallVector<StringRef> stringifyProfile(ArrayRef<T> profiles);

  template <typename T>
  SmallVector<StringRef>
  stringifyProfile(const SmallVector<ArrayRef<T>> &profileSet);

  static llvm::SmallString<32> stringifyTypeInfo(const TypeInfo &typeInfo);

private:
  template <typename T>
  FailureOr<SmallVector<OpComplianceInfo<T>>>
  getOperatorMatchedEntries(Operation *op);

  OperationProfileComplianceMap profileComplianceMap;
  OperationExtensionComplianceMap extensionComplianceMap;
};

#endif // MLIR_DIALECT_TOSA_TRANSFORMS_TOSAPROFILECOMPILANCE_H
