`shape` dialect: add some ops
authorSean Silva <silvasean@google.com>
Thu, 26 Mar 2020 22:11:23 +0000 (15:11 -0700)
committerSean Silva <silvasean@google.com>
Fri, 27 Mar 2020 23:38:42 +0000 (16:38 -0700)
commit569e4f9bc99a755cc30f0102b29b1eefd4fa33b4
treec53cd80049d42ccdf1165fa8e2dc213deaf54ca0
parentcbce88dd3a9ea7161da3c57749cf03873dc7ea79
`shape` dialect: add some ops

- add `to_extent_tensor`
 - rename `create_shape` to `from_extent_tensor` for symmetry
- add `split_at` and `concat` ops for basic shape manipulations

This set of ops is inspired by the requirements of lowering a dynamic-shape-aware batch matmul op. For such an op, the "matrix" dimensions aren't subject to broadcasting but the others are, and so we need to slice, broadcast, and reconstruct the final output shape. Furthermore, the actual broadcasting op used downstream uses a tensor of extents as its preferred shape interface for the actual op that does the broadcasting.

However, this functionality is quite general. It's obvious that `to_extent_tensor` is needed long-term to support many common patterns that involve computations on shapes. We can evolve the shape manipulation ops introduced here. The specific choices made here took into consideration the potentially unranked nature of the !shape.shape type, which means that a simple listing of dimensions to extract isn't possible in general.

Differential Revision: https://reviews.llvm.org/D76817
mlir/include/mlir/Dialect/Shape/IR/Shape.h
mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
mlir/lib/Dialect/Shape/CMakeLists.txt
mlir/lib/Dialect/Shape/IR/Shape.cpp