[Mlir-commits] [mlir] [Linalg] Update transform + vectorization patterns to work with generic convolution ops as well (PR #174196)
Abhishek Varma
llvmlistbot at llvm.org
Sun Jan 4 23:22:49 PST 2026
https://github.com/Abhishek-Varma updated https://github.com/llvm/llvm-project/pull/174196
>From 8ca77c9228105aa10039464c17e11403882660e0 Mon Sep 17 00:00:00 2001
From: Abhishek Varma <abhvarma at amd.com>
Date: Thu, 1 Jan 2026 07:36:31 +0000
Subject: [PATCH] [Linalg] Update Conv Decomposition to work with generic conv
ops
-- This commit updates Conv Decomposition to work with both named as
well as generic convolution ops.
-- This required an update to the `isaConvolutionOfType` API to also
populate dilations/strides info for named convolution ops and since
now a generic LinalgOp is being used as the root op in the pattern
above the assert of the op implementing a ConvolutionOpInterface has
been replaced with an early exit if.
Signed-off-by: Abhishek Varma <abhvarma at amd.com>
---
.../Dialect/Linalg/Transforms/Transforms.h | 27 +-
.../include/mlir/Dialect/Linalg/Utils/Utils.h | 8 +-
.../Dialect/Linalg/Transforms/Transforms.cpp | 107 +--
mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 686 ++++++++++++++----
.../Linalg/transform-op-decompose.mlir | 238 ++++--
5 files changed, 770 insertions(+), 296 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 6678d693719bf..32067358438d3 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1641,16 +1641,16 @@ FailureOr<linalg::GenericOp> deduplicateOperandsAndRemoveDeadResults(
//===----------------------------------------------------------------------===//
/// Rewrites 2-D convolution ops with size-1 window dimensions into 1-D
-/// convolution ops.
+/// convolution ops. Works with both named ops and equivalent generic ops.
template <typename Conv2DOp, typename Conv1DOp>
struct DownscaleSizeOneWindowed2DConvolution final
- : public OpRewritePattern<Conv2DOp> {
- using OpRewritePattern<Conv2DOp>::OpRewritePattern;
+ : public OpInterfaceRewritePattern<LinalgOp> {
+ using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
- FailureOr<Conv1DOp> returningMatchAndRewrite(Conv2DOp convOp,
+ FailureOr<Conv1DOp> returningMatchAndRewrite(LinalgOp convOp,
PatternRewriter &rewriter) const;
- LogicalResult matchAndRewrite(Conv2DOp convOp,
+ LogicalResult matchAndRewrite(LinalgOp convOp,
PatternRewriter &rewriter) const override {
return returningMatchAndRewrite(convOp, rewriter);
}
@@ -1664,29 +1664,28 @@ extern template struct DownscaleSizeOneWindowed2DConvolution<Conv2DNchwFchwOp,
/// Rewrites 2-D depthwise convolution ops with size-1 (w, kw) or (h, kh)
/// dimensions into 1-D depthwise convolution ops.
struct DownscaleDepthwiseConv2DNhwcHwcOp final
- : public OpRewritePattern<DepthwiseConv2DNhwcHwcOp> {
+ : public OpInterfaceRewritePattern<LinalgOp> {
DownscaleDepthwiseConv2DNhwcHwcOp(MLIRContext *context,
PatternBenefit benefit = 1)
- : OpRewritePattern<DepthwiseConv2DNhwcHwcOp>(context, benefit) {}
+ : OpInterfaceRewritePattern<LinalgOp>(context, benefit) {}
FailureOr<DepthwiseConv1DNwcWcOp>
- returningMatchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
- PatternRewriter &rewriter) const;
+ returningMatchAndRewrite(LinalgOp convOp, PatternRewriter &rewriter) const;
- LogicalResult matchAndRewrite(DepthwiseConv2DNhwcHwcOp convOp,
+ LogicalResult matchAndRewrite(LinalgOp convOp,
PatternRewriter &rewriter) const override {
return returningMatchAndRewrite(convOp, rewriter);
}
};
-struct DownscaleConv2DOp final : public OpRewritePattern<Conv2DOp> {
+struct DownscaleConv2DOp final : public OpInterfaceRewritePattern<LinalgOp> {
DownscaleConv2DOp(MLIRContext *context, PatternBenefit benefit = 1)
- : OpRewritePattern<Conv2DOp>(context, benefit) {}
+ : OpInterfaceRewritePattern<LinalgOp>(context, benefit) {}
- FailureOr<Conv1DOp> returningMatchAndRewrite(Conv2DOp convOp,
+ FailureOr<Conv1DOp> returningMatchAndRewrite(LinalgOp convOp,
PatternRewriter &rewriter) const;
- LogicalResult matchAndRewrite(Conv2DOp convOp,
+ LogicalResult matchAndRewrite(LinalgOp convOp,
PatternRewriter &rewriter) const override {
return returningMatchAndRewrite(convOp, rewriter);
}
diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index 9da01f30b52d2..16d557a6ed7fa 100644
--- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
@@ -108,10 +108,12 @@ getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes);
/// Given a linalg `op` this function returns true if it is a convolution op of
/// type `ConvOpTy` and populates `dilations` and `strides` with values inferred
-/// from the indexing maps.
+/// from the indexing maps. If `dilations` or `strides` is nullptr, the
+/// corresponding values are not populated.
template <typename ConvOpTy>
-bool isaConvolutionOpOfType(LinalgOp op, SmallVector<int64_t> *dilations,
- SmallVector<int64_t> *strides);
+bool isaConvolutionOpOfType(LinalgOp op,
+ SmallVector<int64_t> *dilations = nullptr,
+ SmallVector<int64_t> *strides = nullptr);
//===----------------------------------------------------------------------===//
// Fusion / Tiling utilities
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 96cc378f6c21a..7972408318b95 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -32,6 +32,7 @@
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/InterleavedRange.h"
#include "llvm/Support/raw_ostream.h"
+#include <type_traits>
#include <utility>
#define DEBUG_TYPE "linalg-transforms"
@@ -1406,13 +1407,18 @@ LogicalResult DecomposeOuterUnitDimsUnPackOpPattern::matchAndRewrite(
template <typename Conv2DOp, typename Conv1DOp>
FailureOr<Conv1DOp> DownscaleSizeOneWindowed2DConvolution<Conv2DOp, Conv1DOp>::
- returningMatchAndRewrite(Conv2DOp convOp, PatternRewriter &rewriter) const {
+ returningMatchAndRewrite(LinalgOp convOp, PatternRewriter &rewriter) const {
+ // Check if this LinalgOp is of the expected Conv2DOp type (named or generic).
+ SmallVector<int64_t> dilations, strides;
+ if (!linalg::isaConvolutionOpOfType<Conv2DOp>(convOp, &dilations, &strides))
+ return failure();
+
if (convOp.hasPureBufferSemantics())
return failure(); // To be implemented.
- Value input = convOp.getInputs().front();
- Value kernel = convOp.getInputs().back();
- Value output = convOp.getOutputs().front();
+ Value input = convOp.getDpsInputs().front();
+ Value kernel = convOp.getDpsInputs().back();
+ Value output = convOp.getDpsInits().front();
auto inputType = dyn_cast<RankedTensorType>(input.getType());
auto kernelType = dyn_cast<RankedTensorType>(kernel.getType());
@@ -1421,38 +1427,33 @@ FailureOr<Conv1DOp> DownscaleSizeOneWindowed2DConvolution<Conv2DOp, Conv1DOp>::
auto kernelShape = kernelType.getShape();
auto outputShape = outputType.getShape();
- // Get domain indices based on conv2D layout.
- auto [khIndex, kwIndex, ohIndex, owIndex] =
- TypeSwitch<Operation *, std::tuple<int64_t, int64_t, int64_t, int64_t>>(
- convOp)
- .Case([&](linalg::Conv2DNhwcHwcfOp op) {
- return std::make_tuple(0, 1, 1, 2);
- })
- .Case([&](linalg::Conv2DNchwFchwOp op) {
- return std::make_tuple(2, 3, 2, 3);
- })
- .Case([&](linalg::PoolingNhwcSumOp op) {
- return std::make_tuple(0, 1, 1, 2);
- })
- .Case([&](linalg::PoolingNchwSumOp op) {
- return std::make_tuple(0, 1, 2, 3);
- })
- .Case([&](linalg::PoolingNhwcMaxOp op) {
- return std::make_tuple(0, 1, 1, 2);
- })
- .Case([&](linalg::PoolingNhwcMaxUnsignedOp op) {
- return std::make_tuple(0, 1, 1, 2);
- })
- .Case([&](linalg::PoolingNhwcMinOp op) {
- return std::make_tuple(0, 1, 1, 2);
- })
- .Case([&](linalg::PoolingNhwcMinUnsignedOp op) {
- return std::make_tuple(0, 1, 1, 2);
- })
- .Case([&](linalg::PoolingNchwMaxOp op) {
- return std::make_tuple(0, 1, 2, 3);
- })
- .DefaultUnreachable("unexpected conv2d/pool2d operation.");
+ // Get domain indices based on Conv2DOp type. These are known at compile time.
+ int64_t khIndex, kwIndex, ohIndex, owIndex;
+ if constexpr (std::is_same_v<Conv2DOp, linalg::Conv2DNhwcHwcfOp> ||
+ std::is_same_v<Conv2DOp, linalg::PoolingNhwcSumOp> ||
+ std::is_same_v<Conv2DOp, linalg::PoolingNhwcMaxOp> ||
+ std::is_same_v<Conv2DOp, linalg::PoolingNhwcMaxUnsignedOp> ||
+ std::is_same_v<Conv2DOp, linalg::PoolingNhwcMinOp> ||
+ std::is_same_v<Conv2DOp, linalg::PoolingNhwcMinUnsignedOp>) {
+ // NHWC layout: kernel [H, W, ...], output [N, H, W, C]
+ khIndex = 0;
+ kwIndex = 1;
+ ohIndex = 1;
+ owIndex = 2;
+ } else if constexpr (std::is_same_v<Conv2DOp, linalg::Conv2DNchwFchwOp>) {
+ // NCHW_FCHW layout: kernel [..., H, W], output [N, C, H, W]
+ khIndex = 2;
+ kwIndex = 3;
+ ohIndex = 2;
+ owIndex = 3;
+ } else if constexpr (std::is_same_v<Conv2DOp, linalg::PoolingNchwSumOp> ||
+ std::is_same_v<Conv2DOp, linalg::PoolingNchwMaxOp>) {
+ // NCHW pooling layout: kernel [H, W], output [N, C, H, W]
+ khIndex = 0;
+ kwIndex = 1;
+ ohIndex = 2;
+ owIndex = 3;
+ }
// Only handle the case where at least one of the window dimensions is
// of size 1. Other cases can rely on tiling to reduce to such cases.
@@ -1484,13 +1485,9 @@ FailureOr<Conv1DOp> DownscaleSizeOneWindowed2DConvolution<Conv2DOp, Conv1DOp>::
// Rank-reduce strides and dilations too.
// TODO: dropDim 1-liner helper.
- auto strides =
- llvm::to_vector<4>(convOp.getStrides().template getValues<int64_t>());
strides.erase(strides.begin() + (removeH ? 0 : 1));
auto stridesAttr = rewriter.getI64VectorAttr(strides);
- auto dilations =
- llvm::to_vector<4>(convOp.getDilations().template getValues<int64_t>());
dilations.erase(dilations.begin() + (removeH ? 0 : 1));
auto dilationsAttr = rewriter.getI64VectorAttr(dilations);
@@ -1527,13 +1524,19 @@ template struct linalg::DownscaleSizeOneWindowed2DConvolution<PoolingNchwMaxOp,
FailureOr<DepthwiseConv1DNwcWcOp>
DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite(
- DepthwiseConv2DNhwcHwcOp convOp, PatternRewriter &rewriter) const {
+ LinalgOp convOp, PatternRewriter &rewriter) const {
+ // Check if this LinalgOp is a DepthwiseConv2DNhwcHwcOp (named or generic).
+ SmallVector<int64_t> dilations, strides;
+ if (!linalg::isaConvolutionOpOfType<DepthwiseConv2DNhwcHwcOp>(
+ convOp, &dilations, &strides))
+ return failure();
+
if (convOp.hasPureBufferSemantics())
return failure(); // To be implemented.
- Value input = convOp.getInputs().front();
- Value kernel = convOp.getInputs().back();
- Value output = convOp.getOutputs().front();
+ Value input = convOp.getDpsInputs().front();
+ Value kernel = convOp.getDpsInputs().back();
+ Value output = convOp.getDpsInits().front();
auto inputType = dyn_cast<RankedTensorType>(input.getType());
auto kernelType = dyn_cast<RankedTensorType>(kernel.getType());
@@ -1572,12 +1575,9 @@ DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite(
// Rank-reduce strides and dilations too.
// TODO: dropDim 1-liner helper.
- auto strides = llvm::to_vector<4>(convOp.getStrides().getValues<int64_t>());
strides.erase(strides.begin() + (removeH ? 0 : 1));
auto stridesAttr = rewriter.getI64VectorAttr(strides);
- auto dilations =
- llvm::to_vector<4>(convOp.getDilations().getValues<int64_t>());
dilations.erase(dilations.begin() + (removeH ? 0 : 1));
auto dilationsAttr = rewriter.getI64VectorAttr(dilations);
@@ -1594,14 +1594,19 @@ DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite(
}
FailureOr<Conv1DOp>
-DownscaleConv2DOp::returningMatchAndRewrite(Conv2DOp convOp,
+DownscaleConv2DOp::returningMatchAndRewrite(LinalgOp convOp,
PatternRewriter &rewriter) const {
+ // Check if this LinalgOp is a Conv2DOp (named or generic).
+ SmallVector<int64_t> dilations, strides;
+ if (!linalg::isaConvolutionOpOfType<Conv2DOp>(convOp, &dilations, &strides))
+ return failure();
+
if (convOp.hasPureBufferSemantics())
return failure(); // To be implemented.
- Value input = convOp.getInputs().front();
- Value kernel = convOp.getInputs().back();
- Value output = convOp.getOutputs().front();
+ Value input = convOp.getDpsInputs().front();
+ Value kernel = convOp.getDpsInputs().back();
+ Value output = convOp.getDpsInits().front();
auto inputType = dyn_cast<RankedTensorType>(input.getType());
auto kernelType = dyn_cast<RankedTensorType>(kernel.getType());
diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 2718124251c18..1cdd01567c4e7 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -373,11 +373,8 @@ static bool bodyMatcherForMaxSignedPoolOps(Value yieldVal, Block *body) {
body);
}
-// max_unsigned ops should not allow float data type.
-// TODO(#164800): Retire OPDSL logic.
static bool bodyMatcherForMaxUnsignedPoolOps(Value yieldVal, Block *body) {
- return bodyMatcherForPoolOps<arith::MaximumFOp, arith::MaxUIOp>(yieldVal,
- body);
+ return bodyMatcherForPoolOps<arith::MaxUIOp>(yieldVal, body);
}
static bool bodyMatcherForMinSignedPoolOps(Value yieldVal, Block *body) {
@@ -385,11 +382,8 @@ static bool bodyMatcherForMinSignedPoolOps(Value yieldVal, Block *body) {
body);
}
-// min_unsigned ops should not allow float data type.
-// TODO(#164800): Retire OPDSL logic.
static bool bodyMatcherForMinUnsignedPoolOps(Value yieldVal, Block *body) {
- return bodyMatcherForPoolOps<arith::MinimumFOp, arith::MinUIOp>(yieldVal,
- body);
+ return bodyMatcherForPoolOps<arith::MinUIOp>(yieldVal, body);
}
static bool bodyMatcherForSumPoolOps(Value yieldVal, Block *body) {
@@ -601,11 +595,20 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv1DOp>(LinalgOp op,
SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv1DOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (isa<linalg::Conv1DOp>(op)) {
+ // Conv1DOp has no strides/dilations attributes, default to 1.
+ *dilations = SmallVector<int64_t>(1, 1);
+ *strides = SmallVector<int64_t>(1, 1);
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);
AffineExpr W = m.dim(0);
@@ -622,11 +625,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv1DNwcWcfOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv1DNwcWcfOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv1DNwcWcfOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);
AffineExpr N = m.dim(0);
@@ -646,11 +657,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv1DNcwFcwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv1DNcwFcwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv1DNcwFcwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);
AffineExpr N = m.dim(0);
@@ -670,11 +689,20 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DOp>(LinalgOp op,
SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (isa<linalg::Conv2DOp>(op)) {
+ // Conv2DOp has no strides/dilations attributes, default to 1.
+ *dilations = SmallVector<int64_t>(2, 1);
+ *strides = SmallVector<int64_t>(2, 1);
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr H = m.dim(0);
@@ -694,11 +722,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNhwcHwcfOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNhwcHwcfOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNhwcHwcfOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -721,11 +757,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNhwcHwcfQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNhwcHwcfQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNhwcHwcfQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -750,11 +794,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNhwcFhwcOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNhwcFhwcOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNhwcFhwcOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -777,11 +829,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNhwcFhwcQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNhwcFhwcQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNhwcFhwcQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -806,11 +866,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNchwFchwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNchwFchwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNchwFchwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -833,11 +901,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNchwFchwQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNchwFchwQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNchwFchwQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -862,11 +938,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNgchwFgchwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNgchwFgchwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNgchwFgchwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -891,11 +975,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNgchwGfchwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNgchwGfchwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNgchwGfchwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -920,11 +1012,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNgchwGfchwQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNgchwGfchwQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNgchwGfchwQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -951,11 +1051,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNhwgcGfhwcOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNhwgcGfhwcOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNhwgcGfhwcOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -980,11 +1088,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv2DNhwgcGfhwcQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv2DNhwgcGfhwcQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv2DNhwgcGfhwcQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1011,11 +1127,20 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv3DOp>(LinalgOp op,
SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv3DOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (isa<linalg::Conv3DOp>(op)) {
+ // Conv3DOp has no strides/dilations attributes, default to 1.
+ *dilations = SmallVector<int64_t>(3, 1);
+ *strides = SmallVector<int64_t>(3, 1);
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr D = m.dim(0);
@@ -1039,11 +1164,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv3DNdhwcDhwcfOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv3DNdhwcDhwcfOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv3DNdhwcDhwcfOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1070,11 +1203,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv3DNdhwcDhwcfQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv3DNdhwcDhwcfQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv3DNdhwcDhwcfQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1103,11 +1244,19 @@ template <>
bool isaConvolutionOpOfType<linalg::Conv3DNcdhwFcdhwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::Conv3DNcdhwFcdhwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp = dyn_cast<linalg::Conv3DNcdhwFcdhwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1134,11 +1283,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv1DNcwCwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv1DNcwCwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv1DNcwCwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1157,11 +1315,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv1DNwcWcOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv1DNwcWcOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv1DNwcWcOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1180,11 +1347,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv1DNwcWcmOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv1DNwcWcmOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv1DNwcWcmOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1204,11 +1380,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv2DNchwChwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv2DNchwChwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv2DNchwChwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1230,11 +1415,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv2DNhwcHwcOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv2DNhwcHwcOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1256,11 +1450,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv2DNhwcHwcQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv2DNhwcHwcQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1284,11 +1487,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcmOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv2DNhwcHwcmOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv2DNhwcHwcmOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1311,11 +1523,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcmQOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv2DNhwcHwcmQOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv2DNhwcHwcmQOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1340,11 +1561,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv3DNdhwcDhwcOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv3DNdhwcDhwcOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv3DNdhwcDhwcOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1370,11 +1600,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv3DNcdhwCdhwOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv3DNcdhwCdhwOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv3DNcdhwCdhwOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1400,11 +1639,20 @@ template <>
bool isaConvolutionOpOfType<linalg::DepthwiseConv3DNdhwcDhwcmOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::DepthwiseConv3DNdhwcDhwcmOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto convOp =
+ dyn_cast<linalg::DepthwiseConv3DNdhwcDhwcmOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(convOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(convOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides);
AffineExpr N = m.dim(0);
@@ -1431,11 +1679,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNhwcMaxOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNhwcMaxOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNhwcMaxOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::MaxSigned);
@@ -1458,11 +1714,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNhwcMinOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNhwcMinOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNhwcMinOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::MinSigned);
@@ -1485,11 +1749,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNhwcSumOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNhwcSumOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNhwcSumOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::Sum);
@@ -1512,11 +1784,20 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNhwcMaxUnsignedOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNhwcMaxUnsignedOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp =
+ dyn_cast<linalg::PoolingNhwcMaxUnsignedOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::MaxUnsigned);
@@ -1539,11 +1820,20 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNhwcMinUnsignedOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNhwcMinUnsignedOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp =
+ dyn_cast<linalg::PoolingNhwcMinUnsignedOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::MinUnsigned);
@@ -1566,11 +1856,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNchwSumOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNchwSumOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNchwSumOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::Sum);
@@ -1593,11 +1891,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNchwMaxOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNchwMaxOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNchwMaxOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/2, dilations, strides,
PoolingType::MaxSigned);
@@ -1620,11 +1926,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNwcSumOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNwcSumOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNwcSumOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::Sum);
@@ -1644,11 +1958,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNcwSumOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNcwSumOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNcwSumOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::Sum);
@@ -1668,11 +1990,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNwcMaxOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNwcMaxOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNwcMaxOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::MaxSigned);
@@ -1692,11 +2022,20 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNwcMaxUnsignedOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNwcMaxUnsignedOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp =
+ dyn_cast<linalg::PoolingNwcMaxUnsignedOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::MaxUnsigned);
@@ -1716,11 +2055,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNcwMaxOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNcwMaxOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNcwMaxOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::MaxSigned);
@@ -1740,11 +2087,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNwcMinOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNwcMinOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNwcMinOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::MinSigned);
@@ -1764,11 +2119,20 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNwcMinUnsignedOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNwcMinUnsignedOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp =
+ dyn_cast<linalg::PoolingNwcMinUnsignedOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/1, dilations, strides,
PoolingType::MinUnsigned);
@@ -1788,11 +2152,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNdhwcSumOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNdhwcSumOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNdhwcSumOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides,
PoolingType::Sum);
@@ -1819,11 +2191,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNdhwcMaxOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNdhwcMaxOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNdhwcMaxOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides,
PoolingType::MaxSigned);
@@ -1850,11 +2230,19 @@ template <>
bool isaConvolutionOpOfType<linalg::PoolingNdhwcMinOp>(
LinalgOp op, SmallVector<int64_t> *dilations,
SmallVector<int64_t> *strides) {
- if (isa<linalg::PoolingNdhwcMinOp>(op))
+ SmallVector<int64_t> localDilations, localStrides;
+ if (!dilations)
+ dilations = &localDilations;
+ if (!strides)
+ strides = &localStrides;
+ if (auto poolOp = dyn_cast<linalg::PoolingNdhwcMinOp>(op.getOperation())) {
+ *dilations = llvm::to_vector(poolOp.getDilations().getValues<int64_t>());
+ *strides = llvm::to_vector(poolOp.getStrides().getValues<int64_t>());
return true;
+ }
- assert(isaConvolutionOpInterface(op) &&
- "expected op to implement ConvolutionOpInterface");
+ if (!isaConvolutionOpInterface(op))
+ return false;
ConvMatcherBuilder m(op, /*spatialRank=*/3, dilations, strides,
PoolingType::MinSigned);
diff --git a/mlir/test/Dialect/Linalg/transform-op-decompose.mlir b/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
index 60a4c555fa19a..7798cb76e4fb9 100644
--- a/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
@@ -3,113 +3,168 @@
// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map_nhwc_hwcf_input = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>
+#map_nhwc_hwcf_filter = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>
+#map_nhwc_hwcf_output = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+
// CHECK-LABEL: @conv_2d_nhwc_hwcf
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x?xf32>,
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?x?x?xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32>
func.func @conv_2d_nhwc_hwcf(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?x?x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.conv_1d_nwc_wcf
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to conv_1d_nwc_wcf
+ // CHECK-COUNT-2: linalg.conv_1d_nwc_wcf
%0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?x?x?xf32>)
outs (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x1x?x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_nhwc_hwcf_input, #map_nhwc_hwcf_filter, #map_nhwc_hwcf_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%input, %filter : tensor<?x1x?x?xf32>, tensor<1x?x?x?xf32>) outs(%0 : tensor<?x1x?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.mulf %in, %in_0 : f32
+ %3 = arith.addf %out, %2 : f32
+ linalg.yield %3 : f32
+ } -> tensor<?x1x?x?xf32>
+ return %1 : tensor<?x1x?x?xf32>
}
+#map_nchw_fchw_input = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d4, d2 + d5, d3 + d6)>
+#map_nchw_fchw_filter = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d4, d5, d6)>
+#map_nchw_fchw_output = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+
// CHECK-LABEL: @conv_2d_nchw_fchw
// CHECK-SAME: (%[[ARG0:[0-9a-z]+]]: tensor<?x?x1x?xf32>,
// CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<?x?x1x?xf32>,
// CHECK-SAME: %[[ARG2:[0-9a-z]+]]: tensor<?x?x1x?xf32>)
func.func @conv_2d_nchw_fchw(%input: tensor<?x?x1x?xf32>, %filter: tensor<?x?x1x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.conv_1d_ncw_fcw
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to conv_1d_ncw_fcw
+ // CHECK-COUNT-2: linalg.conv_1d_ncw_fcw
%0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x?x1x?xf32>, tensor<?x?x1x?xf32>)
outs (%init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x?x1x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_nchw_fchw_input, #map_nchw_fchw_filter, #map_nchw_fchw_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%input, %filter : tensor<?x?x1x?xf32>, tensor<?x?x1x?xf32>) outs(%0 : tensor<?x?x1x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.mulf %in, %in_0 : f32
+ %3 = arith.addf %out, %2 : f32
+ linalg.yield %3 : f32
+ } -> tensor<?x?x1x?xf32>
+ return %1 : tensor<?x?x1x?xf32>
}
+#map_depthwise_nhwc_hwc_input = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d3)>
+#map_depthwise_nhwc_hwc_filter = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3)>
+#map_depthwise_nhwc_hwc_output = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
// CHECK-LABEL: @depthwise_conv_2d_nhwc_hwc
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x113x96xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x3x96xf32>
func.func @depthwise_conv_2d_nhwc_hwc(%input: tensor<1x1x113x96xf32>, %filter: tensor<1x3x96xf32>) -> tensor<1x1x56x96xf32> {
// CHECK: %[[RES:.+]] = tensor.empty
%init = tensor.empty() : tensor<1x1x56x96xf32>
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICERES:.+]] = tensor.extract_slice %[[RES]]
- // CHECK: %[[OPRES:.+]] = linalg.depthwise_conv_1d_nwc_wc
- // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]]
- // CHECK-SAME: outs(%[[SLICERES]]
- // CHECK: %[[INSERTED:.+]] = tensor.insert_slice %[[OPRES]] into %[[RES]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // Both named and generic ops should decompose to depthwise_conv_1d_nwc_wc
+ // CHECK-COUNT-2: linalg.depthwise_conv_1d_nwc_wc
%0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: tensor<1x1x113x96xf32>, tensor<1x3x96xf32>)
outs(%init: tensor<1x1x56x96xf32>) -> tensor<1x1x56x96xf32>
- // CHECK: %[[INSERTED]]
- return %0: tensor<1x1x56x96xf32>
+ // Generic op version with same semantics (strides = 2).
+ %1 = linalg.generic {indexing_maps = [#map_depthwise_nhwc_hwc_input, #map_depthwise_nhwc_hwc_filter, #map_depthwise_nhwc_hwc_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<1x1x113x96xf32>, tensor<1x3x96xf32>) outs(%0 : tensor<1x1x56x96xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.mulf %in, %in_0 : f32
+ %3 = arith.addf %out, %2 : f32
+ linalg.yield %3 : f32
+ } -> tensor<1x1x56x96xf32>
+ return %1: tensor<1x1x56x96xf32>
}
+#map_conv_2d_input = affine_map<(d0, d1, d2, d3) -> (d0 + d2, d1 + d3)>
+#map_conv_2d_filter = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
+#map_conv_2d_output = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+
// CHECK-LABEL: @conv_2d
// CHECK-SAME: (%[[ARG0:[0-9a-z]+]]: tensor<1x?xf32>,
// CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<1x?xf32>,
// CHECK-SAME: %[[ARG2:[0-9a-z]+]]: tensor<1x?xf32>)
func.func @conv_2d(%input: tensor<1x?xf32>, %filter: tensor<1x?xf32>, %init: tensor<1x?xf32>) -> tensor<1x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.conv_1d
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to conv_1d
+ // CHECK-COUNT-2: linalg.conv_1d
%0 = linalg.conv_2d
ins (%input, %filter: tensor<1x?xf32>, tensor<1x?xf32>)
outs (%init: tensor<1x?xf32>) -> tensor<1x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<1x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_conv_2d_input, #map_conv_2d_filter, #map_conv_2d_output], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<1x?xf32>, tensor<1x?xf32>) outs(%0 : tensor<1x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.mulf %in, %in_0 : f32
+ %3 = arith.addf %out, %2 : f32
+ linalg.yield %3 : f32
+ } -> tensor<1x?xf32>
+ return %1 : tensor<1x?xf32>
}
+#map_pooling_nhwc_input = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>
+#map_pooling_nhwc_filter = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
+#map_pooling_nhwc_output = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
// CHECK-LABEL: @pooling_nhwc_sum
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x?xf32>,
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32>
func.func @pooling_nhwc_sum(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_nwc_sum
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_nwc_sum
+ // CHECK-COUNT-2: linalg.pooling_nwc_sum
%0 = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?xf32>)
outs (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x1x?x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nhwc_input, #map_pooling_nhwc_filter, #map_pooling_nhwc_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x1x?x?xf32>, tensor<1x?xf32>) outs(%0 : tensor<?x1x?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.addf %out, %in : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x1x?x?xf32>
+ return %1 : tensor<?x1x?x?xf32>
}
+#map_pooling_nchw_input = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2 + d4, d3 + d5)>
+#map_pooling_nchw_filter = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
+#map_pooling_nchw_output = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
// CHECK-LABEL: @pooling_nchw_sum
// CHECK-SAME: (%[[ARG0:[0-9a-z]+]]: tensor<?x?x1x?xf32>,
// CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<1x?xf32>,
// CHECK-SAME: %[[ARG2:[0-9a-z]+]]: tensor<?x?x1x?xf32>)
func.func @pooling_nchw_sum(%input: tensor<?x?x1x?xf32>, %filter: tensor<1x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_ncw_sum
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_ncw_sum
+ // CHECK-COUNT-2: linalg.pooling_ncw_sum
%0 = linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x?x1x?xf32>, tensor<1x?xf32>)
outs (%init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x?x1x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nchw_input, #map_pooling_nchw_filter, #map_pooling_nchw_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x?x1x?xf32>, tensor<1x?xf32>) outs(%0 : tensor<?x?x1x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.addf %out, %in : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?x1x?xf32>
+ return %1 : tensor<?x?x1x?xf32>
}
// CHECK-LABEL: @pooling_nhwc_max
@@ -117,17 +172,22 @@ func.func @pooling_nchw_sum(%input: tensor<?x?x1x?xf32>, %filter: tensor<1x?xf32
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32>
func.func @pooling_nhwc_max(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_nwc_max
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_nwc_max
+ // CHECK-COUNT-2: linalg.pooling_nwc_max
%0 = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?xf32>)
outs (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x1x?x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nhwc_input, #map_pooling_nhwc_filter, #map_pooling_nhwc_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x1x?x?xf32>, tensor<1x?xf32>) outs(%0 : tensor<?x1x?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.maximumf %out, %in : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x1x?x?xf32>
+ return %1 : tensor<?x1x?x?xf32>
}
// CHECK-LABEL: @pooling_nhwc_max_unsigned
@@ -135,17 +195,22 @@ func.func @pooling_nhwc_max(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?xf32
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?xi32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xi32>
func.func @pooling_nhwc_max_unsigned(%input: tensor<?x1x?x?xi32>, %filter: tensor<1x?xi32>, %init: tensor<?x1x?x?xi32>) -> tensor<?x1x?x?xi32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_nwc_max_unsigned
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_nwc_max_unsigned
+ // CHECK-COUNT-2: linalg.pooling_nwc_max_unsigned
%0 = linalg.pooling_nhwc_max_unsigned {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xi32>, tensor<1x?xi32>)
outs (%init: tensor<?x1x?x?xi32>) -> tensor<?x1x?x?xi32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x1x?x?xi32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nhwc_input, #map_pooling_nhwc_filter, #map_pooling_nhwc_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x1x?x?xi32>, tensor<1x?xi32>) outs(%0 : tensor<?x1x?x?xi32>) {
+ ^bb0(%in: i32, %in_0: i32, %out: i32):
+ %2 = arith.maxui %out, %in : i32
+ linalg.yield %2 : i32
+ } -> tensor<?x1x?x?xi32>
+ return %1 : tensor<?x1x?x?xi32>
}
// CHECK-LABEL: @pooling_nhwc_min
@@ -153,17 +218,22 @@ func.func @pooling_nhwc_max_unsigned(%input: tensor<?x1x?x?xi32>, %filter: tenso
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32>
func.func @pooling_nhwc_min(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_nwc_min
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_nwc_min
+ // CHECK-COUNT-2: linalg.pooling_nwc_min
%0 = linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?xf32>)
outs (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x1x?x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nhwc_input, #map_pooling_nhwc_filter, #map_pooling_nhwc_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x1x?x?xf32>, tensor<1x?xf32>) outs(%0 : tensor<?x1x?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.minimumf %out, %in : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x1x?x?xf32>
+ return %1 : tensor<?x1x?x?xf32>
}
// CHECK-LABEL: @pooling_nhwc_min_unsigned
@@ -171,17 +241,22 @@ func.func @pooling_nhwc_min(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?xf32
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?xi32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xi32>
func.func @pooling_nhwc_min_unsigned(%input: tensor<?x1x?x?xi32>, %filter: tensor<1x?xi32>, %init: tensor<?x1x?x?xi32>) -> tensor<?x1x?x?xi32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_nwc_min_unsigned
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_nwc_min_unsigned
+ // CHECK-COUNT-2: linalg.pooling_nwc_min_unsigned
%0 = linalg.pooling_nhwc_min_unsigned {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xi32>, tensor<1x?xi32>)
outs (%init: tensor<?x1x?x?xi32>) -> tensor<?x1x?x?xi32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x1x?x?xi32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nhwc_input, #map_pooling_nhwc_filter, #map_pooling_nhwc_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x1x?x?xi32>, tensor<1x?xi32>) outs(%0 : tensor<?x1x?x?xi32>) {
+ ^bb0(%in: i32, %in_0: i32, %out: i32):
+ %2 = arith.minui %out, %in : i32
+ linalg.yield %2 : i32
+ } -> tensor<?x1x?x?xi32>
+ return %1 : tensor<?x1x?x?xi32>
}
// CHECK-LABEL: @pooling_nchw_max
@@ -189,17 +264,22 @@ func.func @pooling_nhwc_min_unsigned(%input: tensor<?x1x?x?xi32>, %filter: tenso
// CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<1x?xf32>,
// CHECK-SAME: %[[ARG2:[0-9a-z]+]]: tensor<?x?x1x?xf32>)
func.func @pooling_nchw_max(%input: tensor<?x?x1x?xf32>, %filter: tensor<1x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> {
- // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
- // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
- // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
- // CHECK: %[[SLICERES:.+]] = linalg.pooling_ncw_max
- // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
+ // CHECK: tensor.extract_slice %[[ARG0]]
+ // CHECK: tensor.extract_slice %[[ARG1]]
+ // CHECK: tensor.extract_slice %[[ARG2]]
+ // Both named and generic ops should decompose to pooling_ncw_max
+ // CHECK-COUNT-2: linalg.pooling_ncw_max
%0 = linalg.pooling_nchw_max {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x?x1x?xf32>, tensor<1x?xf32>)
outs (%init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32>
- // CHECK: return %[[RES]]
- return %0 : tensor<?x?x1x?xf32>
+ // Generic op version with same semantics.
+ %1 = linalg.generic {indexing_maps = [#map_pooling_nchw_input, #map_pooling_nchw_filter, #map_pooling_nchw_output], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<?x?x1x?xf32>, tensor<1x?xf32>) outs(%0 : tensor<?x?x1x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %2 = arith.maximumf %out, %in : f32
+ linalg.yield %2 : f32
+ } -> tensor<?x?x1x?xf32>
+ return %1 : tensor<?x?x1x?xf32>
}
func.func @softmax(%arg0: tensor<2x16x32xf32>, %dst: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
More information about the Mlir-commits
mailing list