This is an archive of the discontinued LLVM Phabricator instance.

Differential D97385

[mlir][linalg] Add pooling ops to Linalg TC ops.
ClosedPublic

Authored by hanchung on Feb 24 2021, 7:13 AM.

Details

Reviewers
antiagainst
nicolasvasilache
Commits
rGbfd3771c9eec: [mlir][linalg] Add pooling ops to Linalg TC ops.
Summary
  • Add EDSC builders for std_cmpf_ogt and std_cmpf_olt.
  • Add pooling_nhwc_min/max/sum ops

Depends On D97384

Diff Detail

Event Timeline

hanchung created this revision.Feb 24 2021, 7:13 AM
Herald added subscribers: mravishankar, teijeong, rdzhabarov and 14 others. · View Herald TranscriptFeb 24 2021, 7:13 AM
hanchung requested review of this revision.Feb 24 2021, 7:13 AM
Herald added a project: Restricted Project. · View Herald TranscriptFeb 24 2021, 7:13 AM
Herald added subscribers: limo1996, stephenneuendorffer. · View Herald Transcript
hanchung updated this revision to Diff 326086.Feb 24 2021, 7:15 AM

remove -canonicalize pass from the test

Harbormaster completed remote builds in B90607: Diff 326085.Feb 24 2021, 8:18 AM
Harbormaster completed remote builds in B90608: Diff 326086.Feb 24 2021, 8:37 AM
hanchung updated this revision to Diff 326153.Feb 24 2021, 11:13 AM

rebase

Harbormaster completed remote builds in B90651: Diff 326153.Feb 24 2021, 1:02 PM
Herald added a subscriber: cota. · View Herald TranscriptFeb 24 2021, 1:02 PM
hanchung edited the summary of this revision. (Show Details)Feb 26 2021, 6:54 AM
nicolasvasilache accepted this revision.Feb 26 2021, 7:15 AM
This revision is now accepted and ready to land.Feb 26 2021, 7:15 AM
Closed by commit rGbfd3771c9eec: [mlir][linalg] Add pooling ops to Linalg TC ops. (authored by hanchung). · Explain WhyFeb 26 2021, 7:18 AM
This revision was automatically updated to reflect the committed changes.
hanchung added a commit: rGbfd3771c9eec: [mlir][linalg] Add pooling ops to Linalg TC ops..

Revision Contents

PathSize
mlir/
include/
mlir/
Dialect/
Linalg/
IR/
41 lines
StandardOps/
EDSC/
10 lines
test/
Dialect/
Linalg/
77 lines
102 lines

Diff 326685

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

Show First 20 Lines • Show All 323 Lines • ▼ Show 20 Lines
""" """
{ {
O(n, f, d, h, w) = std_addf<kd, kh, kw>( O(n, f, d, h, w) = std_addf<kd, kh, kw>(
O(n, f, d, h, w), std_mulf(I(n, c, d * strides[0] + kd * dilations[0], O(n, f, d, h, w), std_mulf(I(n, c, d * strides[0] + kd * dilations[0],
h * strides[1] + kh * dilations[1], h * strides[1] + kh * dilations[1],
w * strides[2] + kw * dilations[2]), w * strides[2] + kw * dilations[2]),
K(kd, kh, kw, c, f))); K(kd, kh, kw, c, f)));
} }
ods_def<PoolingNHWCSum>:
def pooling_nhwc_sum
(I: f32(N, H, W, C), K: f32(KH, KW))
-> (O: f32(N, OH, OW, C))
attr(strides: 2xi64, dilations: 2xi64)
{
O(n, oh, ow, c) = std_addf<kh, kw>(O(n, oh, ow, c),
I(n, oh * strides[0] + kh * dilations[0],
ow * strides[1] + kw * dilations[1], c));
}
ods_def<PoolingNHWCMax>:
def pooling_nhwc_max
(I: f32(N, H, W, C), K: f32(KH, KW))
-> (O: f32(N, OH, OW, C))
attr(strides: 2xi64, dilations: 2xi64)
{
O(n, oh, ow, c) =
std_select<kh, kw>(std_cmpf_ogt(I(n, oh * strides[0] + kh * dilations[0],
ow * strides[1] + kw * dilations[1], c),
O(n, oh, ow, c)),
I(n, oh * strides[0] + kh * dilations[0],
ow * strides[1] + kw * dilations[1], c),
O(n, oh, ow, c));
}
ods_def<PoolingNHWCMin>:
def pooling_nhwc_min
(I: f32(N, H, W, C), K: f32(KH, KW))
-> (O: f32(N, OH, OW, C))
attr(strides: 2xi64, dilations: 2xi64)
{
O(n, oh, ow, c) =
std_select<kh, kw>(std_cmpf_olt(I(n, oh * strides[0] + kh * dilations[0],
ow * strides[1] + kw * dilations[1], c),
O(n, oh, ow, c)),
I(n, oh * strides[0] + kh * dilations[0],
ow * strides[1] + kw * dilations[1], c),
O(n, oh, ow, c));
}

mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h

Show First 20 Lines • Show All 54 Lines • ▼ Show 20 Linesstruct SExtiValueBuilder : public ValueBuilder<SignExtendIOp> {
template <typename... Args> template <typename... Args>
SExtiValueBuilder(Args... args) SExtiValueBuilder(Args... args)
: ValueBuilder<SignExtendIOp>(ScopedContext::getBuilderRef().getI32Type(), : ValueBuilder<SignExtendIOp>(ScopedContext::getBuilderRef().getI32Type(),
args...) {} args...) {}
}; };
using std_sexti32 = SExtiValueBuilder<32>; using std_sexti32 = SExtiValueBuilder<32>;
template <CmpFPredicate Predicate>
struct CmpFValueBuilder : public ValueBuilder<CmpFOp> {
using ValueBuilder<CmpFOp>::ValueBuilder;
template <typename... Args>
CmpFValueBuilder(Args... args) : ValueBuilder<CmpFOp>(Predicate, args...) {}
};
using std_cmpf_ogt = CmpFValueBuilder<CmpFPredicate::OGT>;
using std_cmpf_olt = CmpFValueBuilder<CmpFPredicate::OLT>;
/// Branches into `block` with `operands`. /// Branches into `block` with `operands`.
BranchOp std_br(Block *block, ValueRange operands); BranchOp std_br(Block *block, ValueRange operands);
/// Branches into `trueBranch` with `trueOperands` if `cond` evaluates to `true` /// Branches into `trueBranch` with `trueOperands` if `cond` evaluates to `true`
/// or to `falseBranch` and `falseOperand` if `cond` evaluates to `false`. /// or to `falseBranch` and `falseOperand` if `cond` evaluates to `false`.
CondBranchOp std_cond_br(Value cond, Block *trueBranch, ValueRange trueOperands, CondBranchOp std_cond_br(Value cond, Block *trueBranch, ValueRange trueOperands,
Block *falseBranch, ValueRange falseOperands); Block *falseBranch, ValueRange falseOperands);
/// Provide an index notation around sdt_load and std_store. /// Provide an index notation around sdt_load and std_store.
using StdIndexedValue = using StdIndexedValue =
TemplatedIndexedValue<intrinsics::std_load, intrinsics::std_store>; TemplatedIndexedValue<intrinsics::std_load, intrinsics::std_store>;
} // namespace intrinsics } // namespace intrinsics
} // namespace edsc } // namespace edsc
} // namespace mlir } // namespace mlir
#endif // MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_ #endif // MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_

mlir/test/Dialect/Linalg/generalize-named-ops.mlir

