#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Mesh/Interfaces/ShardingInterface.h"
#include "mlir/Dialect/Quant/IR/Quant.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/IR/BuiltinTypes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Transforms/InliningUtils.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/TypeSwitch.h"
#include <numeric>
#include "mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc"
#include "mlir/Dialect/Tosa/IR/TosaInterfaces.cpp.inc"
#include "mlir/Dialect/Tosa/IR/TosaDialectBytecode.cpp.inc"
#include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
#include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"

Macros
#define	GET_OP_LIST

#define	GET_ATTRDEF_LIST

#define	COMPATIBLE_RETURN_TYPES(OP)

#define	REDUCE_SHAPE_INFER(OP)

#define	NARY_SHAPE_INFER(OP)

#define	GET_ATTRDEF_CLASSES

#define	GET_OP_CLASSES

Functions
template<typename T >
static LogicalResult	verifyConvOp (T op)

static void	buildConvOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input, Value weight, Value bias, DenseI64ArrayAttr pad, DenseI64ArrayAttr stride, DenseI64ArrayAttr dilation)
	This builder is called on all convolution operators except TransposeConv, which has specialized output shape semantics. More...

static void	buildTransConvOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input, Value weight, Value bias, DenseI64ArrayAttr outpad, DenseI64ArrayAttr stride, DenseI64ArrayAttr outputShape)
	Handles tosa.transpose_conv2d which has outpad and output shape attributes. More...

static void	buildFCOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input, Value weight, Value bias)
	The tosa.fully_connected op has its own builder as it does not have strides/dilation/padding. More...

static void	buildMatMulOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value a, Value b)
	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. More...

static void	buildAvgPool2dOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input, DenseArrayAttr kernel, DenseArrayAttr stride, DenseArrayAttr pad, TypeAttr accType)
	Both the tosa.avg_pool2d and unary ops use the same UnaruOpQuantizationAttr but avg_pool operator has its own builder as it has additional parameters not part of the unary ops. More...

static void	buildUnaryOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input)
	This builder is called on single-parameter unary operators that have scale relationship between their input and output, expressed by the UnaryOpQuantizationAttr. More...

static void	buildPadOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input, Value paddings)
	This builder is called on TOSA pad operator that needs to create its own OptionalAttr quantization_attr parameter to scale the padding values correctly. More...

static void	buildExplicitValuePadOpWithQuantInfo (OpBuilder &builder, OperationState &result, Type outputType, Value input, Value paddings, Value padConst)
	This builder is called on TOSA pad operator when an explicit pad_const value is passed in. More...

static LogicalResult	resolveBroadcastShape (const ValueShapeRange &operands, SmallVector< int64_t > &outShape)

static SmallVector< int64_t >	convertToMlirShape (ArrayRef< int64_t > shape)

static LogicalResult	ReduceInferReturnTypes (ShapeAdaptor operandShape, Type inputType, IntegerAttr axis, SmallVectorImpl< ShapedTypeComponents > &inferredReturnShapes)

template<typename T >
static LogicalResult	verifyReduceOp (T op)

static LogicalResult	NAryInferReturnTypes (const ValueShapeRange &operands, SmallVectorImpl< ShapedTypeComponents > &inferredReturnShapes)

static LogicalResult	poolingInferReturnTypes (ShapeAdaptor inputShape, ArrayRef< int64_t > kernel, ArrayRef< int64_t > stride, ArrayRef< int64_t > pad, SmallVectorImpl< ShapedTypeComponents > &inferredReturnShapes)

static void	printInitializationList (OpAsmPrinter &parser, Block::BlockArgListType blocksArgs, ValueRange initializers, StringRef prefix="")

Macro Definition Documentation

◆ COMPATIBLE_RETURN_TYPES

#define COMPATIBLE_RETURN_TYPES ( OP )

Value:

  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]));                       \
  }

Definition at line 1351 of file TosaOps.cpp.

◆ GET_ATTRDEF_CLASSES

#define GET_ATTRDEF_CLASSES

