mlir CollapseShapeOp ExpandShapeOp的构造
作用
CollapseShapeOp把一些维度合并,总Tensor的size不变;ExpandShapeOp把一些维度拆开,总Tensor的size不变。
create方式
SmallVector<int64_t> rhsShapeSimplified;
auto targetRhsType = RankedTensorType::get(rhsShapeSimplified, cast<RankedTensorType>(rhsType).getElementType());
rhsToRewriter =rewriter.create<tensor::CollapseShapeOp>(op->getLoc(), targetRhsType(目标type), op->getOperand(1)(原始value), rhsReassociation);
Reassociation: 输入Tensor到输出Tensor对应的维度坐标信息:把哪些维度合并在一起
rhsReassociation类型是SmallVector本质是SmallVector<SmallVector<int64_t>>
tensor<1x1x1x4x3xf16> reassociation: [[0, 1, 2, 3], [4]];
tensor<1x3x4xf16> reassociation: [[0, 1], [2]]auto targetOutType = RankedTensorType::get(outShape, cast<RankedTensorType>(outType).getElementType());lastToReplace = rewriter.create<tensor::ExpandShapeOp>(op->getLoc(), targetOutType, matmul, outReassociation);
Reassociation: 输出Tensor到输入Tensor对应的维度坐标信息:把哪些维度合并在一起
获取Reassociation的方式,可以直接调三方接口
std::optional<SmallVector<ReassociationIndices>>
mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,ArrayRef<int64_t> targetShape)