1// RUN: mlir-opt %s -test-vector-transfer-full-partial-split | FileCheck %s
2
3// CHECK-DAG: #[[$map_p4:.*]] = affine_map<()[s0] -> (s0 + 4)>
4// CHECK-DAG: #[[$map_p8:.*]] = affine_map<()[s0] -> (s0 + 8)>
5// CHECK-DAG: #[[$map_2d_stride_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
6
7// CHECK-LABEL: split_vector_transfer_read_2d(
8//  CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref
9//  CHECK-SAME: %[[i:[a-zA-Z0-9]*]]: index
10//  CHECK-SAME: %[[j:[a-zA-Z0-9]*]]: index
11func @split_vector_transfer_read_2d(%A: memref<?x8xf32>, %i: index, %j: index) -> vector<4x8xf32> {
12  %c0 = constant 0 : index
13  %f0 = constant 0.0 : f32
14
15  //  CHECK-DAG: %[[c0:.*]] = constant 0 : index
16  //  CHECK-DAG: %[[c8:.*]] = constant 8 : index
17  //  CHECK-DAG: %[[cst:.*]] = constant 0.000000e+00 : f32
18  // alloca for boundary full tile
19  //      CHECK: %[[alloc:.*]] = alloca() {alignment = 32 : i64} : memref<4x8xf32>
20  // %i + 4 <= dim(%A, 0)
21  //      CHECK: %[[idx0:.*]] = affine.apply #[[$map_p4]]()[%[[i]]]
22  //      CHECK: %[[d0:.*]] = dim %[[A]], %[[c0]] : memref<?x8xf32>
23  //      CHECK: %[[cmp0:.*]] = cmpi "sle", %[[idx0]], %[[d0]] : index
24  // %j + 8 <= dim(%A, 1)
25  //      CHECK: %[[idx1:.*]] = affine.apply #[[$map_p8]]()[%[[j]]]
26  //      CHECK: %[[cmp1:.*]] = cmpi "sle", %[[idx1]], %[[c8]] : index
27  // are both conds true
28  //      CHECK: %[[cond:.*]] = and %[[cmp0]], %[[cmp1]] : i1
29  //      CHECK: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32>, index, index) {
30  //               inBounds, just yield %A
31  //      CHECK:   scf.yield %[[A]], %[[i]], %[[j]] : memref<?x8xf32>, index, index
32  //      CHECK: } else {
33  //               slow path, fill tmp alloc and yield a memref_casted version of it
34  //      CHECK:   %[[slow:.*]] = vector.transfer_read %[[A]][%[[i]], %[[j]]], %cst :
35  // CHECK-SAME:     memref<?x8xf32>, vector<4x8xf32>
36  //      CHECK:   %[[cast_alloc:.*]] = vector.type_cast %[[alloc]] :
37  // CHECK-SAME:     memref<4x8xf32> to memref<vector<4x8xf32>>
38  //      CHECK:   store %[[slow]], %[[cast_alloc]][] : memref<vector<4x8xf32>>
39  //      CHECK:   %[[yielded:.*]] = memref_cast %[[alloc]] :
40  // CHECK-SAME:     memref<4x8xf32> to memref<?x8xf32>
41  //      CHECK:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
42  // CHECK-SAME:     memref<?x8xf32>, index, index
43  //      CHECK: }
44  //      CHECK: %[[res:.*]] = vector.transfer_read %[[ifres]]#0[%[[ifres]]#1, %[[ifres]]#2], %[[cst]]
45  // CHECK_SAME:   {masked = [false, false]} : memref<?x8xf32>, vector<4x8xf32>
46  %1 = vector.transfer_read %A[%i, %j], %f0 : memref<?x8xf32>, vector<4x8xf32>
47
48  // CHECK: return %[[res]] : vector<4x8xf32>
49  return %1: vector<4x8xf32>
50}
51
52// CHECK-LABEL: split_vector_transfer_read_strided_2d(
53//  CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref
54//  CHECK-SAME: %[[i:[a-zA-Z0-9]*]]: index
55//  CHECK-SAME: %[[j:[a-zA-Z0-9]*]]: index
56func @split_vector_transfer_read_strided_2d(
57    %A: memref<7x8xf32, offset:?, strides:[?, 1]>,
58    %i: index, %j: index) -> vector<4x8xf32> {
59  %c0 = constant 0 : index
60  %f0 = constant 0.0 : f32
61
62  //  CHECK-DAG: %[[c0:.*]] = constant 0 : index
63  //  CHECK-DAG: %[[c7:.*]] = constant 7 : index
64  //  CHECK-DAG: %[[c8:.*]] = constant 8 : index
65  //  CHECK-DAG: %[[cst:.*]] = constant 0.000000e+00 : f32
66  // alloca for boundary full tile
67  //      CHECK: %[[alloc:.*]] = alloca() {alignment = 32 : i64} : memref<4x8xf32>
68  // %i + 4 <= dim(%A, 0)
69  //      CHECK: %[[idx0:.*]] = affine.apply #[[$map_p4]]()[%[[i]]]
70  //      CHECK: %[[cmp0:.*]] = cmpi "sle", %[[idx0]], %[[c7]] : index
71  // %j + 8 <= dim(%A, 1)
72  //      CHECK: %[[idx1:.*]] = affine.apply #[[$map_p8]]()[%[[j]]]
73  //      CHECK: %[[cmp1:.*]] = cmpi "sle", %[[idx1]], %[[c8]] : index
74  // are both conds true
75  //      CHECK: %[[cond:.*]] = and %[[cmp0]], %[[cmp1]] : i1
76  //      CHECK: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index) {
77  //               inBounds but not cast-compatible: yield a memref_casted form of %A
78  //      CHECK:   %[[casted:.*]] = memref_cast %arg0 :
79  // CHECK-SAME:     memref<7x8xf32, #[[$map_2d_stride_1]]> to memref<?x8xf32, #[[$map_2d_stride_1]]>
80  //      CHECK:   scf.yield %[[casted]], %[[i]], %[[j]] :
81  // CHECK-SAME:     memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index
82  //      CHECK: } else {
83  //               slow path, fill tmp alloc and yield a memref_casted version of it
84  //      CHECK:   %[[slow:.*]] = vector.transfer_read %[[A]][%[[i]], %[[j]]], %cst :
85  // CHECK-SAME:     memref<7x8xf32, #[[$map_2d_stride_1]]>, vector<4x8xf32>
86  //      CHECK:   %[[cast_alloc:.*]] = vector.type_cast %[[alloc]] :
87  // CHECK-SAME:     memref<4x8xf32> to memref<vector<4x8xf32>>
88  //      CHECK:   store %[[slow]], %[[cast_alloc]][] :
89  // CHECK-SAME:     memref<vector<4x8xf32>>
90  //      CHECK:   %[[yielded:.*]] = memref_cast %[[alloc]] :
91  // CHECK-SAME:     memref<4x8xf32> to memref<?x8xf32, #[[$map_2d_stride_1]]>
92  //      CHECK:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
93  // CHECK-SAME:     memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index
94  //      CHECK: }
95  //      CHECK: %[[res:.*]] = vector.transfer_read {{.*}} {masked = [false, false]} :
96  // CHECK-SAME:   memref<?x8xf32, #[[$map_2d_stride_1]]>, vector<4x8xf32>
97  %1 = vector.transfer_read %A[%i, %j], %f0 :
98    memref<7x8xf32, offset:?, strides:[?, 1]>, vector<4x8xf32>
99
100  // CHECK: return %[[res]] : vector<4x8xf32>
101  return %1 : vector<4x8xf32>
102}
103