This is an archive of the discontinued LLVM Phabricator instance.

Differential D113382

[mlir][linalg] Remove padding from tiling options.
ClosedPublic

Authored by gysit on Nov 8 2021, 12:12 AM.

Details

Reviewers
nicolasvasilache
Commits
rG0609eb1b32c2: [mlir][linalg] Remove padding from tiling options.
Summary

Remove the padding options from the tiling options since padding is now implemented by a separate pattern/pass introduced in https://reviews.llvm.org/D112412.

The revsion remove the tile-and-pad-tensors.mlir and replaces it with the pad.mlir that tests padding in isolation (without tiling). Similarly, hoist-padding.mlir is replaced by pad-and-hoist.mlir introduced in https://reviews.llvm.org/D112713.

Depends On D112838

Diff Detail

Revision Contents

PathSize
mlir/
include/
mlir/
Dialect/
Linalg/
Transforms/
25 lines
lib/
Dialect/
Linalg/
Transforms/
26 lines
test/
Dialect/
Linalg/
242 lines
lib/
Dialect/
Linalg/
16 lines
45 lines

Diff 385410

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Show First 20 Lines • Show All 571 Lines • ▼ Show 20 Linesstruct LinalgTilingOptions {
/// Specification markers of how to distribute the `linalg.tiled_loop`. /// Specification markers of how to distribute the `linalg.tiled_loop`.
SmallVector<StringRef, 2> distributionTypes = {}; SmallVector<StringRef, 2> distributionTypes = {};
LinalgTilingOptions &setDistributionTypes(ArrayRef<StringRef> types) { LinalgTilingOptions &setDistributionTypes(ArrayRef<StringRef> types) {
distributionTypes.assign(types.begin(), types.end()); distributionTypes.assign(types.begin(), types.end());
return *this; return *this;
} }
/// Callback returning the padding value to use for a given OpOperand or
/// failure for no padding. Padding operations are introduced if
/// `paddingValueComputationFunction` is set and does not return failure.
/// Padding all operands guarantees the operation is statically shaped and
/// thus can be vectorized.
PaddingValueComputationFunction paddingValueComputationFunction = nullptr;
LinalgTilingOptions &
setPaddingValueComputationFunction(PaddingValueComputationFunction fun) {
paddingValueComputationFunction = std::move(fun);
return *this;
}
/// Callback returning true if the pad tensor operation defining the given
/// OpOperand shall be marked as nofold to enable packing. A padding operation
/// is only marked nofold if `paddingNoFoldComputationFunction` is set and
/// returns true. Otherwise, the nofold attribute is set to false.
PaddingNoFoldComputationFunction paddingNoFoldComputationFunction = nullptr;
LinalgTilingOptions &
setPaddingNoFoldComputationFunction(PaddingNoFoldComputationFunction fun) {
paddingNoFoldComputationFunction = std::move(fun);
return *this;
}
/// Peel the specified loops. /// Peel the specified loops.
SmallVector<int64_t> peeledLoops; SmallVector<int64_t> peeledLoops;
LinalgTilingOptions &setPeeledLoops(ArrayRef<int64_t> loops) { LinalgTilingOptions &setPeeledLoops(ArrayRef<int64_t> loops) {
peeledLoops.clear(); peeledLoops.clear();
peeledLoops.append(loops.begin(), loops.end()); peeledLoops.append(loops.begin(), loops.end());
return *this; return *this;
} }
▲ Show 20 Lines • Show All 764 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

Show First 20 Lines • Show All 329 Lines • ▼ Show 20 LinesLogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
if (!res) if (!res)
return failure(); return failure();
// Clear filter to stop recursive pattern application. // Clear filter to stop recursive pattern application.
filter.replaceLinalgTransformationFilter(rewriter, res->op); filter.replaceLinalgTransformationFilter(rewriter, res->op);
// Peel loops. // Peel loops.
peelLoops(rewriter, *res, options); peelLoops(rewriter, *res, options);
// Consider padding on the fly only if the op has tensor semantics.
if (!options.paddingValueComputationFunction ||
!linalgOp.hasTensorSemantics()) {
result = *res; result = *res;
return success(); return success();
} }
// Try to pad on the fly by rewriting res->op as a padded op. If successful,
// `res.op` is rewritten in static form with padded operands.
LinalgOp paddedOp;
FailureOr<SmallVector<Value>> newResults = rewriteAsPaddedOp(
rewriter, res->op, options.paddingValueComputationFunction,
options.paddingNoFoldComputationFunction, paddedOp);
if (succeeded(newResults)) {
rewriter.replaceOp(res->op, newResults.getValue());
filter.replaceLinalgTransformationFilter(rewriter, paddedOp);
res->op = paddedOp;
result = *res;
// Do not perform replacement of `linalgOp`, let the derived patterns
// do this as they see fit, from the resulting TiledLinalgOp.
return success();
}
// Set so RAII guard does not propagate TiledLinalgOp to `result`.
return failure();
}
static ValueRange getTiledOpResult(TiledLinalgOp tiledOp) { static ValueRange getTiledOpResult(TiledLinalgOp tiledOp) {
if (tiledOp.loops.empty()) if (tiledOp.loops.empty())
return tiledOp.op.getOperation()->getResults(); return tiledOp.op.getOperation()->getResults();
return tiledOp.loops.front()->getResults(); return tiledOp.loops.front()->getResults();
} }
static ValueRange static ValueRange
getTiledAndFusedOpResult(TiledAndFusedLinalgOps tiledAndFusedOp) { getTiledAndFusedOpResult(TiledAndFusedLinalgOps tiledAndFusedOp) {
▲ Show 20 Lines • Show All 566 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/hoist-padding.mlir

  • This file was deleted.
// Specific structural checks are performed on 2-level hoisting
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=2 -canonicalize | FileCheck %s
// IR verification is performed on [0-6]-level hoisting
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=0 | FileCheck %s --check-prefix=VERIFIER-ONLY
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=1 | FileCheck %s --check-prefix=VERIFIER-ONLY
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=3 | FileCheck %s --check-prefix=VERIFIER-ONLY
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=4 | FileCheck %s --check-prefix=VERIFIER-ONLY
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=5 | FileCheck %s --check-prefix=VERIFIER-ONLY
// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-hoist-padding=6 | FileCheck %s --check-prefix=VERIFIER-ONLY
// CHECK-DAG: #[[$DIV3:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 3)>
// CHECK-DAG: #[[$DIV4:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 4)>
// CHECK-DAG: #[[$DIVS3:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 3)>
// CHECK-DAG: #[[$DIVS4:[0-9a-z]+]] = affine_map<()[s0] -> (s0 ceildiv 4)>
#map0 = affine_map<(d0)[s0] -> (2, -d0 + s0)>
#map1 = affine_map<(d0)[s0] -> (4, -d0 + s0)>
#map2 = affine_map<(d0)[s0] -> (3, -d0 + s0)>
#map3 = affine_map<(d0, d1) -> (2, d0 - d1)>
#map4 = affine_map<(d0, d1) -> (3, d0 - d1)>
// CHECK-LABEL: func @matmul_tensors
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor
// VERIFIER-ONLY-LABEL: func @matmul_tensors
func @matmul_tensors(
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
-> tensor<?x?xf32>
{
%c2 = arith.constant 2 : index
%c3 = arith.constant 3 : index
%c4 = arith.constant 4 : index
%cst = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
// CHECK-DAG: %[[dM:.*]] = tensor.dim %[[TA]], %[[C0]] : tensor<?x?xf32>
// CHECK-DAG: %[[dK:.*]] = tensor.dim %[[TA]], %[[C1]] : tensor<?x?xf32>
// CHECK-DAG: %[[dN:.*]] = tensor.dim %[[TB]], %[[C1]] : tensor<?x?xf32>
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
// CHECK: scf.for %[[I:[0-9a-z]+]] =
// First padded tensor is MxKx2x4 under loop M so Kx2x4
// CHECK: %[[SZpad0_K:[0-9]+]] = affine.apply #[[$DIVS4]]()[%[[dK]]]
// CHECK: linalg.init_tensor [%[[SZpad0_K]], 2, 4] : tensor<?x2x4xf32>
// 1-D loop
// CHECK: %[[A:.*]] = scf.for %[[J1:[0-9a-z]+]] =
// Iteration count along J1
// CHECK: %[[IDXpad0_K:[0-9]+]] = affine.apply #[[$DIV4]](%[[J1]])
// CHECK: tensor.extract_slice %{{.*}} [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: linalg.pad_tensor %{{.*}}
// CHECK: : tensor<?x?xf32> to tensor<2x4xf32>
// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%[[IDXpad0_K]], 0, 0]
// CHECK-SAME: [1, 2, 4] [1, 1, 1] : tensor<2x4xf32> into tensor<?x2x4xf32>
// Second tensor is KxN but loop order is (M, N, K) so padded tensor is NxKx4x3
// CHECK: %[[SZpad1_N:[0-9]+]] = affine.apply #[[$DIVS3]]()[%[[dN]]]
// CHECK: %[[SZpad1_K:[0-9]+]] = affine.apply #[[$DIVS4]]()[%[[dK]]]
// CHECK: linalg.init_tensor [%[[SZpad1_N]], %[[SZpad1_K]], 4, 3] : tensor<?x?x4x3xf32>
// 2-D loop
// CHECK: %[[B:.*]] = scf.for %[[K2:[0-9a-z]+]] =
// Iteration count along K2
// CHECK: %[[IDXpad1_K:[0-9]+]] = affine.apply #[[$DIV3]](%[[K2]])
// CHECK: scf.for %[[J2:[0-9a-z]+]] =
// Iteration count along J2
// CHECK: %[[IDXpad1_N:[0-9]+]] = affine.apply #[[$DIV4]](%[[J2]])
// CHECK: tensor.extract_slice %{{.*}} [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: linalg.pad_tensor %{{.*}}
// CHECK: : tensor<?x?xf32> to tensor<4x3xf32>
// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%[[IDXpad1_K]], %[[IDXpad1_N]], 0, 0]
// CHECK-SAME: [1, 1, 4, 3] [1, 1, 1, 1] : tensor<4x3xf32> into tensor<?x?x4x3xf32>
// 2-D loop
// CHECK: scf.for %[[J:[0-9a-zA-Z]+]]
// CHECK: scf.for %[[K:[0-9a-zA-Z]+]]
// Iteration count along K
// CHECK: %[[IDXpad0_K:[0-9]+]] = affine.apply #[[$DIV4]](%[[K]])
// CHECK: %[[stA:.*]] = tensor.extract_slice %[[A]][%[[IDXpad0_K]], 0, 0] [1, 2, 4] [1, 1, 1] :
// CHECK-SAME: tensor<?x2x4xf32> to tensor<2x4xf32>
// Iteration count along J
// CHECK: %[[IDXpad1_N:[0-9]+]] = affine.apply #[[$DIV3]](%[[J]])
// Iteration count along K
// CHECK: %[[IDXpad1_K:[0-9]+]] = affine.apply #[[$DIV4]](%[[K]])
// CHECK: %[[stB:.*]] = tensor.extract_slice %[[B]][%[[IDXpad1_N]], %[[IDXpad1_K]], 0, 0] [1, 1, 4, 3] [1, 1, 1, 1] :
// CHECK-SAME: tensor<?x?x4x3xf32> to tensor<4x3xf32>
// CHECK: %[[stC:.*]] = linalg.pad_tensor %{{.*}}
// CHECK: : tensor<?x?xf32> to tensor<2x3xf32>
// CHECK: linalg.matmul ins(%[[stA]], %[[stB]] : tensor<2x4xf32>, tensor<4x3xf32>)
// CHECK-SAME: outs(%[[stC]] : tensor<2x3xf32>) -> tensor<2x3xf32>
%3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
%4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
%5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
%6 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%7 = affine.min #map0(%arg3)[%6]
%8 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%9 = affine.min #map1(%arg7)[%8]
%10 = tensor.extract_slice %arg0[%arg3, %arg7] [%7, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%11 = tensor.dim %arg1, %c0 : tensor<?x?xf32>
%12 = affine.min #map1(%arg7)[%11]
%13 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%14 = affine.min #map2(%arg5)[%13]
%15 = tensor.extract_slice %arg1[%arg7, %arg5] [%12, %14] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%16 = tensor.dim %arg8, %c0 : tensor<?x?xf32>
%17 = affine.min #map3(%16, %arg3)
%18 = tensor.dim %arg8, %c1 : tensor<?x?xf32>
%19 = affine.min #map4(%18, %arg5)
%20 = tensor.extract_slice %arg8[%arg3, %arg5] [%17, %19] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%21 = arith.subi %c2, %7 : index
%22 = arith.subi %c4, %9 : index
%23 = linalg.pad_tensor %10 low[%c0, %c0] high[%21, %22] {
^bb0(%arg9: index, %arg10: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?x?xf32> to tensor<2x4xf32>
%24 = arith.subi %c4, %12 : index
%25 = arith.subi %c3, %14 : index
%26 = linalg.pad_tensor %15 low[%c0, %c0] high[%24, %25] {
^bb0(%arg9: index, %arg10: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?x?xf32> to tensor<4x3xf32>
%27 = arith.subi %c2, %17 : index
%28 = arith.subi %c3, %19 : index
%29 = linalg.pad_tensor %20 low[%c0, %c0] high[%27, %28] {
^bb0(%arg9: index, %arg10: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?x?xf32> to tensor<2x3xf32>
%30 = linalg.matmul ins(%23, %26 : tensor<2x4xf32>, tensor<4x3xf32>) outs(%29 : tensor<2x3xf32>) -> tensor<2x3xf32>
%31 = tensor.extract_slice %30[0, 0] [%7, %14] [1, 1] : tensor<2x3xf32> to tensor<?x?xf32>
%32 = tensor.insert_slice %31 into %arg8[%arg3, %arg5] [%17, %19] [%c1, %c1] : tensor<?x?xf32> into tensor<?x?xf32>
scf.yield %32 : tensor<?x?xf32>
}
scf.yield %5 : tensor<?x?xf32>
}
scf.yield %4 : tensor<?x?xf32>
}
return %3 : tensor<?x?xf32>
}
// -----
// CHECK-DAG: #[[$MIN_REST8:[0-9a-z]+]] = affine_map<(d0)[s0] -> (8, -d0 + s0)>
// CHECK-DAG: #[[$DIV4:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 4)>
// CHECK-DAG: #[[$DIV2:[0-9a-z]+]] = affine_map<(d0) -> (d0 ceildiv 2)>
#map0 = affine_map<(d0)[s0] -> (8, -d0 + s0)>
#map1 = affine_map<(d0, d1) -> (4, -d0 + d1)>
#map2 = affine_map<(d0, d1) -> (2, -d0 + d1)>
#map3 = affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>
// CHECK-LABEL: func @dot
// VERIFIER-ONLY-LABEL: func @dot
func @dot(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<f32>)
-> tensor<f32>
{
%cst = arith.constant 0.000000e+00 : f32
%c8 = arith.constant 8 : index
%c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index
%c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c0 : tensor<?xf32>
// CHECK: scf.for %[[I:[0-9a-z]+]] =
//
// CHECK: %[[MR8:.*]] = affine.min #[[$MIN_REST8]](%[[I]])
// Init tensor and pack.
// CHECK: %[[INIT_PACKED_A:.*]] = linalg.init_tensor [2, 2, 2] : tensor<2x2x2xf32>
// CHECK: %[[CAST_INIT_PACKED_A:.*]] = tensor.cast %[[INIT_PACKED_A]] : tensor<2x2x2xf32> to tensor<?x?x2xf32>
// CHECK: %[[PACKED_A:.*]] = scf.for %[[II:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[CAST_INIT_PACKED_A]]) -> (tensor<?x?x2xf32>) {
// CHECK: scf.for %[[III:[0-9a-z]+]] =
// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor<?x?x2xf32>
//
// Init tensor and pack.
// CHECK: %[[INIT_PACKED_B:.*]] = linalg.init_tensor [2, 2, 2] : tensor<2x2x2xf32>
// CHECK: %[[CAST_INIT_PACKED_B:.*]] = tensor.cast %[[INIT_PACKED_B]] : tensor<2x2x2xf32> to tensor<?x?x2xf32>
// CHECK: %[[PACKED_B:.*]] = scf.for %[[II_2:[0-9a-z]+]] = {{.*}} iter_args(%{{.*}} = %[[CAST_INIT_PACKED_B]]) -> (tensor<?x?x2xf32>) {
// CHECK: scf.for %[[III_2:[0-9a-z]+]] =
// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}, 0] [1, 1, 2] [1, 1, 1] : tensor<2xf32> into tensor<?x?x2xf32>
// Compute.
// CHECK: scf.for %[[II_3:[0-9a-z]+]] =
// CHECK: scf.for %[[III_3:[0-9a-z]+]] = {{.*}} iter_args(%[[C:.*]] = %{{.*}}) -> (tensor<f32>) {
// CHECK: %[[IDX0:.*]] = affine.apply #[[$DIV4]](%[[II_3]])
// CHECK: %[[IDX1:.*]] = affine.apply #[[$DIV2]](%[[III_3]])
// CHECK: %[[A:.*]] = tensor.extract_slice %[[PACKED_A]][%[[IDX0]], %[[IDX1]], 0] [1, 1, 2] [1, 1, 1] : tensor<?x?x2xf32> to tensor<2xf32>
// CHECK: %[[IDX0_2:.*]] = affine.apply #[[$DIV4]](%[[II_3]])
// CHECK: %[[IDX1_2:.*]] = affine.apply #[[$DIV2]](%[[III_3]])
// CHECK: %[[B:.*]] = tensor.extract_slice %[[PACKED_B]][%[[IDX0_2]], %[[IDX1_2]], 0] [1, 1, 2] [1, 1, 1] : tensor<?x?x2xf32> to tensor<2xf32>
// CHECK: linalg.dot ins(%[[A]], %[[B]] : tensor<2xf32>, tensor<2xf32>) outs(%[[C]] : tensor<f32>) -> tensor<f32>
%1 = scf.for %arg3 = %c0 to %0 step %c8 iter_args(%arg4 = %arg2) -> (tensor<f32>) {
%2 = affine.min #map0(%arg3)[%0]
%3 = scf.for %arg5 = %c0 to %2 step %c4 iter_args(%arg6 = %arg4) -> (tensor<f32>) {
%4 = affine.min #map1(%arg5, %2)
%5 = scf.for %arg7 = %c0 to %4 step %c2 iter_args(%arg8 = %arg6) -> (tensor<f32>) {
%6 = affine.min #map2(%arg7, %4)
%7 = affine.apply #map3(%arg7, %arg5, %arg3)
%8 = tensor.extract_slice %arg0[%7] [%6] [1] : tensor<?xf32> to tensor<?xf32>
%9 = tensor.extract_slice %arg1[%7] [%6] [1] : tensor<?xf32> to tensor<?xf32>
%10 = arith.subi %c2, %6 : index
%11 = linalg.pad_tensor %8 low[%c0] high[%10] {
^bb0(%arg9: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?xf32> to tensor<2xf32>
%12 = linalg.pad_tensor %9 low[%c0] high[%10] {
^bb0(%arg9: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?xf32> to tensor<2xf32>
%13 = linalg.dot ins(%11, %12 : tensor<2xf32>, tensor<2xf32>) outs(%arg8 : tensor<f32>) -> tensor<f32>
scf.yield %13 : tensor<f32>
}
scf.yield %5 : tensor<f32>
}
scf.yield %3 : tensor<f32>
}
return %1 : tensor<f32>
}
// -----
// CHECK-LABEL: func @matmul_2d_tiling
// VERIFIER-ONLY-LABEL: func @matmul_2d_tiling
func @matmul_2d_tiling(%arg0: tensor<32x128xf32>, %arg1: tensor<128x64xf32>, %arg2: tensor<32x64xf32>) -> tensor<32x64xf32> {
%c128 = arith.constant 128 : index
%c64 = arith.constant 64 : index
%c32 = arith.constant 32 : index
%c16 = arith.constant 16 : index
%cst = arith.constant 0.000000e+00 : f32
%c2 = arith.constant 2 : index
%c4 = arith.constant 4 : index
%c0 = arith.constant 0 : index
%1 = scf.for %arg3 = %c0 to %c32 step %c16 iter_args(%arg4 = %arg2) -> (tensor<32x64xf32>) {
%2 = scf.for %arg5 = %c0 to %c64 step %c32 iter_args(%arg6 = %arg4) -> (tensor<32x64xf32>) {
%3 = scf.for %arg7 = %c0 to %c128 step %c32 iter_args(%arg8 = %arg6) -> (tensor<32x64xf32>) {
%4 = tensor.extract_slice %arg0[%arg3, %arg7] [16, 32] [1, 1] : tensor<32x128xf32> to tensor<16x32xf32>
%5 = tensor.extract_slice %arg1[%arg7, %arg5] [32, 32] [1, 1] : tensor<128x64xf32> to tensor<32x32xf32>
%6 = tensor.extract_slice %arg8[%arg3, %arg5] [16, 32] [1, 1] : tensor<32x64xf32> to tensor<16x32xf32>
%7 = scf.for %arg9 = %c0 to %c16 step %c2 iter_args(%arg10 = %6) -> (tensor<16x32xf32>) {
%10 = scf.for %arg11 = %c0 to %c32 step %c4 iter_args(%arg12 = %arg10) -> (tensor<16x32xf32>) {
%11 = scf.for %arg13 = %c0 to %c32 step %c16 iter_args(%arg14 = %arg12) -> (tensor<16x32xf32>) {
%12 = tensor.extract_slice %4[%arg9, %arg13] [2, 16] [1, 1] : tensor<16x32xf32> to tensor<2x16xf32>
%13 = tensor.cast %12 : tensor<2x16xf32> to tensor<?x?xf32>
%14 = tensor.extract_slice %5[%arg13, %arg11] [16, 4] [1, 1] : tensor<32x32xf32> to tensor<16x4xf32>
%15 = tensor.cast %14 : tensor<16x4xf32> to tensor<?x?xf32>
%16 = tensor.extract_slice %arg14[%arg9, %arg11] [2, 4] [1, 1] : tensor<16x32xf32> to tensor<2x4xf32>
%17 = tensor.cast %16 : tensor<2x4xf32> to tensor<?x?xf32>
%18 = linalg.pad_tensor %13 low[%c0, %c0] high[%c0, %c0] {
^bb0(%arg15: index, %arg16: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?x?xf32> to tensor<2x16xf32>
%19 = linalg.pad_tensor %15 low[%c0, %c0] high[%c0, %c0] {
^bb0(%arg15: index, %arg16: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?x?xf32> to tensor<16x4xf32>
%20 = linalg.pad_tensor %17 low[%c0, %c0] high[%c0, %c0] {
^bb0(%arg15: index, %arg16: index): // no predecessors
linalg.yield %cst : f32
} : tensor<?x?xf32> to tensor<2x4xf32>
%21 = linalg.matmul ins(%18, %19 : tensor<2x16xf32>, tensor<16x4xf32>) outs(%20 : tensor<2x4xf32>) -> tensor<2x4xf32>
%22 = tensor.cast %21 : tensor<2x4xf32> to tensor<?x?xf32>
%23 = tensor.insert_slice %22 into %arg14[%arg9, %arg11] [%c2, %c4] [1, 1] : tensor<?x?xf32> into tensor<16x32xf32>
scf.yield %23 : tensor<16x32xf32>
}
scf.yield %11 : tensor<16x32xf32>
}
scf.yield %10 : tensor<16x32xf32>
}
%8 = tensor.cast %7 : tensor<16x32xf32> to tensor<?x?xf32>
%9 = tensor.insert_slice %8 into %arg8[%arg3, %arg5] [%c16, %c32] [1, 1] : tensor<?x?xf32> into tensor<32x64xf32>
scf.yield %9 : tensor<32x64xf32>
}
scf.yield %3 : tensor<32x64xf32>
}
scf.yield %2 : tensor<32x64xf32>
}
return %1 : tensor<32x64xf32>
}

mlir/test/Dialect/Linalg/pad.mlir

  • This file was added.
// RUN: mlir-opt %s -test-linalg-transform-patterns="test-pad-pattern pack-paddings=1,1,0 hoist-paddings=0,0,0" -cse -canonicalize -split-input-file | FileCheck %s
// CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (7, -d0 + 12)>
#map = affine_map<(d0) -> (7, -d0 + 12)>
// CHECK: static_sizes_output_divisible
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<24x12xf32>
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<12x25xf32>
func @static_sizes_output_divisible(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C7:.*]] = arith.constant 7
%c0 = arith.constant 0 : index
%c12 = arith.constant 12 : index
%c25 = arith.constant 25 : index
%c24 = arith.constant 24 : index
%c7 = arith.constant 7 : index
%c5 = arith.constant 5 : index
%c4 = arith.constant 4 : index
// CHECK: scf.for %[[IV0:[0-9a-zA-Z]*]] =
%0 = scf.for %arg3 = %c0 to %c24 step %c4 iter_args(%arg4 = %arg2) -> (tensor<24x25xf32>) {
// CHECK: scf.for %[[IV1:[0-9a-zA-Z]*]] =
%1 = scf.for %arg5 = %c0 to %c25 step %c5 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
// CHECK: scf.for %[[IV2:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG4:.*]] =
%2 = scf.for %arg7 = %c0 to %c12 step %c7 iter_args(%arg8 = %arg6) -> (tensor<24x25xf32>) {
// CHECK: %[[TS2:.*]] = affine.min #[[MAP0]](%[[IV2]])
%3 = affine.min #map(%arg7)
// CHECK: %[[T0:.*]] = tensor.extract_slice %[[ARG0]]
// CHECK: %[[T1:.*]] = tensor.extract_slice %[[ARG1]]
// CHECK: %[[T2:.*]] = tensor.extract_slice %[[ARG4]]
%4 = tensor.extract_slice %arg0[%arg3, %arg7] [4, %3] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
%5 = tensor.extract_slice %arg1[%arg7, %arg5] [%3, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
%6 = tensor.extract_slice %arg8[%arg3, %arg5] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
// Check statically sized matmul inputs with partially divisible sizes are padded.
// CHECK: %[[V0:.*]] = arith.subi %[[C7]], %[[TS2]]
// CHECK: %[[T3:.*]] = linalg.pad_tensor %[[T0]] nofold
// CHECK-SAME: [%[[C0]], %[[C0]]]
// CHECK-SAME: [%[[C0]], %[[V0]]
// CHECK: %[[T4:.*]] = linalg.pad_tensor %[[T1]] nofold
// Check the statically sized matmul output with fully divisible sizes is not padded.
// CHECK: %[[T5:.*]] = linalg.matmul
// CHECK-SAME: ins(%[[T3]], %[[T4]] : tensor<4x7xf32>, tensor<7x5xf32>)
// CHECK-SAME: outs(%[[T2]] : tensor<4x5xf32>)
// CHECK: %[[T6:.*]] = tensor.insert_slice %[[T5]]
%7 = linalg.matmul {__internal_linalg_transform__ = "pad"} ins(%4, %5 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%6 : tensor<4x5xf32>) -> tensor<4x5xf32>
%8 = tensor.insert_slice %7 into %arg8[%arg3, %arg5] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
// CHECK: scf.yield %[[T6]]
scf.yield %8 : tensor<24x25xf32>
}
scf.yield %2 : tensor<24x25xf32>
}
scf.yield %1 : tensor<24x25xf32>
}
return %0 : tensor<24x25xf32>
}
// -----
// CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0) -> (7, -d0 + 25)>
#map = affine_map<(d0) -> (7, -d0 + 25)>
// CHECK: static_sizes_input_divisible
func @static_sizes_input_divisible(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C7:.*]] = arith.constant 7
%c0 = arith.constant 0 : index
%c12 = arith.constant 12 : index
%c25 = arith.constant 25 : index
%c24 = arith.constant 24 : index
%c6 = arith.constant 6 : index
%c7 = arith.constant 7 : index
%c4 = arith.constant 4 : index
// CHECK: scf.for %[[IV0:[0-9a-zA-Z]*]] =
%0 = scf.for %arg3 = %c0 to %c24 step %c4 iter_args(%arg4 = %arg2) -> (tensor<24x25xf32>) {
// CHECK: scf.for %[[IV1:[0-9a-zA-Z]*]] =
%1 = scf.for %arg5 = %c0 to %c25 step %c7 iter_args(%arg6 = %arg4) -> (tensor<24x25xf32>) {
// CHECK: scf.for %[[IV2:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG4:.*]] =
%2 = scf.for %arg7 = %c0 to %c12 step %c6 iter_args(%arg8 = %arg6) -> (tensor<24x25xf32>) {
%3 = tensor.extract_slice %arg0[%arg3, %arg7] [4, 6] [1, 1] : tensor<24x12xf32> to tensor<4x6xf32>
// CHECK: %[[TS1:.*]] = affine.min #[[MAP0]](%[[IV1]])
%4 = affine.min #map(%arg5)
%5 = tensor.extract_slice %arg1[%arg7, %arg5] [6, %4] [1, 1] : tensor<12x25xf32> to tensor<6x?xf32>
// CHECK: %[[T0:.*]] = tensor.extract_slice %[[ARG4]]
%6 = tensor.extract_slice %arg8[%arg3, %arg5] [4, %4] [1, 1] : tensor<24x25xf32> to tensor<4x?xf32>
// Check the statically sized matmul output with partially divisible sizes is padded.
// CHECK: %[[V0:.*]] = arith.subi %[[C7]], %[[TS1]]
// CHECK: %[[T1:.*]] = linalg.pad_tensor %[[T0]] low
// CHECK-SAME: [%[[C0]], %[[C0]]]
// CHECK-SAME: [%[[C0]], %[[V0]]
// CHECK: %[[T2:.*]] = linalg.matmul
// CHECK-SAME: outs(%[[T1]] : tensor<4x7xf32>)
// CHECK: %[[T3:.*]] = tensor.extract_slice %[[T2]]
// CHECK: %[[T4:.*]] = tensor.insert_slice %[[T3]]
%7 = linalg.matmul {__internal_linalg_transform__ = "pad"} ins(%3, %5 : tensor<4x6xf32>, tensor<6x?xf32>) outs(%6 : tensor<4x?xf32>) -> tensor<4x?xf32>
%8 = tensor.insert_slice %7 into %arg8[%arg3, %arg5] [4, %4] [1, 1] : tensor<4x?xf32> into tensor<24x25xf32>
// CHECK: scf.yield %[[T4]]
scf.yield %8 : tensor<24x25xf32>
}
scf.yield %2 : tensor<24x25xf32>
}
scf.yield %1 : tensor<24x25xf32>
}
return %0 : tensor<24x25xf32>
}
// -----
// CHECK-DAG: #[[MAP0:[0-9a-z]+]] = affine_map<(d0)[s0] -> (5, -d0 + s0)>
// CHECK-DAG: #[[MAP1:[0-9a-z]+]] = affine_map<(d0)[s0] -> (7, -d0 + s0)>
// CHECK-DAG: #[[MAP2:[0-9a-z]+]] = affine_map<(d0)[s0] -> (6, -d0 + s0)>
#map0 = affine_map<(d0)[s0] -> (5, -d0 + s0)>
#map1 = affine_map<(d0)[s0] -> (6, -d0 + s0)>
#map2 = affine_map<(d0)[s0] -> (7, -d0 + s0)>
// CHECK: dynamic_sizes
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]*]]: tensor<?x?xf32>
func @dynamic_sizes(%arg0: tensor<?x?xf32>,
%arg1: tensor<?x?xf32>,
%arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1
// CHECK-DAG: %[[C5:.*]] = arith.constant 5
// CHECK-DAG: %[[C6:.*]] = arith.constant 6
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%c6 = arith.constant 6 : index
%c7 = arith.constant 7 : index
%c5 = arith.constant 5 : index
// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D2:.*]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK-DAG: %[[D1:.*]] = tensor.dim %[[ARG1]], %[[C1]]
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
// CHECK: scf.for %[[IV0:[0-9a-zA-Z]*]] =
%3 = scf.for %arg3 = %c0 to %0 step %c5 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
// CHECK: scf.for %[[IV1:[0-9a-zA-Z]*]] =
%4 = scf.for %arg5 = %c0 to %2 step %c7 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
// CHECK: scf.for %[[IV2:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG4:.*]] =
%5 = scf.for %arg7 = %c0 to %1 step %c6 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
// CHECK: %[[TS0:.*]] = affine.min #[[MAP0]](%[[IV0]])[%[[D0]]]
// CHECK: %[[TS2:.*]] = affine.min #[[MAP2]](%[[IV2]])[%[[D2]]]
// CHECK: %[[TS1:.*]] = affine.min #[[MAP1]](%[[IV1]])[%[[D1]]]
%6 = affine.min #map0(%arg3)[%0]
%7 = affine.min #map1(%arg7)[%1]
%8 = tensor.extract_slice %arg0[%arg3, %arg7] [%6, %7] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%9 = affine.min #map2(%arg5)[%2]
%10 = tensor.extract_slice %arg1[%arg7, %arg5] [%7, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%11 = tensor.extract_slice %arg8[%arg3, %arg5] [%6, %9] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
// Check all matmul operands are padded.
// CHECK: %[[V0:.*]] = arith.subi %[[C5]], %[[TS0]]
// CHECK: %[[V1:.*]] = arith.subi %[[C6]], %[[TS2]]
// CHECK: %[[T3:.*]] = linalg.pad_tensor %{{.*}} nofold
// CHECK-SAME: [%[[C0]], %[[C0]]]
// CHECK-SAME: [%[[V0]], %[[V1]]
// CHECK: %[[T4:.*]] = linalg.pad_tensor %{{.*}} nofold
// CHECK: %[[T5:.*]] = linalg.pad_tensor %{{.*}} low
// Check the dynamic matmul has been erased.
// CHECK-NOT: = linalg.matmul {{.*}} tensor<?x?xf32>
// Check all padded matmul operands are statically sized.
// CHECK: %[[T6:.*]] = linalg.matmul
// CHECK-SAME: ins(%[[T3]], %[[T4]] : tensor<5x6xf32>, tensor<6x7xf32>)
// CHECK-SAME: outs(%[[T5]] : tensor<5x7xf32>)
// CHECK: %[[T7:.*]] = tensor.extract_slice %[[T6]][0, 0] [%[[TS0]], %[[TS1]]]
// CHECK: %[[T8:.*]] = tensor.insert_slice %[[T7]]
%12 = linalg.matmul {__internal_linalg_transform__ = "pad"} ins(%8, %10 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%11 : tensor<?x?xf32>) -> tensor<?x?xf32>
%13 = tensor.insert_slice %12 into %arg8[%arg3, %arg5] [%6, %9] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
// CHECK: scf.yield %[[T8]]
scf.yield %13 : tensor<?x?xf32>
}
scf.yield %5 : tensor<?x?xf32>
}
scf.yield %4 : tensor<?x?xf32>
}
return %3 : tensor<?x?xf32>
}
// -----
#map = affine_map<(d0) -> (7, -d0 + 12)>
// CHECK: scalar_operand
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: f32
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<24x12xf32>
func @scalar_operand(%arg0: f32, %arg1: tensor<24x12xf32>) -> tensor<24x12xf32> {
%c0 = arith.constant 0 : index
%c12 = arith.constant 12 : index
%c24 = arith.constant 24 : index
%c7 = arith.constant 7 : index
%c4 = arith.constant 4 : index
// CHECK: scf.for %[[IV0:[0-9a-zA-Z]*]] =
%0 = scf.for %arg2 = %c0 to %c24 step %c4 iter_args(%arg3 = %arg1) -> (tensor<24x12xf32>) {
// CHECK: scf.for %[[IV1:[0-9a-zA-Z]*]] = {{.*}} iter_args(%[[ARG2:.*]] =
%1 = scf.for %arg4 = %c0 to %c12 step %c7 iter_args(%arg5 = %arg3) -> (tensor<24x12xf32>) {
%2 = affine.min #map(%arg4)
// CHECK: %[[T0:.*]] = tensor.extract_slice %[[ARG2]]
// CHECK: %[[T1:.*]] = linalg.pad_tensor %[[T0]] nofold
%3 = tensor.extract_slice %arg5[%arg2, %arg4] [4, %2] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
// Check only the fill output operand is padded.
// CHECK: %[[T6:.*]] = linalg.fill(%[[ARG0]], %[[T1]]
%4 = linalg.fill(%arg0, %3) {__internal_linalg_transform__ = "pad"} : f32, tensor<4x?xf32> -> tensor<4x?xf32>
%5 = tensor.insert_slice %4 into %arg5[%arg2, %arg4] [4, %2] [1, 1] : tensor<4x?xf32> into tensor<24x12xf32>
scf.yield %5 : tensor<24x12xf32>
}
scf.yield %1 : tensor<24x12xf32>
}
return %0 : tensor<24x12xf32>
}

mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir

  • This file was deleted.
// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern padded-operands=0,1,2 nofold-operands=0,1 tile-sizes=2,3,4" -canonicalize | FileCheck %s
// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern padded-operands=0,1 nofold-operands=0,1 tile-sizes=2,3" -canonicalize | FileCheck %s -check-prefix=CHECK-1DIM-TILE
// CHECK-LABEL: func @matmul_tensors(
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xi8>
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xi8>
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
func @matmul_tensors(
%arg0: tensor<?x?xi8>, %arg1: tensor<?x?xi8>, %arg2: tensor<?x?xi32>)
-> tensor<?x?xi32> {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xi32>) {
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xi32>) {
// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xi32>) {
// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xi8> to tensor<?x?xi8>
// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xi8> to tensor<?x?xi8>
// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xi32> to tensor<?x?xi32>
// Dynamic op has been canonicalized away.
// CHECK-NOT: linalg.matmul {{.*}} tensor<?x?xi8>
// Padding injects static information.
// CHECK: %[[pA:.*]] = linalg.pad_tensor %[[sTA]] nofold low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}]
// CHECK: : tensor<?x?xi8> to tensor<2x4xi8>
// CHECK: %[[pB:.*]] = linalg.pad_tensor %[[sTB]] nofold low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}]
// CHECK: : tensor<?x?xi8> to tensor<4x3xi8>
// CHECK: %[[pC:.*]] = linalg.pad_tensor %[[sTC]] low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}]
// CHECK: : tensor<?x?xi32> to tensor<2x3xi32>
// CHECK: %[[pD:.*]] = linalg.matmul ins(%[[pA]], %[[pB]] : tensor<2x4xi8>, tensor<4x3xi8>)
// CHECK-SAME: outs(%[[pC]] : tensor<2x3xi32>) -> tensor<2x3xi32>
// CHECK: %[[sTD:.*]] = tensor.extract_slice %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xi32> to tensor<?x?xi32>
// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xi32> into tensor<?x?xi32>
// CHECK: scf.yield %[[TD]] : tensor<?x?xi32>
// CHECK: scf.yield %[[TD2]] : tensor<?x?xi32>
// CHECK: scf.yield %[[TD1]] : tensor<?x?xi32>
%0 = linalg.matmul {__internal_linalg_transform__ = "tile"}
ins(%arg0, %arg1: tensor<?x?xi8>, tensor<?x?xi8>)
outs(%arg2: tensor<?x?xi32>)
-> tensor<?x?xi32>
// CHECK: return %[[TD0]] : tensor<?x?xi32>
return %0 : tensor<?x?xi32>
}
// CHECK-LABEL: func @generic_scalar_and_tensor(
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?x?xf32>
// CHECK-SAME: %[[VAL:[0-9a-z]+]]: f32) -> tensor<?x?x?xf32> {
func @generic_scalar_and_tensor(
%arg0: tensor<?x?x?xf32>, %arg1: f32)
-> tensor<?x?x?xf32> {
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?x?xf32>) {
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?x?xf32>) {
// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?x?xf32>) {
// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
// Padding injects static information.
// CHECK: %[[pC:.*]] = linalg.pad_tensor %[[sTC]] nofold low[%[[C0]], %[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}, %{{.*}}]
// CHECK: : tensor<?x?x?xf32> to tensor<2x3x4xf32>
// CHECK: %[[pD:.*]] = linalg.generic
// CHECK-SAME: ins(%[[VAL]] : f32) outs(%[[pC]] : tensor<2x3x4xf32>)
// CHECK: %[[sTD:.*]] = tensor.extract_slice %[[pD]][0, 0, 0] [%{{.*}}, %{{.*}}, %{{.*}}] [1, 1, 1] : tensor<2x3x4xf32> to tensor<?x?x?xf32>
// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?x?xf32> into tensor<?x?x?xf32>
// CHECK: scf.yield %[[TD]] : tensor<?x?x?xf32>
// CHECK: scf.yield %[[TD2]] : tensor<?x?x?xf32>
// CHECK: scf.yield %[[TD1]] : tensor<?x?x?xf32>
%0 = linalg.generic {
indexing_maps = [ affine_map<(d0, d1, d2) -> ()>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)> ],
iterator_types = ["parallel", "parallel", "parallel"]}
{__internal_linalg_transform__ = "tile"}
ins(%arg1 : f32)
outs(%arg0: tensor<?x?x?xf32>) {
^bb(%0: f32, %1: f32) :
linalg.yield %0 : f32
} -> tensor<?x?x?xf32>
return %0 : tensor<?x?x?xf32>
}
// CHECK-1DIM-TILE: func @matmul_tensors(
// CHECK-1DIM-TILE: %[[TA:[0-9a-z]+]]: tensor<?x?xi8>
// CHECK-1DIM-TILE: %[[TB:[0-9a-z]+]]: tensor<?x?xi8>
// CHECK-1DIM-TILE: %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
// CHECK-1DIM-TILE-NOT: scf.for
// CHECK-1DIM-TILE: linalg.matmul ins(%[[TA]], %[[TB]] : tensor<?x?xi8>, tensor<?x?xi8>) outs(%[[TC]] : tensor<?x?xi32>) -> tensor<?x?xi32>
func @matmul_partially_padded_tensors(
%arg0: tensor<?x8xi8>, %arg1: tensor<8x?xi8>, %arg2: tensor<?x?xi32>)
-> tensor<?x?xi32> {
%0 = linalg.matmul {__internal_linalg_transform__ = "tile"}
ins(%arg0, %arg1: tensor<?x8xi8>, tensor<8x?xi8>)
outs(%arg2: tensor<?x?xi32>)
-> tensor<?x?xi32>
return %0 : tensor<?x?xi32>
}
// CHECK-LABEL: func @matmul_partially_padded_tensors(
// CHECK: linalg.matmul ins({{.*}}, {{.*}} : tensor<2x4xi8>, tensor<4x3xi8>) outs({{.*}} : tensor<2x3xi32>) -> tensor<2x3xi32>
// Check only the the input operands are padded.
// CHECK-1DIM-TILE: func @matmul_partially_padded_tensors(
// CHECK-1DIM-TILE-SAME: %[[TA:[0-9a-z]+]]: tensor<?x8xi8>
// CHECK-1DIM-TILE-SAME: %[[TB:[0-9a-z]+]]: tensor<8x?xi8>
// CHECK-1DIM-TILE-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
// CHECK-1DIM-TILE: %[[C0:.*]] = arith.constant 0 : index
// CHECK-1DIM-TILE: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xi32>) {
// CHECK-1DIM-TILE: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xi32>) {
// CHECK-1DIM-TILE: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x8xi8> to tensor<?x8xi8>
// CHECK-1DIM-TILE: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<8x?xi8> to tensor<8x?xi8>
// CHECK-1DIM-TILE: %[[sTC:.*]] = tensor.extract_slice %[[TC1]][{{.*}}] : tensor<?x?xi32> to tensor<?x?xi32>
// CHECK-1DIM-TILE: %[[pA:.*]] = linalg.pad_tensor %[[sTA]] nofold low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}]
// CHECK-1DIM-TILE: : tensor<?x8xi8> to tensor<2x8xi8>
// CHECK-1DIM-TILE: %[[pB:.*]] = linalg.pad_tensor %[[sTB]] nofold low[%[[C0]], %[[C0]]] high[%{{.*}}, %{{.*}}]
// CHECK-1DIM-TILE: : tensor<8x?xi8> to tensor<8x3xi8>
// CHECK-1DIM-TILE: %[[pD:.*]] = linalg.matmul ins(%[[pA]], %[[pB]] : tensor<2x8xi8>, tensor<8x3xi8>)
// CHECK-1DIM-TILE: outs(%[[sTC]] : tensor<?x?xi32>) -> tensor<?x?xi32>
// Check that the tile-and-pad transformation actually introduces the padding
// as requested, even if original operation already operates on static
// shapes.
// CHECK-LABEL: @pad_to_same_static_size
func @pad_to_same_static_size(%arg0: tensor<2x3x4xf32>, %arg1: f32) -> tensor<2x3x4xf32> {
// CHECK: %[[c0:.*]] = arith.constant 0 : index
// CHECK-NOT: scf.for
// CHECK: linalg.pad_tensor %{{.*}} nofold low[%[[c0]], %[[c0]], %[[c0]]] high[%[[c0]], %[[c0]], %[[c0]]]
// CHECK: tensor<2x3x4xf32> to tensor<2x3x4xf32>
%0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> ()>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)> ],
iterator_types = ["parallel", "parallel", "parallel"]}
{__internal_linalg_transform__ = "tile"}
ins(%arg1 : f32) outs(%arg0 : tensor<2x3x4xf32>) {
^bb0(%arg2: f32, %arg3: f32): // no predecessors
linalg.yield %arg2 : f32
} -> tensor<2x3x4xf32>
return %0 : tensor<2x3x4xf32>
}
// CHECK-LABEL: @pad_static_divisible_size
func @pad_static_divisible_size(%arg0: tensor<4x6x8xf32>, %arg1: f32) -> tensor<4x6x8xf32> {
// CHECK: %[[c0:.*]] = arith.constant 0 : index
// CHECK-COUNT-3: scf.for
// CHECK: linalg.pad_tensor %{{.*}} nofold low[%[[c0]], %[[c0]], %[[c0]]] high[%[[c0]], %[[c0]], %[[c0]]]
// CHECK: tensor<2x3x4xf32> to tensor<2x3x4xf32>
%0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> ()>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)> ],
iterator_types = ["parallel", "parallel", "parallel"]}
{__internal_linalg_transform__ = "tile"}
ins(%arg1 : f32) outs(%arg0 : tensor<4x6x8xf32>) {
^bb0(%arg2: f32, %arg3: f32): // no predecessors
linalg.yield %arg2 : f32
} -> tensor<4x6x8xf32>
return %0 : tensor<4x6x8xf32>
}

