Page MenuHomePhabricator

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

Authored by hanchung on Mar 18 2020, 11:21 PM.

Details

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.

Diff Detail

Event Timeline

hanchung created this revision.Mar 18 2020, 11:21 PM
hanchung edited the summary of this revision. (Show Details)Mar 18 2020, 11:23 PM
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
277

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

459

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
526–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 added inline comments.Mar 20 2020, 1:04 PM
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
425

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

440

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.

515

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?

Nice. Few comments here

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

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.

459

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

489

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 updated this revision to Diff 252141.Mar 23 2020, 2:30 PM
hanchung marked 14 inline comments as done.

Address comments.

hanchung added inline comments.Mar 23 2020, 2:32 PM
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
425

Yes, this is a good reference. :)

459

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

489

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.

515

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 retitled this revision from [mlir][Linalg] Introduce linalg.pooling op. to [mlir][Linalg] Introduce linalg.pooling_min/max/sum op..Mar 23 2020, 2:33 PM
hanchung edited the summary of this revision. (Show Details)
hanchung updated this revision to Diff 252626.Mar 25 2020, 10:46 AM

Fix clang-tidy.

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

mravishankar accepted this revision.Mar 30 2020, 11:45 AM

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);
...
This revision is now accepted and ready to land.Mar 30 2020, 11:45 AM
nicolasvasilache accepted this revision.Mar 30 2020, 3:06 PM

Thanks for adding this!

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

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 updated this revision to Diff 254018.Mar 31 2020, 3:15 PM
hanchung marked 4 inline comments as done.

Address comments

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

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

hanchung updated this revision to Diff 254020.Mar 31 2020, 3:21 PM

Fix description in SingleInputPoolingBase_Op

Harbormaster completed remote builds in B51202: Diff 254018.
This revision was automatically updated to reflect the committed changes.