maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[MLIR][Linalg] Generate the right type of load/store when lowering max/min pooling ops 

While lowering min/max pooling ops to loops, generate the right kind of
load/stores (std or affine) instead of always generating std
load/stores.

Differential Revision: https://reviews.llvm.org/D83080
This commit is contained in:
Uday Bondhugula 2020-06-11 18:10:53 +05:30
parent 7356b4243a
commit 6d6d5db251
2 changed files with 37 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
View File

@ -333,23 +333,28 @@ static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
template <typename IndexedValueType>
void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingMaxOp op) {
auto indices = getInputAndOutputIndices(allIvs, op);
InputAndOutputIndices indices = getInputAndOutputIndices(allIvs, op);
// Emit scalar form.
Value lhs = std_load(op.output(), indices.outputs);
Value rhs = std_load(op.input(), indices.inputs);
IndexedValueType output(op.output());
IndexedValueType input(op.input());
Value lhs = output(indices.outputs);
Value rhs = input(indices.inputs);
using edsc::op::sgt;
Value maxValue = std_select(sgt(lhs, rhs), lhs, rhs);
std_store(maxValue, op.output(), indices.outputs);
output(indices.outputs) = maxValue;
}
template <typename IndexedValueType>
void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingMinOp op) {
auto indices = getInputAndOutputIndices(allIvs, op);
InputAndOutputIndices indices = getInputAndOutputIndices(allIvs, op);
// Emit scalar form.
Value lhs = std_load(op.output(), indices.outputs);
Value rhs = std_load(op.input(), indices.inputs);
IndexedValueType output(op.output());
IndexedValueType input(op.input());
Value lhs = output(indices.outputs);
Value rhs = input(indices.inputs);
using edsc::op::slt;
Value minValue = std_select(slt(lhs, rhs), lhs, rhs);
std_store(minValue, op.output(), indices.outputs);
output(indices.outputs) = minValue;
}
template <typename IndexedValueType>
void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingSumOp op) {

24
mlir/test/Dialect/Linalg/affine.mlir
View File

@ -123,3 +123,27 @@ func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memre
// CHECK: %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
// CHECK: %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
// CHECK: affine.store %[[res]], %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
// CHECK-LABEL: func @pooling_max_min
func @pooling_max_min(%arg0: memref<?x?xf32>,
%arg1: memref<?x?xi32>,
%arg2: memref<?x?xf32>) {
linalg.pooling_max(%arg0, %arg1, %arg2) { strides = [2, 1] }:
memref<?x?xf32>, memref<?x?xi32>, memref<?x?xf32>
linalg.pooling_min(%arg0, %arg1, %arg2) { strides = [2, 1] }:
memref<?x?xf32>, memref<?x?xi32>, memref<?x?xf32>
return
}
// This is a basic check to make sure the right load/stores are used. loops.mlir
// checks for the rest.
// CHECK: affine.load
// CHECK-NEXT: affine.load
// CHECK-NEXT: cmpf
// CHECK-NEXT: select
// CHECK-NEXT: affine.store
// The min pooling body.
// CHECK: affine.load
// CHECK-NEXT: affine.load
// CHECK-NEXT: cmpf
// CHECK-NEXT: select
// CHECK-NEXT: affine.store