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

[bondhugula]

Summary

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.

[bondhugula]

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

33

Open to suggestions on the best place to move these to.

[bondhugula]

Minor: drop extra space

[ftynse]

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

33

Hmm, I would think that IndexedTy would somehow have the correct included. Would something like

IndexedValueTy output(op.output());
IndexedValueTy input(op.input());
Value lhs = output(indices.outputs);
Value rhs = input(indices.inputs);
//...
output(indices.outputs) = maxValue;

inside the function just work?

[bondhugula]

That works as well, thanks.

[bondhugula]

Update per IndexValueType ctor.

[ftynse]

Thanks for the fix!

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

 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 using edsc::op::operator+;
 
 static SmallVector<Value, 8> makeCanonicalAffineApplies(OpBuilder &b,
                                                         Location loc,
                                                         AffineMap map,
                                                         ArrayRef<Value> vals) {
   if (map.isEmpty())
     return {};
   assert(map.getNumSymbols() == 0);
   assert(map.getNumInputs() == vals.size());
   SmallVector<Value, 8> res;
static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
   } else {
     IndexedValueType I(convOp.input());
     O(oIdx) += F(fIdx) * I(imIdx);
   }
 }
 
 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) {
   auto indices = getInputAndOutputIndices(allIvs, op);
   IndexedValueType input(op.input()), output(op.output());
 
   // Emit scalar form.
   output(indices.outputs) += input(indices.inputs);

mlir/test/Dialect/Linalg/affine.mlir

 //       CHECK:     affine.for %[[n:.*]] = 0 to %[[N]] {
 //       CHECK:       affine.for %[[k:.*]] = 0 to %[[K]] {
 //       CHECK:       %[[va:.*]] = affine.load %[[mA]][%[[b]], %[[m]], %[[k]]] : memref<?x?x?xf32>
 //       CHECK:       %[[vb:.*]] = affine.load %[[mB]][%[[b]], %[[k]], %[[n]]] : memref<?x?x?xf32>
 //       CHECK:       %[[vc:.*]] = affine.load %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
 //       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