Definition at line 2085 of file TosaOps.cpp.

◆ GET_ATTRDEF_LIST

#define GET_ATTRDEF_LIST

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 2092 of file TosaOps.cpp.

◆ GET_OP_LIST

#define GET_OP_LIST

◆ NARY_SHAPE_INFER

#define NARY_SHAPE_INFER ( OP )

Value:

  LogicalResult OP::inferReturnTypeComponents(                                 \
      MLIRContext *context, ::std::optional<Location> location,                \
      ValueShapeRange operands, DictionaryAttr attributes,                     \
      OpaqueProperties properties, RegionRange regions,                        \
      SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
    return NAryInferReturnTypes(operands, inferredReturnShapes);               \
  }

Definition at line 1453 of file TosaOps.cpp.

◆ REDUCE_SHAPE_INFER

#define REDUCE_SHAPE_INFER ( OP )

Value:

  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)

Definition at line 1360 of file TosaOps.cpp.

Function Documentation

◆ buildAvgPool2dOpWithQuantInfo()

static void buildAvgPool2dOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input,
		DenseArrayAttr	kernel,
		DenseArrayAttr	stride,
		DenseArrayAttr	pad,
		TypeAttr	accType
	)

static

Both the tosa.avg_pool2d and unary ops use the same UnaruOpQuantizationAttr but avg_pool operator has its own builder as it has additional parameters not part of the unary ops.

Definition at line 467 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::tosa::buildUnaryOpQuantizationAttr(), and mlir::OperationState::types.

◆ buildConvOpWithQuantInfo()

static void buildConvOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input,
		Value	weight,
		Value	bias,
		DenseI64ArrayAttr	pad,
		DenseI64ArrayAttr	stride,
		DenseI64ArrayAttr	dilation
	)

static

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.

Definition at line 367 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::OperationState::addTypes(), mlir::tosa::buildConvOpQuantizationAttr(), and mlir::tosa::buildConvOpResultTypeInfo().

◆ buildExplicitValuePadOpWithQuantInfo()

static void buildExplicitValuePadOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input,
		Value	paddings,
		Value	padConst
	)

static

This builder is called on TOSA pad operator when an explicit pad_const value is passed in.

It also optionally constructs quantization_attr.

Definition at line 510 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::tosa::buildPadOpQuantizationAttr(), and mlir::OperationState::types.

◆ buildFCOpWithQuantInfo()

static void buildFCOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input,
		Value	weight,
		Value	bias
	)

static

The tosa.fully_connected op has its own builder as it does not have strides/dilation/padding.

Definition at line 411 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::OperationState::addTypes(), mlir::tosa::buildConvOpQuantizationAttr(), and mlir::tosa::buildConvOpResultTypeInfo().

◆ buildMatMulOpWithQuantInfo()

static void buildMatMulOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	a,
		Value	b
	)

static

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.

Definition at line 430 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::OperationState::addTypes(), mlir::tosa::buildMatMulOpQuantizationAttr(), mlir::Builder::getI32Type(), mlir::Builder::getIntegerType(), and mlir::Value::getType().

◆ buildPadOpWithQuantInfo()

static void buildPadOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input,
		Value	paddings
	)

static

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.

Definition at line 498 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::tosa::buildPadOpQuantizationAttr(), and mlir::OperationState::types.

◆ buildTransConvOpWithQuantInfo()

static void buildTransConvOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input,
		Value	weight,
		Value	bias,
		DenseI64ArrayAttr	outpad,
		DenseI64ArrayAttr	stride,
		DenseI64ArrayAttr	outputShape
	)

static

Handles tosa.transpose_conv2d which has outpad and output shape attributes.

Definition at line 390 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::OperationState::addTypes(), mlir::tosa::buildConvOpQuantizationAttr(), and mlir::tosa::buildConvOpResultTypeInfo().

◆ buildUnaryOpWithQuantInfo()

static void buildUnaryOpWithQuantInfo	(	OpBuilder &	builder,
		OperationState &	result,
		Type	outputType,
		Value	input
	)

