Segfault when running shape inference on Tosa ops

MLIR
1

I’m attempting to use the OpTrait::InferTensorType::inferReturnTypes function to infer the return type when creating a Tosa::TransposeOp with operands of known shape.

This results in a segfault in mlir::detail::inferReturnTensorTypes on line:

inferredReturnTypes.push_back(RankedTensorType::get(
          shapeAndType.getDims(), shapeAndType.getElementType()));

because shapeAndType’s element type is nullptr, which results in a segfault in RankedTensorType::get.
In this case shapeAndType is coming from tosa::TransposeOp::inferReturnTypeComponents

Am I misusing the interface or is this a bug? If it’s a bug, what’s the best way to approach fixing it? I assume we could just pick the element type of input1 for the transpose op, but this becomes more complex for other operations where we’d need to combine types.

Is there any plans in the future for the Tosa ops to implement InferTypeOpInterface so that the shape computations happen automatically?

2

InferReturnTypeComponents need not populate all elements. So it may succeed while only populating some of them and in this case that excludes element type. It is weird that element type is not populated though, perhaps just an oversight as interface is only used for shape refinement in the current flow.

3

I think that means that InferTensorType will always result in a segfault whenever InferReturnTypeComponents chooses not to populate element type? If so should there be a check that shapeAndType has an element type to avoid segfaulting?

Instead of using InferTensorType we’re using InferReturnTypeComponents directly which works fine for us and seems to be how it’s used elsewhere.

4

Using InferTensorTypes assumes that InferTypeOpInterface is set on the op along with return type component interface, which is why there is an interface list associated with it that adds both of these. We could (I think) add an explicit check that the op class that calls it has the trait that signifies that it will always populate sufficient information to determine type or fail (with return type components it may not even populate shapes either and pass, it may only populate element types, so there needs the additional verification from the user.