derek/TheRock
1
0 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
1a182e2dac58879f215c0fad0f5f78c76dcafe56
TheRock/compiler/amd-llvm/mlir/test/Dialect/MemRef/canonicalize.mlir
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
Details
CodeQL / Analyze (python) (push) Canceled after 0s
Details
pre-commit / pre-commit (push) Canceled after 0s
Details
Multi-Arch CI / setup (push) Canceled after 0s
Details
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Details
Setup (Multi-Arch) / setup (push) Canceled after 0s
Details
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Details
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
Details
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Details
Multi-Arch CI / CI Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Details
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Details
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

1755 lines
75 KiB
MLIR
Raw Blame History

// RUN: mlir-opt %s -canonicalize="test-convergence" --split-input-file -allow-unregistered-dialect | FileCheck %s
// CHECK-LABEL: collapse_shape_identity_fold
// CHECK-NEXT: return
func.func @collapse_shape_identity_fold(%arg0 : memref<5xi8>) -> memref<5xi8> {
%0 = memref.collapse_shape %arg0 [[0]] : memref<5xi8> into memref<5xi8>
return %0 : memref<5xi8>
}
// -----
// CHECK-LABEL: expand_shape_identity_fold
// CHECK-NEXT: return
func.func @expand_shape_identity_fold(%arg0 : memref<5x4xi8>) -> memref<5x4xi8> {
%0 = memref.expand_shape %arg0 [[0], [1]] output_shape [5, 4] : memref<5x4xi8> into memref<5x4xi8>
return %0 : memref<5x4xi8>
}
// -----
// CHECK-LABEL: collapse_expand_rank0_cancel
// CHECK-NEXT: return
func.func @collapse_expand_rank0_cancel(%arg0 : memref<1x1xi8>) -> memref<1x1xi8> {
%0 = memref.collapse_shape %arg0 [] : memref<1x1xi8> into memref<i8>
%1 = memref.expand_shape %0 [] output_shape [1, 1] : memref<i8> into memref<1x1xi8>
return %1 : memref<1x1xi8>
}
// -----
// CHECK-LABEL: func @subview_of_size_memcast
// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: memref<4x6x16x32xi8>
// CHECK: %[[S:.+]] = memref.subview %[[ARG0]][0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] : memref<4x6x16x32xi8> to memref<16x32xi8, strided{{.*}}>
// CHECK: return %[[S]] : memref<16x32xi8, strided{{.*}}>
func.func @subview_of_size_memcast(%arg : memref<4x6x16x32xi8>) ->
memref<16x32xi8, strided<[32, 1], offset: 512>>{
%0 = memref.cast %arg : memref<4x6x16x32xi8> to memref<?x?x16x32xi8>
%1 = memref.subview %0[0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] :
memref<?x?x16x32xi8> to
memref<16x32xi8, strided<[32, 1], offset: 512>>
return %1 : memref<16x32xi8, strided<[32, 1], offset: 512>>
}
// -----
// CHECK: func @subview_of_strides_memcast
// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: memref<1x1x?xf32, strided{{.*}}>
// CHECK: %[[S:.+]] = memref.subview %[[ARG0]][0, 0, 0] [1, 1, 4]
// CHECK-SAME: to memref<1x4xf32, strided<[35, 1], offset: ?>>
// CHECK: %[[M:.+]] = memref.cast %[[S]]
// CHECK-SAME: to memref<1x4xf32, strided<[?, ?], offset: ?>>
// CHECK: return %[[M]]
func.func @subview_of_strides_memcast(%arg : memref<1x1x?xf32, strided<[35, 7, 1], offset: ?>>) -> memref<1x4xf32, strided<[?, ?], offset: ?>> {
%0 = memref.cast %arg : memref<1x1x?xf32, strided<[35, 7, 1], offset: ?>> to memref<1x1x?xf32, strided<[?, ?, ?], offset: ?>>
%1 = memref.subview %0[0, 0, 0] [1, 1, 4] [1, 1, 1] : memref<1x1x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x4xf32, strided<[?, ?], offset: ?>>
return %1 : memref<1x4xf32, strided<[?, ?], offset: ?>>
}
// -----
// CHECK-LABEL: func @subview_of_static_full_size
// CHECK-SAME: %[[ARG0:.+]]: memref<4x6x16x32xi8>
// CHECK-NOT: memref.subview
// CHECK: return %[[ARG0]] : memref<4x6x16x32xi8>
func.func @subview_of_static_full_size(%arg0 : memref<4x6x16x32xi8>) -> memref<4x6x16x32xi8> {
%0 = memref.subview %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : memref<4x6x16x32xi8> to memref<4x6x16x32xi8>
return %0 : memref<4x6x16x32xi8>
}
// -----
// CHECK-LABEL: func @negative_subview_of_static_full_size
// CHECK-SAME: %[[ARG0:.+]]: memref<16x4xf32, strided<[4, 1], offset: ?>>
// CHECK-SAME: %[[IDX:.+]]: index
// CHECK: %[[S:.+]] = memref.subview %[[ARG0]][%[[IDX]], 0] [16, 4] [1, 1]
// CHECK-SAME: to memref<16x4xf32, strided<[4, 1], offset: ?>>
// CHECK: return %[[S]] : memref<16x4xf32, strided<[4, 1], offset: ?>>
func.func @negative_subview_of_static_full_size(%arg0: memref<16x4xf32, strided<[4, 1], offset: ?>>, %idx: index) -> memref<16x4xf32, strided<[4, 1], offset: ?>> {
%0 = memref.subview %arg0[%idx, 0][16, 4][1, 1] : memref<16x4xf32, strided<[4, 1], offset: ?>> to memref<16x4xf32, strided<[4, 1], offset: ?>>
return %0 : memref<16x4xf32, strided<[4, 1], offset: ?>>
}
// -----
func.func @subview_canonicalize(%arg0 : memref<?x?x?xf32>, %arg1 : index,
%arg2 : index) -> memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
{
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%0 = memref.subview %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : memref<?x?x?xf32> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
return %0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
}
// CHECK-LABEL: func @subview_canonicalize
// CHECK-SAME: %[[ARG0:.+]]: memref<?x?x?xf32>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]
// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]
// CHECK-SAME: : memref<?x?x?xf32> to memref<4x1x?xf32
// CHECK: %[[RESULT:.+]] = memref.cast %[[SUBVIEW]]
// CHECK: return %[[RESULT]]
// -----
func.func @rank_reducing_subview_canonicalize(%arg0 : memref<?x?x?xf32>, %arg1 : index,
%arg2 : index) -> memref<?x?xf32, strided<[?, ?], offset: ?>>
{
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%0 = memref.subview %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : memref<?x?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %0 : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// CHECK-LABEL: func @rank_reducing_subview_canonicalize
// CHECK-SAME: %[[ARG0:.+]]: memref<?x?x?xf32>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]
// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]
// CHECK-SAME: : memref<?x?x?xf32> to memref<4x?xf32
// CHECK: %[[RESULT:.+]] = memref.cast %[[SUBVIEW]]
// CHECK: return %[[RESULT]]
// -----
func.func @multiple_reducing_dims(%arg0 : memref<1x384x384xf32>,
%arg1 : index, %arg2 : index, %arg3 : index) -> memref<?xf32, strided<[1], offset: ?>>
{
%c1 = arith.constant 1 : index
%0 = memref.subview %arg0[0, %arg1, %arg2] [1, %c1, %arg3] [1, 1, 1] : memref<1x384x384xf32> to memref<?x?xf32, strided<[384, 1], offset: ?>>
%1 = memref.subview %0[0, 0] [1, %arg3] [1, 1] : memref<?x?xf32, strided<[384, 1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
return %1 : memref<?xf32, strided<[1], offset: ?>>
}
// CHECK: func @multiple_reducing_dims
// CHECK: %[[REDUCED1:.+]] = memref.subview %{{.+}}[0, %{{.+}}, %{{.+}}] [1, 1, %{{.+}}] [1, 1, 1]
// CHECK-SAME: : memref<1x384x384xf32> to memref<1x?xf32, strided<[384, 1], offset: ?>>
// CHECK: %[[REDUCED2:.+]] = memref.subview %[[REDUCED1]][0, 0] [1, %{{.+}}] [1, 1]
// CHECK-SAME: : memref<1x?xf32, strided<[384, 1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
// -----
func.func @multiple_reducing_dims_dynamic(%arg0 : memref<?x?x?xf32>,
%arg1 : index, %arg2 : index, %arg3 : index) -> memref<?xf32, strided<[1], offset: ?>>
{
%c1 = arith.constant 1 : index
%0 = memref.subview %arg0[0, %arg1, %arg2] [1, %c1, %arg3] [1, 1, 1] : memref<?x?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
%1 = memref.subview %0[0, 0] [1, %arg3] [1, 1] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
return %1 : memref<?xf32, strided<[1], offset: ?>>
}
// CHECK: func @multiple_reducing_dims_dynamic
// CHECK: %[[REDUCED1:.+]] = memref.subview %{{.+}}[0, %{{.+}}, %{{.+}}] [1, 1, %{{.+}}] [1, 1, 1]
// CHECK-SAME: : memref<?x?x?xf32> to memref<1x?xf32, strided<[?, 1], offset: ?>>
// CHECK: %[[REDUCED2:.+]] = memref.subview %[[REDUCED1]][0, 0] [1, %{{.+}}] [1, 1]
// CHECK-SAME: : memref<1x?xf32, strided<[?, 1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
// -----
func.func @multiple_reducing_dims_all_dynamic(%arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>,
%arg1 : index, %arg2 : index, %arg3 : index) -> memref<?xf32, strided<[?], offset: ?>>
{
%c1 = arith.constant 1 : index
%0 = memref.subview %arg0[0, %arg1, %arg2] [1, %c1, %arg3] [1, 1, 1]
: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%1 = memref.subview %0[0, 0] [1, %arg3] [1, 1] : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
return %1 : memref<?xf32, strided<[?], offset: ?>>
}
// CHECK: func @multiple_reducing_dims_all_dynamic
// CHECK: %[[REDUCED1:.+]] = memref.subview %{{.+}}[0, %{{.+}}, %{{.+}}] [1, 1, %{{.+}}] [1, 1, 1]
// CHECK-SAME: : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x?xf32, strided<[?, ?], offset: ?>>
// CHECK: %[[REDUCED2:.+]] = memref.subview %[[REDUCED1]][0, 0] [1, %{{.+}}] [1, 1]
// CHECK-SAME: : memref<1x?xf32, strided<[?, ?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
// -----
func.func @subview_negative_stride1(%arg0 : memref<?xf32>) -> memref<?xf32, strided<[?], offset: ?>>
{
%c0 = arith.constant 0 : index
%c1 = arith.constant -1 : index
%1 = memref.dim %arg0, %c0 : memref<?xf32>
%2 = arith.addi %1, %c1 : index
%3 = memref.subview %arg0[%2] [%1] [%c1] : memref<?xf32> to memref<?xf32, strided<[?], offset: ?>>
return %3 : memref<?xf32, strided<[?], offset: ?>>
}
// CHECK: func @subview_negative_stride1
// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>)
// CHECK: %[[C1:.*]] = arith.constant 0
// CHECK: %[[C2:.*]] = arith.constant -1
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG0]], %[[C1]] : memref<?xf32>
// CHECK: %[[DIM2:.*]] = arith.addi %[[DIM1]], %[[C2]] : index
// CHECK: %[[RES1:.*]] = memref.subview %[[ARG0]][%[[DIM2]]] [%[[DIM1]]] [-1] : memref<?xf32> to memref<?xf32, strided<[-1], offset: ?>>
// CHECK: %[[RES2:.*]] = memref.cast %[[RES1]] : memref<?xf32, strided<[-1], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
// CHECK: return %[[RES2]] : memref<?xf32, strided<[?], offset: ?>>
// -----
func.func @subview_negative_stride2(%arg0 : memref<7xf32>) -> memref<?xf32, strided<[?], offset: ?>>
{
%c0 = arith.constant 0 : index
%c1 = arith.constant -1 : index
%1 = memref.dim %arg0, %c0 : memref<7xf32>
%2 = arith.addi %1, %c1 : index
%3 = memref.subview %arg0[%2] [%1] [%c1] : memref<7xf32> to memref<?xf32, strided<[?], offset: ?>>
return %3 : memref<?xf32, strided<[?], offset: ?>>
}
// CHECK: func @subview_negative_stride2
// CHECK-SAME: (%[[ARG0:.*]]: memref<7xf32>)
// CHECK: %[[RES1:.*]] = memref.subview %[[ARG0]][6] [7] [-1] : memref<7xf32> to memref<7xf32, strided<[-1], offset: 6>>
// CHECK: %[[RES2:.*]] = memref.cast %[[RES1]] : memref<7xf32, strided<[-1], offset: 6>> to memref<?xf32, strided<[?], offset: ?>>
// CHECK: return %[[RES2]] : memref<?xf32, strided<[?], offset: ?>>
// -----
// CHECK-LABEL: func @no_fold_subview_negative_size
// CHECK: %[[SUBVIEW:.+]] = memref.subview
// CHECK: return %[[SUBVIEW]]
func.func @no_fold_subview_negative_size(%input: memref<4x1024xf32>) -> memref<?x256xf32, strided<[1024, 1], offset: 2304>> {
%cst = arith.constant -13 : index
%0 = memref.subview %input[2, 256] [%cst, 256] [1, 1] : memref<4x1024xf32> to memref<?x256xf32, strided<[1024, 1], offset: 2304>>
return %0 : memref<?x256xf32, strided<[1024, 1], offset: 2304>>
}
// -----
// CHECK-LABEL: func @no_fold_subview_zero_stride
// CHECK: %[[SUBVIEW:.+]] = memref.subview
// CHECK: return %[[SUBVIEW]]
func.func @no_fold_subview_zero_stride(%arg0 : memref<10xf32>) -> memref<1xf32, strided<[?], offset: 1>> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%1 = memref.subview %arg0[1] [1] [%c0] : memref<10xf32> to memref<1xf32, strided<[?], offset: 1>>
return %1 : memref<1xf32, strided<[?], offset: 1>>
}
// -----
// CHECK-LABEL: func @no_fold_of_store
// CHECK: %[[cst:.+]] = memref.cast %arg
// CHECK: memref.store %[[cst]]
func.func @no_fold_of_store(%arg : memref<32xi8>, %holder: memref<memref<?xi8>>) {
%0 = memref.cast %arg : memref<32xi8> to memref<?xi8>
memref.store %0, %holder[] : memref<memref<?xi8>>
return
}
// -----
// Test case: Folding of memref.dim(memref.alloca(%size), %idx) -> %size
// CHECK-LABEL: func @dim_of_alloca(
// CHECK-SAME: %[[SIZE:[0-9a-z]+]]: index
// CHECK-NEXT: return %[[SIZE]] : index
func.func @dim_of_alloca(%size: index) -> index {
%0 = memref.alloca(%size) : memref<?xindex>
%c0 = arith.constant 0 : index
%1 = memref.dim %0, %c0 : memref<?xindex>
return %1 : index
}
// -----
// Test case: Folding of memref.dim(memref.alloca(rank(%v)), %idx) -> rank(%v)
// CHECK-LABEL: func @dim_of_alloca_with_dynamic_size(
// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>
// CHECK-NEXT: %[[RANK:.*]] = memref.rank %[[MEM]] : memref<*xf32>
// CHECK-NEXT: return %[[RANK]] : index
func.func @dim_of_alloca_with_dynamic_size(%arg0: memref<*xf32>) -> index {
%0 = memref.rank %arg0 : memref<*xf32>
%1 = memref.alloca(%0) : memref<?xindex>
%c0 = arith.constant 0 : index
%2 = memref.dim %1, %c0 : memref<?xindex>
return %2 : index
}
// -----
// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx]
// CHECK-LABEL: func @dim_of_memref_reshape(
// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>,
// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref<?xindex>
// CHECK-NEXT: %[[IDX:.*]] = arith.constant 3
// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]]
// CHECK-NEXT: memref.store
// CHECK-NOT: memref.dim
// CHECK: return %[[DIM]] : index
func.func @dim_of_memref_reshape(%arg0: memref<*xf32>, %arg1: memref<?xindex>)
-> index {
%c3 = arith.constant 3 : index
%0 = memref.reshape %arg0(%arg1)
: (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
// Update the shape to test that he load ends up in the right place.
memref.store %c3, %arg1[%c3] : memref<?xindex>
%1 = memref.dim %0, %c3 : memref<*xf32>
return %1 : index
}
// -----
// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx]
// CHECK-LABEL: func @dim_of_memref_reshape_i32(
// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>,
// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref<?xi32>
// CHECK-NEXT: %[[IDX:.*]] = arith.constant 3
// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]]
// CHECK-NEXT: %[[CAST:.*]] = arith.index_cast %[[DIM]]
// CHECK-NOT: memref.dim
// CHECK: return %[[CAST]] : index
func.func @dim_of_memref_reshape_i32(%arg0: memref<*xf32>, %arg1: memref<?xi32>)
-> index {
%c3 = arith.constant 3 : index
%0 = memref.reshape %arg0(%arg1)
: (memref<*xf32>, memref<?xi32>) -> memref<*xf32>
%1 = memref.dim %0, %c3 : memref<*xf32>
return %1 : index
}
// -----
// Test case: memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx]
// CHECK-LABEL: func @dim_of_memref_reshape_block_arg_index(
// CHECK-SAME: %[[MEM:[0-9a-z]+]]: memref<*xf32>,
// CHECK-SAME: %[[SHP:[0-9a-z]+]]: memref<?xindex>,
// CHECK-SAME: %[[IDX:[0-9a-z]+]]: index
// CHECK-NEXT: %[[DIM:.*]] = memref.load %[[SHP]][%[[IDX]]]
// CHECK-NOT: memref.dim
// CHECK: return %[[DIM]] : index
func.func @dim_of_memref_reshape_block_arg_index(%arg0: memref<*xf32>, %arg1: memref<?xindex>, %arg2: index) -> index {
%reshape = memref.reshape %arg0(%arg1) : (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
%dim = memref.dim %reshape, %arg2 : memref<*xf32>
return %dim : index
}
// -----
// Test case: memref.dim(memref.reshape %v %shp, %idx) is not folded into memref.load %shp[%idx]
// CHECK-LABEL: func @dim_of_memref_reshape_for(
// CHECK: memref.reshape
// CHECK: memref.dim
// CHECK-NOT: memref.load
func.func @dim_of_memref_reshape_for( %arg0: memref<*xf32>, %arg1: memref<?xindex>) -> index {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%0 = memref.reshape %arg0(%arg1) : (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
%1 = scf.for %arg2 = %c0 to %c4 step %c1 iter_args(%arg3 = %c1) -> (index) {
%2 = memref.dim %0, %arg2 : memref<*xf32>
%3 = arith.muli %arg3, %2 : index
scf.yield %3 : index
}
return %1 : index
}
// -----
// Test case: memref.dim(memref.reshape %v %shp, %idx) is not folded into memref.load %shp[%idx]
// CHECK-LABEL: func @dim_of_memref_reshape_undominated(
// CHECK: memref.reshape
// CHECK: memref.dim
// CHECK-NOT: memref.load
func.func @dim_of_memref_reshape_undominated(%arg0: memref<*xf32>, %arg1: memref<?xindex>, %arg2: index) -> index {
%c4 = arith.constant 4 : index
%reshape = memref.reshape %arg0(%arg1) : (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
%0 = arith.muli %arg2, %c4 : index
%dim = memref.dim %reshape, %0 : memref<*xf32>
return %dim : index
}
// -----
// CHECK-LABEL: func @alloc_const_fold
func.func @alloc_const_fold() -> memref<?xf32> {
// CHECK-NEXT: memref.alloc() : memref<4xf32>
%c4 = arith.constant 4 : index
%a = memref.alloc(%c4) : memref<?xf32>
// CHECK-NEXT: memref.cast %{{.*}} : memref<4xf32> to memref<?xf32>
// CHECK-NEXT: return %{{.*}} : memref<?xf32>
return %a : memref<?xf32>
}
// -----
// CHECK-LABEL: func @alloc_alignment_const_fold
func.func @alloc_alignment_const_fold() -> memref<?xf32> {
// CHECK-NEXT: memref.alloc() {alignment = 4096 : i64} : memref<4xf32>
%c4 = arith.constant 4 : index
%a = memref.alloc(%c4) {alignment = 4096 : i64} : memref<?xf32>
// CHECK-NEXT: memref.cast %{{.*}} : memref<4xf32> to memref<?xf32>
// CHECK-NEXT: return %{{.*}} : memref<?xf32>
return %a : memref<?xf32>
}
// -----
// CHECK-LABEL: func @alloc_const_fold_with_symbols1(
// CHECK: %[[c1:.+]] = arith.constant 1 : index
// CHECK: %[[mem1:.+]] = memref.alloc({{.*}})[%[[c1]], %[[c1]]] : memref<?xi32, strided{{.*}}>
// CHECK: return %[[mem1]] : memref<?xi32, strided{{.*}}>
func.func @alloc_const_fold_with_symbols1(%arg0 : index) -> memref<?xi32, strided<[?], offset: ?>> {
%c1 = arith.constant 1 : index
%0 = memref.alloc(%arg0)[%c1, %c1] : memref<?xi32, strided<[?], offset: ?>>
return %0 : memref<?xi32, strided<[?], offset: ?>>
}
// -----
// CHECK-LABEL: func @alloc_const_fold_with_symbols2(
// CHECK: %[[c1:.+]] = arith.constant 1 : index
// CHECK: %[[mem1:.+]] = memref.alloc()[%[[c1]], %[[c1]]] : memref<1xi32, strided{{.*}}>
// CHECK: %[[mem2:.+]] = memref.cast %[[mem1]] : memref<1xi32, strided{{.*}}> to memref<?xi32, strided{{.*}}>
// CHECK: return %[[mem2]] : memref<?xi32, strided{{.*}}>
func.func @alloc_const_fold_with_symbols2() -> memref<?xi32, strided<[?], offset: ?>> {
%c1 = arith.constant 1 : index
%0 = memref.alloc(%c1)[%c1, %c1] : memref<?xi32, strided<[?], offset: ?>>
return %0 : memref<?xi32, strided<[?], offset: ?>>
}
// -----
// CHECK-LABEL: func @allocator
// CHECK: %[[alloc:.+]] = memref.alloc
// CHECK: memref.store %[[alloc:.+]], %arg0
func.func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index) {
%0 = memref.alloc(%arg1) : memref<?xi32>
memref.store %0, %arg0[] : memref<memref<?xi32>>
return
}
// -----
func.func @compose_collapse_of_collapse_zero_dim(%arg0 : memref<1x1x1xf32>)
-> memref<f32> {
%0 = memref.collapse_shape %arg0 [[0, 1, 2]]
: memref<1x1x1xf32> into memref<1xf32>
%1 = memref.collapse_shape %0 [] : memref<1xf32> into memref<f32>
return %1 : memref<f32>
}
// CHECK-LABEL: func @compose_collapse_of_collapse_zero_dim
// CHECK: memref.collapse_shape %{{.*}} []
// CHECK-SAME: memref<1x1x1xf32> into memref<f32>
// -----
func.func @compose_collapse_of_collapse(%arg0 : memref<?x?x?x?x?xf32>)
-> memref<?x?xf32> {
%0 = memref.collapse_shape %arg0 [[0, 1], [2], [3, 4]]
: memref<?x?x?x?x?xf32> into memref<?x?x?xf32>
%1 = memref.collapse_shape %0 [[0, 1], [2]]
: memref<?x?x?xf32> into memref<?x?xf32>
return %1 : memref<?x?xf32>
}
// CHECK-LABEL: func @compose_collapse_of_collapse
// CHECK: memref.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
// CHECK-NOT: memref.collapse_shape
// -----
func.func @compose_collapse_of_expand_partially_dynamic(%arg0: memref<?xf16>, %arg1: index, %arg2: index) -> memref<8x?x?xf16> {
%expanded = memref.expand_shape %arg0 [[0, 1, 2, 3, 4]] output_shape [4, 2, %arg1, %arg2, 32] : memref<?xf16> into memref<4x2x?x?x32xf16>
%collapsed = memref.collapse_shape %expanded [[0, 1], [2], [3, 4]] : memref<4x2x?x?x32xf16> into memref<8x?x?xf16>
return %collapsed : memref<8x?x?xf16>
}
// CHECK: func @compose_collapse_of_expand_partially_dynamic
// CHECK-SAME: %[[SRC:.[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ORIG_D2:.[a-zA-Z0-9]+]]
// CHECK-SAME: %[[ORIG_D3:.[a-zA-Z0-9]+]]
// CHECK-DAG: %[[C32:.+]] = arith.constant 32
// CHECK: %[[COLLAPSED_D2:.+]] = arith.muli %[[ORIG_D3]], %[[C32]]
// CHECK: %[[RESULT:.+]] = memref.expand_shape %[[SRC]]
// CHECK-SAME: [0, 1, 2]
// CHECK-SAME: output_shape [8, %[[ORIG_D2]], %[[COLLAPSED_D2]]]
// CHECK: return %[[RESULT]]
// -----
func.func @do_not_compose_collapse_of_expand_non_identity_layout(
%arg0: memref<?x?xf32, strided<[?, 1], offset: 0>>, %sz0: index, %sz1: index)
-> memref<?xf32, strided<[?], offset: 0>> {
%1 = memref.expand_shape %arg0 [[0, 1], [2]] output_shape [%sz0, 4, %sz1] :
memref<?x?xf32, strided<[?, 1], offset: 0>> into
memref<?x4x?xf32, strided<[?, ?, 1], offset: 0>>
%2 = memref.collapse_shape %1 [[0, 1, 2]] :
memref<?x4x?xf32, strided<[?, ?, 1], offset: 0>> into
memref<?xf32, strided<[?], offset: 0>>
return %2 : memref<?xf32, strided<[?], offset: 0>>
}
// CHECK-LABEL: func @do_not_compose_collapse_of_expand_non_identity_layout
// CHECK: expand
// CHECK: collapse
// -----
func.func @compose_expand_of_expand(%arg0 : memref<?x?xf32>, %sz0: index, %sz1: index, %sz2: index, %sz3: index)
-> memref<?x6x4x5x?xf32> {
%0 = memref.expand_shape %arg0 [[0, 1], [2]] output_shape [%sz0, 4, %sz1]
: memref<?x?xf32> into memref<?x4x?xf32>
%1 = memref.expand_shape %0 [[0, 1], [2], [3, 4]] output_shape [%sz2, 6, 4, 5, %sz3] : memref<?x4x?xf32> into memref<?x6x4x5x?xf32>
return %1 : memref<?x6x4x5x?xf32>
}
// CHECK-LABEL: func @compose_expand_of_expand
// CHECK: memref.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]] output_shape [%{{.*}}, 6, 4, 5, %{{.*}}]
// CHECK-NOT: memref.expand_shape
// -----
func.func @compose_expand_of_expand_of_zero_dim(%arg0 : memref<f32>)
-> memref<1x1x1xf32> {
%0 = memref.expand_shape %arg0 [] output_shape [1] : memref<f32> into memref<1xf32>
%1 = memref.expand_shape %0 [[0, 1, 2]] output_shape [1, 1, 1]
: memref<1xf32> into memref<1x1x1xf32>
return %1 : memref<1x1x1xf32>
}
// CHECK-LABEL: func @compose_expand_of_expand_of_zero_dim
// CHECK: memref.expand_shape %{{.*}} [] output_shape [1, 1, 1]
// CHECK-SAME: memref<f32> into memref<1x1x1xf32>
// -----
func.func @fold_collapse_of_expand(%arg0 : memref<12x4xf32>) -> memref<12x4xf32> {
%0 = memref.expand_shape %arg0 [[0, 1], [2]] output_shape [3, 4, 4]
: memref<12x4xf32> into memref<3x4x4xf32>
%1 = memref.collapse_shape %0 [[0, 1], [2]]
: memref<3x4x4xf32> into memref<12x4xf32>
return %1 : memref<12x4xf32>
}
// CHECK-LABEL: func @fold_collapse_of_expand
// CHECK-NOT: memref.{{.*}}_shape
// -----
func.func @fold_collapse_collapse_of_expand(%arg0 : memref<?x?xf32>, %sz0: index, %sz1: index)
-> memref<?x?xf32> {
%0 = memref.expand_shape %arg0 [[0, 1], [2]] output_shape [%sz0, 4, %sz1]
: memref<?x?xf32> into memref<?x4x?xf32>
%1 = memref.collapse_shape %0 [[0, 1], [2]]
: memref<?x4x?xf32> into memref<?x?xf32>
return %1 : memref<?x?xf32>
}
// CHECK-LABEL: @fold_collapse_collapse_of_expand
// CHECK-NOT: memref.{{.*}}_shape
// -----
func.func @fold_memref_expand_cast(%arg0 : memref<?x?xf32>) -> memref<2x4x4xf32> {
%0 = memref.cast %arg0 : memref<?x?xf32> to memref<8x4xf32>
%1 = memref.expand_shape %0 [[0, 1], [2]] output_shape [2, 4, 4]
: memref<8x4xf32> into memref<2x4x4xf32>
return %1 : memref<2x4x4xf32>
}
// CHECK-LABEL: @fold_memref_expand_cast
// CHECK: memref.expand_shape
// -----
// CHECK-LABEL: func.func @fold_memref_expand_with_static_to_dynamic_cast(
// CHECK-SAME: %[[ARG0:.*]]: memref<8x4xf32>) -> memref<2x1x4x4xf32> {
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0, 1, 2], [3]]
// CHECK-SAME: output_shape [2, 1, 4, 4] : memref<8x4xf32> into memref<2x1x4x4xf32>
// CHECK: return %[[EXPAND_SHAPE_0]] : memref<2x1x4x4xf32>
// CHECK: }
func.func @fold_memref_expand_with_static_to_dynamic_cast(%arg0: memref<8x4xf32>) -> memref<2x1x4x4xf32> {
%c2 = arith.constant 2 : index
%cast = memref.cast %arg0 : memref<8x4xf32> to memref<?x4xf32>
%expand_shape = memref.expand_shape %cast [[0, 1, 2], [3]] output_shape [%c2, 1, 4, 4]
: memref<?x4xf32> into memref<?x1x4x4xf32>
%cast_0 = memref.cast %expand_shape : memref<?x1x4x4xf32> to memref<2x1x4x4xf32>
return %cast_0 : memref<2x1x4x4xf32>
}
// -----
// CHECK-LABEL: func.func @fold_memref_expand_static_to_dynamic_partial(
// CHECK-SAME: %[[ARG0:.*]]: memref<8x?xf32>) -> memref<1x8x1x?xf32> {
// CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG0]], %[[C1]] : memref<8x?xf32>
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2, 3]]
// CHECK-SAME: output_shape [1, 8, 1, %[[DIM1]]] : memref<8x?xf32> into memref<1x8x1x?xf32>
// CHECK: return %[[EXPAND_SHAPE_0]] : memref<1x8x1x?xf32>
// CHECK: }
func.func @fold_memref_expand_static_to_dynamic_partial(%arg0 : memref<8x?xf32>) -> memref<1x8x1x?xf32> {
%0 = memref.cast %arg0 : memref<8x?xf32> to memref<?x?xf32>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim0 = memref.dim %0, %c0 : memref<?x?xf32>
%dim1 = memref.dim %0, %c1 : memref<?x?xf32>
%1 = memref.expand_shape %0 [[0, 1], [2, 3]] output_shape [1, %dim0, 1, %dim1]
: memref<?x?xf32> into memref<1x?x1x?xf32>
%2 = memref.cast %1 : memref<1x?x1x?xf32> to memref<1x8x1x?xf32>
return %2 : memref<1x8x1x?xf32>
}
// -----
// CHECK-LABEL: func.func @fold_memref_expand_static_to_dynamic_partial_with_arith_const_as_dim(
// CHECK-SAME: %[[ARG0:.*]]: memref<8x?xf32>) -> memref<1x8x1x?xf32> {
// CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG0]], %[[C1]] : memref<8x?xf32>
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2, 3]]
// CHECK-SAME: output_shape [1, 8, 1, %[[DIM1]]] : memref<8x?xf32> into memref<1x8x1x?xf32>
// CHECK: return %[[EXPAND_SHAPE_0]] : memref<1x8x1x?xf32>
// CHECK: }
func.func @fold_memref_expand_static_to_dynamic_partial_with_arith_const_as_dim(%arg0 : memref<8x?xf32>) -> memref<1x8x1x?xf32> {
%0 = memref.cast %arg0 : memref<8x?xf32> to memref<?x?xf32>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim0 = memref.dim %0, %c0 : memref<?x?xf32>
%dim1 = memref.dim %0, %c1 : memref<?x?xf32>
%1 = memref.expand_shape %0 [[0, 1], [2, 3]] output_shape [%c1, %dim0, %c1, %dim1]
: memref<?x?xf32> into memref<?x?x?x?xf32>
%2 = memref.cast %1 : memref<?x?x?x?xf32> to memref<1x8x1x?xf32>
return %2 : memref<1x8x1x?xf32>
}
// -----
// CHECK-LABEL: func.func @fold_memref_expand_static_to_dynamic_multiple(
// CHECK-SAME: %[[ARG0:.*]]: memref<8x?xf32>,
// CHECK-SAME: %[[ARG1:.*]]: index, %[[ARG2:.*]]: index) -> memref<8x1x?x?xf32> {
// CHECK-NOT: memref.cast
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2, 3]]
// CHECK-SAME: output_shape [8, 1, %[[ARG1]], %[[ARG2]]] : memref<8x?xf32> into memref<8x1x?x?xf32>
// CHECK-NOT: memref.cast
// CHECK: return %[[EXPAND_SHAPE_0]] : memref<8x1x?x?xf32>
// CHECK: }
func.func @fold_memref_expand_static_to_dynamic_multiple(%arg0 : memref<8x?xf32>, %arg1 : index, %arg2 : index) -> memref<8x1x?x?xf32> {
%0 = memref.cast %arg0 : memref<8x?xf32> to memref<?x?xf32>
%c0 = arith.constant 0 : index
%dim0 = memref.dim %0, %c0 : memref<?x?xf32>
%1 = memref.expand_shape %0 [[0, 1], [2, 3]] output_shape [%dim0, 1, %arg1, %arg2]
: memref<?x?xf32> into memref<?x1x?x?xf32>
%2 = memref.cast %1 : memref<?x1x?x?xf32> to memref<8x1x?x?xf32>
return %2 : memref<8x1x?x?xf32>
}
// -----
// CHECK-LABEL: func.func @not_fold_memref_expand_with_dynamic_to_static_cast(
// CHECK-SAME: %[[ARG0:.*]]: memref<?x4xf32>) -> memref<2x1x4x4xf32> {
// CHECK: %[[CAST_0:.*]] = memref.cast %[[ARG0]] : memref<?x4xf32> to memref<8x4xf32>
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[CAST_0]] {{\[\[}}0, 1, 2], [3]]
// CHECK-SAME: output_shape [2, 1, 4, 4] : memref<8x4xf32> into memref<2x1x4x4xf32>
// CHECK: return %[[EXPAND_SHAPE_0]] : memref<2x1x4x4xf32>
// CHECK: }
func.func @not_fold_memref_expand_with_dynamic_to_static_cast(%arg0 : memref<?x4xf32>) -> memref<2x1x4x4xf32> {
%0 = memref.cast %arg0 : memref<?x4xf32> to memref<8x4xf32>
%1 = memref.expand_shape %0 [[0, 1, 2], [3]] output_shape [2, 1, 4, 4]
: memref<8x4xf32> into memref<2x1x4x4xf32>
return %1 : memref<2x1x4x4xf32>
}
// -----
// CHECK-LABEL: func.func @not_fold_memref_expand_static_to_dynamic_cast_if_really_dynamic(
// CHECK-SAME: %[[ARG0:.*]]: memref<8x4xf32>,
// CHECK-SAME: %[[ARG1:.*]]: index) -> memref<2x1x4x4xf32> {
// CHECK: %[[C8:.*]] = arith.constant 8 : index
// CHECK: %[[CAST_0:.*]] = memref.cast %[[ARG0]] : memref<8x4xf32> to memref<?x4xf32>
// CHECK: %[[DIVUI_0:.*]] = arith.divui %[[C8]], %[[ARG1]] : index
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[CAST_0]] {{\[\[}}0, 1, 2], [3]]
// CHECK-SAME: output_shape {{\[}}%[[DIVUI_0]], 1, 4, 4] : memref<?x4xf32> into memref<?x1x4x4xf32>
// CHECK: %[[CAST_1:.*]] = memref.cast %[[EXPAND_SHAPE_0]] : memref<?x1x4x4xf32> to memref<2x1x4x4xf32>
// CHECK: return %[[CAST_1]] : memref<2x1x4x4xf32>
// CHECK: }
func.func @not_fold_memref_expand_static_to_dynamic_cast_if_really_dynamic(%arg0 : memref<8x4xf32>, %arg1 : index) -> memref<2x1x4x4xf32> {
%0 = memref.cast %arg0 : memref<8x4xf32> to memref<?x4xf32>
%c0 = arith.constant 0 : index
%dim0 = memref.dim %0, %c0 : memref<?x4xf32>
%dim_ext = arith.divui %dim0 , %arg1: index
%1 = memref.expand_shape %0 [[0, 1, 2], [3]] output_shape [%dim_ext, 1, 4, 4]
: memref<?x4xf32> into memref<?x1x4x4xf32>
%2 = memref.cast %1 : memref<?x1x4x4xf32> to memref<2x1x4x4xf32>
return %2 : memref<2x1x4x4xf32>
}
// -----
// CHECK-LABEL: func.func @fold_memref_expand_static_to_dynamic_layout(
// CHECK-SAME: %[[ARG0:.*]]: memref<8x4xf32>) -> memref<8x1x4xf32> {
// CHECK: %[[EXPAND_SHAPE_0:.*]] = memref.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2]]
// CHECK-SAME: output_shape [8, 1, 4] : memref<8x4xf32> into memref<8x1x4xf32>
// CHECK: return %[[EXPAND_SHAPE_0]] : memref<8x1x4xf32>
// CHECK: }
func.func @fold_memref_expand_static_to_dynamic_layout(%arg0 : memref<8x4xf32>) -> memref<8x1x4xf32> {
%0 = memref.cast %arg0 : memref<8x4xf32> to memref<8x4xf32, strided<[?, ?], offset: ?>>
%1 = memref.expand_shape %0 [[0, 1], [2]] output_shape [8, 1, 4]
: memref<8x4xf32, strided<[?, ?], offset: ?>> into memref<8x1x4xf32, strided<[?,?,?], offset: ?>>
%2 = memref.cast %1 : memref<8x1x4xf32, strided<[?,?,?], offset: ?>> to memref<8x1x4xf32>
return %2 : memref<8x1x4xf32>
}
// -----
// CHECK-LABEL: func @collapse_after_memref_cast_type_change(
// CHECK-SAME: %[[INPUT:.*]]: memref<?x512x1x1xf32>) -> memref<?x?xf32> {
// CHECK: %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
// CHECK-SAME: {{\[\[}}0], [1, 2, 3]] : memref<?x512x1x1xf32> into memref<?x512xf32>
// CHECK: %[[DYNAMIC:.*]] = memref.cast %[[COLLAPSED]] :
// CHECK-SAME: memref<?x512xf32> to memref<?x?xf32>
// CHECK: return %[[DYNAMIC]] : memref<?x?xf32>
// CHECK: }
func.func @collapse_after_memref_cast_type_change(%arg0 : memref<?x512x1x1xf32>) -> memref<?x?xf32> {
%dynamic = memref.cast %arg0: memref<?x512x1x1xf32> to memref<?x?x?x?xf32>
%collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
return %collapsed : memref<?x?xf32>
}
// -----
// CHECK-LABEL: func @collapse_after_memref_cast(
// CHECK-SAME: %[[INPUT:.*]]: memref<?x512x1x?xf32>) -> memref<?x?xf32> {
// CHECK: %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
// CHECK-SAME: {{\[\[}}0], [1, 2, 3]] : memref<?x512x1x?xf32> into memref<?x?xf32>
// CHECK: return %[[COLLAPSED]] : memref<?x?xf32>
func.func @collapse_after_memref_cast(%arg0 : memref<?x512x1x?xf32>) -> memref<?x?xf32> {
%dynamic = memref.cast %arg0: memref<?x512x1x?xf32> to memref<?x?x?x?xf32>
%collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
return %collapsed : memref<?x?xf32>
}
// -----
// CHECK-LABEL: func @collapse_after_memref_cast_type_change_dynamic(
// CHECK-SAME: %[[INPUT:.*]]: memref<1x1x1x?xi64>) -> memref<?x?xi64> {
// CHECK: %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
// CHECK-SAME: {{\[\[}}0, 1, 2], [3]] : memref<1x1x1x?xi64> into memref<1x?xi64>
// CHECK: %[[DYNAMIC:.*]] = memref.cast %[[COLLAPSED]] :
// CHECK-SAME: memref<1x?xi64> to memref<?x?xi64>
// CHECK: return %[[DYNAMIC]] : memref<?x?xi64>
func.func @collapse_after_memref_cast_type_change_dynamic(%arg0: memref<1x1x1x?xi64>) -> memref<?x?xi64> {
%casted = memref.cast %arg0 : memref<1x1x1x?xi64> to memref<1x1x?x?xi64>
%collapsed = memref.collapse_shape %casted [[0, 1, 2], [3]] : memref<1x1x?x?xi64> into memref<?x?xi64>
return %collapsed : memref<?x?xi64>
}
// -----
func.func @reduced_memref(%arg0: memref<2x5x7x1xf32>, %arg1 :index)
-> memref<1x4x1xf32, strided<[35, 7, 1], offset: ?>> {
%c0 = arith.constant 0 : index
%c5 = arith.constant 5 : index
%c4 = arith.constant 4 : index
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%0 = memref.subview %arg0[%arg1, %arg1, %arg1, 0] [%c1, %c4, %c1, 1] [1, 1, 1, 1]
: memref<2x5x7x1xf32> to memref<?x?x?xf32, strided<[35, 7, 1], offset: ?>>
%1 = memref.cast %0
: memref<?x?x?xf32, strided<[35, 7, 1], offset: ?>> to
memref<1x4x1xf32, strided<[35, 7, 1], offset: ?>>
return %1 : memref<1x4x1xf32, strided<[35, 7, 1], offset: ?>>
}
// CHECK-LABEL: func @reduced_memref
// CHECK: %[[RESULT:.+]] = memref.subview
// CHECK-SAME: memref<2x5x7x1xf32> to memref<1x4x1xf32, strided{{.+}}>
// CHECK: return %[[RESULT]]
// -----
// CHECK-LABEL: func @fold_rank_memref
func.func @fold_rank_memref(%arg0 : memref<?x?xf32>) -> (index) {
// Fold a rank into a constant
// CHECK-NEXT: [[C2:%.+]] = arith.constant 2 : index
%rank_0 = memref.rank %arg0 : memref<?x?xf32>
// CHECK-NEXT: return [[C2]]
return %rank_0 : index
}
// -----
func.func @fold_no_op_subview(%arg0 : memref<20x42xf32>) -> memref<20x42xf32, strided<[42, 1]>> {
%0 = memref.subview %arg0[0, 0] [20, 42] [1, 1] : memref<20x42xf32> to memref<20x42xf32, strided<[42, 1]>>
return %0 : memref<20x42xf32, strided<[42, 1]>>
}
// CHECK-LABEL: func @fold_no_op_subview(
// CHECK: %[[ARG0:.+]]: memref<20x42xf32>)
// CHECK: %[[CAST:.+]] = memref.cast %[[ARG0]]
// CHECK: return %[[CAST]]
// -----
func.func @no_fold_subview_with_non_zero_offset(%arg0 : memref<20x42xf32>) -> memref<20x41xf32, strided<[42, 1], offset: 1>> {
%0 = memref.subview %arg0[0, 1] [20, 41] [1, 1] : memref<20x42xf32> to memref<20x41xf32, strided<[42, 1], offset: 1>>
return %0 : memref<20x41xf32, strided<[42, 1], offset: 1>>
}
// CHECK-LABEL: func @no_fold_subview_with_non_zero_offset(
// CHECK: %[[SUBVIEW:.+]] = memref.subview
// CHECK: return %[[SUBVIEW]]
// -----
func.func @no_fold_subview_with_non_unit_stride(%arg0 : memref<20x42xf32>) -> memref<20x5xf32, strided<[42, 2]>> {
%0 = memref.subview %arg0[0, 0] [20, 5] [1, 2] : memref<20x42xf32> to memref<20x5xf32, strided<[42, 2]>>
return %0 : memref<20x5xf32, strided<[42, 2]>>
}
// CHECK-LABEL: func @no_fold_subview_with_non_unit_stride(
// CHECK: %[[SUBVIEW:.+]] = memref.subview
// CHECK: return %[[SUBVIEW]]
// -----
// CHECK-LABEL: func @no_fold_invalid_dynamic_slice
// CHECK: memref.subview %arg0[2] [%{{.*}}] [1] : memref<10xf32> to memref<?xf32, strided<[1], offset: 2>>
func.func @no_fold_invalid_dynamic_slice(%arg0: memref<10xf32>) -> memref<?xf32, strided<[1], offset: 2>> {
%c11 = arith.constant 11 : index
%0 = memref.subview %arg0 [2][%c11][1] : memref<10xf32> to memref<?xf32, strided<[1], offset: 2>>
func.return %0 : memref<?xf32, strided<[1], offset: 2>>
}
// -----
func.func @no_fold_dynamic_no_op_subview(%arg0 : memref<?x?xf32>) -> memref<?x?xf32, strided<[?, 1]>> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%0 = memref.dim %arg0, %c0 : memref<?x?xf32>
%1 = memref.dim %arg0, %c1 : memref<?x?xf32>
%2 = memref.subview %arg0[0, 0] [%0, %1] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1]>>
return %2 : memref<?x?xf32, strided<[?, 1]>>
}
// CHECK-LABEL: func @no_fold_dynamic_no_op_subview(
// CHECK: %[[SUBVIEW:.+]] = memref.subview
// CHECK: return %[[SUBVIEW]]
// -----
func.func @atomicrmw_cast_fold(%arg0 : f32, %arg1 : memref<4xf32>, %c : index) {
%v = memref.cast %arg1 : memref<4xf32> to memref<?xf32>
%a = memref.atomic_rmw addf %arg0, %v[%c] : (f32, memref<?xf32>) -> f32
return
}
// CHECK-LABEL: func @atomicrmw_cast_fold
// CHECK-NEXT: memref.atomic_rmw addf %arg0, %arg1[%arg2] : (f32, memref<4xf32>) -> f32
// -----
func.func @copy_of_cast(%m1: memref<?xf32>, %m2: memref<*xf32>) {
%casted1 = memref.cast %m1 : memref<?xf32> to memref<?xf32, strided<[?], offset: ?>>
%casted2 = memref.cast %m2 : memref<*xf32> to memref<?xf32, strided<[?], offset: ?>>
memref.copy %casted1, %casted2 : memref<?xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
return
}
// CHECK-LABEL: func @copy_of_cast(
// CHECK-SAME: %[[m1:.*]]: memref<?xf32>, %[[m2:.*]]: memref<*xf32>
// CHECK: %[[casted2:.*]] = memref.cast %[[m2]]
// CHECK: memref.copy %[[m1]], %[[casted2]]
// -----
func.func @self_copy(%m1: memref<?xf32>) {
memref.copy %m1, %m1 : memref<?xf32> to memref<?xf32>
return
}
// CHECK-LABEL: func @self_copy
// CHECK-NEXT: return
// -----
// CHECK-LABEL: func @empty_copy
// CHECK-NEXT: return
func.func @empty_copy(%m1: memref<0x10xf32>, %m2: memref<?x10xf32>) {
memref.copy %m1, %m2 : memref<0x10xf32> to memref<?x10xf32>
memref.copy %m2, %m1 : memref<?x10xf32> to memref<0x10xf32>
return
}
// -----
func.func @scopeMerge() {
memref.alloca_scope {
%cnt = "test.count"() : () -> index
%a = memref.alloca(%cnt) : memref<?xi64>
"test.use"(%a) : (memref<?xi64>) -> ()
}
return
}
// CHECK: func @scopeMerge() {
// CHECK-NOT: alloca_scope
// CHECK: %[[cnt:.+]] = "test.count"() : () -> index
// CHECK: %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
// CHECK: "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
// CHECK: return
// -----
func.func @scopeMerge2() {
"test.region"() ({
memref.alloca_scope {
%cnt = "test.count"() : () -> index
%a = memref.alloca(%cnt) : memref<?xi64>
"test.use"(%a) : (memref<?xi64>) -> ()
}
"test.terminator"() : () -> ()
}) : () -> ()
return
}
// CHECK: func @scopeMerge2() {
// CHECK: "test.region"() ({
// CHECK: memref.alloca_scope {
// CHECK: %[[cnt:.+]] = "test.count"() : () -> index
// CHECK: %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
// CHECK: "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
// CHECK: }
// CHECK: "test.terminator"() : () -> ()
// CHECK: }) : () -> ()
// CHECK: return
// CHECK: }
// -----
func.func @scopeMerge3() {
%cnt = "test.count"() : () -> index
"test.region"() ({
memref.alloca_scope {
%a = memref.alloca(%cnt) : memref<?xi64>
"test.use"(%a) : (memref<?xi64>) -> ()
}
"test.terminator"() : () -> ()
}) : () -> ()
return
}
// CHECK: func @scopeMerge3() {
// CHECK: %[[cnt:.+]] = "test.count"() : () -> index
// CHECK: %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
// CHECK: "test.region"() ({
// CHECK: memref.alloca_scope {
// CHECK: "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
// CHECK: }
// CHECK: "test.terminator"() : () -> ()
// CHECK: }) : () -> ()
// CHECK: return
// CHECK: }
// -----
func.func @scopeMerge4() {
%cnt = "test.count"() : () -> index
"test.region"() ({
memref.alloca_scope {
%a = memref.alloca(%cnt) : memref<?xi64>
"test.use"(%a) : (memref<?xi64>) -> ()
}
"test.op"() : () -> ()
"test.terminator"() : () -> ()
}) : () -> ()
return
}
// CHECK: func @scopeMerge4() {
// CHECK: %[[cnt:.+]] = "test.count"() : () -> index
// CHECK: "test.region"() ({
// CHECK: memref.alloca_scope {
// CHECK: %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
// CHECK: "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
// CHECK: }
// CHECK: "test.op"() : () -> ()
// CHECK: "test.terminator"() : () -> ()
// CHECK: }) : () -> ()
// CHECK: return
// CHECK: }
// -----
func.func @scopeMerge5() {
"test.region"() ({
memref.alloca_scope {
affine.parallel (%arg) = (0) to (64) {
%a = memref.alloca(%arg) : memref<?xi64>
"test.use"(%a) : (memref<?xi64>) -> ()
}
}
"test.op"() : () -> ()
"test.terminator"() : () -> ()
}) : () -> ()
return
}
// CHECK: func @scopeMerge5() {
// CHECK: "test.region"() ({
// CHECK: affine.parallel (%[[cnt:.+]]) = (0) to (64) {
// CHECK: %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
// CHECK: "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
// CHECK: }
// CHECK: "test.op"() : () -> ()
// CHECK: "test.terminator"() : () -> ()
// CHECK: }) : () -> ()
// CHECK: return
// CHECK: }
// -----
func.func @scopeInline(%arg : memref<index>) {
%cnt = "test.count"() : () -> index
"test.region"() ({
memref.alloca_scope {
memref.store %cnt, %arg[] : memref<index>
}
"test.terminator"() : () -> ()
}) : () -> ()
return
}
// CHECK: func @scopeInline
// CHECK-NOT: memref.alloca_scope
// -----
// Ensure this case not crash.
// CHECK-LABEL: func.func @scope_merge_without_terminator() {
// CHECK: "test.region"()
// CHECK: memref.alloca_scope
func.func @scope_merge_without_terminator() {
"test.region"() ({
memref.alloca_scope {
%cnt = "test.count"() : () -> index
%a = memref.alloca(%cnt) : memref<?xi64>
"test.use"(%a) : (memref<?xi64>) -> ()
}
}) : () -> ()
return
}
// -----
// Check that we simplify extract_strided_metadata of cast
// when the source of the cast is compatible with what
// `extract_strided_metadata`s accept.
//
// When we apply the transformation the resulting offset, sizes and strides
// should come straight from the inputs of the cast.
// Additionally the folder on extract_strided_metadata should propagate the
// static information.
//
// CHECK-LABEL: func @extract_strided_metadata_of_cast
// CHECK-SAME: %[[ARG:.*]]: memref<3x?xi32, strided<[4, ?], offset: ?>>)
//
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK: %[[BASE:.*]], %[[DYN_OFFSET:.*]], %[[DYN_SIZES:.*]]:2, %[[DYN_STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]]
//
// CHECK: return %[[BASE]], %[[DYN_OFFSET]], %[[C3]], %[[DYN_SIZES]]#1, %[[C4]], %[[DYN_STRIDES]]#1
func.func @extract_strided_metadata_of_cast(
%arg : memref<3x?xi32, strided<[4, ?], offset:?>>)
-> (memref<i32>, index,
index, index,
index, index) {
%cast =
memref.cast %arg :
memref<3x?xi32, strided<[4, ?], offset: ?>> to
memref<?x?xi32, strided<[?, ?], offset: ?>>
%base, %base_offset, %sizes:2, %strides:2 =
memref.extract_strided_metadata %cast:memref<?x?xi32, strided<[?, ?], offset: ?>>
-> memref<i32>, index,
index, index,
index, index
return %base, %base_offset,
%sizes#0, %sizes#1,
%strides#0, %strides#1 :
memref<i32>, index,
index, index,
index, index
}
// -----
// Check that we simplify extract_strided_metadata of cast
// when the source of the cast is compatible with what
// `extract_strided_metadata`s accept.
//
// Same as extract_strided_metadata_of_cast but with constant sizes and strides
// in the destination type.
//
// CHECK-LABEL: func @extract_strided_metadata_of_cast_w_csts
// CHECK-SAME: %[[ARG:.*]]: memref<?x?xi32, strided<[?, ?], offset: ?>>)
//
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[C18:.*]] = arith.constant 18 : index
// CHECK-DAG: %[[C25:.*]] = arith.constant 25 : index
// CHECK: %[[BASE:.*]], %[[DYN_OFFSET:.*]], %[[DYN_SIZES:.*]]:2, %[[DYN_STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]]
//
// CHECK: return %[[BASE]], %[[C25]], %[[C4]], %[[DYN_SIZES]]#1, %[[DYN_STRIDES]]#0, %[[C18]]
func.func @extract_strided_metadata_of_cast_w_csts(
%arg : memref<?x?xi32, strided<[?, ?], offset:?>>)
-> (memref<i32>, index,
index, index,
index, index) {
%cast =
memref.cast %arg :
memref<?x?xi32, strided<[?, ?], offset: ?>> to
memref<4x?xi32, strided<[?, 18], offset: 25>>
%base, %base_offset, %sizes:2, %strides:2 =
memref.extract_strided_metadata %cast:memref<4x?xi32, strided<[?, 18], offset: 25>>
-> memref<i32>, index,
index, index,
index, index
return %base, %base_offset,
%sizes#0, %sizes#1,
%strides#0, %strides#1 :
memref<i32>, index,
index, index,
index, index
}
// -----
// Check that we don't simplify extract_strided_metadata of
// cast when the source of the cast is unranked.
// Unranked memrefs cannot feed into extract_strided_metadata operations.
// Note: Technically we could still fold the sizes and strides.
//
// CHECK-LABEL: func @extract_strided_metadata_of_cast_unranked
// CHECK-SAME: %[[ARG:.*]]: memref<*xi32>)
//
// CHECK: %[[CAST:.*]] = memref.cast %[[ARG]] :
// CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[CAST]]
//
// CHECK: return %[[BASE]], %[[OFFSET]], %[[SIZES]]#0, %[[SIZES]]#1, %[[STRIDES]]#0, %[[STRIDES]]#1
func.func @extract_strided_metadata_of_cast_unranked(
%arg : memref<*xi32>)
-> (memref<i32>, index,
index, index,
index, index) {
%cast =
memref.cast %arg :
memref<*xi32> to
memref<?x?xi32, strided<[?, ?], offset: ?>>
%base, %base_offset, %sizes:2, %strides:2 =
memref.extract_strided_metadata %cast:memref<?x?xi32, strided<[?, ?], offset: ?>>
-> memref<i32>, index,
index, index,
index, index
return %base, %base_offset,
%sizes#0, %sizes#1,
%strides#0, %strides#1 :
memref<i32>, index,
index, index,
index, index
}
// -----
// CHECK-LABEL: func @reinterpret_noop
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3x4xf32>)
// CHECK-NEXT: return %[[ARG]]
func.func @reinterpret_noop(%arg : memref<2x3x4xf32>) -> memref<2x3x4xf32> {
%0 = memref.reinterpret_cast %arg to offset: [0], sizes: [2, 3, 4], strides: [12, 4, 1] : memref<2x3x4xf32> to memref<2x3x4xf32>
return %0 : memref<2x3x4xf32>
}
// -----
// CHECK-LABEL: func @reinterpret_constant_fold
// CHECK-SAME: (%[[ARG:.*]]: memref<f32>)
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [100, 100], strides: [100, 1]
// CHECK: %[[CAST:.*]] = memref.cast %[[RES]]
// CHECK: return %[[CAST]]
func.func @reinterpret_constant_fold(%arg0: memref<f32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c100 = arith.constant 100 : index
%reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%c0], sizes: [%c100, %c100], strides: [%c100, %c1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %reinterpret_cast : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// CHECK-LABEL: func @reinterpret_of_reinterpret
// CHECK-SAME: (%[[ARG:.*]]: memref<?xi8>, %[[SIZE1:.*]]: index, %[[SIZE2:.*]]: index)
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [%[[SIZE2]]], strides: [1]
// CHECK: return %[[RES]]
func.func @reinterpret_of_reinterpret(%arg : memref<?xi8>, %size1: index, %size2: index) -> memref<?xi8> {
%0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size1], strides: [1] : memref<?xi8> to memref<?xi8>
%1 = memref.reinterpret_cast %0 to offset: [0], sizes: [%size2], strides: [1] : memref<?xi8> to memref<?xi8>
return %1 : memref<?xi8>
}
// -----
// CHECK-LABEL: func @reinterpret_of_cast
// CHECK-SAME: (%[[ARG:.*]]: memref<?xi8>, %[[SIZE:.*]]: index)
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [%[[SIZE]]], strides: [1]
// CHECK: return %[[RES]]
func.func @reinterpret_of_cast(%arg : memref<?xi8>, %size: index) -> memref<?xi8> {
%0 = memref.cast %arg : memref<?xi8> to memref<5xi8>
%1 = memref.reinterpret_cast %0 to offset: [0], sizes: [%size], strides: [1] : memref<5xi8> to memref<?xi8>
return %1 : memref<?xi8>
}
// -----
// CHECK-LABEL: func @reinterpret_of_subview
// CHECK-SAME: (%[[ARG:.*]]: memref<?xi8>, %[[SIZE1:.*]]: index, %[[SIZE2:.*]]: index)
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [%[[SIZE2]]], strides: [1]
// CHECK: return %[[RES]]
func.func @reinterpret_of_subview(%arg : memref<?xi8>, %size1: index, %size2: index) -> memref<?xi8> {
%0 = memref.subview %arg[0] [%size1] [1] : memref<?xi8> to memref<?xi8>
%1 = memref.reinterpret_cast %0 to offset: [0], sizes: [%size2], strides: [1] : memref<?xi8> to memref<?xi8>
return %1 : memref<?xi8>
}
// -----
// Check that a reinterpret cast of an equivalent extract strided metadata
// is canonicalized to a plain cast when the destination type is different
// than the type of the original memref.
// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_type_mistach
// CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
// CHECK: %[[CAST:.*]] = memref.cast %[[ARG]] : memref<8x2xf32> to memref<?x?xf32,
// CHECK: return %[[CAST]]
func.func @reinterpret_of_extract_strided_metadata_w_type_mistach(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
%m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Similar to reinterpret_of_extract_strided_metadata_w_type_mistach except that
// we check that the match happen when the static information has been folded.
// E.g., in this case, we know that size of dim 0 is 8 and size of dim 1 is 2.
// So even if we don't use the values sizes#0, sizes#1, as long as they have the
// same constant value, the match is valid.
// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_constants
// CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
// CHECK: %[[CAST:.*]] = memref.cast %[[ARG]] : memref<8x2xf32> to memref<?x?xf32,
// CHECK: return %[[CAST]]
func.func @reinterpret_of_extract_strided_metadata_w_constants(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
%c8 = arith.constant 8: index
%m2 = memref.reinterpret_cast %base to offset: [0], sizes: [%c8, 2], strides: [2, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Check that a reinterpret cast of an equivalent extract strided metadata
// is completely removed when the original memref has the same type.
// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_same_type
// CHECK-SAME: (%[[ARG:.*]]: memref<?x?xf32
// CHECK: return %[[ARG]]
func.func @reinterpret_of_extract_strided_metadata_same_type(%arg0 : memref<?x?xf32, strided<[?,?], offset: ?>>) -> memref<?x?xf32, strided<[?,?], offset: ?>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<?x?xf32, strided<[?,?], offset: ?>> -> memref<f32>, index, index, index, index, index
%m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?,?], offset:?>>
return %m2 : memref<?x?xf32, strided<[?,?], offset:?>>
}
// -----
// Check that we don't simplify reinterpret cast of extract strided metadata
// when the strides don't match.
// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_different_stride
// CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [4, 2, 2], strides: [1, 1, 1]
// CHECK: %[[CAST:.*]] = memref.cast %[[RES]]
// CHECK: return %[[CAST]]
func.func @reinterpret_of_extract_strided_metadata_w_different_stride(%arg0 : memref<8x2xf32>) -> memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
%m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [4, 2, 2], strides: [1, 1, %strides#1] : memref<f32> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
return %m2 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
}
// -----
// Check that we don't simplify reinterpret cast of extract strided metadata
// when the offset doesn't match.
// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_different_offset
// CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [1], sizes: [8, 2], strides: [2, 1]
// CHECK: %[[CAST:.*]] = memref.cast %[[RES]]
// CHECK: return %[[CAST]]
func.func @reinterpret_of_extract_strided_metadata_w_different_offset(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
%m2 = memref.reinterpret_cast %base to offset: [1], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Check that reinterpret_cast with a negative constant size.
// Folding would attempt to create a MemRefType with a negative static dimension,
// which triggers an assertion in MemRefType::get (issue #188407).
// CHECK-LABEL: func @reinterpret_cast_with_negative_size
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3xf32>)
// CHECK: %[[SZ:.*]] = arith.constant -1 : index
// CHECK: memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [1, %[[SZ]]], strides: [-1, 1]
func.func @reinterpret_cast_with_negative_size(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%sz = arith.constant -1 : index
%output = memref.reinterpret_cast %arg0 to
offset: [%c0], sizes: [%c1, %sz], strides: [%sz, %c1]
: memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %output : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Check that reinterpret_cast with a negative constant offset.
// Folding would create an op with a static negative offset, which violates the
// ViewLikeInterface constraint that offsets must be non-negative.
// CHECK-LABEL: func @reinterpret_cast_with_negative_offset
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3xf32>)
// CHECK: %[[NEG:.*]] = arith.constant -1 : index
// CHECK: memref.reinterpret_cast %[[ARG]] to offset: [%[[NEG]]], sizes: [1, 2], strides: [2, 1]
func.func @reinterpret_cast_with_negative_offset(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%neg = arith.constant -1 : index
%output = memref.reinterpret_cast %arg0 to
offset: [%neg], sizes: [%c1, %c2], strides: [%c2, %c1]
: memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %output : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Check that reinterpret_cast with a negative constant size and offset.
// CHECK-LABEL: func @reinterpret_cast_with_negative_size_and_offset
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3xf32>)
// CHECK: %[[NEG:.*]] = arith.constant -1 : index
// CHECK: memref.reinterpret_cast %[[ARG]] to offset: [%[[NEG]]], sizes: [1, %[[NEG]]], strides: [2, 1]
func.func @reinterpret_cast_with_negative_size_and_offset(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%neg = arith.constant -1 : index
%output = memref.reinterpret_cast %arg0 to
offset: [%neg], sizes: [%c1, %neg], strides: [%c2, %c1]
: memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %output : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Check that reinterpret_cast with all negative constant size and offset is not
// folded.
// CHECK-LABEL: func @reinterpret_cast_no_fold_with_all_negative_size_and_offset
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3xf32>)
// CHECK: %[[NEG:.*]] = arith.constant -1 : index
// CHECK: memref.reinterpret_cast %[[ARG]] to offset: [%[[NEG]]], sizes: [%[[NEG]], %[[NEG]]], strides: [2, 1]
func.func @reinterpret_cast_no_fold_with_all_negative_size_and_offset(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%neg = arith.constant -1 : index
%output = memref.reinterpret_cast %arg0 to
offset: [%neg], sizes: [%neg, %neg], strides: [2, 1]
: memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %output : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
// Check that reinterpret_cast with a negative constant stride IS folded.
// Negative strides are valid in MemRef layouts (e.g. reverse iteration),
// and the ViewLikeInterface places no non-negativity constraint on strides.
// CHECK-LABEL: func @reinterpret_cast_fold_negative_stride
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3xf32>)
// CHECK-NOT: arith.constant
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [1, 2], strides: [-1, 1]
// CHECK: memref.cast %[[RC]]
func.func @reinterpret_cast_fold_negative_stride(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%neg = arith.constant -1 : index
%output = memref.reinterpret_cast %arg0 to
offset: [%c0], sizes: [%c1, %c2], strides: [%neg, %c1]
: memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
return %output : memref<?x?xf32, strided<[?, ?], offset: ?>>
}
// -----
func.func @canonicalize_rank_reduced_subview(%arg0 : memref<8x?xf32>,
%arg1 : index) -> memref<?xf32, strided<[?], offset: ?>> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%0 = memref.subview %arg0[%c0, %c0] [1, %arg1] [%c1, %c1] : memref<8x?xf32> to memref<?xf32, strided<[?], offset: ?>>
return %0 : memref<?xf32, strided<[?], offset: ?>>
}
// CHECK: func @canonicalize_rank_reduced_subview
// CHECK-SAME: %[[ARG0:.+]]: memref<8x?xf32>
// CHECK-SAME: %[[ARG1:.+]]: index
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG0]][0, 0] [1, %[[ARG1]]] [1, 1]
// CHECK-SAME: memref<8x?xf32> to memref<?xf32, strided<[1]>>
// -----
// CHECK-LABEL: func @memref_realloc_dead
// CHECK-SAME: %[[SRC:[0-9a-z]+]]: memref<2xf32>
// CHECK-NOT: memref.realloc
// CHECK: return %[[SRC]]
func.func @memref_realloc_dead(%src : memref<2xf32>, %v : f32) -> memref<2xf32>{
%0 = memref.realloc %src : memref<2xf32> to memref<4xf32>
%i2 = arith.constant 2 : index
memref.store %v, %0[%i2] : memref<4xf32>
return %src : memref<2xf32>
}
// -----
// CHECK-LABEL: func @collapse_expand_fold_to_cast(
// CHECK-SAME: %[[m:.*]]: memref<?xf32, strided<[1]>, 3>
// CHECK: %[[casted:.*]] = memref.cast %[[m]] : memref<?xf32, strided<[1]>, 3> to memref<?xf32, 3
// CHECK: return %[[casted]]
func.func @collapse_expand_fold_to_cast(%m: memref<?xf32, strided<[1]>, 3>, %sz0: index)
-> (memref<?xf32, 3>)
{
%0 = memref.expand_shape %m [[0, 1]] output_shape [1, %sz0]
: memref<?xf32, strided<[1]>, 3> into memref<1x?xf32, 3>
%1 = memref.collapse_shape %0 [[0, 1]]
: memref<1x?xf32, 3> into memref<?xf32, 3>
return %1 : memref<?xf32, 3>
}
// -----
// CHECK-LABEL: func @expand_collapse_fold_to_internal_stride_cast(
// CHECK-SAME: %[[m:.*]]: memref<3x1x2x384xui8, strided<[1179648, 768, 384, 1]>>
// CHECK: %[[casted:.*]] = memref.cast %[[m]] : memref<3x1x2x384xui8, strided<[1179648, 768, 384, 1]>> to memref<3x1x2x384xui8, strided<[1179648, 1179648, 384, 1]>>
func.func @expand_collapse_fold_to_internal_stride_cast(%m: memref<3x1x2x384xui8, strided<[1179648, 768, 384, 1]>>)
-> (memref<3x1x2x384xui8, strided<[1179648, 1179648, 384, 1]>>)
{
%0 = memref.collapse_shape %m [[0, 1], [2], [3]]
: memref<3x1x2x384xui8, strided<[1179648, 768, 384, 1]>>
into memref<3x2x384xui8, strided<[1179648, 384, 1]>>
%1 = memref.expand_shape %0 [[0, 1], [2], [3]] output_shape [3, 1, 2, 384]
: memref<3x2x384xui8, strided<[1179648, 384, 1]>>
into memref<3x1x2x384xui8, strided<[1179648, 1179648, 384, 1]>>
return %1 : memref<3x1x2x384xui8, strided<[1179648, 1179648, 384, 1]>>
}
// -----
// CHECK-LABEL: func @expand_collapse_fold_to_outermost_stride_cast(
// CHECK-SAME: %[[m:.*]]: memref<1x3x2x384xui8, strided<[1179648, 768, 384, 1]>>
// CHECK: %[[casted:.*]] = memref.cast %[[m]] : memref<1x3x2x384xui8, strided<[1179648, 768, 384, 1]>> to memref<1x3x2x384xui8, strided<[2304, 768, 384, 1]>>
// CHECK: return %[[casted]]
func.func @expand_collapse_fold_to_outermost_stride_cast(%m: memref<1x3x2x384xui8, strided<[1179648, 768, 384, 1]>>)
-> (memref<1x3x2x384xui8, strided<[2304, 768, 384, 1]>>)
{
%0 = memref.collapse_shape %m [[0, 1], [2], [3]]
: memref<1x3x2x384xui8, strided<[1179648, 768, 384, 1]>>
into memref<3x2x384xui8, strided<[768, 384, 1]>>
%1 = memref.expand_shape %0 [[0, 1], [2], [3]] output_shape [1, 3, 2, 384]
: memref<3x2x384xui8, strided<[768, 384, 1]>>
into memref<1x3x2x384xui8, strided<[2304, 768, 384, 1]>>
return %1 : memref<1x3x2x384xui8, strided<[2304, 768, 384, 1]>>
}
// -----
// CHECK-LABEL: func @expand_collapse_do_not_fold_to_cast(
// CHECK-NOT: memref.cast
func.func @expand_collapse_do_not_fold_to_cast(%m: memref<1x3x2x384xui8, strided<[1179648, 768, 384, 1]>>)
-> (memref<3x1x2x384xui8, strided<[768, 768, 384, 1]>>)
{
%0 = memref.collapse_shape %m [[0, 1], [2], [3]]
: memref<1x3x2x384xui8, strided<[1179648, 768, 384, 1]>>
into memref<3x2x384xui8, strided<[768, 384, 1]>>
%1 = memref.expand_shape %0 [[0, 1], [2], [3]] output_shape [3, 1, 2, 384]
: memref<3x2x384xui8, strided<[768, 384, 1]>>
into memref<3x1x2x384xui8, strided<[768, 768, 384, 1]>>
return %1 : memref<3x1x2x384xui8, strided<[768, 768, 384, 1]>>
}
// -----
// CHECK-LABEL: func @expand_collapse_dynamic_do_not_fold_to_cast(
// CHECK-NOT: memref.cast
func.func @expand_collapse_dynamic_do_not_fold_to_cast(%m: memref<1x?x1x32xsi8, strided<[?, 32, 32, 1]>>, %dyn_size: index)
-> (memref<1x1x?x32xsi8, strided<[?, ?, 32, 1]>>)
{
%0 = memref.collapse_shape %m [[0], [1, 2], [3]]
: memref<1x?x1x32xsi8, strided<[?, 32, 32, 1]>> into memref<1x?x32xsi8, strided<[?, 32, 1]>>
%1 = memref.expand_shape %0 [[0, 1], [2], [3]] output_shape [1, 1, %dyn_size, 32]
: memref<1x?x32xsi8, strided<[?, 32, 1]>> into memref<1x1x?x32xsi8, strided<[?, ?, 32, 1]>>
return %1 : memref<1x1x?x32xsi8, strided<[?, ?, 32, 1]>>
}
// -----
// CHECK-LABEL: func @fold_trivial_subviews(
// CHECK-SAME: %[[m:.*]]: memref<?xf32, strided<[?], offset: ?>>
// CHECK: %[[subview:.*]] = memref.subview %[[m]][5]
// CHECK: return %[[subview]]
func.func @fold_trivial_subviews(%m: memref<?xf32, strided<[?], offset: ?>>,
%sz: index)
-> memref<?xf32, strided<[?], offset: ?>>
{
%0 = memref.subview %m[5] [%sz] [1]
: memref<?xf32, strided<[?], offset: ?>>
to memref<?xf32, strided<[?], offset: ?>>
%1 = memref.subview %0[0] [%sz] [1]
: memref<?xf32, strided<[?], offset: ?>>
to memref<?xf32, strided<[?], offset: ?>>
return %1 : memref<?xf32, strided<[?], offset: ?>>
}
// -----
// CHECK-LABEL: func @load_store_nontemporal(
func.func @load_store_nontemporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>, %output : memref<32xf32, affine_map<(d0) -> (d0)>>) {
%1 = arith.constant 7 : index
// CHECK: memref.load %{{.*}}[%{{.*}}] {nontemporal = true} : memref<32xf32>
%2 = memref.load %input[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
// CHECK: memref.store %{{.*}}, %{{.*}}[%{{.*}}] {nontemporal = true} : memref<32xf32>
memref.store %2, %output[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
func.return
}
// -----
memref.global "private" constant @__constant_32xf32 : memref<32xf32> = dense<1.000000e+00>
// CHECK-LABEL: func @fold_const_splat_global
func.func @fold_const_splat_global() -> memref<32xf32> {
// CHECK-NEXT: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
%0 = memref.get_global @__constant_32xf32 : memref<32xf32>
%alloc = memref.alloc() : memref<32xf32>
%c32 = arith.constant 32 : index
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
scf.for %arg0 = %c0 to %c32 step %c1 {
%1 = memref.load %0[%arg0] : memref<32xf32>
// CHECK: memref.store %[[CST]], %{{.*}}
memref.store %1, %alloc[%arg0] : memref<32xf32>
}
return %alloc : memref<32xf32>
}
// -----
// CHECK-LABEL: func @fold_trivial_memory_space_cast(
// CHECK-SAME: %[[arg:.*]]: memref<?xf32>
// CHECK: return %[[arg]]
func.func @fold_trivial_memory_space_cast(%arg : memref<?xf32>) -> memref<?xf32> {
%0 = memref.memory_space_cast %arg : memref<?xf32> to memref<?xf32>
return %0 : memref<?xf32>
}
// -----
// CHECK-LABEL: func @fold_multiple_memory_space_cast(
// CHECK-SAME: %[[arg:.*]]: memref<?xf32>
// CHECK: %[[res:.*]] = memref.memory_space_cast %[[arg]] : memref<?xf32> to memref<?xf32, 2>
// CHECK: return %[[res]]
func.func @fold_multiple_memory_space_cast(%arg : memref<?xf32>) -> memref<?xf32, 2> {
%0 = memref.memory_space_cast %arg : memref<?xf32> to memref<?xf32, 1>
%1 = memref.memory_space_cast %0 : memref<?xf32, 1> to memref<?xf32, 2>
return %1 : memref<?xf32, 2>
}
// -----
// CHECK-LABEL: func private @ub_negative_alloc_size
func.func private @ub_negative_alloc_size() -> memref<?x?x?xi1> {
%idx1 = index.constant 1
%c-2 = arith.constant -2 : index
%c15 = arith.constant 15 : index
// CHECK: %[[ALLOC:.*]] = memref.alloc(%c-2) : memref<15x?x1xi1>
%alloc = memref.alloc(%c15, %c-2, %idx1) : memref<?x?x?xi1>
return %alloc : memref<?x?x?xi1>
}
// -----
// CHECK-LABEL: func @subview_rank_reduction(
// CHECK-SAME: %[[arg0:.*]]: memref<1x384x384xf32>, %[[arg1:.*]]: index
func.func @subview_rank_reduction(%arg0: memref<1x384x384xf32>, %idx: index)
-> memref<?x?xf32, strided<[384, 1], offset: ?>> {
%c1 = arith.constant 1 : index
// CHECK: %[[subview:.*]] = memref.subview %[[arg0]][0, %[[arg1]], %[[arg1]]] [1, 1, %[[arg1]]] [1, 1, 1] : memref<1x384x384xf32> to memref<1x?xf32, strided<[384, 1], offset: ?>>
// CHECK: %[[cast:.*]] = memref.cast %[[subview]] : memref<1x?xf32, strided<[384, 1], offset: ?>> to memref<?x?xf32, strided<[384, 1], offset: ?>>
%0 = memref.subview %arg0[0, %idx, %idx] [1, %c1, %idx] [1, 1, 1]
: memref<1x384x384xf32> to memref<?x?xf32, strided<[384, 1], offset: ?>>
// CHECK: return %[[cast]]
return %0 : memref<?x?xf32, strided<[384, 1], offset: ?>>
}
// -----
// CHECK-LABEL: func @fold_double_transpose(
// CHECK-SAME: %[[arg0:.*]]: memref<1x2x3x4x5xf32>
func.func @fold_double_transpose(%arg0: memref<1x2x3x4x5xf32>) -> memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>> {
// CHECK: %[[ONETRANSPOSE:.+]] = memref.transpose %[[arg0]] (d0, d1, d2, d3, d4) -> (d4, d2, d1, d3, d0)
%0 = memref.transpose %arg0 (d0, d1, d2, d3, d4) -> (d1, d0, d4, d3, d2) : memref<1x2x3x4x5xf32> to memref<2x1x5x4x3xf32, strided<[60, 120, 1, 5, 20]>>
%1 = memref.transpose %0 (d1, d0, d4, d3, d2) -> (d4, d2, d1, d3, d0) : memref<2x1x5x4x3xf32, strided<[60, 120, 1, 5, 20]>> to memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>>
// CHECK: return %[[ONETRANSPOSE]]
return %1 : memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>>
}
// -----
// CHECK-LABEL: func @fold_double_transpose2(
// CHECK-SAME: %[[arg0:.*]]: memref<1x2x3x4x5xf32>
func.func @fold_double_transpose2(%arg0: memref<1x2x3x4x5xf32>) -> memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>> {
// CHECK: %[[ONETRANSPOSE:.+]] = memref.transpose %[[arg0]] (d0, d1, d2, d3, d4) -> (d4, d2, d1, d3, d0)
%0 = memref.transpose %arg0 (d0, d1, d2, d3, d4) -> (d0, d1, d4, d3, d2) : memref<1x2x3x4x5xf32> to memref<1x2x5x4x3xf32, strided<[120, 60, 1, 5, 20]>>
%1 = memref.transpose %0 (d0, d1, d4, d3, d2) -> (d4, d2, d1, d3, d0) : memref<1x2x5x4x3xf32, strided<[120, 60, 1, 5, 20]>> to memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>>
// CHECK: return %[[ONETRANSPOSE]]
return %1 : memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>>
}
// -----
// CHECK-LABEL: func @fold_identity_transpose(
// CHECK-SAME: %[[arg0:.*]]: memref<1x2x3x4x5xf32>
func.func @fold_identity_transpose(%arg0: memref<1x2x3x4x5xf32>) -> memref<1x2x3x4x5xf32> {
%0 = memref.transpose %arg0 (d0, d1, d2, d3, d4) -> (d1, d0, d4, d3, d2) : memref<1x2x3x4x5xf32> to memref<2x1x5x4x3xf32, strided<[60, 120, 1, 5, 20]>>
%1 = memref.transpose %0 (d1, d0, d4, d3, d2) -> (d0, d1, d2, d3, d4) : memref<2x1x5x4x3xf32, strided<[60, 120, 1, 5, 20]>> to memref<1x2x3x4x5xf32>
// CHECK: return %[[arg0]]
return %1 : memref<1x2x3x4x5xf32>
}
// -----
#transpose_map = affine_map<(d0, d1)[s0] -> (d0 + d1 * s0)>
// CHECK-LABEL: func @cannot_fold_transpose_cast(
// CHECK-SAME: %[[arg0:.*]]: memref<?x4xf32>
func.func @cannot_fold_transpose_cast(%arg0: memref<?x4xf32>) -> memref<?x?xf32, #transpose_map> {
// CHECK: %[[CAST:.*]] = memref.cast %[[arg0]] : memref<?x4xf32> to memref<?x?xf32>
%cast = memref.cast %arg0 : memref<?x4xf32> to memref<?x?xf32>
// CHECK: %[[TRANSPOSE:.*]] = memref.transpose %[[CAST]] (d0, d1) -> (d1, d0) : memref<?x?xf32> to memref<?x?xf32, #{{.*}}>
%transpose = memref.transpose %cast (d0, d1) -> (d1, d0) : memref<?x?xf32> to memref<?x?xf32, #transpose_map>
// CHECK: return %[[TRANSPOSE]]
return %transpose : memref<?x?xf32, #transpose_map>
}
// -----
// CHECK-LABEL: func @fold_assume_alignment_chain
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
func.func @fold_assume_alignment_chain(%0: memref<128xf32>) -> memref<128xf32> {
// CHECK: %[[ALIGN:.+]] = memref.assume_alignment %[[ARG0]], 16
%1 = memref.assume_alignment %0, 16 : memref<128xf32>
// CHECK-NOT: memref.assume_alignment
%2 = memref.assume_alignment %1, 16 : memref<128xf32>
// CHECK: return %[[ALIGN]]
return %2 : memref<128xf32>
}
// -----
// CHECK-LABEL: func @fold_view_cast
// CHECK-SAME: (%[[ARG:.*]]: memref<128xi8>)
func.func @fold_view_cast(%0: memref<128xi8>) -> memref<i32> {
%c0 = arith.constant 0 : index
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[RES:.*]] = memref.view %[[ARG]][%[[C0]]][] : memref<128xi8> to memref<i32>
// CHECK: return %[[RES]]
%1 = memref.cast %0 : memref<128xi8> to memref<?xi8>
%res = memref.view %1[%c0][] : memref<?xi8> to memref<i32>
return %res : memref<i32>
}
// -----
// CHECK-LABEL: func @fold_view_same_source_result_types
// CHECK-SAME: (%[[ARG:.*]]: memref<128xi8>)
func.func @fold_view_same_source_result_types(%0: memref<128xi8>) -> memref<128xi8> {
%c0 = arith.constant 0 : index
// CHECK-NOT: memref.view
// CHECK: return %[[ARG]]
%res = memref.view %0[%c0][] : memref<128xi8> to memref<128xi8>
return %res : memref<128xi8>
}
// -----
// CHECK-LABEL: func @non_fold_view_non_zero_offset
// CHECK-SAME: (%[[ARG:.*]]: memref<128xi8>)
func.func @non_fold_view_non_zero_offset(%0: memref<128xi8>) -> memref<128xi8> {
%c1 = arith.constant 1 : index
// CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[RES:.*]] = memref.view %[[ARG]][%[[C1]]][] : memref<128xi8> to memref<128xi8>
// CHECK: return %[[RES]]
%res = memref.view %0[%c1][] : memref<128xi8> to memref<128xi8>
return %res : memref<128xi8>
}
// -----
// CHECK-LABEL: func @non_fold_view_same_source_dynamic_size
// CHECK-SAME: (%[[ARG:.*]]: memref<?xi8>, %[[SIZE:.*]]: index)
func.func @non_fold_view_same_source_dynamic_size(%0: memref<?xi8>, %arg0 : index) -> memref<?xi8> {
%c0 = arith.constant 0: index
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[RES:.*]] = memref.view %[[ARG]][%[[C0]]][%[[SIZE]]] : memref<?xi8> to memref<?xi8>
// CHECK: return %[[RES]]
%res = memref.view %0[%c0][%arg0] : memref<?xi8> to memref<?xi8>
return %res : memref<?xi8>
}
// -----
// CHECK-LABEL: func.func @replace_view_static_dims(
// CHECK-SAME: %[[ARG0:.*]]: memref<?xi8>, %[[ARG1:.*]]: index) -> memref<?x4xi32> {
// CHECK: %[[VIEW_0:.*]] = memref.view %[[ARG0]]{{\[}}%[[ARG1]]][] : memref<?xi8> to memref<5x4xi32>
// CHECK: %[[CAST_0:.*]] = memref.cast %[[VIEW_0]] : memref<5x4xi32> to memref<?x4xi32>
// CHECK: return %[[CAST_0]] : memref<?x4xi32>
// CHECK: }
func.func @replace_view_static_dims(%src: memref<?xi8>, %offset : index) -> memref<?x4xi32> {
%c5 = arith.constant 5: index
%res = memref.view %src[%offset][%c5] : memref<?xi8> to memref<?x4xi32>
return %res : memref<?x4xi32>
}
// -----
// CHECK-LABEL: func.func @non_replace_view_negative_static_dims(
// CHECK-SAME: %[[ARG0:.*]]: memref<?xi8>, %[[ARG1:.*]]: index) -> memref<?x4xi32> {
// CHECK: %[[CONSTANT_0:.*]] = arith.constant -1 : index
// CHECK: %[[VIEW_0:.*]] = memref.view %[[ARG0]]{{\[}}%[[ARG1]]]{{\[}}%[[CONSTANT_0]]] : memref<?xi8> to memref<?x4xi32>
// CHECK: return %[[VIEW_0]] : memref<?x4xi32>
// CHECK: }
func.func @non_replace_view_negative_static_dims(%src: memref<?xi8>, %offset : index) -> memref<?x4xi32> {
%c-1 = arith.constant -1: index
%res = memref.view %src[%offset][%c-1] : memref<?xi8> to memref<?x4xi32>
return %res : memref<?x4xi32>
}
// -----
// Verify that canonicalization does not crash when a memref.dim is applied to
// a subview with ambiguous dropped dimensions (multiple size-1 source dims with
// all-dynamic strides). The dim should be folded to the corresponding subview
// size operand.
// See: https://github.com/llvm/llvm-project/issues/111244
// CHECK-LABEL: func @no_crash_dim_of_ambiguous_subview
// CHECK-SAME: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: index) -> index
// CHECK-NOT: memref.dim
// CHECK: return %[[ARG1]]
func.func @no_crash_dim_of_ambiguous_subview(
%arg0: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>, %arg1: index) -> index {
%c1 = arith.constant 1 : index
%subview = memref.subview %arg0[0, 0, 0] [1, %arg1, 1] [1, 1, 1]
: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to memref<1x?xf32, strided<[?, ?], offset: ?>>
%dim = memref.dim %subview, %c1 : memref<1x?xf32, strided<[?, ?], offset: ?>>
return %dim : index
}
Reference in new issue View Git Blame Copy Permalink