static

This builder is called on single-parameter unary operators that have scale relationship between their input and output, expressed by the UnaryOpQuantizationAttr.

Definition at line 485 of file TosaOps.cpp.

References mlir::OperationState::addAttribute(), mlir::OperationState::addOperands(), mlir::tosa::buildUnaryOpQuantizationAttr(), and mlir::OperationState::types.

◆ convertToMlirShape()

static SmallVector<int64_t> convertToMlirShape ( ArrayRef< int64_t > shape )

static

Definition at line 837 of file TosaOps.cpp.

◆ NAryInferReturnTypes()

static LogicalResult NAryInferReturnTypes	(	const ValueShapeRange &	operands,
		SmallVectorImpl< ShapedTypeComponents > &	inferredReturnShapes
	)

static

Definition at line 1441 of file TosaOps.cpp.

References resolveBroadcastShape().

◆ poolingInferReturnTypes()

static LogicalResult poolingInferReturnTypes	(	ShapeAdaptor	inputShape,
		ArrayRef< int64_t >	kernel,
		ArrayRef< int64_t >	stride,
		ArrayRef< int64_t >	pad,
		SmallVectorImpl< ShapedTypeComponents > &	inferredReturnShapes
	)

static

Definition at line 1504 of file TosaOps.cpp.

References mlir::ShapeAdaptor::getDimSize().

◆ printInitializationList()

static void printInitializationList	(	OpAsmPrinter &	parser,
		Block::BlockArgListType	blocksArgs,
		ValueRange	initializers,
		StringRef	prefix = `""`
	)

static

Definition at line 2053 of file TosaOps.cpp.

◆ ReduceInferReturnTypes()

static LogicalResult ReduceInferReturnTypes	(	ShapeAdaptor	operandShape,
		Type	inputType,
		IntegerAttr	axis,
		SmallVectorImpl< ShapedTypeComponents > &	inferredReturnShapes
	)

static

Definition at line 1335 of file TosaOps.cpp.

References mlir::ShapeAdaptor::getDims(), mlir::ShapeAdaptor::getRank(), and mlir::ShapeAdaptor::hasRank().

◆ resolveBroadcastShape()

static LogicalResult resolveBroadcastShape	(	const ValueShapeRange &	operands,
		SmallVector< int64_t > &	outShape
	)

static

Definition at line 526 of file TosaOps.cpp.

References mlir::ValueShapeRange::getShape().

Referenced by NAryInferReturnTypes().

◆ verifyConvOp()

template<typename T >

static LogicalResult verifyConvOp ( T op )

static

Definition at line 210 of file TosaOps.cpp.

◆ verifyReduceOp()

template<typename T >

static LogicalResult verifyReduceOp ( T op )

static

Definition at line 1385 of file TosaOps.cpp.

References mlir::TensorType::getShape(), and mlir::TensorType::hasRank().

Macros

Functions

Macro Definition Documentation

◆ COMPATIBLE_RETURN_TYPES

◆ GET_ATTRDEF_CLASSES

◆ GET_ATTRDEF_LIST

◆ GET_OP_CLASSES

◆ GET_OP_LIST

◆ NARY_SHAPE_INFER

◆ REDUCE_SHAPE_INFER

Function Documentation

◆ buildAvgPool2dOpWithQuantInfo()

◆ buildConvOpWithQuantInfo()

◆ buildExplicitValuePadOpWithQuantInfo()

◆ buildFCOpWithQuantInfo()

◆ buildMatMulOpWithQuantInfo()

◆ buildPadOpWithQuantInfo()

◆ buildTransConvOpWithQuantInfo()

◆ buildUnaryOpWithQuantInfo()

◆ convertToMlirShape()

◆ NAryInferReturnTypes()

◆ poolingInferReturnTypes()

◆ printInitializationList()

◆ ReduceInferReturnTypes()

◆ resolveBroadcastShape()

◆ verifyConvOp()

◆ verifyReduceOp()