//===- 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
//===----------------------------------------------------------------------===//

typedef struct {
  mlir::TypeID typeID;
  uint32_t bitWidth;
} TypeInfo;

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) {
    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;
  }

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

  SmallVector<Profile> getCooperativeProfiles(Extension ext) {
    switch (ext) {
    case Extension::int16:
    case Extension::int4:
    case Extension::doubleround:
    case Extension::inexactround:
      return {Profile::pro_int};
    case Extension::bf16:
    case Extension::fp8e4m3:
    case Extension::fp8e5m2:
    case Extension::fft:
    case Extension::mxfp:
      return {Profile::pro_fp};
    case Extension::variable:
    case Extension::controlflow:
    case Extension::dynamic:
    case Extension::int64:
    case Extension::shape:
      return {Profile::pro_fp, Profile::pro_int};
    case Extension::none:
      return {};
    };
    llvm_unreachable("bad Extension type");
  }

  // 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<7> stringifyTypeInfo(const TypeInfo &typeInfo);

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

  OperationProfileComplianceMap profileComplianceMap;
  OperationExtensionComplianceMap extensionComplianceMap;
};

#endif // MLIR_DIALECT_TOSA_TRANSFORMS_TOSAPROFILECOMPILANCE_H
