[RFC] Symbolic Shape Analysis

yaochengji · May 31, 2022, 5:19pm

Background

During a dynamic shape code generation, we always want to find the relationships among the dynamic shapes. The relationships include and not limited to:

Where’s the dimension’s dynamic comes from? For the example below, we’d like to deduce that all of %0, %1, %3’s dynamic is from %arg0.

func.func @several_ops(%arg0: tensor<?x4xf32>, %arg1: tensor<4x4xf32>, %arg2: tensor<4xf32>) -> tensor<?x4xf32> {
  %0 = "mhlo.dot_general"(%arg0, %arg1) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_contracting_dimensions = [1],
      rhs_contracting_dimensions = [0]
    >,
   precision_config = [#mhlo<"precision DEFAULT">, #mhlo<"precision DEFAULT">]
  } : (tensor<?x4xf32>, tensor<4x4xf32>) -> tensor<?x4xf32>
  %1 = shape.shape_of %0 : tensor<?x4xf32> -> tensor<2xindex>
  %2 = "mhlo.dynamic_broadcast_in_dim"(%arg2, %1) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<4xf32>, tensor<2xindex>) -> tensor<?x4xf32>
  %3 = mhlo.add %0, %2 : tensor<?x4xf32>
  return %3 : tensor<?x4xf32>
}

The constraints on the dynamic shapes. Let’s say dimension a and dimension b are both dynamic. We’d like to deduce that whether a equals b, or a + b equals a static dimension, etc.

To address these problems, I’m considering to implement a generic symbolic shape analysis.

Related Work

I found an interesting pass in mlir-hlo repo. It uses affine expression to describe how a dynamic dimension could be calculated. But I’m afraid we’ll do a lot of repetitive work if symbolic shape analysis is based on affine expresion. Since:

We cannot reuse the reifyReturnTypeShapes method in InferShapedTypeOpInterface
If we want to do some simplification or do more analysis on the symbolic expression, we could not reuse the existing pass(Ex. cse, canonicalize) or write new passes based on the MLIR pass framework.

Proposed Solution

Here I drafted a new solution to express symbolic shape calculation. It works in 3 steps:

Create an auxiliary symbolic shape inference function for each original function.
Run shape-reification pass on the created auxiliary functions.
Analyze based on the auxiliary functions.

We still use the IR mentioned in background section as an input IR. After step 1, an auxiliary function will be created as below. Where the arguments of the auxiliary function is the same as the original one, and the results includes two parts. In the example there are 4 intermediate results, %0, %1, %2, %3. Then the return results will be [shape_of %0, shape_of %1, shape_of %2, shape_of %3, %0, %1, %2, %3].

func.func private @_shape_infer_several_ops(%arg0: tensor<?x4xf32>, %arg1: tensor<4x4xf32>, %arg2: tensor<4xf32>) -> (!shape.shape, !shape.shape, !shape.shape, !shape.shape, tensor<?x4xf32>, tensor<2xindex>, tensor<?x4xf32>, tensor<?x4xf32>) attributes {auxiliary_of = "several_ops"} {
  %0 = "mhlo.dot_general"(%arg0, %arg1) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<"precision DEFAULT">, #mhlo<"precision DEFAULT">]} : (tensor<?x4xf32>, tensor<4x4xf32>) -> tensor<?x4xf32>
  %1 = shape.shape_of %0 : tensor<?x4xf32> -> tensor<2xindex>
  %2 = shape.value_as_shape %1 : tensor<2xindex> -> !shape.shape
  %3 = shape.shape_of %0 : tensor<?x4xf32> -> tensor<2xindex>
  %4 = shape.shape_of %3 : tensor<2xindex> -> tensor<1xindex>
  %5 = shape.value_as_shape %4 : tensor<1xindex> -> !shape.shape
  %6 = "mhlo.dynamic_broadcast_in_dim"(%arg2, %3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<4xf32>, tensor<2xindex>) -> tensor<?x4xf32>
  %7 = shape.shape_of %6 : tensor<?x4xf32> -> tensor<2xindex>
  %8 = shape.value_as_shape %7 : tensor<2xindex> -> !shape.shape
  %9 = mhlo.add %0, %6 : tensor<?x4xf32>
  %10 = shape.shape_of %9 : tensor<?x4xf32> -> tensor<2xindex>
  %11 = shape.value_as_shape %10 : tensor<2xindex> -> !shape.shape
  return %2, %5, %8, %11, %0, %3, %6, %9 : !shape.shape, !shape.shape, !shape.shape, !shape.shape, tensor<?x4xf32>, tensor<2xindex>, tensor<?x4xf32>, tensor<?x4xf32>
}

And after step 2, a shape-reification pass and a cse pass will be run on the auxilary function, producing the following IR:

