Skip to content

[mlir][Vector] Remove 0-d corner case condition. #112937

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.

Sign up for GitHub

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

Closed
harrisonGPU wants to merge 8 commits into from
Closed

[mlir][Vector] Remove 0-d corner case condition. #112937

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
5e07a46
[mlir][Vector] Support 0-d vectors natively in VectorStoreToMemrefSto…
harrisonGPU Nov 5, 2024
e5b8fac
Merge branch 'main' into harrison/mlir
harrisonGPU Nov 5, 2024
3f9f7b0
[MLIR] Fix the lit test failure issue.
harrisonGPU Nov 5, 2024
ec6b5ab
Merge branch 'main' into harrison/mlir
harrisonGPU Nov 5, 2024
5f61d16
[MLIR] Update comments.
harrisonGPU Nov 6, 2024
6db82e0
[MLIR] Update comments again.
harrisonGPU Nov 6, 2024
6a5ac3c
[MLIR][Vector] Remove 0-d corner case condition.
harrisonGPU Nov 10, 2024
700b0ea
[MLIR] Remove Transfer vectore lower pattern.
harrisonGPU Nov 15, 2024
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
14 changes: 4 additions & 10 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -518,7 +518,7 @@ struct VectorLoadToMemrefLoadLowering
}
};

/// Replace a 0-d vector.store with a vector.extractelement + memref.store.
/// Replace a vector.store with a vector.extractelement + memref.store.
struct VectorStoreToMemrefStoreLowering
: public OpRewritePattern<vector::StoreOp> {
using OpRewritePattern::OpRewritePattern;
Expand All @@ -530,15 +530,9 @@ struct VectorStoreToMemrefStoreLowering
return rewriter.notifyMatchFailure(storeOp, "not single element vector");

Value extracted;
if (vecType.getRank() == 0) {
// TODO: Unifiy once ExtractOp supports 0-d vectors.
extracted = rewriter.create<vector::ExtractElementOp>(
storeOp.getLoc(), storeOp.getValueToStore());
} else {
SmallVector<int64_t> indices(vecType.getRank(), 0);
extracted = rewriter.create<vector::ExtractOp>(
storeOp.getLoc(), storeOp.getValueToStore(), indices);
}
SmallVector<int64_t> indices(vecType.getRank(), 0);
Copy link
Contributor

@banach-space banach-space Nov 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment above this pattern needs updating.

Copy link
Contributor Author

@harrisonGPU harrisonGPU Nov 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, I have already updated it.

Copy link
Contributor

@banach-space banach-space Nov 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment is still referring to vector.extractelement, which your patch gets rid of the vector.extractelement path.

extracted = rewriter.create<vector::ExtractOp>(
storeOp.getLoc(), storeOp.getValueToStore(), indices);

rewriter.replaceOpWithNewOp<memref::StoreOp>(
storeOp, extracted, storeOp.getBase(), storeOp.getIndices());
Expand Down
5 changes: 2 additions & 3 deletions mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2971,9 +2971,8 @@ func.func @vector_store_op_0d(%memref : memref<200x100xf32>, %i : index, %j : in
// CHECK-LABEL: func @vector_store_op_0d
// CHECK: %[[val:.*]] = arith.constant dense<1.100000e+01> : vector<f32>
// CHECK: %[[cast:.*]] = builtin.unrealized_conversion_cast %[[val]] : vector<f32> to vector<1xf32>
// CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i64
// CHECK: %[[extracted:.*]] = llvm.extractelement %[[cast]][%[[c0]] : i64] : vector<1xf32>
// CHECK: memref.store %[[extracted]], %{{.*}}[%{{.*}}, %{{.*}}]
// CHECK: %[[cast2:.*]] = builtin.unrealized_conversion_cast %[[cast]] : vector<1xf32> to f32
// CHECK: memref.store %[[cast2]], %{{.*}}[%{{.*}}, %{{.*}}]

// -----

Expand Down
8 changes: 3 additions & 5 deletions mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,13 @@ func.func @vector_transfer_ops_0d_memref(%mem: memref<f32>, %vec: vector<1x1x1xf
%f0 = arith.constant 0.0 : f32

// CHECK-NEXT: %[[S:.*]] = memref.load %[[MEM]][] : memref<f32>
// CHECK-NEXT: %[[V:.*]] = vector.broadcast %[[S]] : f32 to vector<f32>
Copy link
Contributor

@banach-space banach-space Nov 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is this broadcast generated? Perhaps that's something that could deleted?

Copy link
Contributor Author

@harrisonGPU harrisonGPU Nov 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because vector::ExtractElementOp requires a vector type input, we need to use vector.broadcast to convert scalars into vectors. We directly use vector::ExtractOp, and even for 0-dimensional vectors, there's no longer a need to perform vector.broadcast. :-)

Copy link
Contributor

@banach-space banach-space Nov 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes sense, but what code-path generated vector.broadcast to begin with? And how do we know it doesn't require updating?

%0 = vector.transfer_read %mem[], %f0 : memref<f32>, vector<f32>

// CHECK-NEXT: %[[SS:.*]] = vector.extractelement %[[V]][] : vector<f32>
// CHECK-NEXT: memref.store %[[SS]], %[[MEM]][] : memref<f32>
// CHECK-NEXT: memref.store %[[S]], %[[MEM]][] : memref<f32>
vector.transfer_write %0, %mem[] : vector<f32>, memref<f32>

// CHECK-NEXT: %[[VV:.*]] = vector.extract %arg1[0, 0, 0] : f32 from vector<1x1x1xf32>
// CHECK-NEXT: memref.store %[[VV]], %[[MEM]][] : memref<f32>
// CHECK-NEXT: %[[V:.*]] = vector.extract %arg1[0, 0, 0] : f32 from vector<1x1x1xf32>
// CHECK-NEXT: memref.store %[[V]], %[[MEM]][] : memref<f32>
vector.store %vec, %mem[] : memref<f32>, vector<1x1x1xf32>

return
Expand Down
Loading