`shape` dialect: add some ops

Authored by silvas on Mar 26 2020, 3:11 PM.


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


silvasMar 27 2020, 4:38 PM
Differential Revision
D76817: `shape` dialect: add some ops
rGcbce88dd3a9e: FunctionRef: Strip cv qualifiers in the converting constructor