1// RUN: mlir-opt %s -test-vector-transfer-flatten-patterns -split-input-file | FileCheck %s
2
3func.func @transfer_read_flattenable_with_offset(
4      %arg : memref<5x4x3x2xi8, offset:?, strides:[24, 6, 2, 1]>) -> vector<5x4x3x2xi8> {
5    %c0 = arith.constant 0 : index
6    %cst = arith.constant 0 : i8
7    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst :
8      memref<5x4x3x2xi8, offset:?, strides:[24, 6, 2, 1]>, vector<5x4x3x2xi8>
9    return %v : vector<5x4x3x2xi8>
10}
11
12// CHECK-LABEL: func @transfer_read_flattenable_with_offset
13// CHECK-SAME:      %[[ARG:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8
14// CHECK:         %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]
15// C-HECK:         %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]
16// C-HECK:         %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>
17// C-HECK:         return %[[VEC2D]]
18
19// -----
20
21func.func @transfer_write_flattenable_with_offset(
22      %arg : memref<5x4x3x2xi8, offset:?, strides:[24, 6, 2, 1]>, %vec : vector<5x4x3x2xi8>) {
23    %c0 = arith.constant 0 : index
24    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] :
25      vector<5x4x3x2xi8>, memref<5x4x3x2xi8, offset:?, strides:[24, 6, 2, 1]>
26    return
27}
28
29// C-HECK-LABEL: func @transfer_write_flattenable_with_offset
30// C-HECK-SAME:      %[[ARG:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8
31// C-HECK-SAME:      %[[VEC:[0-9a-zA-Z]+]]: vector<5x4x3x2xi8>
32// C-HECK-DAG:     %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8, {{.+}}> into memref<120xi8, {{.+}}>
33// C-HECK-DAG:     %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<5x4x3x2xi8> to vector<120xi8>
34// C-HECK:         vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
35
36// -----
37
38func.func @transfer_write_0d(%arg : memref<i8>, %vec : vector<i8>) {
39      vector.transfer_write %vec, %arg[] : vector<i8>, memref<i8>
40      return
41}
42
43// CHECK-LABEL: func @transfer_write_0d
44// CHECK-SAME:       %[[ARG:.+]]: memref<i8>
45// CHECK-SAME:       %[[VEC:.+]]: vector<i8>
46// CHECK:          vector.transfer_write %[[VEC]], %[[ARG]][] : vector<i8>, memref<i8>
47// CHECK:          return
48
49// -----
50
51func.func @transfer_read_0d(%arg : memref<i8>) -> vector<i8> {
52      %cst = arith.constant 0 : i8
53      %0 = vector.transfer_read %arg[], %cst : memref<i8>, vector<i8>
54      return %0 : vector<i8>
55}
56
57// CHECK-LABEL: func @transfer_read_0d
58// CHECK-SAME:       %[[ARG:.+]]: memref<i8>
59// CHECK:            %[[CST:.+]] = arith.constant 0 : i8
60// CHECK:            %[[READ:.+]] = vector.transfer_read %[[ARG]][], %[[CST]] : memref<i8>
61// CHECK:            return %[[READ]]
62
63// -----
64
65#map0 = affine_map<(d0, d1, d2, d3)[s0, s1] -> (d0 * s1 + s0 + d1 * 32 + d2 * 4 + d3)>
66
67func.func @transfer_read_flattenable_with_dynamic_dims_and_indices(%arg0 : memref<?x?x8x4xi8, #map0>, %arg1 : index, %arg2 : index) -> vector<8x4xi8> {
68    %c0_i8 = arith.constant 0 : i8
69    %c0 = arith.constant 0 : index
70    %result = vector.transfer_read %arg0[%arg1, %arg2, %c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<?x?x8x4xi8, #map0>, vector<8x4xi8>
71    return %result : vector<8x4xi8>
72}
73
74// CHECK-LABEL: func @transfer_read_flattenable_with_dynamic_dims_and_indices
75// CHECK-SAME:    %[[ARG0:.+]]: memref<?x?x8x4xi8, {{.+}}>, %[[ARG1:.+]]: index, %[[ARG2:.+]]: index
76// CHECK:       %[[C0_I8:.+]] = arith.constant 0 : i8
77// CHECK:       %[[C0:.+]] = arith.constant 0 : index
78// CHECK:       %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG0]] {{\[}}[0], [1], [2, 3]{{\]}}
79// CHECK-SAME:    : memref<?x?x8x4xi8, {{.+}}> into memref<?x?x32xi8, {{.+}}>
80// CHECK:       %[[VEC1D:.+]] = vector.transfer_read %[[COLLAPSED]]
81// CHECK-SAME:    [%[[ARG1]], %[[ARG2]], %[[C0]]], %[[C0_I8]]
82// CHECK-SAME:    {in_bounds = [true]}
83// CHECK-SAME:    : memref<?x?x32xi8, {{.+}}>, vector<32xi8>
84// CHECK:       %[[VEC2D:.+]] = vector.shape_cast %[[VEC1D]] : vector<32xi8> to vector<8x4xi8>
85// CHECK:       return %[[VEC2D]] : vector<8x4xi8>
86
87// -----
88
89#map0 = affine_map<(d0, d1, d2, d3)[s0, s1] -> (d0 * s1 + s0 + d1 * 32 + d2 * 4 + d3)>
90
91func.func @transfer_write_flattenable_with_dynamic_dims_and_indices(%vec : vector<8x4xi8>, %dst : memref<?x?x8x4xi8, #map0>, %arg1 : index, %arg2 : index) {
92    %c0 = arith.constant 0 : index
93    vector.transfer_write %vec, %dst[%arg1, %arg2, %c0, %c0] {in_bounds = [true, true]} : vector<8x4xi8>, memref<?x?x8x4xi8, #map0>
94    return
95}
96
97// CHECK-LABEL: func @transfer_write_flattenable_with_dynamic_dims_and_indices
98// CHECK-SAME:    %[[ARG0:.+]]: vector<8x4xi8>, %[[ARG1:.+]]: memref<?x?x8x4xi8, {{.+}}>, %[[ARG2:.+]]: index, %[[ARG3:.+]]: index
99// CHECK:       %[[C0:.+]] = arith.constant 0 : index
100// CHECK:       %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}}
101// CHECK-SAME:    : memref<?x?x8x4xi8, {{.+}}> into memref<?x?x32xi8, {{.+}}>
102// CHECK:       %[[VEC1D:.+]] = vector.shape_cast %[[ARG0]] : vector<8x4xi8> to vector<32xi8>
103// CHECK:       vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
104// CHECK-SAME:    [%[[ARG2]], %[[ARG3]], %[[C0]]]
105// CHECK-SAME:    {in_bounds = [true]}
106// CHECK-SAME:    : vector<32xi8>, memref<?x?x32xi8, {{.+}}>
107