[Mlir-commits] [mlir] 6eb0f8e - [mlir][MemRef] Fix a crash when expanding a scalar shape

[Benjamin Kramer]

Author: Benjamin Kramer
Date: 2022-06-08T09:37:40+02:00
New Revision: 6eb0f8e28598658b4f5df27b35a5ea98bad68049

URL: https://github.com/llvm/llvm-project/commit/6eb0f8e28598658b4f5df27b35a5ea98bad68049
DIFF: https://github.com/llvm/llvm-project/commit/6eb0f8e28598658b4f5df27b35a5ea98bad68049.diff

LOG: [mlir][MemRef] Fix a crash when expanding a scalar shape

In this case the reassociation is empty, yielding no strides for the
result type.

Differential Revision: https://reviews.llvm.org/D127232

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/test/Dialect/Tensor/bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index d200fd0485c5f..95f78965b49cd 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -1734,9 +1734,10 @@ computeExpandedLayoutMap(MemRefType srcType, ArrayRef<int64_t> resultShape,
                                   .asStride();
     }
   }
-  return makeStridedLinearLayoutMap(
-      llvm::to_vector<8>(llvm::reverse(reverseResultStrides)), srcOffset,
-      srcType.getContext());
+  auto resultStrides = llvm::to_vector<8>(llvm::reverse(reverseResultStrides));
+  resultStrides.resize(resultShape.size(), 1);
+  return makeStridedLinearLayoutMap(resultStrides, srcOffset,
+                                    srcType.getContext());
 }
 
 static FailureOr<MemRefType>

diff  --git a/mlir/test/Dialect/Tensor/bufferize.mlir b/mlir/test/Dialect/Tensor/bufferize.mlir
index cd88f2fd4afab..5e1f2d25a6fc0 100644
--- a/mlir/test/Dialect/Tensor/bufferize.mlir
+++ b/mlir/test/Dialect/Tensor/bufferize.mlir
@@ -9,6 +9,8 @@
  // CHECK-DAG: #[[$MAP6:.*]] = affine_map<(d0) -> (d0 * 2)>
  // CHECK-DAG: #[[$MAP7:.*]] = affine_map<(d0, d1, d2)[s0] -> (d0 * 8 + s0 + d1 * 4 + d2)>
  // CHECK-DAG: #[[$MAP8:.*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)>
+ // CHECK-DAG: #[[$MAP9:.*]] = affine_map<()[s0] -> (s0)>
+ // CHECK-DAG: #[[$MAP10:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 
 // CHECK-LABEL:   func @dim(
 // CHECK-SAME:              %[[TENSOR:.*]]: tensor<f32>,
@@ -333,6 +335,20 @@ func.func @tensor.expand_shape_of_slice(
   return %1 : tensor<?x7x2x5xf32>
 }
 
+// CHECK-LABEL: func @tensor.expand_shape_of_scalar_slice(
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
+func.func @tensor.expand_shape_of_scalar_slice(
+    %t1: tensor<?xf32>, %o1: index, %s1: index) -> tensor<1xf32> {
+  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?xf32>
+  // CHECK: %[[subview:.*]] = memref.subview %[[m1]][%{{.*}}] [1] [1] :  memref<?xf32> to memref<f32, #[[$MAP9]]>
+  %0 = tensor.extract_slice %t1[%o1][1][1] : tensor<?xf32> to tensor<f32>
+  // CHECK: %[[expanded:.*]] = memref.expand_shape %[[subview]] [] : memref<f32, #[[$MAP9]]> into memref<1xf32, #[[$MAP10]]>
+  %1 = tensor.expand_shape %0 [] : tensor<f32> into tensor<1xf32>
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[expanded]]
+  // CHECK: return %[[r]]
+  return %1 : tensor<1xf32>
+}
+
 // CHECK-LABEL: func @tensor.collapse_shape(
 //  CHECK-SAME:     %[[t1:.*]]: tensor<2x?x?xf32>
 func.func @tensor.collapse_shape(%t1: tensor<2x?x?xf32>) -> tensor<?x?xf32> {