Skip to content

[mlir][tosa] Add acc_type to Tosa-v1.0 Conv Ops #121466

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 8, 2025
Merged

[mlir][tosa] Add acc_type to Tosa-v1.0 Conv Ops #121466

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,11 +126,12 @@ def Tosa_ConvOpQuantInfoBuilder : OpBuilder<
(ins "::mlir::Type":$outputType, "::mlir::Value":$input,
"::mlir::Value":$weight, "::mlir::Value":$bias,
"::mlir::DenseI64ArrayAttr":$pad, "::mlir::DenseI64ArrayAttr":$stride,
"::mlir::DenseI64ArrayAttr":$dilation),
"::mlir::DenseI64ArrayAttr":$dilation,
"::mlir::TypeAttr":$acc_type),
[{
buildConvOpWithQuantInfo($_builder, $_state, outputType,
input, weight, bias,
pad, stride, dilation);
pad, stride, dilation, acc_type);
}]>;

// Handles tosa.transpose_conv2d which has an outpad and output shape attribute.
Expand All @@ -139,12 +140,13 @@ def Tosa_TransConvOpQuantInfoBuilder : OpBuilder<
"::mlir::Value":$weight, "mlir::Value":$bias,
"::mlir::DenseI64ArrayAttr":$outpad,
"::mlir::DenseI64ArrayAttr":$stride,
"::mlir::DenseI64ArrayAttr":$outputShape),
"::mlir::DenseI64ArrayAttr":$outputShape,
"::mlir::TypeAttr":$acc_type),
[{
buildTransConvOpWithQuantInfo($_builder, $_state, outputType,
input, weight, bias,
outpad, stride,
outputShape);
outputShape, acc_type);
}]>;

// The tosa.fully_connected op has its own builder as it does not have
Expand Down
7 changes: 6 additions & 1 deletion mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def Tosa_ArgMaxOp : Tosa_InferShapedTypeOp<"argmax"> {
// Accumulator types.
//===----------------------------------------------------------------------===//

def Tosa_AccType : AnyTypeOf<[I<32>, SI<32>, F16, F32]>;
def Tosa_AccType : AnyTypeOf<[I<32>, I<48>, F16, F32]>;

//===----------------------------------------------------------------------===//
// Operator: avg_pool2d
Expand Down Expand Up @@ -106,6 +106,7 @@ def Tosa_Conv2DOp : Tosa_InferShapedTypeOp<"conv2d"> {
Tosa_IntArrayAttr4:$pad,
Tosa_IntArrayAttr2:$stride,
Tosa_IntArrayAttr2:$dilation,
TypeAttrOf<Tosa_AccType>:$acc_type,
OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info,
DefaultValuedOptionalAttr<BoolAttr, "false">:$local_bound
);
Expand Down Expand Up @@ -135,6 +136,7 @@ def Tosa_Conv3DOp : Tosa_InferShapedTypeOp<"conv3d"> {
Tosa_IntArrayAttr6:$pad,
Tosa_IntArrayAttr3:$stride,
Tosa_IntArrayAttr3:$dilation,
TypeAttrOf<Tosa_AccType>:$acc_type,
OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info,
DefaultValuedOptionalAttr<BoolAttr, "false">:$local_bound
);
Expand Down Expand Up @@ -165,6 +167,7 @@ def Tosa_DepthwiseConv2DOp : Tosa_InferShapedTypeOp<"depthwise_conv2d"> {
Tosa_IntArrayAttr4:$pad,
Tosa_IntArrayAttr2:$stride,
Tosa_IntArrayAttr2:$dilation,
TypeAttrOf<Tosa_AccType>:$acc_type,
OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info,
DefaultValuedOptionalAttr<BoolAttr, "false">:$local_bound
);
Expand Down Expand Up @@ -348,6 +351,7 @@ def Tosa_TransposeConv2DOp : Tosa_InferShapedTypeOp<"transpose_conv2d"> {
Tosa_IntArrayAttr4:$out_pad,
Tosa_IntArrayAttr2:$stride,
Tosa_IntArrayAttr4:$out_shape,
TypeAttrOf<Tosa_AccType>:$acc_type,
OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info,
DefaultValuedOptionalAttr<BoolAttr, "false">:$local_bound
);
Expand All @@ -357,6 +361,7 @@ def Tosa_TransposeConv2DOp : Tosa_InferShapedTypeOp<"transpose_conv2d"> {
);

let builders = [Tosa_TransConvOpQuantInfoBuilder];
let hasVerifier = 1;
}

