1// RUN: mlir-opt -allow-unregistered-dialect %s -test-compose-subview -split-input-file | FileCheck %s
2
3// CHECK: [[MAP:#.*]] = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3456)
4#map0 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 2304)>
5#map1 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3456)>
6
7func.func @main(%input: memref<4x1024xf32>) -> memref<1x128xf32, #map1> {
8  //      CHECK: subview %arg0[3, 384] [1, 128] [1, 1]
9  // CHECK-SAME: memref<4x1024xf32> to memref<1x128xf32, [[MAP]]>
10  %0 = memref.subview %input[2, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, #map0>
11  %1 = memref.subview %0[1, 128] [1, 128] [1, 1] : memref<2x256xf32, #map0> to memref<1x128xf32, #map1>
12  return %1 : memref<1x128xf32, #map1>
13}
14
15// -----
16
17// CHECK: [[MAP:#.*]] = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3745)
18#map0 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 1536)>
19#map1 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 2688)>
20#map2 = affine_map<(d0, d1) -> (d0 * 1024 + d1 + 3745)>
21
22func.func @main(%input: memref<4x1024xf32>) -> memref<1x10xf32, #map2> {
23  //      CHECK: subview %arg0[3, 673] [1, 10] [1, 1]
24  // CHECK-SAME: memref<4x1024xf32> to memref<1x10xf32, [[MAP]]>
25  %0 = memref.subview %input[1, 512] [3, 256] [1, 1] : memref<4x1024xf32> to memref<3x256xf32, #map0>
26  %1 = memref.subview %0[1, 128] [2, 128] [1, 1] : memref<3x256xf32, #map0> to memref<2x128xf32, #map1>
27  %2 = memref.subview %1[1, 33] [1, 10] [1, 1] : memref<2x128xf32, #map1> to memref<1x10xf32, #map2>
28  return %2 : memref<1x10xf32, #map2>
29}
30
31// -----
32
33// CHECK: [[MAP:#.*]] = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)
34#map = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>
35
36func.func @main(%input: memref<4x1024xf32>) -> memref<1x128xf32, #map> {
37  // CHECK: [[CST_3:%.*]] = arith.constant 3 : index
38  %cst_1 = arith.constant 1 : index
39  %cst_2 = arith.constant 2 : index
40  //      CHECK: subview %arg0{{\[}}[[CST_3]], 384] [1, 128] [1, 1]
41  // CHECK-SAME: memref<4x1024xf32> to memref<1x128xf32, [[MAP]]>
42  %0 = memref.subview %input[%cst_2, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, #map>
43  %1 = memref.subview %0[%cst_1, 128] [1, 128] [1, 1] : memref<2x256xf32, #map> to memref<1x128xf32, #map>
44  return %1 : memref<1x128xf32, #map>
45}
46
47// -----
48
49// CHECK: [[MAP:#.*]] = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)
50#map = affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>
51
52func.func @main(%input: memref<4x1024xf32>) -> memref<1x128xf32, #map> {
53  // CHECK: [[CST_3:%.*]] = arith.constant 3 : index
54  %cst_2 = arith.constant 2 : index
55  // CHECK: [[CST_384:%.*]] = arith.constant 384 : index
56  %cst_128 = arith.constant 128 : index
57  //      CHECK: subview %arg0{{\[}}[[CST_3]], [[CST_384]]] [1, 128] [1, 1]
58  // CHECK-SAME: memref<4x1024xf32> to memref<1x128xf32, [[MAP]]>
59  %0 = memref.subview %input[%cst_2, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, #map>
60  %1 = memref.subview %0[1, %cst_128] [1, 128] [1, 1] : memref<2x256xf32, #map> to memref<1x128xf32, #map>
61  return %1 : memref<1x128xf32, #map>
62}
63