[Mlir-commits] [mlir] [mlir][tosa] Convert tosa.transpose_conv2d to linalg.generic directly (PR #79824)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jan 29 05:25:36 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hsiangkai Wang (Hsiangkai)
<details>
<summary>Changes</summary>
Currently, we use reverse, pad, reshape, and conv2d operators, etc, to emulate transpose_conv2d. This patch adds a pattern to convert tosa.transpose_conv2d to linalg.generic directly.
---
Full diff: https://github.com/llvm/llvm-project/pull/79824.diff
3 Files Affected:
- (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp (+77-1)
- (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp (+1)
- (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir (+32)
``````````diff
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 8dc2d27bd545ff8..b9fd03ff67c4e23 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -1010,6 +1010,81 @@ class TransposeConverter : public OpRewritePattern<tosa::TransposeOp> {
return success();
}
};
+
+class TransposeConv2DConverter
+ : public OpConversionPattern<tosa::TransposeConv2DOp> {
+public:
+ using OpConversionPattern<tosa::TransposeConv2DOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(tosa::TransposeConv2DOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const final {
+ Location loc = op->getLoc();
+ Value input = op->getOperand(0);
+ Value weight = op->getOperand(1);
+ Value bias = op->getOperand(2);
+
+ ShapedType inputTy = cast<ShapedType>(input.getType());
+ ShapedType weightTy = cast<ShapedType>(weight.getType());
+ ShapedType biasTy = cast<ShapedType>(bias.getType());
+ ShapedType resultTy = cast<ShapedType>(op->getResult(0).getType());
+
+ if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
+ return rewriter.notifyMatchFailure(
+ op,
+ "tosa.transpose_conv2d requires static shapes for weight and bias");
+
+ Type inputETy = inputTy.getElementType();
+ Type resultETy = resultTy.getElementType();
+
+ if (inputETy.isUnsignedInteger())
+ return rewriter.notifyMatchFailure(
+ op, "tosa.transpose_conv2d does not support unsigned integer input");
+
+ // Broadcast the bias as the starting values for accumulation.
+ auto emptyTensor =
+ rewriter.create<tensor::EmptyOp>(loc, resultTy.getShape(), resultETy);
+
+ Value broadcastBias =
+ linalgBroadcastAndMaybeExtSI(rewriter, loc, bias, emptyTensor);
+
+ auto *context = op->getContext();
+ AffineExpr n, ih, iw, oc, ic, kh, kw;
+ bindDims(context, n, ih, iw, oc, ic, kh, kw);
+
+ constexpr unsigned numDims = 7;
+ auto lhsMap = AffineMap::get(numDims, 0, {n, ih, iw, ic}, context);
+ auto rhsMap = AffineMap::get(numDims, 0, {oc, kh, kw, ic}, context);
+ /* outPad: top, bottom, left, right */
+ ArrayRef<int64_t> outPad = op.getOutPadAttr();
+ ArrayRef<int64_t> stride = op.getStrideAttr();
+ auto resultMap = AffineMap::get(numDims, 0,
+ {n, ih * stride[0] + outPad[0] + kh,
+ iw * stride[1] + outPad[2] + kw, oc},
+ context);
+
+ auto transposeConv2D =
+ rewriter
+ .create<linalg::GenericOp>(
+ loc, resultTy, ValueRange({input, weight}), broadcastBias,
+ ArrayRef<AffineMap>{lhsMap, rhsMap, resultMap},
+ tosa::getNParallelLoopsAttrs(numDims),
+ [&](OpBuilder &nestedBuilder, Location nestedLoc,
+ ValueRange args) {
+ auto mul =
+ nestedBuilder.create<arith::MulFOp>(loc, args[0], args[1])
+ .getResult();
+ auto acc =
+ nestedBuilder.create<arith::AddFOp>(loc, mul, args[2])
+ .getResult();
+ nestedBuilder.create<linalg::YieldOp>(loc, acc);
+ })
+ .getResult(0);
+
+ rewriter.replaceOp(op, transposeConv2D);
+
+ return success();
+ }
+};
} // namespace
void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
@@ -1031,7 +1106,8 @@ void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
MaxPool2dConverter,
AvgPool2dConverter,
FullyConnectedConverter,
- TransposeConverter
+ TransposeConverter,
+ TransposeConv2DConverter
>(patterns->getContext());
// clang-format on
}
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
index 096969391e51b9d..422d1e6189a21e6 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -61,6 +61,7 @@ struct TosaToLinalgNamed
target.addIllegalOp<tosa::MatMulOp>();
target.addIllegalOp<tosa::FullyConnectedOp>();
target.addIllegalOp<tosa::TransposeOp>();
+ target.addIllegalOp<tosa::TransposeConv2DOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index 6616ea7cf699fa5..b39c488e3bbbeb0 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -781,3 +781,35 @@ func.func @test_transpose_dyn_multiple_3d(%arg0: tensor<?x?x?xf32>) {
%1 = "tosa.transpose"(%arg0, %0) : (tensor<?x?x?xf32>, tensor<3xi32>) -> tensor<?x?x?xf32>
return
}
+
+// -----
+
+// CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d4)>
+// CHECK: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d5, d6, d4)>
+// CHECK: #[[$MAP4:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5, d2 + d6, d3)>
+
+// CHECK-LABEL: @transpose_conv2d
+func.func @transpose_conv2d(%arg0: tensor<1x1x2x1xf32>, %arg1: tensor<1x1x2x1xf32>, %arg2: tensor<1xf32>) -> tensor<1x1x3x1xf32> {
+ // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<1x1x3x1xf32>
+ // CHECK: %[[BROADCAST:.+]] = linalg.generic
+ // CHECK-SAME: {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+ // CHECK-SAME: ins(%arg2 : tensor<1xf32>)
+ // CHECK-SAME: outs(%[[INIT]] : tensor<1x1x3x1xf32>) {
+ // CHECK: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32):
+ // CHECK: linalg.yield %[[IN]] : f32
+ // CHECK: } -> tensor<1x1x3x1xf32>
+ // CHECK: %[[RESULT:.+]] = linalg.generic
+ // CHECK-SAME: {indexing_maps = [#[[$MAP2]], #[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<1x1x2x1xf32>, tensor<1x1x2x1xf32>)
+ // CHECK-SAME: outs(%[[BROADCAST]] : tensor<1x1x3x1xf32>) {
+ // CHECK: ^bb0(%[[IN:.+]]: f32, %[[IN_0:.+]]: f32, %[[OUT:.+]]: f32):
+ // CHECK: %[[S3:.+]] = arith.mulf %[[IN]], %[[IN_0]] : f32
+ // CHECK: %[[S4:.+]] = arith.addf %[[S3]], %[[OUT]] : f32
+ // CHECK: linalg.yield %[[S4]] : f32
+ // CHECK: } -> tensor<1x1x3x1xf32>
+ // CHECK: return %[[RESULT]] : tensor<1x1x3x1xf32>
+ %0 = tosa.transpose_conv2d %arg0, %arg1, %arg2 {out_pad = array<i64: 0, 0, 0, 0>, out_shape = array<i64: 1, 1, 3, 1>, stride = array<i64: 1, 1>} : (tensor<1x1x2x1xf32>, tensor<1x1x2x1xf32>, tensor<1xf32>) -> tensor<1x1x3x1xf32>
+ return %0 : tensor<1x1x3x1xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/79824
More information about the Mlir-commits
mailing list