[GlobalOptimization] Make im2col a generic convolution rewrite (#25011) Replaces the three layout-specific im2col patterns (NHWC, NCHW, depthwise NHWC) with one generic rewrite that handles any convolution, named or generic, including grouped, dilated, 1-D/3-D, dynamic and integer ones. The pattern gathers the convolved input with an all-parallel `linalg.generic`, rewires the op to read the gather, keeps the original payload, and collapses the result into a matmul using `DispatchCreation::getCollapsibleLoops`. Add a basic heuristic via a control function that skips convolutions without an output channel dimension (depthwise, pooling), where the gather is never amortized. The pass is renamed to `iree-global-opt-convert-conv-to-im2col`; the pipeline flag is unchanged. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Thomas Ziereis <ziereis@roofline.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/DispatchCreation/BUILD.bazel b/compiler/src/iree/compiler/DispatchCreation/BUILD.bazel index 25cbfaa..7a95d31 100644 --- a/compiler/src/iree/compiler/DispatchCreation/BUILD.bazel +++ b/compiler/src/iree/compiler/DispatchCreation/BUILD.bazel
@@ -56,6 +56,7 @@ "TransposeGenericOps.cpp", ], hdrs = [ + "CollapseDimensions.h", "FusionUtils.h", "MaterializedLayoutTarget.h", "Passes.h",
diff --git a/compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt b/compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt index 1492a9e..e054fb7 100644 --- a/compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt +++ b/compiler/src/iree/compiler/DispatchCreation/CMakeLists.txt
@@ -14,6 +14,7 @@ NAME DispatchCreation HDRS + "CollapseDimensions.h" "FusionUtils.h" "MaterializedLayoutTarget.h" "Passes.h"
diff --git a/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.cpp b/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.cpp index 73d6edb..a5a7145 100644 --- a/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.cpp +++ b/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.cpp
@@ -4,6 +4,8 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +#include "iree/compiler/DispatchCreation/CollapseDimensions.h" + #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" #include "iree/compiler/Dialect/Encoding/Utils/Utils.h" #include "iree/compiler/Dialect/Flow/IR/FlowOps.h" @@ -111,11 +113,7 @@ .getResult(); } -/// Searches the same sequence in all the affine maps and collapses these -/// dimensions. It only applies these to "parallel" loops without mixing them -/// with "reduction" types. It is expected that the `op` has projected -/// permutations only as indexing maps. (Checked using `isEligibleForCollapse`). -static SmallVector<ReassociationIndices> getCollapsibleLoops(Operation *op) { +SmallVector<ReassociationIndices> getCollapsibleLoops(Operation *op) { auto fusionInterfaceOp = cast<LinalgFusionOpInterface>(op); auto tilingInterfaceOp = cast<TilingInterface>(op);
diff --git a/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.h b/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.h new file mode 100644 index 0000000..f43da22 --- /dev/null +++ b/compiler/src/iree/compiler/DispatchCreation/CollapseDimensions.h
@@ -0,0 +1,26 @@ +// Copyright 2026 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#ifndef IREE_COMPILER_DISPATCHCREATION_COLLAPSEDIMENSIONS_H_ +#define IREE_COMPILER_DISPATCHCREATION_COLLAPSEDIMENSIONS_H_ + +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/Operation.h" + +namespace mlir::iree_compiler::DispatchCreation { + +/// Returns the runs of loops of `op` that can be collapsed into one loop. A run +/// consists of loops with the same iterator type whose dimensions appear +/// consecutively, in the same order, in every indexing map that uses any of +/// them. Runs of a single loop are omitted. +/// +/// `op` must implement `LinalgFusionOpInterface` and `TilingInterface`, and all +/// of its indexing maps must be projected permutations. +SmallVector<ReassociationIndices> getCollapsibleLoops(Operation *op); + +} // namespace mlir::iree_compiler::DispatchCreation + +#endif // IREE_COMPILER_DISPATCHCREATION_COLLAPSEDIMENSIONS_H_
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel index 570c141..2d6c95f 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -79,7 +79,7 @@ srcs = [ "CleanupNumericNarrowing.cpp", "ConvertBatchMatmulToMatmul.cpp", - "ConvertConv2DToImg2Col.cpp", + "ConvertConvToIm2Col.cpp", "ConvertQDQToIntegerMath.cpp", "ConvertStridedContractionToContraction.cpp", "DataLayoutPropagation.cpp",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt index d5387a4..3c9ca44 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -73,7 +73,7 @@ SRCS "CleanupNumericNarrowing.cpp" "ConvertBatchMatmulToMatmul.cpp" - "ConvertConv2DToImg2Col.cpp" + "ConvertConvToIm2Col.cpp" "ConvertQDQToIntegerMath.cpp" "ConvertStridedContractionToContraction.cpp" "DataLayoutPropagation.cpp"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/ConvertConv2DToImg2Col.cpp b/compiler/src/iree/compiler/GlobalOptimization/ConvertConv2DToImg2Col.cpp deleted file mode 100644 index 2ce4985..0000000 --- a/compiler/src/iree/compiler/GlobalOptimization/ConvertConv2DToImg2Col.cpp +++ /dev/null
@@ -1,571 +0,0 @@ -// Copyright 2020 The IREE Authors -// -// Licensed under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -#include "iree/compiler/GlobalOptimization/Passes.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" -#include "mlir/Dialect/Utils/StructuredOpsUtils.h" -#include "mlir/IR/AffineExpr.h" -#include "mlir/IR/AffineMap.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -namespace mlir::iree_compiler::GlobalOptimization { - -#define GEN_PASS_DEF_CONVERTCONV2DTOIMG2COLPASS -#include "iree/compiler/GlobalOptimization/Passes.h.inc" // IWYU pragma: export - -static bool hasAllOneValues(DenseIntElementsAttr attr) { - return llvm::all_of( - attr, [](APInt element) { return element.getSExtValue() == 1; }); -} - -static Value createAdd(Location loc, Value x, Value y, bool isInt, - OpBuilder &builder) { - if (isInt) { - return arith::AddIOp::create(builder, loc, x, y); - } - return arith::AddFOp::create(builder, loc, x, y); -} - -static Value createMul(Location loc, Value x, Value y, bool isInt, - OpBuilder &builder) { - if (isInt) { - return arith::MulIOp::create(builder, loc, x, y); - } - return arith::MulFOp::create(builder, loc, x, y); -} - -namespace { - -// Convert linalg.conv_2d_nhwc_hwcf into linalg.generic (for img2col packing) -// and linalg.matmul. -// -// A convolution operation can be written as a matrix-matrix multiplication by -// unfolding the cross correlation between input and filter and explicitly copy -// overlapped sliding window inputs. -// -// Consider 2D input X with single channel input and output and 2x2 filter W: -// [x(0, 0) , x(0, 1) , ..., x(0, n) ] -// [x(1, 0) , x(1, 1) , ..., x(1, n) ] -// [. , . ,. , . ] [w(0, 0), w(0, 1)] -// [. , . , . , . ] (conv) [w(1, 0), w(1, 1)] -// [. , . , ., . ] -// [x(n-1, 0), x(n-1, 1), ..., x(n-1, n-1)] -// -// The packed input data (img2col) is a matrix with |rows| = output spatial -// size, |columns| = filter spatial size. To compute the output Y(i, j) we need -// to calculate the dot product between filter window at input X(x, y)) and the -// filter which will look like the following where r.h.s is the img2col matrix -// and l.h.s is the flattned filter: -// -// clang-format off -// [x(0, 0), x(0, 1), x(1, 0), x(1, 1)] -// [x(0, 1), x(1, 1), x(0, 2), x(1, 2)] (matmul) [w(0, 0), w(0, 1), w(1, 0), w(1, 1)] -// [x(0, 1), x(1, 1), x(0, 2), x(1, 2)] -// [ . , . , . , . ] -// clang-format on -// -// In general for 2D case with (N, H, W, C) input and (Kh, Kw, C, D) filter -// and output (N, Ho, Wo, D) the convolutin is the following matrix-matrix -// multiplication (Ho x Wo, Kh x Kw x C) * (Kh x Kw x C, D) for each input in -// the N input. For the case where N > 1 its a batched matrxi-matrix -// multiplication. -class ConvertConv2DNhwcHwcf final - : public OpRewritePattern<linalg::Conv2DNhwcHwcfOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::Conv2DNhwcHwcfOp convOp, - PatternRewriter &rewriter) const override { - auto inputType = cast<ShapedType>(convOp.getInputs()[0].getType()); - auto filterType = cast<ShapedType>(convOp.getInputs()[1].getType()); - auto outputType = cast<ShapedType>(convOp.getOutputs()[0].getType()); - - if (!filterType.hasStaticShape() || !inputType.hasStaticShape()) { - return rewriter.notifyMatchFailure(convOp, [](Diagnostic &diag) { - diag << "[unimplemented] " - << "expected 'filterType' and 'inputType' to have static shapes."; - }); - } - - // TODO: Support dilation. - if (!hasAllOneValues(convOp.getDilations())) { - return rewriter.notifyMatchFailure(convOp, [](Diagnostic &diag) { - diag << "[unimplemented] " - << "expected no dilations (expected dilations to all be one)."; - }); - } - - Value input = convOp.getInputs()[0]; - Value filter = convOp.getInputs()[1]; - Value output = convOp.getOutputs()[0]; - - auto filterShape = filterType.getShape(); - auto outputShape = outputType.getShape(); - - const int n = outputShape[0]; - const int oh = outputShape[1]; - const int ow = outputShape[2]; - const int oc = outputShape[3]; - const int fh = filterShape[0]; - const int fw = filterShape[1]; - const int ic = filterShape[2]; - - auto loc = convOp.getLoc(); - - SmallVector<int64_t> colTensorShape = {n, oh, ow, fh, fw, ic}; - - Value colTensor = tensor::EmptyOp::create(rewriter, loc, colTensorShape, - inputType.getElementType()); - - AffineExpr nDim, ohDim, owDim, khDim, kwDim, icDim; - bindDims(getContext(), nDim, ohDim, owDim, khDim, kwDim, icDim); - - auto shSym = rewriter.getAffineConstantExpr( - convOp.getStrides().getValues<int64_t>()[0]); - auto swSym = rewriter.getAffineConstantExpr( - convOp.getStrides().getValues<int64_t>()[1]); - - SmallVector<AffineExpr> inputExprs = {nDim, ohDim * shSym + khDim, - owDim * swSym + kwDim, icDim}; - - auto nloops = colTensorShape.size(); - - auto parallel = utils::IteratorType::parallel; - auto reduction = utils::IteratorType::reduction; - SmallVector<utils::IteratorType, 3> img2colIterators(nloops, parallel); - - SmallVector<AffineMap> img2colIndexingMaps = { - AffineMap::get(nloops, 0, inputExprs, rewriter.getContext()), - AffineMap::getMultiDimIdentityMap(nloops, rewriter.getContext())}; - - auto img2ColTensor = linalg::GenericOp::create( - rewriter, loc, colTensor.getType(), - /*inputs=*/input, /*outputs=*/colTensor, img2colIndexingMaps, - img2colIterators, - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { - linalg::YieldOp::create(nestedBuilder, nestedLoc, args[0]); - }); - - SmallVector<ReassociationIndices> img2ColTensorReassocIndices; - SmallVector<ReassociationIndices> outputReassocIndices; - RankedTensorType reshapedImg2ColTensorType, reshapedOutputType; - if (n == 1) { - img2ColTensorReassocIndices = {{0, 1, 2}, {3, 4, 5}}; - outputReassocIndices = {{0, 1, 2}, {3}}; - - reshapedImg2ColTensorType = RankedTensorType::get( - {oh * ow, fh * fw * ic}, inputType.getElementType()); - reshapedOutputType = - RankedTensorType::get({oh * ow, oc}, outputType.getElementType()); - } else { - img2ColTensorReassocIndices = {{0}, {1, 2}, {3, 4, 5}}; - outputReassocIndices = {{0}, {1, 2}, {3}}; - - reshapedImg2ColTensorType = RankedTensorType::get( - {n, oh * ow, fh * fw * ic}, inputType.getElementType()); - reshapedOutputType = - RankedTensorType::get({n, oh * ow, oc}, outputType.getElementType()); - } - - SmallVector<ReassociationIndices> filterReassocIndices = {{0, 1, 2}, {3}}; - auto reshapedFilterType = - RankedTensorType::get({fh * fw * ic, oc}, inputType.getElementType()); - - Value reshapedImg2ColTensor = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedImg2ColTensorType, img2ColTensor.getResult(0), - img2ColTensorReassocIndices); - - Value reshapedFilter = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedFilterType, filter, filterReassocIndices); - - Value reshapedOutput = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedOutputType, output, outputReassocIndices); - - Value result; - if (n == 1) { - auto matmulOp = linalg::MatmulOp::create( - rewriter, loc, reshapedOutputType, - ArrayRef<Value>{reshapedImg2ColTensor, reshapedFilter}, - ArrayRef<Value>{reshapedOutput}); - result = matmulOp.getResults().front(); - } else { - // For cases where batch is not 1, we need to keep the batch dimension - // separate. Because the filter does not share the same batch dimension, - // the batch dimension is only used in indexing the input and output. Thus - // we cannot use existing linalg named ops like linalg.batch_matmul. - // i.e. (B x) M x K * K x N = (B x) M x N - AffineExpr bDim, mDim, nDim, kDim; - bindDims(getContext(), bDim, mDim, nDim, kDim); - auto lhsMap = AffineMap::get(4, 0, {bDim, mDim, kDim}, getContext()); - auto rhsMap = AffineMap::get(4, 0, {kDim, nDim}, getContext()); - auto resultMap = AffineMap::get(4, 0, {bDim, mDim, nDim}, getContext()); - SmallVector<utils::IteratorType> genericIterators = {parallel, parallel, - parallel, reduction}; - bool isInt = isa<IntegerType>(outputType.getElementType()); - auto genericOp = linalg::GenericOp::create( - rewriter, loc, reshapedOutputType, - /*inputs=*/ValueRange{reshapedImg2ColTensor, reshapedFilter}, - /*outputs=*/ValueRange{reshapedOutput}, - ArrayRef<AffineMap>{lhsMap, rhsMap, resultMap}, genericIterators, - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { - Value mul = createMul(loc, args[0], args[1], isInt, nestedBuilder); - Value add = createAdd(loc, mul, args[2], isInt, nestedBuilder); - linalg::YieldOp::create(nestedBuilder, nestedLoc, add); - }); - result = genericOp.getResults().front(); - } - - auto reshapedResult = tensor::ExpandShapeOp::create( - rewriter, loc, outputType, result, outputReassocIndices); - - rewriter.replaceOp(convOp, ArrayRef<Value>{reshapedResult}); - - return success(); - } -}; - -// Similar to the conv pattern above except there is no reduction among the -// input channels so each convolution can be a matrix-vector product and -// by transposing both input filter so channels are outer most the computation -// is a batched matrix-vector product. -class ConvertDepthwiseConv2DNhwcHwc final - : public OpRewritePattern<linalg::DepthwiseConv2DNhwcHwcOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::DepthwiseConv2DNhwcHwcOp convOp, - PatternRewriter &rewriter) const override { - auto inputType = cast<RankedTensorType>(convOp.getInputs()[0].getType()); - auto filterType = cast<RankedTensorType>(convOp.getInputs()[1].getType()); - auto outputType = cast<RankedTensorType>(convOp.getOutputs()[0].getType()); - - if (!filterType.hasStaticShape() || !inputType.hasStaticShape()) { - return rewriter.notifyMatchFailure(convOp, [](Diagnostic &diag) { - diag << "[unimplemented] " - << "expected 'filterType' and 'inputType' to have static shapes."; - }); - } - - // TODO: Support dilation. - if (!hasAllOneValues(convOp.getDilations())) { - return rewriter.notifyMatchFailure(convOp, [](Diagnostic &diag) { - diag << "[unimplemented] " - << "expected no dilations (expected dilations to all be one)."; - }); - } - - auto loc = convOp.getLoc(); - - auto transposeOperand = [&](Value operand, ArrayRef<int64_t> indices) { - auto operandTensorType = cast<RankedTensorType>(operand.getType()); - auto nloops = indices.size(); - auto inputShape = operandTensorType.getShape(); - - SmallVector<AffineExpr> exprs = - llvm::map_to_vector(indices, [&](int64_t index) -> AffineExpr { - return rewriter.getAffineDimExpr(index); - }); - - SmallVector<int64_t> targetShape = llvm::map_to_vector( - indices, [&](int64_t index) -> int64_t { return inputShape[index]; }); - - Value outputTensor = tensor::EmptyOp::create( - rewriter, loc, targetShape, operandTensorType.getElementType()); - - SmallVector<utils::IteratorType> loopAttributeTypes( - nloops, utils::IteratorType::parallel); - - SmallVector<AffineMap> indexingMaps = { - inversePermutation( - AffineMap::get(nloops, 0, exprs, rewriter.getContext())), - AffineMap::getMultiDimIdentityMap(nloops, rewriter.getContext())}; - - auto transposedOp = linalg::GenericOp::create( - rewriter, loc, outputTensor.getType(), - /*inputs=*/operand, /*outputs=*/outputTensor, indexingMaps, - loopAttributeTypes, - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { - linalg::YieldOp::create(nestedBuilder, nestedLoc, args[0]); - }); - - return transposedOp.getResult(0); - }; - - Value input = convOp.getInputs()[0]; - Value filter = convOp.getInputs()[1]; - Value output = convOp.getOutputs()[0]; - - // Transpose input, filter so channels are outermost - auto inputT = transposeOperand(input, {0, 3, 1, 2}); - auto filterT = transposeOperand(filter, {2, 0, 1}); - auto filterTShape = cast<RankedTensorType>(filterT.getType()).getShape(); - auto outputShape = outputType.getShape(); - - const int n = outputShape[0]; - const int oh = outputShape[1]; - const int ow = outputShape[2]; - const int c = outputShape[3]; - const int fh = filterTShape[1]; - const int fw = filterTShape[2]; - - SmallVector<int64_t> colTensorShape = {n, c, oh, ow, fh, fw}; - Value transposedOutputTensor = transposeOperand(output, {0, 3, 1, 2}); - - AffineExpr nDim, cDim, ohDim, owDim, khDim, kwDim; - bindDims(getContext(), nDim, cDim, ohDim, owDim, khDim, kwDim); - - auto shSym = rewriter.getAffineConstantExpr( - convOp.getStrides().getValues<int64_t>()[0]); - auto swSym = rewriter.getAffineConstantExpr( - convOp.getStrides().getValues<int64_t>()[1]); - - SmallVector<AffineExpr> inputExprs = {nDim, cDim, ohDim * shSym + khDim, - owDim * swSym + kwDim}; - - auto nloops = colTensorShape.size(); - - SmallVector<utils::IteratorType> loopAttributeTypes( - nloops, utils::IteratorType::parallel); - - SmallVector<AffineMap> indexingMaps = { - AffineMap::get(nloops, 0, inputExprs, rewriter.getContext()), - AffineMap::getMultiDimIdentityMap(nloops, rewriter.getContext())}; - - Value colTensor = tensor::EmptyOp::create(rewriter, loc, colTensorShape, - inputType.getElementType()); - - auto img2ColTensor = linalg::GenericOp::create( - rewriter, loc, colTensor.getType(), - /*inputs=*/inputT, /*outputs=*/colTensor, indexingMaps, - loopAttributeTypes, - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { - linalg::YieldOp::create(nestedBuilder, nestedLoc, args[0]); - }); - - SmallVector<ReassociationIndices> img2ColTensorReassocIndices = { - {0, 1}, {2, 3}, {4, 5}}; - SmallVector<ReassociationIndices> filterReassociationIndice = {{0}, {1, 2}}; - SmallVector<ReassociationIndices> outputReassociationIndice = {{0, 1}, - {2, 3}}; - - auto reshapedImg2ColTensorType = RankedTensorType::get( - {n * c, oh * ow, fh * fw}, inputType.getElementType()); - auto reshapedFilterTensorType = - RankedTensorType::get({c, fh * fw}, filterType.getElementType()); - auto reshapedOutputTensorType = - RankedTensorType::get({n * c, oh * ow}, outputType.getElementType()); - - Value reshapedImg2ColTensor = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedImg2ColTensorType, img2ColTensor.getResult(0), - img2ColTensorReassocIndices); - Value reshapedFilterTensor = - tensor::CollapseShapeOp::create(rewriter, loc, reshapedFilterTensorType, - filterT, filterReassociationIndice); - Value reshapedoutputTensor = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedOutputTensorType, transposedOutputTensor, - outputReassociationIndice); - - auto batchMatVecResult = linalg::BatchMatvecOp::create( - rewriter, loc, TypeRange{reshapedoutputTensor.getType()}, - ValueRange{reshapedImg2ColTensor, reshapedFilterTensor}, - ValueRange{reshapedoutputTensor}); - - SmallVector<ReassociationIndices> batchMatVecReassociationIndice = {{0, 1}, - {2, 3}}; - - Value batchMatVecResultReshaped = tensor::ExpandShapeOp::create( - rewriter, loc, transposedOutputTensor.getType(), - batchMatVecResult.getResult(0), batchMatVecReassociationIndice); - - auto transposedResult = - transposeOperand(batchMatVecResultReshaped, {0, 2, 3, 1}); - - rewriter.replaceOp(convOp, ArrayRef<Value>{transposedResult}); - return success(); - } -}; - -// For nchw, because the channels are to the left of the image shape dimensions, -// the position of the contraction dimension in the resulting matmul is -// reversed. This swaps the LHS and RHS of the matmul when compared with nhwc -// (i.e. (D, C x Kh x Kw) * (C x Kh x Kw, Ho x Wo)) -class ConvertConv2DNchwFchw final - : public OpRewritePattern<linalg::Conv2DNchwFchwOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::Conv2DNchwFchwOp convOp, - PatternRewriter &rewriter) const override { - auto inputType = cast<ShapedType>(convOp.getInputs()[0].getType()); - auto filterType = cast<ShapedType>(convOp.getInputs()[1].getType()); - auto outputType = cast<ShapedType>(convOp.getOutputs()[0].getType()); - - if (!filterType.hasStaticShape() || !inputType.hasStaticShape()) { - return rewriter.notifyMatchFailure(convOp, [](Diagnostic &diag) { - diag << "[unimplemented] " - << "expected 'filterType' and 'inputType' to have static shapes."; - }); - } - - // TODO: Support dilation. - if (!hasAllOneValues(convOp.getDilations())) { - return rewriter.notifyMatchFailure(convOp, [](Diagnostic &diag) { - diag << "[unimplemented] " - << "expected no dilations (expected dilations to all be one)."; - }); - } - - Value input = convOp.getInputs()[0]; - Value filter = convOp.getInputs()[1]; - Value output = convOp.getOutputs()[0]; - - auto filterShape = filterType.getShape(); - auto outputShape = outputType.getShape(); - - const int n = outputShape[0]; - const int oc = outputShape[1]; - const int oh = outputShape[2]; - const int ow = outputShape[3]; - const int ic = filterShape[1]; - const int fh = filterShape[2]; - const int fw = filterShape[3]; - - auto loc = convOp.getLoc(); - - SmallVector<int64_t> colTensorShape = {n, ic, fh, fw, oh, ow}; - - Value colTensor = tensor::EmptyOp::create(rewriter, loc, colTensorShape, - inputType.getElementType()); - - AffineExpr nDim, icDim, khDim, kwDim, ohDim, owDim; - bindDims(getContext(), nDim, icDim, khDim, kwDim, ohDim, owDim); - - auto shSym = rewriter.getAffineConstantExpr( - convOp.getStrides().getValues<int64_t>()[0]); - auto swSym = rewriter.getAffineConstantExpr( - convOp.getStrides().getValues<int64_t>()[1]); - - SmallVector<AffineExpr> inputExprs = {nDim, icDim, ohDim * shSym + khDim, - owDim * swSym + kwDim}; - - auto nloops = colTensorShape.size(); - - auto parallel = utils::IteratorType::parallel; - auto reduction = utils::IteratorType::reduction; - SmallVector<utils::IteratorType, 3> img2colIterators(nloops, parallel); - - SmallVector<AffineMap> img2colIndexingMaps = { - AffineMap::get(nloops, 0, inputExprs, rewriter.getContext()), - AffineMap::getMultiDimIdentityMap(nloops, rewriter.getContext())}; - - auto img2ColTensor = linalg::GenericOp::create( - rewriter, loc, colTensor.getType(), - /*inputs=*/input, /*outputs=*/colTensor, img2colIndexingMaps, - img2colIterators, - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { - linalg::YieldOp::create(nestedBuilder, nestedLoc, args[0]); - }); - - SmallVector<ReassociationIndices> filterReassocIndices = {{0}, {1, 2, 3}}; - auto reshapedFilterType = - RankedTensorType::get({oc, fh * fw * ic}, inputType.getElementType()); - Value reshapedFilter = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedFilterType, filter, filterReassocIndices); - - SmallVector<ReassociationIndices> img2ColTensorReassocIndices; - SmallVector<ReassociationIndices> outputReassocIndices; - RankedTensorType reshapedImg2ColTensorType, reshapedOutputType; - if (n == 1) { - img2ColTensorReassocIndices = {{0, 1, 2, 3}, {4, 5}}; - outputReassocIndices = {{0, 1}, {2, 3}}; - - reshapedImg2ColTensorType = RankedTensorType::get( - {fh * fw * ic, oh * ow}, inputType.getElementType()); - reshapedOutputType = - RankedTensorType::get({oc, oh * ow}, outputType.getElementType()); - } else { - img2ColTensorReassocIndices = {{0}, {1, 2, 3}, {4, 5}}; - outputReassocIndices = {{0}, {1}, {2, 3}}; - - reshapedImg2ColTensorType = RankedTensorType::get( - {n, fh * fw * ic, oh * ow}, inputType.getElementType()); - reshapedOutputType = - RankedTensorType::get({n, oc, oh * ow}, outputType.getElementType()); - } - - Value reshapedImg2ColTensor = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedImg2ColTensorType, img2ColTensor.getResult(0), - img2ColTensorReassocIndices); - - Value reshapedOutput = tensor::CollapseShapeOp::create( - rewriter, loc, reshapedOutputType, output, outputReassocIndices); - - Value result; - if (n == 1) { - auto matmulOp = linalg::MatmulOp::create( - rewriter, loc, reshapedOutputType, - ArrayRef<Value>{reshapedFilter, reshapedImg2ColTensor}, - ArrayRef<Value>{reshapedOutput}); - result = matmulOp.getResults().front(); - } else { - // For cases where batch is not 1, we need to keep the batch dimension - // separate. Because the filter does not share the same batch dimension, - // the batch dimension is only used in indexing the input and output. Thus - // we cannot use existing linalg named ops like linalg.batch_matmul. - // i.e. M x K * (B x) K x N = (B x) M x N - AffineExpr bDim, mDim, nDim, kDim; - bindDims(getContext(), bDim, mDim, nDim, kDim); - auto lhsMap = AffineMap::get(4, 0, {mDim, kDim}, getContext()); - auto rhsMap = AffineMap::get(4, 0, {bDim, kDim, nDim}, getContext()); - auto resultMap = AffineMap::get(4, 0, {bDim, mDim, nDim}, getContext()); - SmallVector<utils::IteratorType> genericIterators = {parallel, parallel, - parallel, reduction}; - bool isInt = isa<IntegerType>(outputType.getElementType()); - auto genericOp = linalg::GenericOp::create( - rewriter, loc, reshapedOutputType, - /*inputs=*/ValueRange{reshapedFilter, reshapedImg2ColTensor}, - /*outputs=*/ValueRange{reshapedOutput}, - ArrayRef<AffineMap>{lhsMap, rhsMap, resultMap}, genericIterators, - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { - Value mul = createMul(loc, args[0], args[1], isInt, nestedBuilder); - Value add = createAdd(loc, mul, args[2], isInt, nestedBuilder); - linalg::YieldOp::create(nestedBuilder, nestedLoc, add); - }); - result = genericOp.getResults().front(); - } - - auto reshapedResult = tensor::ExpandShapeOp::create( - rewriter, loc, outputType, result, outputReassocIndices); - - rewriter.replaceOp(convOp, ArrayRef<Value>{reshapedResult}); - - return success(); - } -}; - -class ConvertConv2DToImg2ColPass - : public impl::ConvertConv2DToImg2ColPassBase<ConvertConv2DToImg2ColPass> { - void runOnOperation() override { - MLIRContext *context = &getContext(); - RewritePatternSet patterns(&getContext()); - patterns.insert<ConvertConv2DNhwcHwcf, ConvertDepthwiseConv2DNhwcHwc, - ConvertConv2DNchwFchw>(context); - if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { - return signalPassFailure(); - } - } -}; - -} // namespace - -} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/ConvertConvToIm2Col.cpp b/compiler/src/iree/compiler/GlobalOptimization/ConvertConvToIm2Col.cpp new file mode 100644 index 0000000..9676ed0 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/ConvertConvToIm2Col.cpp
@@ -0,0 +1,288 @@ +// Copyright 2020 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/DispatchCreation/CollapseDimensions.h" +#include "iree/compiler/GlobalOptimization/Passes.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/Transforms.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/AffineExpr.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include <functional> + +namespace mlir::iree_compiler::GlobalOptimization { + +#define GEN_PASS_DEF_CONVERTCONVTOIM2COLPASS +#include "iree/compiler/GlobalOptimization/Passes.h.inc" // IWYU pragma: export + +namespace { + +/// Returns the loops of `convOp` that its input indexing map can use: depth +/// loops, then batch and output image loops, in the order they appear in the +/// output map, followed by input channel and filter window loops in the order +/// they appear in the filter map. +static SmallVector<unsigned> +getInputLoopOrder(linalg::LinalgOp convOp, + const linalg::ConvolutionDimensions &convDims) { + AffineMap filterMap = + convOp.getMatchingIndexingMap(convOp.getDpsInputOperand(/*i=*/1)); + AffineMap outputMap = + convOp.getMatchingIndexingMap(convOp.getDpsInitOperand(/*i=*/0)); + SmallVector<unsigned> loopOrder; + auto appendInMapOrder = [&](AffineMap map, ArrayRef<unsigned> loops) { + for (AffineExpr expr : map.getResults()) { + unsigned loop = cast<AffineDimExpr>(expr).getPosition(); + if (llvm::is_contained(loops, loop)) { + loopOrder.push_back(loop); + } + } + }; + appendInMapOrder(outputMap, convDims.depth); + appendInMapOrder(outputMap, llvm::to_vector(llvm::concat<const unsigned>( + convDims.batch, convDims.outputImage))); + appendInMapOrder(filterMap, llvm::to_vector(llvm::concat<const unsigned>( + convDims.inputChannel, convDims.filterLoop))); + return loopOrder; +} + +/// Returns the input indexing map of `convOp` with each loop in `filterLoops` +/// that has a static extent of one replaced by zero. A loop of extent one only +/// takes the value zero, so the returned map addresses the same input element +/// as the original one in every iteration of `convOp`. +static AffineMap getInputMapWithoutUnitWindows(linalg::LinalgOp convOp, + ArrayRef<unsigned> filterLoops) { + MLIRContext *context = convOp.getContext(); + unsigned numLoops = convOp.getNumLoops(); + SmallVector<int64_t> loopSizes = convOp.getStaticLoopRanges(); + SmallVector<AffineExpr> replacements = + llvm::map_to_vector(llvm::seq<unsigned>(numLoops), [&](unsigned loop) { + return getAffineDimExpr(loop, context); + }); + for (unsigned loop : filterLoops) { + if (loopSizes[loop] == 1) { + replacements[loop] = getAffineConstantExpr(0, context); + } + } + AffineMap inputMap = + convOp.getMatchingIndexingMap(convOp.getDpsInputOperand(0)); + return simplifyAffineMap(inputMap.replaceDimsAndSymbols( + replacements, /*symReplacements=*/{}, numLoops, /*numResultSyms=*/0)); +} + +/// A convolution input materialized by `gatherInput`. +struct GatheredInput { + // The gathered tensor. It has one dimension per loop the input access uses, + // in contraction loop order. + Value value; + // Indexing map from the convolution's loops to `value`: a projected + // permutation selecting the gathered loops. + AffineMap contractionMap; +}; + +/// Materializes the access `inputMap` makes into the input of `convOp`. +/// +/// The gather is an all-parallel copy whose iteration space consists of the +/// loops `inputMap` uses, ordered as in `loopOrder`. For an NCHW convolution +/// with `inputMap = (n, oc, oh, ow, c, kh, kw) -> (n, c, oh + kh, ow + kw)` and +/// loop order (n, oh, ow, c, kh, kw), the gather loops are +/// (n, oh, ow, c, kh, kw), and +/// - the gather reads the input with `inputMap` rewritten onto its own +/// loops: `(d0, d1, d2, d3, d4, d5) -> (d0, d3, d1 + d4, d2 + d5)`; +/// - the convolution then reads the gathered tensor with +/// `(n, oc, oh, ow, c, kh, kw) -> (n, oh, ow, c, kh, kw)`. +static GatheredInput gatherInput(RewriterBase &rewriter, + linalg::LinalgOp convOp, AffineMap inputMap, + ArrayRef<unsigned> loopOrder) { + MLIRContext *context = rewriter.getContext(); + Location loc = convOp.getLoc(); + Value input = convOp.getDpsInputs().front(); + SmallVector<OpFoldResult> loopSizes = + llvm::map_to_vector(convOp.createLoopRanges(rewriter, loc), + [](Range range) { return range.size; }); + + SmallVector<unsigned> gatherLoops = llvm::filter_to_vector( + loopOrder, [&](unsigned loop) { return inputMap.isFunctionOfDim(loop); }); + unsigned gatherRank = gatherLoops.size(); + + // `loopToGatherDim` renames each gathered loop to its gather dimension; the + // remaining loops do not occur in `inputMap`, so their placeholder entry is + // never used. `gatherDimToLoop` is the reverse selection. + SmallVector<AffineExpr> loopToGatherDim(convOp.getNumLoops(), + getAffineConstantExpr(0, context)); + SmallVector<AffineExpr> gatherDimToLoop; + SmallVector<OpFoldResult> gatherSizes; + for (auto [gatherDim, loop] : llvm::enumerate(gatherLoops)) { + loopToGatherDim[loop] = getAffineDimExpr(gatherDim, context); + gatherDimToLoop.push_back(getAffineDimExpr(loop, context)); + gatherSizes.push_back(loopSizes[loop]); + } + + AffineMap readMap = inputMap.replaceDimsAndSymbols( + loopToGatherDim, /*symReplacements=*/{}, gatherRank, + /*numResultSyms=*/0); + Value init = tensor::EmptyOp::create(rewriter, loc, gatherSizes, + getElementTypeOrSelf(input.getType())); + auto gatherOp = linalg::GenericOp::create( + rewriter, loc, init.getType(), input, init, + ArrayRef<AffineMap>{readMap, rewriter.getMultiDimIdentityMap(gatherRank)}, + SmallVector<utils::IteratorType>(gatherRank, + utils::IteratorType::parallel), + [](OpBuilder &b, Location nestedLoc, ValueRange args) { + linalg::YieldOp::create(b, nestedLoc, args[0]); + }); + return {gatherOp.getResult(0), + AffineMap::get(convOp.getNumLoops(), /*symbolCount=*/0, + gatherDimToLoop, context)}; +} + +/// Rewrites a convolution as a contraction over an explicitly gathered input. +/// +/// The convolved input access (e.g. `(n, c, oh * s + kh * d, ow * s + kw * d)`) +/// is the only part of a convolution that a contraction cannot express. This +/// pattern materializes it with an all-parallel copy indexed by the loops the +/// access uses, then rewires the convolution to read that copy through a +/// projected permutation. The payload is kept as is, so floating-point and +/// integer bodies, including operand extensions, are supported unchanged, and +/// any layout, stride, dilation, group or depth multiplier works without +/// per-layout handling. Loops that every operand keeps contiguous are then +/// collapsed, which turns the standard layouts into plain (batch) matmuls. +/// +/// Filter window loops of extent one are substituted with zero before deciding +/// what to gather. When that leaves the input access a projected permutation, +/// as for an unstrided 1x1 convolution, the input is read in place. +/// +/// The pattern only checks that the rewrite is valid; whether it is worth +/// applying to a given convolution is decided by `controlFn`. +/// +/// TODO: Upstream this pattern and `getCollapsibleLoops` to MLIR, where they +/// would generalize `linalg::populateConvertConv2DToImg2ColPatterns` beyond +/// its four named 2-D convolutions. +struct ConvertConvToIm2Col final : OpInterfaceRewritePattern<linalg::LinalgOp> { + using ControlFn = std::function<bool(linalg::LinalgOp, + const linalg::ConvolutionDimensions &)>; + + ConvertConvToIm2Col(MLIRContext *context, ControlFn controlFn) + : OpInterfaceRewritePattern(context), controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp, + PatternRewriter &rewriter) const override { + if (!linalg::isaConvolutionOpInterface(linalgOp)) { + return rewriter.notifyMatchFailure(linalgOp, "not a convolution"); + } + // The gather is built on tensors, and the rewrite moves the op's body + // region into the contraction. + if (!linalgOp.hasPureTensorSemantics() || linalgOp->getNumRegions() != 1) { + return rewriter.notifyMatchFailure( + linalgOp, "expected tensor semantics and a body region"); + } + // A preset compilation info or lowering config refers to the loops of the + // original op, which the rewrite replaces. + if (getCompilationInfo(linalgOp) || getLoweringConfig(linalgOp)) { + return rewriter.notifyMatchFailure( + linalgOp, "has a preset compilation info or lowering config"); + } + // Requiring exactly an input and a filter excludes quantized convolutions + // with zero-point operands; they are lowered before im2col runs. + FailureOr<linalg::ConvolutionDimensions> convDims = + linalg::inferConvolutionDims(linalgOp); + if (failed(convDims)) { + return rewriter.notifyMatchFailure(linalgOp, + "failed to infer convolution dims"); + } + if (!controlFn(linalgOp, *convDims)) { + return rewriter.notifyMatchFailure(linalgOp, + "rejected by the control function"); + } + + AffineMap inputMap = + getInputMapWithoutUnitWindows(linalgOp, convDims->filterLoop); + SmallVector<unsigned> loopOrder = getInputLoopOrder(linalgOp, *convDims); + + // Gather the input unless its access is already a projected permutation. + Value newInput = linalgOp.getDpsInputs().front(); + AffineMap newInputMap = inputMap; + if (!inputMap.isProjectedPermutation()) { + GatheredInput gathered = + gatherInput(rewriter, linalgOp, inputMap, loopOrder); + newInput = gathered.value; + newInputMap = gathered.contractionMap; + } + + // Rebuild the op as a generic around the new input. Named and generic + // convolutions alike keep their body region, so the payload is moved + // unchanged. + SmallVector<Value> inputs = linalgOp.getDpsInputs(); + inputs.front() = newInput; + SmallVector<AffineMap> contractionMaps = linalgOp.getIndexingMapsArray(); + contractionMaps.front() = newInputMap; + auto contractionOp = linalg::GenericOp::create( + rewriter, linalgOp.getLoc(), linalgOp->getResultTypes(), inputs, + linalgOp.getDpsInits(), contractionMaps, + linalgOp.getIteratorTypesArray()); + rewriter.inlineRegionBefore(linalgOp->getRegion(0), + contractionOp.getRegion(), + contractionOp.getRegion().end()); + rewriter.replaceOp(linalgOp, contractionOp->getResults()); + + // Collapse the contraction's contiguous loops, so standard layouts become + // (batch) matmuls with one loop per dimension kind. The reshapes this + // introduces also keep the gather from being fused back into the + // contraction, which would recreate the convolution. + SmallVector<ReassociationIndices> collapsibleLoops = + DispatchCreation::getCollapsibleLoops(contractionOp); + if (collapsibleLoops.empty()) { + return success(); + } + // The loops are contiguous in every operand by construction, which is + // the only precondition of collapsing. + FailureOr<linalg::CollapseResult> collapsed = + linalg::collapseOpIterationDims(contractionOp, collapsibleLoops, + rewriter); + assert(succeeded(collapsed) && "collapsible groups are not contiguous"); + rewriter.replaceOp(contractionOp, collapsed->results); + return success(); + } + +private: + // Decides which valid convolutions are rewritten. + ControlFn controlFn; +}; + +/// Returns true if im2col can pay off for `convOp`, independent of the target. +/// The gather copies every input element once per filter window position it +/// falls into, and the contraction only amortizes that copy by reusing each +/// gathered element across output channels. Without an output channel loop of +/// extent other than one, i.e. for depthwise convolutions with a unit +/// multiplier and for pooling, nothing reuses the gather and the rewrite only +/// adds memory traffic. A dynamic output channel extent may exceed one, so it +/// counts as reuse. +static bool +hasOutputChannelReuse(linalg::LinalgOp convOp, + const linalg::ConvolutionDimensions &convDims) { + SmallVector<int64_t> loopSizes = convOp.getStaticLoopRanges(); + return llvm::any_of(convDims.outputChannel, + [&](unsigned loop) { return loopSizes[loop] != 1; }); +} + +struct ConvertConvToIm2ColPass final + : impl::ConvertConvToIm2ColPassBase<ConvertConvToIm2ColPass> { + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + patterns.add<ConvertConvToIm2Col>(context, hasOutputChannelReuse); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { + return signalPassFailure(); + } + } +}; + +} // namespace + +} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp index bd99e43..d3bbde2 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp +++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -138,7 +138,7 @@ .addPredicatedPass(clConvertQDQToIntegerMath, createConvertQDQToIntegerMathPass) .addPredicatedPass(transformOptions.useIm2colForConvs, - createConvertConv2DToImg2ColPass) + createConvertConvToIm2ColPass) .addPass(IREE::Flow::createCanonicalizePass) .addPass(createRemoveZeroExtentTensorsPass) .addPass(createDetachElementwiseFromNamedOpsPass)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td index 6289471..0c807b5 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/Passes.td +++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -24,9 +24,23 @@ ]; } -def ConvertConv2DToImg2ColPass : - Pass<"iree-global-opt-convert-conv2d-to-img2col", ""> { - let summary = "Convert linalg convolution ops to matmul img2col based implementation"; +def ConvertConvToIm2ColPass : + Pass<"iree-global-opt-convert-conv-to-im2col", ""> { + let summary = "Rewrites convolutions as contractions over an im2col gather"; + let description = [{ + Rewrites every linalg op that implements a convolution, named or generic, + as an all-parallel `linalg.generic` that gathers the convolved input + windows, followed by a contraction over the gathered tensor with the + original payload. Contraction loops are collapsed wherever the operand + layouts allow, so standard layouts become plain or batched matmuls. + Convolutions without an output channel dimension, i.e. depthwise + convolutions with a unit multiplier and pooling ops, are left unchanged: + nothing reuses the gathered input, so the rewrite never pays off. + + The pass expects quantization to be lowered already: quantized + convolutions such as `linalg.conv_2d_nhwc_hwcf_q` are not matched, and + dequantize producers are not looked through. + }]; let dependentDialects = [ "mlir::linalg::LinalgDialect", "mlir::tensor::TensorDialect",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel index ee4b4aa..31107d5 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -18,7 +18,7 @@ # keep sorted [ "cleanup_numeric_narrowing.mlir", - "conv2d_to_img2col.mlir", + "conv_to_im2col.mlir", "convert_batch_matmul_to_matmul.mlir", "convert_qdq_to_integer_math_algebra.mlir", "convert_qdq_to_integer_math_maps.mlir",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt index 3773729..8c5654a 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -15,7 +15,7 @@ lit SRCS "cleanup_numeric_narrowing.mlir" - "conv2d_to_img2col.mlir" + "conv_to_im2col.mlir" "convert_batch_matmul_to_matmul.mlir" "convert_qdq_to_integer_math_algebra.mlir" "convert_qdq_to_integer_math_maps.mlir"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/conv2d_to_img2col.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/conv2d_to_img2col.mlir deleted file mode 100644 index 33a7a14..0000000 --- a/compiler/src/iree/compiler/GlobalOptimization/test/conv2d_to_img2col.mlir +++ /dev/null
@@ -1,184 +0,0 @@ -// RUN: iree-opt --split-input-file --iree-global-opt-convert-conv2d-to-img2col %s | FileCheck %s - -func.func @conv_16433136(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> { - %0 = linalg.conv_2d_nhwc_hwcf - {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> } - ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<3x3x4x16xf32>) - outs(%arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> - return %0 : tensor<1x14x14x16xf32> -} -// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d3, d2 + d4, d5)> -// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)> -// CHECK: @conv_16433136 -// CHECK: %[[INPUT:.+]]: tensor<1x16x16x4xf32> -// CHECK: %[[FILTER:.+]]: tensor<3x3x4x16xf32> -// CHECK: %[[OUTPUT:.+]]: tensor<1x14x14x16xf32> -// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x14x14x3x3x4xf32> -// CHECK: %[[COL_TENSOR:.+]] = linalg.generic -// CHECK-SAME: #[[MAP0]] -// CHECK-SAME: #[[MAP1]] -// CHECK: ^bb0(%[[IN_DATA:.+]]: f32, %[[OUT_DATA:.+]]: f32) -// CHECK: linalg.yield %[[IN_DATA]] : f32 -// CHECK-DAG: %[[RESHAPED_INIT_COL_TENSOR:.+]] = tensor.collapse_shape %[[COL_TENSOR]] -// CHECK-SAME: [0, 1, 2], [3, 4, 5] -// CHECK-SAME: tensor<1x14x14x3x3x4xf32> into tensor<196x36xf32> -// CHECK-DAG: %[[RESHAPED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] -// CHECK-SAME: [0, 1, 2], [3] -// CHECK-SAME: tensor<3x3x4x16xf32> into tensor<36x16xf32> -// CHECK-DAG: %[[RESHAPED_OUTPUT:.+]] = tensor.collapse_shape %[[OUTPUT]] -// CHECK-SAME: [0, 1, 2], [3] -// CHECK: %[[MATMUL_RESULT:.+]] = linalg.matmul ins(%[[RESHAPED_INIT_COL_TENSOR]], %[[RESHAPED_FILTER]] : tensor<196x36xf32>, tensor<36x16xf32>) outs(%[[RESHAPED_OUTPUT]] : tensor<196x16xf32>) -// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0, 1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<196x16xf32> into tensor<1x14x14x16xf32> -// CHECK: return %[[RESULT]] - -// ----- - -func.func @depthwise_conv_hwc_114x16x3(%input: tensor<1x114x114x16xf32>, %filter: tensor<3x3x16xf32>, %output: tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32> { - %0 = linalg.depthwise_conv_2d_nhwc_hwc { - dilations = dense<1> : tensor<2xi64>, - strides = dense<1> : tensor<2xi64> - } ins(%input, %filter : tensor<1x114x114x16xf32>, tensor<3x3x16xf32>) outs(%output : tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32> - return %0 : tensor<1x112x112x16xf32> -} - -// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)> -// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> -// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)> -// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> -// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2 + d4, d3 + d5)> -// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)> -// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1, d2)> -// CHECK: @depthwise_conv_hwc_114x16x3 -// CHECK-SAME: %[[INPUT:.+]]: tensor<1x114x114x16xf32> -// CHECK-SAME: %[[FILTER:.+]]: tensor<3x3x16xf32> -// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x112x112x16xf32> -// CHECK: %[[INPUT_T_INIT:.+]] = tensor.empty() : tensor<1x16x114x114xf32> -// CHECK: %[[INPUT_T:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[INPUT]] : tensor<1x114x114x16xf32>) outs(%[[INPUT_T_INIT]] : tensor<1x16x114x114xf32>) { -// CHECK-NEXT: ^bb0(%[[ARG3:.+]]: f32, %[[ARG4:.+]]: f32): -// CHECK-NEXT: linalg.yield %[[ARG3]] : f32 -// CHECK-NEXT: } -> tensor<1x16x114x114xf32> -// CHECK: %[[FILTER_T_INIT:.+]] = tensor.empty() : tensor<16x3x3xf32> -// CHECK: %[[FILTER_T:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[FILTER]] : tensor<3x3x16xf32>) outs(%[[FILTER_T_INIT]] : tensor<16x3x3xf32>) { -// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32): -// CHECK: linalg.yield -// CHECK: } -> tensor<16x3x3xf32> -// CHECK: %[[INIT_OUTPUT_TENSOR:.+]] = tensor.empty() : tensor<1x16x112x112xf32> -// CHECK: %[[OUTPUT_T:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[OUTPUT]] : tensor<1x112x112x16xf32>) outs(%[[INIT_OUTPUT_TENSOR]] : tensor<1x16x112x112xf32>) { -// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32): -// CHECK-NEXT: linalg.yield -// CHECK-NEXT: } -> tensor<1x16x112x112xf32> -// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x16x112x112x3x3xf32> -// CHECK: %[[COL_TENSOR:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP4]], #[[MAP5]]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[INPUT_T]] : tensor<1x16x114x114xf32>) outs(%[[INIT_COL_TENSOR]] : tensor<1x16x112x112x3x3xf32>) { -// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32): -// CHECK-NEXT: linalg.yield -// CHECK-NEXT: } -> tensor<1x16x112x112x3x3xf32> -// CHECK: %[[COL_TENSOR_R:.+]] = tensor.collapse_shape %[[COL_TENSOR]] -// CHECK-SAME: tensor<1x16x112x112x3x3xf32> into tensor<16x12544x9xf32> -// CHECK: %[[FILTER_T_R:.+]] = tensor.collapse_shape %[[FILTER_T]] -// CHECK-SAME: tensor<16x3x3xf32> into tensor<16x9xf32> -// CHECK: %[[OUTPUT_T_R:.+]] = tensor.collapse_shape %[[OUTPUT_T]] -// CHECK-SAME: tensor<1x16x112x112xf32> into tensor<16x12544xf32> -// CHECK: %[[BMV_RESULT:.+]] = linalg.batch_matvec ins(%[[COL_TENSOR_R]], %[[FILTER_T_R]] : tensor<16x12544x9xf32>, tensor<16x9xf32>) outs(%[[OUTPUT_T_R]] : tensor<16x12544xf32>) -> tensor<16x12544xf32> -// CHECK: %[[RESULT_R:.+]] = tensor.expand_shape %[[BMV_RESULT]] -// CHECK-SAME: tensor<16x12544xf32> into tensor<1x16x112x112xf32> -// CHECK: %[[RESULT_INIT:.+]] = tensor.empty() : tensor<1x112x112x16xf32> -// CHECK: %[[RESULT:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP6]], #[[MAP1]]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[RESULT_R]] : tensor<1x16x112x112xf32>) outs(%[[RESULT_INIT]] : tensor<1x112x112x16xf32>) { -// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32): -// CHECK-NEXT: linalg.yield -// CHECK-NEXT: } -> tensor<1x112x112x16xf32> -// CHECK: return %[[RESULT]] : tensor<1x112x112x16xf32> - -// ----- - -func.func @batch_nhwc_conv(%arg0: tensor<8x16x16x4xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<8x14x14x16xf32>) -> tensor<8x14x14x16xf32> { - %0 = linalg.conv_2d_nhwc_hwcf - {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> } - ins(%arg0, %arg1: tensor<8x16x16x4xf32>, tensor<3x3x4x16xf32>) - outs(%arg2: tensor<8x14x14x16xf32>) -> tensor<8x14x14x16xf32> - return %0 : tensor<8x14x14x16xf32> -} - -// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d3, d2 + d4, d5)> -// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)> -// CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> -// CHECK-DAG: #[[RHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> -// CHECK-DAG: #[[RESMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> - -// CHECK: func.func @batch_nhwc_conv -// CHECK-SAME: (%[[INPUT:.+]]: tensor<8x16x16x4xf32>, %[[FILTER:.+]]: tensor<3x3x4x16xf32>, %[[INIT:.+]]: tensor<8x14x14x16xf32>) -// CHECK: %[[IT:.+]] = tensor.empty() : tensor<8x14x14x3x3x4xf32> -// CHECK: %[[IMG2COL:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[INPUT]] : tensor<8x16x16x4xf32>) -// CHECK-SAME: outs(%[[IT]] : tensor<8x14x14x3x3x4xf32>) -// CHECK: %[[CS_INPUT:.+]] = tensor.collapse_shape %[[IMG2COL]] {{\[}}[0], [1, 2], [3, 4, 5]] : tensor<8x14x14x3x3x4xf32> into tensor<8x196x36xf32> -// CHECK: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2], [3]] : tensor<3x3x4x16xf32> into tensor<36x16xf32> -// CHECK: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2], [3]] : tensor<8x14x14x16xf32> into tensor<8x196x16xf32> -// CHECK: %[[MATMUL:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[LHSMAP]], #[[RHSMAP]], #[[RESMAP]]], -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] -// CHECK-SAME: ins(%[[CS_INPUT]], %[[CS_FILTER]] : tensor<8x196x36xf32>, tensor<36x16xf32>) -// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x196x16xf32>) -// CHECK: ^bb0(%[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: f32): -// CHECK: %[[MUL:.+]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32 -// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]], %[[ARG2]] : f32 -// CHECK: linalg.yield %[[ADD]] : f32 -// CHECK: } -> tensor<8x196x16xf32> -// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1, 2], [3]] output_shape [8, 14, 14, 16] : tensor<8x196x16xf32> into tensor<8x14x14x16xf32> -// CHECK: return %[[CS_FINAL]] - -// ----- - -func.func @batch_nchw_conv(%arg0: tensor<8x4x16x16xf32>, %arg1: tensor<16x4x3x3xf32>, %arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> { - %0 = linalg.conv_2d_nchw_fchw - {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> } - ins(%arg0, %arg1: tensor<8x4x16x16xf32>, tensor<16x4x3x3xf32>) - outs(%arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> - return %0 : tensor<8x16x14x14xf32> -} - -// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2 + d4, d3 + d5)> -// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)> -// CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> -// CHECK-DAG: #[[RHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)> -// CHECK-DAG: #[[RESMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> - -// CHECK: func.func @batch_nchw_conv -// CHECK-SAME: (%[[INPUT:.+]]: tensor<8x4x16x16xf32>, %[[FILTER:.+]]: tensor<16x4x3x3xf32>, %[[INIT:.+]]: tensor<8x16x14x14xf32>) -// CHECK: %[[IT:.+]] = tensor.empty() : tensor<8x4x3x3x14x14xf32> -// CHECK: %[[IMG2COL:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"] -// CHECK-SAME: ins(%[[INPUT]] : tensor<8x4x16x16xf32>) -// CHECK-SAME: outs(%[[IT]] : tensor<8x4x3x3x14x14xf32>) -// CHECK: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x4x3x3xf32> into tensor<16x36xf32> -// CHECK: %[[CS_INPUT:.+]] = tensor.collapse_shape %[[IMG2COL]] {{\[}}[0], [1, 2, 3], [4, 5]] : tensor<8x4x3x3x14x14xf32> into tensor<8x36x196xf32> -// CHECK: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]] : tensor<8x16x14x14xf32> into tensor<8x16x196xf32> -// CHECK: %[[MATMUL:.+]] = linalg.generic -// CHECK-SAME: indexing_maps = [#[[LHSMAP]], #[[RHSMAP]], #[[RESMAP]]], -// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] -// CHECK-SAME: ins(%[[CS_FILTER]], %[[CS_INPUT]] : tensor<16x36xf32>, tensor<8x36x196xf32>) -// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x16x196xf32>) -// CHECK: ^bb0(%[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: f32): -// CHECK: %[[MUL:.+]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32 -// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]], %[[ARG2]] : f32 -// CHECK: linalg.yield %[[ADD]] : f32 -// CHECK: } -> tensor<8x16x196xf32> -// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 14, 14] : tensor<8x16x196xf32> into tensor<8x16x14x14xf32> -// CHECK: return %[[CS_FINAL]]
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/conv_to_im2col.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/conv_to_im2col.mlir new file mode 100644 index 0000000..58430e0 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/test/conv_to_im2col.mlir
@@ -0,0 +1,540 @@ +// RUN: iree-opt --split-input-file --iree-global-opt-convert-conv-to-im2col %s | FileCheck %s + +// Channels-last: the batch and output image loops stay contiguous in the +// gathered input and the output, so the convolution becomes a plain matmul. +func.func @conv_nhwc_hwcf(%input: tensor<1x16x16x4xf32>, %filter: tensor<3x3x4x16xf32>, + %init: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> { + %0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x16x16x4xf32>, tensor<3x3x4x16xf32>) + outs(%init : tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> + return %0 : tensor<1x14x14x16xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d3, d2 + d4, d5)> +// CHECK-DAG: #[[$ID6:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @conv_nhwc_hwcf( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<1x16x16x4xf32> +// CHECK-SAME: %[[FILTER:[a-zA-Z0-9_]+]]: tensor<3x3x4x16xf32> +// CHECK-SAME: %[[INIT:[a-zA-Z0-9_]+]]: tensor<1x14x14x16xf32> +// CHECK: %[[GATHER_INIT:.+]] = tensor.empty() : tensor<1x14x14x3x3x4xf32> +// CHECK: %[[GATHER:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], #[[$ID6]]] +// CHECK-SAME: ins(%[[INPUT]] : tensor<1x16x16x4xf32>) outs(%[[GATHER_INIT]] : tensor<1x14x14x3x3x4xf32>) +// CHECK: %[[LHS_COLLAPSED:.+]] = tensor.collapse_shape %[[GATHER]] {{\[}}[0, 1, 2], [3, 4, 5]] : tensor<1x14x14x3x3x4xf32> into tensor<196x36xf32> +// CHECK: %[[RHS_COLLAPSED:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2], [3]] : tensor<3x3x4x16xf32> into tensor<36x16xf32> +// CHECK: %[[OUT_COLLAPSED:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1, 2], [3]] : tensor<1x14x14x16xf32> into tensor<196x16xf32> +// CHECK: %[[MATMUL:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[LHS_COLLAPSED]], %[[RHS_COLLAPSED]] : tensor<196x36xf32>, tensor<36x16xf32>) +// CHECK-SAME: outs(%[[OUT_COLLAPSED]] : tensor<196x16xf32>) +// CHECK: arith.mulf +// CHECK: arith.addf +// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1, 2], [3]] output_shape [1, 14, 14, 16] +// CHECK: return %[[RESULT]] + +// ----- + +// With an output-channel-first filter, the filter becomes the transposed +// matmul operand. +func.func @conv_nhwc_fhwc(%input: tensor<1x16x16x4xf32>, %filter: tensor<16x3x3x4xf32>, + %init: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> { + %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x16x16x4xf32>, tensor<16x3x3x4xf32>) + outs(%init : tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> + return %0 : tensor<1x14x14x16xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d3, d2 + d4, d5)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @conv_nhwc_fhwc( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<1x14x14x3x3x4xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<196x36xf32>, tensor<16x36xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<196x16xf32>) + +// ----- + +// Channels-first: the output channel separates the batch from the output image +// loops, so the batch stays a separate loop. The K loops follow the filter +// order, which keeps the filter collapsible without a transpose. +func.func @conv_nchw_fchw_batch(%input: tensor<8x4x16x16xf32>, %filter: tensor<16x4x3x3xf32>, + %init: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> { + %0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<8x4x16x16xf32>, tensor<16x4x3x3xf32>) + outs(%init : tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> + return %0 : tensor<8x16x14x14xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d3, d1 + d4, d2 + d5)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @conv_nchw_fchw_batch( +// CHECK: %[[GATHER:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<8x14x14x4x3x3xf32>) +// CHECK-DAG: tensor.collapse_shape %[[GATHER]] {{\[}}[0], [1, 2], [3, 4, 5]] : tensor<8x14x14x4x3x3xf32> into tensor<8x196x36xf32> +// CHECK-DAG: tensor.collapse_shape %{{.+}} {{\[}}[0], [1, 2, 3]] : tensor<16x4x3x3xf32> into tensor<16x36xf32> +// CHECK-DAG: tensor.collapse_shape %{{.+}} {{\[}}[0], [1], [2, 3]] : tensor<8x16x14x14xf32> into tensor<8x16x196xf32> +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] +// CHECK: tensor.expand_shape %{{.+}} {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 14, 14] + +// ----- + +// Strides and dilations only appear in the gather's read map. +func.func @conv_nchw_strided_dilated(%input: tensor<1x4x16x16xf32>, %filter: tensor<16x4x3x3xf32>, + %init: tensor<1x16x6x6xf32>) -> tensor<1x16x6x6xf32> { + %0 = linalg.conv_2d_nchw_fchw {dilations = dense<2> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} + ins(%input, %filter : tensor<1x4x16x16xf32>, tensor<16x4x3x3xf32>) + outs(%init : tensor<1x16x6x6xf32>) -> tensor<1x16x6x6xf32> + return %0 : tensor<1x16x6x6xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d3, d1 * 2 + d4 * 2, d2 * 2 + d5 * 2)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @conv_nchw_strided_dilated( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<1x6x6x4x3x3xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x36x36xf32>, tensor<16x36xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x16x36xf32>) + +// ----- + +// An unstrided 1x1 convolution reads its input in place, without a gather. +func.func @pointwise_conv_nchw(%input: tensor<1x32x8x8xf32>, %filter: tensor<16x32x1x1xf32>, + %init: tensor<1x16x8x8xf32>) -> tensor<1x16x8x8xf32> { + %0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x32x8x8xf32>, tensor<16x32x1x1xf32>) + outs(%init : tensor<1x16x8x8xf32>) -> tensor<1x16x8x8xf32> + return %0 : tensor<1x16x8x8xf32> +} +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @pointwise_conv_nchw( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<1x32x8x8xf32> +// CHECK-NOT: tensor.empty +// CHECK: %[[LHS:.+]] = tensor.collapse_shape %[[INPUT]] {{\[}}[0], [1], [2, 3]] : tensor<1x32x8x8xf32> into tensor<1x32x64xf32> +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%[[LHS]], %{{.+}} : tensor<1x32x64xf32>, tensor<16x32x1xf32>) + +// ----- + +// A strided 1x1 convolution still needs the gather to subsample its input. +func.func @strided_pointwise_conv_nchw(%input: tensor<1x32x8x8xf32>, %filter: tensor<16x32x1x1xf32>, + %init: tensor<1x16x4x4xf32>) -> tensor<1x16x4x4xf32> { + %0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} + ins(%input, %filter : tensor<1x32x8x8xf32>, tensor<16x32x1x1xf32>) + outs(%init : tensor<1x16x4x4xf32>) -> tensor<1x16x4x4xf32> + return %0 : tensor<1x16x4x4xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1 * 2, d2 * 2)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @strided_pointwise_conv_nchw( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<1x4x4x32xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x16x32xf32>, tensor<16x32x1xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x16x16xf32>) + +// ----- + +// Groups are batch loops of the contraction and come first. +func.func @grouped_conv_ngchw_gfchw(%input: tensor<1x2x4x10x10xf32>, %filter: tensor<2x8x4x3x3xf32>, + %init: tensor<1x2x8x8x8xf32>) -> tensor<1x2x8x8x8xf32> { + %0 = linalg.conv_2d_ngchw_gfchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x2x4x10x10xf32>, tensor<2x8x4x3x3xf32>) + outs(%init : tensor<1x2x8x8x8xf32>) -> tensor<1x2x8x8x8xf32> + return %0 : tensor<1x2x8x8x8xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d0, d4, d2 + d5, d3 + d6)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d0, d3, d4)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d4)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)> +// CHECK-LABEL: func.func @grouped_conv_ngchw_gfchw( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<2x1x8x8x4x3x3xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x1x64x36xf32>, tensor<2x8x36xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x2x8x64xf32>) + +// ----- + +// Channels-last groups also become the leading batch loop of the gather. +func.func @grouped_conv_nhwgc_gfhwc(%input: tensor<1x10x10x2x4xf32>, %filter: tensor<2x8x3x3x4xf32>, + %init: tensor<1x8x8x2x8xf32>) -> tensor<1x8x8x2x8xf32> { + %0 = linalg.conv_2d_nhwgc_gfhwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x10x10x2x4xf32>, tensor<2x8x3x3x4xf32>) + outs(%init : tensor<1x8x8x2x8xf32>) -> tensor<1x8x8x2x8xf32> + return %0 : tensor<1x8x8x2x8xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d2 + d4, d3 + d5, d0, d6)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @grouped_conv_nhwgc_gfhwc( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<2x1x8x8x3x3x4xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x64x36xf32>, tensor<2x8x36xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<64x2x8xf32>) + +// ----- + +// A channel multiplier is an output channel that reuses the gathered input, so +// depthwise convolutions with a multiplier are rewritten with the channel as a +// batch loop. +func.func @depthwise_conv_multiplier(%input: tensor<1x10x10x4xf32>, %filter: tensor<3x3x4x2xf32>, + %init: tensor<1x8x8x4x2xf32>) -> tensor<1x8x8x4x2xf32> { + %0 = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x10x10x4xf32>, tensor<3x3x4x2xf32>) + outs(%init : tensor<1x8x8x4x2xf32>) -> tensor<1x8x8x4x2xf32> + return %0 : tensor<1x8x8x4x2xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2 + d4, d3 + d5, d0)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @depthwise_conv_multiplier( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<4x1x8x8x3x3xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<4x64x9xf32>, tensor<9x4x2xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<64x4x2xf32>) + +// ----- + +// A unit channel multiplier still has an output channel loop, but nothing +// reuses the gathered input across it, so the convolution is left alone. +func.func @depthwise_conv_unit_multiplier(%input: tensor<1x10x10x4xf32>, %filter: tensor<3x3x4x1xf32>, + %init: tensor<1x8x8x4x1xf32>) -> tensor<1x8x8x4x1xf32> { + %0 = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x10x10x4xf32>, tensor<3x3x4x1xf32>) + outs(%init : tensor<1x8x8x4x1xf32>) -> tensor<1x8x8x4x1xf32> + return %0 : tensor<1x8x8x4x1xf32> +} +// CHECK-LABEL: func.func @depthwise_conv_unit_multiplier( +// CHECK: linalg.depthwise_conv_2d_nhwc_hwcm +// CHECK-NOT: linalg.generic + +// ----- + +// The gather copies the narrow input type; the extensions stay in the +// contraction payload. +func.func @conv_mixed_precision(%input: tensor<1x4x10x10xbf16>, %filter: tensor<16x4x3x3xbf16>, + %init: tensor<1x16x8x8xf32>) -> tensor<1x16x8x8xf32> { + %0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x4x10x10xbf16>, tensor<16x4x3x3xbf16>) + outs(%init : tensor<1x16x8x8xf32>) -> tensor<1x16x8x8xf32> + return %0 : tensor<1x16x8x8xf32> +} +// CHECK-LABEL: func.func @conv_mixed_precision( +// CHECK: %[[GATHER_INIT:.+]] = tensor.empty() : tensor<1x8x8x4x3x3xbf16> +// CHECK: linalg.generic +// CHECK-SAME: outs(%[[GATHER_INIT]] : tensor<1x8x8x4x3x3xbf16>) +// CHECK: linalg.generic +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x64x36xbf16>, tensor<16x36xbf16>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x16x64xf32>) +// CHECK: ^bb0(%[[IN:.+]]: bf16, %[[WEIGHT:.+]]: bf16, %[[ACC:.+]]: f32): +// CHECK: %[[IN_F32:.+]] = arith.extf %[[IN]] : bf16 to f32 +// CHECK: %[[WEIGHT_F32:.+]] = arith.extf %[[WEIGHT]] : bf16 to f32 +// CHECK: %[[PRODUCT:.+]] = arith.mulf %[[IN_F32]], %[[WEIGHT_F32]] : f32 +// CHECK: arith.addf %[[ACC]], %[[PRODUCT]] : f32 + +// ----- + +// Generic convolutions are rewritten as well, with their payload unchanged. +#input_map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d1 + d4, d2 + d5)> +#filter_map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d3, d4, d5)> +#output_map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)> +func.func @generic_integer_conv_chw(%input: tensor<16x162x162xi8>, %filter: tensor<16x16x3x3xi8>, + %init: tensor<16x160x160xi32>) -> tensor<16x160x160xi32> { + %0 = linalg.generic {indexing_maps = [#input_map, #filter_map, #output_map], + iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} + ins(%input, %filter : tensor<16x162x162xi8>, tensor<16x16x3x3xi8>) outs(%init : tensor<16x160x160xi32>) { + ^bb0(%in: i8, %weight: i8, %acc: i32): + %in_i32 = arith.extsi %in : i8 to i32 + %weight_i32 = arith.extsi %weight : i8 to i32 + %product = arith.muli %in_i32, %weight_i32 : i32 + %sum = arith.addi %acc, %product : i32 + linalg.yield %sum : i32 + } -> tensor<16x160x160xi32> + return %0 : tensor<16x160x160xi32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d0 + d3, d1 + d4)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @generic_integer_conv_chw( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<160x160x16x3x3xi8>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<25600x144xi8>, tensor<16x144xi8>) +// CHECK-SAME: outs(%{{.+}} : tensor<16x25600xi32>) +// CHECK: ^bb0(%[[IN:.+]]: i8, %[[WEIGHT:.+]]: i8, %[[ACC:.+]]: i32): +// CHECK: %[[IN_I32:.+]] = arith.extsi %[[IN]] : i8 to i32 +// CHECK: %[[WEIGHT_I32:.+]] = arith.extsi %[[WEIGHT]] : i8 to i32 +// CHECK: %[[PRODUCT:.+]] = arith.muli %[[IN_I32]], %[[WEIGHT_I32]] : i32 +// CHECK: arith.addi %[[ACC]], %[[PRODUCT]] : i32 + +// ----- + +func.func @conv_1d_ncw_fcw(%input: tensor<2x4x18xf32>, %filter: tensor<8x4x3xf32>, + %init: tensor<2x8x8xf32>) -> tensor<2x8x8xf32> { + %0 = linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>} + ins(%input, %filter : tensor<2x4x18xf32>, tensor<8x4x3xf32>) + outs(%init : tensor<2x8x8xf32>) -> tensor<2x8x8xf32> + return %0 : tensor<2x8x8xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1 * 2 + d3)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @conv_1d_ncw_fcw( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<2x4x18xf32> +// CHECK-SAME: %[[FILTER:[a-zA-Z0-9_]+]]: tensor<8x4x3xf32> +// CHECK-SAME: %[[INIT:[a-zA-Z0-9_]+]]: tensor<2x8x8xf32> +// CHECK: %[[GATHER:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: ins(%[[INPUT]] : tensor<2x4x18xf32>) outs(%{{.+}} : tensor<2x8x4x3xf32>) +// CHECK: %[[LHS_COLLAPSED:.+]] = tensor.collapse_shape %[[GATHER]] {{\[}}[0], [1], [2, 3]] : tensor<2x8x4x3xf32> into tensor<2x8x12xf32> +// CHECK: %[[RHS_COLLAPSED:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2]] : tensor<8x4x3xf32> into tensor<8x12xf32> +// CHECK: %[[RESULT:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[LHS_COLLAPSED]], %[[RHS_COLLAPSED]] : tensor<2x8x12xf32>, tensor<8x12xf32>) +// CHECK-SAME: outs(%[[INIT]] : tensor<2x8x8xf32>) +// CHECK: return %[[RESULT]] + +// ----- + +// Dynamic loop extents size the gather from the operands. +func.func @conv_1d_dynamic(%input: tensor<?x?x4xf32>, %filter: tensor<3x4x8xf32>, + %init: tensor<?x?x8xf32>) -> tensor<?x?x8xf32> { + %0 = linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} + ins(%input, %filter : tensor<?x?x4xf32>, tensor<3x4x8xf32>) + outs(%init : tensor<?x?x8xf32>) -> tensor<?x?x8xf32> + return %0 : tensor<?x?x8xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 + d2, d3)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @conv_1d_dynamic( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<?x?x4xf32> +// CHECK-SAME: %[[FILTER:[a-zA-Z0-9_]+]]: tensor<3x4x8xf32> +// CHECK-SAME: %[[INIT:[a-zA-Z0-9_]+]]: tensor<?x?x8xf32> +// CHECK-DAG: %[[BATCH:.+]] = tensor.dim %[[INPUT]], %c0 +// CHECK-DAG: %[[WIDTH:.+]] = tensor.dim %[[INIT]], %c1 +// CHECK: %[[GATHER_INIT:.+]] = tensor.empty(%[[BATCH]], %[[WIDTH]]) : tensor<?x?x3x4xf32> +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%[[GATHER_INIT]] : tensor<?x?x3x4xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x12xf32>, tensor<12x8xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<?x8xf32>) + +// ----- + +// A channels-last 3x3x3 convolution gathers all three window loops and becomes +// a plain matmul. +func.func @conv_3d_ndhwc_dhwcf(%input: tensor<1x6x10x10x4xf32>, %filter: tensor<3x3x3x4x16xf32>, + %init: tensor<1x4x8x8x16xf32>) -> tensor<1x4x8x8x16xf32> { + %0 = linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} + ins(%input, %filter : tensor<1x6x10x10x4xf32>, tensor<3x3x3x4x16xf32>) + outs(%init : tensor<1x4x8x8x16xf32>) -> tensor<1x4x8x8x16xf32> + return %0 : tensor<1x4x8x8x16xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1 + d4, d2 + d5, d3 + d6, d7)> +// CHECK-DAG: #[[$ID8:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> +// CHECK-LABEL: func.func @conv_3d_ndhwc_dhwcf( +// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<1x6x10x10x4xf32> +// CHECK-SAME: %[[FILTER:[a-zA-Z0-9_]+]]: tensor<3x3x3x4x16xf32> +// CHECK-SAME: %[[INIT:[a-zA-Z0-9_]+]]: tensor<1x4x8x8x16xf32> +// CHECK: %[[GATHER_INIT:.+]] = tensor.empty() : tensor<1x4x8x8x3x3x3x4xf32> +// CHECK: %[[GATHER:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], #[[$ID8]]] +// CHECK-SAME: ins(%[[INPUT]] : tensor<1x6x10x10x4xf32>) outs(%[[GATHER_INIT]] : tensor<1x4x8x8x3x3x3x4xf32>) +// CHECK: %[[LHS_COLLAPSED:.+]] = tensor.collapse_shape %[[GATHER]] {{\[}}[0, 1, 2, 3], [4, 5, 6, 7]] : tensor<1x4x8x8x3x3x3x4xf32> into tensor<256x108xf32> +// CHECK: %[[RHS_COLLAPSED:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2, 3], [4]] : tensor<3x3x3x4x16xf32> into tensor<108x16xf32> +// CHECK: %[[OUT_COLLAPSED:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1, 2, 3], [4]] : tensor<1x4x8x8x16xf32> into tensor<256x16xf32> +// CHECK: %[[MATMUL:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] +// CHECK-SAME: ins(%[[LHS_COLLAPSED]], %[[RHS_COLLAPSED]] : tensor<256x108xf32>, tensor<108x16xf32>) +// CHECK-SAME: outs(%[[OUT_COLLAPSED]] : tensor<256x16xf32>) +// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1, 2, 3], [4]] output_shape [1, 4, 8, 8, 16] +// CHECK: return %[[RESULT]] + +// ----- + +// Channels-first 3x3x3: the K loops follow the filter order, so the filter +// collapses without a transpose and the batch stays a separate loop. +func.func @conv_3d_ncdhw_fcdhw(%input: tensor<1x4x6x10x10xf32>, %filter: tensor<16x4x3x3x3xf32>, + %init: tensor<1x16x4x8x8xf32>) -> tensor<1x16x4x8x8xf32> { + %0 = linalg.conv_3d_ncdhw_fcdhw {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} + ins(%input, %filter : tensor<1x4x6x10x10xf32>, tensor<16x4x3x3x3xf32>) + outs(%init : tensor<1x16x4x8x8xf32>) -> tensor<1x16x4x8x8xf32> + return %0 : tensor<1x16x4x8x8xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d4, d1 + d5, d2 + d6, d3 + d7)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @conv_3d_ncdhw_fcdhw( +// CHECK: %[[GATHER:.+]] = linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<1x4x8x8x4x3x3x3xf32>) +// CHECK-DAG: tensor.collapse_shape %[[GATHER]] {{\[}}[0], [1, 2, 3], [4, 5, 6, 7]] : tensor<1x4x8x8x4x3x3x3xf32> into tensor<1x256x108xf32> +// CHECK-DAG: tensor.collapse_shape %{{.+}} {{\[}}[0], [1, 2, 3, 4]] : tensor<16x4x3x3x3xf32> into tensor<16x108xf32> +// CHECK-DAG: tensor.collapse_shape %{{.+}} {{\[}}[0], [1], [2, 3, 4]] : tensor<1x16x4x8x8xf32> into tensor<1x16x256xf32> +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] +// CHECK: tensor.expand_shape %{{.+}} {{\[}}[0], [1], [2, 3, 4]] output_shape [1, 16, 4, 8, 8] + +// ----- + +// Only the depth window loop has a non-unit extent, so it is the only window +// loop gathered. The unit window loops remain as a unit reduction loop of the +// contraction. +func.func @conv_3d_unit_spatial_window(%input: tensor<1x4x6x10x10xf32>, %filter: tensor<16x4x3x1x1xf32>, + %init: tensor<1x16x4x10x10xf32>) -> tensor<1x16x4x10x10xf32> { + %0 = linalg.conv_3d_ncdhw_fcdhw {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} + ins(%input, %filter : tensor<1x4x6x10x10xf32>, tensor<16x4x3x1x1xf32>) + outs(%init : tensor<1x16x4x10x10xf32>) -> tensor<1x16x4x10x10xf32> + return %0 : tensor<1x16x4x10x10xf32> +} +// CHECK-DAG: #[[$GATHER_READ:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d4, d1 + d5, d2, d3)> +// CHECK-DAG: #[[$LHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d3)> +// CHECK-DAG: #[[$RHS:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4)> +// CHECK-DAG: #[[$OUT:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)> +// CHECK-LABEL: func.func @conv_3d_unit_spatial_window( +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$GATHER_READ]], +// CHECK-SAME: outs(%{{.+}} : tensor<1x4x10x10x4x3xf32>) +// CHECK: linalg.generic +// CHECK-SAME: indexing_maps = [#[[$LHS]], #[[$RHS]], #[[$OUT]]] +// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"] +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x400x12xf32>, tensor<16x12x1xf32>) +// CHECK-SAME: outs(%{{.+}} : tensor<1x16x400xf32>) + +// ----- + +// A depthwise convolution has no output channel to reuse the gathered input +// across, so it is left alone. +func.func @depthwise_conv_nhwc_hwc(%input: tensor<1x114x114x16xf32>, %filter: tensor<3x3x16xf32>, + %init: tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32> { + %0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x114x114x16xf32>, tensor<3x3x16xf32>) + outs(%init : tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32> + return %0 : tensor<1x112x112x16xf32> +} +// CHECK-LABEL: func.func @depthwise_conv_nhwc_hwc( +// CHECK: linalg.depthwise_conv_2d_nhwc_hwc +// CHECK-NOT: linalg.generic + +// ----- + +// Pooling ops have no output channel either and are left alone. +func.func @pooling_nhwc_sum(%input: tensor<1x16x16x4xf32>, %window: tensor<3x3xf32>, + %init: tensor<1x14x14x4xf32>) -> tensor<1x14x14x4xf32> { + %0 = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %window : tensor<1x16x16x4xf32>, tensor<3x3xf32>) + outs(%init : tensor<1x14x14x4xf32>) -> tensor<1x14x14x4xf32> + return %0 : tensor<1x14x14x4xf32> +} +// CHECK-LABEL: func.func @pooling_nhwc_sum( +// CHECK: linalg.pooling_nhwc_sum +// CHECK-NOT: linalg.generic + +// ----- + +// Quantized convolutions carry zero-point operands and are expected to be +// lowered before im2col runs. +func.func @quantized_conv_nhwc_hwcf(%input: tensor<1x10x10x4xi8>, %filter: tensor<3x3x4x16xi8>, + %input_zp: i32, %filter_zp: i32, + %init: tensor<1x8x8x16xi32>) -> tensor<1x8x8x16xi32> { + %0 = linalg.conv_2d_nhwc_hwcf_q {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter, %input_zp, %filter_zp : tensor<1x10x10x4xi8>, tensor<3x3x4x16xi8>, i32, i32) + outs(%init : tensor<1x8x8x16xi32>) -> tensor<1x8x8x16xi32> + return %0 : tensor<1x8x8x16xi32> +} +// CHECK-LABEL: func.func @quantized_conv_nhwc_hwcf( +// CHECK: linalg.conv_2d_nhwc_hwcf_q +// CHECK-NOT: linalg.generic + +// ----- + +// A preset lowering config refers to the original loops, so configured +// convolutions are left alone. +#input_map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d1 + d4, d2 + d5)> +#filter_map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d3, d4, d5)> +#output_map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)> +func.func @generic_conv_lowering_config(%input: tensor<4x10x10xf32>, %filter: tensor<16x4x3x3xf32>, + %init: tensor<16x8x8xf32>) -> tensor<16x8x8xf32> { + %0 = linalg.generic {indexing_maps = [#input_map, #filter_map, #output_map], + iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction", "reduction"], + lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0, 0, 0, 0, 0, 0]]>} + ins(%input, %filter : tensor<4x10x10xf32>, tensor<16x4x3x3xf32>) outs(%init : tensor<16x8x8xf32>) { + ^bb0(%in: f32, %weight: f32, %acc: f32): + %product = arith.mulf %in, %weight : f32 + %sum = arith.addf %acc, %product : f32 + linalg.yield %sum : f32 + } -> tensor<16x8x8xf32> + return %0 : tensor<16x8x8xf32> +} +// CHECK-LABEL: func.func @generic_conv_lowering_config( +// CHECK: linalg.generic +// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<4x10x10xf32>, tensor<16x4x3x3xf32>) +// CHECK-SAME: lowering_config = +// CHECK-NOT: linalg.generic + +// ----- + +// The same holds for a preset compilation info. +#compilation = #iree_codegen.compilation_info< + lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0, 0, 0, 0, 0, 0, 0]]>, + translation_info = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<Default>>> +func.func @conv_compilation_info(%input: tensor<1x16x16x4xf32>, %filter: tensor<3x3x4x16xf32>, + %init: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> { + %0 = linalg.conv_2d_nhwc_hwcf {compilation_info = #compilation, + dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} + ins(%input, %filter : tensor<1x16x16x4xf32>, tensor<3x3x4x16xf32>) + outs(%init : tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> + return %0 : tensor<1x14x14x16xf32> +} +// CHECK-LABEL: func.func @conv_compilation_info( +// CHECK: linalg.conv_2d_nhwc_hwcf +// CHECK-SAME: compilation_info = +// CHECK-NOT: linalg.generic