Questions regarding TOSA shape type

Hey,
I have few questions regarding recent changes took place at TOSA dialect.
There is a new type that has been declared into TOSA dialect which is called
!tosa.shape in this commit [TOSA] Add Tosa_Shape type and ConstShapeOp by Jerry-Ge · Pull Request #122547 · llvm/llvm-project · GitHub

For example, !tosa.shape<3> indicates a shape with rank 3, my first question is what the
motivation behind such a type, why the old RankedTensorType is not enough to indicate such
a requirement?

Also, having the MLIR below

func.func @test_const_shape() -> !tosa.shape<6> {
  %cst = tosa.const_shape {value = dense<[[3, 1, 2] , [2, 3, 1]> : tensor<2x3xindex>} : () -> !tosa.shape<6>
  return %cst : !tosa.shape<6>
}

We are trying to create a shape of rank 6 from six values of type tensor<2x3xindex> of rank 2:

  • Isn’t it the same operation of the shape.const_shape, why not using it instead of declaring additional operation?
  • Is there any meaning of not having a flat/ one-dimensional values in the const_shape?
  • It feels a little bit overloading for the term of rank, we verify that the number of the elements in the DenseElementsAttribute equals the TosaShape’s rank. why not calling it num of elements instead of rank?

we introduced !tosa.shape type in order to allow step-by-step introduction of shape expressions into TOSA dialect.

The trait, TosaResolvaleShapeOperands, placed on all TOSA operators, checks that any operand of !tosa.shape type has a producer (ie, function inputs not allowed), and that producer is a tosa shape operator. Initially, only one shape operator exists, namely const_shape, therefore we only allow constant shapes for now.

The next shape operator (ie, operators that produce !tosa.shape values) we plan to upstream is tosa.dim, which extracts a specific dimension of the shape of a value and returns it as !tosa.shape<1>.

Additional shape operators like add_shape, sub_shape, mul_shape, div_shape, concat_shape, slice_shape, etc, will then allow construction of shape expressions based on const_shape and dim operator values.

The goal is to be able to lower dynamic network into TOSA by constructing shape expressions based on dim operators, then when the dynamic shapes are specialized to static shapes, all dim operator outputs become constants, and we can fold all shape expressions into constant values.

Now, on the specific example you had for const_shape. The !tosa.shape values are shapes, which are 1 dimensional vector of N elements, where N is the rank of the shapes. So !tosa.shape<6> is the rank-6 shape type. The attribute value in your example really should be { value = dense<[3, 1, 2, 2, 3, 1]> : tensor<6xindex> }. We will add a verifier to check for that.

Hope I’ve answered your questions.

2 Likes

Thanks for the detailed answer, It has answered my question :raised_hands:

1 Like

Thank you for your interest and we apologize for the delay in responding. Please feel free to tag us by name or you can catch us on Discord as well!