sophgo / tpu-mlir

Machine learning compiler based on MLIR for Sophgo TPU.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Implement shape_inference for each operator

charlesxzb opened this issue · comments

Background:

Sometimes we may encounter situations where the output types of some model layers are UnrankedTensorType (i.e., <* x dtype>, there is no rank and dimension info in Type) or the shapes of the original model layers are incorrect. To handle these situations, we need to perform shape inference in TPU-MLIR.

Currently, the shape inference mechanism has been implemented in TPU-MLIR. The output shape of each op will be checked. If there is an op without output shape info, we will do shape inference for it and set its output shape as the inference result. If shape info exists, the shape will be verified by comparing it with the inference result.

However, this mechanism is not completed yet because we still have to implement the shape inference function for each operator, so we hope you can finish it with us.

Tasks:

For each operator, implement its inference function in the corresponding cpp file under the directory: tpu-mlir/lib/Top/Interfaces/.

Example: shape_inference function of ConvOp
image

In addition, if there is a general function for a certain type of operators, please implement it in tpu-mlir/lib/Interfaces/ShapeInterface.cpp, then call it in the shape_inference function of these operators. For example, for activation operators, we use the common_shape_inference function, which simply set the output shape as the input shape, to do the shape inference.

Example of general shape inference function
image

To Test:

  1. Make sure you have module::setShapeOrVerify(out, in_shape) (which is used to set the output shape or check the correctness of the current shape) at the end of the shape inference function.
  2. test_onnx.py --case op_name
  3. The shape_inference function is valid if the test success.

Note:

A PR can be submitted once every shape_inference function is completed.

中文版------------------------------------------------------------------------------------------------

背景:

有时我们可能会遇到某些模型层的输出类型为 UnrankedTensorType(即 <* x dtype>,Type中不包含维度与维数信息)或原始模型层的shape不正确的情况。 为了应对这些情况,我们需要在 TPU-MLIR 中进行shape推理。

目前,TPU-MLIR 已经实现了Shape推理的机制。 即遍历每个算子的输出形状, 如果某个算子没有输出shape信息,我们就对其进行shape推理,并将它的输出shape设置为推理的结果,如果存在输出shape信息,则会通过将其与推理结果进行比较来验证shape的正确性。

但是目前这个机制还没有完成,因为我们还需要为每个算子实现对应的shape_inference方法,所以我们希望能获得各位的帮助来一起完善这个机制。

任务:

对于每一个算子,在**tpu-mlir/lib/Top/Interfaces/**目录下对应的cpp文件中实现其shape_inference方法。

示例:ConvOp 的 shape_inference 函数
image

另外,如果有针对某类算子的通用shape_infernece方法,请tpu-mlir/lib/Interfaces/ShapeInterface.cpp下实现,然后在这些算子的shape_inference方法中调用。 例如,对于激活算子,我们统一使用common_shape_inference函数,它只是单纯地将输出shape设置为输入shape。

通用shape推理方法示例
image

测试方式:

  1. 确保在shape_inference方法的末尾有 module::setShapeOrVerify(out, in_shape) (用于设置输出shape或检查当前shape的正确性)。
  2. test_onnx.py --case op_name
  3. 如果测试成功,则shape_inference方法正确

提示:

每完成一个算子的shape_inference方法都可以提交一次PR