Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir][tosa] Work around GCC bug in tosa-to-tensor #91521

Merged
merged 1 commit into from
May 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 5 additions & 6 deletions mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -55,20 +55,17 @@ TensorType inferReshapeExpandedType(TensorType inputType,
// Check if the input is static, and if so, get its total size
bool inputIsStatic = inputType.hasStaticShape();
int64_t totalSize = inputIsStatic ? inputType.getNumElements() : -1;

// Compute result shape
bool resultIsStatic = true;
auto resultShape = llvm::map_to_vector(newShape, [&](int64_t size) -> int64_t {
// If this is not a placeholder, do not change it
if (size >= 0)
return size;

// If we do not know the total size of the tensor, keep this dimension
// dynamic in the result shape.
if (!inputIsStatic) {
resultIsStatic = false;
if (!inputIsStatic)
return ShapedType::kDynamic;
}

// Calculate the product of all elements in 'newShape' except for the -1
// placeholder, which we discard by negating the result.
Expand All @@ -84,12 +81,14 @@ TensorType inferReshapeExpandedType(TensorType inputType,
return totalSize / totalSizeNoPlaceholder;
});

bool resultIsStatic = !ShapedType::isDynamicShape(resultShape);

// A syntactic restriction in 'tensor.expand_shape' forbids a dynamically
// shaped input from being reshaped into a statically shaped result. We may
// simply turn the first result dimension dynamic to address this.
if (!inputIsStatic && resultIsStatic)
resultShape[0] = ShapedType::kDynamic;

// The 'tensor.expand_shape' op also forbids a statically shaped input from
// being reshaped into a dynamically shaped result, but the placeholder
// inference algorithm above guarantees that this will never be the case.
Expand Down
15 changes: 15 additions & 0 deletions mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -394,6 +394,21 @@ func.func @test_reshape_6d_down_s2s_auto(%arg0: tensor<1x2x3x5x7x11xf32>) -> ten

// -----

// This test would previously fail on GCC with certain compiler flags.
// The GCC issue would cause invalid IR after tosa-to-tensor, so this test
// locks down that the code goes through tosa-to-tensor and verifies.
//
// See https://github.com/llvm/llvm-project/pull/91521 for a full description.

// CHECK-LABEL: reshape_bug_fix
// CHECK: tensor.expand_shape
func.func @reshape_bug_fix(%arg0: tensor<?xf32>) -> tensor<1x1x1x?xf32> {
%0 = tosa.reshape %arg0 {new_shape = array<i64: 1, 1, 1, -1>} : (tensor<?xf32>) -> tensor<1x1x1x?xf32>
return %0 : tensor<1x1x1x?xf32>
}

// -----

// CHECK-LABEL: test_reshape_6d_down_s2s_explicit
// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<1x2x3x5x7x11xf32>
// CHECK: %[[VAL_0:.*]] = tensor.collapse_shape %[[ARG_0]] {{\[\[}}0, 1, 2], [3], [4, 5]] : tensor<1x2x3x5x7x11xf32> into tensor<6x5x77xf32>
Expand Down
Loading