//===----------------------------------------------------------------------===//
Expand Down
78 changes: 71 additions & 7 deletions mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,15 +210,26 @@ template <typename T>
static LogicalResult verifyConvOp(T op) {
// All TOSA conv ops have an input() and weight().
auto inputType = llvm::dyn_cast<RankedTensorType>(op.getInput().getType());
auto weightType = llvm::dyn_cast<RankedTensorType>(op.getWeight().getType());

RankedTensorType weightType;
if constexpr (std::is_same_v<T, tosa::TransposeConv2DOp>)
weightType = llvm::dyn_cast<RankedTensorType>(op.getFilter().getType());
else
weightType = llvm::dyn_cast<RankedTensorType>(op.getWeight().getType());

// Must be ranked tensor types
if (!inputType) {
op.emitOpError("expect a ranked tensor for input, got ") << op.getInput();
return failure();
}
if (!weightType) {
op.emitOpError("expect a ranked tensor for weight, got ") << op.getWeight();
if constexpr (std::is_same_v<T, tosa::TransposeConv2DOp>) {
op.emitOpError("expect a ranked tensor for filter, got ")
<< op.getFilter();
} else {
op.emitOpError("expect a ranked tensor for weight, got ")
<< op.getWeight();
}
return failure();
}

Expand Down Expand Up @@ -271,6 +282,38 @@ LogicalResult tosa::ConstOp::verify() {
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::UniformQuantizedType>(inputEType))
inputEType = quantType.getStorageType();

auto accType = op.getAccType();
if (inputEType.isInteger(8) && !accType.isInteger(32))
return op.emitOpError("accumulator type for i8 tensor is not i32");

if (inputEType.isInteger(16) && !accType.isInteger(48))
return op.emitOpError("accumulator type for i16 tensor is not i48");

if ((inputEType.isFloat8E5M2() || inputEType.isFloat8E4M3()) &&
!accType.isF16())
return op.emitOpError("accumulator type for f8 tensor is not f16");

if (inputEType.isF16() && !(accType.isF16() || accType.isF32()))
return op.emitOpError("accumulator type for f16 tensor is not f16/f32");

if (inputEType.isBF16() && !accType.isF32())
return op.emitOpError("accumulator type for bf16 tensor is not f32");

if (inputEType.isF32() && !accType.isF32())
return op.emitOpError("accumulator type for f32 tensor is not f32");

return success();
}

