[Mlir-commits] [mlir] [mlir][memref] Add a `ViewOp::getMixedSizes` (PR #176561)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This PR adds a useful `getMixedSizes` method for memref.view.

---
Full diff: https://github.com/llvm/llvm-project/pull/176561.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+3) 
- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+16) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index bd96bace7994f..6f8f1481725fc 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -2459,6 +2459,9 @@ def MemRef_ViewOp : MemRef_Op<"view", [
     /// The result of a view is always a memref.
     MemRefType getType() { return ::llvm::cast<MemRefType>(getResult().getType()); }
 
+    // Return both static and dynamic sizes as a list of `OpFoldResult`.
+    SmallVector<OpFoldResult> getMixedSizes();
+
     /// Returns the dynamic sizes for this view operation. This is redundant
     /// with `sizes` but needed in template implementations. More specifically:
     /// ```
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index e0f7a8b452a1d..aea2bc58dd5ee 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -3762,6 +3762,22 @@ OpFoldResult ViewOp::fold(FoldAdaptor adaptor) {
   return {};
 }
 
+SmallVector<OpFoldResult> ViewOp::getMixedSizes() {
+  SmallVector<OpFoldResult> result;
+  unsigned ctr = 0;
+  OpBuilder b(getContext());
+  MemRefType resultType = getType();
+  auto resultShape = resultType.getShape();
+  for (int64_t i = 0; i < resultType.getRank(); ++i) {
+    if (resultType.isDynamicDim(i)) {
+      result.push_back(getSizes()[ctr++]);
+    } else {
+      result.push_back(b.getIndexAttr(resultShape[i]));
+    }
+  }
+  return result;
+}
+
 namespace {
 
 struct ViewOpShapeFolder : public OpRewritePattern<ViewOp> {

``````````

</details>


https://github.com/llvm/llvm-project/pull/176561