mlir/test/lib/Dialect/Linalg/TestLinalgCodegenStrategy.cpp

Show First 20 Lines • Show All 137 Lines • ▼ Show 20 Linesstruct TestLinalgCodegenStrategy
Option<std::string> anchorFuncOpName{ Option<std::string> anchorFuncOpName{
*this, "anchor-func", *this, "anchor-func",
llvm::cl::desc( llvm::cl::desc(
"Which single func op is the anchor for the codegen strategy to " "Which single func op is the anchor for the codegen strategy to "
"latch on."), "latch on."),
llvm::cl::init("")}; llvm::cl::init("")};
}; };
// For now, just assume it is the zero of type.
// In the future, it should be the zero of type + op.
static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
auto t = getElementTypeOrSelf(op.get());
return b.create<arith::ConstantOp>(op.getOwner()->getLoc(), t,
b.getZeroAttr(t));
}
void TestLinalgCodegenStrategy::runStrategy( void TestLinalgCodegenStrategy::runStrategy(
LinalgTilingOptions tilingOptions, LinalgTilingOptions tilingOptions,
LinalgTilingOptions registerTilingOptions, LinalgTilingOptions registerTilingOptions,
LinalgPaddingOptions paddingOptions, LinalgPaddingOptions paddingOptions,
vector::VectorContractLowering vectorContractLowering, vector::VectorContractLowering vectorContractLowering,
vector::VectorTransferSplit vectorTransferSplit) { vector::VectorTransferSplit vectorTransferSplit) {
assert(!anchorOpName.empty()); assert(!anchorOpName.empty());
CodegenStrategy strategy; CodegenStrategy strategy;
Show All 29 Linesvoid TestLinalgCodegenStrategy::runStrategy(
FuncOp funcOp = getFunction(); FuncOp funcOp = getFunction();
OpPassManager dynamicPM("builtin.func"); OpPassManager dynamicPM("builtin.func");
strategy.configurePassPipeline(dynamicPM, funcOp.getContext()); strategy.configurePassPipeline(dynamicPM, funcOp.getContext());
if (failed(runPipeline(dynamicPM, funcOp))) if (failed(runPipeline(dynamicPM, funcOp)))
return signalPassFailure(); return signalPassFailure();
} }
} // end anonymous namespace } // end anonymous namespace
// For now, just assume it is the zero of type.
// In the future, it should be the zero of type + op.
static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
auto t = getElementTypeOrSelf(op.get());
return b.create<arith::ConstantOp>(op.getOwner()->getLoc(), t,
b.getZeroAttr(t));
}
/// Apply transformations specified as patterns. /// Apply transformations specified as patterns.
void TestLinalgCodegenStrategy::runOnFunction() { void TestLinalgCodegenStrategy::runOnFunction() {
if (!anchorFuncOpName.empty() && anchorFuncOpName != getFunction().getName()) if (!anchorFuncOpName.empty() && anchorFuncOpName != getFunction().getName())
return; return;
LinalgTilingOptions tilingOptions; LinalgTilingOptions tilingOptions;
if (!tileSizes.empty()) if (!tileSizes.empty())
tilingOptions = tilingOptions.setTileSizes(tileSizes); tilingOptions = tilingOptions.setTileSizes(tileSizes);
▲ Show 20 Lines • Show All 49 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

Show First 20 Lines • Show All 87 Lines • ▼ Show 20 Lines Option<bool> testGenericToVectorPattern{
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testTilePattern{*this, "test-tile-pattern", Option<bool> testTilePattern{*this, "test-tile-pattern",
llvm::cl::desc("Test tile pattern"), llvm::cl::desc("Test tile pattern"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testTileScalarizeDynamicDims{ Option<bool> testTileScalarizeDynamicDims{
*this, "test-tile-scalarize-dynamic-dims", *this, "test-tile-scalarize-dynamic-dims",
llvm::cl::desc("Test tiling of dynamic dims by 1"), llvm::cl::desc("Test tiling of dynamic dims by 1"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<int> testHoistPadding{*this, "test-hoist-padding",
llvm::cl::desc("Test hoist padding"),
llvm::cl::init(0)};
Option<bool> testPadPattern{*this, "test-pad-pattern", Option<bool> testPadPattern{*this, "test-pad-pattern",
llvm::cl::desc("Test pad pattern"), llvm::cl::desc("Test pad pattern"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testTransformPadTensor{ Option<bool> testTransformPadTensor{
*this, "test-transform-pad-tensor", *this, "test-transform-pad-tensor",
llvm::cl::desc("Test transform pad tensor by copying with generic ops"), llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testGeneralizePadTensor{ Option<bool> testGeneralizePadTensor{
*this, "test-generalize-pad-tensor", *this, "test-generalize-pad-tensor",
llvm::cl::desc("Test transform pad tensor by copying with generic ops"), llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
llvm::cl::init(false)}; llvm::cl::init(false)};
Option<bool> testSwapSubTensorPadTensor{ Option<bool> testSwapSubTensorPadTensor{
*this, "test-swap-subtensor-padtensor", *this, "test-swap-subtensor-padtensor",
llvm::cl::desc("Test rewrite of subtensor(pad_tensor) into " llvm::cl::desc("Test rewrite of subtensor(pad_tensor) into "
"pad_tensor(subtensor)"), "pad_tensor(subtensor)"),
llvm::cl::init(false)}; llvm::cl::init(false)};
ListOption<int64_t> paddedOperands{
*this, "padded-operands",
llvm::cl::desc("Operands to pad when test-tile-pattern"),
llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
ListOption<int64_t> nofoldOperands{
*this, "nofold-operands",
llvm::cl::desc("Operands to set nofold when test-tile-pattern"),
llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
ListOption<int64_t> packPaddings{ ListOption<int64_t> packPaddings{
*this, "pack-paddings", *this, "pack-paddings",
llvm::cl::desc("Operand packing flags when test-pad-pattern"), llvm::cl::desc("Operand packing flags when test-pad-pattern"),
llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated}; llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
ListOption<int64_t> hoistPaddings{ ListOption<int64_t> hoistPaddings{
*this, "hoist-paddings", *this, "hoist-paddings",
llvm::cl::desc("Operand hoisting depths when test-pad-pattern"), llvm::cl::desc("Operand hoisting depths when test-pad-pattern"),
llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated}; llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
▲ Show 20 Lines • Show All 479 Lines • ▼ Show 20 Lines
static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) { static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
auto t = getElementTypeOrSelf(op.get()); auto t = getElementTypeOrSelf(op.get());
return b.create<arith::ConstantOp>(op.getOwner()->getLoc(), t, return b.create<arith::ConstantOp>(op.getOwner()->getLoc(), t,
b.getZeroAttr(t)); b.getZeroAttr(t));
} }
static void applyTilePattern(FuncOp funcOp, std::string loopType, static void applyTilePattern(FuncOp funcOp, std::string loopType,
ArrayRef<int64_t> tileSizes, ArrayRef<int64_t> tileSizes,
ArrayRef<int64_t> paddedOperands,
ArrayRef<int64_t> nofoldOperands,
ArrayRef<int64_t> peeledLoops, ArrayRef<int64_t> peeledLoops,
bool scalarizeDynamicDims) { bool scalarizeDynamicDims) {
MLIRContext *context = funcOp.getContext(); MLIRContext *context = funcOp.getContext();
RewritePatternSet tilingPattern(context); RewritePatternSet tilingPattern(context);
LinalgTilingLoopType type = LinalgTilingLoopType type =
llvm::StringSwitch<LinalgTilingLoopType>(loopType) llvm::StringSwitch<LinalgTilingLoopType>(loopType)
.Case("for", LinalgTilingLoopType::Loops) .Case("for", LinalgTilingLoopType::Loops)
.Case("affine", LinalgTilingLoopType::AffineLoops) .Case("affine", LinalgTilingLoopType::AffineLoops)
.Case("parallel", LinalgTilingLoopType::ParallelLoops) .Case("parallel", LinalgTilingLoopType::ParallelLoops)
.Case("tiled_loop", LinalgTilingLoopType::TiledLoops); .Case("tiled_loop", LinalgTilingLoopType::TiledLoops);
auto linalgTilingOptions = linalg::LinalgTilingOptions() auto linalgTilingOptions = linalg::LinalgTilingOptions()
.setPeeledLoops(peeledLoops) .setPeeledLoops(peeledLoops)
.setLoopType(type); .setLoopType(type);
if (scalarizeDynamicDims) { if (scalarizeDynamicDims) {
linalgTilingOptions.scalarizeDynamicDims(); linalgTilingOptions.scalarizeDynamicDims();
assert(tileSizes.empty() && assert(tileSizes.empty() &&
"tileSizes and scalarizeDynamicDims is mutually exclusive"); "tileSizes and scalarizeDynamicDims is mutually exclusive");
} else { } else {
linalgTilingOptions.setTileSizes(tileSizes); linalgTilingOptions.setTileSizes(tileSizes);
} }
if (!paddedOperands.empty()) {
auto paddingFunc = [&](OpBuilder &b,
OpOperand &opOperand) -> FailureOr<Value> {
if (llvm::count(paddedOperands, opOperand.getOperandNumber()) == 0)
return failure();
return getNeutralOfLinalgOp(b, opOperand);
};
auto nofoldFunc = [&](OpOperand &opOperand) {
if (llvm::count(nofoldOperands, opOperand.getOperandNumber()) != 0)
return true;
return false;
};
linalgTilingOptions.setPaddingValueComputationFunction(paddingFunc);
linalgTilingOptions.setPaddingNoFoldComputationFunction(nofoldFunc);
}
tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulOp>, tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulOp>,
linalg::LinalgTilingPattern<linalg::GenericOp>>( linalg::LinalgTilingPattern<linalg::GenericOp>>(
context, linalgTilingOptions, context, linalgTilingOptions,
linalg::LinalgTransformationFilter(Identifier::get("tile", context))); linalg::LinalgTransformationFilter(Identifier::get("tile", context)));
(void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern)); (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
} }
static void applyPadPattern(FuncOp funcOp, ArrayRef<int64_t> packPaddings, static void applyPadPattern(FuncOp funcOp, ArrayRef<int64_t> packPaddings,
▲ Show 20 Lines • Show All 140 Lines • ▼ Show 20 Linesvoid TestLinalgTransforms::runOnFunction() {
if (testGeneralizePadTensor) if (testGeneralizePadTensor)
return applyGeneralizePadTensorPatterns(getFunction()); return applyGeneralizePadTensorPatterns(getFunction());
if (testSwapSubTensorPadTensor) if (testSwapSubTensorPadTensor)
return applyExtractSliceOfPadTensorSwapPattern(getFunction()); return applyExtractSliceOfPadTensorSwapPattern(getFunction());
if (testTiledLoopPeeling.hasValue()) if (testTiledLoopPeeling.hasValue())
return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling, return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling,
skipPartial); skipPartial);
if (testTilePattern) if (testTilePattern)
return applyTilePattern(getFunction(), loopType, tileSizes, paddedOperands, return applyTilePattern(getFunction(), loopType, tileSizes, peeledLoops,
nofoldOperands, peeledLoops,
/*scalarizeDynamicDims=*/false); /*scalarizeDynamicDims=*/false);
if (testTileScalarizeDynamicDims) if (testTileScalarizeDynamicDims)
return applyTilePattern(getFunction(), loopType, tileSizes, paddedOperands, return applyTilePattern(getFunction(), loopType, tileSizes,
nofoldOperands,
/*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true); /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
if (testHoistPadding) {
getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
PadTensorOp hoistedOp;
FailureOr<Value> newResult = linalg::hoistPaddingOnTensors(
padTensorOp, testHoistPadding, hoistedOp);
if (succeeded(newResult)) {
padTensorOp.getResult().replaceAllUsesWith(newResult.getValue());
padTensorOp->erase();
}
});
}
if (testPadPattern) if (testPadPattern)
return applyPadPattern(getFunction(), packPaddings, hoistPaddings); return applyPadPattern(getFunction(), packPaddings, hoistPaddings);
if (testInterchangePattern.hasValue()) if (testInterchangePattern.hasValue())
return applyInterchangePattern(getFunction(), testInterchangePattern); return applyInterchangePattern(getFunction(), testInterchangePattern);
if (testDecomposeConvolutionPattern) if (testDecomposeConvolutionPattern)
return applyDecomposeConvolutionPatterns(getFunction()); return applyDecomposeConvolutionPatterns(getFunction());
} }
namespace mlir { namespace mlir {
namespace test { namespace test {
void registerTestLinalgTransforms() { void registerTestLinalgTransforms() {
PassRegistration<TestLinalgTransforms>(); PassRegistration<TestLinalgTransforms>();
} }
} // namespace test } // namespace test
} // namespace mlir } // namespace mlir
mlir/test/Dialect/Linalg/pad.mlir