[mlir][Linalg] Introduce linalg.pooling_min/max/sum op.

[hanchung]

Summary

Performs an N-D pooling operation similarly to the description in the TF
documentation:
https://www.tensorflow.org/api_docs/python/tf/nn/pool

Different from the description, this operation doesn't perform on batch and
channel. It only takes tensors of rank N.

output[x[0], ..., x[N-1]] =
  REDUCE_{z[0], ..., z[N-1]}
    input[
          x[0] * strides[0] - pad_before[0] + dilation_rate[0]*z[0],
          ...
          x[N-1]*strides[N-1] - pad_before[N-1] + dilation_rate[N-1]*z[N-1]
          ],

The required optional arguments are:

• strides: an i64 array specifying the stride (i.e. step) for window loops.
• dilations: an i64 array specifying the filter upsampling/input downsampling rate
• padding: an i64 array of pairs (low, high) specifying the number of elements to pad along a dimension.

If strides or dilations attributes are missing then the default value is
one for each of the input dimensions. Similarly, padding values are zero
for both low and high in each of the dimensions, if not specified.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

273

I would keep the libraryCallName #.
If / when it changes in the future we won't have to track it down.

444

If we go through the trouble of creating a new "named op" for pooling, I would not specify a region.

Instead I would have a linalg.pooling_max, linalg.pooling_min, linalg.pooling.sum, ... with no regions that just apply the proper operation when lowering to loops.

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

527

let's use

//clang-format off
...
//clang-format on

and 1 pattern per line if you don't mind, so we have fewer reflowing diff in the longer term.

[mehdi_amini]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

410

The doc should define what is a "pooling function" and what does it mean to "pool"

425

Please remove username from TODO (this is a google internal practice).

Also note that this doc will be automatically published on the website when you push this.

500

This is quite a block of C++, can you move this all out of ODS?

mlir/test/Dialect/Linalg/loops.mlir

259

Have you looked into the syntax of other operations (like the GPU launch or for loops) which elide the block and integrate the block arg in the syntax itself?

[mravishankar]

Nice. Few comments here

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

410

You should probably use this description (https://www.tensorflow.org/api_docs/python/tf/nn/pool) . This is similar to how the conv op is defined.

444

wouldnt that be a new named op for every pooling operation. This seems better to me.

474

That seems strange. Looking at the documentation (https://www.tensorflow.org/api_docs/python/tf/nn/pool) this doesnt seem to be the case.

[hanchung]

Address comments.

[hanchung]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

410

Yes, this is a good reference. :)

444

Both are okay to me. I think maybe a "named op" should have more concrete behavior. Otherwise it's like a generic op.

474

I think it's should be the case:
window_shape: Sequence of N ints >= 1

Also, if this is not the case, we must have another attributes to specify the corresponding dimensions.

500

I'm following other ops style here. I'm happy to move all of them out of ODS, but can we keep it as this way for now?

Nicolas, do you have any comment/suggestion here?

mlir/test/Dialect/Linalg/loops.mlir

259

Good to know it! I looked into the syntax and thought it's much better than this way. However, the op doesn't take a region now. I'll follow it if I need a region next time. Thank you!

[hanchung]

Fix clang-tidy.

I was following the naming sytle around but it seemed that the namings should all be updated...

[mravishankar]

LGTM. Please wait for @nicolasvasilache to accept.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

958

NIT: Just add a macro for the indirection.

You could achieve the same thing by making the weightedPoolingInputIndex function extern template function.

So make the definition

template <typename PoolingOp>
extern SmallVector<AffineExpr, 4> weightedPoolingInputIndex(PoolingOp op, ...);

and make this

template <typename PoolingOp>
SmallVector<AffineExpr, 4> weightedPolingInputIndex(PoolingOp op, ...) { .. }

#define INSTANTIATE_WEIGHTED_POOLING_INPUT_INDEX(opname) \
template SmallVector<AffineExpr, 4> weightedPoolingInputIndex<opname>(opname op, ...);

INSTANTIATE_WEIGHTED_POOLING_INPUT_INDEX(PoolingMinOp);
...

[nicolasvasilache]

Thanks for adding this!

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

500

A lot of these things will be auto-generated soon.
I have been personally moving things out of C++ and into Tablegen to reduce the number of places we need to modify.
I don't have a particular suggestion.

[hanchung]

Address comments

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

500

This might should be auto-generated. Let's keep it this way, I'm happy to help on it.

[hanchung]

Fix description in SingleInputPoolingBase_Op

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/SideEffects.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 namespace linalg {
 
 class ConvOp;
+class PoolingOp;
 
 /// Returns the name mangled library call name to disambiguate between different
 /// overloads at the C level. The name mangling scheme is basic and uses MLIR
 /// type names:
 ///   1. form a string which is the concatenation of the linalg op name with all
 ///      the operand type names, separate by underscores;
 ///   2. drop the `linalg.` prefix, and the `<`, `>`, `?` symbols from the type.
 /// Assumes `op` is a LinalgOp.

 ///   name mangles into `linalg_matmul_viewxxf32_viewxxf32_viewxxf32_impl`
 std::string generateLibraryCallName(Operation *op);
 
 /// Returns `num` AffineDimExpr dimensions at positions
 ///   [startIdx, startIdx + num) and increments `startIdx` to `startIdx + num`.
 SmallVector<AffineExpr, 4> makeAffineDimExprs(unsigned num, unsigned &startIdx,
                                               MLIRContext *context);
 
-/// Builds the indexing expressions for a ConvOp `op`. Returns the vector of
-/// AffineMaps representing:
-///   `stride[i] * xs[i] + dilation[i] * zs[i] - pad_low[i]`
-SmallVector<AffineExpr, 4> weightedConvInputIndex(ConvOp op,
-                                                  ArrayRef<AffineExpr> xs,
-                                                  ArrayRef<AffineExpr> zs);
+/// Builds the indexing expressions for a ConvOp/PoolingOp `op`. Returns the
+/// vector of AffineMaps representing:
+///   `stride[i] * outputDims[i] + dilation[i] * windowDims[i] - pad_low[i]`
+SmallVector<AffineExpr, 4>
+weightedPoolingInputIndex(ConvOp op, ArrayRef<AffineExpr> outputDims,
+                          ArrayRef<AffineExpr> windowDims);
+SmallVector<AffineExpr, 4>
+weightedPoolingInputIndex(PoolingOp op, ArrayRef<AffineExpr> outputDims,
+                          ArrayRef<AffineExpr> windowDims);
 
 /// Returns `maybeMap.get()` if `maybeMap` is set, otherwise returns the
 /// symbol-less identity map of `rank`.
 AffineMap extractOrIdentityMap(Optional<AffineMap> maybeMap, unsigned rank,
                                MLIRContext *context);
 
 /// Return the vector that is the concatenation of `a` and `b`.
 SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
                                        AffineMap::get(3, 0, {r_k, j}),
                                        AffineMap::get(3, 0, {i, j})};
     }
   }];
 
   let hasFolder = 1;
 }
 
