maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][linalg] Fixed issue generating reassociation map with Rank-0 types 

Rank-0 case causes a graph during linalg reshape operation.

Differential Revision: https://reviews.llvm.org/D102282
This commit is contained in:
Rob Suderman 2021-05-11 15:31:20 -07:00
parent 1470b8587f
commit 7b57517507
2 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
View File

@ -1134,10 +1134,12 @@ mlir::linalg::getReassociationIndicesForReshape(ShapedType sourceType,
return llvm::None;
currIndices.push_back(sourceDim++);
// If there are no dimensions in the target to match, then append the
// `currIndices` to the last element of the reassociationMap.
// If the reassociation is empty but the currIndices is not, this by
// definition is folding unit-dimensions with the result being scalar type.
// So only append the `currIndices` if reassociation map is not empty.
if (targetDim == targetShape.size()) {
reassociationMap.back().append(currIndices.begin(), currIndices.end());
if (!reassociationMap.empty() && !currIndices.empty())
reassociationMap.back().append(currIndices.begin(), currIndices.end());
// Break out of the loops. We should be done here.
break;
}

11
mlir/test/Dialect/Linalg/canonicalize.mlir
View File

@ -43,6 +43,17 @@ func @memref_cast_into_tiled_loop(%arg0: memref<192xf32>) {
// -----
// CHECK-LABEL: zero_rank_reshape_multi
func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: return %arg0
%0 = linalg.tensor_reshape %arg0 [] : tensor<f32> into tensor<1xf32>
%1 = linalg.tensor_reshape %0 [[0, 1]] : tensor<1xf32> into tensor<1x1xf32>
%2 = linalg.tensor_reshape %1 [] : tensor<1x1xf32> into tensor<f32>
return %2 : tensor<f32>
}
// -----
func @collapsing_tensor_reshapes(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32>
{
%0 = linalg.tensor_reshape %arg0 [[0, 1], [2], [3, 4]]