TosaProfileCompliance.h

Go to the documentation of this file.

//===- 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 <initializer_list>
#include <unordered_map>
 
#include "mlir/Dialect/Tosa/IR/TargetEnv.h"
#include "mlir/Dialect/Tosa/Transforms/Passes.h"
 
#include "mlir/Support/TypeID.h"
 
namespace mlir {
namespace 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
};
 
struct TypeInfoRef {
  const TypeInfo *ptr;
  TypeInfoRef(const TypeInfo &info) : ptr(&info) {}
};
 
struct VersionedTypeInfo {
  SmallVector<TypeInfo> first;
  SpecificationVersion second;
 
  VersionedTypeInfo() = default;
  VersionedTypeInfo(std::initializer_list<TypeInfoRef> types,
                    SpecificationVersion version);
  VersionedTypeInfo(SmallVector<TypeInfo> types, SpecificationVersion version);
  VersionedTypeInfo(const VersionedTypeInfo &other);
  VersionedTypeInfo(VersionedTypeInfo &&other);
  VersionedTypeInfo &operator=(const VersionedTypeInfo &other);
  VersionedTypeInfo &operator=(VersionedTypeInfo &&other);
  ~VersionedTypeInfo();
};
 
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;
 
  OpComplianceInfo() = default;
  OpComplianceInfo(std::initializer_list<T> mode,
                   std::initializer_list<VersionedTypeInfo> operandTypeInfoSet,
                   CheckCondition condition = CheckCondition::anyOf);
  OpComplianceInfo(SmallVector<T> mode,
                   SmallVector<VersionedTypeInfo> operandTypeInfoSet,
                   CheckCondition condition = CheckCondition::anyOf);
  OpComplianceInfo(const OpComplianceInfo &other);
  OpComplianceInfo(OpComplianceInfo &&other);
  OpComplianceInfo &operator=(const OpComplianceInfo &other);
  OpComplianceInfo &operator=(OpComplianceInfo &&other);
  ~OpComplianceInfo();
};
 
extern template struct OpComplianceInfo<Profile>;
extern template struct OpComplianceInfo<Extension>;
 
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)); }
  ArrayRef<TypeInfo> getInfo() const { return tyInfo; }
 
private:
  TypeInfo convertTypeToInfo(Type type) {
    // Treat the deprecated TOSA type as builtin mxi8 for compliance checks.
    const auto getTypeID = [](Type type) {
      if (isa<tosa::mxint8Type>(type))
        return MXI8Type::getTypeID();
      return type.getTypeID();
    };
    if (auto blockScaledTy = dyn_cast<tosa::BlockScaledType>(type)) {
      Type valueTy = blockScaledTy.getValueType();
      Type scaleTy = blockScaledTy.getScaleType();
      return {type.getTypeID(), tosa::getBitWidth(valueTy), getTypeID(valueTy),
              scaleTy.getTypeID(), blockScaledTy.getBlockShape()};
    }
    return {getTypeID(type), 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>
  const std::unordered_map<std::string, SmallVector<OpComplianceInfo<T>>> &
  getProfileComplianceMap() const;
 
  // 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);
 
  bool isSameTypeInfo(const TypeInfo &a, const 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, ArrayRef<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;
};
 
template <>
const OperationProfileComplianceMap &
TosaProfileCompliance::getProfileComplianceMap<Profile>() const;
 
template <>
const OperationExtensionComplianceMap &
TosaProfileCompliance::getProfileComplianceMap<Extension>() const;
} // namespace tosa
} // namespace mlir
 
#endif // MLIR_DIALECT_TOSA_TRANSFORMS_TOSAPROFILECOMPILANCE_H