LogicalResult tosa::ArgMaxOp::verify() {
// Ensure output is of 32-bit integer
const auto resultETy = llvm::cast<ShapedType>(getType()).getElementType();
Expand Down Expand Up @@ -368,12 +411,14 @@ static void buildConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
Type outputType, Value input, Value weight,
Value bias, DenseI64ArrayAttr pad,
DenseI64ArrayAttr stride,
DenseI64ArrayAttr dilation) {
DenseI64ArrayAttr dilation,
TypeAttr accType) {

result.addOperands({input, weight, bias});
result.addAttribute("pad", pad);
result.addAttribute("stride", stride);
result.addAttribute("dilation", dilation);
result.addAttribute("acc_type", accType);

auto quantAttr = buildConvOpQuantizationAttr(builder, input, weight);
if (quantAttr) {
Expand All @@ -390,11 +435,12 @@ static void buildConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
static void buildTransConvOpWithQuantInfo(
OpBuilder &builder, OperationState &result, Type outputType, Value input,
Value weight, Value bias, DenseI64ArrayAttr outpad,
DenseI64ArrayAttr stride, DenseI64ArrayAttr outputShape) {
DenseI64ArrayAttr stride, DenseI64ArrayAttr outputShape, TypeAttr accType) {
result.addOperands({input, weight, bias});
result.addAttribute("out_pad", outpad);
result.addAttribute("stride", stride);
result.addAttribute("out_shape", outputShape);
result.addAttribute("acc_type", accType);
auto quantAttr = ::buildConvOpQuantizationAttr(builder, input, weight);

if (quantAttr) {
Expand Down Expand Up @@ -1595,7 +1641,11 @@ LogicalResult Conv2DOp::inferReturnTypeComponents(
return success();
}

LogicalResult Conv2DOp::verify() { return verifyConvOp(*this); }
LogicalResult Conv2DOp::verify() {
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed())
return failure();
return success();
}

LogicalResult Conv3DOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
Expand Down Expand Up @@ -1667,7 +1717,11 @@ LogicalResult Conv3DOp::inferReturnTypeComponents(
return success();
}

LogicalResult Conv3DOp::verify() { return verifyConvOp(*this); }
LogicalResult Conv3DOp::verify() {
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed())
return failure();
return success();
}

LogicalResult AvgPool2dOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
Expand Down Expand Up @@ -1762,7 +1816,11 @@ LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
return success();
}

LogicalResult DepthwiseConv2DOp::verify() { return verifyConvOp(*this); }
LogicalResult DepthwiseConv2DOp::verify() {
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed())
return failure();
return success();
}

LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
Expand Down Expand Up @@ -1828,6 +1886,12 @@ LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
return success();
}

LogicalResult TransposeConv2DOp::verify() {
if (verifyConvOp(*this).failed() || verifyConvOpModes(*this).failed())
return failure();
return success();
}

LogicalResult IfOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
IfOp::Adaptor adaptor,
Expand Down
11 changes: 7 additions & 4 deletions mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeTransposeConv.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,13 +75,15 @@ class TransposeConvNonStridedConverter
loc, resultTy, input, reverse2, bias,
rewriter.getDenseI64ArrayAttr(convPad),
rewriter.getDenseI64ArrayAttr(stride),
rewriter.getDenseI64ArrayAttr({1, 1}), *op.getQuantizationInfo());
rewriter.getDenseI64ArrayAttr({1, 1}),
/* acc_type = */ op.getAccType(), *op.getQuantizationInfo());
} else {
conv2d = rewriter.create<tosa::Conv2DOp>(
loc, resultTy, input, reverse2, bias,
rewriter.getDenseI64ArrayAttr(convPad),
rewriter.getDenseI64ArrayAttr(stride),
rewriter.getDenseI64ArrayAttr({1, 1}));
rewriter.getDenseI64ArrayAttr({1, 1}),
/* acc_type = */ op.getAccTypeAttr());
}

rewriter.replaceOp(op, conv2d);
Expand Down Expand Up @@ -238,15 +240,16 @@ class TransposeConvStridedConverter
/*pad=*/rewriter.getDenseI64ArrayAttr({0, 0, 0, 0}),
/*stride=*/rewriter.getDenseI64ArrayAttr({1, 1}),
/*dilation=*/rewriter.getDenseI64ArrayAttr({1, 1}),
*op.getQuantizationInfo())
/* acc_type = */ op.getAccType(), *op.getQuantizationInfo())
.getResult();
} else {
conv2d = CreateOpAndInferShape<tosa::Conv2DOp>(
rewriter, loc, UnrankedTensorType::get(resultETy), input,
weight, zeroBias,
/*pad=*/rewriter.getDenseI64ArrayAttr({0, 0, 0, 0}),
/*stride=*/rewriter.getDenseI64ArrayAttr({1, 1}),
/*dilation=*/rewriter.getDenseI64ArrayAttr({1, 1}))
/*dilation=*/rewriter.getDenseI64ArrayAttr({1, 1}),
/* acc_type = */ op.getAccTypeAttr())
.getResult();
}

Expand Down
Loading
Loading