[PATCH] D76817: `shape` dialect: add some ops

[Jacques Pienaar via Phabricator via llvm-commits]

jpienaar added inline comments.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:189
+    in the tensor equals the rank of the shape, and the elements equal the
+    extents of the shape.
+  }];
----------------
silvas wrote:
> jpienaar wrote:
> > silvas wrote:
> > > jpienaar wrote:
> > > > What happens for unranked and or dynamic tensors?
> > > it return the shape. Notice that the return type is not static shaped. For an unranked tensor it will return a `tensor<?xi32>` (or index someday).
> > Thus far FromExtentTensorOp has as input a tensor with known values and produces a fixed shape shape (e.g., there is no special numeric values that correspond to unknown). But another problem is that ToExtentTensor now needs to convert from a shape which may be an error to a tensor of ints. The intention is to have all shape ops be side-effect free and I don't see how that can be maintained for this op.
> > 
> > What I had though of here was something like:
> > 
> > ```
> > shape.if_static(%s : shape) {
> >   ^bb(%t: tensor<?xi32>):
> >      ... now use as 
> > } else {
> > }
> > ```
> > 
> > so that you only convert when it is safe,
> There is no such thing as a tensor with unknown or dynamic or even partially specified shape at runtime.
> 
> One can consider the abstract !shape.shape having a runtime backing representation which represents a dataflow lattice value (such as tracking extent upper bounds to allow buffer allocation before a data-dependent shape op (like "unique") completes), but that would require specialized ops that are aware of the underlying backing representation for the abstract !shape.shape in order to construct such a value, since the choice of lattice is arbitrary.
> 
> There's only a small number of ops on !shape.shape that reify the concept of the shape being a lattice value (since there just isn't much you can abstractly do with a lattice, by definition). `shape.join` is one of them (and the complementary "most general shape" is the only other one I can think of). All other ops should be defined as if they are runtime computations on concrete shapes, and any particular lattice that a user intends to use will need to substitute in appropriate transfer functions either via dataflow analysis statically, or reified in the ir during lowering (in their own dialect), or literally have the runtime manifestation of !shape.shape have virtual methods they can override to substitute the transfer functions for the primitives.
> 
> As far as side effects in the error case, you can't dodge them at least how you've defined them here, at least in my use case. The result of the shape.concat/shape.slice_at/etc. is used to broadcast a tensor to a particular concrete runtime shape. If the !shape.shape is an error, then now you've just passed the error onto the "broadcast this tensor to this shape op". In general you could pass the resulting shape to "{numpy,tf,torch}.zeros(shape)" which to your point should be pure.
> 
> I feel strongly that we need a dialect that can be used to model runtime shape computations. If you want this dialect to be used purely for abstract dataflow then that's okay and I can create a different dialect.
> 
> I think this error handling situation needs more thought. Neither tf dialect nor xla_hlo dialects (or any other dialects I'm aware of that would use this) have any actual IR manifestation for what happens in error cases. E.g. most "matmul" ops that I'm aware of claim to be side effect free. Hmm....
I'm not sure I follow, shape here is a effectively an error monad, the ops operating on it propagates the error (as mentioned this could end up merging multiple errors together), but the ops are side effect free and can be folded, else many of these ops are effectively asserts that can fail (e.g., they are adding constraints on the number of elements, the maximum extent of dimensions, ...). Once lowered to a runtime executable form, then those errors would need to be considered. E.g., a tf.assert or some such needs to be inserted before spanning the gap to tf.zeros or what not that operates on tensor of ints. There are ops in those platforms which can handle those and so in lowering handle it - there are actual runtime ops there with cancellation or abort mechanisms.

The goal of this dialect is to describe shape computations for multiple different consumers at analysis, compilation or for lowering to runtime. I don't see there being a single consumer/runtime here as the way that different consumers/runtimes could handle computations, failures, error reporting, would be different and keeping it abstract means the lowerings to those targets can consider what they need (e.g., perhaps we are in "assume there is no errors" mode and all checks can be elided, or a target is massively distributed and you want error reporting to be totally ordered).

These are defined by runtime behavior, it is just one of the "runtime" is symbolic analysis (e.g., runtime behavior includes how to create a constraint programming problem). Being able to reuse these in multiple contexts is important, else we end up defining the shapes of ops for each use case. For the runtime case with all the values known, you just need 1D tensor of ints, and you could reuse normal ops to implement the computations (e.g., you could use tf.slice or xla_hlo.mul or linalg), the result of these then just gets consumed by runtime specific functions to allocate/deallocate buffers, so you don't need any special types or ops at that point. So once you are using tensor of ints you are free to use arbitrary ops and you've jumped to a new domain, but given bridging these two could require handling errors, I want the lowering for a use to handle that.

So specified as runtime behavior on shapes makes sense to me, and it makes sense to have this op to convert from `!shape.shape` to `tensor<?xindex>` in this dialect, but not for use in a shape functions (else you are just using `tensor<?x index>` here as a cast to appease the type system).

I might be wrong though, I was trying to sketch multiple cases in my head (and refining the comment every now and again :)) and while I can see cases where you have benefit and approaches to use such an op, I don't know how many of them are really real. But one thing that is true, if we have such an if instead, you get the conversion and if we find it is not needed, it amounts to replacing the if with a .ToTensorExtend and folding the else [which could also be the lowering pattern towards runtime you use and that would give you the exact same shape function you have today]


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D76817/new/

https://reviews.llvm.org/D76817