Skip to content

[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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

linuxlonelyeagle
Copy link
Member

@linuxlonelyeagle linuxlonelyeagle commented Jul 3, 2025

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.

@llvmbot
Copy link
Member

llvmbot commented Jul 3, 2025

@llvm/pr-subscribers-mlir-memref

Author: lonely eagle (linuxlonelyeagle)

Changes

supports 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:

  • (modified) mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp (-4)
  • (modified) mlir/test/Transforms/compose-subview.mlir (+26)
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: ?>>
+}
+

@llvmbot
Copy link
Member

llvmbot commented Jul 3, 2025

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

Changes

supports 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:

  • (modified) mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp (-4)
  • (modified) mlir/test/Transforms/compose-subview.mlir (+26)
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: ?>>
+}
+

@linuxlonelyeagle linuxlonelyeagle force-pushed the advance-support-compose-subview branch from df898d1 to 8c5ea51 Compare July 3, 2025 13:13
@linuxlonelyeagle linuxlonelyeagle changed the title [mlir][memref]support test-compose-subview dynamic size [mlir][memref] Support test-compose-subview dynamic size Jul 4, 2025
@ftynse
Copy link
Member

ftynse commented Jul 5, 2025

Could you please describe why are you making this change in the commit message?

@linuxlonelyeagle
Copy link
Member Author

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.

for xxx step %tile_m { // tile m
   %subview_m = subview [][%tile_m, size0][]
  for xxx { // tie n
     %subview_n = subview [][size1, %tile_n][]
    for xxx step %tile_k { // tile k
        %subview_subview_m = subview %subview_m[][%tile_m, %tike_k][]
         ....
        gemm
    }
 }
}

@linuxlonelyeagle
Copy link
Member Author

Could you please describe why are you making this change in the commit message?

I made a detailed addition to the commit.

@linuxlonelyeagle
Copy link
Member Author

@ftynse Have you forgotten about this PR?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants