636 lines · plain
1// RUN: mlir-opt %s | mlir-opt | FileCheck %s2// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s3 4// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>5 6// CHECK-LABEL: func @alloc() {7func.func @alloc() {8^bb0:9 // Test simple alloc.10 // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>11 %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>12 13 %c0 = "arith.constant"() {value = 0: index} : () -> index14 %c1 = "arith.constant"() {value = 1: index} : () -> index15 16 // Test alloc with dynamic dimensions.17 // CHECK: %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>18 %1 = memref.alloc(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>19 20 // Test alloc with no dynamic dimensions and one symbol.21 // CHECK: %{{.*}} = memref.alloc()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>22 %2 = memref.alloc()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>23 24 // Test alloc with dynamic dimensions and one symbol.25 // CHECK: %{{.*}} = memref.alloc(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>26 %3 = memref.alloc(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>27 28 // Alloc with no mappings.29 // b/116054838 Parser crash while parsing ill-formed AllocOp30 // CHECK: %{{.*}} = memref.alloc() : memref<2xi32>31 %4 = memref.alloc() : memref<2 x i32>32 33 // CHECK: return34 return35}36 37// CHECK-LABEL: func @alloca() {38func.func @alloca() {39^bb0:40 // Test simple alloc.41 // CHECK: %{{.*}} = memref.alloca() : memref<1024x64xf32, 1>42 %0 = memref.alloca() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>43 44 %c0 = "arith.constant"() {value = 0: index} : () -> index45 %c1 = "arith.constant"() {value = 1: index} : () -> index46 47 // Test alloca with dynamic dimensions.48 // CHECK: %{{.*}} = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>49 %1 = memref.alloca(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>50 51 // Test alloca with no dynamic dimensions and one symbol.52 // CHECK: %{{.*}} = memref.alloca()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>53 %2 = memref.alloca()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>54 55 // Test alloca with dynamic dimensions and one symbol.56 // CHECK: %{{.*}} = memref.alloca(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>57 %3 = memref.alloca(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>58 59 // Alloca with no mappings, but with alignment.60 // CHECK: %{{.*}} = memref.alloca() {alignment = 64 : i64} : memref<2xi32>61 %4 = memref.alloca() {alignment = 64} : memref<2 x i32>62 63 return64}65 66// CHECK-LABEL: func @dealloc() {67func.func @dealloc() {68^bb0:69 // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32>70 %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>71 72 // CHECK: memref.dealloc %{{.*}} : memref<1024x64xf32>73 memref.dealloc %0 : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>74 return75}76 77// CHECK-LABEL: func @load_store78func.func @load_store() {79^bb0:80 // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>81 %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>82 83 %1 = arith.constant 0 : index84 %2 = arith.constant 1 : index85 86 // CHECK: %{{.*}} = memref.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x64xf32, 1>87 %3 = memref.load %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>88 89 // CHECK: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x64xf32, 1>90 memref.store %3, %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>91 92 return93}94 95// CHECK-LABEL: func @dma_ops()96func.func @dma_ops() {97 %c0 = arith.constant 0 : index98 %stride = arith.constant 32 : index99 %elt_per_stride = arith.constant 16 : index100 101 %A = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>102 %Ah = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 1>103 %tag = memref.alloc() : memref<1 x f32>104 105 %num_elements = arith.constant 256 : index106 107 memref.dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0] : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>108 memref.dma_wait %tag[%c0], %num_elements : memref<1 x f32>109 // CHECK: dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}[%{{.*}}] : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>110 // CHECK-NEXT: dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>111 112 // DMA with strides113 memref.dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0], %stride, %elt_per_stride : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>114 memref.dma_wait %tag[%c0], %num_elements : memref<1 x f32>115 // CHECK-NEXT: dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}} : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>116 // CHECK-NEXT: dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>117 118 return119}120 121// CHECK-LABEL: func @memref_reinterpret_cast122func.func @memref_reinterpret_cast(%in: memref<?xf32>)123 -> memref<10x?xf32, strided<[?, 1], offset: ?>> {124 %c0 = arith.constant 0 : index125 %c10 = arith.constant 10 : index126 %out = memref.reinterpret_cast %in to127 offset: [%c0], sizes: [10, %c10], strides: [%c10, 1]128 : memref<?xf32> to memref<10x?xf32, strided<[?, 1], offset: ?>>129 return %out : memref<10x?xf32, strided<[?, 1], offset: ?>>130}131 132// CHECK-LABEL: func @memref_reinterpret_cast_static_to_dynamic_sizes133func.func @memref_reinterpret_cast_static_to_dynamic_sizes(%in: memref<?xf32>)134 -> memref<10x?xf32, strided<[?, 1], offset: ?>> {135 %out = memref.reinterpret_cast %in to136 offset: [1], sizes: [10, 10], strides: [1, 1]137 : memref<?xf32> to memref<10x?xf32, strided<[?, 1], offset: ?>>138 return %out : memref<10x?xf32, strided<[?, 1], offset: ?>>139}140 141// CHECK-LABEL: func @memref_reinterpret_cast_dynamic_offset142func.func @memref_reinterpret_cast_dynamic_offset(%in: memref<?xf32>, %offset: index)143 -> memref<10x?xf32, strided<[?, 1], offset: ?>> {144 %out = memref.reinterpret_cast %in to145 offset: [%offset], sizes: [10, 10], strides: [1, 1]146 : memref<?xf32> to memref<10x?xf32, strided<[?, 1], offset: ?>>147 return %out : memref<10x?xf32, strided<[?, 1], offset: ?>>148}149 150// CHECK-LABEL: func @memref_reshape(151func.func @memref_reshape(%unranked: memref<*xf32>, %shape1: memref<1xi32>,152 %shape2: memref<2xi32>, %shape3: memref<?xi32>) -> memref<*xf32> {153 %dyn_vec = memref.reshape %unranked(%shape1)154 : (memref<*xf32>, memref<1xi32>) -> memref<?xf32>155 %dyn_mat = memref.reshape %dyn_vec(%shape2)156 : (memref<?xf32>, memref<2xi32>) -> memref<?x?xf32>157 %new_unranked = memref.reshape %dyn_mat(%shape3)158 : (memref<?x?xf32>, memref<?xi32>) -> memref<*xf32>159 return %new_unranked : memref<*xf32>160}161 162// CHECK-LABEL: memref.global @memref0 : memref<2xf32>163memref.global @memref0 : memref<2xf32>164 165// CHECK-LABEL: memref.global constant @memref1 : memref<2xf32> = dense<[0.000000e+00, 1.000000e+00]>166memref.global constant @memref1 : memref<2xf32> = dense<[0.0, 1.0]>167 168// CHECK-LABEL: memref.global @memref2 : memref<2xf32> = uninitialized169memref.global @memref2 : memref<2xf32> = uninitialized170 171// CHECK-LABEL: memref.global "private" @memref3 : memref<2xf32> = uninitialized172memref.global "private" @memref3 : memref<2xf32> = uninitialized173 174// CHECK-LABEL: memref.global "private" constant @memref4 : memref<2xf32> = uninitialized175memref.global "private" constant @memref4 : memref<2xf32> = uninitialized176 177// CHECK-LABEL: memref.global "private" constant @memref5 : memref<1xf16, 42 : i32> = dense<1.000000e+00>178"memref.global"() <{constant, initial_value = dense<1.000000e+00> : tensor<1xf16>, sym_name = "memref5", sym_visibility = "private", type = memref<1xf16, 42 : i32>}> : () -> ()179 180// CHECK-LABEL: func @read_global_memref181func.func @read_global_memref() {182 %0 = memref.get_global @memref0 : memref<2xf32>183 return184}185 186// CHECK-LABEL: func @memref_copy187func.func @memref_copy() {188 %0 = memref.alloc() : memref<2xf32>189 %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>190 %2 = memref.alloc() : memref<2xf32>191 %3 = memref.cast %0 : memref<2xf32> to memref<*xf32>192 memref.copy %1, %3 : memref<*xf32> to memref<*xf32>193 return194}195 196// CHECK-LABEL: func @memref_dealloc197func.func @memref_dealloc() {198 %0 = memref.alloc() : memref<2xf32>199 %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>200 memref.dealloc %1 : memref<*xf32>201 return202}203 204 205// CHECK-LABEL: func @memref_alloca_scope206func.func @memref_alloca_scope() {207 memref.alloca_scope {208 memref.alloca_scope.return209 }210 return211}212 213// CHECK-LABEL: func @memref_cast(%arg0214func.func @memref_cast(%arg0: memref<4xf32>, %arg1 : memref<?xf32>, %arg2 : memref<64x16x4xf32, strided<[64, 4, 1], offset: 0>>) {215 // CHECK: memref.cast %{{.*}} : memref<4xf32> to memref<?xf32>216 %0 = memref.cast %arg0 : memref<4xf32> to memref<?xf32>217 218 // CHECK: memref.cast %{{.*}} : memref<?xf32> to memref<4xf32>219 %1 = memref.cast %arg1 : memref<?xf32> to memref<4xf32>220 221 // CHECK: memref.cast %{{.*}} : memref<64x16x4xf32, strided<[64, 4, 1]>> to memref<64x16x4xf32, strided<[?, ?, ?], offset: ?>>222 %2 = memref.cast %arg2 : memref<64x16x4xf32, strided<[64, 4, 1], offset: 0>> to memref<64x16x4xf32, strided<[?, ?, ?], offset: ?>>223 224 // CHECK: memref.cast {{%.*}} : memref<64x16x4xf32, strided<[?, ?, ?], offset: ?>> to memref<64x16x4xf32, strided<[64, 4, 1]>>225 %3 = memref.cast %2 : memref<64x16x4xf32, strided<[?, ?, ?], offset: ?>> to memref<64x16x4xf32, strided<[64, 4, 1], offset: 0>>226 227 // CHECK: memref.cast %{{.*}} : memref<4xf32> to memref<*xf32>228 %4 = memref.cast %1 : memref<4xf32> to memref<*xf32>229 230 // CHECK: memref.cast %{{.*}} : memref<*xf32> to memref<4xf32>231 %5 = memref.cast %4 : memref<*xf32> to memref<4xf32>232 return233}234 235// Check that unranked memrefs with non-default memory space roundtrip236// properly.237// CHECK-LABEL: @unranked_memref_roundtrip(memref<*xf32, 4>)238func.func private @unranked_memref_roundtrip(memref<*xf32, 4>)239 240// CHECK-LABEL: func @load_store_prefetch241func.func @load_store_prefetch(memref<4x4xi32>, index) {242^bb0(%0: memref<4x4xi32>, %1: index):243 // CHECK: %0 = memref.load %arg0[%arg1, %arg1] : memref<4x4xi32>244 %2 = "memref.load"(%0, %1, %1) : (memref<4x4xi32>, index, index)->i32245 246 // CHECK: %{{.*}} = memref.load %arg0[%arg1, %arg1] : memref<4x4xi32>247 %3 = memref.load %0[%1, %1] : memref<4x4xi32>248 249 // CHECK: memref.prefetch %arg0[%arg1, %arg1], write, locality<1>, data : memref<4x4xi32>250 memref.prefetch %0[%1, %1], write, locality<1>, data : memref<4x4xi32>251 252 // CHECK: memref.prefetch %arg0[%arg1, %arg1], read, locality<3>, instr : memref<4x4xi32>253 memref.prefetch %0[%1, %1], read, locality<3>, instr : memref<4x4xi32>254 255 return256}257 258// Test with zero-dimensional operands using no index in load/store.259// CHECK-LABEL: func @zero_dim_no_idx260func.func @zero_dim_no_idx(%arg0 : memref<i32>, %arg1 : memref<i32>, %arg2 : memref<i32>) {261 %0 = memref.load %arg0[] : memref<i32>262 memref.store %0, %arg1[] : memref<i32>263 return264 // CHECK: %0 = memref.load %{{.*}}[] : memref<i32>265 // CHECK: memref.store %{{.*}}, %{{.*}}[] : memref<i32>266}267 268 269// CHECK-LABEL: func @load_store_alignment270func.func @load_store_alignment(%memref: memref<4xi32>) {271 %c0 = arith.constant 0 : index272 // CHECK: memref.load {{.*}} {alignment = 16 : i64}273 %val = memref.load %memref[%c0] { alignment = 16 } : memref<4xi32>274 // CHECK: memref.store {{.*}} {alignment = 16 : i64}275 memref.store %val, %memref[%c0] { alignment = 16 } : memref<4xi32>276 return277}278 279// CHECK-LABEL: func @memref_view(%arg0280func.func @memref_view(%arg0 : index, %arg1 : index, %arg2 : index) {281 %0 = memref.alloc() : memref<2048xi8>282 // Test two dynamic sizes and dynamic offset.283 // CHECK: memref.view {{.*}} : memref<2048xi8> to memref<?x?xf32>284 %1 = memref.view %0[%arg2][%arg0, %arg1] : memref<2048xi8> to memref<?x?xf32>285 286 // Test one dynamic size and dynamic offset.287 // CHECK: memref.view {{.*}} : memref<2048xi8> to memref<4x?xf32>288 %3 = memref.view %0[%arg2][%arg1] : memref<2048xi8> to memref<4x?xf32>289 290 // Test static sizes and static offset.291 // CHECK: memref.view {{.*}} : memref<2048xi8> to memref<64x4xf32>292 %c0 = arith.constant 0: index293 %5 = memref.view %0[%c0][] : memref<2048xi8> to memref<64x4xf32>294 return295}296 297// CHECK-LABEL: func @assume_alignment298// CHECK-SAME: %[[MEMREF:.*]]: memref<4x4xf16>299func.func @assume_alignment(%0: memref<4x4xf16>) {300 // CHECK: %{{.*}} = memref.assume_alignment %[[MEMREF]], 16 : memref<4x4xf16>301 %1 = memref.assume_alignment %0, 16 : memref<4x4xf16>302 return303}304 305// CHECK-LABEL: func @distinct_objects306// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf16>, %[[ARG1:.*]]: memref<?xf32>, %[[ARG2:.*]]: memref<?xf64>)307func.func @distinct_objects(%arg0: memref<?xf16>, %arg1: memref<?xf32>, %arg2: memref<?xf64>) -> (memref<?xf16>, memref<?xf32>, memref<?xf64>) {308 // CHECK: %[[RES:.*]]:3 = memref.distinct_objects %[[ARG0]], %[[ARG1]], %[[ARG2]] : memref<?xf16>, memref<?xf32>, memref<?xf64>309 %1, %2, %3 = memref.distinct_objects %arg0, %arg1, %arg2 : memref<?xf16>, memref<?xf32>, memref<?xf64>310 // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2 : memref<?xf16>, memref<?xf32>, memref<?xf64>311 return %1, %2, %3 : memref<?xf16>, memref<?xf32>, memref<?xf64>312}313 314// CHECK-LABEL: func @expand_collapse_shape_static315func.func @expand_collapse_shape_static(316 %arg0: memref<3x4x5xf32>,317 %arg1: tensor<3x4x5xf32>,318 %arg2: tensor<3x?x5xf32>,319 %arg3: memref<30x20xf32, strided<[4000, 2], offset: 100>>,320 %arg4: memref<1x5xf32, strided<[5, 1], offset: ?>>,321 %arg5: memref<f32>,322 %arg6: memref<3x4x5xf32, strided<[240, 60, 10], offset: 0>>,323 %arg7: memref<1x2049xi64, strided<[?, ?], offset: ?>>,324 %arg8: memref<1x1x1024xi8, strided<[40960, 4096, 1], offset: 0>>,325 %arg9: memref<24x1x1x1024xi8, strided<[40960, 40960, 4096, 1], offset: 0>>) {326 // Reshapes that collapse and expand back a contiguous buffer.327// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]328// CHECK-SAME: memref<3x4x5xf32> into memref<12x5xf32>329 %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :330 memref<3x4x5xf32> into memref<12x5xf32>331 332// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2]] output_shape [3, 4, 5]333// CHECK-SAME: memref<12x5xf32> into memref<3x4x5xf32>334 %r0 = memref.expand_shape %0 [[0, 1], [2]] output_shape [3, 4, 5] :335 memref<12x5xf32> into memref<3x4x5xf32>336 337// CHECK: memref.collapse_shape {{.*}} {{\[}}[0], [1, 2]]338// CHECK-SAME: memref<3x4x5xf32> into memref<3x20xf32>339 %1 = memref.collapse_shape %arg0 [[0], [1, 2]] :340 memref<3x4x5xf32> into memref<3x20xf32>341 342// CHECK: memref.expand_shape {{.*}} {{\[}}[0], [1, 2]] output_shape [3, 4, 5]343// CHECK-SAME: memref<3x20xf32> into memref<3x4x5xf32>344 %r1 = memref.expand_shape %1 [[0], [1, 2]] output_shape [3, 4, 5] :345 memref<3x20xf32> into memref<3x4x5xf32>346 347// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1, 2]]348// CHECK-SAME: memref<3x4x5xf32> into memref<60xf32>349 %2 = memref.collapse_shape %arg0 [[0, 1, 2]] :350 memref<3x4x5xf32> into memref<60xf32>351 352// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1, 2]] output_shape [3, 4, 5]353// CHECK-SAME: memref<60xf32> into memref<3x4x5xf32>354 %r2 = memref.expand_shape %2 [[0, 1, 2]] output_shape [3, 4, 5] :355 memref<60xf32> into memref<3x4x5xf32>356 357// CHECK: memref.expand_shape {{.*}} [] output_shape [1, 1]358// CHECK-SAME: memref<f32> into memref<1x1xf32>359 %r5 = memref.expand_shape %arg5 [] output_shape [1, 1] :360 memref<f32> into memref<1x1xf32>361 362// Reshapes with a custom layout map.363// CHECK: memref.expand_shape {{.*}} {{\[}}[0], [1, 2]] output_shape [30, 4, 5]364 %l0 = memref.expand_shape %arg3 [[0], [1, 2]] output_shape [30, 4, 5] :365 memref<30x20xf32, strided<[4000, 2], offset: 100>>366 into memref<30x4x5xf32, strided<[4000, 10, 2], offset: 100>>367 368// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2]] output_shape [2, 15, 20]369 %l1 = memref.expand_shape %arg3 [[0, 1], [2]] output_shape [2, 15, 20] :370 memref<30x20xf32, strided<[4000, 2], offset: 100>>371 into memref<2x15x20xf32, strided<[60000, 4000, 2], offset: 100>>372 373// CHECK: memref.expand_shape {{.*}} {{\[}}[0], [1, 2]] output_shape [1, 1, 5]374 %r4 = memref.expand_shape %arg4 [[0], [1, 2]] output_shape [1, 1, 5] :375 memref<1x5xf32, strided<[5, 1], offset: ?>> into376 memref<1x1x5xf32, strided<[5, 5, 1], offset: ?>>377 378 // Note: Only the collapsed two shapes are contiguous in the follow test case.379// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]380 %r6 = memref.collapse_shape %arg6 [[0, 1], [2]] :381 memref<3x4x5xf32, strided<[240, 60, 10], offset: 0>> into382 memref<12x5xf32, strided<[60, 10], offset: 0>>383 384// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1]]385 %r7 = memref.collapse_shape %arg7 [[0, 1]] :386 memref<1x2049xi64, strided<[?, ?], offset: ?>> into387 memref<2049xi64, strided<[?], offset: ?>>388 389 // %arg8: memref<1x1x1024xi8, strided<[40960, 4096, 1], offset: 0>>,390 // %arg9: memref<24x1x1x1024xi8, strided<[40960, 40960, 4096, 1], offset: 0>>) {391 392// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1, 2]]393 %r8 = memref.collapse_shape %arg8 [[0, 1, 2]] :394 memref<1x1x1024xi8, strided<[40960, 4096, 1], offset: 0>> into395 memref<1024xi8, strided<[1], offset: 0>>396 397// CHECK: memref.collapse_shape {{.*}} {{\[}}[0], [1, 2, 3]]398 %r9 = memref.collapse_shape %arg9 [[0], [1, 2, 3]] :399 memref<24x1x1x1024xi8, strided<[40960, 40960, 4096, 1], offset: 0>> into400 memref<24x1024xi8, strided<[40960, 1], offset: 0>>401 402 // Reshapes that expand and collapse back a contiguous buffer with some 1's.403// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]] output_shape [1, 3, 4, 1, 5]404// CHECK-SAME: memref<3x4x5xf32> into memref<1x3x4x1x5xf32>405 %3 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]] output_shape [1, 3, 4, 1, 5]:406 memref<3x4x5xf32> into memref<1x3x4x1x5xf32>407 408// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]409// CHECK-SAME: memref<1x3x4x1x5xf32> into memref<3x4x5xf32>410 %r3 = memref.collapse_shape %3 [[0, 1], [2], [3, 4]] :411 memref<1x3x4x1x5xf32> into memref<3x4x5xf32>412 413 // Reshapes on tensors.414// CHECK: tensor.expand_shape {{.*}}: tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>415 %t0 = tensor.expand_shape %arg1 [[0, 1], [2], [3, 4]] output_shape [1, 3, 4, 1, 5] :416 tensor<3x4x5xf32> into tensor<1x3x4x1x5xf32>417 418// CHECK: tensor.collapse_shape {{.*}}: tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>419 %rt0 = tensor.collapse_shape %t0 [[0, 1], [2], [3, 4]] :420 tensor<1x3x4x1x5xf32> into tensor<3x4x5xf32>421 422// CHECK: tensor.dim %arg2, {{.*}} : tensor<3x?x5xf32>423// CHECK: tensor.expand_shape {{.*}}: tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>424 %c1 = arith.constant 1 : index425 %sz1 = tensor.dim %arg2, %c1 : tensor<3x?x5xf32>426 %t1 = tensor.expand_shape %arg2 [[0, 1], [2], [3, 4]] output_shape [1, 3, %sz1, 1, 5] :427 tensor<3x?x5xf32> into tensor<1x3x?x1x5xf32>428 429// CHECK: tensor.collapse_shape {{.*}}: tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>430 %rt1 = tensor.collapse_shape %t1 [[0], [1, 2], [3, 4]] :431 tensor<1x3x?x1x5xf32> into tensor<1x?x5xf32>432 return433}434 435// CHECK-LABEL: func @expand_collapse_shape_dynamic436func.func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,437 %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: 0>>,438 %arg2: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>,439 %arg3: memref<?x42xf32, strided<[42, 1], offset: 0>>,440 %arg4: index,441 %arg5: index,442 %arg6: index,443 %arg7: memref<4x?x4xf32>) {444// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]445// CHECK-SAME: memref<?x?x?xf32> into memref<?x?xf32>446 %0 = memref.collapse_shape %arg0 [[0, 1], [2]] :447 memref<?x?x?xf32> into memref<?x?xf32>448 449// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2]] output_shape [%arg4, 4, %arg5]450// CHECK-SAME: memref<?x?xf32> into memref<?x4x?xf32>451 %r0 = memref.expand_shape %0 [[0, 1], [2]] output_shape [%arg4, 4, %arg5] :452 memref<?x?xf32> into memref<?x4x?xf32>453 454// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]455// CHECK-SAME: memref<?x?x?xf32, strided<[?, ?, 1]>> into memref<?x?xf32, strided<[?, 1]>>456 %1 = memref.collapse_shape %arg1 [[0, 1], [2]] :457 memref<?x?x?xf32, strided<[?, ?, 1], offset: 0>> into458 memref<?x?xf32, strided<[?, 1], offset: 0>>459 460// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2]] output_shape [%arg4, 4, %arg5]461// CHECK-SAME: memref<?x?xf32, strided<[?, 1]>> into memref<?x4x?xf32, strided<[?, ?, 1]>>462 %r1 = memref.expand_shape %1 [[0, 1], [2]] output_shape [%arg4, 4, %arg5] :463 memref<?x?xf32, strided<[?, 1], offset: 0>> into464 memref<?x4x?xf32, strided<[?, ?, 1], offset: 0>>465 466// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1], [2]]467// CHECK-SAME: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> into memref<?x?xf32, strided<[?, 1], offset: ?>>468 %2 = memref.collapse_shape %arg2 [[0, 1], [2]] :469 memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> into470 memref<?x?xf32, strided<[?, 1], offset: ?>>471 472// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2]] output_shape [%arg4, 4, %arg5]473// CHECK-SAME: memref<?x?xf32, strided<[?, 1], offset: ?>> into memref<?x4x?xf32, strided<[?, ?, 1], offset: ?>>474 %r2 = memref.expand_shape %2 [[0, 1], [2]] output_shape [%arg4, 4, %arg5] :475 memref<?x?xf32, strided<[?, 1], offset: ?>> into476 memref<?x4x?xf32, strided<[?, ?, 1], offset: ?>>477 478// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1]]479// CHECK-SAME: memref<?x42xf32, strided<[42, 1]>> into memref<?xf32, strided<[1]>>480 %3 = memref.collapse_shape %arg3 [[0, 1]] :481 memref<?x42xf32, strided<[42, 1], offset: 0>> into482 memref<?xf32, strided<[1]>>483 484// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1]] output_shape [%arg6, 42]485// CHECK-SAME: memref<?xf32, strided<[1]>> into memref<?x42xf32>486 %r3 = memref.expand_shape %3 [[0, 1]] output_shape [%arg6, 42] :487 memref<?xf32, strided<[1]>> into memref<?x42xf32>488 489// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2], [3, 4]]490 %4 = memref.expand_shape %arg7 [[0, 1], [2], [3, 4]] output_shape [2, 2, %arg4, 2, 2]491 : memref<4x?x4xf32> into memref<2x2x?x2x2xf32>492 return493}494 495func.func @expand_collapse_shape_zero_dim(%arg0 : memref<1x1xf32>, %arg1 : memref<f32>)496 -> (memref<f32>, memref<1x1xf32>) {497 %0 = memref.collapse_shape %arg0 [] : memref<1x1xf32> into memref<f32>498 %1 = memref.expand_shape %0 [] output_shape [1, 1] : memref<f32> into memref<1x1xf32>499 return %0, %1 : memref<f32>, memref<1x1xf32>500}501// CHECK-LABEL: func @expand_collapse_shape_zero_dim502// CHECK: memref.collapse_shape %{{.*}} [] : memref<1x1xf32> into memref<f32>503// CHECK: memref.expand_shape %{{.*}} [] output_shape [1, 1] : memref<f32> into memref<1x1xf32>504 505func.func @collapse_shape_to_dynamic506 (%arg0: memref<?x?x?x4x?xf32>) -> memref<?x?x?xf32> {507 %0 = memref.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :508 memref<?x?x?x4x?xf32> into memref<?x?x?xf32>509 return %0 : memref<?x?x?xf32>510}511// CHECK: func @collapse_shape_to_dynamic512// CHECK: memref.collapse_shape513// CHECK-SAME: [0], [1], [2, 3, 4]514 515// -----516 517// CHECK-LABEL: func @expand_collapse_shape_transposed_layout518func.func @expand_collapse_shape_transposed_layout(519 %m0: memref<?x?xf32, strided<[1, 10], offset: 0>>,520 %m1: memref<4x5x6xf32, strided<[1, ?, 1000], offset: 0>>,521 %sz0: index,522 %sz1: index) {523 524 %r0 = memref.expand_shape %m0 [[0], [1, 2]] output_shape [%sz0, %sz1, 5] :525 memref<?x?xf32, strided<[1, 10], offset: 0>> into526 memref<?x?x5xf32, strided<[1, 50, 10], offset: 0>>527 %rr0 = memref.collapse_shape %r0 [[0], [1, 2]] :528 memref<?x?x5xf32, strided<[1, 50, 10], offset: 0>> into529 memref<?x?xf32, strided<[1, 10], offset: 0>>530 531 %r1 = memref.expand_shape %m1 [[0, 1], [2], [3, 4]] output_shape [2, 2, 5, 2, 3] :532 memref<4x5x6xf32, strided<[1, ?, 1000], offset: 0>> into533 memref<2x2x5x2x3xf32, strided<[2, 1, ?, 3000, 1000], offset: 0>>534 %rr1 = memref.collapse_shape %r1 [[0, 1], [2], [3, 4]] :535 memref<2x2x5x2x3xf32, strided<[2, 1, ?, 3000, 1000], offset: 0>> into536 memref<4x5x6xf32, strided<[1, ?, 1000], offset: 0>>537 return538}539 540// -----541 542func.func @rank(%t : memref<4x4x?xf32>) {543 // CHECK: %{{.*}} = memref.rank %{{.*}} : memref<4x4x?xf32>544 %0 = "memref.rank"(%t) : (memref<4x4x?xf32>) -> index545 546 // CHECK: %{{.*}} = memref.rank %{{.*}} : memref<4x4x?xf32>547 %1 = memref.rank %t : memref<4x4x?xf32>548 return549}550 551// ------552 553// CHECK-LABEL: func @atomic_rmw554// CHECK-SAME: ([[BUF:%.*]]: memref<10xf32>, [[VAL:%.*]]: f32, [[I:%.*]]: index)555func.func @atomic_rmw(%I: memref<10xf32>, %val: f32, %i : index) {556 %x = memref.atomic_rmw addf %val, %I[%i] : (f32, memref<10xf32>) -> f32557 // CHECK: memref.atomic_rmw addf [[VAL]], [[BUF]]{{\[}}[[I]]]558 return559}560 561// CHECK-LABEL: func @generic_atomic_rmw562// CHECK-SAME: ([[BUF:%.*]]: memref<1x2xf32>, [[I:%.*]]: index, [[J:%.*]]: index)563func.func @generic_atomic_rmw(%I: memref<1x2xf32>, %i : index, %j : index) {564 %x = memref.generic_atomic_rmw %I[%i, %j] : memref<1x2xf32> {565 // CHECK-NEXT: memref.generic_atomic_rmw [[BUF]]{{\[}}[[I]], [[J]]] : memref566 ^bb0(%old_value : f32):567 %c1 = arith.constant 1.0 : f32568 %out = arith.addf %c1, %old_value : f32569 memref.atomic_yield %out : f32570 // CHECK: index_attr = 8 : index571 } { index_attr = 8 : index }572 return573}574 575// -----576 577func.func @extract_strided_metadata(%memref : memref<10x?xf32>)578 -> memref<?x?xf32, strided<[?, ?], offset: ?>> {579 580 %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %memref581 : memref<10x?xf32> -> memref<f32>, index, index, index, index, index582 583 %m2 = memref.reinterpret_cast %base to584 offset: [%offset],585 sizes: [%sizes#0, %sizes#1],586 strides: [%strides#0, %strides#1]587 : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>588 589 return %m2: memref<?x?xf32, strided<[?, ?], offset: ?>>590}591 592// -----593 594// CHECK-LABEL: func @memref_realloc_ss595func.func @memref_realloc_ss(%src : memref<2xf32>) -> memref<4xf32>{596 %0 = memref.realloc %src : memref<2xf32> to memref<4xf32>597 return %0 : memref<4xf32>598}599 600// CHECK-LABEL: func @memref_realloc_sd601func.func @memref_realloc_sd(%src : memref<2xf32>, %d : index) -> memref<?xf32>{602 %0 = memref.realloc %src(%d) : memref<2xf32> to memref<?xf32>603 return %0 : memref<?xf32>604}605 606// CHECK-LABEL: func @memref_realloc_ds607func.func @memref_realloc_ds(%src : memref<?xf32>) -> memref<4xf32>{608 %0 = memref.realloc %src: memref<?xf32> to memref<4xf32>609 return %0 : memref<4xf32>610}611 612// CHECK-LABEL: func @memref_realloc_dd613func.func @memref_realloc_dd(%src : memref<?xf32>, %d: index)614 -> memref<?xf32>{615 %0 = memref.realloc %src(%d) : memref<?xf32> to memref<?xf32>616 return %0 : memref<?xf32>617}618 619// CHECK-LABEL: func @memref_extract_aligned_pointer620func.func @memref_extract_aligned_pointer(%src : memref<?xf32>) -> index {621 %0 = memref.extract_aligned_pointer_as_index %src : memref<?xf32> -> index622 return %0 : index623}624 625// CHECK-LABEL: func @memref_memory_space_cast626func.func @memref_memory_space_cast(%src : memref<?xf32>) -> memref<?xf32, 1> {627 %dst = memref.memory_space_cast %src : memref<?xf32> to memref<?xf32, 1>628 return %dst : memref<?xf32, 1>629}630 631// CHECK-LABEL: func @memref_transpose_map632func.func @memref_transpose_map(%src : memref<?x?xf32>) -> memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d1 * s0 + d0)>> {633 %dst = memref.transpose %src (i, j) -> (j, i) : memref<?x?xf32> to memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d1 * s0 + d0)>>634 return %dst : memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d1 * s0 + d0)>>635}636