derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
5756 lines
217 KiB
C++
5756 lines
217 KiB
C++
//===- TosaOps.cpp - MLIR Dialect for TOSA --------------------------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// \file
|
|
// This file implements the TOSA Specification:
|
|
// https://www.mlplatform.org/tosa/tosa_spec.html
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
|
#include "mlir/Dialect/Quant/IR/Quant.h"
|
|
#include "mlir/Dialect/Shard/Interfaces/ShardingInterface.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
|
|
#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
|
|
#include "mlir/Dialect/Utils/IndexingUtils.h"
|
|
#include "mlir/Dialect/Utils/VerificationUtils.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/DialectImplementation.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/IR/TypeUtilities.h"
|
|
#include "mlir/Interfaces/InferTypeOpInterface.h"
|
|
#include "mlir/Transforms/InliningUtils.h"
|
|
#include "llvm/ADT/APFloat.h"
|
|
#include "llvm/ADT/SmallVectorExtras.h"
|
|
#include "llvm/ADT/TypeSwitch.h"
|
|
|
|
#include <numeric>
|
|
#include <type_traits>
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::tosa;
|
|
|
|
#include "mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc"
|
|
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Tosa dialect interface includes.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Tosa/IR/TosaAvailability.cpp.inc"
|
|
#include "mlir/Dialect/Tosa/IR/TosaEnums.cpp.inc"
|
|
#include "mlir/Dialect/Tosa/IR/TosaInterfaces.cpp.inc"
|
|
#include "mlir/Dialect/Tosa/IR/TosaOpAvailabilityImpl.inc"
|
|
|
|
namespace {
|
|
#include "mlir/Dialect/Tosa/IR/TosaDialectBytecode.cpp.inc"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Dialect Function Inliner Interface.
|
|
//===----------------------------------------------------------------------===//
|
|
struct TosaInlinerInterface : public DialectInlinerInterface {
|
|
using DialectInlinerInterface::DialectInlinerInterface;
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Analysis Hooks.
|
|
//===--------------------------------------------------------------------===//
|
|
|
|
/// All operations can be inlined by default.
|
|
bool isLegalToInline(Operation *op, Region *region, bool wouldBeCloned,
|
|
IRMapping &map) const final {
|
|
return true;
|
|
}
|
|
|
|
/// All regions with If and While parent operators can be inlined.
|
|
bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
|
|
IRMapping &map) const final {
|
|
return (isa<tosa::IfOp>(dest->getParentOp()) ||
|
|
isa<tosa::WhileOp>(dest->getParentOp()));
|
|
}
|
|
};
|
|
|
|
/// This class implements the bytecode interface for the Tosa dialect.
|
|
struct TosaDialectBytecodeInterface : public BytecodeDialectInterface {
|
|
TosaDialectBytecodeInterface(Dialect *dialect)
|
|
: BytecodeDialectInterface(dialect) {}
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Attributes
|
|
|
|
Attribute readAttribute(DialectBytecodeReader &reader) const override {
|
|
return ::readAttribute(getContext(), reader);
|
|
}
|
|
|
|
LogicalResult writeAttribute(Attribute attr,
|
|
DialectBytecodeWriter &writer) const override {
|
|
return ::writeAttribute(attr, writer);
|
|
}
|
|
|
|
//===--------------------------------------------------------------------===//
|
|
// Types
|
|
|
|
Type readType(DialectBytecodeReader &reader) const override {
|
|
return ::readType(getContext(), reader);
|
|
}
|
|
|
|
LogicalResult writeType(Type type,
|
|
DialectBytecodeWriter &writer) const override {
|
|
return ::writeType(type, writer);
|
|
}
|
|
|
|
void writeVersion(DialectBytecodeWriter &writer) const final {
|
|
// TODO: Populate.
|
|
}
|
|
|
|
std::unique_ptr<DialectVersion>
|
|
readVersion(DialectBytecodeReader &reader) const final {
|
|
// TODO: Populate
|
|
reader.emitError("Dialect does not support versioning");
|
|
return nullptr;
|
|
}
|
|
|
|
LogicalResult upgradeFromVersion(Operation *topLevelOp,
|
|
const DialectVersion &version) const final {
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA control flow support.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Returns the while loop body.
|
|
SmallVector<Region *> tosa::WhileOp::getLoopRegions() {
|
|
return {&getBodyGraph()};
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA variable operator support.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
static SmallVector<int64_t> convertToMlirShape(ArrayRef<int64_t> shape) {
|
|
return map_to_vector(shape, [](int64_t dim) {
|
|
return dim == -1 ? ShapedType::kDynamic : dim;
|
|
});
|
|
}
|
|
|
|
// returns type of variable op
|
|
RankedTensorType mlir::tosa::getVariableType(tosa::VariableOp variableOp) {
|
|
Type elementType = variableOp.getType();
|
|
DenseIntElementsAttr varShapeAttr = variableOp.getVarShape();
|
|
auto shape = convertToMlirShape(to_vector(varShapeAttr.getValues<int64_t>()));
|
|
return RankedTensorType::get(shape, elementType);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Tosa dialect initialization.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void TosaDialect::initialize() {
|
|
addTypes<
|
|
#define GET_TYPEDEF_LIST
|
|
#include "mlir/Dialect/Tosa/IR/TosaOpsTypesBase.cpp.inc"
|
|
>();
|
|
addOperations<
|
|
#define GET_OP_LIST
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
|
|
>();
|
|
addAttributes<
|
|
#define GET_ATTRDEF_LIST
|
|
#include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"
|
|
>();
|
|
addInterfaces<TosaDialectBytecodeInterface, TosaInlinerInterface>();
|
|
declarePromisedInterfaces<
|
|
shard::ShardingInterface, ClampOp, SigmoidOp, TanhOp, AddOp,
|
|
ArithmeticRightShiftOp, BitwiseAndOp, BitwiseOrOp, BitwiseXorOp, IntDivOp,
|
|
LogicalAndOp, LogicalLeftShiftOp, LogicalRightShiftOp, LogicalOrOp,
|
|
LogicalXorOp, MaximumOp, MinimumOp, MulOp, PowOp, SubOp, AbsOp,
|
|
BitwiseNotOp, CeilOp, ClzOp, ExpOp, FloorOp, LogOp, LogicalNotOp,
|
|
NegateOp, ReciprocalOp, RsqrtOp, SelectOp, EqualOp, GreaterOp,
|
|
GreaterEqualOp, MatMulOp>();
|
|
}
|
|
|
|
Operation *TosaDialect::materializeConstant(OpBuilder &builder, Attribute value,
|
|
Type type, Location loc) {
|
|
// Tosa dialect constants only support ElementsAttr unlike standard dialect
|
|
// constant which supports all attributes.
|
|
if (llvm::isa<shapeType>(type) && llvm::isa<DenseIntElementsAttr>(value)) {
|
|
return tosa::ConstShapeOp::create(builder, loc, type,
|
|
llvm::cast<DenseIntElementsAttr>(value));
|
|
}
|
|
if (llvm::isa<ElementsAttr>(value))
|
|
return tosa::ConstOp::create(builder, loc, type,
|
|
llvm::cast<ElementsAttr>(value));
|
|
return nullptr;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Parsers and printers
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
ParseResult getShapeAndElementType(OpAsmParser &parser, Type parsedType,
|
|
DenseElementsAttr &varShapeAttr,
|
|
TypeAttr &typeAttr) {
|
|
if (auto shapedType = dyn_cast<ShapedType>(parsedType)) {
|
|
if (!shapedType.hasRank())
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected ranked type";
|
|
|
|
auto elementType = shapedType.getElementType();
|
|
typeAttr = TypeAttr::get(elementType);
|
|
ArrayRef<int64_t> shape = shapedType.getShape();
|
|
Builder builder(parser.getContext());
|
|
varShapeAttr = builder.getIndexTensorAttr(convertFromMlirShape(shape));
|
|
return success();
|
|
}
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected shaped type";
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// parses the optional initial value or type for a tosa variable
|
|
// with initial value:
|
|
// tosa.variable @name = dense<0.0> : tensor<1x8xf32>
|
|
//
|
|
// without initial value:
|
|
// tosa.variable @name : tensor<1x8xf32>
|
|
ParseResult mlir::tosa::parseVariableOpTypeOrInitialValue(
|
|
OpAsmParser &parser, DenseElementsAttr &varShapeAttr, TypeAttr &typeAttr,
|
|
Attribute &initialValueAttr) {
|
|
if (succeeded(parser.parseOptionalEqual())) {
|
|
if (failed(parser.parseAttribute(initialValueAttr))) {
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected attribute";
|
|
}
|
|
if (auto typedAttr = dyn_cast<TypedAttr>(initialValueAttr)) {
|
|
return getShapeAndElementType(parser, typedAttr.getType(), varShapeAttr,
|
|
typeAttr);
|
|
}
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected Typed attr";
|
|
}
|
|
|
|
initialValueAttr = nullptr;
|
|
Type parsedType;
|
|
if (failed(parser.parseColonType(parsedType))) {
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected type after colon";
|
|
}
|
|
return getShapeAndElementType(parser, parsedType, varShapeAttr, typeAttr);
|
|
}
|
|
|
|
void mlir::tosa::printVariableOpTypeOrInitialValue(
|
|
OpAsmPrinter &p, Operation *op, DenseElementsAttr varShapeAttr,
|
|
TypeAttr typeAttr, Attribute initialValueAttr) {
|
|
bool needsSpace = false;
|
|
if (!dyn_cast_or_null<TypedAttr>(initialValueAttr)) {
|
|
auto shape =
|
|
convertToMlirShape(to_vector(varShapeAttr.getValues<int64_t>()));
|
|
Type elementType = typeAttr.getValue();
|
|
RankedTensorType tensorType =
|
|
RankedTensorType::get(ArrayRef<int64_t>(shape), elementType);
|
|
auto tensorTypeAttr = TypeAttr::get(tensorType);
|
|
p << ": ";
|
|
p.printAttribute(tensorTypeAttr);
|
|
needsSpace = true; // subsequent attr value needs a space separator
|
|
}
|
|
if (initialValueAttr) {
|
|
if (needsSpace)
|
|
p << ' ';
|
|
p << "= ";
|
|
p.printAttribute(initialValueAttr);
|
|
}
|
|
}
|
|
|
|
namespace {
|
|
|
|
// parse attributes with special handling for tosa enum attributes
|
|
template <typename EnumType>
|
|
ParseResult parseAttrEntryWithEnumHandling(OpAsmParser &parser,
|
|
NamedAttrList &outAttrs) {
|
|
llvm::StringRef name;
|
|
if (parser.parseOptionalKeyword(&name) || parser.parseEqual())
|
|
return failure();
|
|
|
|
// special handling: rounding_mode accepts a *bare* RoundingMode enum
|
|
// keyword.
|
|
llvm::StringRef kw;
|
|
if constexpr (std::is_same_v<EnumType, tosa::RoundingMode>) {
|
|
if (name == "rounding_mode" &&
|
|
succeeded(parser.parseOptionalKeyword(&kw))) {
|
|
auto sym = symbolizeRoundingMode(kw);
|
|
if (!sym)
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "invalid rounding_mode value: " << kw;
|
|
auto attr = RoundingModeAttr::get(parser.getContext(), sym.value());
|
|
outAttrs.push_back(NamedAttribute(name, attr));
|
|
return success();
|
|
}
|
|
}
|
|
// special handling: mode accepts a *bare* ResizeMode enum keyword.
|
|
if constexpr (std::is_same_v<EnumType, tosa::ResizeMode>) {
|
|
if (name == "mode" && succeeded(parser.parseOptionalKeyword(&kw))) {
|
|
auto sym = symbolizeResizeMode(kw);
|
|
if (!sym)
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "invalid resize mode value: " << kw;
|
|
auto attr = ResizeModeAttr::get(parser.getContext(), sym.value());
|
|
outAttrs.push_back(NamedAttribute(name, attr));
|
|
return success();
|
|
}
|
|
}
|
|
// special handling: nan_mode accepts a *bare* NanPropagationMode enum
|
|
// keyword.
|
|
if constexpr (std::is_same_v<EnumType, tosa::NanPropagationMode>) {
|
|
if (name == "nan_mode" && succeeded(parser.parseOptionalKeyword(&kw))) {
|
|
auto sym = symbolizeNanPropagationMode(kw);
|
|
if (!sym)
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "invalid nan_mode value: " << kw;
|
|
auto attr = NanPropagationModeAttr::get(parser.getContext(), sym.value());
|
|
outAttrs.push_back(NamedAttribute(name, attr));
|
|
return success();
|
|
}
|
|
}
|
|
|
|
// special handling: block_size accepts a *bare* BlockSizeMode enum
|
|
if constexpr (std::is_same_v<EnumType, tosa::BlockSize>) {
|
|
if (name == "block_size" && succeeded(parser.parseOptionalKeyword(&kw))) {
|
|
auto sym = symbolizeBlockSize(kw);
|
|
if (!sym)
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "invalid block_size value: " << kw;
|
|
auto attr = BlockSizeAttr::get(parser.getContext(), sym.value());
|
|
outAttrs.push_back(NamedAttribute(name, attr));
|
|
return success();
|
|
}
|
|
}
|
|
|
|
// Default path: parse any normal attribute literal, including fully qualified
|
|
// enum keyword
|
|
Attribute attr;
|
|
return parser.parseAttribute(attr, name, outAttrs);
|
|
}
|
|
|
|
template <typename EnumType>
|
|
ParseResult parseWithEnumHandling(OpAsmParser &parser, OperationState &result) {
|
|
// parse operands
|
|
SmallVector<OpAsmParser::UnresolvedOperand, 5> operands;
|
|
if (parser.parseCommaSeparatedList(
|
|
[&]() { return parser.parseOperand(operands.emplace_back()); }))
|
|
return failure();
|
|
|
|
// Parse { attr-dict } with special handling for enum bare token
|
|
NamedAttrList attrs;
|
|
if (succeeded(parser.parseOptionalLBrace()) &&
|
|
failed(parser.parseOptionalRBrace())) {
|
|
do {
|
|
if (parseAttrEntryWithEnumHandling<EnumType>(parser, attrs))
|
|
return failure();
|
|
} while (succeeded(parser.parseOptionalComma()));
|
|
if (parser.parseRBrace())
|
|
return failure();
|
|
}
|
|
|
|
FunctionType fnTy;
|
|
if (parser.parseColonType(fnTy))
|
|
return failure();
|
|
|
|
// Resolve operands and types
|
|
if (failed(parser.resolveOperands(operands, fnTy.getInputs(),
|
|
parser.getCurrentLocation(),
|
|
result.operands)))
|
|
return failure();
|
|
|
|
result.addTypes(fnTy.getResults());
|
|
result.addAttributes(attrs);
|
|
|
|
return success();
|
|
}
|
|
|
|
void printNamedAttr(OpAsmPrinter &parser, const NamedAttribute namedAttr) {
|
|
parser << namedAttr.getName().strref() << " = ";
|
|
auto attr = namedAttr.getValue();
|
|
if (auto roundingModeAttr = dyn_cast<tosa::RoundingModeAttr>(attr)) {
|
|
parser << roundingModeAttr.getValue();
|
|
} else if (auto resizeModeAttr = dyn_cast<tosa::ResizeModeAttr>(attr)) {
|
|
parser << resizeModeAttr.getValue();
|
|
} else if (auto nanPropagationModeAttr =
|
|
dyn_cast<tosa::NanPropagationModeAttr>(attr)) {
|
|
parser << nanPropagationModeAttr.getValue();
|
|
} else if (auto blockSizeAttr = dyn_cast<tosa::BlockSizeAttr>(attr)) {
|
|
parser << blockSizeAttr.getValue();
|
|
} else {
|
|
parser.printAttribute(attr);
|
|
}
|
|
}
|
|
|
|
// print with special handling for default valued NanPropagationMode attribute
|
|
void printWithNanPropagationHandling(OpAsmPrinter &parser, Operation *op) {
|
|
parser << " ";
|
|
parser.printOperands(op->getOperands());
|
|
|
|
NamedAttrList toPrint(op->getAttrs());
|
|
// remove default NanPropagate attribute
|
|
const auto kDefaultNanValue = NanPropagationMode::PROPAGATE;
|
|
for (auto attr : op->getAttrs()) {
|
|
if (auto nanAttr = dyn_cast<NanPropagationModeAttr>(attr.getValue())) {
|
|
if (nanAttr.getValue() == kDefaultNanValue) {
|
|
// elide from toPrint
|
|
toPrint.erase(attr.getName());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!toPrint.empty()) {
|
|
parser << " {";
|
|
llvm::interleaveComma(toPrint, parser, [&](const NamedAttribute namedAttr) {
|
|
printNamedAttr(parser, namedAttr);
|
|
});
|
|
parser << "}";
|
|
}
|
|
|
|
parser << " : ";
|
|
parser.printFunctionalType(op);
|
|
}
|
|
|
|
// print with special handling for enums: RoundingMode, ResizeMode
|
|
void printWithEnumHandling(OpAsmPrinter &parser, Operation *op) {
|
|
parser << " ";
|
|
parser.printOperands(op->getOperands());
|
|
|
|
if (!op->getAttrs().empty()) {
|
|
parser << " {";
|
|
llvm::interleaveComma(op->getAttrs(), parser,
|
|
[&](const NamedAttribute namedAttr) {
|
|
printNamedAttr(parser, namedAttr);
|
|
});
|
|
parser << "}";
|
|
}
|
|
|
|
parser << " : ";
|
|
parser.printFunctionalType(op);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
ParseResult RescaleOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::RoundingMode>(parser, result);
|
|
}
|
|
|
|
void RescaleOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult ApplyScaleOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::RoundingMode>(parser, result);
|
|
}
|
|
|
|
void ApplyScaleOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult ResizeOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::ResizeMode>(parser, result);
|
|
}
|
|
|
|
void ResizeOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult ArgMaxOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void ArgMaxOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult MaxPool2dOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void MaxPool2dOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult MaxPool2dAdaptiveOp::parse(OpAsmParser &parser,
|
|
OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void MaxPool2dAdaptiveOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult ClampOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void ClampOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult MaximumOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void MaximumOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult MinimumOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void MinimumOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult ReduceMaxOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void ReduceMaxOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult ReduceMinOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
return parseWithEnumHandling<tosa::NanPropagationMode>(parser, result);
|
|
}
|
|
|
|
void ReduceMinOp::print(OpAsmPrinter &parser) {
|
|
printWithNanPropagationHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult MatmulTBlockScaledOp::parse(OpAsmParser &parser,
|
|
OperationState &result) {
|
|
return parseWithEnumHandling<tosa::BlockSize>(parser, result);
|
|
}
|
|
|
|
void MatmulTBlockScaledOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult CastFromBlockScaledOp::parse(OpAsmParser &parser,
|
|
OperationState &result) {
|
|
return parseWithEnumHandling<tosa::BlockSize>(parser, result);
|
|
}
|
|
|
|
void CastFromBlockScaledOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult CastToBlockScaledOp::parse(OpAsmParser &parser,
|
|
OperationState &result) {
|
|
return parseWithEnumHandling<tosa::BlockSize>(parser, result);
|
|
}
|
|
|
|
void CastToBlockScaledOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
ParseResult Conv2DBlockScaledOp::parse(OpAsmParser &parser,
|
|
OperationState &result) {
|
|
return parseWithEnumHandling<tosa::BlockSize>(parser, result);
|
|
}
|
|
|
|
void Conv2DBlockScaledOp::print(OpAsmPrinter &parser) {
|
|
printWithEnumHandling(parser, *this);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Tosa utilities.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
static std::optional<int64_t> idivCheck(const int64_t lhs, const int64_t rhs) {
|
|
if (lhs % rhs != 0)
|
|
return std::nullopt;
|
|
return lhs / rhs;
|
|
}
|
|
|
|
static Type getStorageElementTypeOrSelf(Type type) {
|
|
auto srcType = getElementTypeOrSelf(type);
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(srcType))
|
|
srcType = getStorageElementTypeFromQuantized(quantType);
|
|
return srcType;
|
|
}
|
|
|
|
static Type getStorageElementTypeOrSelf(Value value) {
|
|
return getStorageElementTypeOrSelf(value.getType());
|
|
}
|
|
|
|
static LogicalResult verifyRescaleValueAndZpTypes(Operation *op, Value val,
|
|
Value valZp, StringRef name) {
|
|
Type eType = getStorageElementTypeOrSelf(val.getType());
|
|
Type eZpType = getStorageElementTypeOrSelf(valZp.getType());
|
|
|
|
bool bothInts =
|
|
mlir::isa<IntegerType>(eType) && mlir::isa<IntegerType>(eZpType);
|
|
bool sameBitWidth =
|
|
(eType.getIntOrFloatBitWidth() == eZpType.getIntOrFloatBitWidth());
|
|
|
|
if (!bothInts || !sameBitWidth) {
|
|
return op->emitOpError()
|
|
<< "expected " << name << " and " << name
|
|
<< "_zp to both be integer of the same bitwidth, but got " << eType
|
|
<< " vs. " << eZpType;
|
|
}
|
|
return success();
|
|
}
|
|
|
|
// Create a pad-const const tensor with value of `val` of required data-type
|
|
Value mlir::tosa::createPadConstTensor(OpBuilder &builder, Location loc,
|
|
Value src, int32_t val) {
|
|
const auto srcType = getElementTypeOrSelf(src);
|
|
const auto srcElemType = getStorageElementTypeOrSelf(src);
|
|
const auto padConstType = mlir::RankedTensorType::get({1}, srcType);
|
|
const auto padConstEType = mlir::RankedTensorType::get({1}, srcElemType);
|
|
const auto padConstAttr{
|
|
llvm::isa<FloatType>(srcElemType)
|
|
? DenseElementsAttr::get(padConstEType,
|
|
builder.getFloatAttr(srcElemType, val))
|
|
: DenseElementsAttr::get(padConstEType,
|
|
builder.getIntegerAttr(srcElemType, val))};
|
|
return tosa::ConstOp::create(builder, loc, padConstType, padConstAttr);
|
|
}
|
|
|
|
unsigned mlir::tosa::getBitWidth(Type type) {
|
|
if (dyn_cast<tosa::mxint8Type>(type))
|
|
return 8;
|
|
return type.getIntOrFloatBitWidth();
|
|
}
|
|
|
|
// Update dim size if current dim is dynamic, otherwise raise an error if sizes
|
|
// do not match
|
|
LogicalResult tryUpdateDimOrFailure(Operation *op, int64_t &currDim,
|
|
const int64_t newDim,
|
|
const StringRef operandName,
|
|
const StringRef dimName) {
|
|
if (ShapedType::isDynamic(currDim)) {
|
|
currDim = newDim;
|
|
return success();
|
|
} else if (ShapedType::isStatic(newDim) && currDim != newDim) {
|
|
return op->emitOpError("expected ")
|
|
<< dimName << " of " << operandName << " to match size " << currDim
|
|
<< ", got " << newDim;
|
|
}
|
|
return success();
|
|
}
|
|
|
|
static void printShapeToDiagnostic(InFlightDiagnostic &diag,
|
|
ArrayRef<int64_t> shape) {
|
|
auto printDim = [&](int64_t dim) {
|
|
if (ShapedType::isDynamic(dim))
|
|
diag << "?";
|
|
else
|
|
diag << dim;
|
|
};
|
|
|
|
llvm::interleaveComma(shape, diag, printDim);
|
|
}
|
|
|
|
static LogicalResult
|
|
verifyOutputShapeCompatibleWithExpected(Operation *op, ShapedType outputType,
|
|
ArrayRef<int64_t> expectedShape,
|
|
StringRef outputName = "output") {
|
|
assert(outputType.hasRank() && "expected output type to be ranked");
|
|
|
|
if (succeeded(verifyCompatibleShape(outputType.getShape(), expectedShape)))
|
|
return success();
|
|
|
|
InFlightDiagnostic diag = op->emitOpError("expected ");
|
|
diag << outputName << " shape ";
|
|
printShapeToDiagnostic(diag, outputType.getShape());
|
|
diag << " to be compatible with inferred shape ";
|
|
printShapeToDiagnostic(diag, expectedShape);
|
|
return diag;
|
|
}
|
|
|
|
LogicalResult verifyConvOutputSize(
|
|
Operation *op, const int64_t inputSize, const int64_t kernelSize,
|
|
const int64_t outputSize, const int64_t padBefore, const int64_t padAfter,
|
|
const int64_t stride, const int64_t dilation, const llvm::StringRef dimName,
|
|
const llvm::StringRef dimAxis, const llvm::StringRef padBeforeName,
|
|
const llvm::StringRef padAfterName) {
|
|
if (inputSize == ShapedType::kDynamic || kernelSize == ShapedType::kDynamic)
|
|
return success();
|
|
|
|
// ERROR_IF: O != idiv_check(I - 1 + pa + pb - (K - 1) * d, s) + 1
|
|
|
|
const std::optional<int64_t> calculatedOutSizeMinusOne = idivCheck(
|
|
inputSize - 1 + padBefore + padAfter - (kernelSize - 1) * dilation,
|
|
stride);
|
|
if (!calculatedOutSizeMinusOne.has_value())
|
|
return op->emitOpError("expected input_")
|
|
<< dimName << " - 1 + pad_" << padBeforeName << " + pad_"
|
|
<< padAfterName << " - (kernel_" << dimName << " - 1) * dilation_"
|
|
<< dimAxis << " to be wholly divisible by stride_" << dimAxis
|
|
<< ", got (" << inputSize << " - 1 + " << padBefore << " + "
|
|
<< padAfter << " - (" << kernelSize << " - 1) * " << dilation
|
|
<< ") / " << stride;
|
|
|
|
const int64_t calculatedOutSize = calculatedOutSizeMinusOne.value() + 1;
|
|
if (outputSize != ShapedType::kDynamic && calculatedOutSize != outputSize)
|
|
return op->emitOpError("calculated output ")
|
|
<< dimName << " did not match expected: "
|
|
<< "calculated=" << calculatedOutSize << ", expected=" << outputSize;
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Operator Verifiers.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyConvOp(T op) {
|
|
const auto inputType = llvm::dyn_cast<TensorType>(op.getInput().getType());
|
|
const auto weightType = llvm::dyn_cast<TensorType>(op.getWeight().getType());
|
|
|
|
auto inputEType = inputType.getElementType();
|
|
auto weightEType = weightType.getElementType();
|
|
auto biasEType =
|
|
llvm::cast<ShapedType>(op.getBias().getType()).getElementType();
|
|
auto resultEType =
|
|
llvm::cast<ShapedType>(op.getResult().getType()).getElementType();
|
|
bool biasIsFloat = llvm::isa<FloatType>(biasEType);
|
|
bool resultIsFloat = llvm::isa<FloatType>(resultEType);
|
|
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(inputEType))
|
|
inputEType = getStorageElementTypeFromQuantized(quantType);
|
|
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(weightEType))
|
|
weightEType = getStorageElementTypeFromQuantized(quantType);
|
|
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(biasEType))
|
|
biasEType = getStorageElementTypeFromQuantized(quantType);
|
|
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(resultEType))
|
|
resultEType = getStorageElementTypeFromQuantized(quantType);
|
|
|
|
if (biasIsFloat && resultIsFloat && (biasEType != resultEType)) {
|
|
// for now, only enforce bias element type == result element type for
|
|
// float types.
|
|
op.emitOpError(
|
|
"expect both bias and result to have same element type, got ")
|
|
<< biasEType << " and " << resultEType;
|
|
return failure();
|
|
}
|
|
|
|
if (isa<Float8E5M2Type>(inputEType) || isa<Float8E4M3FNType>(inputEType) ||
|
|
isa<Float8E5M2Type>(weightEType) || isa<Float8E4M3FNType>(weightEType)) {
|
|
if (inputEType != weightEType) {
|
|
op.emitOpError(
|
|
"expect both input and weight to have same element type, got ")
|
|
<< inputEType << " and " << weightEType;
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
bool inputIsFloat = llvm::isa<FloatType>(inputEType);
|
|
bool weightIsFloat = llvm::isa<FloatType>(weightEType);
|
|
|
|
// Either both must be float or both non-float.
|
|
if (inputIsFloat != weightIsFloat) {
|
|
op.emitOpError(
|
|
"expect both input and weight to be float or not together, got ")
|
|
<< inputEType << " and " << weightEType;
|
|
return failure();
|
|
}
|
|
|
|
auto inputZpEType = getStorageElementTypeOrSelf(op.getInputZp().getType());
|
|
if (inputEType != inputZpEType) {
|
|
return op.emitOpError("expect both input and its zero point are the same "
|
|
"element type, got ")
|
|
<< inputEType << " and " << inputZpEType;
|
|
}
|
|
|
|
auto weightZpEType = getStorageElementTypeOrSelf(op.getWeightZp().getType());
|
|
if (weightEType != weightZpEType) {
|
|
return op.emitOpError("expect both weight and its zero point are the same "
|
|
"element type, got ")
|
|
<< weightEType << " and " << weightZpEType;
|
|
}
|
|
|
|
FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
|
|
if (succeeded(maybeIZp) && op.verifyInputZeroPoint(*maybeIZp).failed())
|
|
return failure();
|
|
|
|
FailureOr<int64_t> maybeWZp = op.getWeightZeroPoint();
|
|
if (succeeded(maybeWZp) && op.verifyWeightZeroPoint(*maybeWZp).failed())
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ConstOp::verify() {
|
|
|
|
auto attrType = llvm::dyn_cast<TensorType>(getValuesAttr().getType());
|
|
auto outputType = llvm::dyn_cast<TensorType>(getOutput().getType());
|
|
|
|
if (!attrType || !outputType) {
|
|
emitOpError("expected tensors for attr/result type");
|
|
return failure();
|
|
}
|
|
|
|
if (auto result = llvm::dyn_cast<mlir::quant::QuantizedType>(
|
|
outputType.getElementType())) {
|
|
if (getStorageElementTypeFromQuantized(result) == attrType.getElementType())
|
|
return success();
|
|
}
|
|
|
|
if (attrType.getElementType() != outputType.getElementType()) {
|
|
emitOpError("expected same attr/result element types");
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyConvOpModes(T op) {
|
|
auto inputEType =
|
|
llvm::cast<ShapedType>(op.getInput().getType()).getElementType();
|
|
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(inputEType))
|
|
inputEType = getStorageElementTypeFromQuantized(quantType);
|
|
|
|
auto accType = op.getAccType();
|
|
if (inputEType.isInteger(8) && !accType.isInteger(32))
|
|
return op.emitOpError("accumulator type for i8 tensor is not i32, got ")
|
|
<< accType;
|
|
|
|
if (inputEType.isInteger(16) && !accType.isInteger(48))
|
|
return op.emitOpError("accumulator type for i16 tensor is not i48, got ")
|
|
<< accType;
|
|
|
|
if (isa<Float8E5M2Type, Float8E4M3Type>(inputEType) &&
|
|
!(accType.isF16() || accType.isF32()))
|
|
return op.emitOpError("accumulator type for f8 tensor is not f16/f32, got ")
|
|
<< accType;
|
|
|
|
if (inputEType.isF16() && !(accType.isF16() || accType.isF32()))
|
|
return op.emitOpError(
|
|
"accumulator type for f16 tensor is not f16/f32, got ")
|
|
<< accType;
|
|
|
|
if (inputEType.isBF16() && !accType.isF32())
|
|
return op.emitOpError("accumulator type for bf16 tensor is not f32, got ")
|
|
<< accType;
|
|
|
|
if (inputEType.isF32() && !accType.isF32())
|
|
return op.emitOpError("accumulator type for f32 tensor is not f32, got ")
|
|
<< accType;
|
|
|
|
auto resultEType =
|
|
llvm::cast<ShapedType>(op.getResult().getType()).getElementType();
|
|
|
|
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(resultEType))
|
|
resultEType = getStorageElementTypeFromQuantized(quantType);
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ERROR_IF functions.
|
|
// ERROR_IF is a predicate that must set an error if the condition holds.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyConvOpErrorIf(T op) {
|
|
llvm::ArrayRef<int64_t> padding = op.getPad();
|
|
if (llvm::any_of(padding, [](int64_t p) { return p < 0; }))
|
|
return op.emitOpError("expect all padding values to be >= 0, got ")
|
|
<< padding;
|
|
|
|
llvm::ArrayRef<int64_t> strides = op.getStride();
|
|
if (llvm::any_of(strides, [](int64_t s) { return s < 1; }))
|
|
return op.emitOpError("expect all stride values to be >= 1, got ")
|
|
<< strides;
|
|
|
|
llvm::ArrayRef<int64_t> dilations = op.getDilation();
|
|
if (llvm::any_of(dilations, [](int64_t d) { return d < 1; }))
|
|
return op.emitOpError("expect all dilation values to be >= 1, got ")
|
|
<< dilations;
|
|
|
|
const RankedTensorType outputType =
|
|
llvm::dyn_cast<RankedTensorType>(op.getOutput().getType());
|
|
if (!outputType)
|
|
// Skip following checks if output is not ranked
|
|
return success();
|
|
|
|
const RankedTensorType inputType =
|
|
llvm::dyn_cast<RankedTensorType>(op.getInput().getType());
|
|
const RankedTensorType weightType =
|
|
llvm::dyn_cast<RankedTensorType>(op.getWeight().getType());
|
|
|
|
if (inputType && weightType) {
|
|
// input = [_,IH,IW,_], weight = [_,KH,KW,_], output = [_,OH,OW,_]
|
|
if constexpr (std::is_same<T, tosa::Conv2DOp>::value) {
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(1), weightType.getDimSize(1),
|
|
outputType.getDimSize(1), padding[0], padding[1], strides[0],
|
|
dilations[0], "height", "y", "top", "bottom")))
|
|
return failure();
|
|
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(2), weightType.getDimSize(2),
|
|
outputType.getDimSize(2), padding[2], padding[3], strides[1],
|
|
dilations[1], "width", "x", "left", "right")))
|
|
return failure();
|
|
}
|
|
|
|
// input = [_,IH,IW,_], weight = [KH,KW,_,_], output = [_,OH,OW,_]
|
|
if constexpr (std::is_same<T, tosa::DepthwiseConv2DOp>::value) {
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(1), weightType.getDimSize(0),
|
|
outputType.getDimSize(1), padding[0], padding[1], strides[0],
|
|
dilations[0], "height", "y", "top", "bottom")))
|
|
return failure();
|
|
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(2), weightType.getDimSize(1),
|
|
outputType.getDimSize(2), padding[2], padding[3], strides[1],
|
|
dilations[1], "width", "x", "left", "right")))
|
|
return failure();
|
|
}
|
|
|
|
// input = [_,ID,IH,IW,_], weight = [_,KD,KH,KW,_], output = [_,OD,OH,OW,_]
|
|
if constexpr (std::is_same<T, tosa::Conv3DOp>::value) {
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(1), weightType.getDimSize(1),
|
|
outputType.getDimSize(1), padding[0], padding[1], strides[0],
|
|
dilations[0], "depth", "d", "front", "back")))
|
|
return failure();
|
|
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(2), weightType.getDimSize(2),
|
|
outputType.getDimSize(2), padding[2], padding[3], strides[1],
|
|
dilations[1], "height", "y", "top", "bottom")))
|
|
return failure();
|
|
|
|
if (failed(verifyConvOutputSize(
|
|
op, inputType.getDimSize(3), weightType.getDimSize(3),
|
|
outputType.getDimSize(3), padding[4], padding[5], strides[2],
|
|
dilations[2], "width", "x", "left", "right")))
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
const RankedTensorType biasType =
|
|
llvm::dyn_cast<RankedTensorType>(op.getBias().getType());
|
|
if (!biasType)
|
|
// Skip following checks if bias is not ranked
|
|
return success();
|
|
|
|
const int64_t biasChannels = biasType.getDimSize(0);
|
|
const int64_t outputChannels =
|
|
outputType.getDimSize(outputType.getRank() - 1);
|
|
if (biasChannels == ShapedType::kDynamic ||
|
|
outputChannels == ShapedType::kDynamic)
|
|
// Skip following checks if biasChannels or outputChannels is dynamic dim
|
|
return success();
|
|
|
|
if (biasChannels != outputChannels && biasChannels != 1)
|
|
return op.emitOpError(
|
|
"bias channels expected to be equal to output channels (")
|
|
<< outputChannels << ") or 1, got " << biasChannels;
|
|
|
|
return success();
|
|
}
|
|
|
|
// Verify whether same type and shape of the given two types.
|
|
static LogicalResult errorIfTypeOrShapeMismatch(Operation *op, Type type1,
|
|
StringRef name1, Type type2,
|
|
StringRef name2) {
|
|
auto shapeType1 = dyn_cast<ShapedType>(type1);
|
|
auto shapeType2 = dyn_cast<ShapedType>(type2);
|
|
if (!shapeType1 || !shapeType2)
|
|
return failure();
|
|
|
|
auto elemType1 = shapeType1.getElementType();
|
|
auto elemType2 = shapeType2.getElementType();
|
|
if (elemType1 != elemType2)
|
|
return op->emitOpError()
|
|
<< "require same element type for " << name1 << " (" << elemType1
|
|
<< ") and " << name2 << " (" << elemType2 << ")";
|
|
|
|
if (failed(verifyCompatibleShape(type1, type2)))
|
|
return op->emitOpError()
|
|
<< "require same shapes for " << name1 << " (" << type1 << ") and "
|
|
<< name2 << " (" << type2 << ")";
|
|
|
|
return success();
|
|
}
|
|
|
|
// Verify whether same length, type, and shape of the given two tensor lists.
|
|
static LogicalResult errorIfTypeOrShapeMismatch(Operation *op, ValueRange list1,
|
|
StringRef name1,
|
|
ValueRange list2,
|
|
StringRef name2) {
|
|
if (list1.size() != list2.size())
|
|
return op->emitOpError()
|
|
<< "require same number of values in " << name1 << " ("
|
|
<< list1.size() << ") and " << name2 << " (" << list2.size() << ")";
|
|
|
|
for (auto [type1, type2] :
|
|
llvm::zip_equal(list1.getTypes(), list2.getTypes())) {
|
|
if (errorIfTypeOrShapeMismatch(op, type1, name1, type2, name2).failed())
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
static inline LogicalResult errorIfShapeNotSizeOne(Operation *op, Type type) {
|
|
ShapeAdaptor shapeAdaptor(type);
|
|
if (!shapeAdaptor.hasRank() || !shapeAdaptor.hasStaticShape())
|
|
return success();
|
|
|
|
return shapeAdaptor.getNumElements() == 1 ? success() : failure();
|
|
}
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyVariableOpErrorIf(T op, Type type, StringRef name) {
|
|
Operation *symTableOp =
|
|
op->template getParentWithTrait<OpTrait::SymbolTable>();
|
|
if (!symTableOp)
|
|
// If the operation is not the scope of a symbol table, we cannot
|
|
// verify it against it's declaration.
|
|
return success();
|
|
|
|
SymbolTable symTable(symTableOp);
|
|
const auto varOp = symTable.lookup<tosa::VariableOp>(op.getName());
|
|
|
|
// Verify prior declaration
|
|
if (!varOp)
|
|
return op->emitOpError("'")
|
|
<< op.getName() << "' has not been declared by 'tosa.variable'";
|
|
|
|
// Verify type and shape
|
|
auto variableType = getVariableType(varOp);
|
|
if (errorIfTypeOrShapeMismatch(op, type, name, variableType,
|
|
"the input tensor")
|
|
.failed())
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
// verify that inType and outType have same element types
|
|
template <typename T>
|
|
static LogicalResult verifySameElementTypes(T op, Type aType, Type bType,
|
|
StringRef aName = "input",
|
|
StringRef bName = "output") {
|
|
auto aTType = llvm::dyn_cast<TensorType>(aType);
|
|
auto bTType = llvm::dyn_cast<TensorType>(bType);
|
|
if (!aTType) {
|
|
op.emitOpError("expect shaped tensor for") << aName << ", got " << aType;
|
|
return failure();
|
|
}
|
|
if (!bTType) {
|
|
op.emitOpError("expect shaped tensor for") << bName << ", got" << bType;
|
|
return failure();
|
|
}
|
|
auto aElementType = aTType.getElementType();
|
|
auto bElementType = bTType.getElementType();
|
|
auto aQuantType =
|
|
llvm::dyn_cast<mlir::quant::UniformQuantizedType>(aElementType);
|
|
auto bQuantType =
|
|
llvm::dyn_cast<mlir::quant::UniformQuantizedType>(bElementType);
|
|
if ((aElementType.isIntOrIndexOrFloat() || aQuantType) &&
|
|
(bElementType.isIntOrIndexOrFloat() || bQuantType) &&
|
|
aElementType != bElementType) {
|
|
// only check if both element types are int/index/float/UniformQuantized
|
|
// eg, not sure how to check quant::QuantizedType
|
|
// this happens in test_conv2d_q_grouped_convolution in
|
|
// tfl-to-tosa-pipeline.mlir
|
|
op.emitOpError("expect ")
|
|
<< aName << " and " << bName << " to have same element type, got "
|
|
<< aElementType << " and " << bElementType;
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ArgMaxOp::verify() {
|
|
const ShapedType resultType = llvm::cast<ShapedType>(getType());
|
|
|
|
// Ensure output is of 32-bit integer
|
|
if (const auto resultETy = resultType.getElementType();
|
|
!resultETy.isIntOrIndex())
|
|
return emitOpError("result tensor is not of integer type");
|
|
|
|
const auto inputType = llvm::cast<ShapedType>(getInput().getType());
|
|
if (!inputType.hasRank())
|
|
return success();
|
|
|
|
// Ensure axis is within the tensor rank
|
|
const int64_t axis = getAxisAttr().getInt();
|
|
if (((axis < 0) || axis >= inputType.getRank()))
|
|
return emitOpError("specified axis is outside the rank of the tensor");
|
|
|
|
if (!resultType.hasRank())
|
|
return success();
|
|
|
|
const ArrayRef<int64_t> inputShape = inputType.getShape();
|
|
const ArrayRef<int64_t> outputShape = resultType.getShape();
|
|
llvm::SmallVector<int64_t> expectedOutputShape(inputShape);
|
|
expectedOutputShape.erase(expectedOutputShape.begin() + axis);
|
|
if (failed(verifyCompatibleShape(expectedOutputShape, outputShape)))
|
|
return emitOpError("expected output shape '")
|
|
<< expectedOutputShape << "', got '" << outputShape << "'";
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyPoolingOpImpl(Operation *op,
|
|
ArrayRef<int64_t> kernel,
|
|
ArrayRef<int64_t> strides,
|
|
ArrayRef<int64_t> padding, Value input,
|
|
Value output) {
|
|
const bool hasKernel = kernel.size() > 0;
|
|
const bool hasStrides = strides.size() > 0;
|
|
const bool hasPad = padding.size() > 0;
|
|
|
|
if (hasKernel && llvm::any_of(kernel, [](int64_t s) { return s < 1; }))
|
|
return op->emitOpError("expect all kernel values to be >= 1, got ")
|
|
<< kernel;
|
|
|
|
if (hasStrides && llvm::any_of(strides, [](int64_t s) { return s < 1; }))
|
|
return op->emitOpError("expect all stride values to be >= 1, got ")
|
|
<< strides;
|
|
|
|
if (hasPad && llvm::any_of(padding, [](int64_t p) { return p < 0; }))
|
|
return op->emitOpError("expect all padding values to be >= 0, got ")
|
|
<< padding;
|
|
|
|
if (hasKernel && hasPad) {
|
|
// Padding must be less than kernel size to avoid a divide-by-zero
|
|
const int64_t kernelX = kernel[1];
|
|
const int64_t padLeft = padding[2];
|
|
const int64_t padRight = padding[3];
|
|
if (padRight >= kernelX || padLeft >= kernelX)
|
|
return op->emitOpError("expected left/right padding to be less than the "
|
|
"width of the kernel, got pad_left=")
|
|
<< padLeft << ", pad_right=" << padRight
|
|
<< ", kernel_x=" << kernelX;
|
|
|
|
const int64_t kernelY = kernel[0];
|
|
const int64_t padTop = padding[0];
|
|
const int64_t padBottom = padding[1];
|
|
if (padTop >= kernelY || padBottom >= kernelY)
|
|
return op->emitOpError("expected top/bottom padding to be less than the "
|
|
"height of the kernel, got pad_top=")
|
|
<< padTop << ", pad_bottom=" << padBottom
|
|
<< ", kernel_y=" << kernelY;
|
|
}
|
|
|
|
const auto inputType = llvm::dyn_cast<RankedTensorType>(input.getType());
|
|
const auto outputType = llvm::dyn_cast<RankedTensorType>(output.getType());
|
|
if (!inputType || !outputType)
|
|
return success();
|
|
|
|
if (hasKernel && hasStrides && hasPad) {
|
|
const auto verifyOutputSize =
|
|
[op](const int64_t inputSize, const int64_t outputSize,
|
|
const int64_t kernelSize, const int64_t strideSize,
|
|
const int64_t padBefore, const int64_t padAfter,
|
|
const llvm::StringRef dimName, const llvm::StringRef dimAxis,
|
|
const llvm::StringRef padBeforeName,
|
|
const llvm::StringRef padAfterName) -> LogicalResult {
|
|
if (ShapedType::isDynamic(inputSize))
|
|
return success();
|
|
|
|
const std::optional<int64_t> calculatedOutSizeMinusOne =
|
|
idivCheck(inputSize + padBefore + padAfter - kernelSize, strideSize);
|
|
if (!calculatedOutSizeMinusOne.has_value())
|
|
return op->emitOpError("expected input_")
|
|
<< dimName << " + pad_" << padBeforeName << " + pad_"
|
|
<< padAfterName << " - kernel_" << dimAxis
|
|
<< " to be wholly divisible by stride_" << dimAxis << ", got ("
|
|
<< inputSize << " + " << padBefore << " + " << padAfter << " - "
|
|
<< kernelSize << ") / " << strideSize;
|
|
|
|
const int64_t calculatedOutSize = calculatedOutSizeMinusOne.value() + 1;
|
|
if (ShapedType::isStatic(outputSize) && calculatedOutSize != outputSize)
|
|
return op->emitOpError("calculated output ")
|
|
<< dimName << " did not match expected: " << "calculated="
|
|
<< calculatedOutSize << ", expected=" << outputSize;
|
|
|
|
return success();
|
|
};
|
|
|
|
if (failed(verifyOutputSize(inputType.getDimSize(1),
|
|
outputType.getDimSize(1), kernel[0], strides[0],
|
|
padding[0], padding[1], "height", "y", "top",
|
|
"bottom")))
|
|
return failure();
|
|
|
|
if (failed(verifyOutputSize(
|
|
inputType.getDimSize(2), outputType.getDimSize(2), kernel[1],
|
|
strides[1], padding[2], padding[3], "width", "x", "left", "right")))
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyPoolingOp(T op) {
|
|
return verifyPoolingOpImpl(op.getOperation(), op.getKernel(), op.getStride(),
|
|
op.getPad(), op.getInput(), op.getOutput());
|
|
}
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyAvgPoolCommonTypeAndZpChecks(T op) {
|
|
const Type inputETy = getStorageElementTypeOrSelf(op.getInput().getType());
|
|
const Type resultETy = getStorageElementTypeOrSelf(op.getOutput().getType());
|
|
const Type inputZpETy =
|
|
getStorageElementTypeOrSelf(op.getInputZp().getType());
|
|
const Type outputZpETy =
|
|
getStorageElementTypeOrSelf(op.getOutputZp().getType());
|
|
|
|
auto accType = op.getAccType();
|
|
if (llvm::isa<IntegerType>(inputETy) && !accType.isInteger(32))
|
|
return op.emitOpError("accumulator type for integer tensor is not i32");
|
|
|
|
if (inputETy.isF16() && !(accType.isF16() || accType.isF32()))
|
|
return op.emitOpError("accumulator type for f16 tensor is not f16/f32");
|
|
|
|
if (inputETy.isBF16() && !accType.isF32())
|
|
return op.emitOpError("accumulator type for bf16 tensor is not f32");
|
|
|
|
if (inputETy.isF32() && !accType.isF32())
|
|
return op.emitOpError("accumulator type for f32 tensor is not f32");
|
|
|
|
if (inputETy != inputZpETy)
|
|
return op.emitOpError("expect both input and its zero point are the same "
|
|
"element type, got ")
|
|
<< inputETy << " and " << inputZpETy;
|
|
|
|
if (resultETy != outputZpETy)
|
|
return op.emitOpError("expect both output and its zero point are the same "
|
|
"element type, got ")
|
|
<< resultETy << " and " << outputZpETy;
|
|
|
|
FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
|
|
if (succeeded(maybeIZp) && op.verifyInputZeroPoint(*maybeIZp).failed())
|
|
return failure();
|
|
|
|
FailureOr<int64_t> maybeOZp = op.getOutputZeroPoint();
|
|
if (succeeded(maybeOZp) && op.verifyOutputZeroPoint(*maybeOZp).failed())
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
namespace {
|
|
struct AdaptivePoolingConstShapeValues {
|
|
llvm::SmallVector<int64_t> kernel;
|
|
llvm::SmallVector<int64_t> stride;
|
|
llvm::SmallVector<int64_t> pad;
|
|
};
|
|
} // namespace
|
|
|
|
template <typename T>
|
|
static constexpr bool IsSupportedAdaptivePoolConstShapeVerifyOp =
|
|
std::is_same_v<T, tosa::AvgPool2dAdaptiveOp> ||
|
|
std::is_same_v<T, tosa::MaxPool2dAdaptiveOp>;
|
|
|
|
template <typename T,
|
|
typename std::enable_if<IsSupportedAdaptivePoolConstShapeVerifyOp<T>,
|
|
int>::type = 0>
|
|
static void extractAdaptivePoolingConstShapeOperands(
|
|
T op, AdaptivePoolingConstShapeValues &values) {
|
|
tosa::getConstShapeValues(op.getKernel().getDefiningOp(), values.kernel);
|
|
tosa::getConstShapeValues(op.getStride().getDefiningOp(), values.stride);
|
|
tosa::getConstShapeValues(op.getPad().getDefiningOp(), values.pad);
|
|
}
|
|
|
|
LogicalResult tosa::AvgPool2dOp::verify() {
|
|
if (failed(verifyPoolingOp(*this)))
|
|
return failure();
|
|
if (failed(verifyAvgPoolCommonTypeAndZpChecks(*this)))
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::AvgPool2dAdaptiveOp::verify() {
|
|
AdaptivePoolingConstShapeValues values;
|
|
extractAdaptivePoolingConstShapeOperands(*this, values);
|
|
|
|
// If pad/stride/kernel are not constant, this is okay, we just can't check
|
|
// their values. extractAdaptivePoolingConstShapeOperands will return an empty
|
|
// list for each non CTC input. verifyPoolingOpImpl will need to handle values
|
|
// not being present, and return success if they cannot be checked.
|
|
|
|
if (failed(verifyPoolingOpImpl(getOperation(), values.kernel, values.stride,
|
|
values.pad, getInput(), getOutput())))
|
|
return failure();
|
|
|
|
if (failed(verifyAvgPoolCommonTypeAndZpChecks(*this)))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ClampOp::verify() {
|
|
mlir::Type inputETy =
|
|
llvm::cast<ShapedType>(getInput().getType()).getElementType();
|
|
if (auto quantType =
|
|
llvm::dyn_cast<mlir::quant::UniformQuantizedType>(inputETy)) {
|
|
inputETy = getStorageElementTypeFromQuantized(quantType);
|
|
}
|
|
mlir::Type outputETy =
|
|
llvm::cast<ShapedType>(getOutput().getType()).getElementType();
|
|
if (auto quantType =
|
|
llvm::dyn_cast<mlir::quant::UniformQuantizedType>(outputETy)) {
|
|
outputETy = getStorageElementTypeFromQuantized(quantType);
|
|
}
|
|
if (inputETy != outputETy)
|
|
return emitOpError("input/output element types are incompatible.");
|
|
|
|
auto maxValAttr = getMaxValAttr();
|
|
auto minValAttr = getMinValAttr();
|
|
|
|
unsigned dataTypeBitWidth = inputETy.getIntOrFloatBitWidth();
|
|
|
|
if (inputETy.isInteger(dataTypeBitWidth)) {
|
|
// if input datatype is integer, check that the min_val/max_val attributes
|
|
// are integer attributes, and that their type is the same as the input's
|
|
// datatype
|
|
auto intMaxValAttr = mlir::dyn_cast<mlir::IntegerAttr>(maxValAttr);
|
|
auto intMinValAttr = mlir::dyn_cast<mlir::IntegerAttr>(minValAttr);
|
|
if (!intMaxValAttr || !intMinValAttr ||
|
|
(intMaxValAttr.getType() != intMinValAttr.getType()) ||
|
|
(intMaxValAttr.getType() != inputETy))
|
|
return emitOpError("min/max attributes types are incompatible with "
|
|
"input/output element types.");
|
|
|
|
const bool isUnsigned = inputETy.isUnsignedInteger();
|
|
const bool isBoolean = inputETy.isInteger(1);
|
|
const APInt minVal = intMinValAttr.getValue();
|
|
const APInt maxVal = intMaxValAttr.getValue();
|
|
if ((isUnsigned || isBoolean) ? maxVal.ult(minVal) : maxVal.slt(minVal))
|
|
return emitOpError("expected min_val <= max_val, got min_val=")
|
|
<< minValAttr << ", max_val=" << maxValAttr;
|
|
} else {
|
|
// otherwise, input datatype is float, check that the min_val/max_val
|
|
// attributes share the same type and that their type is the same as the
|
|
// input's datatype
|
|
auto floatMaxValAttr = mlir::dyn_cast<mlir::FloatAttr>(maxValAttr);
|
|
auto floatMinValAttr = mlir::dyn_cast<mlir::FloatAttr>(minValAttr);
|
|
if (!floatMaxValAttr || !floatMinValAttr ||
|
|
(floatMaxValAttr.getType() != floatMinValAttr.getType()) ||
|
|
(floatMaxValAttr.getType() != inputETy))
|
|
return emitOpError("min/max attributes types are incompatible with "
|
|
"input/output element types.");
|
|
|
|
const APFloat minVal = floatMinValAttr.getValue();
|
|
const APFloat maxVal = floatMaxValAttr.getValue();
|
|
if (minVal.isNaN() || maxVal.isNaN())
|
|
return emitOpError("min/max attributes should not be 'NaN', got min_val=")
|
|
<< minValAttr << ", max_val=" << maxValAttr;
|
|
|
|
if (maxVal < minVal)
|
|
return emitOpError("expected min_val <= max_val, got min_val=")
|
|
<< minValAttr << ", max_val=" << maxValAttr;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Operator Quantization Builders.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// This builder is called on all convolution operators except TransposeConv,
|
|
/// which has specialized output shape semantics. The builder also defines the
|
|
/// bitwidth of the output given the bit width of the input & weight content.
|
|
static void buildConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
|
|
Type outputType, Value input, Value weight,
|
|
Value bias, DenseI64ArrayAttr pad,
|
|
DenseI64ArrayAttr stride,
|
|
DenseI64ArrayAttr dilation,
|
|
TypeAttr accType) {
|
|
auto zps = createZPsAsConst(builder, input, weight);
|
|
result.addOperands({input, weight, bias, zps.first, zps.second});
|
|
result.addAttribute("pad", pad);
|
|
result.addAttribute("stride", stride);
|
|
result.addAttribute("dilation", dilation);
|
|
result.addAttribute("acc_type", accType);
|
|
Type finalOutputType = outputType;
|
|
auto quantAttr = buildConvOpQuantizationAttr(builder, input, weight);
|
|
if (quantAttr) {
|
|
finalOutputType =
|
|
buildConvOpResultTypeInfo(builder, outputType, input, weight);
|
|
}
|
|
result.addTypes(finalOutputType);
|
|
}
|
|
|
|
/// Handles tosa.transpose_conv2d which has outpad and output shape
|
|
/// attributes.
|
|
static void
|
|
buildTransConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
|
|
Type outputType, Value input, Value weight,
|
|
Value bias, DenseI64ArrayAttr outpad,
|
|
DenseI64ArrayAttr stride, TypeAttr accType) {
|
|
auto zps = createZPsAsConst(builder, input, weight);
|
|
result.addOperands({input, weight, bias, zps.first, zps.second});
|
|
result.addAttribute("out_pad", outpad);
|
|
result.addAttribute("stride", stride);
|
|
result.addAttribute("acc_type", accType);
|
|
Type finalOutputType = outputType;
|
|
auto quantAttr = buildConvOpQuantizationAttr(builder, input, weight);
|
|
if (quantAttr) {
|
|
finalOutputType =
|
|
buildConvOpResultTypeInfo(builder, outputType, input, weight);
|
|
}
|
|
result.addTypes(finalOutputType);
|
|
}
|
|
|
|
/// The tosa.matmul op is also intended to be generated where a fully_connected
|
|
/// op must be constructed where the weight is not a constant. In this case,
|
|
/// the fully_connected op must be expressed using matmul.
|
|
/// TODO: Add link to the leglization document explaining this.
|
|
static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
|
|
OperationState &result, Type outputType,
|
|
Value a, Value b) {
|
|
auto zps = createZPsAsConst(builder, a, b);
|
|
result.addOperands({a, b, zps.first, zps.second});
|
|
|
|
Type finalOutputType{outputType};
|
|
if (auto quantAttr = buildMatMulOpQuantizationAttr(builder, a, b)) {
|
|
auto eType = getStorageElementTypeOrSelf(a.getType());
|
|
auto inputBits = eType.getIntOrFloatBitWidth();
|
|
|
|
auto outputShapedType = llvm::dyn_cast<ShapedType>(outputType);
|
|
assert(outputShapedType && "Output must be a shaped type");
|
|
|
|
IntegerType accElementType;
|
|
if (inputBits == 16)
|
|
accElementType = builder.getIntegerType(48);
|
|
else
|
|
accElementType = builder.getI32Type();
|
|
|
|
finalOutputType = outputShapedType.clone(accElementType);
|
|
}
|
|
result.addTypes(finalOutputType);
|
|
}
|
|
|
|
/// Both the tosa.avg_pool2d and unary ops use the same
|
|
/// UnaryOpQuantizationAttr but avg_pool operator has its own builder as it
|
|
/// has additional parameters not part of the unary ops.
|
|
static void
|
|
buildAvgPool2dOpWithQuantInfo(OpBuilder &builder, OperationState &result,
|
|
Type outputType, Value input,
|
|
DenseArrayAttr kernel, DenseArrayAttr stride,
|
|
DenseArrayAttr pad, TypeAttr accType) {
|
|
const Location loc{result.location};
|
|
int64_t inputZp{0};
|
|
int64_t outputZp{0};
|
|
|
|
if (auto quantAttr =
|
|
buildUnaryOpQuantizationAttr(builder, input, outputType)) {
|
|
inputZp = quantAttr.getInputZp();
|
|
outputZp = quantAttr.getOutputZp();
|
|
}
|
|
const std::optional<Value> inputZpOp =
|
|
createZeroPointTensor(builder, loc, input.getType(), inputZp);
|
|
if (!inputZpOp) {
|
|
(void)emitError(
|
|
loc,
|
|
"Failed to create input zero point tensor for quantized AVG_POOL2D op");
|
|
}
|
|
const std::optional<Value> outputZpOp =
|
|
createZeroPointTensor(builder, loc, outputType, outputZp);
|
|
if (!outputZpOp) {
|
|
(void)emitError(loc, "Failed to create output zero point tensor for "
|
|
"quantized AVG_POOL2D op");
|
|
}
|
|
|
|
if (inputZpOp && outputZpOp) {
|
|
result.addOperands({input, inputZpOp.value(), outputZpOp.value()});
|
|
} else {
|
|
// failed to create one or more zero points above: just add input as
|
|
// operands this will trigger error in building the op because of missing
|
|
// zero points
|
|
result.addOperands({input});
|
|
}
|
|
result.addAttribute("kernel", kernel);
|
|
result.addAttribute("stride", stride);
|
|
result.addAttribute("pad", pad);
|
|
result.addAttribute("acc_type", accType);
|
|
result.types.push_back(outputType);
|
|
}
|
|
|
|
/// This builder mirrors avg_pool2d quant-info handling and materializes
|
|
/// kernel/stride/pad as const_shape operands for avg_pool2d_adaptive.
|
|
static void buildAvgPool2dAdaptiveOpWithQuantInfo(
|
|
OpBuilder &builder, OperationState &result, Type outputType, Value input,
|
|
DenseI64ArrayAttr kernel, DenseI64ArrayAttr stride, DenseI64ArrayAttr pad,
|
|
TypeAttr accType) {
|
|
const Location loc{result.location};
|
|
int64_t inputZp{0};
|
|
int64_t outputZp{0};
|
|
|
|
if (auto quantAttr =
|
|
buildUnaryOpQuantizationAttr(builder, input, outputType)) {
|
|
inputZp = quantAttr.getInputZp();
|
|
outputZp = quantAttr.getOutputZp();
|
|
}
|
|
const std::optional<Value> inputZpOp =
|
|
createZeroPointTensor(builder, loc, input.getType(), inputZp);
|
|
if (!inputZpOp) {
|
|
(void)emitError(loc,
|
|
"Failed to create input zero point tensor for quantized "
|
|
"AVG_POOL2D_ADAPTIVE op");
|
|
}
|
|
const std::optional<Value> outputZpOp =
|
|
createZeroPointTensor(builder, loc, outputType, outputZp);
|
|
if (!outputZpOp) {
|
|
(void)emitError(loc, "Failed to create output zero point tensor for "
|
|
"quantized AVG_POOL2D_ADAPTIVE op");
|
|
}
|
|
|
|
if (inputZpOp && outputZpOp) {
|
|
ImplicitLocOpBuilder b(loc, builder);
|
|
Value kernelShape = getTosaConstShape(b, kernel.asArrayRef());
|
|
Value strideShape = getTosaConstShape(b, stride.asArrayRef());
|
|
Value padShape = getTosaConstShape(b, pad.asArrayRef());
|
|
result.addOperands({input, inputZpOp.value(), outputZpOp.value(),
|
|
kernelShape, strideShape, padShape});
|
|
} else {
|
|
// Failed to create one or more zero points above: just add input as
|
|
// operands. This will trigger error in building the op because of missing
|
|
// operands.
|
|
result.addOperands({input});
|
|
}
|
|
result.addAttribute("acc_type", accType);
|
|
result.types.push_back(outputType);
|
|
}
|
|
|
|
/// This builder is called on single-parameter negate operator
|
|
/// to construct input and output zero points based on their
|
|
/// types.
|
|
static void buildNegateOpWithQuantInfo(OpBuilder &builder,
|
|
OperationState &result, Type outputType,
|
|
Value input) {
|
|
const Location loc{result.location};
|
|
int64_t input1Zp{0};
|
|
int64_t outputZp{0};
|
|
auto quantAttr = buildUnaryOpQuantizationAttr(builder, input, outputType);
|
|
if (quantAttr) {
|
|
input1Zp = quantAttr.getInputZp();
|
|
outputZp = quantAttr.getOutputZp();
|
|
}
|
|
const std::optional<Value> input1ZpOp =
|
|
createZeroPointTensor(builder, loc, input.getType(), input1Zp);
|
|
if (!input1ZpOp) {
|
|
(void)emitError(
|
|
loc, "Failed to create input1 zero point for quantized NEGATE op");
|
|
}
|
|
|
|
const std::optional<Value> outputZpOp =
|
|
createZeroPointTensor(builder, loc, input.getType(), outputZp);
|
|
if (!outputZpOp) {
|
|
(void)emitError(
|
|
loc, "Failed to create output zero point for quantized NEGATE op");
|
|
}
|
|
|
|
if (input1ZpOp && outputZpOp) {
|
|
result.addOperands({input, input1ZpOp.value(), outputZpOp.value()});
|
|
} else {
|
|
// failed to create one or more zero points above: just add input as
|
|
// operands. This will trigger error in building the op because of
|
|
// missing zero points
|
|
result.addOperands({input});
|
|
}
|
|
|
|
result.types.push_back(outputType);
|
|
}
|
|
|
|
/// This builder is called on TOSA pad operator that needs to create its own
|
|
/// OptionalAttr quantization_attr parameter to scale the padding values
|
|
/// correctly. No pad_const is interpreted as zero-padding.
|
|
static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
|
|
Type outputType, Value input,
|
|
Value paddings) {
|
|
const Location loc{result.location};
|
|
int32_t zp{0};
|
|
const auto quantAttr = buildPadOpQuantizationAttr(builder, input);
|
|
if (quantAttr) {
|
|
zp = static_cast<int32_t>(quantAttr.getInputZp());
|
|
}
|
|
const auto padConstOp{createPadConstTensor(builder, loc, input, zp)};
|
|
result.addOperands({input, paddings, padConstOp});
|
|
result.types.push_back(outputType);
|
|
}
|
|
|
|
static void buildVariableOp(OpBuilder &builder, OperationState &result,
|
|
StringRef name, Type variableType,
|
|
Attribute initialValue) {
|
|
const Location loc{result.location};
|
|
auto nameAttr = builder.getStringAttr(name);
|
|
|
|
auto shapedType = dyn_cast<ShapedType>(variableType);
|
|
if (!shapedType) {
|
|
(void)emitError(loc, "variable type must be a shaped type");
|
|
return;
|
|
}
|
|
if (!shapedType.hasRank()) {
|
|
(void)emitError(loc, "variable type must be a ranked type");
|
|
return;
|
|
}
|
|
|
|
auto elementType = shapedType.getElementType();
|
|
auto elementTypeAttr = TypeAttr::get(elementType);
|
|
ArrayRef<int64_t> shape = shapedType.getShape();
|
|
auto varShapeAttr = builder.getIndexTensorAttr(convertFromMlirShape(shape));
|
|
|
|
result.addAttribute("sym_name", nameAttr);
|
|
result.addAttribute("var_shape", varShapeAttr);
|
|
result.addAttribute("type", elementTypeAttr);
|
|
result.addAttribute("initial_value", initialValue);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Operator Return Type Inference.
|
|
//===----------------------------------------------------------------------===//
|
|
static FailureOr<int64_t> resolveBroadcastDim(const int64_t dim1,
|
|
const int64_t dim2) {
|
|
if (dim1 == 1)
|
|
return dim2;
|
|
if (dim2 == 1)
|
|
return dim1;
|
|
|
|
if (ShapedType::isStatic(dim1) && ShapedType::isStatic(dim2) && dim1 != dim2)
|
|
return failure();
|
|
|
|
// Prefer static dimension over dynamic
|
|
return ShapedType::isDynamic(dim1) ? dim2 : dim1;
|
|
}
|
|
|
|
static LogicalResult resolveBroadcastShape(const ValueShapeRange &operands,
|
|
SmallVector<int64_t> &outShape) {
|
|
int64_t outRank = 0;
|
|
for (int i = 0, e = operands.size(); i != e; ++i) {
|
|
auto shape = operands.getShape(i);
|
|
if (!shape.hasRank()) {
|
|
// TODO(jennik): Update function to have better case handling for
|
|
// invalid operands and for ranked tensors.
|
|
return failure();
|
|
}
|
|
outRank = std::max<int64_t>(outRank, shape.getRank());
|
|
}
|
|
|
|
outShape.resize(outRank, 1);
|
|
|
|
for (int i = 0, e = operands.size(); i != e; ++i) {
|
|
auto shape = operands.getShape(i);
|
|
auto rankDiff = outShape.size() - shape.getRank();
|
|
|
|
for (size_t i = 0, e = shape.getRank(); i < e; ++i) {
|
|
auto dim1 = outShape[i + rankDiff];
|
|
auto dim2 = shape.getDimSize(i);
|
|
|
|
const FailureOr<int64_t> maybeResolvedDim =
|
|
resolveBroadcastDim(dim1, dim2);
|
|
if (failed(maybeResolvedDim))
|
|
return failure();
|
|
const int64_t resolvedDim = *maybeResolvedDim;
|
|
outShape[i + rankDiff] = resolvedDim;
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ArgMaxOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
IntegerAttr axis = adaptor.getProperties().axis;
|
|
int32_t axisVal = axis.getValue().getSExtValue();
|
|
|
|
if (!inputShape.hasRank()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents());
|
|
return success();
|
|
}
|
|
|
|
SmallVector<int64_t> outShape;
|
|
outShape.reserve(inputShape.getRank() - 1);
|
|
for (int i = 0, s = inputShape.getRank(); i < s; i++) {
|
|
if (i == axisVal)
|
|
continue;
|
|
outShape.push_back(inputShape.getDimSize(i));
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::RFFT2dOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
RFFT2dOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInputReal().getType());
|
|
|
|
if (!inputShape.hasRank())
|
|
return failure();
|
|
|
|
llvm::SmallVector<int64_t> outputShape;
|
|
outputShape.resize(3, ShapedType::kDynamic);
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
outputShape[1] = inputShape.getDimSize(1);
|
|
int64_t inWidth = inputShape.getDimSize(2);
|
|
|
|
// Note that we can support this calculation symbolically
|
|
// in the future e.g. [x, y, z] -> [x, y, z / 2 + 1]
|
|
if (inWidth != ShapedType::kDynamic)
|
|
outputShape[2] = inWidth / 2 + 1;
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyDimIsPowerOfTwo(Operation *op, const int64_t dimSize,
|
|
const llvm::StringRef dimName) {
|
|
const bool isPowerOfTwo = (dimSize & (dimSize - 1)) == 0 && dimSize > 0;
|
|
if (!isPowerOfTwo)
|
|
return op->emitOpError("expected ")
|
|
<< dimName << " to be a power of two, got " << dimSize;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::RFFT2dOp::verify() {
|
|
const auto outputTypes = getResultTypes();
|
|
if (failed(verifyCompatibleShapes(outputTypes)))
|
|
return emitOpError("expected output shapes to match, got ") << outputTypes;
|
|
|
|
const auto inputType =
|
|
llvm::dyn_cast<RankedTensorType>(getInputReal().getType());
|
|
if (!inputType)
|
|
return success();
|
|
|
|
const int64_t height = inputType.getDimSize(1);
|
|
if (ShapedType::isStatic(height) &&
|
|
failed(verifyDimIsPowerOfTwo(*this, height, "height")))
|
|
return failure();
|
|
|
|
const int64_t width = inputType.getDimSize(2);
|
|
if (ShapedType::isStatic(width) &&
|
|
failed(verifyDimIsPowerOfTwo(*this, width, "width")))
|
|
return failure();
|
|
|
|
const auto outputType = llvm::dyn_cast<RankedTensorType>(outputTypes[0]);
|
|
if (!outputType)
|
|
return success();
|
|
|
|
// Batch and height input/output dimensions should match
|
|
if (failed(verifyCompatibleShape(inputType.getShape().drop_back(),
|
|
outputType.getShape().drop_back())))
|
|
return emitOpError("expected batch and height dimensions of input/output "
|
|
"to match, got input=")
|
|
<< inputType << " output=" << outputType;
|
|
|
|
// Output width dimension expected to be input_width / 2 + 1
|
|
const int64_t outputWidth = outputType.getDimSize(2);
|
|
if (ShapedType::isStatic(width) && ShapedType::isStatic(outputWidth) &&
|
|
(outputWidth != (width / 2) + 1))
|
|
return emitOpError(
|
|
"expected output width to be equal to input_width / 2 + 1, got ")
|
|
<< outputWidth;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::FFT2dOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
FFT2dOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
inferredReturnShapes.push_back(
|
|
ShapedTypeComponents(ShapeAdaptor(adaptor.getInputReal().getType())));
|
|
inferredReturnShapes.push_back(
|
|
ShapedTypeComponents(ShapeAdaptor(adaptor.getInputImag().getType())));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::FFT2dOp::verify() {
|
|
const auto inputRealType =
|
|
llvm::dyn_cast<RankedTensorType>(getInputReal().getType());
|
|
const auto inputImagType =
|
|
llvm::dyn_cast<RankedTensorType>(getInputImag().getType());
|
|
if (!inputRealType || !inputImagType)
|
|
return success();
|
|
|
|
const auto trySelectStaticDim = [](const int64_t a, const int64_t b) {
|
|
return ShapedType::isDynamic(a) ? a : b;
|
|
};
|
|
|
|
const int64_t height = trySelectStaticDim(inputRealType.getDimSize(1),
|
|
inputImagType.getDimSize(1));
|
|
if (ShapedType::isStatic(height) &&
|
|
failed(verifyDimIsPowerOfTwo(*this, height, "height")))
|
|
return failure();
|
|
|
|
const int64_t width = trySelectStaticDim(inputRealType.getDimSize(2),
|
|
inputImagType.getDimSize(2));
|
|
if (ShapedType::isStatic(width) &&
|
|
failed(verifyDimIsPowerOfTwo(*this, width, "width")))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ConcatOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
// Infer all dimension sizes by reducing based on inputs.
|
|
const Properties &prop = adaptor.getProperties();
|
|
int32_t axis = prop.axis.getValue().getSExtValue();
|
|
llvm::SmallVector<int64_t> outputShape;
|
|
bool hasRankedInput = false;
|
|
for (auto operand : adaptor.getOperands()) {
|
|
ShapeAdaptor operandShape(operand.getType());
|
|
if (!operandShape.hasRank())
|
|
continue;
|
|
|
|
// Copy the Operand's rank.
|
|
if (!hasRankedInput)
|
|
outputShape.resize(operandShape.getRank(), ShapedType::kDynamic);
|
|
|
|
// Copy shapes until the dim is non-dynamic.
|
|
for (int i = 0, s = operandShape.getRank(); i < s; i++) {
|
|
if (i == axis || operandShape.isDynamicDim(i))
|
|
continue;
|
|
if (outputShape[i] == ShapedType::kDynamic)
|
|
outputShape[i] = operandShape.getDimSize(i);
|
|
if (outputShape[i] != operandShape.getDimSize(i))
|
|
return emitOptionalError(location,
|
|
"Cannot concat tensors with different sizes"
|
|
" on the non-axis dimension ",
|
|
i);
|
|
}
|
|
|
|
hasRankedInput = true;
|
|
}
|
|
|
|
if (adaptor.getInput1().empty())
|
|
return failure();
|
|
|
|
Type inputType =
|
|
llvm::cast<TensorType>(adaptor.getInput1().getType()[0]).getElementType();
|
|
if (!hasRankedInput) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(inputType));
|
|
return success();
|
|
}
|
|
|
|
// Determine the dimension size along the concatenation axis.
|
|
int64_t concatDimSize = 0;
|
|
for (auto operand : adaptor.getOperands()) {
|
|
ShapeAdaptor operandShape(operand.getType());
|
|
|
|
// We need to know the length of the concatenation axis of all inputs to
|
|
// determine the dimension size of the output shape.
|
|
if (!operandShape.hasRank() || operandShape.isDynamicDim(axis)) {
|
|
concatDimSize = ShapedType::kDynamic;
|
|
break;
|
|
}
|
|
|
|
concatDimSize += operandShape.getDimSize(axis);
|
|
}
|
|
|
|
outputShape[axis] = concatDimSize;
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ConcatOp::verify() {
|
|
// check that each input has same element type as output
|
|
auto outType = getOutput().getType();
|
|
const Operation::operand_range inputList = getInput1();
|
|
|
|
// Check there is at least one input
|
|
if (inputList.empty())
|
|
return emitOpError("expect at least one input");
|
|
|
|
if (!llvm::all_of(inputList, [&](auto input) {
|
|
return succeeded(verifySameElementTypes(
|
|
*this, /* inType = */ input.getType(), outType));
|
|
})) {
|
|
return failure();
|
|
}
|
|
|
|
const int32_t axis = getAxis();
|
|
ShapeAdaptor firstRankedInputShape = nullptr;
|
|
for (const auto &input : inputList) {
|
|
const Type inputType = input.getType();
|
|
ShapeAdaptor currShape(inputType);
|
|
if (currShape.hasRank()) {
|
|
firstRankedInputShape = currShape;
|
|
// Check axis is in expected range
|
|
if (axis < 0 || axis >= firstRankedInputShape.getRank())
|
|
return emitOpError("expect axis to be within range 0 < axis < "
|
|
"rank(input1[firstRankedTensorIdx]), got ")
|
|
<< axis;
|
|
break;
|
|
}
|
|
}
|
|
|
|
const auto allOperandsHasRank = [](const Value input) {
|
|
return ShapeAdaptor(input.getType()).hasRank();
|
|
};
|
|
if (llvm::all_of(inputList, allOperandsHasRank)) {
|
|
const int64_t firstInputRank = firstRankedInputShape.getRank();
|
|
|
|
for (const auto &[index, input] : llvm::enumerate(inputList.drop_front())) {
|
|
const ShapeAdaptor inputShape(input.getType());
|
|
const int64_t inputRank = inputShape.getRank();
|
|
const size_t operandNum = index + 1;
|
|
|
|
// Check that each operand has the same rank
|
|
if (inputRank != firstInputRank)
|
|
return emitOpError(
|
|
"expect all operands to have the same rank, but got ")
|
|
<< firstInputRank << " vs " << inputRank << " on operands 0 and "
|
|
<< operandNum;
|
|
|
|
// Check non-axis dims match
|
|
for (int i = 0; i < inputRank; i++) {
|
|
const int64_t inputDim = inputShape.getDimSize(i);
|
|
const int64_t firstInputDim = firstRankedInputShape.getDimSize(i);
|
|
if (i == axis || firstRankedInputShape.isDynamicDim(i) ||
|
|
inputShape.isDynamicDim(i))
|
|
continue;
|
|
if (inputDim != firstInputDim)
|
|
return emitOpError("expect all operand shapes to have the same sizes "
|
|
"on non-axis dimensions, but got ")
|
|
<< inputDim << " vs " << firstInputDim << " at index " << i
|
|
<< " on operands 0 and " << operandNum;
|
|
}
|
|
}
|
|
|
|
const ShapeAdaptor outputShape(outType);
|
|
if (outputShape.hasRank() && outputShape.getRank() != firstInputRank)
|
|
return emitOpError("expect output rank to match inputs rank, got ")
|
|
<< outputShape.getRank() << " vs " << firstInputRank;
|
|
|
|
// ERROR_IF(axis_sum != shape[axis]);
|
|
int64_t axisSum = 0;
|
|
for (const auto &input : inputList) {
|
|
const ShapeAdaptor inputShape(input.getType());
|
|
if (inputShape.isDynamicDim(axis)) {
|
|
// make axisSum negative to indicate invalid value
|
|
axisSum = -1;
|
|
break;
|
|
}
|
|
axisSum += inputShape.getDimSize(axis);
|
|
}
|
|
|
|
if (axisSum >= 0 && outputShape.hasRank() &&
|
|
!outputShape.isDynamicDim(axis) &&
|
|
axisSum != outputShape.getDimSize(axis))
|
|
return emitOpError("requires sum of axis dimensions of input1 "
|
|
"equal to output axis dimension, got ")
|
|
<< axisSum << " and " << outputShape.getDimSize(axis);
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::EqualOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ValueShapeRange operands, DictionaryAttr attributes, PropertyRef properties,
|
|
RegionRange regions,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
auto elementType = IntegerType::get(context, /*width=*/1);
|
|
|
|
llvm::SmallVector<int64_t> outShape;
|
|
if (resolveBroadcastShape(operands, outShape).failed()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(elementType));
|
|
return success();
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outShape, elementType));
|
|
return success();
|
|
}
|
|
|
|
bool tosa::EqualOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
|
|
if (l.size() != r.size() || l.size() != 1)
|
|
return false;
|
|
return succeeded(verifyCompatibleShape(l[0], r[0]));
|
|
}
|
|
|
|
LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
MatMulOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor lhsShape(adaptor.getA().getType());
|
|
ShapeAdaptor rhsShape(adaptor.getB().getType());
|
|
|
|
// All shapes are dynamic.
|
|
SmallVector<int64_t> outShape;
|
|
outShape.resize(3, ShapedType::kDynamic);
|
|
|
|
if (lhsShape.hasRank()) {
|
|
outShape[0] = lhsShape.getDimSize(0);
|
|
outShape[1] = lhsShape.getDimSize(1);
|
|
}
|
|
|
|
if (rhsShape.hasRank()) {
|
|
outShape[0] = outShape[0] == ShapedType::kDynamic ? rhsShape.getDimSize(0)
|
|
: outShape[0];
|
|
outShape[2] = rhsShape.getDimSize(2);
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult MatMulOp::verify() {
|
|
const ShapeAdaptor aShape(getA().getType());
|
|
const ShapeAdaptor bShape(getB().getType());
|
|
const Type aElementType = aShape.getElementType();
|
|
const Type bElementType = bShape.getElementType();
|
|
|
|
const auto aQuantizedEType =
|
|
llvm::dyn_cast<quant::UniformQuantizedType>(aElementType);
|
|
const auto bQuantizedEType =
|
|
llvm::dyn_cast<quant::UniformQuantizedType>(bElementType);
|
|
|
|
if (aQuantizedEType || bQuantizedEType) {
|
|
if (!aQuantizedEType || !bQuantizedEType) {
|
|
return emitOpError("expect operands to be both quantized or both not "
|
|
"quantized, got ")
|
|
<< aElementType << " and " << bElementType;
|
|
}
|
|
// both a and b have quantized element types
|
|
auto aQuantWidth = aQuantizedEType.getStorageTypeIntegralWidth();
|
|
auto bQuantWidth = bQuantizedEType.getStorageTypeIntegralWidth();
|
|
if (aQuantWidth != bQuantWidth) {
|
|
return emitOpError("expect quantized operands to have same widths, got ")
|
|
<< aQuantWidth << " and " << bQuantWidth;
|
|
}
|
|
}
|
|
|
|
// check a_zp and b_zp
|
|
auto aEType = getStorageElementTypeOrSelf(aElementType);
|
|
auto aZpEType = getStorageElementTypeOrSelf(getAZp().getType());
|
|
if (aEType != aZpEType)
|
|
return emitOpError("expect input a and a_zp have the same "
|
|
"element type, got ")
|
|
<< aEType << " and " << aZpEType;
|
|
|
|
const Type bEType = getStorageElementTypeOrSelf(bElementType);
|
|
const Type bZpEType = getStorageElementTypeOrSelf(getBZp().getType());
|
|
if (bEType != bZpEType)
|
|
return emitOpError("expect input b and b_zp have the same "
|
|
"element type, got ")
|
|
<< bEType << " and " << bZpEType;
|
|
|
|
FailureOr<int64_t> maybeAZp = getAZeroPoint();
|
|
if (succeeded(maybeAZp) && verifyAZeroPoint(*maybeAZp).failed())
|
|
return failure();
|
|
|
|
FailureOr<int64_t> maybeBZp = getBZeroPoint();
|
|
if (succeeded(maybeBZp) && verifyBZeroPoint(*maybeBZp).failed())
|
|
return failure();
|
|
|
|
// Verify input/output shapes
|
|
int64_t N = ShapedType::kDynamic;
|
|
int64_t H = ShapedType::kDynamic;
|
|
int64_t W = ShapedType::kDynamic;
|
|
int64_t C = ShapedType::kDynamic;
|
|
|
|
if (aShape.hasRank()) {
|
|
N = aShape.getDimSize(0);
|
|
H = aShape.getDimSize(1);
|
|
C = aShape.getDimSize(2);
|
|
}
|
|
|
|
if (bShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, N, bShape.getDimSize(0), "b",
|
|
"batch")) ||
|
|
failed(tryUpdateDimOrFailure(*this, C, bShape.getDimSize(1), "b",
|
|
"channels")))
|
|
return failure();
|
|
W = bShape.getDimSize(2);
|
|
}
|
|
|
|
const SmallVector<int64_t, 3> expectedOutputShape = {N, H, W};
|
|
const auto outputType = cast<ShapedType>(getResult().getType());
|
|
if (outputType.hasRank() &&
|
|
failed(
|
|
verifyCompatibleShape(outputType.getShape(), expectedOutputShape))) {
|
|
InFlightDiagnostic opError = emitOpError("expected output shape ");
|
|
printShapeToDiagnostic(opError, outputType.getShape());
|
|
opError << " to be compatible with expected output shape ";
|
|
printShapeToDiagnostic(opError, expectedOutputShape);
|
|
return opError;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::MatmulTBlockScaledOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
MatmulTBlockScaledOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
SmallVector<int64_t, 3> outShape(3, ShapedType::kDynamic);
|
|
|
|
const auto aDataShape = cast<ShapedType>(adaptor.getAData().getType());
|
|
if (aDataShape.hasRank()) {
|
|
outShape[0] = aDataShape.getDimSize(0);
|
|
outShape[1] = aDataShape.getDimSize(1);
|
|
}
|
|
|
|
const auto aScaleShape = cast<ShapedType>(adaptor.getAScale().getType());
|
|
if (aScaleShape.hasRank()) {
|
|
outShape[0] = ShapedType::isDynamic(outShape[0]) ? aScaleShape.getDimSize(0)
|
|
: outShape[0];
|
|
outShape[1] = ShapedType::isDynamic(outShape[1]) ? aScaleShape.getDimSize(1)
|
|
: outShape[1];
|
|
}
|
|
|
|
// If B batch size is 1, it is broadcast across A's batch size
|
|
const auto bDataShape = cast<ShapedType>(adaptor.getBData().getType());
|
|
if (bDataShape.hasRank()) {
|
|
const int64_t bDataBatchSize = bDataShape.getDimSize(0);
|
|
if (bDataBatchSize != 1)
|
|
outShape[0] =
|
|
ShapedType::isDynamic(outShape[0]) ? bDataBatchSize : outShape[0];
|
|
outShape[2] = bDataShape.getDimSize(1);
|
|
}
|
|
|
|
const auto bScaleShape = cast<ShapedType>(adaptor.getBScale().getType());
|
|
if (bScaleShape.hasRank()) {
|
|
const int64_t bScaleBatchSize = bScaleShape.getDimSize(0);
|
|
if (bScaleBatchSize != 1)
|
|
outShape[0] =
|
|
ShapedType::isDynamic(outShape[0]) ? bScaleBatchSize : outShape[0];
|
|
outShape[2] = ShapedType::isDynamic(outShape[2]) ? bScaleShape.getDimSize(1)
|
|
: outShape[2];
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult MatmulTBlockScaledOp::verify() {
|
|
// Verify same input data types
|
|
const Type aDataType = getAData().getType();
|
|
const Type bDataType = getBData().getType();
|
|
if (failed(verifySameElementTypes(*this, aDataType, bDataType, "A_data",
|
|
"B_data")))
|
|
return failure();
|
|
|
|
// Verify input shape compatibility
|
|
int64_t N = ShapedType::kDynamic;
|
|
int64_t D = ShapedType::kDynamic;
|
|
int64_t H = ShapedType::kDynamic;
|
|
int64_t W = ShapedType::kDynamic;
|
|
int64_t C = ShapedType::kDynamic;
|
|
int64_t multiplesOfC = ShapedType::kDynamic;
|
|
|
|
const ShapeAdaptor aDataShape = ShapeAdaptor(aDataType);
|
|
if (aDataShape.hasRank()) {
|
|
N = aDataShape.getDimSize(0);
|
|
H = aDataShape.getDimSize(1);
|
|
C = aDataShape.getDimSize(2);
|
|
}
|
|
|
|
const ShapeAdaptor aScaleShape = ShapeAdaptor(getAScale().getType());
|
|
if (aScaleShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, N, aScaleShape.getDimSize(0),
|
|
"a_scale", "batch")) ||
|
|
failed(tryUpdateDimOrFailure(*this, H, aScaleShape.getDimSize(1),
|
|
"a_scale", "height")))
|
|
return failure();
|
|
multiplesOfC = aScaleShape.getDimSize(2);
|
|
}
|
|
|
|
const ShapeAdaptor bDataShape = ShapeAdaptor(bDataType);
|
|
if (bDataShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, D, bDataShape.getDimSize(0),
|
|
"b_data", "batch")) ||
|
|
failed(tryUpdateDimOrFailure(*this, C, bDataShape.getDimSize(2),
|
|
"b_data", "channels")))
|
|
return failure();
|
|
W = bDataShape.getDimSize(1);
|
|
}
|
|
|
|
const ShapeAdaptor bScaleShape = ShapeAdaptor(getBScale().getType());
|
|
if (bScaleShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, D, bScaleShape.getDimSize(0),
|
|
"b_scale", "batch")) ||
|
|
failed(tryUpdateDimOrFailure(*this, W, bScaleShape.getDimSize(1),
|
|
"b_scale", "width")) ||
|
|
failed(tryUpdateDimOrFailure(*this, multiplesOfC,
|
|
bScaleShape.getDimSize(2), "b_scale",
|
|
"C/block_size")))
|
|
return failure();
|
|
}
|
|
|
|
// Verify batch size is broadcast compatible
|
|
if (ShapedType::isStatic(N) && ShapedType::isStatic(D) && N != D && D != 1)
|
|
return emitOpError("expect B matrix batch size to be broadcast compatible "
|
|
"with A, got D=")
|
|
<< D << " vs N=" << N;
|
|
|
|
// Verify C is a multiple of block size
|
|
const uint32_t blockSize = BlockSizeAttr::getBlockSizeValue(getBlockSize());
|
|
if (blockSize != BlockSizeAttr::getBlockSizeValue(BlockSize::BLOCK_SIZE_32))
|
|
return emitOpError("expect block size to be 32, got ") << blockSize;
|
|
if (ShapedType::isStatic(C) && C % blockSize != 0)
|
|
return emitOpError("expect C to be a multiple of block size, got C=")
|
|
<< C << ", block_size=" << blockSize;
|
|
|
|
// Verify multiplesOfC is C / block size
|
|
if (ShapedType::isStatic(C) && ShapedType::isStatic(multiplesOfC) &&
|
|
multiplesOfC != C / blockSize)
|
|
return emitOpError(
|
|
"expect scale operands dimension 2 to equal C/block_size (")
|
|
<< C << "/" << blockSize << ")" << ", got " << multiplesOfC;
|
|
|
|
// Verify output shape
|
|
N = ShapedType::isDynamic(N) ? D : N;
|
|
const SmallVector<int64_t, 3> expectedOutputShape = {N, H, W};
|
|
const auto outputType = cast<ShapedType>(getResult().getType());
|
|
if (outputType.hasRank() &&
|
|
failed(
|
|
verifyCompatibleShape(outputType.getShape(), expectedOutputShape))) {
|
|
InFlightDiagnostic opError = emitOpError("expected output shape ");
|
|
printShapeToDiagnostic(opError, outputType.getShape());
|
|
opError << " to be compatible with expected output shape ";
|
|
printShapeToDiagnostic(opError, expectedOutputShape);
|
|
return opError;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::PadOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
PadOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
auto paddingRank =
|
|
cast<tosa::shapeType>(adaptor.getPadding().getType()).getRank();
|
|
SmallVector<int64_t> outputShape;
|
|
|
|
// If the input rank is unknown, we can infer the output rank using the
|
|
// padding shape's rank divided by 2.
|
|
if (!inputShape.hasRank()) {
|
|
outputShape.resize(paddingRank / 2, ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
SmallVector<int64_t> paddingValues;
|
|
// If the paddings value is not a constant, all dimensions must be dynamic.
|
|
if (!tosa::getConstShapeValues(adaptor.getPadding().getDefiningOp(),
|
|
paddingValues)) {
|
|
outputShape.resize(inputShape.getRank(), ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
outputShape.reserve(inputShape.getRank());
|
|
for (int i = 0, s = inputShape.getRank(); i < s; i++) {
|
|
if (inputShape.isDynamicDim(i)) {
|
|
outputShape.push_back(ShapedType::kDynamic);
|
|
continue;
|
|
}
|
|
auto padFront = paddingValues[i * 2];
|
|
auto padBack = paddingValues[i * 2 + 1];
|
|
if (padFront < 0 || padBack < 0) {
|
|
// if either padding for dim i is -1, output dim is unknown
|
|
outputShape.push_back(ShapedType::kDynamic);
|
|
continue;
|
|
}
|
|
|
|
outputShape.push_back(inputShape.getDimSize(i) + padFront + padBack);
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::PadOp::verify() {
|
|
if (verifySameElementTypes(*this, /* inType = */ getInput1().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
|
|
if (auto padConst = getPadConst()) {
|
|
if (verifySameElementTypes(*this, /* inType = */ padConst.getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
RankedTensorType inputType =
|
|
llvm::dyn_cast<RankedTensorType>(getInput1().getType());
|
|
RankedTensorType outputType =
|
|
llvm::dyn_cast<RankedTensorType>(getOutput().getType());
|
|
if (!inputType || !outputType)
|
|
return success();
|
|
|
|
if (failed(verifyRanksMatch(getOperation(), inputType, outputType, "input",
|
|
"output")))
|
|
return failure();
|
|
|
|
auto inputRank = inputType.getRank();
|
|
DenseIntElementsAttr paddingAttr;
|
|
if (!matchPattern(getPadding(), m_Constant(&paddingAttr)))
|
|
return success();
|
|
|
|
auto paddingValues = paddingAttr.getValues<APInt>();
|
|
if (paddingValues.size() != static_cast<size_t>(inputRank * 2))
|
|
return emitOpError() << "padding tensor must have " << inputRank
|
|
<< " * 2 = " << inputRank * 2 << " elements, but got "
|
|
<< paddingValues.size();
|
|
|
|
auto inputShape = inputType.getShape();
|
|
auto outputShape = outputType.getShape();
|
|
|
|
for (int64_t i = 0; i < inputRank; ++i) {
|
|
int64_t padStart = paddingValues[i * 2].getSExtValue();
|
|
int64_t padEnd = paddingValues[i * 2 + 1].getSExtValue();
|
|
|
|
if ((padStart < 0 && padStart != -1) || (padEnd < 0 && padEnd != -1)) {
|
|
return emitOpError()
|
|
<< "invalid padding values at dimension " << i
|
|
<< ": values must be non-negative or -1 for dynamic padding, got ["
|
|
<< padStart << ", " << padEnd << "]";
|
|
}
|
|
|
|
// Skip shape verification for dynamic input/output
|
|
if (inputShape[i] == ShapedType::kDynamic ||
|
|
outputShape[i] == ShapedType::kDynamic)
|
|
continue;
|
|
|
|
if (outputShape[i] != inputShape[i] + padStart + padEnd) {
|
|
return emitOpError() << "mismatch in output shape at dimension " << i
|
|
<< ": expected " << inputShape[i] << " + "
|
|
<< padStart << " + " << padEnd << " = "
|
|
<< (inputShape[i] + padStart + padEnd)
|
|
<< ", but got " << outputShape[i];
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::SliceOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
SliceOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
|
|
Type inputType = getElementTypeOrSelf(adaptor.getInput1().getType());
|
|
SmallVector<int64_t> start;
|
|
SmallVector<int64_t> size;
|
|
|
|
if (!tosa::getConstShapeValues(adaptor.getStart().getDefiningOp(), start) ||
|
|
!tosa::getConstShapeValues(adaptor.getSize().getDefiningOp(), size)) {
|
|
auto rank = cast<tosa::shapeType>(adaptor.getSize().getType()).getRank();
|
|
SmallVector<int64_t> fallback(rank, ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(fallback, inputType));
|
|
return success();
|
|
}
|
|
|
|
// if size[i] is -1, all remaining elements in dimension i are included
|
|
// in the slice, similar to TF.
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
// initialize outputShape to all unknown
|
|
SmallVector<int64_t> outputShape(size.size(), ShapedType::kDynamic);
|
|
if (inputShape.hasRank()) {
|
|
for (size_t i = 0; i < size.size(); i++) {
|
|
if (size[i] != 0 && size[i] >= -1 && start[i] >= 0 &&
|
|
(ShapedType::isDynamic(inputShape.getDimSize(i)) ||
|
|
start[i] < inputShape.getDimSize(i))) {
|
|
// size[i] is not 0 and not < -1, and start[i] is in valid range
|
|
if (ShapedType::isDynamic(inputShape.getDimSize(i))) {
|
|
// input shape has unknown dim[i] - only valid if size[i] > 0
|
|
if (size[i] > 0) {
|
|
outputShape[i] = size[i];
|
|
}
|
|
} else {
|
|
// input shape has known dim[i]
|
|
if (size[i] == -1) {
|
|
outputShape[i] = inputShape.getDimSize(i) - start[i];
|
|
} else if (start[i] + size[i] <= inputShape.getDimSize(i)) {
|
|
// start[i] + size[i] is within bound of input shape's dim[i]
|
|
outputShape[i] = size[i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
outputShape = convertToMlirShape(size);
|
|
}
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::SliceOp::verify() {
|
|
const Value input = getInput1();
|
|
const Value output = getOutput();
|
|
if (verifySameElementTypes(*this, /* inType = */ input.getType(),
|
|
/* outType = */ output.getType())
|
|
.failed())
|
|
return failure();
|
|
|
|
const Value start = getStart();
|
|
const Value size = getSize();
|
|
const ShapeAdaptor inputShape(input.getType());
|
|
const ShapeAdaptor outputShape(output.getType());
|
|
|
|
if (inputShape.hasRank()) {
|
|
const auto inputRank = inputShape.getRank();
|
|
if (outputShape.hasRank() && inputRank != outputShape.getRank())
|
|
return emitOpError(
|
|
"expect input1 and output to have the same ranks, got ")
|
|
<< inputRank << " and " << outputShape.getRank();
|
|
|
|
const auto startShapeRank =
|
|
llvm::cast<tosa::shapeType>(start.getType()).getRank();
|
|
if (inputRank != startShapeRank)
|
|
return emitOpError("length of start is not equal to rank of input shape");
|
|
|
|
const auto sizeShapeRank =
|
|
llvm::cast<tosa::shapeType>(size.getType()).getRank();
|
|
if (inputRank != sizeShapeRank)
|
|
return emitOpError("length of size is not equal to rank of input shape");
|
|
}
|
|
|
|
SmallVector<int64_t> startValues;
|
|
tosa::getConstShapeValues(start.getDefiningOp(), startValues);
|
|
if (startValues.size()) {
|
|
if (llvm::any_of(startValues, [](const int64_t v) {
|
|
return v < 0 && v != kInferableDimSize;
|
|
}))
|
|
return emitOpError("start values must be non-negative, got [")
|
|
<< startValues << "]";
|
|
}
|
|
|
|
SmallVector<int64_t> sizeValues;
|
|
if (!tosa::getConstShapeValues(size.getDefiningOp(), sizeValues))
|
|
return success();
|
|
|
|
if (llvm::any_of(sizeValues, [](const int64_t v) {
|
|
return v <= 0 && v != kInferableDimSize;
|
|
}))
|
|
return emitOpError("size values must be > 0, got [") << sizeValues << "]";
|
|
if (outputShape.hasRank()) {
|
|
SmallVector<int64_t> outputDims;
|
|
outputShape.getDims(outputDims);
|
|
const bool hasNoInferableDims = llvm::all_of(
|
|
sizeValues, [](const int64_t v) { return v != kInferableDimSize; });
|
|
if (hasNoInferableDims &&
|
|
failed(verifyCompatibleShape(outputDims, sizeValues)))
|
|
return emitOpError("expected output shape to match size values, got ")
|
|
<< output.getType() << " vs [" << sizeValues << "]";
|
|
}
|
|
|
|
if (inputShape.hasRank() && startValues.size()) {
|
|
SmallVector<int64_t> inputDims;
|
|
inputShape.getDims(inputDims);
|
|
for (const auto &[index, vals] :
|
|
llvm::enumerate(llvm::zip_equal(startValues, sizeValues, inputDims))) {
|
|
const auto &[start, size, inputDim] = vals;
|
|
if (start == kInferableDimSize || size == kInferableDimSize ||
|
|
ShapedType::isDynamic(inputDim))
|
|
continue;
|
|
if (start + size > inputDim)
|
|
return emitOpError("start + size must be less than or equal to input "
|
|
"dimension size, got start=")
|
|
<< start << ", size=" << size
|
|
<< " vs input dim size=" << inputDim << " at dimension "
|
|
<< index;
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::MulOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ValueShapeRange operands, DictionaryAttr attributes, PropertyRef properties,
|
|
RegionRange regions,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
// mul op's output shape only depend on input1 and input2, not on shift
|
|
ValueShapeRange twoInputs = operands.drop_back();
|
|
llvm::SmallVector<int64_t> outShape;
|
|
if (resolveBroadcastShape(twoInputs, outShape).failed()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents());
|
|
} else {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::MulOp::verify() {
|
|
const Value output = getOutput();
|
|
auto resElemType = getElementTypeOrSelf(output);
|
|
|
|
// Verify if the element type among operands and result match tosa
|
|
// specification.
|
|
if (auto resIntType = dyn_cast<IntegerType>(resElemType)) {
|
|
IntegerType lhsIntType =
|
|
dyn_cast<IntegerType>(getElementTypeOrSelf(getInput1()));
|
|
IntegerType rhsIntType =
|
|
dyn_cast<IntegerType>(getElementTypeOrSelf(getInput2()));
|
|
if (!lhsIntType || !rhsIntType || lhsIntType != rhsIntType)
|
|
return emitOpError("requires the same element type for all operands");
|
|
|
|
// Though the spec requires the element type of result to be i32, a more
|
|
// relaxed way is provided at dialect level for easier cooperating with
|
|
// other dialects.
|
|
if (lhsIntType.getWidth() > resIntType.getWidth())
|
|
return emitOpError("invalid data type size for operands or result");
|
|
|
|
} else {
|
|
// For other supported type, the spec requires requires the same element
|
|
// type for all operands (excludes `shift` operand) and results.
|
|
for (int i = 0; i < 2; ++i) {
|
|
if (getElementTypeOrSelf(getOperand(i)) != resElemType)
|
|
return emitOpError(
|
|
"requires the same element type for all operands and results");
|
|
}
|
|
|
|
// verify shift has value 0 for non-integer types
|
|
ElementsAttr shiftElem;
|
|
if (matchPattern(getShift(), m_Constant(&shiftElem))) {
|
|
int32_t shift = shiftElem.getValues<IntegerAttr>()[0].getInt();
|
|
if (shift != 0) {
|
|
return emitOpError() << "require shift to be 0 for float type";
|
|
}
|
|
}
|
|
}
|
|
|
|
// Verify the op has same ranks for all main operands (excludes extra operands
|
|
// such as shift of mul op, so this is the only difference with the built-in
|
|
// `SameOperandsAndResultRank` trait) and results types, if known.
|
|
TypeRange operandTypes = getOperandTypes();
|
|
ShapedType aType = cast<ShapedType>(operandTypes[0]);
|
|
ShapedType bType = cast<ShapedType>(operandTypes[1]);
|
|
|
|
const bool aHasRank = aType.hasRank();
|
|
const bool bHasRank = bType.hasRank();
|
|
|
|
bool hasExpectedOutputShape = false;
|
|
SmallVector<int64_t> expectedOutputShape;
|
|
|
|
if (aHasRank && bHasRank) {
|
|
const int64_t aRank = aType.getRank();
|
|
const int64_t bRank = bType.getRank();
|
|
if (aRank != bRank)
|
|
return emitOpError("a and b operands don't have matching ranks, got ")
|
|
<< aRank << " and " << bRank;
|
|
|
|
// check for broadcast compatible shapes
|
|
if (!mlir::OpTrait::util::getBroadcastedShape(
|
|
aType.getShape(), bType.getShape(), expectedOutputShape))
|
|
return emitOpError("a and b operands don't have broadcast-compatible "
|
|
"shapes, got ")
|
|
<< aType << " and " << bType;
|
|
hasExpectedOutputShape = true;
|
|
}
|
|
|
|
ShapedType resultType = cast<ShapedType>(output.getType());
|
|
if (!resultType.hasRank())
|
|
return success();
|
|
|
|
const int64_t resultRank = resultType.getRank();
|
|
if (aHasRank && resultRank != aType.getRank())
|
|
return emitOpError("result type has different rank than a, got ")
|
|
<< resultRank << " vs " << aType.getRank();
|
|
if (bHasRank && resultRank != bType.getRank())
|
|
return emitOpError("result type has different rank than b, got ")
|
|
<< resultRank << " vs " << bType.getRank();
|
|
|
|
if (hasExpectedOutputShape &&
|
|
failed(verifyOutputShapeCompatibleWithExpected(getOperation(), resultType,
|
|
expectedOutputShape)))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::TableOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
TableOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
|
|
if (!inputShape.hasRank()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents());
|
|
return success();
|
|
}
|
|
|
|
inferredReturnShapes.resize(1);
|
|
inputShape.getDims(inferredReturnShapes[0]);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::TableOp::verify() {
|
|
const TensorType inputType = getInput1().getType();
|
|
const TensorType outputType = getOutput().getType();
|
|
|
|
if (!inputType.hasRank() || !outputType.hasRank())
|
|
return success();
|
|
|
|
if (failed(verifyRanksMatch(getOperation(), inputType, outputType, "input",
|
|
"result")))
|
|
return failure();
|
|
|
|
auto inputDims = inputType.getShape();
|
|
auto outputDims = outputType.getShape();
|
|
for (auto it : llvm::enumerate(llvm::zip(inputDims, outputDims))) {
|
|
int64_t dim = it.index();
|
|
auto [inputDim, outputDim] = it.value();
|
|
if (ShapedType::isStatic(outputDim) && outputDim != inputDim) {
|
|
return emitOpError() << "dim(result, " << dim << ") = " << outputDim
|
|
<< " doesn't match dim(input, " << dim
|
|
<< ") = " << inputDim;
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
tosa::TileOp::getConstantMultiples(SmallVector<int64_t> &multiples) {
|
|
// Multiples must be constants.
|
|
DenseIntElementsAttr multiplesAttr;
|
|
if (!matchPattern(getMultiples(), m_Constant(&multiplesAttr)))
|
|
return failure();
|
|
multiples =
|
|
llvm::map_to_vector(multiplesAttr.getValues<APInt>(),
|
|
[](const APInt &val) { return val.getSExtValue(); });
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::TileOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
TileOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
Type inputType = getElementTypeOrSelf(adaptor.getInput1().getType());
|
|
SmallVector<int64_t> multiples;
|
|
if (!tosa::getConstShapeValues(adaptor.getMultiples().getDefiningOp(),
|
|
multiples)) {
|
|
auto rank =
|
|
cast<tosa::shapeType>(adaptor.getMultiples().getType()).getRank();
|
|
SmallVector<int64_t> fallback(rank, ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(fallback, inputType));
|
|
return success();
|
|
}
|
|
multiples = convertToMlirShape(multiples);
|
|
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
SmallVector<int64_t> outputShape;
|
|
if (!inputShape.hasRank()) {
|
|
outputShape.resize(multiples.size(), ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(
|
|
ShapedTypeComponents(outputShape, inputType));
|
|
return success();
|
|
}
|
|
if (static_cast<size_t>(inputShape.getRank()) != multiples.size())
|
|
return failure();
|
|
|
|
// Any non dynamic dimension can be multiplied to a known size.
|
|
outputShape.reserve(multiples.size());
|
|
for (int i = 0, s = inputShape.getRank(); i < s; i++) {
|
|
if (multiples[i] == ShapedType::kDynamic) {
|
|
outputShape.push_back(ShapedType::kDynamic);
|
|
} else {
|
|
int64_t dim = inputShape.getDimSize(i);
|
|
if (dim != ShapedType::kDynamic)
|
|
dim *= multiples[i];
|
|
outputShape.push_back(dim);
|
|
}
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::TileOp::verify() {
|
|
if (verifySameElementTypes(*this, /* intype = */ getInput1().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
ShapedType inputType = llvm::cast<ShapedType>(getInput1().getType());
|
|
ShapedType outputType = llvm::cast<ShapedType>(getType());
|
|
|
|
shapeType multiplesType =
|
|
llvm::cast<tosa::shapeType>(getMultiples().getType());
|
|
|
|
auto multiplesRank = multiplesType.getRank();
|
|
|
|
if (inputType.hasRank()) {
|
|
if (inputType.getRank() != multiplesRank)
|
|
return emitOpError("expect 'multiples' to have rank ")
|
|
<< inputType.getRank() << " but got " << multiplesRank << ".";
|
|
if (outputType.hasRank() &&
|
|
failed(verifyRanksMatch(getOperation(), inputType, outputType, "input",
|
|
"output")))
|
|
return failure();
|
|
} else if (outputType.hasRank() && outputType.getRank() != multiplesRank)
|
|
return emitOpError("expect 'multiples' array to have length ")
|
|
<< outputType.getRank() << " but got " << multiplesRank << ".";
|
|
|
|
SmallVector<int64_t> multiples;
|
|
if (getConstantMultiples(multiples).succeeded() &&
|
|
llvm::any_of(multiples, [](int64_t v) { return v <= 0 && v != -1; }))
|
|
return emitOpError(
|
|
"expect element of 'multiples' to be positive integer or -1.");
|
|
|
|
return success();
|
|
}
|
|
|
|
bool tosa::ReshapeOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
|
|
if (l.size() != r.size() || l.size() != 1)
|
|
return false;
|
|
return getElementTypeOrSelf(l[0]) == getElementTypeOrSelf(r[0]);
|
|
}
|
|
|
|
LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ReshapeOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
Type inputType = getElementTypeOrSelf(adaptor.getInput1().getType());
|
|
llvm::SmallVector<int64_t> newShapeValue;
|
|
if (!tosa::getConstShapeValues(adaptor.getShape().getDefiningOp(),
|
|
newShapeValue)) {
|
|
auto rank = cast<tosa::shapeType>(adaptor.getShape().getType()).getRank();
|
|
SmallVector<int64_t> fallback(rank, ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(fallback, inputType));
|
|
return success();
|
|
}
|
|
newShapeValue = convertToMlirShape(newShapeValue);
|
|
|
|
// We cannot infer from the total number of elements so we must take the
|
|
// shape attribute as exact.
|
|
if (!inputShape.hasRank() || !inputShape.hasStaticShape()) {
|
|
inferredReturnShapes.push_back(
|
|
ShapedTypeComponents(newShapeValue, inputType));
|
|
return success();
|
|
}
|
|
|
|
// Determine the number of elements covered by the slice of all static
|
|
// dimensions. This allows us to infer the length of the remaining dynamic
|
|
// dimension.
|
|
int64_t numElements = inputShape.getNumElements();
|
|
int64_t staticMul = 1;
|
|
for (auto val : newShapeValue) {
|
|
if (ShapedType::isStatic(val)) {
|
|
staticMul *= val;
|
|
}
|
|
}
|
|
|
|
// Determine the length of the dynamic dimension.
|
|
for (auto &val : newShapeValue) {
|
|
if (ShapedType::isDynamic(val))
|
|
val = numElements / staticMul;
|
|
}
|
|
|
|
inferredReturnShapes.push_back(
|
|
ShapedTypeComponents(newShapeValue, inputType));
|
|
return success();
|
|
}
|
|
|
|
llvm::LogicalResult tosa::ReshapeOp::verify() {
|
|
if (verifySameElementTypes(*this, /* inType = */ getInput1().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
TensorType inputType = getInput1().getType();
|
|
|
|
SmallVector<int64_t> shapeValues;
|
|
if (!tosa::getConstShapeValues(getShape().getDefiningOp(), shapeValues)) {
|
|
// skip following checks if shape is not constant
|
|
return mlir::success();
|
|
}
|
|
|
|
int missingDims = llvm::count(shapeValues, kInferableDimSize);
|
|
if (missingDims > 1)
|
|
return emitOpError() << "expected at most one target dimension to be "
|
|
<< kInferableDimSize;
|
|
|
|
const auto outputType = dyn_cast<RankedTensorType>(getType());
|
|
if (!outputType)
|
|
return success();
|
|
|
|
if ((int64_t)shapeValues.size() != outputType.getRank())
|
|
return emitOpError() << "new shape does not match result rank";
|
|
|
|
for (auto [newShapeDim, outputShapeDim] :
|
|
zip(shapeValues, outputType.getShape())) {
|
|
if (newShapeDim != kInferableDimSize &&
|
|
newShapeDim != ShapedType::kDynamic &&
|
|
outputShapeDim != ShapedType::kDynamic && newShapeDim != outputShapeDim)
|
|
return emitOpError() << "new shape is inconsistent with result shape";
|
|
|
|
if (newShapeDim != ShapedType::kDynamic && newShapeDim < kInferableDimSize)
|
|
return emitOpError() << "new shape has invalid tensor dimension size "
|
|
<< newShapeDim;
|
|
}
|
|
|
|
if (inputType.hasStaticShape()) {
|
|
int64_t inputElementsNum = inputType.getNumElements();
|
|
if (outputType.hasStaticShape()) {
|
|
int64_t outputElementsNum = outputType.getNumElements();
|
|
if (inputElementsNum != outputElementsNum) {
|
|
return emitOpError() << "cannot reshape " << inputElementsNum
|
|
<< " elements into " << outputElementsNum;
|
|
}
|
|
}
|
|
|
|
int64_t newShapeElementsNum =
|
|
llvm::accumulate(shapeValues, int64_t(1), [](int64_t acc, int64_t dim) {
|
|
return (dim > 0) ? acc * dim : acc;
|
|
});
|
|
bool isStaticNewShape =
|
|
llvm::all_of(shapeValues, [](int64_t s) { return s > 0; });
|
|
if ((isStaticNewShape && inputElementsNum != newShapeElementsNum) ||
|
|
(!isStaticNewShape && newShapeElementsNum > inputElementsNum)) {
|
|
return emitOpError() << "cannot reshape " << inputElementsNum
|
|
<< " elements into " << newShapeElementsNum;
|
|
}
|
|
}
|
|
|
|
return mlir::success();
|
|
}
|
|
|
|
bool tosa::ReshapeBlockScaledOp::isCompatibleReturnTypes(TypeRange l,
|
|
TypeRange r) {
|
|
if (l.size() != r.size() || l.size() < 1 || l.size() > 2)
|
|
return false;
|
|
bool ok = (getElementTypeOrSelf(l[0]) == getElementTypeOrSelf(r[0]));
|
|
if (l.size() == 2)
|
|
ok = ok && (getElementTypeOrSelf(l[1]) == getElementTypeOrSelf(r[1]));
|
|
return ok;
|
|
}
|
|
|
|
LogicalResult tosa::ReshapeBlockScaledOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ReshapeBlockScaledOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
|
|
const auto numInputs = adaptor.getInput().size();
|
|
ShapeAdaptor inputShape(adaptor.getInput()[0].getType());
|
|
Type inputType = getElementTypeOrSelf(adaptor.getInput()[0].getType());
|
|
llvm::SmallVector<int64_t> newShapeValue;
|
|
const auto newShape = adaptor.getNewValueShape();
|
|
if (!tosa::getConstShapeValues(newShape.getDefiningOp(), newShapeValue)) {
|
|
auto rank = cast<tosa::shapeType>(newShape.getType()).getRank();
|
|
SmallVector<int64_t> fallback(rank, ShapedType::kDynamic);
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(fallback, inputType));
|
|
if (numInputs == 2)
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(
|
|
fallback, getElementTypeOrSelf(adaptor.getInput()[1].getType())));
|
|
return success();
|
|
}
|
|
|
|
const uint32_t blockSize =
|
|
BlockSizeAttr::getBlockSizeValue(adaptor.getBlockSize());
|
|
|
|
llvm::SmallVector<int64_t> newScaleShapeValue;
|
|
if (numInputs == 2) {
|
|
newScaleShapeValue.assign(newShapeValue.begin(), newShapeValue.end());
|
|
if (ShapedType::isStatic(newScaleShapeValue.back()))
|
|
newScaleShapeValue.back() /= blockSize;
|
|
}
|
|
|
|
inferredReturnShapes.push_back(
|
|
ShapedTypeComponents(newShapeValue, inputType));
|
|
if (numInputs == 2) {
|
|
// Fix up scale shape - with special case for last dimension
|
|
for (size_t idx = 0; idx < newShapeValue.size(); idx++) {
|
|
if (ShapedType::isDynamic(newScaleShapeValue[idx])) {
|
|
newScaleShapeValue[idx] = newShapeValue[idx];
|
|
if (idx == (newShapeValue.size() - 1))
|
|
newScaleShapeValue[idx] /= blockSize;
|
|
}
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(
|
|
newScaleShapeValue,
|
|
getElementTypeOrSelf(adaptor.getInput()[1].getType())));
|
|
}
|
|
return success();
|
|
}
|
|
|
|
llvm::LogicalResult tosa::ReshapeBlockScaledOp::verify() {
|
|
const Operation::operand_range inputList = getInput();
|
|
const Operation::result_range outputList = getResults();
|
|
|
|
if (inputList.size() == 0)
|
|
return emitOpError("requires at least one input");
|
|
|
|
if (inputList.size() > 2)
|
|
return emitOpError("requires at most two inputs");
|
|
|
|
if (inputList.size() != outputList.size())
|
|
return emitOpError("requires number of results to match inputs");
|
|
|
|
if (verifySameElementTypes(*this, /* inType = */ inputList[0].getType(),
|
|
/* outType = */ outputList[0].getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
|
|
const auto inputType = llvm::cast<ShapedType>(inputList[0].getType());
|
|
if (!inputType.hasRank())
|
|
return success();
|
|
const uint32_t blockSize = BlockSizeAttr::getBlockSizeValue(getBlockSize());
|
|
|
|
if (inputList.size() == 2) {
|
|
if (blockSize != BlockSizeAttr::getBlockSizeValue(BlockSize::BLOCK_SIZE_32))
|
|
return emitOpError("expect block size to be 32, got ") << blockSize;
|
|
if (llvm::any_of(inputList, [](Value v) {
|
|
const auto input = cast<ShapedType>(v.getType());
|
|
return input.hasRank() && input.getRank() == 0;
|
|
}))
|
|
return emitOpError(
|
|
"requires all input shapes have a rank greater than 0");
|
|
if (llvm::any_of(outputList, [](Value v) {
|
|
const auto output = cast<ShapedType>(v.getType());
|
|
return output.hasRank() && output.getRank() == 0;
|
|
}))
|
|
return emitOpError(
|
|
"requires all result shapes have a rank greater than 0");
|
|
|
|
if (verifySameElementTypes(*this, /* inType = */ inputList[1].getType(),
|
|
/* outType = */ outputList[1].getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
|
|
const auto inputScaleType = llvm::cast<ShapedType>(inputList[1].getType());
|
|
if (inputScaleType.hasRank()) {
|
|
if (inputType.getRank() != inputScaleType.getRank())
|
|
return emitOpError("input shapes do not have same rank");
|
|
|
|
// Check all but the last dimension that the input shape dimensions match
|
|
for (auto dimIdx = 0; dimIdx < inputType.getRank() - 1; dimIdx++) {
|
|
const int64_t inputValueDim = inputType.getDimSize(dimIdx);
|
|
const int64_t inputScaleDim = inputScaleType.getShape()[dimIdx];
|
|
if (ShapedType::isStatic(inputValueDim) &&
|
|
ShapedType::isStatic(inputScaleDim) &&
|
|
inputValueDim != inputScaleDim)
|
|
return emitOpError("input shapes for data and scale do not match on "
|
|
"dimension ")
|
|
<< dimIdx;
|
|
}
|
|
|
|
// Verify last dimension of input is a multiple of block size
|
|
const int64_t lastValueDim =
|
|
inputType.getDimSize(inputType.getRank() - 1);
|
|
if (ShapedType::isStatic(lastValueDim)) {
|
|
if (lastValueDim % blockSize != 0)
|
|
return emitOpError("expect last dimension of input_data (")
|
|
<< lastValueDim << ") to be divisible by block_size ("
|
|
<< blockSize << ")";
|
|
|
|
const int64_t lastScaleDim =
|
|
inputScaleType.getDimSize(inputScaleType.getRank() - 1);
|
|
// Verify last dimension of scale is lastValueDim / block size
|
|
if (ShapedType::isStatic(lastScaleDim) &&
|
|
lastScaleDim != lastValueDim / blockSize)
|
|
return emitOpError("expect last dimension of scale_data (")
|
|
<< lastScaleDim << ") to be " << lastValueDim << "/"
|
|
<< blockSize;
|
|
}
|
|
}
|
|
} else {
|
|
if (blockSize != BlockSizeAttr::getBlockSizeValue(BlockSize::BLOCK_SIZE_1))
|
|
return emitOpError("expect block size to be 1, got ") << blockSize;
|
|
}
|
|
|
|
// Get the new value shape dimension values
|
|
SmallVector<int64_t> shapeValues;
|
|
if (!tosa::getConstShapeValues(getNewValueShape().getDefiningOp(),
|
|
shapeValues)) {
|
|
// skip following checks if shape is not constant
|
|
return mlir::success();
|
|
}
|
|
|
|
if (inputList.size() == 2) {
|
|
if (static_cast<int64_t>(shapeValues.size()) == 0)
|
|
return emitOpError("requires new shape to have a rank greater than 0");
|
|
|
|
const int64_t lastShapeDim = shapeValues.back();
|
|
if (ShapedType::isStatic(lastShapeDim) && lastShapeDim % blockSize != 0)
|
|
return emitOpError("expect last dimension of new shape (")
|
|
<< lastShapeDim << ") to be divisible by block_size (" << blockSize
|
|
<< ")";
|
|
}
|
|
|
|
const auto outputType = llvm::cast<ShapedType>(outputList[0].getType());
|
|
if (!outputType.hasRank())
|
|
return success();
|
|
|
|
if (static_cast<int64_t>(shapeValues.size()) != outputType.getRank())
|
|
return emitOpError() << "result does not match new shape rank";
|
|
|
|
for (auto [newShapeDim, outputShapeDim] :
|
|
zip(shapeValues, outputType.getShape())) {
|
|
if (ShapedType::isStatic(newShapeDim) &&
|
|
ShapedType::isStatic(outputShapeDim) && newShapeDim != outputShapeDim)
|
|
return emitOpError() << "result shape is inconsistent with new shape";
|
|
}
|
|
|
|
if (outputList.size() == 2) {
|
|
// Set up scale shape from new shape given
|
|
SmallVector<int64_t> scaleShapeValues(shapeValues.begin(),
|
|
shapeValues.end());
|
|
scaleShapeValues.back() /= blockSize;
|
|
|
|
const auto outputScaleType =
|
|
llvm::cast<ShapedType>(outputList[1].getType());
|
|
if (outputScaleType.hasRank()) {
|
|
if ((int64_t)scaleShapeValues.size() != outputScaleType.getRank())
|
|
return emitOpError() << "result scale does not match new shape rank";
|
|
|
|
for (auto [newScaleShapeDim, outputScaleShapeDim] :
|
|
zip(scaleShapeValues, outputScaleType.getShape())) {
|
|
if (ShapedType::isStatic(newScaleShapeDim) &&
|
|
ShapedType::isStatic(outputScaleShapeDim) &&
|
|
newScaleShapeDim != outputScaleShapeDim)
|
|
return emitOpError()
|
|
<< "result scale shape is inconsistent with new shape";
|
|
}
|
|
}
|
|
}
|
|
|
|
if (inputType.hasStaticShape()) {
|
|
int64_t inputElementsNum = inputType.getNumElements();
|
|
if (outputType.hasStaticShape()) {
|
|
int64_t outputElementsNum = outputType.getNumElements();
|
|
if (inputElementsNum != outputElementsNum) {
|
|
return emitOpError() << "cannot reshape " << inputElementsNum
|
|
<< " elements into " << outputElementsNum;
|
|
}
|
|
}
|
|
|
|
int64_t newShapeElementsNum =
|
|
llvm::accumulate(shapeValues, int64_t(1), [](int64_t acc, int64_t dim) {
|
|
return (dim > 0) ? acc * dim : acc;
|
|
});
|
|
bool isStaticNewShape =
|
|
llvm::all_of(shapeValues, [](int64_t s) { return s > 0; });
|
|
if ((isStaticNewShape && inputElementsNum != newShapeElementsNum) ||
|
|
(!isStaticNewShape && newShapeElementsNum > inputElementsNum)) {
|
|
return emitOpError() << "cannot reshape " << inputElementsNum
|
|
<< " elements into " << newShapeElementsNum;
|
|
}
|
|
}
|
|
|
|
return mlir::success();
|
|
}
|
|
|
|
// return failure if val is not a constant
|
|
// set zp to -1 if val is non-zero float or val is not integer nor float
|
|
// otherwise set zp to val's constant value
|
|
static FailureOr<int64_t> getZeroPoint(Value val, bool signExtend) {
|
|
ElementsAttr zpAttr;
|
|
if (!matchPattern(val, m_Constant(&zpAttr))) {
|
|
return failure();
|
|
}
|
|
|
|
Type zpElemType = zpAttr.getElementType();
|
|
|
|
if (llvm::isa<FloatType>(zpElemType)) {
|
|
if (zpAttr.getValues<APFloat>()[0].isZero()) {
|
|
return 0;
|
|
}
|
|
// return non-zero value to trigger error check
|
|
return -1;
|
|
}
|
|
|
|
if (llvm::isa<IntegerType>(zpElemType)) {
|
|
if (signExtend)
|
|
return zpAttr.getValues<APInt>()[0].getSExtValue();
|
|
return zpAttr.getValues<APInt>()[0].getZExtValue();
|
|
}
|
|
|
|
// return non-zero value to trigger error check
|
|
return -1;
|
|
}
|
|
|
|
static FailureOr<int64_t> getConstantScalarIntValue(Value val) {
|
|
ElementsAttr attr;
|
|
if (!matchPattern(val, m_Constant(&attr)))
|
|
return failure();
|
|
|
|
if (!llvm::isa<IntegerType>(attr.getElementType()) ||
|
|
attr.getNumElements() != 1)
|
|
return failure();
|
|
|
|
return attr.getValues<APInt>()[0].getSExtValue();
|
|
}
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyZeroPoint(T op, Value val, const int64_t &zp,
|
|
const std::string &operand) {
|
|
Type zpElemType = getElementTypeOrSelf(val);
|
|
|
|
if (!zpElemType.isInteger(8) && zp != 0) {
|
|
// convert operand to lower case for error message
|
|
std::string lower = operand;
|
|
llvm::transform(lower, lower.begin(), ::tolower);
|
|
return op.emitOpError()
|
|
<< lower << " zero point must be zero for non-int8 integer types";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyZeroPoint(tosa::RescaleOp op, Value zpVal,
|
|
const int64_t &zp,
|
|
const std::string &operand) {
|
|
bool isInputZp = (operand == "Input");
|
|
|
|
bool tensorUnsigned =
|
|
isInputZp ? op.getInputUnsigned() : op.getOutputUnsigned();
|
|
StringRef tensorName = isInputZp ? "input" : "output";
|
|
|
|
Type zpElemType = getElementTypeOrSelf(zpVal);
|
|
|
|
if (zp != 0) {
|
|
if (!zpElemType.isInteger(8) &&
|
|
!(zpElemType.isInteger(16) && tensorUnsigned)) {
|
|
return op.emitOpError()
|
|
<< "expect " << tensorName << "_zp of 0, got " << zp;
|
|
}
|
|
if (zpElemType.isInteger(16) && tensorUnsigned && zp != 32768) {
|
|
return op.emitOpError() << "expect " << tensorName
|
|
<< "_zp of 0 or 32768 for unsigned int16 "
|
|
<< tensorName << ", got " << zp;
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
#define ZERO_POINT_HELPER(OP, OPERAND_NAME, SIGN_EXTEND) \
|
|
FailureOr<int64_t> tosa::OP::get##OPERAND_NAME##ZeroPoint() { \
|
|
return getZeroPoint(get##OPERAND_NAME##Zp(), SIGN_EXTEND); \
|
|
} \
|
|
LogicalResult tosa::OP::verify##OPERAND_NAME##ZeroPoint(int64_t zp) { \
|
|
return verifyZeroPoint(*this, get##OPERAND_NAME##Zp(), zp, #OPERAND_NAME); \
|
|
}
|
|
|
|
ZERO_POINT_HELPER(Conv2DOp, Input, true)
|
|
ZERO_POINT_HELPER(Conv2DOp, Weight, true)
|
|
ZERO_POINT_HELPER(Conv3DOp, Input, true)
|
|
ZERO_POINT_HELPER(Conv3DOp, Weight, true)
|
|
ZERO_POINT_HELPER(DepthwiseConv2DOp, Input, true)
|
|
ZERO_POINT_HELPER(DepthwiseConv2DOp, Weight, true)
|
|
ZERO_POINT_HELPER(TransposeConv2DOp, Input, true)
|
|
ZERO_POINT_HELPER(TransposeConv2DOp, Weight, true)
|
|
ZERO_POINT_HELPER(AvgPool2dOp, Input, true)
|
|
ZERO_POINT_HELPER(AvgPool2dOp, Output, true)
|
|
ZERO_POINT_HELPER(AvgPool2dAdaptiveOp, Input, true)
|
|
ZERO_POINT_HELPER(AvgPool2dAdaptiveOp, Output, true)
|
|
ZERO_POINT_HELPER(MatMulOp, A, true)
|
|
ZERO_POINT_HELPER(MatMulOp, B, true)
|
|
ZERO_POINT_HELPER(NegateOp, Input1, true)
|
|
ZERO_POINT_HELPER(NegateOp, Output, true)
|
|
ZERO_POINT_HELPER(RescaleOp, Input, !getInputUnsigned())
|
|
ZERO_POINT_HELPER(RescaleOp, Output, !getOutputUnsigned())
|
|
#undef ZERO_POINT_HELPER
|
|
|
|
LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
TransposeOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
|
|
// If input rank and permutation length is unknown, the output rank is
|
|
// unknown.
|
|
if (!inputShape.hasRank()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents());
|
|
return success();
|
|
}
|
|
|
|
const auto inputRank = inputShape.getRank();
|
|
|
|
// This would imply the number of permutations does not match the rank of
|
|
// the input which is illegal.
|
|
if (adaptor.getPerms().size() != static_cast<size_t>(inputRank)) {
|
|
return failure();
|
|
}
|
|
|
|
SmallVector<int64_t> outputShape;
|
|
// Rank-0 means no permutations matter.
|
|
if (inputRank == 0) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
// Check whether the input dimensions are all the same.
|
|
bool allTheSame = true;
|
|
for (int i = 1, s = inputRank; i < s; i++) {
|
|
if (inputShape.getDimSize(0) != inputShape.getDimSize(i)) {
|
|
allTheSame = false;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// If all of the input dimensions are the same we don't care about the
|
|
// permutation.
|
|
if (allTheSame) {
|
|
outputShape.resize(inputRank, inputShape.getDimSize(0));
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
outputShape.resize(inputRank, ShapedType::kDynamic);
|
|
|
|
// Constant permutation values must be within the input rank.
|
|
if (llvm::any_of(adaptor.getPerms(),
|
|
[inputRank](const auto i) { return i >= inputRank; }))
|
|
return failure();
|
|
|
|
outputShape.reserve(inputRank);
|
|
for (int i = 0, s = inputRank; i < s; i++) {
|
|
outputShape[i] = inputShape.getDimSize(adaptor.getPerms()[i]);
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::TransposeOp::verify() {
|
|
if (verifySameElementTypes(*this, /* inType = */ getInput1().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
|
|
const ShapeAdaptor inputShape(getInput1().getType());
|
|
const ShapeAdaptor outputShape(getOutput().getType());
|
|
|
|
const llvm::ArrayRef<int32_t> constantPerms = getPerms();
|
|
|
|
if (inputShape.hasRank() &&
|
|
constantPerms.size() != static_cast<size_t>(inputShape.getRank()))
|
|
return emitOpError() << "expected perms attribute to have size "
|
|
<< inputShape.getRank()
|
|
<< " (input rank) but got size "
|
|
<< constantPerms.size();
|
|
|
|
if (inputShape.hasRank() && outputShape.hasRank() &&
|
|
inputShape.getRank() != outputShape.getRank())
|
|
return emitOpError()
|
|
<< "expected input tensor rank to equal result tensor rank";
|
|
|
|
if (outputShape.hasRank() &&
|
|
constantPerms.size() != static_cast<size_t>(outputShape.getRank()))
|
|
return emitOpError() << "expected perms attribute to have size "
|
|
<< outputShape.getRank()
|
|
<< " (output rank) but got size "
|
|
<< constantPerms.size();
|
|
|
|
if (!llvm::all_of(constantPerms,
|
|
[&constantPerms](int32_t s) {
|
|
return s >= 0 &&
|
|
static_cast<size_t>(s) < constantPerms.size();
|
|
}) ||
|
|
!isPermutationVector(llvm::map_to_vector(
|
|
constantPerms, [](int32_t v) -> int64_t { return v; })))
|
|
return emitOpError() << "expected valid permutation indices";
|
|
|
|
// ERROR_IF(tensor_size(shape1) != tensor_size(shape))
|
|
if (inputShape.hasStaticShape() && outputShape.hasStaticShape() &&
|
|
inputShape.getNumElements() != outputShape.getNumElements())
|
|
return emitOpError() << "expected input1 and output to have same numbers "
|
|
"of elements, got "
|
|
<< inputShape.getNumElements() << " and "
|
|
<< outputShape.getNumElements();
|
|
|
|
// Verify that the types of the input and output tensors are properly
|
|
// permuted.
|
|
if (inputShape.hasRank() && outputShape.hasRank()) {
|
|
for (auto i = 0; i < outputShape.getRank(); i++) {
|
|
if (inputShape.isDynamicDim(constantPerms[i]) ||
|
|
outputShape.isDynamicDim(i))
|
|
continue;
|
|
|
|
if (inputShape.getDimSize(constantPerms[i]) != outputShape.getDimSize(i))
|
|
return emitOpError()
|
|
<< "expected output tensor dim " << i << " to match "
|
|
<< "input dim " << constantPerms[i] << " with value of "
|
|
<< inputShape.getDimSize(constantPerms[i]);
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult TransposeOp::reifyResultShapes(
|
|
OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
|
|
|
|
const llvm::ArrayRef<int32_t> transposePerms = getPerms();
|
|
|
|
Value input = getInput1();
|
|
auto inputType = cast<TensorType>(input.getType());
|
|
|
|
SmallVector<OpFoldResult> returnedDims(inputType.getRank());
|
|
for (auto dim : transposePerms) {
|
|
int32_t dimInInput = transposePerms[dim];
|
|
if (inputType.isDynamicDim(dimInInput))
|
|
returnedDims[dim] =
|
|
tensor::DimOp::create(builder, getLoc(), input, dimInInput)
|
|
.getResult();
|
|
else
|
|
returnedDims[dim] =
|
|
builder.getIndexAttr(inputType.getDimSize(dimInInput));
|
|
}
|
|
|
|
reifiedReturnShapes.emplace_back(std::move(returnedDims));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::GatherOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
GatherOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t> outputShape;
|
|
outputShape.resize(3, ShapedType::kDynamic);
|
|
|
|
ShapeAdaptor valuesShape(adaptor.getValues().getType());
|
|
if (valuesShape.hasRank()) {
|
|
outputShape[0] = valuesShape.getDimSize(0);
|
|
outputShape[2] = valuesShape.getDimSize(2);
|
|
}
|
|
|
|
ShapeAdaptor indicesShape(adaptor.getIndices().getType());
|
|
if (indicesShape.hasRank()) {
|
|
if (outputShape[0] == ShapedType::kDynamic)
|
|
outputShape[0] = indicesShape.getDimSize(0);
|
|
if (outputShape[1] == ShapedType::kDynamic)
|
|
outputShape[1] = indicesShape.getDimSize(1);
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::RowGatherBlockScaledOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
RowGatherBlockScaledOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
const auto values = adaptor.getValues();
|
|
if (values.empty())
|
|
return failure();
|
|
|
|
SmallVector<int64_t> dataShape(3, ShapedType::kDynamic);
|
|
const ShapeAdaptor valuesShape(values.front().getType());
|
|
if (valuesShape.hasRank()) {
|
|
dataShape[0] = valuesShape.getDimSize(0);
|
|
dataShape[2] = valuesShape.getDimSize(2);
|
|
}
|
|
|
|
const ShapeAdaptor indicesShape(adaptor.getIndices().getType());
|
|
if (indicesShape.hasRank()) {
|
|
if (dataShape[0] == ShapedType::kDynamic)
|
|
dataShape[0] = indicesShape.getDimSize(0);
|
|
|
|
if (auto rowCount = getConstantScalarIntValue(adaptor.getRowCount());
|
|
succeeded(rowCount) && rowCount.value() > 0) {
|
|
const int64_t indicesW = indicesShape.getDimSize(1);
|
|
if (ShapedType::isStatic(indicesW))
|
|
dataShape[1] = indicesW * rowCount.value();
|
|
}
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(dataShape));
|
|
if (values.size() == 1)
|
|
return success();
|
|
|
|
SmallVector<int64_t> scaleShape = dataShape;
|
|
const uint32_t blockSize =
|
|
BlockSizeAttr::getBlockSizeValue(adaptor.getBlockSize());
|
|
if (ShapedType::isStatic(dataShape[2]))
|
|
scaleShape[2] = dataShape[2] / blockSize;
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(scaleShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::GatherOp::verify() {
|
|
if (verifySameElementTypes(*this, /* inType = */ getValues().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
|
|
const ShapeAdaptor valuesShape(getValues().getType());
|
|
const ShapeAdaptor indicesShape(getIndices().getType());
|
|
const ShapeAdaptor outputShape(getOutput().getType());
|
|
|
|
int64_t n = ShapedType::kDynamic;
|
|
int64_t w = ShapedType::kDynamic;
|
|
int64_t c = ShapedType::kDynamic;
|
|
|
|
if (valuesShape.hasRank()) {
|
|
n = valuesShape.getDimSize(0);
|
|
c = valuesShape.getDimSize(2);
|
|
}
|
|
if (indicesShape.hasRank()) {
|
|
const int64_t indicesN = indicesShape.getDimSize(0);
|
|
w = indicesShape.getDimSize(1);
|
|
if (n == ShapedType::kDynamic)
|
|
n = indicesN;
|
|
else if (indicesN != ShapedType::kDynamic && n != indicesN)
|
|
return emitOpError() << "requires indices dimension 0 to have size " << n
|
|
<< ", got " << indicesN;
|
|
}
|
|
if (outputShape.hasRank()) {
|
|
const int64_t outputN = outputShape.getDimSize(0);
|
|
const int64_t outputW = outputShape.getDimSize(1);
|
|
const int64_t outputC = outputShape.getDimSize(2);
|
|
if (n != ShapedType::kDynamic && outputN != ShapedType::kDynamic &&
|
|
n != outputN)
|
|
return emitOpError() << "requires output dimension 0 to have size " << n
|
|
<< ", got " << outputN;
|
|
|
|
if (w != ShapedType::kDynamic && outputW != ShapedType::kDynamic &&
|
|
w != outputW)
|
|
return emitOpError() << "requires output dimension 1 to have size " << w
|
|
<< ", got " << outputW;
|
|
if (c != ShapedType::kDynamic && outputC != ShapedType::kDynamic &&
|
|
c != outputC)
|
|
return emitOpError() << "requires output dimension 2 to have size " << c
|
|
<< ", got " << outputC;
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::RowGatherBlockScaledOp::verify() {
|
|
const OperandRange values = getValues();
|
|
const ResultRange output = getOutput();
|
|
if (values.empty() || values.size() > 2)
|
|
return emitOpError()
|
|
<< "expects values tensor list length to be 1 or 2, got "
|
|
<< values.size();
|
|
if (output.size() != values.size())
|
|
return emitOpError()
|
|
<< "expects output tensor list length to match values tensor list "
|
|
"length, got "
|
|
<< output.size() << " results for " << values.size()
|
|
<< " input tensors";
|
|
|
|
const uint32_t blockSize = BlockSizeAttr::getBlockSizeValue(getBlockSize());
|
|
if (values.size() == 1 && blockSize != 1)
|
|
return emitOpError()
|
|
<< "requires block_size to be BLOCK_SIZE_1 when values tensor list "
|
|
"length is 1";
|
|
if (values.size() == 2 && blockSize == 1)
|
|
return emitOpError()
|
|
<< "requires block_size to not be BLOCK_SIZE_1 when values tensor "
|
|
"list length is 2";
|
|
|
|
if (failed(verifySameElementTypes(*this, values[0].getType(),
|
|
output[0].getType(), "values[0]",
|
|
"output[0]")))
|
|
return failure();
|
|
if (values.size() == 2 && failed(verifySameElementTypes(
|
|
*this, values[1].getType(), output[1].getType(),
|
|
"values[1]", "output[1]")))
|
|
return failure();
|
|
|
|
if (auto rowCount = getConstantScalarIntValue(getRowCount());
|
|
succeeded(rowCount) && rowCount.value() <= 0)
|
|
return emitOpError() << "requires row_count to be > 0, got "
|
|
<< rowCount.value();
|
|
|
|
int64_t n = ShapedType::kDynamic;
|
|
int64_t k = ShapedType::kDynamic;
|
|
int64_t c = ShapedType::kDynamic;
|
|
int64_t w = ShapedType::kDynamic;
|
|
int64_t multiplesOfC = ShapedType::kDynamic;
|
|
|
|
const ShapeAdaptor valuesDataShape(values[0].getType());
|
|
if (valuesDataShape.hasRank()) {
|
|
n = valuesDataShape.getDimSize(0);
|
|
k = valuesDataShape.getDimSize(1);
|
|
c = valuesDataShape.getDimSize(2);
|
|
}
|
|
|
|
if (ShapedType::isStatic(c) && c % blockSize != 0)
|
|
return emitOpError() << "expects channels of values[0] (" << c
|
|
<< ") to be divisible by block_size (" << blockSize
|
|
<< ")";
|
|
|
|
const ShapeAdaptor indicesShape(getIndices().getType());
|
|
if (indicesShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, n, indicesShape.getDimSize(0),
|
|
"indices", "batch")))
|
|
return failure();
|
|
w = indicesShape.getDimSize(1);
|
|
}
|
|
|
|
const ShapeAdaptor outputDataShape(output[0].getType());
|
|
if (outputDataShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, n, outputDataShape.getDimSize(0),
|
|
"output[0]", "batch")) ||
|
|
failed(tryUpdateDimOrFailure(*this, c, outputDataShape.getDimSize(2),
|
|
"output[0]", "channels")))
|
|
return failure();
|
|
|
|
if (auto rowCount = getConstantScalarIntValue(getRowCount());
|
|
succeeded(rowCount) && rowCount.value() > 0 &&
|
|
ShapedType::isStatic(w)) {
|
|
const int64_t expectedOutputRows = w * rowCount.value();
|
|
if (ShapedType::isStatic(outputDataShape.getDimSize(1)) &&
|
|
outputDataShape.getDimSize(1) != expectedOutputRows)
|
|
return emitOpError() << "requires output[0] dimension 1 to have size "
|
|
<< expectedOutputRows << ", got "
|
|
<< outputDataShape.getDimSize(1);
|
|
}
|
|
}
|
|
|
|
if (values.size() == 2) {
|
|
const ShapeAdaptor valuesScaleShape(values[1].getType());
|
|
if (valuesScaleShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, n, valuesScaleShape.getDimSize(0),
|
|
"values[1]", "batch")) ||
|
|
failed(tryUpdateDimOrFailure(*this, k, valuesScaleShape.getDimSize(1),
|
|
"values[1]", "rows")))
|
|
return failure();
|
|
multiplesOfC = valuesScaleShape.getDimSize(2);
|
|
}
|
|
|
|
const ShapeAdaptor outputScaleShape(output[1].getType());
|
|
if (outputScaleShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, n, outputScaleShape.getDimSize(0),
|
|
"output[1]", "batch")))
|
|
return failure();
|
|
|
|
if (auto rowCount = getConstantScalarIntValue(getRowCount());
|
|
succeeded(rowCount) && rowCount.value() > 0 &&
|
|
ShapedType::isStatic(w)) {
|
|
const int64_t expectedOutputRows = w * rowCount.value();
|
|
if (ShapedType::isStatic(outputScaleShape.getDimSize(1)) &&
|
|
outputScaleShape.getDimSize(1) != expectedOutputRows)
|
|
return emitOpError() << "requires output[1] dimension 1 to have size "
|
|
<< expectedOutputRows << ", got "
|
|
<< outputScaleShape.getDimSize(1);
|
|
}
|
|
|
|
if (ShapedType::isDynamic(multiplesOfC))
|
|
multiplesOfC = outputScaleShape.getDimSize(2);
|
|
else if (ShapedType::isStatic(outputScaleShape.getDimSize(2)) &&
|
|
multiplesOfC != outputScaleShape.getDimSize(2))
|
|
return emitOpError()
|
|
<< "expected channels of output[1] to match size "
|
|
<< multiplesOfC << ", got " << outputScaleShape.getDimSize(2);
|
|
}
|
|
|
|
if (ShapedType::isStatic(c) && ShapedType::isStatic(multiplesOfC) &&
|
|
multiplesOfC != c / blockSize)
|
|
return emitOpError()
|
|
<< "expects channels of scale tensors to equal C/block_size (" << c
|
|
<< "/" << blockSize << "), got " << multiplesOfC;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ResizeOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t, 4> outputShape;
|
|
outputShape.resize(4, ShapedType::kDynamic);
|
|
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
if (!inputShape.hasRank())
|
|
return failure();
|
|
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
outputShape[3] = inputShape.getDimSize(3);
|
|
int64_t inputHeight = inputShape.getDimSize(1);
|
|
int64_t inputWidth = inputShape.getDimSize(2);
|
|
|
|
if ((inputHeight == ShapedType::kDynamic) ||
|
|
(inputWidth == ShapedType::kDynamic))
|
|
return failure();
|
|
|
|
SmallVector<int64_t> scaleInt, offsetInt, borderInt;
|
|
if (!tosa::getConstShapeValues(adaptor.getScale().getDefiningOp(),
|
|
scaleInt) ||
|
|
!tosa::getConstShapeValues(adaptor.getOffset().getDefiningOp(),
|
|
offsetInt) ||
|
|
!tosa::getConstShapeValues(adaptor.getBorder().getDefiningOp(),
|
|
borderInt)) {
|
|
return failure();
|
|
}
|
|
|
|
// Compute the output shape based on attributes: scale, offset, and border.
|
|
const int64_t outputHeight =
|
|
(((inputHeight - 1) * scaleInt[0] - offsetInt[0] + borderInt[0]) /
|
|
scaleInt[1]) +
|
|
1;
|
|
|
|
const int64_t outputWidth =
|
|
(((inputWidth - 1) * scaleInt[2] - offsetInt[1] + borderInt[1]) /
|
|
scaleInt[3]) +
|
|
1;
|
|
|
|
if (outputHeight < 0 || outputWidth < 0) {
|
|
return emitOptionalError(
|
|
location,
|
|
"calculated output height and width must be non-negative, "
|
|
"got height = ",
|
|
outputHeight, ", width = ", outputWidth);
|
|
}
|
|
|
|
outputShape[1] = outputHeight;
|
|
outputShape[2] = outputWidth;
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ResizeOp::verify() {
|
|
const Value input = getInput();
|
|
const Value output = getOutput();
|
|
const RankedTensorType inputType =
|
|
llvm::dyn_cast<RankedTensorType>(input.getType());
|
|
const RankedTensorType outputType =
|
|
llvm::dyn_cast<RankedTensorType>(output.getType());
|
|
|
|
SmallVector<int64_t> scaleValues;
|
|
SmallVector<int64_t> offsetValues;
|
|
SmallVector<int64_t> borderValues;
|
|
if (!tosa::getConstShapeValues(getScale().getDefiningOp(), scaleValues) ||
|
|
!tosa::getConstShapeValues(getOffset().getDefiningOp(), offsetValues) ||
|
|
!tosa::getConstShapeValues(getBorder().getDefiningOp(), borderValues)) {
|
|
// Skip following checks if shape is not constant
|
|
return success();
|
|
}
|
|
|
|
if (llvm::any_of(scaleValues, [](int64_t s) { return s <= 0; }))
|
|
return emitOpError("expect all scale values to be > 0, got ")
|
|
<< scaleValues;
|
|
|
|
const int64_t scaleYN = scaleValues[0];
|
|
const int64_t scaleYD = scaleValues[1];
|
|
const int64_t scaleXN = scaleValues[2];
|
|
const int64_t scaleXD = scaleValues[3];
|
|
|
|
const int64_t offsetY = offsetValues[0];
|
|
const int64_t offsetX = offsetValues[1];
|
|
|
|
const int64_t borderY = borderValues[0];
|
|
const int64_t borderX = borderValues[1];
|
|
|
|
if (!inputType)
|
|
return success();
|
|
if (!outputType)
|
|
return success();
|
|
|
|
const int64_t oh = outputType.getDimSize(1);
|
|
const int64_t ow = outputType.getDimSize(2);
|
|
const int64_t ih = inputType.getDimSize(1);
|
|
const int64_t iw = inputType.getDimSize(2);
|
|
|
|
// Don't check with input height that could be broadcast (ih != 1)
|
|
// since Linalg, a consumer of TOSA, expects broadcasting support
|
|
// in resize to be available. Taking the cautious approach for now,
|
|
// we can consider removing support for broadcasting later.
|
|
if (ih != ShapedType::kDynamic && ih != 1) {
|
|
const std::optional<int64_t> calculatedOutHeightMinusOne =
|
|
idivCheck((ih - 1) * scaleYN - offsetY + borderY, scaleYD);
|
|
if (!calculatedOutHeightMinusOne.has_value())
|
|
return emitOpError("expected (input_height - 1) * scale_y_n - offset_y + "
|
|
"border_y ")
|
|
<< "to be wholly divisible by scale_y_d, got ((" << ih
|
|
<< " - 1) * " << scaleYN << " - " << offsetY << " + " << borderY
|
|
<< ") / " << scaleYD;
|
|
const int64_t calculatedOutHeight = calculatedOutHeightMinusOne.value() + 1;
|
|
if (oh != ShapedType::kDynamic && calculatedOutHeight != oh)
|
|
return emitOpError("calculated output height did not match expected: ")
|
|
<< "calculated=" << calculatedOutHeight << ", expected=" << oh;
|
|
}
|
|
|
|
// Don't check with input width that could be broadcast (iw != 1)
|
|
// since Linalg, a consumer of TOSA, expects broadcasting support
|
|
// in resize to be available. Taking the cautious approach for now,
|
|
// we can consider removing support for broadcasting later.
|
|
if (iw != ShapedType::kDynamic && iw != 1) {
|
|
const int64_t scaledInWidth = (iw - 1) * scaleXN - offsetX + borderX;
|
|
const std::optional<int64_t> calculatedOutWidthMinusOne =
|
|
idivCheck(scaledInWidth, scaleXD);
|
|
if (!calculatedOutWidthMinusOne.has_value())
|
|
return emitOpError("expected (input_width - 1) * scale_x_n - offset_x + "
|
|
"border_x ")
|
|
<< "to be wholly divisible by scale_x_d, got ((" << iw
|
|
<< " - 1) * " << scaleXN << " - " << offsetX << " + " << borderX
|
|
<< ") / " << scaleXD;
|
|
const int64_t calculatedOutWidth = calculatedOutWidthMinusOne.value() + 1;
|
|
if (ow != ShapedType::kDynamic && calculatedOutWidth != ow)
|
|
return emitOpError("calculated output width did not match expected: ")
|
|
<< "calculated=" << calculatedOutWidth << ", expected=" << ow;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ScatterOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
ScatterOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t> outputShape;
|
|
outputShape.resize(3, ShapedType::kDynamic);
|
|
|
|
ShapeAdaptor valuesInShape(adaptor.getValuesIn().getType());
|
|
if (valuesInShape.hasRank()) {
|
|
outputShape[0] = valuesInShape.getDimSize(0);
|
|
outputShape[1] = valuesInShape.getDimSize(1);
|
|
outputShape[2] = valuesInShape.getDimSize(2);
|
|
}
|
|
|
|
ShapeAdaptor indicesShape(adaptor.getIndices().getType());
|
|
if (indicesShape.hasRank()) {
|
|
if (outputShape[0] == ShapedType::kDynamic)
|
|
outputShape[0] = indicesShape.getDimSize(0);
|
|
}
|
|
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
if (inputShape.hasRank()) {
|
|
if (outputShape[0] == ShapedType::kDynamic)
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
if (outputShape[2] == ShapedType::kDynamic)
|
|
outputShape[2] = inputShape.getDimSize(2);
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ScatterOp::verify() {
|
|
if (verifySameElementTypes(*this, /* inType = */ getValuesIn().getType(),
|
|
/* outType = */ getValuesOut().getType())
|
|
.failed() ||
|
|
verifySameElementTypes(*this, /* inType = */ getInput().getType(),
|
|
/* outType = */ getValuesOut().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
|
|
const ShapeAdaptor valuesInShape(getValuesIn().getType());
|
|
const ShapeAdaptor indicesShape(getIndices().getType());
|
|
const ShapeAdaptor inputShape(getInput().getType());
|
|
const ShapeAdaptor outputShape(getValuesOut().getType());
|
|
|
|
int64_t n = ShapedType::kDynamic;
|
|
int64_t k = ShapedType::kDynamic;
|
|
int64_t w = ShapedType::kDynamic;
|
|
int64_t c = ShapedType::kDynamic;
|
|
if (valuesInShape.hasRank()) {
|
|
n = valuesInShape.getDimSize(0);
|
|
k = valuesInShape.getDimSize(1);
|
|
c = valuesInShape.getDimSize(2);
|
|
}
|
|
if (indicesShape.hasRank()) {
|
|
const int64_t indicesN = indicesShape.getDimSize(0);
|
|
w = indicesShape.getDimSize(1);
|
|
if (n == ShapedType::kDynamic)
|
|
n = indicesN;
|
|
else if (indicesN != ShapedType::kDynamic && n != indicesN)
|
|
return emitOpError() << "requires indices dimension 0 to have size " << n
|
|
<< ", got " << indicesN;
|
|
}
|
|
if (inputShape.hasRank()) {
|
|
const int64_t inputN = inputShape.getDimSize(0);
|
|
const int64_t inputW = inputShape.getDimSize(1);
|
|
const int64_t inputC = inputShape.getDimSize(2);
|
|
if (n == ShapedType::kDynamic)
|
|
n = inputN;
|
|
else if (inputN != ShapedType::kDynamic && n != inputN)
|
|
return emitOpError() << "requires input dimension 0 to have size " << n
|
|
<< ", got " << inputN;
|
|
if (w == ShapedType::kDynamic)
|
|
w = inputW;
|
|
else if (inputW != ShapedType::kDynamic && w != inputW)
|
|
return emitOpError() << "requires input dimension 1 to have size " << w
|
|
<< ", got " << inputW;
|
|
|
|
if (c == ShapedType::kDynamic)
|
|
c = inputC;
|
|
else if (inputC != ShapedType::kDynamic && c != inputC)
|
|
return emitOpError() << "requires input dimension 2 to have size " << c
|
|
<< ", got " << inputC;
|
|
}
|
|
if (outputShape.hasRank()) {
|
|
const int64_t outputN = outputShape.getDimSize(0);
|
|
const int64_t outputK = outputShape.getDimSize(1);
|
|
const int64_t outputC = outputShape.getDimSize(2);
|
|
if (n != ShapedType::kDynamic && outputN != ShapedType::kDynamic &&
|
|
n != outputN)
|
|
return emitOpError() << "requires values_out dimension 0 to have size "
|
|
<< n << ", got " << outputN;
|
|
if (k == ShapedType::kDynamic)
|
|
k = outputK;
|
|
else if (outputK != ShapedType::kDynamic && k != outputK)
|
|
return emitOpError() << "requires values_out dimension 1 to have size "
|
|
<< k << ", got " << outputK;
|
|
if (c != ShapedType::kDynamic && outputC != ShapedType::kDynamic &&
|
|
c != outputC)
|
|
return emitOpError() << "requires values_out dimension 2 to have size "
|
|
<< c << ", got " << outputC;
|
|
}
|
|
if (k != ShapedType::kDynamic && w != ShapedType::kDynamic && !(k >= w))
|
|
return emitOpError() << "requires dimensions K >= W, got K=" << k
|
|
<< " and W=" << w;
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult ReduceInferReturnTypes(
|
|
ShapeAdaptor operandShape, Type inputType, IntegerAttr axis,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
int64_t axisVal = axis.getValue().getSExtValue();
|
|
if (!operandShape.hasRank() || operandShape.getRank() <= axisVal) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(inputType));
|
|
return success();
|
|
}
|
|
|
|
SmallVector<int64_t> outputShape;
|
|
operandShape.getDims(outputShape);
|
|
outputShape[axisVal] = 1;
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape, inputType));
|
|
return success();
|
|
}
|
|
|
|
#define COMPATIBLE_RETURN_TYPES(OP) \
|
|
bool OP::isCompatibleReturnTypes(TypeRange l, TypeRange r) { \
|
|
if (l.size() != r.size() || l.size() != 1) \
|
|
return false; \
|
|
if (getElementTypeOrSelf(l[0]) != getElementTypeOrSelf(r[0])) \
|
|
return false; \
|
|
return succeeded(verifyCompatibleShape(l[0], r[0])); \
|
|
}
|
|
|
|
#define REDUCE_SHAPE_INFER(OP) \
|
|
LogicalResult OP::inferReturnTypeComponents( \
|
|
MLIRContext *context, ::std::optional<Location> location, \
|
|
OP::Adaptor adaptor, \
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) { \
|
|
Type inputType = \
|
|
llvm::cast<TensorType>(adaptor.getInput().getType()).getElementType(); \
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType()); \
|
|
const Properties &prop = adaptor.getProperties(); \
|
|
return ReduceInferReturnTypes(inputShape, inputType, prop.axis, \
|
|
inferredReturnShapes); \
|
|
} \
|
|
COMPATIBLE_RETURN_TYPES(OP)
|
|
|
|
REDUCE_SHAPE_INFER(tosa::ReduceAllOp)
|
|
REDUCE_SHAPE_INFER(tosa::ReduceAnyOp)
|
|
REDUCE_SHAPE_INFER(tosa::ReduceMaxOp)
|
|
REDUCE_SHAPE_INFER(tosa::ReduceMinOp)
|
|
REDUCE_SHAPE_INFER(tosa::ReduceProductOp)
|
|
REDUCE_SHAPE_INFER(tosa::ReduceSumOp)
|
|
#undef REDUCE_SHAPE_INFER
|
|
COMPATIBLE_RETURN_TYPES(tosa::ConcatOp)
|
|
#undef COMPATIBLE_RETURN_TYPES
|
|
|
|
template <typename T>
|
|
static LogicalResult verifyReduceOp(T op) {
|
|
// All TOSA reduce Ops have input, output and axis.
|
|
TensorType inputType = op.getInput().getType();
|
|
TensorType outputType = op.getOutput().getType();
|
|
int32_t reduceAxis = op.getAxis();
|
|
|
|
if (reduceAxis < 0) {
|
|
op.emitOpError("reduce axis must not be negative");
|
|
return failure();
|
|
}
|
|
if (inputType.hasRank()) {
|
|
int64_t inputRank = inputType.getRank();
|
|
// We allow for a special case where the input/output shape has rank 0 and
|
|
// axis is also 0.
|
|
if (reduceAxis >= inputRank && (reduceAxis != 0 || inputRank != 0)) {
|
|
op.emitOpError("expect input tensor rank (")
|
|
<< inputRank << ") to be larger than reduce axis (" << reduceAxis
|
|
<< ")";
|
|
return failure();
|
|
}
|
|
}
|
|
if (outputType.hasRank()) {
|
|
int64_t outputRank = outputType.getRank();
|
|
if (inputType.hasRank() && outputRank != inputType.getRank()) {
|
|
op.emitOpError(
|
|
"expect output tensor rank to be equal to input tensor rank");
|
|
return failure();
|
|
}
|
|
if (reduceAxis >= outputRank && (reduceAxis != 0 || outputRank != 0)) {
|
|
op.emitOpError("expect output tensor rank (")
|
|
<< outputRank << ") to be larger than reduce axis (" << reduceAxis
|
|
<< ")";
|
|
return failure();
|
|
}
|
|
// We can only verify the reduced dimension size to be 1 if this is not
|
|
// the special case of output rank == 0.
|
|
if (outputRank != 0) {
|
|
auto outputShape = outputType.getShape();
|
|
if (!outputType.isDynamicDim(reduceAxis) &&
|
|
outputShape[reduceAxis] != 1) {
|
|
op.emitOpError("expect reduced dimension size to be 1, got ")
|
|
<< outputShape[reduceAxis];
|
|
return failure();
|
|
}
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ReduceAllOp::verify() { return verifyReduceOp(*this); }
|
|
LogicalResult tosa::ReduceAnyOp::verify() { return verifyReduceOp(*this); }
|
|
LogicalResult tosa::ReduceMaxOp::verify() { return verifyReduceOp(*this); }
|
|
LogicalResult tosa::ReduceMinOp::verify() { return verifyReduceOp(*this); }
|
|
LogicalResult tosa::ReduceProductOp::verify() { return verifyReduceOp(*this); }
|
|
LogicalResult tosa::ReduceSumOp::verify() { return verifyReduceOp(*this); }
|
|
|
|
static LogicalResult NAryInferReturnTypes(
|
|
const ValueShapeRange &operands,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t> outShape;
|
|
if (resolveBroadcastShape(operands, outShape).failed()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents());
|
|
} else {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
|
|
}
|
|
return success();
|
|
}
|
|
|
|
#define NARY_SHAPE_INFER(OP) \
|
|
LogicalResult OP::inferReturnTypeComponents( \
|
|
MLIRContext *context, ::std::optional<Location> location, \
|
|
ValueShapeRange operands, DictionaryAttr attributes, \
|
|
PropertyRef properties, RegionRange regions, \
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) { \
|
|
return NAryInferReturnTypes(operands, inferredReturnShapes); \
|
|
}
|
|
|
|
NARY_SHAPE_INFER(tosa::AbsOp)
|
|
NARY_SHAPE_INFER(tosa::AddOp)
|
|
NARY_SHAPE_INFER(tosa::ArithmeticRightShiftOp)
|
|
NARY_SHAPE_INFER(tosa::BitwiseAndOp)
|
|
NARY_SHAPE_INFER(tosa::BitwiseOrOp)
|
|
NARY_SHAPE_INFER(tosa::BitwiseXorOp)
|
|
NARY_SHAPE_INFER(tosa::BitwiseNotOp)
|
|
NARY_SHAPE_INFER(tosa::CastOp)
|
|
NARY_SHAPE_INFER(tosa::CeilOp)
|
|
NARY_SHAPE_INFER(tosa::ClampOp)
|
|
NARY_SHAPE_INFER(tosa::ClzOp)
|
|
NARY_SHAPE_INFER(tosa::CosOp)
|
|
NARY_SHAPE_INFER(tosa::ExpOp)
|
|
NARY_SHAPE_INFER(tosa::FloorOp)
|
|
NARY_SHAPE_INFER(tosa::GreaterEqualOp)
|
|
NARY_SHAPE_INFER(tosa::GreaterOp)
|
|
NARY_SHAPE_INFER(tosa::IdentityOp)
|
|
NARY_SHAPE_INFER(tosa::IntDivOp)
|
|
NARY_SHAPE_INFER(tosa::LogOp)
|
|
NARY_SHAPE_INFER(tosa::LogicalAndOp)
|
|
NARY_SHAPE_INFER(tosa::LogicalLeftShiftOp)
|
|
NARY_SHAPE_INFER(tosa::LogicalNotOp)
|
|
NARY_SHAPE_INFER(tosa::LogicalOrOp)
|
|
NARY_SHAPE_INFER(tosa::LogicalRightShiftOp)
|
|
NARY_SHAPE_INFER(tosa::LogicalXorOp)
|
|
NARY_SHAPE_INFER(tosa::MaximumOp)
|
|
NARY_SHAPE_INFER(tosa::MinimumOp)
|
|
NARY_SHAPE_INFER(tosa::PowOp)
|
|
NARY_SHAPE_INFER(tosa::ReciprocalOp)
|
|
NARY_SHAPE_INFER(tosa::ReverseOp)
|
|
NARY_SHAPE_INFER(tosa::RsqrtOp)
|
|
NARY_SHAPE_INFER(tosa::SinOp)
|
|
NARY_SHAPE_INFER(tosa::SelectOp)
|
|
NARY_SHAPE_INFER(tosa::SubOp)
|
|
NARY_SHAPE_INFER(tosa::TanhOp)
|
|
NARY_SHAPE_INFER(tosa::ErfOp)
|
|
NARY_SHAPE_INFER(tosa::SigmoidOp)
|
|
#undef PRED_SHAPE_INFER
|
|
|
|
LogicalResult tosa::NegateOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
NegateOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput1().getType());
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(inputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::NegateOp::verify() {
|
|
// Verify same element type
|
|
const Type input1Type = getInput1().getType();
|
|
const Type outputType = getOutput().getType();
|
|
if (verifySameElementTypes(*this, input1Type, outputType).failed())
|
|
return failure();
|
|
|
|
// Verify same shape
|
|
const SmallVector<Type, 2> types = {input1Type, outputType};
|
|
if (failed(verifyCompatibleShapes(types)))
|
|
return emitOpError() << "requires the same shape for input1 and output";
|
|
|
|
const Type input1EType = getStorageElementTypeOrSelf(getInput1().getType());
|
|
const Type input1ZpEType =
|
|
getStorageElementTypeOrSelf(getInput1Zp().getType());
|
|
if (input1EType != input1ZpEType) {
|
|
return emitOpError("expect both input1 and its zero point are the same "
|
|
"element type, got ")
|
|
<< input1EType << " and " << input1ZpEType;
|
|
}
|
|
const Type outputEType = getStorageElementTypeOrSelf(getOutput().getType());
|
|
const Type outputZpEType =
|
|
getStorageElementTypeOrSelf(getOutputZp().getType());
|
|
if (outputEType != outputZpEType) {
|
|
return emitOpError("expect both output and its zero point are the same "
|
|
"element type, got ")
|
|
<< outputEType << " and " << outputZpEType;
|
|
}
|
|
|
|
FailureOr<int64_t> maybeIZp = getInput1ZeroPoint();
|
|
if (succeeded(maybeIZp) && verifyInput1ZeroPoint(*maybeIZp).failed())
|
|
return failure();
|
|
|
|
FailureOr<int64_t> maybeOZp = getOutputZeroPoint();
|
|
if (succeeded(maybeOZp) && verifyOutputZeroPoint(*maybeOZp).failed())
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult poolingInferReturnTypes(
|
|
ShapeAdaptor inputShape, ArrayRef<int64_t> kernel, ArrayRef<int64_t> stride,
|
|
ArrayRef<int64_t> pad,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t> outputShape;
|
|
outputShape.resize(4, ShapedType::kDynamic);
|
|
|
|
// We only know the rank if the input type is unranked.
|
|
if (!inputShape.hasRank()) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
// Batch and number of channels are identical for pooling layer.
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
outputShape[3] = inputShape.getDimSize(3);
|
|
|
|
int64_t height = inputShape.getDimSize(1);
|
|
int64_t width = inputShape.getDimSize(2);
|
|
|
|
if (ShapedType::isStatic(height)) {
|
|
int64_t padded = height + pad[0] + pad[1] - kernel[0];
|
|
outputShape[1] = padded / stride[0] + 1;
|
|
}
|
|
|
|
if (ShapedType::isStatic(width)) {
|
|
int64_t padded = width + pad[2] + pad[3] - kernel[1];
|
|
outputShape[2] = padded / stride[1] + 1;
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
template <typename AdaptorT>
|
|
class ConvInferShapeAdaptor;
|
|
|
|
class ConvInferShapeAdaptorBase {
|
|
protected:
|
|
static void updateIfDynamic(int64_t ¤t, int64_t candidate) {
|
|
if (ShapedType::isDynamic(current))
|
|
current = candidate;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
class ConvInferShapeAdaptor<Conv2DOp::Adaptor>
|
|
: public ConvInferShapeAdaptorBase {
|
|
public:
|
|
explicit ConvInferShapeAdaptor(Conv2DOp::Adaptor adaptor)
|
|
: adaptor(adaptor) {}
|
|
|
|
void inferInputShape(SmallVectorImpl<int64_t> &outputShape,
|
|
SmallVectorImpl<int64_t> &inputSpatial) {
|
|
const ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
if (!inputShape.hasRank())
|
|
return;
|
|
|
|
const int64_t outputBatch = inputShape.getDimSize(0);
|
|
const int64_t inputHeight = inputShape.getDimSize(1);
|
|
const int64_t inputWidth = inputShape.getDimSize(2);
|
|
|
|
outputShape[0] = outputBatch;
|
|
inputSpatial[0] = inputHeight;
|
|
inputSpatial[1] = inputWidth;
|
|
}
|
|
|
|
void inferWeightShape(SmallVectorImpl<int64_t> &outputShape,
|
|
SmallVectorImpl<int64_t> &weightSpatial) {
|
|
const ShapeAdaptor weightShape(adaptor.getWeight().getType());
|
|
if (!weightShape.hasRank())
|
|
return;
|
|
|
|
const int64_t outputChannels = weightShape.getDimSize(0);
|
|
const int64_t kernelHeight = weightShape.getDimSize(1);
|
|
const int64_t kernelWidth = weightShape.getDimSize(2);
|
|
|
|
outputShape[3] = outputChannels;
|
|
weightSpatial[0] = kernelHeight;
|
|
weightSpatial[1] = kernelWidth;
|
|
}
|
|
|
|
int64_t getNumSpatialDims() const { return 2; }
|
|
int64_t getOutputRank() const { return 4; }
|
|
|
|
LogicalResult getSpatialParameters(SmallVector<int64_t> &padValues,
|
|
SmallVector<int64_t> &strideValues,
|
|
SmallVector<int64_t> &dilationValues) {
|
|
padValues.assign(adaptor.getPad().begin(), adaptor.getPad().end());
|
|
strideValues.assign(adaptor.getStride().begin(), adaptor.getStride().end());
|
|
dilationValues.assign(adaptor.getDilation().begin(),
|
|
adaptor.getDilation().end());
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
Conv2DOp::Adaptor adaptor;
|
|
};
|
|
|
|
template <>
|
|
class ConvInferShapeAdaptor<Conv2DBlockScaledOp::Adaptor>
|
|
: public ConvInferShapeAdaptorBase {
|
|
public:
|
|
explicit ConvInferShapeAdaptor(Conv2DBlockScaledOp::Adaptor adaptor)
|
|
: adaptor(adaptor) {}
|
|
|
|
void inferInputShape(SmallVectorImpl<int64_t> &outputShape,
|
|
SmallVectorImpl<int64_t> &inputSpatial) {
|
|
const ShapeAdaptor inputDataShape(adaptor.getInputData().getType());
|
|
if (inputDataShape.hasRank()) {
|
|
const int64_t outputBatch = inputDataShape.getDimSize(0);
|
|
const int64_t inputHeight = inputDataShape.getDimSize(1);
|
|
const int64_t inputWidth = inputDataShape.getDimSize(2);
|
|
|
|
outputShape[0] = outputBatch;
|
|
inputSpatial[0] = inputHeight;
|
|
inputSpatial[1] = inputWidth;
|
|
}
|
|
|
|
const ShapeAdaptor inputScaleShape(adaptor.getInputScale().getType());
|
|
if (!inputScaleShape.hasRank())
|
|
return;
|
|
|
|
const int64_t scaleBatch = inputScaleShape.getDimSize(0);
|
|
const int64_t scaleHeight = inputScaleShape.getDimSize(1);
|
|
const int64_t scaleWidth = inputScaleShape.getDimSize(2);
|
|
|
|
updateIfDynamic(outputShape[0], scaleBatch);
|
|
updateIfDynamic(inputSpatial[0], scaleHeight);
|
|
updateIfDynamic(inputSpatial[1], scaleWidth);
|
|
}
|
|
|
|
void inferWeightShape(SmallVectorImpl<int64_t> &outputShape,
|
|
SmallVectorImpl<int64_t> &weightSpatial) {
|
|
const ShapeAdaptor weightDataShape(adaptor.getWeightData().getType());
|
|
if (weightDataShape.hasRank()) {
|
|
const int64_t outputChannels = weightDataShape.getDimSize(0);
|
|
const int64_t kernelHeight = weightDataShape.getDimSize(1);
|
|
const int64_t kernelWidth = weightDataShape.getDimSize(2);
|
|
|
|
outputShape[3] = outputChannels;
|
|
weightSpatial[0] = kernelHeight;
|
|
weightSpatial[1] = kernelWidth;
|
|
}
|
|
|
|
const ShapeAdaptor weightScaleShape(adaptor.getWeightScale().getType());
|
|
if (!weightScaleShape.hasRank())
|
|
return;
|
|
|
|
const int64_t scaleOutputChannels = weightScaleShape.getDimSize(0);
|
|
const int64_t scaleKernelHeight = weightScaleShape.getDimSize(1);
|
|
const int64_t scaleKernelWidth = weightScaleShape.getDimSize(2);
|
|
|
|
updateIfDynamic(outputShape[3], scaleOutputChannels);
|
|
updateIfDynamic(weightSpatial[0], scaleKernelHeight);
|
|
updateIfDynamic(weightSpatial[1], scaleKernelWidth);
|
|
}
|
|
|
|
int64_t getNumSpatialDims() const { return 2; }
|
|
int64_t getOutputRank() const { return 4; }
|
|
|
|
LogicalResult getSpatialParameters(SmallVector<int64_t> &padValues,
|
|
SmallVector<int64_t> &strideValues,
|
|
SmallVector<int64_t> &dilationValues) {
|
|
if (!tosa::getConstShapeValues(adaptor.getPad().getDefiningOp(),
|
|
padValues) ||
|
|
!tosa::getConstShapeValues(adaptor.getStride().getDefiningOp(),
|
|
strideValues) ||
|
|
!tosa::getConstShapeValues(adaptor.getDilation().getDefiningOp(),
|
|
dilationValues))
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
Conv2DBlockScaledOp::Adaptor adaptor;
|
|
};
|
|
|
|
template <>
|
|
class ConvInferShapeAdaptor<Conv3DOp::Adaptor>
|
|
: public ConvInferShapeAdaptorBase {
|
|
public:
|
|
explicit ConvInferShapeAdaptor(Conv3DOp::Adaptor adaptor)
|
|
: adaptor(adaptor) {}
|
|
|
|
void inferInputShape(SmallVectorImpl<int64_t> &outputShape,
|
|
SmallVectorImpl<int64_t> &inputSpatial) {
|
|
const ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
if (!inputShape.hasRank())
|
|
return;
|
|
|
|
const int64_t outputBatch = inputShape.getDimSize(0);
|
|
const int64_t inputDepth = inputShape.getDimSize(1);
|
|
const int64_t inputHeight = inputShape.getDimSize(2);
|
|
const int64_t inputWidth = inputShape.getDimSize(3);
|
|
|
|
outputShape[0] = outputBatch;
|
|
inputSpatial[0] = inputDepth;
|
|
inputSpatial[1] = inputHeight;
|
|
inputSpatial[2] = inputWidth;
|
|
}
|
|
|
|
void inferWeightShape(SmallVectorImpl<int64_t> &outputShape,
|
|
SmallVectorImpl<int64_t> &weightSpatial) {
|
|
const ShapeAdaptor weightShape(adaptor.getWeight().getType());
|
|
if (!weightShape.hasRank())
|
|
return;
|
|
|
|
const int64_t outputChannels = weightShape.getDimSize(0);
|
|
const int64_t kernelDepth = weightShape.getDimSize(1);
|
|
const int64_t kernelHeight = weightShape.getDimSize(2);
|
|
const int64_t kernelWidth = weightShape.getDimSize(3);
|
|
|
|
outputShape[4] = outputChannels;
|
|
weightSpatial[0] = kernelDepth;
|
|
weightSpatial[1] = kernelHeight;
|
|
weightSpatial[2] = kernelWidth;
|
|
}
|
|
|
|
int64_t getNumSpatialDims() const { return 3; }
|
|
int64_t getOutputRank() const { return 5; }
|
|
|
|
LogicalResult getSpatialParameters(SmallVector<int64_t> &padValues,
|
|
SmallVector<int64_t> &strideValues,
|
|
SmallVector<int64_t> &dilationValues) {
|
|
padValues.assign(adaptor.getPad().begin(), adaptor.getPad().end());
|
|
strideValues.assign(adaptor.getStride().begin(), adaptor.getStride().end());
|
|
dilationValues.assign(adaptor.getDilation().begin(),
|
|
adaptor.getDilation().end());
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
Conv3DOp::Adaptor adaptor;
|
|
};
|
|
|
|
template <typename AdaptorT>
|
|
LogicalResult inferConvReturnTypeComponents(
|
|
AdaptorT adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ConvInferShapeAdaptor<AdaptorT> convShapeAdaptor(adaptor);
|
|
llvm::SmallVector<int64_t> outputShape(convShapeAdaptor.getOutputRank(),
|
|
ShapedType::kDynamic);
|
|
llvm::SmallVector<int64_t> inputSpatial(convShapeAdaptor.getNumSpatialDims(),
|
|
ShapedType::kDynamic);
|
|
llvm::SmallVector<int64_t> weightSpatial(convShapeAdaptor.getNumSpatialDims(),
|
|
ShapedType::kDynamic);
|
|
|
|
convShapeAdaptor.inferInputShape(outputShape, inputSpatial);
|
|
convShapeAdaptor.inferWeightShape(outputShape, weightSpatial);
|
|
|
|
const ShapeAdaptor biasShape = adaptor.getBias().getType();
|
|
if (biasShape.hasRank()) {
|
|
const int64_t biasSize = biasShape.getDimSize(0);
|
|
if (biasSize != 1) {
|
|
const size_t outputChannelDim = convShapeAdaptor.getOutputRank() - 1;
|
|
outputShape[outputChannelDim] =
|
|
ShapedType::isDynamic(outputShape[outputChannelDim])
|
|
? biasSize
|
|
: outputShape[outputChannelDim];
|
|
}
|
|
}
|
|
|
|
SmallVector<int64_t> padValues;
|
|
SmallVector<int64_t> strideValues;
|
|
SmallVector<int64_t> dilationValues;
|
|
if (failed(convShapeAdaptor.getSpatialParameters(padValues, strideValues,
|
|
dilationValues))) {
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
for (int64_t dim = 0; dim < convShapeAdaptor.getNumSpatialDims(); ++dim) {
|
|
if (!ShapedType::isStatic(inputSpatial[dim]) ||
|
|
!ShapedType::isStatic(weightSpatial[dim]))
|
|
continue;
|
|
const int64_t inputSize =
|
|
inputSpatial[dim] + padValues[2 * dim] + padValues[2 * dim + 1];
|
|
const int64_t filterSize =
|
|
(weightSpatial[dim] - 1) * dilationValues[dim] + 1;
|
|
const int64_t unstridedResult = inputSize - filterSize + 1;
|
|
outputShape[dim + 1] = (unstridedResult - 1) / strideValues[dim] + 1;
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult Conv2DOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
Conv2DOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
return inferConvReturnTypeComponents(adaptor, inferredReturnShapes);
|
|
}
|
|
|
|
LogicalResult Conv2DOp::verify() {
|
|
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed() ||
|
|
verifyConvOpErrorIf(*this).failed())
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
LogicalResult Conv2DBlockScaledOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
Conv2DBlockScaledOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
return inferConvReturnTypeComponents(adaptor, inferredReturnShapes);
|
|
}
|
|
|
|
LogicalResult Conv2DBlockScaledOp::verify() {
|
|
if (failed(verifySameElementTypes(*this, getInputData().getType(),
|
|
getWeightData().getType(), "input_data",
|
|
"weight_data")) ||
|
|
failed(verifySameElementTypes(*this, getInputScale().getType(),
|
|
getWeightScale().getType(), "input_scale",
|
|
"weight_scale")) ||
|
|
failed(verifySameElementTypes(*this, getBias().getType(),
|
|
getOutput().getType(), "bias", "output")))
|
|
return failure();
|
|
|
|
// Verify input shape compatibility
|
|
int64_t N = ShapedType::kDynamic;
|
|
int64_t IH = ShapedType::kDynamic;
|
|
int64_t IW = ShapedType::kDynamic;
|
|
int64_t IC = ShapedType::kDynamic;
|
|
int64_t multiplesOfIC = ShapedType::kDynamic;
|
|
int64_t OC = ShapedType::kDynamic;
|
|
int64_t KH = ShapedType::kDynamic;
|
|
int64_t KW = ShapedType::kDynamic;
|
|
|
|
const ShapeAdaptor inputDataShape(getInputData().getType());
|
|
if (inputDataShape.hasRank()) {
|
|
N = inputDataShape.getDimSize(0);
|
|
IH = inputDataShape.getDimSize(1);
|
|
IW = inputDataShape.getDimSize(2);
|
|
IC = inputDataShape.getDimSize(3);
|
|
}
|
|
|
|
const ShapeAdaptor inputScaleShape(getInputScale().getType());
|
|
if (inputScaleShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, N, inputScaleShape.getDimSize(0),
|
|
"input_scale", "batch size")) ||
|
|
failed(tryUpdateDimOrFailure(*this, IH, inputScaleShape.getDimSize(1),
|
|
"input_scale", "input height")) ||
|
|
failed(tryUpdateDimOrFailure(*this, IW, inputScaleShape.getDimSize(2),
|
|
"input_scale", "input width")))
|
|
return failure();
|
|
multiplesOfIC = inputScaleShape.getDimSize(3);
|
|
}
|
|
|
|
const ShapeAdaptor weightDataShape(getWeightData().getType());
|
|
if (weightDataShape.hasRank()) {
|
|
OC = weightDataShape.getDimSize(0);
|
|
KH = weightDataShape.getDimSize(1);
|
|
KW = weightDataShape.getDimSize(2);
|
|
if (failed(tryUpdateDimOrFailure(*this, IC, weightDataShape.getDimSize(3),
|
|
"weight_data", "input channels")))
|
|
return failure();
|
|
}
|
|
|
|
const ShapeAdaptor weightScaleShape(getWeightScale().getType());
|
|
if (weightScaleShape.hasRank()) {
|
|
if (failed(tryUpdateDimOrFailure(*this, OC, weightScaleShape.getDimSize(0),
|
|
"weight_scale", "output channels")) ||
|
|
failed(tryUpdateDimOrFailure(*this, KH, weightScaleShape.getDimSize(1),
|
|
"weight_scale", "kernel height")) ||
|
|
failed(tryUpdateDimOrFailure(*this, KW, weightScaleShape.getDimSize(2),
|
|
"weight_scale", "kernel width")) ||
|
|
failed(tryUpdateDimOrFailure(*this, multiplesOfIC,
|
|
weightScaleShape.getDimSize(3),
|
|
"weight_scale", "input channel blocks")))
|
|
return failure();
|
|
}
|
|
|
|
const uint32_t blockSize = BlockSizeAttr::getBlockSizeValue(getBlockSize());
|
|
if (blockSize != BlockSizeAttr::getBlockSizeValue(BlockSize::BLOCK_SIZE_32))
|
|
return emitOpError("expect block size to be 32, got ") << blockSize;
|
|
// Verify IC is a multiple of block size
|
|
if (ShapedType::isStatic(IC) && IC % blockSize != 0)
|
|
return emitOpError("expect IC to be a multiple of block size, got IC=")
|
|
<< IC << ", block_size=" << blockSize;
|
|
|
|
// Verify multiplesOfIC is IC / block size
|
|
if (ShapedType::isStatic(IC) && ShapedType::isStatic(multiplesOfIC) &&
|
|
multiplesOfIC != IC / blockSize)
|
|
return emitOpError(
|
|
"expect scale operands dimension 2 to equal IC/block_size (")
|
|
<< IC << "/" << blockSize << ")"
|
|
<< ", got " << multiplesOfIC;
|
|
|
|
// Verify pad/stride/dilation values
|
|
SmallVector<int64_t> padValues;
|
|
if (tosa::getConstShapeValues(getPad().getDefiningOp(), padValues)) {
|
|
if (llvm::any_of(padValues, [](int64_t p) { return p < 0; }))
|
|
return emitOpError("expect all padding values to be >= 0, got ")
|
|
<< padValues;
|
|
}
|
|
|
|
SmallVector<int64_t> strideValues;
|
|
if (tosa::getConstShapeValues(getStride().getDefiningOp(), strideValues)) {
|
|
if (llvm::any_of(strideValues, [](int64_t s) { return s < 1; }))
|
|
return emitOpError("expect all stride values to be >= 1, got ")
|
|
<< strideValues;
|
|
}
|
|
|
|
SmallVector<int64_t> dilationValues;
|
|
if (tosa::getConstShapeValues(getDilation().getDefiningOp(),
|
|
dilationValues)) {
|
|
if (llvm::any_of(dilationValues, [](int64_t d) { return d < 1; }))
|
|
return emitOpError("expect all dilation values to be >= 1, got ")
|
|
<< dilationValues;
|
|
}
|
|
|
|
// Verify output shape compatibility
|
|
const ShapeAdaptor outputShape(getOutput().getType());
|
|
if (!padValues.empty() && !strideValues.empty() && !dilationValues.empty() &&
|
|
outputShape.hasRank()) {
|
|
if (failed(verifyConvOutputSize(*this, IH, KH, outputShape.getDimSize(1),
|
|
padValues[0], padValues[1], strideValues[0],
|
|
dilationValues[0], "height", "y", "top",
|
|
"bottom")) ||
|
|
failed(verifyConvOutputSize(*this, IW, KW, outputShape.getDimSize(2),
|
|
padValues[2], padValues[3], strideValues[1],
|
|
dilationValues[1], "width", "x", "left",
|
|
"right")))
|
|
return failure();
|
|
}
|
|
|
|
// Verify bias
|
|
const ShapeAdaptor biasShape(getBias().getType());
|
|
if (biasShape.hasRank() && outputShape.hasRank()) {
|
|
const int64_t biasChannels = biasShape.getDimSize(0);
|
|
const int64_t outputChannels =
|
|
outputShape.getDimSize(outputShape.getRank() - 1);
|
|
if (biasChannels == ShapedType::kDynamic ||
|
|
outputChannels == ShapedType::kDynamic)
|
|
// Skip following checks if biasChannels or outputChannels is dynamic dim
|
|
return success();
|
|
|
|
if (biasChannels != outputChannels && biasChannels != 1)
|
|
return emitOpError(
|
|
"bias channels expected to be equal to output channels (")
|
|
<< outputChannels << ") or 1, got " << biasChannels;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult Conv3DOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
Conv3DOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
return inferConvReturnTypeComponents(adaptor, inferredReturnShapes);
|
|
}
|
|
|
|
LogicalResult Conv3DOp::verify() {
|
|
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed() ||
|
|
verifyConvOpErrorIf(*this).failed())
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
LogicalResult AvgPool2dOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
AvgPool2dOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
const Properties &prop = adaptor.getProperties();
|
|
return poolingInferReturnTypes(inputShape, prop.kernel, prop.stride, prop.pad,
|
|
inferredReturnShapes);
|
|
}
|
|
|
|
LogicalResult AvgPool2dAdaptiveOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
AvgPool2dAdaptiveOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
|
|
llvm::SmallVector<int64_t> kernelValues;
|
|
llvm::SmallVector<int64_t> strideValues;
|
|
llvm::SmallVector<int64_t> padValues;
|
|
if (tosa::getConstShapeValues(adaptor.getKernel().getDefiningOp(),
|
|
kernelValues) &&
|
|
tosa::getConstShapeValues(adaptor.getStride().getDefiningOp(),
|
|
strideValues) &&
|
|
tosa::getConstShapeValues(adaptor.getPad().getDefiningOp(), padValues)) {
|
|
return poolingInferReturnTypes(inputShape, kernelValues, strideValues,
|
|
padValues, inferredReturnShapes);
|
|
}
|
|
|
|
llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamic);
|
|
if (inputShape.hasRank()) {
|
|
// Keep N & C as pooling only changes H & W.
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
outputShape[3] = inputShape.getDimSize(3);
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult MaxPool2dOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
MaxPool2dOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
const Properties &prop = adaptor.getProperties();
|
|
return poolingInferReturnTypes(inputShape, prop.kernel, prop.stride, prop.pad,
|
|
inferredReturnShapes);
|
|
}
|
|
|
|
LogicalResult MaxPool2dAdaptiveOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
MaxPool2dAdaptiveOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
|
|
llvm::SmallVector<int64_t> kernelValues;
|
|
llvm::SmallVector<int64_t> strideValues;
|
|
llvm::SmallVector<int64_t> padValues;
|
|
if (tosa::getConstShapeValues(adaptor.getKernel().getDefiningOp(),
|
|
kernelValues) &&
|
|
tosa::getConstShapeValues(adaptor.getStride().getDefiningOp(),
|
|
strideValues) &&
|
|
tosa::getConstShapeValues(adaptor.getPad().getDefiningOp(), padValues)) {
|
|
return poolingInferReturnTypes(inputShape, kernelValues, strideValues,
|
|
padValues, inferredReturnShapes);
|
|
}
|
|
|
|
llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamic);
|
|
if (inputShape.hasRank()) {
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
outputShape[3] = inputShape.getDimSize(3);
|
|
}
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult MaxPool2dOp::verify() {
|
|
if (failed(verifySameElementTypes(*this, /* intype = */ getInput().getType(),
|
|
/* outType = */ getOutput().getType())))
|
|
return failure();
|
|
|
|
if (failed(verifyPoolingOp(*this)))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult MaxPool2dAdaptiveOp::verify() {
|
|
if (failed(verifySameElementTypes(*this, /* intype = */ getInput().getType(),
|
|
/* outType = */ getOutput().getType())))
|
|
return failure();
|
|
|
|
AdaptivePoolingConstShapeValues values;
|
|
extractAdaptivePoolingConstShapeOperands(*this, values);
|
|
|
|
if (failed(verifyPoolingOpImpl(getOperation(), values.kernel, values.stride,
|
|
values.pad, getInput(), getOutput())))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
DepthwiseConv2DOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamic);
|
|
|
|
int64_t inputWidth = ShapedType::kDynamic;
|
|
int64_t inputHeight = ShapedType::kDynamic;
|
|
int64_t inputChannels = ShapedType::kDynamic;
|
|
|
|
int64_t weightWidth = ShapedType::kDynamic;
|
|
int64_t weightHeight = ShapedType::kDynamic;
|
|
int64_t depthChannels = ShapedType::kDynamic;
|
|
|
|
// Input shape describes input width/height and batch.
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
if (inputShape.hasRank()) {
|
|
outputShape[0] = inputShape.getDimSize(0);
|
|
inputHeight = inputShape.getDimSize(1);
|
|
inputWidth = inputShape.getDimSize(2);
|
|
inputChannels = inputShape.getDimSize(3);
|
|
}
|
|
|
|
// Weight shapes describes the filter width/height and the output channels.
|
|
ShapeAdaptor weightShape(adaptor.getWeight().getType());
|
|
if (weightShape.hasRank()) {
|
|
weightHeight = weightShape.getDimSize(0);
|
|
weightWidth = weightShape.getDimSize(1);
|
|
inputChannels = ShapedType::isDynamic(inputChannels)
|
|
? weightShape.getDimSize(2)
|
|
: inputChannels;
|
|
depthChannels = weightShape.getDimSize(3);
|
|
}
|
|
|
|
// If both inputChannels and depthChannels are available we can determine
|
|
// the output channels.
|
|
if (ShapedType::isStatic(inputChannels) &&
|
|
ShapedType::isStatic(depthChannels)) {
|
|
outputShape[3] = inputChannels * depthChannels;
|
|
}
|
|
|
|
// Bias shape can describe the output channels.
|
|
ShapeAdaptor biasShape(adaptor.getBias().getType());
|
|
if (biasShape.hasRank() && ShapedType::isDynamic(outputShape[3])) {
|
|
int64_t bc = biasShape.getDimSize(0);
|
|
if (bc != ShapedType::kDynamic && bc != 1)
|
|
outputShape[3] = bc;
|
|
}
|
|
|
|
llvm::ArrayRef<int64_t> dilation = adaptor.getDilation();
|
|
llvm::ArrayRef<int64_t> padding = adaptor.getPad();
|
|
llvm::ArrayRef<int64_t> stride = adaptor.getStride();
|
|
|
|
if (ShapedType::isStatic(inputHeight) && ShapedType::isStatic(weightHeight)) {
|
|
int64_t inputSize = inputHeight + padding[0] + padding[1];
|
|
int64_t filterSize = (weightHeight - 1) * dilation[0] + 1;
|
|
int64_t unstridedResult = inputSize - filterSize + 1;
|
|
outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
|
|
}
|
|
|
|
if (ShapedType::isStatic(inputWidth) && ShapedType::isStatic(weightWidth)) {
|
|
int64_t inputSize = inputWidth + padding[2] + padding[3];
|
|
int64_t filterSize = (weightWidth - 1) * dilation[1] + 1;
|
|
int64_t unstridedResult = inputSize - filterSize + 1;
|
|
outputShape[2] = (unstridedResult - 1) / stride[1] + 1;
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult DepthwiseConv2DOp::verify() {
|
|
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed() ||
|
|
verifyConvOpErrorIf(*this).failed())
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
TransposeConv2DOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamic);
|
|
|
|
int64_t inputWidth = ShapedType::kDynamic;
|
|
int64_t inputHeight = ShapedType::kDynamic;
|
|
int64_t weightWidth = ShapedType::kDynamic;
|
|
int64_t weightHeight = ShapedType::kDynamic;
|
|
|
|
// Input shape describes input width/height and batch.
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
if (inputShape.hasRank()) {
|
|
outputShape[0] = ShapedType::isDynamic(outputShape[0])
|
|
? inputShape.getDimSize(0)
|
|
: outputShape[0];
|
|
inputHeight = inputShape.getDimSize(1);
|
|
inputWidth = inputShape.getDimSize(2);
|
|
}
|
|
|
|
// Weight shapes describes the filter width/height and the output channels.
|
|
ShapeAdaptor weightShape(adaptor.getWeight().getType());
|
|
if (weightShape.hasRank()) {
|
|
outputShape[3] = ShapedType::isDynamic(outputShape[3])
|
|
? weightShape.getDimSize(0)
|
|
: outputShape[3];
|
|
weightHeight = weightShape.getDimSize(1);
|
|
weightWidth = weightShape.getDimSize(2);
|
|
}
|
|
|
|
// Bias shape can describe the output channels.
|
|
ShapeAdaptor biasShape(adaptor.getBias().getType());
|
|
if (biasShape.hasRank() && ShapedType::isDynamic(outputShape[3])) {
|
|
int64_t bc = biasShape.getDimSize(0);
|
|
if (bc != ShapedType::kDynamic && bc != 1)
|
|
outputShape[3] = bc;
|
|
}
|
|
|
|
llvm::ArrayRef<int64_t> padding = adaptor.getOutPad();
|
|
llvm::ArrayRef<int64_t> stride = adaptor.getStride();
|
|
|
|
if (ShapedType::isStatic(inputHeight) && ShapedType::isStatic(weightHeight)) {
|
|
int64_t calculateSize =
|
|
(inputHeight - 1) * stride[0] + padding[0] + padding[1] + weightHeight;
|
|
outputShape[1] =
|
|
ShapedType::isDynamic(outputShape[1]) ? calculateSize : outputShape[1];
|
|
}
|
|
|
|
if (ShapedType::isStatic(inputWidth) && ShapedType::isStatic(weightWidth)) {
|
|
int64_t calculateSize =
|
|
(inputWidth - 1) * stride[1] + padding[2] + padding[3] + weightWidth;
|
|
outputShape[2] =
|
|
ShapedType::isDynamic(outputShape[2]) ? calculateSize : outputShape[2];
|
|
}
|
|
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult TransposeConv2DOp::verify() {
|
|
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed())
|
|
return failure();
|
|
|
|
const llvm::ArrayRef<int64_t> strides = getStride();
|
|
const int64_t strideY = strides[0];
|
|
const int64_t strideX = strides[1];
|
|
|
|
if (strideY < 1 || strideX < 1)
|
|
return emitOpError("expect all stride values to be >= 1, got [")
|
|
<< strides << "]";
|
|
|
|
const auto checkPadAgainstKernelDim =
|
|
[this](int64_t padValue, int64_t kernelDimSize, llvm::StringRef padName,
|
|
llvm::StringRef kernelDimName) -> LogicalResult {
|
|
if (padValue <= -kernelDimSize)
|
|
return emitOpError("expected ")
|
|
<< padName << " > -" << kernelDimName << ", but got: " << padName
|
|
<< "=" << padValue << " and " << kernelDimName << "="
|
|
<< kernelDimSize;
|
|
return success();
|
|
};
|
|
|
|
const llvm::ArrayRef<int64_t> padding = getOutPad();
|
|
const int64_t outPadTop = padding[0];
|
|
const int64_t outPadBottom = padding[1];
|
|
const int64_t outPadLeft = padding[2];
|
|
const int64_t outPadRight = padding[3];
|
|
|
|
const auto weightType =
|
|
llvm::dyn_cast<RankedTensorType>(getWeight().getType());
|
|
|
|
if (weightType) {
|
|
const int64_t kernelHeight = weightType.getDimSize(1);
|
|
if (ShapedType::isStatic(kernelHeight)) {
|
|
if (failed(checkPadAgainstKernelDim(outPadTop, kernelHeight,
|
|
"out_pad_top", "KH")))
|
|
return failure();
|
|
|
|
if (failed(checkPadAgainstKernelDim(outPadBottom, kernelHeight,
|
|
"out_pad_bottom", "KH")))
|
|
return failure();
|
|
}
|
|
|
|
const int64_t kernelWidth = weightType.getDimSize(2);
|
|
if (ShapedType::isStatic(kernelWidth)) {
|
|
if (failed(checkPadAgainstKernelDim(outPadLeft, kernelWidth,
|
|
"out_pad_left", "KW")))
|
|
return failure();
|
|
|
|
if (failed(checkPadAgainstKernelDim(outPadRight, kernelWidth,
|
|
"out_pad_right", "KW")))
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
// Rest of the checks depend on the output type being a RankedTensorType
|
|
const auto outputType =
|
|
llvm::dyn_cast<RankedTensorType>(getOutput().getType());
|
|
if (!outputType)
|
|
return success();
|
|
|
|
const auto inputType = llvm::dyn_cast<RankedTensorType>(getInput().getType());
|
|
if (inputType && weightType) {
|
|
const int64_t inputHeight = inputType.getDimSize(1);
|
|
const int64_t kernelHeight = weightType.getDimSize(1);
|
|
const int64_t outputHeight = outputType.getDimSize(1);
|
|
|
|
if (ShapedType::isStatic(inputHeight) &&
|
|
ShapedType::isStatic(outputHeight)) {
|
|
if (outputHeight !=
|
|
(inputHeight - 1) * strideY + outPadTop + outPadBottom + kernelHeight)
|
|
return emitOpError(
|
|
"dimension mismatch: expected OH == (IH - 1) * stride_y "
|
|
"+ out_pad_top + out_pad_bottom + KH, but got ")
|
|
<< outputHeight << " != (" << inputHeight << " - 1) * "
|
|
<< strideY << " + " << outPadTop << " + " << outPadBottom
|
|
<< " + " << kernelHeight;
|
|
}
|
|
|
|
const int64_t inputWidth = inputType.getDimSize(2);
|
|
const int64_t kernelWidth = weightType.getDimSize(2);
|
|
const int64_t outputWidth = outputType.getDimSize(2);
|
|
|
|
if (ShapedType::isStatic(inputWidth) && ShapedType::isStatic(outputWidth)) {
|
|
if (outputWidth !=
|
|
(inputWidth - 1) * strideX + outPadLeft + outPadRight + kernelWidth)
|
|
return emitOpError(
|
|
"dimension mismatch: expected OW == (IW - 1) * stride_x "
|
|
"+ out_pad_left + out_pad_right + KW, but got ")
|
|
<< outputWidth << " != (" << inputWidth << " - 1) * " << strideX
|
|
<< " + " << outPadLeft << " + " << outPadRight << " + "
|
|
<< kernelWidth;
|
|
}
|
|
}
|
|
|
|
const auto biasType = llvm::dyn_cast<RankedTensorType>(getBias().getType());
|
|
|
|
if (!biasType)
|
|
return success();
|
|
|
|
const int64_t biasChannels = biasType.getDimSize(0);
|
|
|
|
// Skip further checks if bias is dynamic
|
|
if (biasChannels == ShapedType::kDynamic)
|
|
return success();
|
|
|
|
const int64_t outputChannels = outputType.getDimSize(3);
|
|
if (!ShapedType::isDynamic(outputChannels) &&
|
|
biasChannels != outputChannels && biasChannels != 1)
|
|
return emitOpError(
|
|
"bias channels expected to be equal to output channels (")
|
|
<< outputChannels << ") or 1, got " << biasChannels;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult RescaleOp::verify() {
|
|
const auto inputType = llvm::cast<ShapedType>(getInput().getType());
|
|
auto inputElementType =
|
|
getStorageElementTypeOrSelf(inputType.getElementType());
|
|
if (!mlir::isa<IntegerType>(inputElementType)) {
|
|
emitOpError("expect input to have integer element type, got ")
|
|
<< inputElementType;
|
|
return failure();
|
|
}
|
|
|
|
const auto outputType = llvm::cast<ShapedType>(getOutput().getType());
|
|
auto outputElementType =
|
|
getStorageElementTypeOrSelf(outputType.getElementType());
|
|
if (!mlir::isa<IntegerType>(outputElementType)) {
|
|
emitOpError("expect output to have integer element type, got ")
|
|
<< outputElementType;
|
|
return failure();
|
|
}
|
|
|
|
if (verifyRescaleValueAndZpTypes(*this, getInput(), getInputZp(), "input")
|
|
.failed())
|
|
return failure();
|
|
|
|
if (verifyRescaleValueAndZpTypes(*this, getOutput(), getOutputZp(), "output")
|
|
.failed())
|
|
return failure();
|
|
|
|
FailureOr<int64_t> maybeIZp = getInputZeroPoint();
|
|
if (succeeded(maybeIZp) && verifyInputZeroPoint(*maybeIZp).failed())
|
|
return failure();
|
|
|
|
FailureOr<int64_t> maybeOZp = getOutputZeroPoint();
|
|
if (succeeded(maybeOZp) && verifyOutputZeroPoint(*maybeOZp).failed())
|
|
return failure();
|
|
|
|
const auto multiplierType = llvm::cast<ShapedType>(getMultiplier().getType());
|
|
// multiplier element type must be i32 for scale32 = true
|
|
if (getScale32() && !multiplierType.getElementType().isInteger(32)) {
|
|
emitOpError("expect i32 element type for multiplier for scale32=true, got ")
|
|
<< multiplierType.getElementType();
|
|
return failure();
|
|
}
|
|
|
|
// multiplier element type must be i16 for scale32 = false
|
|
if (!getScale32() && !multiplierType.getElementType().isInteger(16)) {
|
|
emitOpError(
|
|
"expect i16 element type for multiplier for scale32=false, got ")
|
|
<< multiplierType.getElementType();
|
|
return failure();
|
|
}
|
|
|
|
if (!inputType.hasRank())
|
|
return success();
|
|
|
|
// multiplier/shift must have shape = {numChannels},
|
|
// where numChannel is 1 if per_channel = false
|
|
// otherwise numChannel is dimension in input shape's last axis
|
|
int64_t numChannels = 1;
|
|
if (getPerChannel()) {
|
|
if (inputType.getRank() < 1) {
|
|
emitOpError("requires input to be at least rank 1 when per_channel is "
|
|
"true, but got rank ")
|
|
<< inputType.getRank();
|
|
return failure();
|
|
}
|
|
numChannels = inputType.getDimSize(inputType.getRank() - 1);
|
|
}
|
|
|
|
if (outputType.hasRank()) {
|
|
if (failed(verifyOutputShapeCompatibleWithExpected(
|
|
getOperation(), outputType, inputType.getShape())))
|
|
return failure();
|
|
}
|
|
|
|
if (multiplierType.hasRank()) {
|
|
ArrayRef<int64_t> multiplierShape = multiplierType.getShape();
|
|
// multiplier input has rank 1 by dialect definition
|
|
if (multiplierShape[0] != ShapedType::kDynamic &&
|
|
multiplierShape[0] != numChannels) {
|
|
emitOpError("expect shape of { ")
|
|
<< numChannels << " } for multiplier input, got { "
|
|
<< multiplierShape[0] << " }";
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
const auto shiftType = llvm::cast<ShapedType>(getShift().getType());
|
|
if (shiftType.hasRank()) {
|
|
ArrayRef<int64_t> shiftShape = shiftType.getShape();
|
|
// shift input has rank 1 by dialect definition
|
|
if (shiftShape[0] != ShapedType::kDynamic && shiftShape[0] != numChannels) {
|
|
emitOpError("expect shape of { ")
|
|
<< numChannels << " } for shift input, got { " << shiftShape[0]
|
|
<< " }";
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult RescaleOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
RescaleOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
ShapeAdaptor inputShape(adaptor.getInput().getType());
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(inputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult CastFromBlockScaledOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
CastFromBlockScaledOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
const ShapeAdaptor inputShape(adaptor.getInputData().getType());
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(inputShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult CastFromBlockScaledOp::verify() {
|
|
const Type inputDataType = getInputData().getType();
|
|
const Type outputDataType = getResult().getType();
|
|
if (failed(verifyCompatibleShape(inputDataType, outputDataType)))
|
|
return emitOpError() << "require compatible shapes for input_data ("
|
|
<< inputDataType << ") and " << "output_data ("
|
|
<< outputDataType << ")";
|
|
|
|
const ShapeAdaptor inputDataShape = ShapeAdaptor(inputDataType);
|
|
|
|
if (inputDataShape.hasRank()) {
|
|
const unsigned int blockSize =
|
|
BlockSizeAttr::getBlockSizeValue(getBlockSize());
|
|
if (blockSize != BlockSizeAttr::getBlockSizeValue(BlockSize::BLOCK_SIZE_32))
|
|
return emitOpError("expect block size to be 32, got ") << blockSize;
|
|
const int64_t inputDataLastDim =
|
|
inputDataShape.getDimSize(inputDataShape.getRank() - 1);
|
|
if (inputDataLastDim % blockSize != 0)
|
|
return emitOpError() << "expect last dimension of input_data ("
|
|
<< inputDataLastDim
|
|
<< ") to be divisible by block_size (" << blockSize
|
|
<< ")";
|
|
|
|
const Type inputScaleType = getInputScale().getType();
|
|
const ShapeAdaptor inputScaleShape = ShapeAdaptor(inputScaleType);
|
|
|
|
if (inputScaleShape.hasRank()) {
|
|
SmallVector<int64_t> inputDataDims, inputScaleDims;
|
|
inputDataShape.getDims(inputDataDims);
|
|
inputScaleShape.getDims(inputScaleDims);
|
|
|
|
if (inputDataDims.size() != inputScaleDims.size() ||
|
|
failed(verifyCompatibleShape(
|
|
ArrayRef<int64_t>(inputDataDims).drop_back(1),
|
|
ArrayRef<int64_t>(inputScaleDims).drop_back(1))))
|
|
return emitOpError()
|
|
<< "require compatible shapes for input_data (" << inputDataType
|
|
<< ") and " << "input_scale (" << inputScaleType
|
|
<< ") except for the last dimension";
|
|
|
|
const SmallVector<int64_t, 2> dimsToCheck{inputDataLastDim / blockSize,
|
|
inputScaleDims.back()};
|
|
if (ShapedType::isStatic(inputDataLastDim) &&
|
|
failed(verifyCompatibleDims(dimsToCheck)))
|
|
return emitOpError()
|
|
<< "expect last dimension of input_scale ("
|
|
<< inputScaleDims.back()
|
|
<< ") to be equal to last dimension of input_data / block_size ("
|
|
<< inputDataDims.back() / blockSize << ")";
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult CastToBlockScaledOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
CastToBlockScaledOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
const ShapeAdaptor inputShape(adaptor.getInputData().getType());
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(inputShape));
|
|
if (!inputShape.hasRank())
|
|
return success();
|
|
|
|
// Calculate output_scale shape if ranked input provided
|
|
SmallVector<int64_t> outputScaleShape;
|
|
inputShape.getDims(outputScaleShape);
|
|
const int64_t lastDimLoc = inputShape.getRank() - 1;
|
|
const int64_t lastDimSize = inputShape.getDimSize(lastDimLoc);
|
|
if (ShapedType::isStatic(lastDimSize)) {
|
|
const unsigned int blockSize =
|
|
BlockSizeAttr::getBlockSizeValue(adaptor.getBlockSize());
|
|
outputScaleShape[lastDimLoc] = lastDimSize / blockSize;
|
|
}
|
|
inferredReturnShapes.push_back(ShapedTypeComponents(outputScaleShape));
|
|
return success();
|
|
}
|
|
|
|
LogicalResult CastToBlockScaledOp::verify() {
|
|
const Type inputDataType = getInputData().getType();
|
|
const Type outputDataType = getResult(0).getType();
|
|
if (failed(verifyCompatibleShape(inputDataType, outputDataType)))
|
|
return emitOpError() << "require compatible shapes for input_data ("
|
|
<< inputDataType << ") and " << "output_data ("
|
|
<< outputDataType << ")";
|
|
|
|
const unsigned int blockSize =
|
|
BlockSizeAttr::getBlockSizeValue(getBlockSize());
|
|
if (blockSize != BlockSizeAttr::getBlockSizeValue(BlockSize::BLOCK_SIZE_32))
|
|
return emitOpError("expect block size to be 32, got ") << blockSize;
|
|
const ShapeAdaptor inputDataShape = ShapeAdaptor(inputDataType);
|
|
if (inputDataShape.hasRank()) {
|
|
const int64_t inputDataLastDim =
|
|
inputDataShape.getDimSize(inputDataShape.getRank() - 1);
|
|
if (ShapedType::isStatic(inputDataLastDim) &&
|
|
inputDataLastDim % blockSize != 0)
|
|
return emitOpError() << "expect last dimension of input_data ("
|
|
<< inputDataLastDim
|
|
<< ") to be divisible by block_size (" << blockSize
|
|
<< ")";
|
|
}
|
|
|
|
const ShapeAdaptor outputDataShape = ShapeAdaptor(outputDataType);
|
|
const Type outputScaleType = getResult(1).getType();
|
|
const ShapeAdaptor outputScaleShape = ShapeAdaptor(outputScaleType);
|
|
if (outputDataShape.hasRank() && outputScaleShape.hasRank()) {
|
|
SmallVector<int64_t> outputDataDims, outputScaleDims;
|
|
outputDataShape.getDims(outputDataDims);
|
|
outputScaleShape.getDims(outputScaleDims);
|
|
|
|
if (outputDataDims.size() != outputScaleDims.size() ||
|
|
failed(verifyCompatibleShape(
|
|
ArrayRef<int64_t>(outputDataDims).drop_back(1),
|
|
ArrayRef<int64_t>(outputScaleDims).drop_back(1))))
|
|
return emitOpError() << "require compatible shapes for output_data ("
|
|
<< outputDataType << ") and " << "output_scale ("
|
|
<< outputScaleType
|
|
<< ") except for the last dimension";
|
|
|
|
const int64_t outputDataLastDim = outputDataDims.back();
|
|
const SmallVector<int64_t, 2> dimsToCheck{outputDataLastDim / blockSize,
|
|
outputScaleDims.back()};
|
|
if (ShapedType::isStatic(outputDataLastDim) &&
|
|
failed(verifyCompatibleDims(dimsToCheck)))
|
|
return emitOpError()
|
|
<< "expect last dimension of output_scale ("
|
|
<< outputScaleDims.back()
|
|
<< ") to be equal to last dimension of output_data / block_size ("
|
|
<< outputDataDims.back() / blockSize << ")";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult IfOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
IfOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<tosa::YieldOp> yieldOps;
|
|
for (Region *region : adaptor.getRegions()) {
|
|
for (auto &block : *region)
|
|
if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
|
|
yieldOps.push_back(returnOp);
|
|
}
|
|
|
|
if (yieldOps.empty())
|
|
return failure();
|
|
|
|
// Get the initial type information for the yield op.
|
|
llvm::SmallVector<ValueKnowledge> resultKnowledge;
|
|
resultKnowledge.reserve(yieldOps.front().getNumOperands());
|
|
for (auto operand : yieldOps.front().getOperands()) {
|
|
resultKnowledge.push_back(
|
|
ValueKnowledge::getKnowledgeFromType(operand.getType()));
|
|
}
|
|
|
|
for (auto yieldOp : yieldOps) {
|
|
if (resultKnowledge.size() != yieldOp.getNumOperands())
|
|
return failure();
|
|
|
|
for (const auto &it : llvm::enumerate(yieldOp.getOperands())) {
|
|
int32_t index = it.index();
|
|
auto meet = ValueKnowledge::meet(
|
|
resultKnowledge[index],
|
|
ValueKnowledge::getKnowledgeFromType(it.value().getType()));
|
|
if (!meet)
|
|
continue;
|
|
resultKnowledge[index] = meet;
|
|
}
|
|
}
|
|
|
|
for (const ValueKnowledge &result : resultKnowledge) {
|
|
inferredReturnShapes.push_back(result.getShapedTypeComponents());
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult WhileOp::inferReturnTypeComponents(
|
|
MLIRContext *context, ::std::optional<Location> location,
|
|
WhileOp::Adaptor adaptor,
|
|
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
|
|
llvm::SmallVector<tosa::YieldOp> yieldOps;
|
|
for (auto &block : adaptor.getBodyGraph())
|
|
if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
|
|
yieldOps.push_back(returnOp);
|
|
|
|
// TOSA's while must have a tosa.yield as its terminator. If not found this
|
|
// tosa.while is invalid.
|
|
if (yieldOps.empty())
|
|
return failure();
|
|
|
|
// Get the initial type information from the operand types.
|
|
llvm::SmallVector<ValueKnowledge> resultKnowledge;
|
|
resultKnowledge.reserve(yieldOps.front().getNumOperands());
|
|
for (auto operand : yieldOps.front().getOperands()) {
|
|
resultKnowledge.push_back(
|
|
ValueKnowledge::getKnowledgeFromType(operand.getType()));
|
|
}
|
|
|
|
for (auto yieldOp : yieldOps) {
|
|
if (resultKnowledge.size() != yieldOp.getNumOperands())
|
|
return failure();
|
|
|
|
for (const auto &it : llvm::enumerate(yieldOp.getOperands())) {
|
|
int32_t index = it.index();
|
|
if (auto meet = ValueKnowledge::meet(
|
|
resultKnowledge[index],
|
|
ValueKnowledge::getKnowledgeFromType(it.value().getType()))) {
|
|
resultKnowledge[index] = meet;
|
|
}
|
|
}
|
|
}
|
|
|
|
for (const ValueKnowledge &result : resultKnowledge) {
|
|
inferredReturnShapes.push_back(result.getShapedTypeComponents());
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
std::optional<SmallVector<int64_t, 4>> ApplyScaleOp::getShapeForUnroll() {
|
|
if (auto vt = llvm::dyn_cast<VectorType>(getType()))
|
|
return llvm::to_vector<4>(vt.getShape());
|
|
return std::nullopt;
|
|
}
|
|
|
|
static void printInitializationList(OpAsmPrinter &parser,
|
|
Block::BlockArgListType blocksArgs,
|
|
ValueRange initializers,
|
|
StringRef prefix = "") {
|
|
assert(blocksArgs.size() == initializers.size() &&
|
|
"expected same length of arguments and initializers");
|
|
if (initializers.empty())
|
|
return;
|
|
|
|
parser << prefix << '(';
|
|
llvm::interleaveComma(
|
|
llvm::zip(blocksArgs, initializers), parser,
|
|
[&](auto it) { parser << std::get<0>(it) << " = " << std::get<1>(it); });
|
|
parser << ")";
|
|
}
|
|
|
|
// parse and print of IfOp refer to the implementation of SCF dialect.
|
|
ParseResult IfOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
// Create the regions for 'then'.
|
|
result.regions.reserve(2);
|
|
Region *thenRegion = result.addRegion();
|
|
Region *elseRegion = result.addRegion();
|
|
|
|
OpAsmParser::UnresolvedOperand cond;
|
|
|
|
if (parser.parseOperand(cond))
|
|
return failure();
|
|
|
|
SmallVector<OpAsmParser::Argument, 4> regionArgs;
|
|
SmallVector<OpAsmParser::UnresolvedOperand, 4> operands;
|
|
|
|
// Parse the optional block arguments
|
|
OptionalParseResult listResult =
|
|
parser.parseOptionalAssignmentList(regionArgs, operands);
|
|
if (listResult.has_value() && failed(listResult.value()))
|
|
return failure();
|
|
|
|
// Parse a colon.
|
|
if (failed(parser.parseColon()))
|
|
return parser.emitError(parser.getCurrentLocation(),
|
|
"expected type for condition operand");
|
|
|
|
// Parse the type of the condition operand
|
|
Type condType;
|
|
if (failed(parser.parseType(condType)))
|
|
return parser.emitError(parser.getCurrentLocation(),
|
|
"expected type for condition operand");
|
|
|
|
// Resolve operand with provided type
|
|
if (failed(parser.resolveOperand(cond, condType, result.operands)))
|
|
return failure();
|
|
|
|
// Parse optional block arg types
|
|
if (listResult.has_value()) {
|
|
FunctionType functionType;
|
|
|
|
if (failed(parser.parseType(functionType)))
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected list of types for block arguments "
|
|
<< "followed by arrow type and list of return types";
|
|
|
|
result.addTypes(functionType.getResults());
|
|
|
|
if (functionType.getNumInputs() != operands.size()) {
|
|
return parser.emitError(parser.getCurrentLocation())
|
|
<< "expected as many input types as operands " << "(expected "
|
|
<< operands.size() << " got " << functionType.getNumInputs()
|
|
<< ")";
|
|
}
|
|
|
|
// Resolve input operands.
|
|
if (failed(parser.resolveOperands(operands, functionType.getInputs(),
|
|
parser.getCurrentLocation(),
|
|
result.operands)))
|
|
return failure();
|
|
} else {
|
|
// Parse optional results type list.
|
|
if (parser.parseOptionalArrowTypeList(result.types))
|
|
return failure();
|
|
}
|
|
|
|
// Parse the 'then' region.
|
|
if (parser.parseRegion(*thenRegion, /*arguments=*/{}, /*argTypes=*/{}))
|
|
return failure();
|
|
|
|
// If we find an 'else' keyword then parse the 'else' region.
|
|
if (!parser.parseOptionalKeyword("else")) {
|
|
if (parser.parseRegion(*elseRegion, /*arguments=*/{}, /*argTypes=*/{}))
|
|
return failure();
|
|
}
|
|
|
|
// Parse the optional attribute list.
|
|
if (parser.parseOptionalAttrDict(result.attributes))
|
|
return failure();
|
|
return success();
|
|
}
|
|
|
|
void IfOp::print(OpAsmPrinter &p) {
|
|
p << " " << getCondition();
|
|
|
|
printInitializationList(p, getThenGraph().front().getArguments(),
|
|
getInputList(), " ");
|
|
p << " : ";
|
|
p << getCondition().getType();
|
|
|
|
if (!getInputList().empty()) {
|
|
p << " (";
|
|
llvm::interleaveComma(getInputList().getTypes(), p);
|
|
p << ")";
|
|
}
|
|
p.printArrowTypeList(getResultTypes());
|
|
p << " ";
|
|
|
|
p.printRegion(getThenGraph());
|
|
|
|
// Print the 'else' regions if it exists and has a block.
|
|
auto &elseRegion = getElseGraph();
|
|
if (!elseRegion.empty()) {
|
|
p << " else ";
|
|
p.printRegion(elseRegion);
|
|
}
|
|
|
|
p.printOptionalAttrDict((*this)->getAttrs());
|
|
}
|
|
|
|
LogicalResult IfOp::verify() {
|
|
if (errorIfTypeOrShapeMismatch(*this, getThenGraph().front().getArguments(),
|
|
"'then_graph' arguments", getInputList(),
|
|
"'input_list'")
|
|
.failed())
|
|
return failure();
|
|
|
|
if (errorIfTypeOrShapeMismatch(*this, getElseGraph().front().getArguments(),
|
|
"'else_graph' arguments", getInputList(),
|
|
"'input_list'")
|
|
.failed())
|
|
return failure();
|
|
|
|
// MLIR will verify the absence of the terminator for us if otherwise.
|
|
if (getThenGraph().front().mightHaveTerminator()) {
|
|
auto thenYield =
|
|
dyn_cast<tosa::YieldOp>(getThenGraph().front().getTerminator());
|
|
if (thenYield && errorIfTypeOrShapeMismatch(
|
|
*this, thenYield.getInputs(), "'then_graph' results",
|
|
getOutputList(), "'output_list'")
|
|
.failed())
|
|
return failure();
|
|
}
|
|
|
|
// MLIR will verify the absence of the terminator for us if otherwise.
|
|
if (getElseGraph().front().mightHaveTerminator()) {
|
|
auto elseYield =
|
|
dyn_cast<tosa::YieldOp>(getElseGraph().front().getTerminator());
|
|
if (elseYield && errorIfTypeOrShapeMismatch(
|
|
*this, elseYield.getInputs(), "'else_graph' results",
|
|
getOutputList(), "'output_list'")
|
|
.failed())
|
|
return failure();
|
|
}
|
|
|
|
auto condType = getCondition().getType();
|
|
if (errorIfShapeNotSizeOne(*this, condType).failed())
|
|
return emitOpError() << "'condition' must be a size 1 tensor, got "
|
|
<< condType;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult WhileOp::verify() {
|
|
if (errorIfTypeOrShapeMismatch(*this, getInputList(), "'input_list'",
|
|
getOutputList(), "'output_list'")
|
|
.failed())
|
|
return failure();
|
|
|
|
if (errorIfTypeOrShapeMismatch(*this, getCondGraph().front().getArguments(),
|
|
"'cond_graph' arguments", getInputList(),
|
|
"'input_list'")
|
|
.failed())
|
|
return failure();
|
|
|
|
if (errorIfTypeOrShapeMismatch(*this, getBodyGraph().front().getArguments(),
|
|
"'body_graph' arguments", getInputList(),
|
|
"'input_list'")
|
|
.failed())
|
|
return failure();
|
|
|
|
if (getBodyGraph().front().mightHaveTerminator()) {
|
|
auto bodyYield =
|
|
dyn_cast<tosa::YieldOp>(getBodyGraph().front().getTerminator());
|
|
if (bodyYield && errorIfTypeOrShapeMismatch(*this, bodyYield.getInputs(),
|
|
"'body_graph' results",
|
|
getInputList(), "'input_list'")
|
|
.failed())
|
|
return failure();
|
|
}
|
|
|
|
// Condition block output must be a single element tensor with a single bool
|
|
// value.
|
|
if (!getCondGraph().front().mightHaveTerminator())
|
|
return success();
|
|
|
|
auto condYield =
|
|
dyn_cast<tosa::YieldOp>(getCondGraph().front().getTerminator());
|
|
if (!condYield)
|
|
return success();
|
|
|
|
if (condYield.getInputs().size() != 1)
|
|
return emitOpError() << "require 'cond_graph' only have one result";
|
|
|
|
auto condOutType = condYield.getInputs()[0].getType();
|
|
if (errorIfShapeNotSizeOne(*this, condOutType).failed())
|
|
return emitOpError() << "'cond_graph' result must be a size 1 tensor, got "
|
|
<< condOutType;
|
|
|
|
if (!getElementTypeOrSelf(condOutType).isInteger(1))
|
|
return emitOpError() << "'cond_graph' result must be a boolean tensor, got "
|
|
<< condOutType;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult ReverseOp::verify() {
|
|
TensorType inputType = getInput1().getType();
|
|
int32_t reverseAxis = getAxis();
|
|
|
|
if (reverseAxis < 0)
|
|
return emitOpError("expected non-negative reverse axis");
|
|
if (inputType.hasRank()) {
|
|
int64_t inputRank = inputType.getRank();
|
|
// We allow for a special case where the input/output shape has rank 0 and
|
|
// axis is also 0.
|
|
if (reverseAxis >= inputRank && (reverseAxis != 0 || inputRank != 0))
|
|
return emitOpError("expect input tensor rank (")
|
|
<< inputRank << ") to be larger than reverse axis (" << reverseAxis
|
|
<< ")";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::SelectOp::verify() {
|
|
// verify input2 and input3 have same element type as output
|
|
if (verifySameElementTypes(*this, /* inType = */ getOnTrue().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed() ||
|
|
verifySameElementTypes(*this, /* inType = */ getOnFalse().getType(),
|
|
/* outType = */ getOutput().getType())
|
|
.failed()) {
|
|
return failure();
|
|
}
|
|
// verify input1 has element type of bool
|
|
auto predicateType = llvm::dyn_cast<ShapedType>(getPred().getType());
|
|
if (!predicateType) {
|
|
return emitOpError("expect shaped tensor for input1, got ")
|
|
<< getInput1().getType();
|
|
}
|
|
auto predicateElementType = predicateType.getElementType();
|
|
if (!predicateElementType.isInteger(1)) {
|
|
return emitOpError("expect element type of bool for input1, got ")
|
|
<< predicateElementType;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::VariableReadOp::verify() {
|
|
if (verifyVariableOpErrorIf(*this, getOutput1().getType(), "'output1'")
|
|
.failed())
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::VariableWriteOp::verify() {
|
|
if (verifyVariableOpErrorIf(*this, getInput1().getType(), "'input1'")
|
|
.failed())
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
// parse and print of WhileOp refer to the implementation of SCF dialect.
|
|
ParseResult WhileOp::parse(OpAsmParser &parser, OperationState &result) {
|
|
SmallVector<OpAsmParser::Argument, 4> regionArgs;
|
|
SmallVector<OpAsmParser::UnresolvedOperand, 4> operands;
|
|
Region *cond = result.addRegion();
|
|
Region *body = result.addRegion();
|
|
|
|
OptionalParseResult listResult =
|
|
parser.parseOptionalAssignmentList(regionArgs, operands);
|
|
if (listResult.has_value() && failed(listResult.value()))
|
|
return failure();
|
|
|
|
FunctionType functionType;
|
|
SMLoc typeLoc = parser.getCurrentLocation();
|
|
if (failed(parser.parseColonType(functionType)))
|
|
return failure();
|
|
|
|
result.addTypes(functionType.getResults());
|
|
|
|
if (functionType.getNumInputs() != operands.size()) {
|
|
return parser.emitError(typeLoc)
|
|
<< "expected as many input types as operands " << "(expected "
|
|
<< operands.size() << " got " << functionType.getNumInputs() << ")";
|
|
}
|
|
|
|
// Resolve input operands.
|
|
if (failed(parser.resolveOperands(operands, functionType.getInputs(),
|
|
parser.getCurrentLocation(),
|
|
result.operands)))
|
|
return failure();
|
|
|
|
// Propagate the types into the region arguments.
|
|
for (size_t i = 0, e = regionArgs.size(); i != e; ++i)
|
|
regionArgs[i].type = functionType.getInput(i);
|
|
|
|
return failure(parser.parseRegion(*cond, regionArgs) ||
|
|
parser.parseKeyword("do") || parser.parseRegion(*body) ||
|
|
parser.parseOptionalAttrDictWithKeyword(result.attributes));
|
|
}
|
|
|
|
void WhileOp::print(OpAsmPrinter &parser) {
|
|
printInitializationList(parser, getCondGraph().front().getArguments(),
|
|
getInputList(), " ");
|
|
parser << " : ";
|
|
parser.printFunctionalType(getInputList().getTypes(),
|
|
getResults().getTypes());
|
|
parser << ' ';
|
|
parser.printRegion(getCondGraph(), /*printEntryBlockArgs=*/false);
|
|
parser << " do ";
|
|
parser.printRegion(getBodyGraph());
|
|
parser.printOptionalAttrDictWithKeyword((*this)->getAttrs());
|
|
}
|
|
|
|
// Create a rank-1 const tensor for zero point of the source tensor.
|
|
std::optional<Value> mlir::tosa::createZeroPointTensor(OpBuilder &builder,
|
|
Location loc,
|
|
Type srcElemType,
|
|
int64_t zp) {
|
|
srcElemType = getStorageElementTypeOrSelf(srcElemType);
|
|
auto zpType = mlir::RankedTensorType::get({1}, srcElemType);
|
|
if (llvm::isa<FloatType>(srcElemType)) {
|
|
auto zpAttr = DenseElementsAttr::get(
|
|
zpType, builder.getFloatAttr(srcElemType, static_cast<double>(zp)));
|
|
return tosa::ConstOp::create(builder, loc, zpType, zpAttr);
|
|
}
|
|
if (llvm::isa<IntegerType>(srcElemType)) {
|
|
auto zpAttr =
|
|
DenseElementsAttr::get(zpType, builder.getIntegerAttr(srcElemType, zp));
|
|
return tosa::ConstOp::create(builder, loc, zpType, zpAttr);
|
|
}
|
|
llvm::errs() << "zero point is not allowed for unsupported data types\n";
|
|
return std::nullopt;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Shape and Shape Operators Helper functions.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
bool mlir::tosa::isa_tosa_shape_type(mlir::Type t) {
|
|
return mlir::isa<tosa::shapeType>(t);
|
|
}
|
|
|
|
LogicalResult
|
|
mlir::tosa::shapeType::verify(function_ref<InFlightDiagnostic()> emitError,
|
|
int rank) {
|
|
if (rank < 0)
|
|
return emitError() << "invalid rank (must be >= 0): " << rank;
|
|
return success();
|
|
}
|
|
|
|
LogicalResult OpTrait::tosa::verifyTosaResolvableShapeOperands(Operation *op) {
|
|
for (auto v : op->getOperands()) {
|
|
if (mlir::isa<::mlir::tosa::shapeType>(v.getType())) {
|
|
Operation *definingOp = v.getDefiningOp();
|
|
if (!definingOp || !definingOp->hasTrait<TosaShapeOperator>()) {
|
|
return op->emitOpError("shape operand is not compile time resolvable");
|
|
}
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
OpTrait::tosa::verifyTosaShapeOperatorWithSameRanks(Operation *op) {
|
|
if (failed(OpTrait::impl::verifyAtLeastNOperands(op, 1)))
|
|
return failure();
|
|
|
|
// delegate function that returns rank of shape type
|
|
auto getRank = [](const Type type) {
|
|
return mlir::cast<mlir::tosa::shapeType>(type).getRank();
|
|
};
|
|
auto operandTypes = op->getOperandTypes();
|
|
auto resultTypes = op->getResultTypes();
|
|
|
|
auto rank = getRank(*op->getOperandTypes().begin());
|
|
for (auto type : operandTypes) {
|
|
if (getRank(type) != rank) {
|
|
return op->emitOpError("operands don't have matching ranks");
|
|
}
|
|
}
|
|
for (auto type : resultTypes) {
|
|
if (getRank(type) != rank) {
|
|
return op->emitOpError("result shape has different rank than operands");
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Shape Operators verify functions.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult tosa::ConstShapeOp::verify() {
|
|
// check one dimensional rank
|
|
auto valuesRank = getValues().getType().getRank();
|
|
if (valuesRank != 1)
|
|
return emitOpError("expect elements in attribute values with rank 1");
|
|
// check that number of elements in values attr equal to rank of result shape
|
|
auto count = getValues().getNumElements();
|
|
auto rank = (cast<tosa::shapeType>(getResult().getType())).getRank();
|
|
if (count != rank && (count != 1 || rank != 0)) {
|
|
return emitOpError("expect number of elements in attribute values (")
|
|
<< count << ") to be equal to the rank (" << rank
|
|
<< ") for the result shape type";
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::DimOp::verify() {
|
|
const tosa::shapeType outShapeType =
|
|
cast<tosa::shapeType>(getResult().getType());
|
|
if (outShapeType.getRank() != 1)
|
|
return emitOpError("expect output shape type to contain one element, got ")
|
|
<< outShapeType;
|
|
|
|
const ShapeAdaptor inputType(getInput1().getType());
|
|
if (inputType.hasRank()) {
|
|
const int64_t inputRank = inputType.getRank();
|
|
const int64_t axis = getAxisAttr().getInt();
|
|
if (axis < 0 || axis >= inputRank)
|
|
return emitOpError("expect axis to be in the range [0, ")
|
|
<< inputRank << "), got " << axis;
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::ConcatShapeOp::verify() {
|
|
const tosa::shapeType outShapeType =
|
|
cast<tosa::shapeType>(getResult().getType());
|
|
const int64_t outputRank = outShapeType.getRank();
|
|
const Operation::operand_range inputList = getInput();
|
|
|
|
if (inputList.size() == 0)
|
|
return emitOpError("requires at least one input shape");
|
|
|
|
if (llvm::any_of(inputList, [](Value v) {
|
|
return cast<tosa::shapeType>(v.getType()).getRank() == 0;
|
|
}))
|
|
return emitOpError("requires all inputs shapes have a rank greater than 0");
|
|
|
|
const int64_t inputsRank =
|
|
llvm::accumulate(inputList, 0, [](int64_t acc, const Value &input) {
|
|
const tosa::shapeType inShapeType =
|
|
cast<tosa::shapeType>(input.getType());
|
|
return acc + inShapeType.getRank();
|
|
});
|
|
if (outputRank != inputsRank)
|
|
return emitOpError("requires output shape rank to be equal to the sum of "
|
|
"the input shape ranks (")
|
|
<< inputsRank << "), got " << outputRank;
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult tosa::SliceShapeOp::verify() {
|
|
std::optional<int32_t> start;
|
|
DenseIntElementsAttr startAttr;
|
|
if (matchPattern(getStart(), m_Constant(&startAttr)))
|
|
start = startAttr.getValues<int32_t>()[0];
|
|
if (start && start.value() < 0)
|
|
return emitOpError("expected non-negative start index, got ")
|
|
<< start.value();
|
|
|
|
std::optional<int32_t> size;
|
|
DenseIntElementsAttr sizeAttr;
|
|
if (matchPattern(getSize(), m_Constant(&sizeAttr)))
|
|
size = sizeAttr.getValues<int32_t>()[0];
|
|
if (size && size.value() <= 0)
|
|
return emitOpError("expected positive size, got ") << size.value();
|
|
|
|
if (!size)
|
|
return success();
|
|
|
|
const tosa::shapeType outShapeType =
|
|
cast<tosa::shapeType>(getResult().getType());
|
|
const int64_t outputRank = outShapeType.getRank();
|
|
if (outputRank != size)
|
|
return emitOpError(
|
|
"expected output type size to be equal to size attribute, got ")
|
|
<< outputRank << " vs " << size.value();
|
|
|
|
if (!start)
|
|
return success();
|
|
|
|
const tosa::shapeType inShapeType =
|
|
cast<tosa::shapeType>(getInput().getType());
|
|
const int64_t inputRank = inShapeType.getRank();
|
|
const int64_t sliceSize = start.value() + size.value();
|
|
if (sliceSize > inputRank)
|
|
return emitOpError("expected start + size to be less than or equal to "
|
|
"input shape rank (")
|
|
<< inputRank << "), got " << sliceSize;
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Attribute Definitions.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#define GET_ATTRDEF_CLASSES
|
|
#include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Type Definitions.
|
|
//===----------------------------------------------------------------------===//
|
|
#define GET_TYPEDEF_CLASSES
|
|
#include "mlir/Dialect/Tosa/IR/TosaOpsTypesBase.cpp.inc"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TOSA Operator Definitions.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
|