NOTE: Effort is being made to make this facility generic so that it can be eventually upstreamed to MLIR in some fashion. However, because MLIR lacks a set of frontend ops and generally does not currently have any frontend oriented infrastructure, it is being prototyped within IREE in order to find a working set of ops and algorithms.
In general, there are three levels of shape information that can be present in the input IR (or trivially derived by applying some form of shape inferencing algorithm). Each additional one imposes more work on the compiler and runtime, so generally, the implementation progresses by addressing each once the former is well established:
- Fully static shapes: No tensors have dynamic dimensions. All tensors are ranked.
- Ranked Dynamicism: All tensors have ranks, but some dimensions may be unspecified.
- Unranked Dynamicism: Some tensors have indeterminate ranks.
At this stage, Dynamic Shapes in IREE refers to supporting dynamic ranked dynamic tensors, where some dimensions are left unspecified at public function boundaries. It is expected that once this is solid, some support can be considered for unranked dynamicism, and it is further expected that will entail new ops, algorithms and runtime support, apart from what is needed for ranked dynamicism.
Within the category of Ranked Dynamicism, it is well known that some dynamic dimensions are easier to deal with than others: in common DNN use, outer dimensions are much easier and more common with respect to code generation and kernel fanout than dynamic inner dimensions.
While the shape handling machinery is relatively generic, we expect that real backends will be limited with respect to how much they support all combinations of dynamic dimensions. Eventually, IREE intends to solve this by having relatively robust CPU fallback for fully dynamic cases and actionable warnings that pinpoint when more specificity could increase performance.
In general, the IREE compiler frontend should accept modules containing
functions with operands/results that have dynamic dimensions. Such functions may
also have runtime dependent shapes in the form of GetShape-style ops which get
a shape from an arbitrary tensor, perform some arithmetic on it and use the
results elsewhere.
IREE is introducing a shape dialect with a handful of ops and transformations
that are useful for materializing dynamic shape computations in terms of high
level ops on tensors.
ranked_shape: This value type represents the dynamic dimensions of a partially known, ranked shape. It is used early in the compiler to represent anywhere that dynamic dimensions need to be passed (i.e. function args/results, etc). At lower levels of the compiler, it will generally be dis-aggregated into loose SSA values. This type also carries the datatype used to represent the dimensions. This is currently fixed to i32 but may be leveraged eventually to use smaller integer when such things are known to be legal.
get_ranked_shape: Takes a tensor SSA value and returns a correspondingranked_shape. Early in the compilation flow, anything that needs a ranked shape should add such ops so that the compiler can later determine which shape computations to materialize. Getting theranked_shapeof a static tensor should yield a constant.tie_shape: Takes tensor and ranked_shape SSA values and returns the tensor. This is used as a junction point by the shape materialization passes to know at various points precisely what the shape is.- ... TODO: need
get_shape_dimand conversions to/from 1D tensors and loose SSA values.
Early in the process, all functions should have their arguments and results
expanded so that any dynamic tensors in their signature will gain a new
argument/result for the corresponding ranked_shape. This is done by expanding
the signatures and for arguments, inserting placeholder tie_shape ops which
preserve the association for later materialization. For results,
get_ranked_shape ops are inserted.
This is carried out by the iree-shape-expand-function-dynamic-dims pass, which
uses the conversion framework under the hood to perform type expansion.
This pass is typically done early in the compiler flow.
A lot of scheduling logic will need to access shapes (i.e. allocation, workgroup
size calculation, etc). In general, this should all be done based on a
get_ranked_shape op and corresponding get_shape_dim ops. For fully static
cases, these should reduce down to constants. For dynamic dimensions, the
get_ranked_shape ops serve as anchors where later parts of the compiler know
they need to materialize shape values.
TODO: We have a sketch of this but are still proving it out.
Generally, it should be possible, for any get_ranked_shape op, to trace up the
use-def chain and materialize shape manipulation arithmetic. Once materialized,
a tie_shape op should be inserted to memorialize the junction. Eventually,
every get_ranked_shape op should be follow a tie_shape op, and the
canonicalization rules will elide the get_ranked_shape. There is complexity
around blocks, control flow, etc, but this basic algorithm should be workable.
Work is ongoing upstream to provide a facility to register shape functions with ops, which would provide a dynamic, dialect independent way to know what arithmetic to materialize. However, in most cases this is not necessary. The built-in traits around types and sizes will allow most propagation to happen without shape functions. We intend to start with a static set of cases for the rest in order to prove the concept.
TODO: We have a sketch of this but are still proving it out.
It is quite common in real-world DNN usage to get the 1D tensor representing a shape and perform arbitrary tensor ops on it (usually basic arithmetic, slicing, concating, tiling, etc). While this is perfectly acceptable from a correctness standpoint, it is usually not performant: shapes are typically very small one dimensional vectors, and computations on them are usually trivial to reduce to small sequences of scalar machine code of a form that CPUs are very good at executing. Further, we often want to run these calculations eagerly when dispatching functions, etc (i.e. to pre-allocate buffers) and having them isolated (versus treating them as arbitrary dense computations) can be quite valuable.
We expect that the type bracketing that happens with ranked_shape and the
corresponding ops will make it easy to write some simple DRR patterns to
identify such shape manipulation sequences and lower them directly to regions of
vm ops operating on scalars. Such regions can be retained and directly emitted
when lowering to the vm dialect and/or CPU code generation and would run with
low overhead along with any other scheduling code.
While an optimization, we suspect this is an important one.
TODO: This is mostly placeholder
There is work happening upstream to implement MLIR-integrated shape inference. In the mean-time, IREE expects that the input to the compiler has already had some shape inference performed on it. In practice, for TensorFlow, there is a pass which applies TensorFlow's pre-MLIR shape inference mechanisms to derive such things. This has limitations but is a reasonable starting point.
TODO: This is mostly placeholder.
Much of the newer structured-ops based codegen is capable of working (within bounds) with ranked dynamic shapes without much work. Given the lack of an e2e story, much of this has been done "by way of code review" and there are certainly issues to be resolved.
In addition, there are several ABI issues and negotiations with the backend that still need to be fleshed out.