[Mlir-commits] [mlir] 6eb0f8e - [mlir][MemRef] Fix a crash when expanding a scalar shape
Benjamin Kramer
llvmlistbot at llvm.org
Wed Jun 8 00:42:56 PDT 2022
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> {
More information about the Mlir-commits
mailing list