TosaToLinalg: Support unsigned tosa.max_pool2d (FXML-4556)

mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h

@@ -52,7 +52,10 @@ void populateTosaToLinalgConversionPatterns(TypeConverter &converter,
 
 /// Populates conversion passes from TOSA dialect to Linalg named operations.
 void populateTosaToLinalgNamedConversionPatterns(
-    RewritePatternSet *patterns, const TosaToLinalgNamedOptions &options);
+    TypeConverter &converter, RewritePatternSet *patterns,
+    const TosaToLinalgNamedOptions &options);
+
+void populateTosaToLinalgTypeConversion(TypeConverter &converter);
 
 } // namespace tosa
 } // namespace mlir

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

@@ -2756,3 +2756,37 @@ void mlir::tosa::populateTosaToLinalgConversionPatterns(
 
   // clang-format on
 }
+
+void mlir::tosa::populateTosaToLinalgTypeConversion(TypeConverter &converter) {
+  converter.addConversion([&](Type type) -> std::optional<Type> {
+    if (type.isUnsignedInteger()) {
+      return IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth(),
+                              IntegerType::SignednessSemantics::Signless);
+    }
+    return type;
+  });
+  converter.addConversion([&](TensorType type) -> std::optional<Type> {
+    auto converted = converter.convertType(type.getElementType());
+    if (!converted)
+      return {};
+    return type.clone(converted);
+  });
+  converter.addSourceMaterialization([&](OpBuilder &builder, Type resultType,
+                                         ValueRange inputs,
+                                         Location loc) -> std::optional<Value> {
+    if (inputs.size() != 1)
+      return std::nullopt;
+
+    return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
+        .getResult(0);
+  });
+  converter.addTargetMaterialization([&](OpBuilder &builder, Type resultType,
+                                         ValueRange inputs,
+                                         Location loc) -> std::optional<Value> {
+    if (inputs.size() != 1)
+      return std::nullopt;
+
+    return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
+        .getResult(0);
+  });
+}

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

@@ -760,17 +760,23 @@ class FullyConnectedConverter
   }
 };
 
-class MaxPool2dConverter : public OpRewritePattern<tosa::MaxPool2dOp> {
+class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
 public:
-  using OpRewritePattern<tosa::MaxPool2dOp>::OpRewritePattern;
+  using OpConversionPattern::OpConversionPattern;
 
-  LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
-                                PatternRewriter &rewriter) const final {
+  LogicalResult
+  matchAndRewrite(tosa::MaxPool2dOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const final {
     Location loc = op.getLoc();
-    Value input = op.getInput();
+    Value input = adaptor.getInput();
     ShapedType inputTy = cast<ShapedType>(input.getType());
 
-    ShapedType resultTy = cast<ShapedType>(op.getType());
+    bool isUnsigned =
+        cast<ShapedType>(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();
 
     auto dynamicDimsOr =
@@ -786,7 +792,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()));
@@ -823,9 +832,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();
   }
 };
@@ -1091,7 +1106,8 @@ class TransposeConverter : public OpRewritePattern<tosa::TransposeOp> {
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
-    RewritePatternSet *patterns, const TosaToLinalgNamedOptions &options) {
+    TypeConverter &converter, RewritePatternSet *patterns,
+    const TosaToLinalgNamedOptions &options) {
   if (options.preferConv2DKernelLayoutHWCF) {
     patterns->add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcHwcfOp,
                                 linalg::Conv2DNhwcHwcfQOp>>(
@@ -1105,11 +1121,13 @@ void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
       // clang-format off
       ConvConverter<tosa::Conv3DOp, linalg::Conv3DNdhwcDhwcfOp, linalg::Conv3DNdhwcDhwcfQOp>,
       DepthwiseConvConverter,
-      MaxPool2dConverter,
       AvgPool2dConverter,
       FullyConnectedConverter,
       TransposeConverter
   >(patterns->getContext());
+  patterns->add<
+      MaxPool2dConverter
+    >(converter, patterns->getContext());
   patterns->add<
       MatMulConverter>(patterns->getContext(), options.useMatmulForSingleBatch);
   // clang-format on

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp

@@ -47,6 +47,9 @@ struct TosaToLinalgNamed
   }
 
   void runOnOperation() override {
+    TypeConverter converter;
+    mlir::tosa::populateTosaToLinalgTypeConversion(converter);
+
     RewritePatternSet patterns(&getContext());
     ConversionTarget target(getContext());
     target.addLegalDialect<linalg::LinalgDialect, tosa::TosaDialect,
@@ -68,7 +71,8 @@ struct TosaToLinalgNamed
     TosaToLinalgNamedOptions options;
     options.preferConv2DKernelLayoutHWCF = preferConv2DKernelLayoutHWCF;
     options.useMatmulForSingleBatch = useMatmulForSingleBatch;
-    tosa::populateTosaToLinalgNamedConversionPatterns(&patterns, options);
+    tosa::populateTosaToLinalgNamedConversionPatterns(converter, &patterns,
+                                                      options);
     if (failed(applyFullConversion(func, target, std::move(patterns))))
       signalPassFailure();
   }

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp

@@ -46,31 +46,7 @@ struct TosaToLinalg : public impl::TosaToLinalgBase<TosaToLinalg> {
 
   void runOnOperation() override {
     TypeConverter converter;
-    converter.addConversion([&](Type type) -> std::optional<Type> {
-      if (type.isUnsignedInteger()) {
-        return IntegerType::get(&getContext(), type.getIntOrFloatBitWidth(),
-                                IntegerType::SignednessSemantics::Signless);
-      }
-      return type;
-    });
-    converter.addConversion([&](TensorType type) -> std::optional<Type> {
-      auto converted = converter.convertType(type.getElementType());
-      if (!converted)
-        return {};
-      return type.clone(converted);
-    });
-    converter.addConversion(
-        [&converter](FunctionType ty) -> std::optional<Type> {
-          SmallVector<Type> inputs;
-          if (failed(converter.convertTypes(ty.getInputs(), inputs)))
-            return std::nullopt;
-
-          SmallVector<Type> results;
-          if (failed(converter.convertTypes(ty.getResults(), results)))
-            return std::nullopt;
-
-          return FunctionType::get(ty.getContext(), inputs, results);
-        });
+    mlir::tosa::populateTosaToLinalgTypeConversion(converter);
 
     RewritePatternSet patterns(&getContext());
     ConversionTarget target(getContext());

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir

@@ -199,6 +199,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