-def ConvOp : LinalgStructured_Op<"conv", [NInputs<2>, NOutputs<1>]> {
+/// A base class for pooling operation such as conv. The arguments must contain
+/// optional arguments `strides`, `dilations` and `padding` with following type:
+///   OptionalAttr<I64ArrayAttr>:$strides
+///   OptionalAttr<I64ArrayAttr>:$dilations
+///   OptionalAttr<I64ElementsAttr>:$padding
+/// `stirdes` denotes the step of each window along the dimension.
+class PoolingBase_Op<string mnemonic, list<OpTrait> props>
+  : LinalgStructuredBase_Op<mnemonic, props> {
+  let description = [{
+    The optional arguments are:
+      - strides: an i64 array specifying the stride (i.e. step) for window
+        loops.
+      - dilations: an i64 array specifying the filter upsampling/input
+        downsampling rate
+      - padding: an i64 array of pairs (low, high) specifying the number of
+        elements to pad along a dimension.
+
+    If strides or dilations attributes are missing then the default value is
+    one for each of the input dimensions. Similarly, padding values are zero
+    for both low and high in each of the dimensions, if not specified.
+  }];
+
+  code commonUtils = [{
+    std::string getLibraryCallName() {
+      return generateLibraryCallName(getOperation());
+    }
+
+    int64_t getStride(unsigned i) {
+      assert(i < getNumWindowLoops());
+      if (!strides().hasValue()) return 1;
+      return strides()->getValue()[i]
+        .cast<IntegerAttr>().getValue().getSExtValue();
+    }
+
+    int64_t getDilation(unsigned i) {
+      assert(i < getNumWindowLoops());
+      if (!dilations().hasValue()) return 1;
+      return dilations()->getValue()[i]
+        .cast<IntegerAttr>().getValue().getSExtValue();
+    }
+
+    int64_t getLowPad(unsigned i) {
+      assert(i < getNumWindowLoops());
+      if (!padding().hasValue()) return 0;
+      return padding().getValue().getValue<int64_t>({i, 0});
+    }
+  }];
+}
+
+def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
 
   let description = [{
     Generic n-D convolution as described in the TF documentation:
     https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/nn/convolution
 
     ```
       output[b, x[0], ..., x[N-1], k] =
       sum_{z[0], ..., z[N-1], q}
def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
   // element. Using padding is equivalent to inserting those same zero values
   // into the input before doing the convolution.
   let arguments = (ins AnyStridedMemRef:$filter, AnyStridedMemRef:$input,
                    AnyStridedMemRef:$output,
                    OptionalAttr<I64ArrayAttr>:$strides,
                    OptionalAttr<I64ArrayAttr>:$dilations,
                    OptionalAttr<I64ElementsAttr>:$padding);
 
-  let extraClassDeclaration = libraryCallName # [{
+  let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
+
+  let extraClassDeclaration = commonUtils # [{
     // TODO(ntv) extend to support more than 1 dimensions and potentially
     // grouping too.
     unsigned getNumBatchDimensions() { return 1; }
 
     unsigned getNumInputFeatureDimensions() { return 1; }
 
     unsigned getNumOutputFeatureDimensions() { return 1; }
 
llvm::Optional<SmallVector<StringRef, 8>> referenceIterators() {
         nPar - getNumBatchDimensions() - getNumInputFeatureDimensions();
       SmallVector<StringRef, 8> iters(nPar, getParallelIteratorTypeName());
       iters.reserve(nPar + nRed + nWin);
       iters.append(nRed, getReductionIteratorTypeName());
       iters.append(nWin, getWindowIteratorTypeName());
       return iters;
     }
 
-    int64_t getStride(unsigned i) {
-      assert(i < getNumWindowLoops());
-      if (!strides().hasValue()) return 1;
-      return strides()->getValue()[i]
-        .cast<IntegerAttr>().getValue().getSExtValue();
-    }
-
-    int64_t getDilation(unsigned i) {
-      assert(i < getNumWindowLoops());
-      if (!dilations().hasValue()) return 1;
-      return dilations()->getValue()[i]
-        .cast<IntegerAttr>().getValue().getSExtValue();
-    }
-
-    int64_t getLowPad(unsigned i) {
-      assert(i < getNumWindowLoops());
-      if (!padding().hasValue()) return 0;
-      return padding().getValue().getValue<int64_t>({i, 0});
-    }
-
     //   F(z0, ..., zN-1, q, k) *
     //     I(b, x0 + z0 - pad_low_0, ..., xN-1 + zN-1 - pad_low_N-1, q)
     //   ->  O(b, x0, ..., xN-1, k)
     // for N equal to `nWindow`. If there is no padding attirbute, it will be
     // ignored.
     llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       MLIRContext *context = getContext();
       auto nWin = getNumWindowLoops();
llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       auto ks = makeAffineDimExprs(
         getNumOutputFeatureDimensions(), idx, context);
       // Non-window reduction dim: sum_{z[0], ..., z[N-1], q}
       auto qs = makeAffineDimExprs(
         getNumInputFeatureDimensions(), idx, context);
       // Window reduction dims: sum_{z[0], ..., z[N-1], q}
       auto zs = makeAffineDimExprs(nWin, idx, context);
       // Construct the weighedSum expression.
-      auto ws = weightedConvInputIndex(*this, xs, zs);
+      auto ws = weightedPoolingInputIndex(*this, xs, zs);
       return SmallVector<AffineMap, 8>{
         // filter[z[0], ..., z[N-1], q, k]
         AffineMap::get(idx, 0, concat(concat(zs, qs), ks)),
         // input[b,
         //       x[0]*s[0] + d[0]*z[0] - pad_low[0],
         //       ...
         //       x[N-1]*s[N-1] + d[N-1]*z[N-1] - pad_low[N-1],
         //       q]
         AffineMap::get(idx, 0, concat(concat(bs, ws), qs)),
         // output[b, x[0], ..., x[N-1], k]
         AffineMap::get(idx, 0, concat(concat(bs, xs), ks))};
     }
   }];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
+def PoolingOp: PoolingBase_Op<"pooling", [NInputs<2>, NOutputs<1>]> {
+  let description = [{
+    Applies a pooling function to all elements in each window of the input
+    multi-dimensional array, producing an output multi-dimensional array with
+    the same number of elements as the number of valid positions of the window.
+
+    A linalg.pooling is written as:
+    ```mlir
+        linalg.pooling #trait_attribute %input, %windowDims, %output {other-attributes} :
+          memref<?x?xf32, stride_specification>,
+          memref<?x?xf32, stride_specification>,
+          memref<?x?xf32, stride_specification>
+    ```
+
+    Where #trait_attributes is an alias of a dictionary attribute described in
+    PoolingBase_Op.
+
+    TODO(hanchung): Figure out a better way to handle window dimensions, i.e.,
+    eliminate the fake memref.
+    The window dimensions are specified by argument `windowDims`. The i-th
+    dimension in the shape of `windowDims` denotes the size of the window along
+    dimension i. For example, if the window size is 2x3, then a memref<2x3>
+    should be passed to the operation as `windowDims`.
+
+
+    linalg.pooling takes a region (like generic op) as the pooling function. The
+    signature of the block must be:
+
+    ```
+      ^bb0([input view element type], [output view element type])
+                  -> ([output view element type])
+    ```
+
+    Example:
+
+    ```mlir
+      linalg.pooling { strides = [2, 1, 2] } %arg0, %arg1, %arg2 {
+        ^bb(%a: f32, %b: f32) :
+          %0 = addf %a, %b : f32
+          linalg.yield %0 : f32
+      }: memref<?x?x?xf32>, memref<?x?x?xi32>, memref<?x?x?xf32>
+    ```
+
+  }];
+
+  // TODO(hanchung): Consider to make pooling op take a function as argument.
+  let arguments = (ins AnyStridedMemRef:$input,
+                   AnyStridedMemRef:$windowDims,
+                   AnyStridedMemRef:$output,
+                   OptionalAttr<I64ArrayAttr>:$strides,
+                   OptionalAttr<I64ArrayAttr>:$dilations,
+                   OptionalAttr<I64ElementsAttr>:$padding);
+
+  let regions = (region AnyRegion:$region);
+
+  let extraClassDeclaration = commonUtils# [{
+    SmallVector<StringRef, 8> linalgTraitAttrNames() {
+      return SmallVector<StringRef, 8>{
+        getStridesAttrName(), getDilationsAttrName(), getPaddingAttrName()
+      };
+    }
+
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIterators() {
+      // Outer parallel loops are always the number of output dimensions.
+      unsigned nPar = getOutputShapedType(0).getRank();
+      // The window loops has the same number loops with output dimensions.
+      unsigned nWin = nPar;
+      SmallVector<StringRef, 8> iters(nPar, getParallelIteratorTypeName());
+      iters.reserve(nPar + nWin);
+      iters.append(nWin, getWindowIteratorTypeName());
+      return iters;
+    }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      MLIRContext *context = getContext();
+      auto nPar = getNumParallelLoops();
+      auto nWin = getNumWindowLoops();
+      assert(nWin > 0 && "expected at least one window dimension");
+      unsigned idx = 0;
+      auto outputDims = makeAffineDimExprs(nPar, idx, context);
+      auto windowDims = makeAffineDimExprs(nWin, idx, context);
+      // Construct the weighedSum expression.
+      auto inputDims =
+          weightedPoolingInputIndex(*this, outputDims, windowDims);
+      return SmallVector<AffineMap, 8>{
+        // input
+        AffineMap::get(idx, 0, inputDims),
+        // windowDims
+        AffineMap::get(idx, 0, windowDims),
+        // output
+        AffineMap::get(idx, 0, outputDims)
+        };
+    }
+  }];
+
+  let verifier = [{ return ::verify(*this); }];
+
+  let printer = [{ return ::print(p, *this); }];
+
+  let parser = [{ return ::parseTraitAndRegion(parser, result); }];
+
+  let hasFolder = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Generic Linalg ops.
 //===----------------------------------------------------------------------===//
 def LinalgOperand: Type<
   Or<[AnyRankedTensor.predicate, AnyStridedMemRef.predicate]>>;
 
 class LinalgOperandOfRank<int rank>: Type<
   And<[
let extraClassDeclaration = [{
      }
 
     llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       llvm_unreachable(
         "No such thing as reference indexing maps for a generic op.");
      }
   }];
   let printer = [{ return ::print(p, *this); }];
-  let parser = [{ return ::parseGenericOp(parser, result); }];
+  let parser = [{ return ::parseTraitAndRegion(parser, result); }];
 }
 
 /// Index-free GenericOp.
 def GenericOp : GenericOpBase<"generic"> {
   let description = [{
     Generic Linalg op form where the key properties of the computation are
     specified as attributes. In pretty form, a linalg.generic op is written as:
 

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 /// Attribute name for the StrArrayAttr which encodes the SymbolAttr for the
 /// MLIR function that implements the body of the structured op.
 constexpr StringRef getFunAttrName() { return "fun"; }
 
 /// Attribute name for the StrArrayAttr which encodes the external library
 /// function that implements the structured op.
 constexpr StringRef getLibraryCallAttrName() { return "library_call"; }
 
+/// Attribute name for the StrArrayAttr which encodes the value of strides.
+constexpr StringRef getStridesAttrName() { return "strides"; }
+
+/// Attribute name for the StrArrayAttr which encodes the value of dilations.
+constexpr StringRef getDilationsAttrName() { return "dilations"; }
+
+/// Attribute name for the StrArrayAttr which encodes the value of paddings.
+constexpr StringRef getPaddingAttrName() { return "padding"; }
+
 /// Use to encode that a particular iterator type has parallel semantics.
 constexpr StringRef getParallelIteratorTypeName() { return "parallel"; }
 
 /// Use to encode that a particular iterator type has reduction semantics.
 constexpr StringRef getReductionIteratorTypeName() { return "reduction"; }
 
 /// Use to encode that a particular iterator type has window semantics.
 constexpr StringRef getWindowIteratorTypeName() { return "window"; }

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 };
 
 /// Populate the given list with patterns that convert from Linalg to Standard.
 static void
 populateLinalgToStandardConversionPatterns(OwningRewritePatternList &patterns,
                                            MLIRContext *ctx) {
   // TODO(ntv) ConvOp conversion needs to export a descriptor with relevant
   // attribute values such as kernel striding and dilation.
   patterns.insert<CopyTransposeConversion, LinalgOpConversion<ConvOp>,
-                  LinalgOpConversion<CopyOp>, LinalgOpConversion<DotOp>,
-                  LinalgOpConversion<FillOp>, LinalgOpConversion<GenericOp>,
+                  LinalgOpConversion<PoolingOp>, LinalgOpConversion<CopyOp>,
+                  LinalgOpConversion<DotOp>, LinalgOpConversion<FillOp>,
+                  LinalgOpConversion<GenericOp>,
                   LinalgOpConversion<IndexedGenericOp>,
                   LinalgOpConversion<MatmulOp>, LinalgOpConversion<MatvecOp>>(
       ctx);
 }
 
 } // namespace
 
 /// Populate the given list with patterns that convert from Linalg to LLVM.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 // For such operations that do not correspond to library calls (i.e. defined in
 // LinalgOps.td), we define an overloaded `print` function and a
 // parse`className` function.
 
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 
-template <typename GenericOpType>
-static void printGenericOp(OpAsmPrinter &p, GenericOpType op) {
+template <typename OpType>
+static void printTraitAndRegion(OpAsmPrinter &p, OpType op) {
   auto attrNames = op.linalgTraitAttrNames();
   llvm::StringSet<> linalgTraitAttrsSet;
   linalgTraitAttrsSet.insert(attrNames.begin(), attrNames.end());
   SmallVector<NamedAttribute, 8> attrs;
   for (auto attr : op.getAttrs())
     if (linalgTraitAttrsSet.count(attr.first.strref()) > 0)
       attrs.push_back(attr);
 
   auto dictAttr = DictionaryAttr::get(attrs, op.getContext());
   p << op.getOperationName() << " " << dictAttr << " " << op.getOperands();
   if (!op.region().empty())
     p.printRegion(op.region());
   p.printOptionalAttrDict(op.getAttrs(), attrNames);
   p << ": " << op.getOperandTypes();
+}
 
+template <typename GenericOpType>
+static void printGenericOp(OpAsmPrinter &p, GenericOpType op) {
+  printTraitAndRegion(p, op);
   auto outputTensorTypes = op.getResultTypes();
   if (!outputTensorTypes.empty())
     p << " -> " << outputTensorTypes;
 }
 
 static void print(OpAsmPrinter &p, GenericOp op) { printGenericOp(p, op); }
 
 static void print(OpAsmPrinter &p, IndexedGenericOp op) {
   printGenericOp(p, op);
 }
 
-static ParseResult parseGenericOp(OpAsmParser &parser, OperationState &result) {
+static void print(OpAsmPrinter &p, PoolingOp op) { printTraitAndRegion(p, op); }
+
+static ParseResult parseTraitAndRegion(OpAsmParser &parser,
+                                       OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 8> operandsInfo, regionOperandsInfo;
   DictionaryAttr dictAttr;
   // Parse the core linalg traits that must check into a dictAttr.
   // The name is unimportant as we will overwrite result.attributes.
   // The core linalg traits must contain the information necessary to pass the
   // verifier.
   if (parser.parseAttribute(dictAttr, "_", result.attributes) ||
       parser.parseOperandList(operandsInfo))
static LogicalResult verify(YieldOp op) {
   auto genericOp = dyn_cast<GenericOp>(parentOp);
   if (genericOp)
     return verifyYield(op, genericOp);
 
   auto indexedGenericOp = dyn_cast<IndexedGenericOp>(parentOp);
   if (indexedGenericOp)
     return verifyYield(op, indexedGenericOp);
 
+  if (auto poolingOp = dyn_cast<PoolingOp>(parentOp))
+    return verifyYield(op, poolingOp);
+
   return op.emitOpError("expected '")
          << GenericOp::getOperationName() << "' or '"
          << IndexedGenericOp::getOperationName() << "' parent op";
 }
 
 /////// Operations corresponding to library calls defined with Tablegen ////////
 
 static LogicalResult verify(FillOp op) {
static LogicalResult verify(CopyOp op) {
   }
   if (rank == 0 && inputPermutationMap)
     return op.emitOpError("expected no input permutation when rank == 0");
   if (rank == 0 && outputPermutationMap)
     return op.emitOpError("expected no output permutation when rank == 0");
   return success();
 }
 
-static LogicalResult
-verifyStrideOrDilation(ConvOp op, ArrayRef<Attribute> attrs, bool isStride) {
+template <typename LinalgPoolingOp>
+static LogicalResult verifyStrideOrDilation(LinalgPoolingOp op,
+                                            ArrayRef<Attribute> attrs,
+                                            bool isStride) {
   auto strideOrDilation = isStride ? "stride" : "dilation";
   if (attrs.size() != op.getNumWindowLoops())
     return op.emitOpError("expects num ")
            << strideOrDilation
            << "s equal to number of window dimensions: " << attrs.size()
            << " vs " << op.getNumWindowLoops();
   return success();
 }
static LogicalResult verify(ConvOp op) {
   if (auto dilations = op.dilations()) {
     if (failed(verifyStrideOrDilation(op, dilations->getValue(),
                                       /*isStride=*/false)))
       return failure();
   }
   return success();
 }
 
+static LogicalResult verify(PoolingOp op) {
+  auto inputType = op.input().getType().cast<MemRefType>();
+  auto outputType = op.output().getType().cast<MemRefType>();
+  if (outputType.getElementType() != inputType.getElementType())
+    return op.emitOpError("expects memref elemental types to match");
+
+  auto windowDimsType = op.windowDims().getType().cast<MemRefType>();
+  if (outputType.getRank() != inputType.getRank() ||
+      outputType.getRank() != windowDimsType.getRank())
+    return op.emitOpError("expects memref ranks to match");
+
+  if (auto strides = op.strides()) {
+    if (failed(
+            verifyStrideOrDilation(op, strides->getValue(), /*isStride=*/true)))
+      return failure();
+  }
+
+  if (op.region().empty())
+    return op.emitOpError("expects a region");
+  if (op.region().front().getNumArguments() != 2)
+    return op.emitOpError("expected number of block arguments to be 2");
+
+  return success();
+}
+
 namespace mlir {
 namespace linalg {
 
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
 
mlir::linalg::makeAffineDimExprs(unsigned num, unsigned &startIdx,
                                  MLIRContext *context) {
   SmallVector<AffineExpr, 4> res;
   res.reserve(num);
   for (unsigned i = 0; i < num; ++i)
     res.push_back(getAffineDimExpr(startIdx++, context));
   return res;
 }
 
-SmallVector<AffineExpr, 4>
-mlir::linalg::weightedConvInputIndex(ConvOp op, ArrayRef<AffineExpr> xs,
-                                     ArrayRef<AffineExpr> zs) {
-  assert(xs.size() == zs.size());
+template <typename LinalgOp>
+static SmallVector<AffineExpr, 4>
+weightedPoolingInputIndexImpl(LinalgOp op, ArrayRef<AffineExpr> outputDims,
+                              ArrayRef<AffineExpr> windowDims) {
+  assert(outputDims.size() == windowDims.size());
   SmallVector<AffineExpr, 4> res;
-  res.reserve(xs.size());
-  for (unsigned i = 0, e = xs.size(); i < e; ++i) {
+  res.reserve(outputDims.size());
+  for (unsigned i = 0, e = outputDims.size(); i < e; ++i) {
     // TODO(ntv): add a level of indirection to linalg.generic.
-    auto expr =
-        op.getStride(i) * xs[i] + op.getDilation(i) * zs[i] - op.getLowPad(i);
+    auto expr = op.getStride(i) * outputDims[i] +
+                op.getDilation(i) * windowDims[i] - op.getLowPad(i);
     res.push_back(expr);
   }
   return res;
 }
 
+SmallVector<AffineExpr, 4>
+mlir::linalg::weightedPoolingInputIndex(ConvOp op,
+                                        ArrayRef<AffineExpr> outputDims,
+                                        ArrayRef<AffineExpr> windowDims) {
+  return weightedPoolingInputIndexImpl(op, outputDims, windowDims);
+}
+
+SmallVector<AffineExpr, 4>
+mlir::linalg::weightedPoolingInputIndex(PoolingOp op,
+                                        ArrayRef<AffineExpr> outputDims,
+                                        ArrayRef<AffineExpr> windowDims) {
+  return weightedPoolingInputIndexImpl(op, outputDims, windowDims);
+}
+
 SmallVector<AffineExpr, 4> mlir::linalg::concat(ArrayRef<AffineExpr> a,
                                                 ArrayRef<AffineExpr> b) {
   auto rangeA = llvm::make_range(a.begin(), a.end());
   auto rangeB = llvm::make_range(b.begin(), b.end());
   auto concatRanges = llvm::concat<const AffineExpr>(rangeA, rangeB);
   return llvm::to_vector<4>(concatRanges);
 }
 

 
 // TODO(ntv, rriddle): Consider making all this boilerplate easy to autogenerate
 // with Tablegen. This seems a desirable property in the context of OpInterfaces
 // where a Linalg "named" op **isa** LinalgOp.
 LogicalResult ConvOp::fold(ArrayRef<Attribute>,
                            SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
+LogicalResult PoolingOp::fold(ArrayRef<Attribute>,
+                              SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}
 LogicalResult CopyOp::fold(ArrayRef<Attribute>,
                            SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 LogicalResult DotOp::fold(ArrayRef<Attribute>,
                           SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }

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

static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
     IndexedValueType F(convOp.filter()), I(convOp.input()), O(convOp.output());
 
     // Emit scalar form.
     ValueHandle paddedInput = getConvOpInput(convOp, I, imIdx);
     O(oIdx) += F(fIdx) * paddedInput;
   }
 };
 
+template <typename IndexedValueType>
+class LinalgScopedEmitter<IndexedValueType, PoolingOp> {
+public:
+  static void emitScalarImplementation(ArrayRef<Value> allIvs,
+                                       PoolingOp poolingOp) {
+    auto b = ScopedContext::getBuilder();
+    auto loc = ScopedContext::getLocation();
+    auto mapsRange = poolingOp.indexing_maps().getAsRange<AffineMapAttr>();
+    auto maps = functional::map([](AffineMapAttr a) { return a.getValue(); },
+                                mapsRange);
+    SmallVector<ValueHandle, 8> iIdx(
+        makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
+    SmallVector<ValueHandle, 8> oIdx(
+        makeCanonicalAffineApplies(b, loc, maps[2], allIvs));
+
+    // 1. Emit std_load from operands.
+    IndexedValueType input(poolingOp.input()), output(poolingOp.output());
+    ValueHandle arg0 = input(iIdx);
+    ValueHandle arg1 = output(oIdx);
+    SmallVector<Value, 2> indexedValues = {arg0.getValue(), arg1.getValue()};
+
+    // 2. Inline region, currently only works for a single basic block.
+    ValueHandleArray indexing = {ValueHandleArray(oIdx)};
+    Value outputBuffer = poolingOp.output();
+    inlineRegionAndEmitStdStore(poolingOp, indexedValues, indexing,
+                                outputBuffer);
+  }
+};
+
 // Emits the MLIR for the scalar part of the generic op by:
 //   1. Emitting std_load and std_store ops for each input and output
 //      view in order. This is achieved by applying the appropriate input or
 //      output map to the enclosing induction variables.
 //   2. Emitting a call to `op.fun()` that takes as arguments the scalars
 //      from point 1. above.
 //   3. Emitting std_store to store the results of 2. to the output
 //      views.
template LogicalResult mlir::linalg::linalgOpToAffineLoops<OP_TYPE>(         \
       PatternRewriter & rewriter, Operation * op);
 
 INSTANTIATE_LINALG_OP_TO_LOOPS(CopyOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(FillOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(DotOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(MatvecOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(MatmulOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(ConvOp)
+INSTANTIATE_LINALG_OP_TO_LOOPS(PoolingOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(GenericOp)
 INSTANTIATE_LINALG_OP_TO_LOOPS(IndexedGenericOp)
 
 // TODO(pifon): Enable lowering to parallel loops for ops other than
 // linalg.generic for now to be on the safe side.
 template LogicalResult
 mlir::linalg::linalgOpToParallelLoops<GenericOp>(PatternRewriter &rewriter,
                                                  Operation *op);

mlir/test/Dialect/Linalg/invalid.mlir

 
 // -----
 
 func @reshape(%arg0: memref<?x?x?xf32>) {
   // expected-error @+1 {{expected collapsed type to be 'memref<?x?xf32>', but got 'memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>'}}
   %0 = linalg.reshape %arg0 [affine_map<(i, j, k) -> (i, j)>, affine_map<(i, j, k) -> (k)>] :
     memref<?x?x?xf32> into memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>>
 }
+
+// -----
+
+func @pooling_rank_mismatch(%arg0: memref<?x?x?xf32>,
+                            %arg1: memref<2x3xf32>,
+                            %arg2: memref<?x?x?xf32>) {
+  // expected-error @+1 {{expects memref ranks to match}}
+  linalg.pooling { strides = [2, 1, 2] } %arg0, %arg1, %arg2 {
+    ^bb(%a: f32, %b: f32) :
+      %0 = addf %a, %b : f32
+      linalg.yield %0 : f32
+  }: memref<?x?x?xf32>, memref<2x3xf32>, memref<?x?x?xf32>
+  return
+}
+
+// -----
+
+func @pooling_missing_region(%arg0: memref<?x?x?xf32>,
+                             %arg1: memref<2x3x2xf32>,
+                             %arg2: memref<?x?x?xf32>) {
+  // expected-error @+1 {{expects a region}}
+  linalg.pooling { strides = [2, 1, 2] } %arg0, %arg1, %arg2 :
+    memref<?x?x?xf32>, memref<2x3x2xf32>, memref<?x?x?xf32>
+  return
+}

mlir/test/Dialect/Linalg/loops.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops | FileCheck %s
 
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
 // RUN: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
 
 // CHECK-DAG: #[[strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 // CHECK-DAG: #[[strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
 // CHECK-DAG: #[[clampMinMap:.*]] = affine_map<(d0) -> (d0, 0)>
 
+// CHECK-DAG: #[[Stride1Dilation1:.*]] = affine_map<(d0, d1) -> (d0  + d1)>
 // CHECK-DAG: #[[Stride2Dilation1:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1)>
 // CHECK-DAG: #[[Stride2Dilation4:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1 * 4)>
 // CHECK-DAG: #[[Stride3Dilation5:.*]] = affine_map<(d0, d1) -> (d0 * 3 + d1 * 5)>
 
 
 func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index

 //       CHECK:                 %{{.*}} = load %{{.*}}[%{{.*}}, %[[IDX]], %[[IDY]], %{{.*}}] : memref<?x?x?x?xf32>
 //       CHECK:                 %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : f32
 //       CHECK:                 %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?x?xf32>
 //       CHECK:                 %{{.*}} = mulf %{{.*}}, %{{.*}} : f32
 //       CHECK:                 %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?x?xf32>
 //       CHECK:                 %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 //       CHECK:                 store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<?x?x?x?xf32>
 
+func @pooling_sum(%arg0: memref<?x?xf32>,
+                  %arg1: memref<?x?xi32>,
+                  %arg2: memref<?x?xf32>) {
+  linalg.pooling { strides = [2, 1] } %arg0, %arg1, %arg2 {
+    ^bb(%a: f32, %b: f32) :
+      %0 = addf %a, %b : f32
+      linalg.yield %0 : f32
+  }: memref<?x?xf32>, memref<?x?xi32>, memref<?x?xf32>
+  return
+}
+// CHECK-LABEL: func @pooling_sum
+//       CHECK:   %[[WX:.*]] = dim %arg1, 0 : memref<?x?xi32>
+//       CHECK:   %[[WY:.*]] = dim %arg1, 1 : memref<?x?xi32>
+//       CHECK:   %[[OX:.*]] = dim %arg2, 0 : memref<?x?xf32>
+//       CHECK:   %[[OY:.*]] = dim %arg2, 1 : memref<?x?xf32>
+//       CHECK:   loop.for %{{.*}} = %{{.*}} to %[[OX]] step %{{.*}} {
+//       CHECK:     loop.for %{{.*}} = %{{.*}} to %[[OY]] step %{{.*}} {
+//       CHECK:       loop.for %{{.*}} = %{{.*}} to %[[WX]] step %{{.*}} {
+//       CHECK:         loop.for %{{.*}} = %{{.*}} to %[[WY]] step %{{.*}} {
+//       CHECK:           %[[IX:.*]] = affine.apply #[[Stride2Dilation1]](%{{.*}}, %{{.*}})
+//       CHECK:           %[[IY:.*]] = affine.apply #[[Stride1Dilation1]](%{{.*}}, %{{.*}})
+//       CHECK:           %[[A:.*]] = load %{{.*}}[%[[IX]], %[[IY]]] : memref<?x?xf32>
+//       CHECK:           %[[B:.*]] = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
+//       CHECK:           %[[RES:.*]] = addf %[[A]], %[[B]] : f32
+//       CHECK:           store %[[RES]], %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
+
 func @foo(%0: f32, %1: f32, %2: f32) -> (f32, f32) {
   %f0 = constant 0.0 : f32
   return %f0, %f0 : f32, f32
 }
 #accesses = [
   affine_map<(i, j, k) -> (i, j)>,
   affine_map<(i, j, k) -> (i, j, k)>,
   affine_map<(i, j, k) -> (i, k, j)>

mlir/test/Dialect/Linalg/roundtrip.mlir

 //  CHECK-SAME:                      [0, 1], [1, 1]]> : tensor<2x2xi64>,
 //  CHECK-SAME:     strides = [1, 1]} :
 //  CHECK-SAME:     memref<?x?x?x?xf32>,
 //  CHECK-SAME:     memref<?x?x?x?xf32>,
 //  CHECK-SAME:     memref<?x?x?x?xf32>
 
 // -----
 
+func @pooling_sum(%arg0: memref<?x?x?xf32>,
+                  %arg1: memref<?x?x?xi32>,
+                  %arg2: memref<?x?x?xf32>) {
+  linalg.pooling { strides = [2, 1, 2] } %arg0, %arg1, %arg2 {
+    ^bb(%a: f32, %b: f32) :
+      %0 = addf %a, %b : f32
+      linalg.yield %0 : f32
+  }: memref<?x?x?xf32>, memref<?x?x?xi32>, memref<?x?x?xf32>
+  return
+}
+// CHECK-LABEL: func @pooling_sum
+//       CHECK:   linalg.pooling {strides = [2, 1, 2]}
+//  CHECK-SAME:     %{{.*}}, %{{.*}}, %{{.*}} {
+//       CHECK:    ^{{.*}}(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
+//       CHECK:      %[[RES:.*]] = addf %[[LHS]], %[[RHS]] : f32
+//       CHECK:      linalg.yield %[[RES]] : f32
+//       CHECK:   }: memref<?x?x?xf32>, memref<?x?x?xi32>, memref<?x?x?xf32>
+
+// -----
+
 // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECK-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 
 func @foo(%0: vector<3x4xi4>, %1: f32) -> f32 {
   %f0 = constant 0.0 : f32
   return %f0 : f32
 }