[mlir][Bufferization] Support cast from ranked to unranked in canonic…

[akuegel]

#150511 changed the canonicalization pattern to not allow casts from ranked to unranked anymore. This patch restores this functionality, while still keeping the fix to preserve memory space and layout.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Adrian Kuegel (akuegel)

Changes

#150511 changed the canonicalization pattern to not allow casts from ranked to unranked anymore. This patch restores this functionality, while still keeping the fix to preserve memory space and layout.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+3-5)
• (modified) mlir/test/Dialect/Bufferization/canonicalize.mlir (+13)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 7eb729f349638..f1f12f4bca70e 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -806,14 +806,12 @@ struct ToBufferOfCast : public OpRewritePattern<ToBufferOp> {
     if (!srcTensorType)
       return failure();
     auto currentOutputMemRefType =
-        dyn_cast<MemRefType>(toBuffer.getResult().getType());
+        dyn_cast<BaseMemRefType>(toBuffer.getResult().getType());
     if (!currentOutputMemRefType)
       return failure();
 
-    auto memrefType = MemRefType::get(srcTensorType.getShape(),
-                                      srcTensorType.getElementType(),
-                                      currentOutputMemRefType.getLayout(),
-                                      currentOutputMemRefType.getMemorySpace());
+    auto memrefType = currentOutputMemRefType.cloneWith(
+        srcTensorType.getShape(), srcTensorType.getElementType());
     Value memref = ToBufferOp::create(rewriter, toBuffer.getLoc(), memrefType,
                                       tensorCastOperand.getOperand(),
                                       toBuffer.getReadOnly());
diff --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index 2acd19453a04d..ae1d1fcfc19dc 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -263,6 +263,19 @@ func.func @tensor_cast_to_buffer(%arg0 : tensor<4x6x16x32xi8>) ->
 // CHECK-SAME: memref<4x6x16x32xi8> to memref<?x?x16x32xi8>
 // CHECK:   return %[[M1]] : memref<?x?x16x32xi8>
 
+// CHECK-LABEL: func @tensor_cast_to_unranked_buffer
+//  CHECK-SAME:   %[[ARG0:.+]]: tensor<4x6x16x32xi8>
+func.func @tensor_cast_to_unranked_buffer(%arg0 : tensor<4x6x16x32xi8>) ->
+  memref<*xi8> {
+  %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor<*xi8>
+  %1 = bufferization.to_buffer %0 read_only : tensor<*xi8> to memref<*xi8>
+  return %1 : memref<*xi8>
+}
+// CHECK:   %[[M:.+]] = bufferization.to_buffer %[[ARG0]] read_only : tensor<4x6x16x32xi8>
+// CHECK:   %[[M1:.+]] = memref.cast %[[M]]
+// CHECK-SAME: memref<4x6x16x32xi8> to memref<*xi8>
+// CHECK:   return %[[M1]] : memref<*xi8>
+
 // -----
 
 // CHECK-LABEL: func @tensor_cast_to_buffer

[llvm-ci]

LLVM Buildbot has detected a new failure on builder amdgpu-offload-ubuntu-22-cmake-build-only running on rocm-docker-ubu-22 while building mlir at step 3 "checkout".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/203/builds/19086

Here is the relevant piece of the build log for the reference

Step 3 (checkout) failure: update (failure)
git version 2.34.1
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': Could not resolve host: github.com
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': Could not resolve host: github.com

[llvm-ci]

LLVM Buildbot has detected a new failure on builder amdgpu-offload-rhel-9-cmake-build-only running on rocm-docker-rhel-9 while building mlir at step 3 "checkout".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/205/builds/17875

Here is the relevant piece of the build log for the reference

Step 3 (checkout) failure: update (failure)
git version 2.43.5
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': Could not resolve host: github.com
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': Could not resolve host: github.com

[llvm-ci]

LLVM Buildbot has detected a new failure on builder amdgpu-offload-rhel-8-cmake-build-only running on rocm-docker-rhel-8 while building mlir at step 3 "checkout".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/204/builds/17898

Here is the relevant piece of the build log for the reference

Step 3 (checkout) failure: update (failure)
git version 2.43.5
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': Could not resolve host: github.com
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': Could not resolve host: github.com