func.func private @_shape_infer_several_ops(%arg0: tensor<?x4xf32>, %arg1: tensor<4x4xf32>, %arg2: tensor<4xf32>) -> (!shape.shape, !shape.shape, !shape.shape, !shape.shape, tensor<?x4xf32>, tensor<2xindex>, tensor<?x4xf32>, tensor<?x4xf32>) attributes {auxiliary_of = "several_ops"} {
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  %0 = shape.const_shape [2] : tensor<1xindex>
  %1 = "mhlo.dot_general"(%arg0, %arg1) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = [#mhlo<"precision DEFAULT">, #mhlo<"precision DEFAULT">]} : (tensor<?x4xf32>, tensor<4x4xf32>) -> tensor<?x4xf32>
  %2 = tensor.dim %arg0, %c0 : tensor<?x4xf32>
  %3 = tensor.from_elements %2, %c4 : tensor<2xindex>
  %4 = shape.value_as_shape %3 : tensor<2xindex> -> !shape.shape
  %5 = shape.value_as_shape %0 : tensor<1xindex> -> !shape.shape
  %6 = "mhlo.dynamic_broadcast_in_dim"(%arg2, %3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<4xf32>, tensor<2xindex>) -> tensor<?x4xf32>
  %7 = mhlo.add %1, %6 : tensor<?x4xf32>
  return %4, %5, %4, %4, %1, %3, %6, %7 : !shape.shape, !shape.shape, !shape.shape, !shape.shape, tensor<?x4xf32>, tensor<2xindex>, tensor<?x4xf32>, tensor<?x4xf32>
}

Finally, we could add some analysis based on the auxiliary function. Note that here the symbolic shape sources should be Values in the original function, not ones in the auxiliary function.

============= symbolic expr sources table for @several_ops =============
original value: %0 = "mhlo.dot_general"...
symbolic shape sources: 
<block argument> of type 'tensor<?x4xf32>' at index: 0

original value: %1 = shape.shape_of %0 : tensor<?x4xf32> -> tensor<2xindex>
symbolic shape sources: 

original value: %2 = "mhlo.dynamic_broadcast_in_dim"...
symbolic shape sources: 
<block argument> of type 'tensor<?x4xf32>' at index: 0

original value: %3 = mhlo.add %0, %2 : tensor<?x4xf32>
symbolic shape sources: 
<block argument> of type 'tensor<?x4xf32>' at index: 0

Questions

How to express constraints on the dynamic shapes? Since the auxiliary function in my current solution only address the calculation of a dynamic shape. The constraints could be useful if we meet cases like this:

a : [?, 100]
b : [?, 100]
c = concat(a, b, axis=0)

And if we know a.shape[0] + b.shape[0] = 1024 from annotation or deduction on the previous operations. Then we could know c has a static shape of [1024, 100].

I found the reifyReturnTypeShapes could be implemented in a very flexible way. Instead of shape.shape_of, developers could use tensor.dim op (or any other ops I don’t know) to get the actual dimension of the operands’ shapes. Therefore I have to a new DimOp pattern to Shape Reification pass. Is it possible to limit the flexibility of reifyReturnTypeShapes interface?
How to map the Value in the auxiliary function back to the original function? Since some passes will run on the auxiliary functions, which adds difficulty to the mapping.

jpienaar · May 31, 2022, 5:37pm

This is very timely! We have just been talking about this (well in general context, not MHLO one, with application to TOSA initially) the last couple of weeks. I’ll give a quick answers now, but will respond more completely later.

This is where meet comes in, but here you are asking (I think) more how we could take advantage of inference context info here. That seems more on the pass that uses the shape lib functions (well shape lib is the association mechanism we currently have upstream, I haven’t looked at yours yet).

It is more generic than it needs to be indeed and we should restrict it more, but I don’t think you need to add a new op to your reification, that op should be inserted in lowering to arith/tensor etc.

That is one where there are a couple of variants, in many of our uses it is currently done implicitly (we don’t just inline without context and so things are already strung along like needed so that we don’t need to keep that mapping in the IR), and there is an existing “shape associate” op that one can use (basically it ties the shape to the result). But I’m planning to send another op as suggestion “soon” and it will take inspiration from torch-mlir/TorchOps.td at 06750815d14371a7b8f0e964f2380517583e745f · llvm/torch-mlir · GitHub (now I’m still wondering a bit about overlap as now one has multiple ways of expressing it, so one could just be the higher level one)

jpienaar · May 31, 2022, 5:48pm

(llvm-project/test-shape-fn-report.mlir at main · llvm/llvm-project · GitHub for example of usage of shape.function_library op with shape association, here I had the association to ops in the mapping attribute of it while functions defined their own shape function association)

yaochengji · May 31, 2022, 6:03pm

Thanks for your reply @jpienaar . Of course, the analysis should be a generic one, not bound to a single dialect. Here I choose mhlo dialect is only because many ops in mhlo implement reifyReturnTypeShapes interface.

And as you’re also considering this work, is there any design doc could be shared publicly?

Pls take a look at the DotGeneral’s reifyReturnTypeShapes, it directly use tensor.dim here. That’s why I added the DimOp pattern.

yaochengji · May 31, 2022, 10:58pm

It seems the shape.function_library is mainly used to associate a shape inference function to an op. Still take this IR below as an example. Here a and b could be from the same op, or different ops. And we could infer that c.shape[0] = a.shape[0] + b.shape[0] according to the shape inference function. But the question is that how could we express we know a.shape[0] + b.shape[0] actually equals 1024, and how could we replace c.shape[0] with the static dimension size 1024?

a : [?, 100]
b : [?, 100]
c = concat(a, b, axis=0)

wyzero · June 1, 2022, 6:16am

Hello all, we currently are also working on this topic. Followings are some our thoughts about this problem and preliminary code:

github.com/alibaba/BladeDISC

[Design Doc] Shape constraint IR and related optimizations

opened 03:39AM - 16 May 22 UTC

closed 06:28AM - 28 Jul 22 UTC

wyzero

feature Shape Constraint IR

# Background Dynamic shape problem is ubiquitous when deploying AI workloads. A…lmost all our in-house models have dynamic shape problem in some degree. Dynamic shape problem can be divided into two categories. - Input shape of the model is dynamic * different batch size in inference scenarios * different image size, sequence length, etc - Shape is dynamic by design * For seq2seq model, the number of decoding step is dynamic * For CV detection model, the number of detected instants is dynamic * For recommendation model, shape is affected by the value of its inputs (e.g. `tf.UniqueOp`) To solve the above problem, BladeDISC, an e2e dynamic shape compiler, is proposed. We already have seen many advantages of using BladeDISC over the previous static-shape-based compiler (e.g. XLA) in many pratical scenarios. However, there are still many rooms to improvement, especially in terms of the performance. In dynamic shape semantics, the concrete shape is out of reach, making many optimization strategies used in static shape semantics invalid. According to our previous experience in optimizing dynamic shape AI workloads, we found that shape constraint, the relationship between different symbolic shape (or dimension size), is very helpful in terms of performance optimization. For example: - (partial) shape inference, propagating some known information to make best use of the existing information. * e.g. `mhlo.add(tensor<?x10xf32>, tensor<10x?xf32>) -> tensor<?x?xf32>` can be simplifier to `mhlo.add(tensor<10x10xf32>, tensor<10x10xf32>) -> tensor<10x10xf32>`. - graph optimization. For example, removing redundant `mhlo.dynamic_broadcast_in_dim` op if we can infer that the output and input of the broadcast op have the same shape. - layout optimization. For example, grouping shape-compatible ops, using same layout inside the group, saving redundant layout conversion between ops inside the groups. - Fusion decision. Making use of the shape constraint information, we can do better fusion decision in dynamic shape semantics, e.g., only fusing shape-compatible ops into one kernel. - Codegen optimization. For example, removing many redundant index calculation if we know some symbolic dimensions have the same size, which is crucial when the fusion pattern is relative large. - and many others ... In this post, we will first give a brief introduction to the definition of the shape constraint we used in the post and where we can get the shape constraint information. Then, we will show the limitations of the current strategy of making use of the shape constraint information and why we need to use shape constraint IR. At last, we will describe the details of shape constraint IR design and the related optimizations. # What is shape constraint? There are two kinds of shape constraint: - Structure shape constraint, the predicate (relationship) between different symbolic dimensions. * dimension size equality + the size of one dimension of a tensor is equal to the size of one dimension of another tensor. + the size of one dimension of a tensor is equal to the size of another dimension of the same tensor. * Number elements equality + the number of elements of a tensor is equal to the number of elements of another tensor. * group multiplication equality + e.g. reshape([a, b, c, d]) -> [a*b, c*d] * and many others ... - Shape distribution constraint * dim size % 4 = 0; * Likely values of the symbolic dimension (shape); * Possible shape ranges of the symbolic dimension (shape); * and many others ... # Where to get shape constraint information? ## Semantics of the definition of operation * The inputs and outputs of an element op of `mhlo` dialect should have same shape * The input and output of a `mhlo.dynamic_reshape` op should have same number of elements * The inputs and output of a `mhlo.concat` op should have same dim size for each non-concat dim * and many others... ![image](https://user-images.githubusercontent.com/15364516/168256459-5af1e105-b0d5-4879-9bc3-52e55ad49c13.png) ## Injected by frontend converter BladeDISC Frontend converter lowers a coarse grained op to a bunch of fine grained MHLO ops. Not only the data computation but also shape constraint should be lowered during dialect conversions. For example, `tf.SplitOp` will be lowered into a series of mho.RealDynamicSliceOp. According to the definition of the `tf.SplitOp`, we can know that all the outputs of the split op should have same shape and `the split dimension of the input % (number of outputs) == 0`. ![image](https://user-images.githubusercontent.com/15364516/168256068-272c6c45-fec9-450d-886e-1988a651cd68.png) In the current implementation, we only do data computation IR lowering, losing some shape constraint information, which in turn leads to losing some opportunities to do further optimizations. ## Symbolic Shape Analysis There are two kinds of symbolic shape analysis ways. - Propagation existing shape constraint to find more shape constraint * Input -> output, input -> input, output -> input - Shape computation IR analysis * As shown in the below example, we can know more about the dynamic reshape op by analyzing the shape computation IR (to be concrete, the IR used to calculate the target shape of the `dynamic_reshape` op). ![image](https://user-images.githubusercontent.com/15364516/168256138-ac3d62fe-9db4-47a4-ae68-77e16483891c.png) ## Injected at JIT compilation time The input shapes are actually known at JIT compilation time. We may choose to inject such information as likely shape values of the inputs of the graph and then propagate such information. ## Provided by end users User may provide more information about the distribution of the shape of inputs (e.g. shape range) or relationships of the shape of inputs (e.g. two inputs have same shape). And then, the compiler can make use of such information to do better optimizations. # Why shape constraint IR Although we have already made use of the shape constraint in some degree, there are still many limitations of the current implementation. In the current implementation, The shape constraint information itself is not stored using IR and we obtains the shape constraint information by doing global analysis on the data computation IR each time we need it. The above design choice leads to following problems: - shape constraint information loss during dialect conversion. During the whole optimization pipeline, the data computation IR is lowered many times (e.g. tf/torch dialect -> mhlo, mhlo -> lmhlo). For each lowering, we only do data computation IR lowering currently, losing some shape constraint information. - the result of shape constraint analysis is not stable after IR mutation (e.g. after running canonicalization pass). Such inconsistent result may lead to wired behaviour when making use of shape constraint information across passes. - doing analysis on the data computation IR globally to get shape constraint information is more challenging or even impossible after lowering to buffer level (e.g. lmhlo dialect). It's very hard to do analysis by chasing a series of load/store ops, especially when there exists control flow ops. - shape distribution constraint can not always be analyzed from the data computation IR. And we do not have a unify way to describe such information across different dialects. To solve the above problem, we propose using shape constraint IR to describe shape constraint information explicitly. We'll give more details about the design in following sections. # Shape constraint IR design ## Basic idea - using type instead of dedicated shape constraint op to encode shape constraint information into data computation IR. * Type-based schema not introduce any other ops, which means existing pattern-based passes (e.g. `matmul+BiadAdd` -> `FusedMatmulBiasAdd` ) not do need to take shape constraint ops into consideration. ``` // type-based encoding schema // each unique reference attribute (e.g. @S0, @S1) represents a unique symbolic dim. %2 = mhlo.add(%0, %1) : (tensor<?x?xf32, [@S0, @S1]>, tensor<?x?xf32, [@S0, @S1]>) -> tensor<?x?xf32, [@S0, @S1]> // op-based encoding schema %d0 = tensor.dim %0, %c0 : tensor<?x?xf32> %d1 = tensor.dim %0, %c1 : tensor<?x?xf32> // explicitly binding each dimension size with the corresponding tensor %n0 = disc_shape.tie_shape(%0, %d0, %d1) : tensor<?x?xf32> %n1 = disc_shape.tie_shape(%1, %d0, %d1) : tensor<?x?xf32> %2 = mhlo.add(%n1, %n2) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32> %n2 = disc_shape.tie_shape(%2, %d0, %d1) : tensor<?x?xf32> ``` - Separating shape constraint IR and data computation IR. * Some shape constraint information (e.g. shape distribution constraint) is not always possible to be bound with data computation IR. * A unified way to describe shape constraint information, no matter what dialects the data computation IR used. And thus, the shape constraint information will no longer be lost when lowering data computation IR during the whole optimization pipeline. ``` func @main() { ... %0 = any_dialect.any_operation(...) : tensor<?x?xf32, [@S0, @S1]> ... } disc_shape.SymbolDim @S0 { range list : [[...], [...], ...] likely_values : [...] ... symbolic_shape_graph: @shape_constraint_graph } disc_shape.SymbolDim @S1 { range list : [[...], [...], ...] likely_values : [...] ... symbolic_shape_graph: @shape_constraint_graph } // A separated function to store shape constraint predicates between different symbolic dimensions. // Each symbolic dim is either bound to a `disc_shape.dim` op or `disc_shape.bind_dim` func @shape_constraint_graph(...) { %0 = disc_shape.dim() {ref: @S0} : index %1 = disc_shape.dim() {ref: @S1} : index disc_shape.tie_predicate_divisible(d0, d1) // d0 % d1 == 0 // other tie_* ops // disc_shape.tie_predicate_eq(d0, d1) // d0 == d1 // disc_shape.tie_predicate_lt(d0, d1) // dim less than // disc_shape.tie_predicate_mul_eq(d0, d1, d2, ...) // d0 = d1 * d2 * ... // // d0 * d1 = s0 * s1 // disc_shape.tie_predicate_group_mul_eq([d0, d1, ..], [s0, s1, ...]) // // d0 = affine.apply(d1, d2, ...) {affine_attr = ...} // disc_shape.tie_predicate_affine_eq(d0, d1, d2, ...) {affine_attr = ...} } ``` - explicitly materialize all shape computation IR on tensor level instead of on buffer level. Currently, we only materialize all shape computation IR after bufferization (mhlo -> lmhlo). It means we have to do some shape computation IR analysis and optimization on buffer level, which could be much difficult than doing same thing on tensor level since we have to chasing each load/store ops and buffer alias to do analysis. ## Basic flow Most logic is implemented inside the shape optimization pass which is shown as following diagram. The shape optimization pass is tensor-level only pass, and is reenterable. The whole pass can be divided into two stages. - stage one: Explicitly materialize shape computation IR on tensor level. - stage two: shape constraint analysis and shape computation IR optimization on tensor level. ![image](https://user-images.githubusercontent.com/15364516/168513885-ea810cbe-10f9-4080-aa8b-46479c40e948.png) ### stage one: explicitly materialize shape computation IR on tensor level It's main functions: - (partial) shape inference, (input -> output only in the phase) * no need to do shape inference manually anymore. * taking advantage of existing `canonicalize pattern & cse pass` * example: + `concat([tensor<?x10xf32>, tensor<?x10xf32>], axis=1) -> tensor<?x20xf32>` - enable more opportunities to optimize shape computation IR * again, taking advantage of existing `canonicalize pattern & cse pass` * examples: ``` // An unified solution to handle partial known shape information - mhlo.real_dynamic_slice: - %out = %in[:,1:] - mhlo.dynamic_conv - if stride = 1, padding = kernel size, then output spatial size == input spatial size - dynamic pad - if only padding along some of the axes, - then output and input have some size along the other axes - ... // take `mhlo.real_dynamicm_slice: %out = %in[:,1:]` as an example %d0 = %tensor.dim %in, %c0 : tensor<?x?xf32> %d1 = %tensor.dim %in, %c1 : tensor<?x?xf32> %begin = tensor.from_elements %c0, %c1 : tensor<2xindex> %end = tensor.from_elements %d0, %d1 : tensor<2xindex> %strides = tensor.from_elements %c1, %c1 : tensor<2xindex> %out = mhlo.real_dynamic_slice(%in, %begin, %end, %strides) // we only need to provide a `tensor.extract_elements + tensor.from_elements` // rewrite pattern, and then we can know that the output and input have same // dim size along the first axis by taking advantage of existing // `canonicalize pattern & cse pass` ``` It's basic flow: - For a given data computation function (e.g. the main function), we insert a `disc_shape.tie_shape` for each value having `RankedTensorType` type. An example is shown as below. ``` // before inserting tie_shape op func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> { %0 = "mhlo.some_ops"(%arg0) : tensor<?xf32> return %0 : tensor<?xf32> } // after inserting tie_shape op func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> { %c0 = constant 0 : index // insert tie_shape op for %arg0 %d0_arg0 = tensor.dim %arg0, %c0 : tensor<?xf32> %new_arg0 = disc_shape.tie_shape(%arg0, %d0_arg0) %result = "mhlo.some_ops"(%new_arg0) : tensor<?xf32> // insert tie_shape op for %result %d0_result = tensor.dim %result, %c0 : tensor<?xf32> %new_result = disc_shape.tie_shape(%result, %d0_result) return %new_result : tensor<?xf32> } ``` - recursively resolve tensor.dim/shape.shape_of op based on shape interface implementation. basic idea is shown as below. ``` // take `tensor.dim` as an example, same to shape.shape_of for each tensor.dim op: operand = dimOp.getOperand(0) definingOp = operand.getDefiningOp() if (definingOp has shape interface implementation) materialize the shape interface implementation // reference implementation pass: // mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp ``` ### stage two: shape constraint analysis and shape computation IR optimization on tensor level - It’s an iterative process - The whole process can be divided into five steps: * canonicalization stage + including normal canonicalization pattern & CSE + and a bunch of new rewrite patterns - Scalarize shape computation IR whenever possible + shape.shape_of -> a bunch of tensor.dim ops + shape.broadcast(shape_a, shape_b) - Dim-wise broadcast + tensor.from_elements + bcast/reshape simplifier - ... * load existing shape constraint IR * global shape analysis + analysis shape computation IR -> find more constraint + inject shape constraint implied by `mhlo` op definition + ... * global shape optimization + using same SSA value for symbolic equal dims + refine types with partial known information globally * save the updated shape constraint into IR

The basic idea is to use shape constraint IR to model both structure shape constraint (e.g. two symbolic dimensions are equal) and shape distribution constraint (e.g. shape range/likely value).

We found that both these two types of shape constraint would be useful for performance optimization in dynamic shape scenarios.

Simply relying on shape constraint analysis is not enough to explore both of these two kinds of shape constraints.

jpienaar · June 1, 2022, 9:40pm

A way (not necessarily most concise) is

%c = shape.const_size 1024
%a0 = shape.get_extent %a, 0 : !shape.shape, index -> !shape.size
%b0 = shape.get_extent %b, 0 : !shape.shape, index -> !shape.size
%sum = shape.add %a0, %b0 : !shape.size, !shape.size -> !shape.size
%cc = shape.meet %c, %sum : !shape.size, !shape.size -> !shape.size

that establishes the equality in the IR using the constraints, this would be injected by the analysis. Now as to replacement, that would be a propagation pass. The one I was playing around with what in the context of dependent type example I was working on and was quite simple just walking along meet nodes signifying equality. Coupling that simple propagation with canonicalization & CSE worked quite well, but I was not trying to capture algebraic terms in the propagation, for that you’d probably want map key’d of affine expressions in canonical form (commutative matching etc you may need additional work on, so potentially custom mapping would work better). You can look at JITRT in TF repo for a full fledge case I believe. Although if you know they are equal, then I’d expect you’d capture it in the inference context as you are propagating it and may not even materialize it, you’d just replace shape of concat with constant and be done.

That’s interesting, can you expand on this? And do you use the same info for both or do you have dedicated patterns for each?

(@yaochengji the tie_shape example that is in the Alibaba example is effectively shape.with_shape and was equivalent to shapex.tie_shape until it was removed , and a way of result association to shape that )

Not yet, we are approaching it from a different angle and in particular I’d want to avoid reification where possible. But I want to do a couple of prototypes before larger design.

wyzero · June 2, 2022, 2:59am

The basic idea is shown as following.

func @main() {
  ...
  %0 = any_dialect.any_operation(...) : tensor<?x?xf32, [@S0, @S1]>
  ...
}

disc_shape.SymbolDim @S0 {
  range list : [[...], [...], ...]
  likely_values : [...]
  ...
  symbolic_shape_graph: @shape_constraint_graph
}

disc_shape.SymbolDim @S1 {
  range list : [[...], [...], ...]
  likely_values : [...]
  ...
  symbolic_shape_graph: @shape_constraint_graph
}

// A separated function to store shape constraint predicates between different symbolic dimensions.
// Each symbolic dim is either bound to a `disc_shape.dim` op or `disc_shape.bind_dim`
func @shape_constraint_graph(...) {
  %0 = disc_shape.dim() {ref: @S0} : index
  %1 = disc_shape.dim() {ref: @S1} : index
  disc_shape.tie_predicate_divisible(d0, d1) // d0 % d1 == 0
  
  // other tie_* ops
  //   disc_shape.tie_predicate_eq(d0, d1)  // d0 == d1
  //   disc_shape.tie_predicate_lt(d0, d1)  // dim less than
  //   disc_shape.tie_predicate_mul_eq(d0, d1, d2, ...) // d0 = d1 * d2 * ...
  //   // d0 * d1 = s0 * s1
  //   disc_shape.tie_predicate_group_mul_eq([d0, d1, ..], [s0, s1, ...])
  //   // d0 = affine.apply(d1, d2, ...) {affine_attr = ...}
  //   disc_shape.tie_predicate_affine_eq(d0, d1, d2, ...) {affine_attr = ...}
}

As shown in above IR, there are two kinds of shape constraints.

Structure shape constraint. It is the predicate (relationship) between different symbolic dimensions. For example, the size of one dimension of a tensor is equal to the size of another dimension of the same tensor.
Shape distribution constraint. For example,
d0 % 4 == 0, d0 != 1, d0 in range[2, 100], d0 is more likely to be 4

We use shape constraint IR to explicitly encode these shape constraint information instead of relying on shape constraint analysis on data computation IR each time we need it due to following reasons.

some structure shape constraint or shape distribution constraints are injected by users. Such information can not be deduced from data computation IR directly. Thus we need dedicated IR to encode such information.
some shape constraint information may be lost during dialect conversion. During dialect conversion, we may decompose a high level op into some more fine-grained low level ops. The original shape constraint information captured by the semantic of the high level op may be lost. For example, tf.SplitOp will be lowered into a series of mho.RealDynamicSliceOp. According to the definition of the tf.SplitOp , we can know that all the outputs of the split op should have same shape and the split dimension of the input % (number of outputs) == 0. Depending on the concrete lowering strategy, we may lose one of the above two shape constraints. Using shape constraint IR, we can lower not only the data computation semantic but also the corresponding shape constraint information.

jpienaar · June 2, 2022, 5:42pm

These are expressible in shape dialect today. likely isn’t. But there it is in operation, for these you have global symbolic constants here that you add into the tensor type as attribute, effectively serializing the inference context JITRT uses. Where all do you use this?

jpienaar · June 2, 2022, 5:42pm

Does this mean it can have multiple closed intervals?

ezhulenev · June 3, 2022, 6:20pm

FYI: Few links from JitRt symbolic shapes:

We do care only about equalities to remove broadcasts (and also always materialize dim of size 1), but also interested in a more generic shape analysis.

wyzero · June 6, 2022, 7:05am

Yes, it could be multiple intervals. For example, a special case is a list of discrete numbers，[[2,2], [5,5]]

wyzero · June 6, 2022, 7:40am

Some use cases that we are interested:

1, broadcast/reshape/… simplification. For example, remove redundant broadcast. Simplify mhlo.real_dynamic_slice op (e.g. find some dimensions that are actually fully sliced.)

2, Fusion decision. Making use of the shape constraint information, we can do better fusion decision in dynamic shape semantics, e.g., only fusing known shape-compatible ops into one kernel.

3, CodeGen optimization. For example, removing many redundant index calculation if we know some symbolic dimensions have the same size, which is important when the fusion pattern is relative large. Tiling/vectorization bound-check could be simplified if we know s0 % tile_size/vector_size == 0

4, using likely values to help us do better speculation decision (e.g. placement, layout, codegen schedule, implicit broadcast speculation).

herhut · June 9, 2022, 2:41am

We do this based on the output of the shape analysis that was linked earlier I believe. The rewrite patterns query the shape analysis. The main difference in approach is that we do not materialize the result of the shape analysis itself in IR but only properties of operations that we have inferred (like whether a broadcast is expanding, a reshape is just a collapse/expand, etc.).

We have seen that we need some caching of shape analysis results to avoid rerunning it but shied away from using the IR for this, as it requires that we keep that information correct when we rewrite. @wyzero have you seen this as an issue that arises in practice?

Our approach here has been that we assume dimensions line up for all HLO operations. We use an encoding with shape.assuming regions to guard operations with their prerequisites and then propagate those to form fusion islands.

Interesting. Do you have a code pointer for this? I am curious to see how this works. Do you make index computations explicit at the IR level and then cse dim operations (conceptually) where you know they are the same?

wyzero · June 10, 2022, 1:44pm

1, I agree that explicitly using IR to store shape constraint information may introduce more work when doing transformation to make the IR consistent. However, some kinds of shape constraints, for example, shape distribution constraints, can not be naturally captured by the semantics of data computation IR. These kinds of shape constraints are also useful in terms of performance optimization. On the other hand, we are not trying to store the calculation logic for one symbolic shape in another place again, instead we are just trying to store some high level predicates between different symbolics (e.g. a < b, a % b == 0, a * b = c * d). To reduce the overhead of maintaining the shape constraint IR, we also provide a shape optimization pass to automatically infer and propagate shape constraint information captured by the semantics of data computation IR (e.g. mhlo ops), thus the authors of the rewrite pattern only need to insert those lost shape constraints (like the above example: tf.split → milo).

2, index calculation simplification. The basic ideas is shown as following.

- Binding all symbolic-equal dimensions to the same SSA value

// Step #1: original input graph
func @main(%arg0 : tensor<?x?xf32, [@S0, @S0]>, %arg1 : tensor<?x?xf32, [@S0, @S0]>) -> tensor<?x?xf32, [@S0, @S0]> {
  // a specialized version of mhlo.add: squared add
  %ret = mhlo.add(%arg0, %arg1) : tensor<?x?xf32, [@S0, @S0]>
  return %ret : tensor<?x?xf32, [@S0, @S0]>
}

// Step #2: Explicitly binding all known symbolic-equal dimensions to the same SSA value right before going from tensor world to buffer world.

func @main(%arg0 : tensor<?x?xf32, [@S0, @S0]>, %arg1 : tensor<?x?xf32, [@S0, @S0]>) -> tensor<?x?xf32, [@S0, @S0]> {
  %c0 = constant 0: index
  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32, [@S0, @S0]>
  %new_arg0 = disc_shape.tie_shape(%arg0, %d0, %d0) : tensor<?x?xf32, [@S0, @S0]>
  %new_arg1 = disc_shape.tie_shape(%arg1, %d0, %d0) : tensor<?x?xf32, [@S0, @S0]>
  %ret = mhlo.add(%new_arg0, %new_arg1) : tensor<?x?xf32, [@S0, @S0]>
  %new_ret = disc_shape.tie_shape(%ret, %d0, %d0) : tensor<?x?xf32, [@S0, @S0]>
  return %new_ret : tensor<?x?xf32, [@S0, @S0]>
}

// bufferize to lmhlo
//   - Symbolic dims are resolved into the same SSA value
//   - Enable more opportunities for normal CSE & canonicalization patterns.

func @main(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>) -> memref<?x?xf32> {
  %c0 = constant 0 : index
  %d0 = memref.dim %arg0, %c0 : memref<?x?xf32>
  %new_arg0 = memref.reinterpret_cast %arg0 to offset: [0], sizes: [%d0, %d0], …
  %new_arg1 = memref.reinterpret_cast %arg1 to offset: [0], sizes: [%d0, %d0], …
  %ret = memref.alloc(%d0, %d0) : memref<?xf32>
  "lmhlo.add"(%new_arg0, %new_arg1, %ret)
  return %ret : memref<?x?xf32>
}

Example use case: eliminate redundant linearize/de-linearize pair in the fusion pattern (relying on the first step).

Suppose we have a fusion pattern abs -> reshape -> transpose -> exp and L(op) denotes the linear index of a op, multi-dim(op) denotes the multi-dim indices for the op. We can remove the last pair of linearize-and-delinearize (suppose start from the output).

The feature is still WIP, part of the code is listed as following:

Using the same SSA value for all known symbolic-equal dimensions whenever possible, link
Flatten mermen accesses explicitly, link
Linearize and Delinearize pair elimination, link

`

yaochengji · June 11, 2022, 7:03pm

@jpienaar
Here %c is a const, therefore the propagation pass could know that it’s better to replace %sum with %c.

But for a generic case, shape.meet %x, %y, how could we decide %x should be replace by %y, or vise versa?

_sean_silva · June 15, 2022, 5:54pm

Hey folks, I’m a little late to this thread, but I think you folks might find interesting what we are doing in Torch-MLIR. We have a pretty sophisticated shape refinement system (doc). It is similar to a lot of the approaches in this thread, but has a few characteristics of note:

The shape functions are written in a restricted Python subset, which has many benefits:
a. It is extremely easy to author, test, and debug the shape functions.
b. They are shared with upstream PyTorch here.
c. It is trivially easy to insert certain kinds of shape guards and assertions (and test them!)
We found that a nontrivial amount of “list computation” and control flow is needed when dealing with operators which are allowed to have multiple ranks – for example iterating over all the batch dimensions of an operator and building the output shape.
We reuse our TorchScript compiler to compile the shape functions themselves.
Due to how we model our reified shape IR, identifying symbolically identical dimensions comes naturally from CSE after simplification.

Here example IR for a trivial unary function gist. Note that the shape function might look scary in the IR, but it is really just this trivial Python function link. Upstream PyTorch decided that they would like to make sure that the shape functions never return aliased lists, so we make a copy. Handling this falls out naturally from the more general types of shape functions we need to handle.

Here are some examples of some pretty interesting shape functions that we handle:
max_pool_2d - very complex logic and error checks
matmul - handles many different rank combinations
linear - demonstrates composition of other shape functions
transpose - conditional logic inside a loop is the natural way to express this shape function

In our case, we have to deal with a large and wild set of operators, and our shape functions are the Torch-level shape functions which have to be defensive against arbitrary user inputs. For a system like MHLO which is more orthogonalized and has more verifier-enforced properties, other choices probably make sense.

herhut · June 20, 2022, 5:11pm

Thanks @wyzero for sharing the details.

Have you seen cases where a rewrite would make inferring constraints impossible? A case where the shape constraint propagation would need to coarsen a constraint, as it can no longer be derived from the program? Or would you, in such case, keep the previous, more precise constraints and just have weaker ones for rewritten IR?

wyzero · June 23, 2022, 6:02am

I think the above mentioned tf.Split could be an example? we need to inject split_dim % num_outputs == 0 when lowering it to a bunch of mhlo real_dynamic_slice ops.

jpienaar · July 10, 2022, 7:13pm

Sorry missed this question, these two are equivalent/both members of same equivalence class post the meet, in this case you’d replace both x and y with the result of the meet (true in const or dynamic case). But you may be asking inside the generated shape method/I may be answering a different question.