-
Notifications
You must be signed in to change notification settings - Fork 14.4k
[mlir][memref] Support test-compose-subview dynamic size #146881
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
base: main
Are you sure you want to change the base?
[mlir][memref] Support test-compose-subview dynamic size #146881
Conversation
@llvm/pr-subscribers-mlir-memref Author: lonely eagle (linuxlonelyeagle) Changessupports the case where the sizes of the subview op is dynamic. Full diff: https://github.com/llvm/llvm-project/pull/146881.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp b/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp
index d25ddb41aa4eb..020d08ddda408 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp
@@ -81,10 +81,6 @@ struct ComposeSubViewOpPattern : public OpRewritePattern<memref::SubViewOp> {
for (auto &&[opOffset, sourceOffset, sourceStride, opSize] :
llvm::zip(op.getMixedOffsets(), sourceOp.getMixedOffsets(),
sourceOp.getMixedStrides(), op.getMixedSizes())) {
- // We only support static sizes.
- if (isa<Value>(opSize)) {
- return failure();
- }
sizes.push_back(opSize);
Attribute opOffsetAttr = llvm::dyn_cast_if_present<Attribute>(opOffset),
sourceOffsetAttr =
diff --git a/mlir/test/Transforms/compose-subview.mlir b/mlir/test/Transforms/compose-subview.mlir
index 53fbb8a356def..8377820821744 100644
--- a/mlir/test/Transforms/compose-subview.mlir
+++ b/mlir/test/Transforms/compose-subview.mlir
@@ -101,3 +101,29 @@ func.func @subview_strided(%input: memref<4x1024xf32>) -> memref<1x64xf32, strid
%1 = memref.subview %0[%cst_1, 64] [1, 64] [2, 2] : memref<2x256xf32, strided<[2048, 2], offset: ?>> to memref<1x64xf32, strided<[4096, 4], offset: ?>>
return %1 : memref<1x64xf32, strided<[4096, 4], offset: ?>>
}
+
+// -----
+
+// CHECK-LABEL: func.func @single_dynamic_size_subview(
+// CHECK-SAME: %[[SRC:.*]]: memref<256x?xf32, strided<[?, ?], offset: ?>>,
+// CHECK-SAME: %{{.*}}: index,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> memref<8x?xf32, strided<[?, ?], offset: ?>> {
+func.func @single_dynamic_size_subview(%arg0: memref<256x?xf32, strided<[?, ?], offset: ?>>, %arg1 : index, %arg2 : index) -> memref<8x?xf32, strided<[?, ?], offset: ?>>{
+ %subview = memref.subview %arg0[0, 0][8, %arg1][1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<8x?xf32, strided<[?, ?], offset: ?>>
+ %subview_1 = memref.subview %subview[0, 0][8, %arg2][1, 1] : memref<8x?xf32, strided<[?, ?], offset: ?>> to memref<8x?xf32, strided<[?, ?], offset: ?>>
+ // CHECK: %{{.*}} = memref.subview %[[SRC]][0, 0] [8, %[[SIZE]]] [1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<8x?xf32, strided<[?, ?], offset: ?>>
+ return %subview_1 : memref<8x?xf32, strided<[?, ?], offset: ?>>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @all_dynamic_size_subview(
+// CHECK-SAME: %[[SRC:.*]]: memref<256x?xf32, strided<[?, ?], offset: ?>>,
+// CHECK-SAME: %{{.*}}: index,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+func.func @all_dynamic_size_subview(%arg0: memref<256x?xf32, strided<[?, ?], offset: ?>>, %arg1: index, %arg2: index) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+ %subview = memref.subview %arg0[0, 0] [%arg2, %arg2] [1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
+ // CHECK: %{{.*}} = memref.subview %[[SRC]][0, 0] {{\[}}%[[SIZE]], %[[SIZE]]] [1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
+ return %subview : memref<?x?xf32, strided<[?, ?], offset: ?>>
+}
+
|
@llvm/pr-subscribers-mlir Author: lonely eagle (linuxlonelyeagle) Changessupports the case where the sizes of the subview op is dynamic. Full diff: https://github.com/llvm/llvm-project/pull/146881.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp b/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp
index d25ddb41aa4eb..020d08ddda408 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp
@@ -81,10 +81,6 @@ struct ComposeSubViewOpPattern : public OpRewritePattern<memref::SubViewOp> {
for (auto &&[opOffset, sourceOffset, sourceStride, opSize] :
llvm::zip(op.getMixedOffsets(), sourceOp.getMixedOffsets(),
sourceOp.getMixedStrides(), op.getMixedSizes())) {
- // We only support static sizes.
- if (isa<Value>(opSize)) {
- return failure();
- }
sizes.push_back(opSize);
Attribute opOffsetAttr = llvm::dyn_cast_if_present<Attribute>(opOffset),
sourceOffsetAttr =
diff --git a/mlir/test/Transforms/compose-subview.mlir b/mlir/test/Transforms/compose-subview.mlir
index 53fbb8a356def..8377820821744 100644
--- a/mlir/test/Transforms/compose-subview.mlir
+++ b/mlir/test/Transforms/compose-subview.mlir
@@ -101,3 +101,29 @@ func.func @subview_strided(%input: memref<4x1024xf32>) -> memref<1x64xf32, strid
%1 = memref.subview %0[%cst_1, 64] [1, 64] [2, 2] : memref<2x256xf32, strided<[2048, 2], offset: ?>> to memref<1x64xf32, strided<[4096, 4], offset: ?>>
return %1 : memref<1x64xf32, strided<[4096, 4], offset: ?>>
}
+
+// -----
+
+// CHECK-LABEL: func.func @single_dynamic_size_subview(
+// CHECK-SAME: %[[SRC:.*]]: memref<256x?xf32, strided<[?, ?], offset: ?>>,
+// CHECK-SAME: %{{.*}}: index,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> memref<8x?xf32, strided<[?, ?], offset: ?>> {
+func.func @single_dynamic_size_subview(%arg0: memref<256x?xf32, strided<[?, ?], offset: ?>>, %arg1 : index, %arg2 : index) -> memref<8x?xf32, strided<[?, ?], offset: ?>>{
+ %subview = memref.subview %arg0[0, 0][8, %arg1][1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<8x?xf32, strided<[?, ?], offset: ?>>
+ %subview_1 = memref.subview %subview[0, 0][8, %arg2][1, 1] : memref<8x?xf32, strided<[?, ?], offset: ?>> to memref<8x?xf32, strided<[?, ?], offset: ?>>
+ // CHECK: %{{.*}} = memref.subview %[[SRC]][0, 0] [8, %[[SIZE]]] [1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<8x?xf32, strided<[?, ?], offset: ?>>
+ return %subview_1 : memref<8x?xf32, strided<[?, ?], offset: ?>>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @all_dynamic_size_subview(
+// CHECK-SAME: %[[SRC:.*]]: memref<256x?xf32, strided<[?, ?], offset: ?>>,
+// CHECK-SAME: %{{.*}}: index,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+func.func @all_dynamic_size_subview(%arg0: memref<256x?xf32, strided<[?, ?], offset: ?>>, %arg1: index, %arg2: index) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+ %subview = memref.subview %arg0[0, 0] [%arg2, %arg2] [1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
+ // CHECK: %{{.*}} = memref.subview %[[SRC]][0, 0] {{\[}}%[[SIZE]], %[[SIZE]]] [1, 1] : memref<256x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
+ return %subview : memref<?x?xf32, strided<[?, ?], offset: ?>>
+}
+
|
df898d1
to
8c5ea51
Compare
Could you please describe why are you making this change in the commit message? |
For the following gemm algorithm, I do tiles on m, n, and k.I would like to be able to compose-subview, but I find that test-compose-subview doesn't work. So I made a modification to the pass, and so far it looks like there is no problem with such a modification.
|
I made a detailed addition to the commit. |
@ftynse Have you forgotten about this PR? |
Supports the case where the sizes of the subview op is dynamic.When there are more for loops in the tile algorithm, multiple subviews are performed and test-compose-subview does not work when the size operand of the subview ops is dynamic value.