Skip to content

[MLIR] TosaToLinalgNamed: Lower unsigned tosa.max_pool2d #123290

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 20, 2025
Merged

[MLIR] TosaToLinalgNamed: Lower unsigned tosa.max_pool2d #123290

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
3 changes: 2 additions & 1 deletion mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,8 @@ void populateTosaToLinalgConversionPatterns(const TypeConverter &converter,

/// Populates conversion passes from TOSA dialect to Linalg named operations.
void populateTosaToLinalgNamedConversionPatterns(
RewritePatternSet *patterns, const TosaToLinalgNamedOptions &options);
const TypeConverter &converter, RewritePatternSet *patterns,
const TosaToLinalgNamedOptions &options);

} // namespace tosa
} // namespace mlir
Expand Down
52 changes: 36 additions & 16 deletions mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -695,17 +695,18 @@ class FullyConnectedConverter
}
};

class MaxPool2dConverter : public OpRewritePattern<tosa::MaxPool2dOp> {
class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
public:
using OpRewritePattern<tosa::MaxPool2dOp>::OpRewritePattern;
using OpConversionPattern::OpConversionPattern;

// Compute the dynamic output sizes of the maxpool operation.
static SmallVector<Value>
computeDynamicOutputSizes(tosa::MaxPool2dOp op, PatternRewriter &rewriter) {
computeDynamicOutputSizes(tosa::MaxPool2dOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) {
TensorType resultTy = op.getType();
Location loc = op.getLoc();

TypedValue<TensorType> input = op.getInput();
Value input = adaptor.getInput();
ArrayRef<int64_t> kernel = op.getKernel();
ArrayRef<int64_t> pad = op.getPad();
ArrayRef<int64_t> stride = op.getStride();
Expand Down Expand Up @@ -744,16 +745,22 @@ class MaxPool2dConverter : public OpRewritePattern<tosa::MaxPool2dOp> {
return dynamicDims;
}

LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
PatternRewriter &rewriter) const final {
LogicalResult
matchAndRewrite(tosa::MaxPool2dOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const final {
Location loc = op.getLoc();
TypedValue<TensorType> input = op.getInput();
ShapedType inputTy = input.getType();
Value input = adaptor.getInput();
ShapedType inputTy = cast<ShapedType>(input.getType());

ShapedType resultTy = op.getType();
bool isUnsigned = op.getType().getElementType().isUnsignedInteger();
ShapedType resultTy =
cast<ShapedType>(getTypeConverter()->convertType(op.getType()));
if (!resultTy)
return rewriter.notifyMatchFailure(op, "failed to convert type");
Type resultETy = inputTy.getElementType();

SmallVector<Value> dynamicDims = computeDynamicOutputSizes(op, rewriter);
SmallVector<Value> dynamicDims =
computeDynamicOutputSizes(op, adaptor, rewriter);

// Determine what the initial value needs to be for the max pool op.
TypedAttr initialAttr;
Expand All @@ -762,7 +769,10 @@ class MaxPool2dConverter : public OpRewritePattern<tosa::MaxPool2dOp> {
resultETy, APFloat::getLargest(
cast<FloatType>(resultETy).getFloatSemantics(), true));

if (isa<IntegerType>(resultETy))
else if (isUnsigned)
initialAttr = rewriter.getIntegerAttr(
resultETy, APInt::getZero(resultETy.getIntOrFloatBitWidth()));
else if (isa<IntegerType>(resultETy))
initialAttr = rewriter.getIntegerAttr(
resultETy,
APInt::getSignedMinValue(resultETy.getIntOrFloatBitWidth()));
Expand Down Expand Up @@ -798,9 +808,15 @@ class MaxPool2dConverter : public OpRewritePattern<tosa::MaxPool2dOp> {
Value fakeWindowDims =
rewriter.create<tensor::EmptyOp>(loc, kernel, resultETy);

rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxOp>(
op, ArrayRef<Type>{resultTy}, ValueRange{paddedInput, fakeWindowDims},
filledEmptyTensor, strideAttr, dilationAttr);
if (isUnsigned) {
rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxUnsignedOp>(
op, ArrayRef<Type>{resultTy}, ValueRange{paddedInput, fakeWindowDims},
filledEmptyTensor, strideAttr, dilationAttr);
} else {
rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxOp>(
op, ArrayRef<Type>{resultTy}, ValueRange{paddedInput, fakeWindowDims},
filledEmptyTensor, strideAttr, dilationAttr);
}
return success();
}
};
Expand Down Expand Up @@ -1070,7 +1086,8 @@ class TransposeConverter : public OpRewritePattern<tosa::TransposeOp> {
} // namespace

void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
RewritePatternSet *patterns, const TosaToLinalgNamedOptions &options) {
const TypeConverter &converter, RewritePatternSet *patterns,
const TosaToLinalgNamedOptions &options) {
if (options.preferConv2DKernelLayoutHWCF) {
patterns->add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcHwcfOp,
linalg::Conv2DNhwcHwcfQOp>>(
Expand All @@ -1085,10 +1102,13 @@ void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
ConvConverter<tosa::Conv3DOp, linalg::Conv3DNdhwcDhwcfOp, linalg::Conv3DNdhwcDhwcfQOp>,
DepthwiseConvConverter,
MatMulConverter,
MaxPool2dConverter,
AvgPool2dConverter,
FullyConnectedConverter,
TransposeConverter
>(patterns->getContext());

patterns->add<
MaxPool2dConverter
>(converter, patterns->getContext());
// clang-format on
}
6 changes: 5 additions & 1 deletion mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ struct TosaToLinalgNamed
}

void runOnOperation() override {
TypeConverter converter;
tosa::populateTosaTypeConversion(converter);

RewritePatternSet patterns(&getContext());
ConversionTarget target(getContext());
target.addLegalDialect<linalg::LinalgDialect, tosa::TosaDialect,
Expand All @@ -67,7 +70,8 @@ struct TosaToLinalgNamed
FunctionOpInterface func = getOperation();
TosaToLinalgNamedOptions options;
options.preferConv2DKernelLayoutHWCF = preferConv2DKernelLayoutHWCF;
tosa::populateTosaToLinalgNamedConversionPatterns(&patterns, options);
tosa::populateTosaToLinalgNamedConversionPatterns(converter, &patterns,
options);
if (failed(applyFullConversion(func, target, std::move(patterns))))
signalPassFailure();
}
Expand Down
13 changes: 13 additions & 0 deletions mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,19 @@ func.func @max_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> () {
return
}

// CHECK-LABEL: @max_pool_ui8
func.func @max_pool_ui8(%arg0: tensor<1x6x34x62xui8>) -> tensor<1x4x32x62xui8> {
// CHECK: builtin.unrealized_conversion_cast {{.*}} : tensor<1x6x34x62xui8> to tensor<1x6x34x62xi8>
// CHECK: arith.constant 0
// CHECK: linalg.pooling_nhwc_max_unsigned
// CHECK-SAME: ins({{.*}} : tensor<1x6x34x62xi8>, tensor<3x3xi8>)
// CHECK-SAME: outs({{.*}} : tensor<1x4x32x62xi8>)
// CHECK-SAME: -> tensor<1x4x32x62xi8>
// CHECK: builtin.unrealized_conversion_cast {{.*}} : tensor<1x4x32x62xi8> to tensor<1x4x32x62xui8>
%0 = tosa.max_pool2d %arg0 {pad = array<i64: 0, 0, 0, 0>, kernel = array<i64: 3, 3>, stride = array<i64: 1, 1>} : (tensor<1x6x34x62xui8>) -> tensor<1x4x32x62xui8>
return %0 : tensor<1x4x32x62xui8>
}

// CHECK-LABEL: @max_pool_i16
func.func @max_pool_i16(%arg0: tensor<1x6x34x62xi16>) -> () {
// CHECK: arith.constant -32768
Expand Down
Loading