Show First 20 Lines • Show All 253 Lines • ▼ Show 20 Lines
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>) // CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK-NEXT: %[[MUL:.+]] = mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK-NEXT: %[[ADD:.+]] = addf %[[BBARG2]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = addf %[[BBARG2]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// -----
func @pooling_nhwc_sum(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
outs(%init: memref<?x?x?x?xf32>)
return
}
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
// CHECK: func @pooling_nhwc_sum
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[RES:.+]] = addf %[[BBARG2]], %[[BBARG0]] : f32
// CHECK-NEXT: linalg.yield %[[RES]] : f32
// -----
func @pooling_nhwc_max(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>}
ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
outs(%init: memref<?x?x?x?xf32>)
return
}
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 3 + d5, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
// CHECK: func @pooling_nhwc_max
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[CMP:.+]] = cmpf ogt, %[[BBARG0]], %[[BBARG2]] : f32
// CHECK-NEXT: %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
// CHECK-NEXT: linalg.yield %[[RES]] : f32
// -----
func @pooling_nhwc_min(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
linalg.pooling_nhwc_min {dilations = dense<3> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
outs(%init: memref<?x?x?x?xf32>)
return
}
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4 * 3, d2 * 2 + d5 * 3, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
// CHECK: func @pooling_nhwc_min
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[CMP:.+]] = cmpf olt, %[[BBARG0]], %[[BBARG2]] : f32
// CHECK-NEXT: %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
// CHECK-NEXT: linalg.yield %[[RES]] : f32

mlir/test/Dialect/Linalg/named-ops.mlir

Show First 20 Lines • Show All 240 Lines • ▼ Show 20 Linesfunc @conv_3d_input_ncdhw_filter_dhwcf(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) {
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>) // CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>)
linalg.conv_3d_input_ncdhw_filter_dhwcf {dilations = dense<1> : tensor<3xi64>, linalg.conv_3d_input_ncdhw_filter_dhwcf {dilations = dense<1> : tensor<3xi64>,
strides = dense<1> : tensor<3xi64>} strides = dense<1> : tensor<3xi64>}
ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
outs (%output: memref<?x?x?x?x?xf32>) outs (%output: memref<?x?x?x?x?xf32>)
return return
} }
// -----
// CHECK-LABEL: func @pooling_nhwc_sum_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_sum
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x6x6x1xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
func @pooling_nhwc_sum_tensor(%input: tensor<1x6x6x1xf32>) -> tensor<1x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32>
%cst = constant 0.000000e+00 : f32
%fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xf32>, f32 -> tensor<1x2x2x1xf32>
%res = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x6x6x1xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32>
}
// -----
// CHECK-LABEL: func @pooling_nhwc_sum
// CHECK: linalg.pooling_nhwc_sum
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x6x6x1xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xf32>)
func @pooling_nhwc_sum(%input: memref<1x6x6x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x6x6x1xf32>, memref<3x3xf32>)
outs(%output: memref<1x2x2x1xf32>)
return
}
// -----
// CHECK-LABEL: func @pooling_nhwc_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x6x6x1xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
func @pooling_nhwc_max_tensor(%input: tensor<1x6x6x1xf32>) -> tensor<1x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32>
%cst = constant 0.000000e+00 : f32
%fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xf32>, f32 -> tensor<1x2x2x1xf32>
%res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x6x6x1xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32>
}
// -----
// CHECK-LABEL: func @pooling_nhwc_max
// CHECK: linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x6x6x1xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xf32>)
func @pooling_nhwc_max(%input: memref<1x6x6x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x6x6x1xf32>, memref<3x3xf32>)
outs(%output: memref<1x2x2x1xf32>)
return
}
// -----
// CHECK-LABEL: func @pooling_nhwc_min_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_min
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x6x6x1xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
func @pooling_nhwc_min_tensor(%input: tensor<1x6x6x1xf32>) -> tensor<1x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32>
%cst = constant 0.000000e+00 : f32
%fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xf32>, f32 -> tensor<1x2x2x1xf32>
%res = linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x6x6x1xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32>
}
// -----
// CHECK-LABEL: func @pooling_nhwc_min
// CHECK: linalg.pooling_nhwc_min
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x6x6x1xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xf32>)
func @pooling_nhwc_min(%input: memref<1x6x6x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x6x6x1xf32>, memref<3x3xf32>)
outs(%output: memref<1x2x2x1xf32>)
return
}