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